@illuma-ai/agents 1.0.90 → 1.0.94

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

@@ -0,0 +1,1100 @@
+import { describe, it, expect, beforeEach, jest } from '@jest/globals';
+import type { ToolCall, ToolCallChunk } from '@langchain/core/messages/tool';
+import type { StandardGraph } from '@/graphs';
+import type { AgentContext } from '@/agents/AgentContext';
+import type * as t from '@/types';
+import { StepTypes, ToolCallTypes, Providers, GraphEvents } from '@/common';
+import {
+  handleToolCallChunks,
+  handleToolCalls,
+  handleServerToolResult,
+} from '../handlers';
+
+type MockGraph = {
+  getStepKey: jest.Mock;
+  getStepIdByKey: jest.Mock;
+  getRunStep: jest.Mock;
+  dispatchRunStep: jest.Mock;
+  dispatchRunStepDelta: jest.Mock;
+  toolCallStepIds: Map<string, string>;
+  messageStepHasToolCalls: Map<string, boolean>;
+  messageIdsByStepKey: Map<string, string>;
+  prelimMessageIdsByStepKey: Map<string, string>;
+  invokedToolIds?: Set<string>;
+  handlerRegistry?: {
+    getHandler: jest.Mock;
+  };
+  streamingToolCallBuffer: {
+    append: jest.Mock;
+    setToolName: jest.Mock;
+    getToolName: jest.Mock;
+    getRawArgs: jest.Mock;
+    extractFieldValue: jest.Mock;
+    has: jest.Mock;
+    clear: jest.Mock;
+    clearAll: jest.Mock;
+    setIndexMapping: jest.Mock;
+    getIdByIndex: jest.Mock;
+  };
+};
+
+function createMockGraph(overrides?: Partial<MockGraph>): MockGraph {
+  let stepCounter = 0;
+  return {
+    getStepKey: jest.fn<() => string>().mockReturnValue('step-key'),
+    getStepIdByKey: jest.fn<() => string>().mockReturnValue('prev-step-id'),
+    getRunStep: jest
+      .fn<() => t.RunStep | undefined>()
+      .mockReturnValue(undefined),
+    dispatchRunStep: jest
+      .fn<() => Promise<string>>()
+      .mockImplementation(async () => `new-step-${++stepCounter}`),
+    dispatchRunStepDelta: jest
+      .fn<() => Promise<void>>()
+      .mockResolvedValue(undefined),
+    toolCallStepIds: new Map(),
+    messageStepHasToolCalls: new Map(),
+    messageIdsByStepKey: new Map(),
+    prelimMessageIdsByStepKey: new Map(),
+    invokedToolIds: undefined,
+    handlerRegistry: undefined,
+    streamingToolCallBuffer: {
+      append: jest.fn(),
+      setToolName: jest.fn(),
+      getToolName: jest.fn(),
+      getRawArgs: jest.fn(),
+      extractFieldValue: jest.fn(),
+      has: jest.fn().mockReturnValue(false),
+      clear: jest.fn(),
+      clearAll: jest.fn(),
+      setIndexMapping: jest.fn(),
+      getIdByIndex: jest.fn().mockReturnValue(undefined),
+    },
+    ...overrides,
+  };
+}
+
+function makeRunStep(
+  type: StepTypes,
+  opts?: { tool_calls?: t.AgentToolCall[]; id?: string; index?: number }
+): t.RunStep {
+  const stepDetails: t.StepDetails =
+    type === StepTypes.MESSAGE_CREATION
+      ? {
+          type: StepTypes.MESSAGE_CREATION,
+          message_creation: { message_id: 'msg-1' },
+        }
+      : { type: StepTypes.TOOL_CALLS, tool_calls: opts?.tool_calls ?? [] };
+  return {
+    type,
+    id: opts?.id ?? 'run-step-1',
+    index: opts?.index ?? 0,
+    stepDetails,
+    usage: null,
+  };
+}
+
+function makeToolCall(id: string, name = 'calculator'): ToolCall {
+  return { id, name, args: {}, type: 'tool_call' };
+}
+
+function makeToolCallChunk(opts?: {
+  id?: string;
+  name?: string;
+  index?: number;
+}): ToolCallChunk {
+  return {
+    id: opts?.id,
+    name: opts?.name,
+    args: '',
+    index: opts?.index ?? 0,
+    type: 'tool_call_chunk',
+  };
+}
+
+const defaultMetadata = { run_id: 'test-run' };
+
+describe('handleToolCallChunks', () => {
+  let graph: MockGraph;
+
+  beforeEach(() => {
+    graph = createMockGraph();
+  });
+
+  it('creates TOOL_CALLS step when previous step is MESSAGE_CREATION', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    const chunks = [makeToolCallChunk({ index: 2 })];
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: chunks,
+      metadata: defaultMetadata,
+    });
+
+    const dispatchCalls = graph.dispatchRunStep.mock.calls;
+    expect(dispatchCalls).toHaveLength(1);
+    expect(dispatchCalls[0][1]).toEqual(
+      expect.objectContaining({ type: StepTypes.TOOL_CALLS })
+    );
+    expect(graph.messageStepHasToolCalls.has('prev-step-id')).toBe(true);
+    expect(graph.dispatchRunStepDelta).toHaveBeenCalledTimes(1);
+  });
+
+  it('reuses existing TOOL_CALLS step without dispatching a new one', async () => {
+    const toolStep = makeRunStep(StepTypes.TOOL_CALLS);
+    graph.getRunStep.mockReturnValue(toolStep);
+
+    const chunks = [makeToolCallChunk({ index: 2 })];
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: chunks,
+      metadata: defaultMetadata,
+    });
+
+    expect(graph.dispatchRunStep).not.toHaveBeenCalled();
+    expect(graph.dispatchRunStepDelta).toHaveBeenCalledTimes(1);
+    expect(graph.dispatchRunStepDelta).toHaveBeenCalledWith(
+      'prev-step-id',
+      expect.objectContaining({ type: StepTypes.TOOL_CALLS })
+    );
+  });
+
+  it('creates MESSAGE_CREATION when no previous step exists', async () => {
+    let callCount = 0;
+    graph.getStepIdByKey.mockImplementation(() => {
+      callCount++;
+      if (callCount === 1) {
+        throw new Error('No step found');
+      }
+      return 'new-step-1';
+    });
+    graph.getRunStep.mockReturnValue(makeRunStep(StepTypes.MESSAGE_CREATION));
+
+    const chunks = [makeToolCallChunk({ index: 0 })];
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: chunks,
+      metadata: defaultMetadata,
+    });
+
+    const dispatchCalls = graph.dispatchRunStep.mock.calls;
+    expect(dispatchCalls.length).toBeGreaterThanOrEqual(2);
+    expect(dispatchCalls[0][1]).toEqual(
+      expect.objectContaining({ type: StepTypes.MESSAGE_CREATION })
+    );
+    expect(dispatchCalls[1][1]).toEqual(
+      expect.objectContaining({ type: StepTypes.TOOL_CALLS })
+    );
+  });
+
+  it('skips TOOL_CALLS dispatch when already dispatched for this MESSAGE_CREATION', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+    graph.messageStepHasToolCalls.set('prev-step-id', true);
+
+    const chunks = [makeToolCallChunk({ index: 2 })];
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: chunks,
+      metadata: defaultMetadata,
+    });
+
+    expect(graph.dispatchRunStep).not.toHaveBeenCalled();
+    expect(graph.dispatchRunStepDelta).toHaveBeenCalledTimes(1);
+  });
+
+  it('sanitizes empty string id and name to undefined', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    const chunk = makeToolCallChunk({ id: '', name: '' });
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: [chunk],
+      metadata: defaultMetadata,
+    });
+
+    expect(chunk.id).toBeUndefined();
+    expect(chunk.name).toBeUndefined();
+  });
+
+  it('populates tool_calls when chunk has valid id and name', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    const chunks = [
+      makeToolCallChunk({ id: 'tooluse_abc', name: 'calculator', index: 2 }),
+    ];
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: chunks,
+      metadata: defaultMetadata,
+    });
+
+    const toolCallsArg = graph.dispatchRunStep.mock
+      .calls[0][1] as t.ToolCallsDetails;
+    expect(toolCallsArg.tool_calls).toEqual([
+      expect.objectContaining({
+        id: 'tooluse_abc',
+        name: 'calculator',
+        type: ToolCallTypes.TOOL_CALL,
+      }),
+    ]);
+  });
+
+  it('never dispatches empty text block alongside TOOL_CALLS step', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    const chunks = [
+      makeToolCallChunk({ id: 'tooluse_abc', name: 'calculator', index: 2 }),
+    ];
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: chunks,
+      metadata: defaultMetadata,
+    });
+
+    const allDispatches = graph.dispatchRunStep.mock.calls;
+    expect(allDispatches).toHaveLength(1);
+    const stepDetails = allDispatches[0][1] as t.StepDetails;
+    expect(stepDetails.type).toBe(StepTypes.TOOL_CALLS);
+    expect(stepDetails).not.toHaveProperty('content');
+    expect(stepDetails).not.toHaveProperty('text');
+  });
+
+  it('feeds buffer with id+args on same chunk (OpenAI/Anthropic pattern)', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    const chunk: ToolCallChunk = {
+      id: 'tc_1',
+      name: 'content_tool',
+      args: '{"action":"wri',
+      index: 0,
+      type: 'tool_call_chunk',
+    };
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: [chunk],
+      metadata: defaultMetadata,
+    });
+
+    expect(graph.streamingToolCallBuffer.setIndexMapping).toHaveBeenCalledWith(
+      0,
+      'tc_1'
+    );
+    expect(graph.streamingToolCallBuffer.setToolName).toHaveBeenCalledWith(
+      'tc_1',
+      'content_tool'
+    );
+    expect(graph.streamingToolCallBuffer.append).toHaveBeenCalledWith(
+      'tc_1',
+      '{"action":"wri'
+    );
+  });
+
+  it('feeds buffer via index lookup for Bedrock DELTA chunks (no id)', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+    // Simulate: index 0 was previously mapped to 'tc_1'
+    graph.streamingToolCallBuffer.getIdByIndex.mockReturnValue('tc_1');
+
+    const deltaChunk: ToolCallChunk = {
+      id: undefined,
+      name: undefined,
+      args: 'te","content":"hello',
+      index: 0,
+      type: 'tool_call_chunk',
+    };
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: [deltaChunk],
+      metadata: defaultMetadata,
+    });
+
+    expect(graph.streamingToolCallBuffer.getIdByIndex).toHaveBeenCalledWith(0);
+    expect(graph.streamingToolCallBuffer.append).toHaveBeenCalledWith(
+      'tc_1',
+      'te","content":"hello'
+    );
+  });
+
+  it('does not feed buffer when DELTA chunk has no index mapping', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+    // No index mapping exists
+    graph.streamingToolCallBuffer.getIdByIndex.mockReturnValue(undefined);
+
+    const deltaChunk: ToolCallChunk = {
+      id: undefined,
+      name: undefined,
+      args: 'some args',
+      index: 5,
+      type: 'tool_call_chunk',
+    };
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: [deltaChunk],
+      metadata: defaultMetadata,
+    });
+
+    expect(graph.streamingToolCallBuffer.append).not.toHaveBeenCalled();
+  });
+
+  it('dispatches delta even when chunks lack id/name (Bedrock pattern)', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    const chunks = [makeToolCallChunk({ index: 2 })];
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey: 'step-key',
+      toolCallChunks: chunks,
+      metadata: defaultMetadata,
+    });
+
+    expect(graph.dispatchRunStepDelta).toHaveBeenCalledTimes(1);
+    const toolCallsArg = graph.dispatchRunStep.mock
+      .calls[0][1] as t.ToolCallsDetails;
+    expect(toolCallsArg.tool_calls).toEqual([]);
+  });
+});
+
+describe('handleToolCalls', () => {
+  let graph: MockGraph;
+
+  beforeEach(() => {
+    graph = createMockGraph();
+  });
+
+  it('returns early when metadata is missing', async () => {
+    await handleToolCalls(
+      [makeToolCall('id-1')],
+      undefined,
+      graph as unknown as StandardGraph
+    );
+    expect(graph.dispatchRunStep).not.toHaveBeenCalled();
+  });
+
+  it('returns early when toolCalls is undefined', async () => {
+    await handleToolCalls(
+      undefined,
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+    expect(graph.dispatchRunStep).not.toHaveBeenCalled();
+  });
+
+  it('returns early when toolCalls is empty', async () => {
+    await handleToolCalls(
+      [],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+    expect(graph.dispatchRunStep).not.toHaveBeenCalled();
+  });
+
+  it('skips tool call when id already in toolCallStepIds', async () => {
+    graph.toolCallStepIds.set('id-1', 'existing-step');
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    await handleToolCalls(
+      [makeToolCall('id-1')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    expect(graph.dispatchRunStep).not.toHaveBeenCalled();
+  });
+
+  it('assigns fallback id when tool_call.id is undefined', async () => {
+    const tc: ToolCall = {
+      id: undefined as unknown as string,
+      name: 'calc',
+      args: {},
+      type: 'tool_call',
+    };
+    graph.getStepIdByKey.mockImplementation(() => {
+      throw new Error('no step');
+    });
+
+    await handleToolCalls(
+      [tc],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    expect(tc.id).toBeDefined();
+    expect(tc.id!.startsWith('toolu_')).toBe(true);
+    expect(graph.dispatchRunStep).toHaveBeenCalled();
+  });
+
+  it('flags messageStepHasToolCalls and dispatches TOOL_CALLS when prev step is MESSAGE_CREATION', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    await handleToolCalls(
+      [makeToolCall('id-1')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    expect(graph.messageStepHasToolCalls.get('prev-step-id')).toBe(true);
+    const calls = graph.dispatchRunStep.mock.calls;
+    expect(calls).toHaveLength(1);
+    expect(calls[0][1]).toEqual(
+      expect.objectContaining({ type: StepTypes.TOOL_CALLS })
+    );
+  });
+
+  it('creates MESSAGE_CREATION when no previous step exists', async () => {
+    graph.getStepIdByKey.mockImplementation(() => {
+      throw new Error('no step');
+    });
+
+    await handleToolCalls(
+      [makeToolCall('id-1')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    const calls = graph.dispatchRunStep.mock.calls;
+    expect(calls).toHaveLength(2);
+    expect(calls[0][1]).toEqual(
+      expect.objectContaining({ type: StepTypes.MESSAGE_CREATION })
+    );
+    expect(calls[1][1]).toEqual(
+      expect.objectContaining({ type: StepTypes.TOOL_CALLS })
+    );
+  });
+
+  it('reuses empty TOOL_CALLS step exactly once', async () => {
+    const emptyToolStep = makeRunStep(StepTypes.TOOL_CALLS, {
+      id: 'empty-step',
+      tool_calls: [],
+    });
+    graph.getRunStep.mockReturnValue(emptyToolStep);
+    graph.getStepIdByKey.mockReturnValue('empty-step-id');
+
+    await handleToolCalls(
+      [makeToolCall('id-1'), makeToolCall('id-2')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    expect(graph.toolCallStepIds.get('id-1')).toBe('empty-step-id');
+
+    const calls = graph.dispatchRunStep.mock.calls;
+    expect(calls).toHaveLength(1);
+    expect(calls[0][1]).toEqual(
+      expect.objectContaining({
+        type: StepTypes.TOOL_CALLS,
+        tool_calls: [expect.objectContaining({ id: 'id-2' })],
+      })
+    );
+  });
+
+  it('gives each parallel tool call its own step (3 tool calls)', async () => {
+    const emptyToolStep = makeRunStep(StepTypes.TOOL_CALLS, { tool_calls: [] });
+    graph.getStepIdByKey.mockReturnValue('chunk-step-id');
+
+    let callCount = 0;
+    graph.getRunStep.mockImplementation(() => {
+      if (callCount === 0) {
+        callCount++;
+        return emptyToolStep;
+      }
+      return makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'prev',
+            name: 'calc',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      });
+    });
+
+    await handleToolCalls(
+      [makeToolCall('id-1'), makeToolCall('id-2'), makeToolCall('id-3')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    expect(graph.toolCallStepIds.get('id-1')).toBe('chunk-step-id');
+
+    const calls = graph.dispatchRunStep.mock.calls;
+    expect(calls).toHaveLength(2);
+    expect((calls[0][1] as t.ToolCallsDetails).tool_calls![0].id).toBe('id-2');
+    expect((calls[1][1] as t.ToolCallsDetails).tool_calls![0].id).toBe('id-3');
+  });
+
+  it('never creates MESSAGE_CREATION for parallel tool calls after TOOL_CALLS prev', async () => {
+    const emptyToolStep = makeRunStep(StepTypes.TOOL_CALLS, { tool_calls: [] });
+    graph.getStepIdByKey.mockReturnValue('chunk-step-id');
+
+    let callCount = 0;
+    graph.getRunStep.mockImplementation(() => {
+      if (callCount === 0) {
+        callCount++;
+        return emptyToolStep;
+      }
+      return makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'prev',
+            name: 'calc',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      });
+    });
+
+    await handleToolCalls(
+      [makeToolCall('id-1'), makeToolCall('id-2'), makeToolCall('id-3')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    const msgCreationCalls = graph.dispatchRunStep.mock.calls.filter(
+      (call) => (call[1] as t.StepDetails).type === StepTypes.MESSAGE_CREATION
+    );
+    expect(msgCreationCalls).toHaveLength(0);
+  });
+
+  it('dispatches new TOOL_CALLS directly when prev TOOL_CALLS has existing data', async () => {
+    const populatedToolStep = makeRunStep(StepTypes.TOOL_CALLS, {
+      tool_calls: [
+        {
+          id: 'existing',
+          name: 'calc',
+          args: {},
+          type: 'tool_call',
+        } as t.AgentToolCall,
+      ],
+    });
+    graph.getRunStep.mockReturnValue(populatedToolStep);
+
+    await handleToolCalls(
+      [makeToolCall('id-1')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    const calls = graph.dispatchRunStep.mock.calls;
+    expect(calls).toHaveLength(1);
+    expect(calls[0][1]).toEqual(
+      expect.objectContaining({ type: StepTypes.TOOL_CALLS })
+    );
+  });
+
+  it('never dispatches empty text block with tool_call_ids (MESSAGE_CREATION path)', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    await handleToolCalls(
+      [makeToolCall('id-1')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    for (const call of graph.dispatchRunStep.mock.calls) {
+      const stepDetails = call[1] as t.StepDetails;
+      if (stepDetails.type === StepTypes.TOOL_CALLS) {
+        expect(stepDetails).not.toHaveProperty('content');
+        expect(stepDetails).not.toHaveProperty('text');
+      }
+    }
+  });
+
+  it('never dispatches empty text block with tool_call_ids (no prev step path)', async () => {
+    graph.getStepIdByKey.mockImplementation(() => {
+      throw new Error('no step');
+    });
+
+    await handleToolCalls(
+      [makeToolCall('id-1')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    for (const call of graph.dispatchRunStep.mock.calls) {
+      const stepDetails = call[1] as t.StepDetails;
+      if (stepDetails.type === StepTypes.TOOL_CALLS) {
+        expect(stepDetails).not.toHaveProperty('content');
+        expect(stepDetails).not.toHaveProperty('text');
+      }
+      if (stepDetails.type === StepTypes.MESSAGE_CREATION) {
+        const msgDetails = stepDetails as t.MessageCreationDetails;
+        expect(msgDetails.message_creation.message_id).toBeDefined();
+        expect(msgDetails).not.toHaveProperty('tool_call_ids');
+      }
+    }
+  });
+});
+
+describe('handleToolCallChunks + handleToolCalls integration', () => {
+  let graph: MockGraph;
+  const stepKey = 'step-key';
+
+  beforeEach(() => {
+    graph = createMockGraph();
+  });
+
+  it('Bedrock single tool: chunks create empty TOOL_CALLS, then handleToolCalls reuses it', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey,
+      toolCallChunks: [makeToolCallChunk({ index: 2 })],
+      metadata: defaultMetadata,
+    });
+
+    const chunkStepId = graph.dispatchRunStep.mock.results[0].value as string;
+    const resolvedChunkStepId = await chunkStepId;
+
+    const emptyToolStep = makeRunStep(StepTypes.TOOL_CALLS, { tool_calls: [] });
+    graph.getRunStep.mockReturnValue(emptyToolStep);
+    graph.getStepIdByKey.mockReturnValue(resolvedChunkStepId);
+    graph.dispatchRunStep.mockClear();
+
+    await handleToolCalls(
+      [makeToolCall('tooluse_abc')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    expect(graph.toolCallStepIds.get('tooluse_abc')).toBe(resolvedChunkStepId);
+    expect(graph.dispatchRunStep).not.toHaveBeenCalled();
+  });
+
+  it('Bedrock parallel: 3 chunks then 3 tool calls yields 3 unique step IDs', async () => {
+    const msgStep = makeRunStep(StepTypes.MESSAGE_CREATION);
+    graph.getRunStep.mockReturnValue(msgStep);
+
+    await handleToolCallChunks({
+      graph: graph as unknown as StandardGraph,
+      stepKey,
+      toolCallChunks: [makeToolCallChunk({ index: 2 })],
+      metadata: defaultMetadata,
+    });
+
+    const chunkStepId = await (graph.dispatchRunStep.mock.results[0]
+      .value as Promise<string>);
+
+    const emptyToolStep = makeRunStep(StepTypes.TOOL_CALLS, { tool_calls: [] });
+    graph.getStepIdByKey.mockReturnValue(chunkStepId);
+
+    let callIdx = 0;
+    graph.getRunStep.mockImplementation(() => {
+      if (callIdx === 0) {
+        callIdx++;
+        return emptyToolStep;
+      }
+      return makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'prev',
+            name: 'calc',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      });
+    });
+    graph.dispatchRunStep.mockClear();
+
+    let newStepCounter = 10;
+    graph.dispatchRunStep.mockImplementation(
+      async () => `new-step-${++newStepCounter}`
+    );
+
+    await handleToolCalls(
+      [makeToolCall('id-1'), makeToolCall('id-2'), makeToolCall('id-3')],
+      defaultMetadata,
+      graph as unknown as StandardGraph
+    );
+
+    expect(graph.toolCallStepIds.get('id-1')).toBe(chunkStepId);
+
+    const dispatchedIds = graph.dispatchRunStep.mock.calls.map(
+      (_, i) => graph.dispatchRunStep.mock.results[i].value
+    );
+    expect(dispatchedIds).toHaveLength(2);
+
+    const allStepIds = new Set([
+      chunkStepId,
+      graph.toolCallStepIds.get('id-1'),
+      ...graph.dispatchRunStep.mock.calls.map((call) => {
+        const tc = (call[1] as t.ToolCallsDetails).tool_calls;
+        return tc?.[0]?.id;
+      }),
+    ]);
+
+    expect(graph.toolCallStepIds.get('id-1')).toBe(chunkStepId);
+    expect(allStepIds.size).toBeGreaterThanOrEqual(2);
+
+    const msgCreationCalls = graph.dispatchRunStep.mock.calls.filter(
+      (call) => (call[1] as t.StepDetails).type === StepTypes.MESSAGE_CREATION
+    );
+    expect(msgCreationCalls).toHaveLength(0);
+  });
+});
+
+describe('handleServerToolResult', () => {
+  let graph: MockGraph;
+  const anthropicContext = { provider: Providers.ANTHROPIC } as AgentContext;
+
+  beforeEach(() => {
+    graph = createMockGraph();
+  });
+
+  it('returns false when provider is not Anthropic', async () => {
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [{ type: 'tool_result', tool_use_id: 'tu-1', content: 'ok' }],
+      agentContext: { provider: Providers.OPENAI } as AgentContext,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('returns false when content is a string', async () => {
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: 'plain text',
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('returns false when content is null/undefined', async () => {
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: undefined,
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('returns false when content is empty array', async () => {
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('returns false when single content item has no tool_use_id', async () => {
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [{ type: 'tool_result', content: 'ok' } as t.ToolResultContent],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('skips content parts with empty tool_use_id', async () => {
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [
+        { type: 'tool_result', tool_use_id: '', content: 'ok' },
+        { type: 'tool_result', tool_use_id: 'tu-valid', content: 'ok' },
+      ] as t.MessageContentComplex[],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('warns and skips when toolCallStepIds has no mapping for tool_use_id', async () => {
+    const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => {});
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [
+        { type: 'tool_result', tool_use_id: 'tu-missing', content: 'ok' },
+      ] as t.MessageContentComplex[],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+    expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('tu-missing'));
+    warnSpy.mockRestore();
+  });
+
+  it('warns when run step does not exist for stepId', async () => {
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(undefined);
+    const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => {});
+
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [
+        { type: 'tool_result', tool_use_id: 'tu-1', content: 'ok' },
+      ] as t.MessageContentComplex[],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+    expect(warnSpy).toHaveBeenCalledWith(
+      expect.stringContaining('does not exist')
+    );
+    warnSpy.mockRestore();
+  });
+
+  it('warns when run step is not a TOOL_CALLS type', async () => {
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(makeRunStep(StepTypes.MESSAGE_CREATION));
+    const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => {});
+
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [
+        { type: 'tool_result', tool_use_id: 'tu-1', content: 'ok' },
+      ] as t.MessageContentComplex[],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+    expect(warnSpy).toHaveBeenCalledWith(
+      expect.stringContaining('not a tool call step')
+    );
+    warnSpy.mockRestore();
+  });
+
+  it('skips when no matching tool call found in step details', async () => {
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(
+      makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'tu-other',
+            name: 'calc',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      })
+    );
+
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [
+        { type: 'tool_result', tool_use_id: 'tu-1', content: 'ok' },
+      ] as t.MessageContentComplex[],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('returns true and sets skipHandling when a valid tool result is found', async () => {
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(
+      makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'tu-1',
+            name: 'calc',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      })
+    );
+
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [
+        { type: 'tool_result', tool_use_id: 'tu-1', content: 'ok' },
+      ] as t.MessageContentComplex[],
+      agentContext: anthropicContext,
+    });
+    expect(result).toBe(true);
+  });
+
+  it('calls handleAnthropicSearchResults for web_search_result type', async () => {
+    const mockToolEndHandle = jest
+      .fn<(...args: unknown[]) => Promise<void>>()
+      .mockResolvedValue(undefined);
+    graph.handlerRegistry = {
+      getHandler: jest.fn().mockReturnValue({ handle: mockToolEndHandle }),
+    };
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(
+      makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'tu-1',
+            name: 'web_search',
+            args: { query: 'test' },
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      })
+    );
+
+    const webSearchContent: t.ToolResultContent = {
+      type: 'web_search_result',
+      tool_use_id: 'tu-1',
+      content: [
+        {
+          type: 'web_search_result',
+          url: 'https://example.com',
+          title: 'Example',
+          encrypted_index: 'abc',
+          page_age: '2024-01-01',
+        },
+      ],
+    };
+
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [webSearchContent] as t.MessageContentComplex[],
+      metadata: defaultMetadata,
+      agentContext: anthropicContext,
+    });
+
+    expect(result).toBe(true);
+    expect(mockToolEndHandle).toHaveBeenCalledWith(
+      GraphEvents.TOOL_END,
+      expect.objectContaining({ input: { query: 'test' } }),
+      defaultMetadata,
+      graph
+    );
+    expect(graph.invokedToolIds).toBeDefined();
+    expect(graph.invokedToolIds!.has('tu-1')).toBe(true);
+  });
+
+  it('initializes invokedToolIds set when null', async () => {
+    const mockToolEndHandle = jest
+      .fn<(...args: unknown[]) => Promise<void>>()
+      .mockResolvedValue(undefined);
+    graph.handlerRegistry = {
+      getHandler: jest.fn().mockReturnValue({ handle: mockToolEndHandle }),
+    };
+    graph.invokedToolIds = undefined;
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(
+      makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'tu-1',
+            name: 'web_search',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      })
+    );
+
+    const webSearchContent: t.ToolResultContent = {
+      type: 'web_search_tool_result',
+      tool_use_id: 'tu-1',
+      content: [
+        {
+          type: 'web_search_result',
+          url: 'https://example.com',
+          title: 'Test',
+          encrypted_index: 'x',
+        },
+      ],
+    };
+
+    await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [webSearchContent] as t.MessageContentComplex[],
+      metadata: defaultMetadata,
+      agentContext: anthropicContext,
+    });
+
+    expect(graph.invokedToolIds).toBeInstanceOf(Set);
+    expect(graph.invokedToolIds!.has('tu-1')).toBe(true);
+  });
+
+  it('warns when web search content is not an array', async () => {
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(
+      makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'tu-1',
+            name: 'web_search',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      })
+    );
+    const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => {});
+
+    const webSearchContent: t.ToolResultContent = {
+      type: 'web_search_result',
+      tool_use_id: 'tu-1',
+      content: 'not an array',
+    };
+
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [webSearchContent] as t.MessageContentComplex[],
+      metadata: defaultMetadata,
+      agentContext: anthropicContext,
+    });
+
+    expect(result).toBe(true);
+    expect(warnSpy).toHaveBeenCalledWith(
+      expect.stringContaining('Expected content to be an array')
+    );
+    warnSpy.mockRestore();
+  });
+
+  it('warns when content is not an Anthropic web search result', async () => {
+    graph.toolCallStepIds.set('tu-1', 'step-1');
+    graph.getRunStep.mockReturnValue(
+      makeRunStep(StepTypes.TOOL_CALLS, {
+        tool_calls: [
+          {
+            id: 'tu-1',
+            name: 'web_search',
+            args: {},
+            type: 'tool_call',
+          } as t.AgentToolCall,
+        ],
+      })
+    );
+    const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => {});
+
+    const webSearchContent: t.ToolResultContent = {
+      type: 'web_search_result',
+      tool_use_id: 'tu-1',
+      content: [{ type: 'text', text: 'not a search result' }],
+    };
+
+    const result = await handleServerToolResult({
+      graph: graph as unknown as StandardGraph,
+      content: [webSearchContent] as t.MessageContentComplex[],
+      metadata: defaultMetadata,
+      agentContext: anthropicContext,
+    });
+
+    expect(result).toBe(true);
+    expect(warnSpy).toHaveBeenCalledWith(
+      expect.stringContaining(
+        'Expected content to be an Anthropic web search result'
+      )
+    );
+    warnSpy.mockRestore();
+  });
+});