@illuma-ai/agents 1.0.90 → 1.0.93

package/src/tools/__tests__/truncation-recovery.integration.test.ts

@@ -0,0 +1,560 @@
+/**
+ * Integration test: Full Truncation Recovery Pipeline
+ *
+ * Tests the end-to-end flow that would catch the Bedrock buffer bug:
+ *   Streaming chunks → handleToolCallChunks → StreamingToolCallBuffer → handleToolCalls → ToolNode.recoverTruncatedArgs → tool receives recovered args
+ *
+ * Uses real StreamingToolCallBuffer (not mocked) and real ToolNode with real tool definitions.
+ * Only the graph dispatch layer is mocked (since we're not testing SSE dispatch).
+ *
+ * Verified against live Bedrock API output (see bedrock-live-poc.ts):
+ *   START chunk: { id, name, index, args: "" }
+ *   DELTA chunks: { args, index } — NO id, NO name
+ */
+
+import { z } from 'zod';
+import { tool } from '@langchain/core/tools';
+import { AIMessage } from '@langchain/core/messages';
+import { describe, it, expect, beforeEach, jest } from '@jest/globals';
+import type { ToolCallChunk } from '@langchain/core/messages/tool';
+import type { StructuredToolInterface } from '@langchain/core/tools';
+import type { StandardGraph } from '@/graphs';
+import type * as t from '@/types';
+import { StepTypes } from '@/common';
+import { StreamingToolCallBuffer } from '../StreamingToolCallBuffer';
+import { handleToolCallChunks } from '../handlers';
+import { ToolNode } from '../ToolNode';
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/** Creates a mock graph that delegates buffer operations to a REAL StreamingToolCallBuffer */
+// eslint-disable-next-line @typescript-eslint/explicit-function-return-type
+function createIntegrationGraph(buffer: StreamingToolCallBuffer) {
+  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({
+      type: StepTypes.MESSAGE_CREATION,
+      id: 'prev-step-id',
+      index: 0,
+      stepDetails: {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: { message_id: 'msg-1' },
+      },
+      usage: null,
+    }),
+    dispatchRunStep: jest
+      .fn<() => Promise<string>>()
+      .mockImplementation(async () => `step-${++stepCounter}`),
+    dispatchRunStepDelta: jest
+      .fn<() => Promise<void>>()
+      .mockResolvedValue(undefined),
+    toolCallStepIds: new Map<string, string>(),
+    messageStepHasToolCalls: new Map<string, boolean>(),
+    messageIdsByStepKey: new Map<string, string>(),
+    prelimMessageIdsByStepKey: new Map<string, string>(),
+    streamingToolCallBuffer: buffer, // REAL buffer, not mocked
+  };
+}
+
+/** Creates a tool that captures received args for assertion */
+function createCaptureTool(
+  name: string,
+  schema: z.ZodObject<Record<string, z.ZodTypeAny>>,
+  capturedArgs: Record<string, unknown>[]
+): StructuredToolInterface {
+  return tool(
+    async (input) => {
+      capturedArgs.push({ ...input });
+      return 'OK';
+    },
+    { name, description: `Test tool: ${name}`, schema }
+  ) as unknown as StructuredToolInterface;
+}
+
+/** Simulates Bedrock streaming: START chunk with {id, name, index, args:""}, then DELTA chunks with {args, index} */
+function simulateBedrockChunks(
+  toolCallId: string,
+  toolName: string,
+  index: number,
+  argsJson: string,
+  chunkSize = 30
+): ToolCallChunk[] {
+  const chunks: ToolCallChunk[] = [];
+
+  // START chunk: id + name + index, args="" (verified from live Bedrock API)
+  chunks.push({
+    id: toolCallId,
+    name: toolName,
+    index,
+    args: '',
+    type: 'tool_call_chunk',
+  });
+
+  // DELTA chunks: args + index, NO id, NO name (verified from live Bedrock API)
+  for (let i = 0; i < argsJson.length; i += chunkSize) {
+    chunks.push({
+      id: undefined,
+      name: undefined,
+      index,
+      args: argsJson.substring(i, i + chunkSize),
+      type: 'tool_call_chunk',
+    });
+  }
+
+  return chunks;
+}
+
+/** Simulates OpenAI/Anthropic streaming: every chunk has {id, args} */
+function simulateOpenAIChunks(
+  toolCallId: string,
+  toolName: string,
+  index: number,
+  argsJson: string,
+  chunkSize = 30
+): ToolCallChunk[] {
+  const chunks: ToolCallChunk[] = [];
+
+  // First chunk: id + name + first args fragment
+  const firstArgs = argsJson.substring(0, chunkSize);
+  chunks.push({
+    id: toolCallId,
+    name: toolName,
+    index,
+    args: firstArgs,
+    type: 'tool_call_chunk',
+  });
+
+  // Subsequent chunks: id + args (OpenAI sends id on every chunk)
+  for (let i = chunkSize; i < argsJson.length; i += chunkSize) {
+    chunks.push({
+      id: toolCallId,
+      name: undefined,
+      index,
+      args: argsJson.substring(i, i + chunkSize),
+      type: 'tool_call_chunk',
+    });
+  }
+
+  return chunks;
+}
+
+const contentToolSchema = z.object({
+  action: z.string(),
+  filename: z.string().optional(),
+  content: z.string().optional(),
+});
+
+const codeExecutorSchema = z.object({
+  code: z.string(),
+  language: z.string().optional(),
+});
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe('Truncation Recovery Pipeline (Integration)', () => {
+  let buffer: StreamingToolCallBuffer;
+
+  beforeEach(() => {
+    buffer = new StreamingToolCallBuffer();
+  });
+
+  it('Bedrock truncated write — full pipeline recovery', async () => {
+    const capturedArgs: Record<string, unknown>[] = [];
+    const graph = createIntegrationGraph(buffer);
+
+    // Full args JSON that Bedrock would stream (including content)
+    const fullArgsJson =
+      '{"action":"write","filename":"app.tsx","content":"import React from \\"react\\";\\nexport default function App() {\\n  return <div>Hello World</div>;\\n}"}';
+
+    // Simulate Bedrock streaming: START + DELTA chunks
+    const chunks = simulateBedrockChunks(
+      'tooluse_abc',
+      'content_tool',
+      0,
+      fullArgsJson,
+      20
+    );
+
+    // Feed all chunks through handleToolCallChunks (real buffer accumulation)
+    for (const chunk of chunks) {
+      await handleToolCallChunks({
+        graph: graph as unknown as StandardGraph,
+        stepKey: 'step-key',
+        toolCallChunks: [chunk],
+        metadata: { run_id: 'test' },
+      });
+    }
+
+    // Verify buffer was populated
+    expect(buffer.has('tooluse_abc')).toBe(true);
+    expect(buffer.getRawArgs('tooluse_abc')).toBe(fullArgsJson);
+
+    // Simulate truncation: parsePartialJson lost the content field
+    const truncatedArgs = { action: 'write', filename: 'app.tsx' }; // content MISSING
+
+    // Create ToolNode with real buffer
+    const contentTool = createCaptureTool(
+      'content_tool',
+      contentToolSchema,
+      capturedArgs
+    );
+    const toolNode = new ToolNode({
+      tools: [contentTool],
+      streamingToolCallBuffer: buffer,
+    });
+
+    // Invoke ToolNode with truncated args
+    const aiMsg = new AIMessage({
+      content: '',
+      tool_calls: [
+        { id: 'tooluse_abc', name: 'content_tool', args: truncatedArgs },
+      ],
+    });
+    await toolNode.invoke({ messages: [aiMsg] });
+
+    // Assert: tool received RECOVERED content, not truncated
+    expect(capturedArgs).toHaveLength(1);
+    expect(capturedArgs[0].action).toBe('write');
+    expect(capturedArgs[0].filename).toBe('app.tsx');
+    expect(capturedArgs[0].content).toBe(
+      'import React from "react";\nexport default function App() {\n  return <div>Hello World</div>;\n}'
+    );
+  });
+
+  it('OpenAI/Anthropic chunked write — full pipeline recovery', async () => {
+    const capturedArgs: Record<string, unknown>[] = [];
+    const graph = createIntegrationGraph(buffer);
+
+    const fullArgsJson =
+      '{"action":"write","filename":"test.tsx","content":"const x = 42;\\nconst y = 100;"}';
+    const chunks = simulateOpenAIChunks(
+      'tc_openai',
+      'content_tool',
+      0,
+      fullArgsJson,
+      15
+    );
+
+    for (const chunk of chunks) {
+      await handleToolCallChunks({
+        graph: graph as unknown as StandardGraph,
+        stepKey: 'step-key',
+        toolCallChunks: [chunk],
+        metadata: { run_id: 'test' },
+      });
+    }
+
+    expect(buffer.has('tc_openai')).toBe(true);
+    expect(buffer.getRawArgs('tc_openai')).toBe(fullArgsJson);
+
+    // Simulate truncation: content missing
+    const contentTool = createCaptureTool(
+      'content_tool',
+      contentToolSchema,
+      capturedArgs
+    );
+    const toolNode = new ToolNode({
+      tools: [contentTool],
+      streamingToolCallBuffer: buffer,
+    });
+
+    const aiMsg = new AIMessage({
+      content: '',
+      tool_calls: [
+        {
+          id: 'tc_openai',
+          name: 'content_tool',
+          args: { action: 'write', filename: 'test.tsx' },
+        },
+      ],
+    });
+    await toolNode.invoke({ messages: [aiMsg] });
+
+    expect(capturedArgs[0].content).toBe('const x = 42;\nconst y = 100;');
+  });
+
+  it('multiple parallel tool calls — no cross-contamination', async () => {
+    const capturedArgs: Record<string, unknown>[] = [];
+    const graph = createIntegrationGraph(buffer);
+
+    // Two parallel tool calls with different indexes
+    const args1 =
+      '{"action":"write","filename":"a.tsx","content":"file A content"}';
+    const args2 =
+      '{"action":"write","filename":"b.tsx","content":"file B content"}';
+
+    const chunks1 = simulateBedrockChunks('tc_1', 'content_tool', 0, args1, 15);
+    const chunks2 = simulateBedrockChunks('tc_2', 'content_tool', 1, args2, 15);
+
+    // Interleave chunks (realistic concurrent streaming)
+    const interleaved: ToolCallChunk[] = [];
+    const maxLen = Math.max(chunks1.length, chunks2.length);
+    for (let i = 0; i < maxLen; i++) {
+      if (i < chunks1.length) interleaved.push(chunks1[i]);
+      if (i < chunks2.length) interleaved.push(chunks2[i]);
+    }
+
+    for (const chunk of interleaved) {
+      await handleToolCallChunks({
+        graph: graph as unknown as StandardGraph,
+        stepKey: 'step-key',
+        toolCallChunks: [chunk],
+        metadata: { run_id: 'test' },
+      });
+    }
+
+    expect(buffer.getRawArgs('tc_1')).toBe(args1);
+    expect(buffer.getRawArgs('tc_2')).toBe(args2);
+
+    const contentTool = createCaptureTool(
+      'content_tool',
+      contentToolSchema,
+      capturedArgs
+    );
+    const toolNode = new ToolNode({
+      tools: [contentTool],
+      streamingToolCallBuffer: buffer,
+    });
+
+    // Only tc_1 is truncated (missing content), tc_2 has all fields
+    const aiMsg = new AIMessage({
+      content: '',
+      tool_calls: [
+        {
+          id: 'tc_1',
+          name: 'content_tool',
+          args: { action: 'write', filename: 'a.tsx' },
+        }, // truncated
+        {
+          id: 'tc_2',
+          name: 'content_tool',
+          args: {
+            action: 'write',
+            filename: 'b.tsx',
+            content: 'file B content',
+          },
+        }, // complete
+      ],
+    });
+    await toolNode.invoke({ messages: [aiMsg] });
+
+    expect(capturedArgs).toHaveLength(2);
+    // tc_1: content recovered from buffer
+    expect(capturedArgs[0].content).toBe('file A content');
+    // tc_2: content unchanged (already present in parsed args)
+    expect(capturedArgs[1].content).toBe('file B content');
+  });
+
+  it('buffer cleanup after recovery — no stale data', async () => {
+    const capturedArgs: Record<string, unknown>[] = [];
+    const graph = createIntegrationGraph(buffer);
+
+    const fullArgs = '{"action":"write","filename":"x.tsx","content":"test"}';
+    const chunks = simulateBedrockChunks(
+      'tc_cleanup',
+      'content_tool',
+      0,
+      fullArgs
+    );
+
+    for (const chunk of chunks) {
+      await handleToolCallChunks({
+        graph: graph as unknown as StandardGraph,
+        stepKey: 'step-key',
+        toolCallChunks: [chunk],
+        metadata: { run_id: 'test' },
+      });
+    }
+
+    expect(buffer.has('tc_cleanup')).toBe(true);
+
+    const contentTool = createCaptureTool(
+      'content_tool',
+      contentToolSchema,
+      capturedArgs
+    );
+    const toolNode = new ToolNode({
+      tools: [contentTool],
+      streamingToolCallBuffer: buffer,
+    });
+
+    const aiMsg = new AIMessage({
+      content: '',
+      tool_calls: [
+        { id: 'tc_cleanup', name: 'content_tool', args: { action: 'write' } },
+      ],
+    });
+    await toolNode.invoke({ messages: [aiMsg] });
+
+    // Buffer should be cleared after processing
+    expect(buffer.has('tc_cleanup')).toBe(false);
+    expect(buffer.getRawArgs('tc_cleanup')).toBeUndefined();
+  });
+
+  it('no truncation — buffer populated but not used for overwrite', async () => {
+    const capturedArgs: Record<string, unknown>[] = [];
+    const graph = createIntegrationGraph(buffer);
+
+    const fullArgs = '{"action":"read","filename":"existing.tsx"}';
+    const chunks = simulateBedrockChunks(
+      'tc_complete',
+      'content_tool',
+      0,
+      fullArgs
+    );
+
+    for (const chunk of chunks) {
+      await handleToolCallChunks({
+        graph: graph as unknown as StandardGraph,
+        stepKey: 'step-key',
+        toolCallChunks: [chunk],
+        metadata: { run_id: 'test' },
+      });
+    }
+
+    const contentTool = createCaptureTool(
+      'content_tool',
+      contentToolSchema,
+      capturedArgs
+    );
+    const toolNode = new ToolNode({
+      tools: [contentTool],
+      streamingToolCallBuffer: buffer,
+    });
+
+    // Complete args — nothing missing
+    const aiMsg = new AIMessage({
+      content: '',
+      tool_calls: [
+        {
+          id: 'tc_complete',
+          name: 'content_tool',
+          args: { action: 'read', filename: 'existing.tsx' },
+        },
+      ],
+    });
+    await toolNode.invoke({ messages: [aiMsg] });
+
+    expect(capturedArgs[0].action).toBe('read');
+    expect(capturedArgs[0].filename).toBe('existing.tsx');
+    // No content field in args or buffer — should stay absent
+    expect(capturedArgs[0]).not.toHaveProperty('content');
+  });
+
+  it('generic tool recovery — non-content_tool (code_executor)', async () => {
+    const capturedArgs: Record<string, unknown>[] = [];
+    const graph = createIntegrationGraph(buffer);
+
+    const largeCode =
+      'function fibonacci(n) {\\n  if (n <= 1) return n;\\n  return fibonacci(n-1) + fibonacci(n-2);\\n}\\nconsole.log(fibonacci(10));';
+    const fullArgs = `{"code":"${largeCode}","language":"javascript"}`;
+    const chunks = simulateBedrockChunks(
+      'tc_code',
+      'execute_code',
+      0,
+      fullArgs,
+      20
+    );
+
+    for (const chunk of chunks) {
+      await handleToolCallChunks({
+        graph: graph as unknown as StandardGraph,
+        stepKey: 'step-key',
+        toolCallChunks: [chunk],
+        metadata: { run_id: 'test' },
+      });
+    }
+
+    const codeTool = createCaptureTool(
+      'execute_code',
+      codeExecutorSchema,
+      capturedArgs
+    );
+    const toolNode = new ToolNode({
+      tools: [codeTool],
+      streamingToolCallBuffer: buffer,
+    });
+
+    // Truncated: only has language, code is missing
+    const aiMsg = new AIMessage({
+      content: '',
+      tool_calls: [
+        {
+          id: 'tc_code',
+          name: 'execute_code',
+          args: { language: 'javascript' },
+        },
+      ],
+    });
+    await toolNode.invoke({ messages: [aiMsg] });
+
+    expect(capturedArgs[0].code).toContain('fibonacci');
+    expect(capturedArgs[0].language).toBe('javascript');
+  });
+
+  it('Bedrock large content (5000+ chars) — recovery preserves full content', async () => {
+    const capturedArgs: Record<string, unknown>[] = [];
+    const graph = createIntegrationGraph(buffer);
+
+    // Generate large content (simulating a real dashboard component)
+    const lines: string[] = [];
+    lines.push('import React from \\"react\\";');
+    lines.push('import { useState, useEffect } from \\"react\\";');
+    for (let i = 0; i < 100; i++) {
+      lines.push(`// Component section ${i}`);
+      lines.push(`const Section${i} = () => <div>Section ${i}</div>;`);
+    }
+    lines.push('export default function Dashboard() {');
+    lines.push('  return <div>Dashboard</div>;');
+    lines.push('}');
+    const content = lines.join('\\n');
+
+    const fullArgs = `{"action":"write","filename":"dashboard.tsx","content":"${content}"}`;
+    const chunks = simulateBedrockChunks(
+      'tc_large',
+      'content_tool',
+      0,
+      fullArgs,
+      50
+    );
+
+    for (const chunk of chunks) {
+      await handleToolCallChunks({
+        graph: graph as unknown as StandardGraph,
+        stepKey: 'step-key',
+        toolCallChunks: [chunk],
+        metadata: { run_id: 'test' },
+      });
+    }
+
+    const contentTool = createCaptureTool(
+      'content_tool',
+      contentToolSchema,
+      capturedArgs
+    );
+    const toolNode = new ToolNode({
+      tools: [contentTool],
+      streamingToolCallBuffer: buffer,
+    });
+
+    // Completely empty args (truncation wiped everything)
+    const aiMsg = new AIMessage({
+      content: '',
+      tool_calls: [{ id: 'tc_large', name: 'content_tool', args: {} }],
+    });
+    await toolNode.invoke({ messages: [aiMsg] });
+
+    expect(capturedArgs[0].action).toBe('write');
+    expect(capturedArgs[0].filename).toBe('dashboard.tsx');
+    const recoveredContent = capturedArgs[0].content as string;
+    expect(recoveredContent).toContain('import React from "react"');
+    expect(recoveredContent).toContain('Section 99');
+    expect(recoveredContent).toContain('export default function Dashboard()');
+  });
+});