illuma-agents 1.0.4 → 1.0.6

package/src/messages/cache.ts

@@ -9,6 +9,9 @@ type MessageWithContent = {
 
 /**
  * Anthropic API: Adds cache control to the appropriate user messages in the payload.
+ * Strips ALL existing cache control (both Anthropic and Bedrock formats) from all messages,
+ * then adds fresh cache control to the last 2 user messages in a single backward pass.
+ * This ensures we don't accumulate stale cache points across multiple turns.
  * @param messages - The array of message objects.
  * @returns - The updated array of message objects with cache control added.
  */
@@ -22,15 +25,26 @@ export function addCacheControl<T extends AnthropicMessage | BaseMessage>(
   const updatedMessages = [...messages];
   let userMessagesModified = 0;
 
-  for (
-    let i = updatedMessages.length - 1;
-    i >= 0 && userMessagesModified < 2;
-    i--
-  ) {
+  for (let i = updatedMessages.length - 1; i >= 0; i--) {
     const message = updatedMessages[i];
-    if ('getType' in message && message.getType() !== 'human') {
-      continue;
-    } else if ('role' in message && message.role !== 'user') {
+    const isUserMessage =
+      ('getType' in message && message.getType() === 'human') ||
+      ('role' in message && message.role === 'user');
+
+    if (Array.isArray(message.content)) {
+      message.content = message.content.filter(
+        (block) => !isCachePoint(block as MessageContentComplex)
+      ) as typeof message.content;
+
+      for (let j = 0; j < message.content.length; j++) {
+        const block = message.content[j] as Record<string, unknown>;
+        if ('cache_control' in block) {
+          delete block.cache_control;
+        }
+      }
+    }
+
+    if (userMessagesModified >= 2 || !isUserMessage) {
       continue;
     }
 
@@ -60,10 +74,77 @@ export function addCacheControl<T extends AnthropicMessage | BaseMessage>(
   return updatedMessages;
 }
 
+/**
+ * Checks if a content block is a cache point
+ */
+function isCachePoint(block: MessageContentComplex): boolean {
+  return 'cachePoint' in block && !('type' in block);
+}
+
+/**
+ * Removes all Anthropic cache_control fields from messages
+ * Used when switching from Anthropic to Bedrock provider
+ */
+export function stripAnthropicCacheControl<T extends MessageWithContent>(
+  messages: T[]
+): T[] {
+  if (!Array.isArray(messages)) {
+    return messages;
+  }
+
+  const updatedMessages = [...messages];
+
+  for (let i = 0; i < updatedMessages.length; i++) {
+    const message = updatedMessages[i];
+    const content = message.content;
+
+    if (Array.isArray(content)) {
+      for (let j = 0; j < content.length; j++) {
+        const block = content[j] as Record<string, unknown>;
+        if ('cache_control' in block) {
+          delete block.cache_control;
+        }
+      }
+    }
+  }
+
+  return updatedMessages;
+}
+
+/**
+ * Removes all Bedrock cachePoint blocks from messages
+ * Used when switching from Bedrock to Anthropic provider
+ */
+export function stripBedrockCacheControl<T extends MessageWithContent>(
+  messages: T[]
+): T[] {
+  if (!Array.isArray(messages)) {
+    return messages;
+  }
+
+  const updatedMessages = [...messages];
+
+  for (let i = 0; i < updatedMessages.length; i++) {
+    const message = updatedMessages[i];
+    const content = message.content;
+
+    if (Array.isArray(content)) {
+      message.content = content.filter(
+        (block) => !isCachePoint(block as MessageContentComplex)
+      ) as typeof content;
+    }
+  }
+
+  return updatedMessages;
+}
+
 /**
  * Adds Bedrock Converse API cache points to the last two messages.
  * Inserts `{ cachePoint: { type: 'default' } }` as a separate content block
  * immediately after the last text block in each targeted message.
+ * Strips ALL existing cache control (both Bedrock and Anthropic formats) from all messages,
+ * then adds fresh cache points to the last 2 messages in a single backward pass.
+ * This ensures we don't accumulate stale cache points across multiple turns.
  * @param messages - The array of message objects.
  * @returns - The updated array of message objects with cache points added.
  */
@@ -77,23 +158,32 @@ export function addBedrockCacheControl<
   const updatedMessages: T[] = messages.slice();
   let messagesModified = 0;
 
-  for (
-    let i = updatedMessages.length - 1;
-    i >= 0 && messagesModified < 2;
-    i--
-  ) {
+  for (let i = updatedMessages.length - 1; i >= 0; i--) {
     const message = updatedMessages[i];
-
-    if (
+    const isToolMessage =
       'getType' in message &&
       typeof message.getType === 'function' &&
-      message.getType() === 'tool'
-    ) {
-      continue;
-    }
+      message.getType() === 'tool';
 
     const content = message.content;
 
+    if (Array.isArray(content)) {
+      message.content = content.filter(
+        (block) => !isCachePoint(block)
+      ) as typeof content;
+
+      for (let j = 0; j < message.content.length; j++) {
+        const block = message.content[j] as Record<string, unknown>;
+        if ('cache_control' in block) {
+          delete block.cache_control;
+        }
+      }
+    }
+
+    if (messagesModified >= 2 || isToolMessage) {
+      continue;
+    }
+
     if (typeof content === 'string' && content === '') {
       continue;
     }
@@ -107,9 +197,9 @@ export function addBedrockCacheControl<
       continue;
     }
 
-    if (Array.isArray(content)) {
+    if (Array.isArray(message.content)) {
       let hasCacheableContent = false;
-      for (const block of content) {
+      for (const block of message.content) {
         if (block.type === ContentTypes.TEXT) {
           if (typeof block.text === 'string' && block.text !== '') {
             hasCacheableContent = true;
@@ -123,15 +213,15 @@ export function addBedrockCacheControl<
       }
 
       let inserted = false;
-      for (let j = content.length - 1; j >= 0; j--) {
-        const block = content[j] as MessageContentComplex;
+      for (let j = message.content.length - 1; j >= 0; j--) {
+        const block = message.content[j] as MessageContentComplex;
         const type = (block as { type?: string }).type;
         if (type === ContentTypes.TEXT || type === 'text') {
           const text = (block as { text?: string }).text;
           if (text === '' || text === undefined) {
             continue;
           }
-          content.splice(j + 1, 0, {
+          message.content.splice(j + 1, 0, {
             cachePoint: { type: 'default' },
           } as MessageContentComplex);
           inserted = true;
@@ -139,7 +229,7 @@ export function addBedrockCacheControl<
         }
       }
       if (!inserted) {
-        content.push({
+        message.content.push({
           cachePoint: { type: 'default' },
         } as MessageContentComplex);
       }

package/src/messages/ensureThinkingBlock.test.ts

@@ -0,0 +1,393 @@
+import { AIMessage, HumanMessage, ToolMessage } from '@langchain/core/messages';
+import type { ExtendedMessageContent } from '@/types';
+import { ensureThinkingBlockInMessages } from './format';
+import { Providers, ContentTypes } from '@/common';
+
+describe('ensureThinkingBlockInMessages', () => {
+  describe('messages with thinking blocks (should not be modified)', () => {
+    test('should not modify AI message that already has thinking block', () => {
+      const messages = [
+        new HumanMessage({ content: 'Hello' }),
+        new AIMessage({
+          content: [
+            { type: ContentTypes.THINKING, thinking: 'Let me think...' },
+            { type: 'text', text: 'Hi there!' },
+          ],
+        }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[1]).toBeInstanceOf(AIMessage);
+      expect((result[1].content as ExtendedMessageContent[])[0].type).toBe(
+        ContentTypes.THINKING
+      );
+    });
+
+    test('should not modify AI message that has redacted_thinking block', () => {
+      const messages = [
+        new HumanMessage({ content: 'Hello' }),
+        new AIMessage({
+          content: [
+            { type: 'redacted_thinking', data: 'redacted' },
+            { type: 'text', text: 'Hi there!' },
+          ],
+        }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[1]).toBeInstanceOf(AIMessage);
+      expect((result[1].content as ExtendedMessageContent[])[0].type).toBe(
+        'redacted_thinking'
+      );
+    });
+  });
+
+  describe('messages with tool_calls (should be converted)', () => {
+    test('should convert AI message with tool_calls to HumanMessage', () => {
+      const messages = [
+        new HumanMessage({ content: 'What is the weather?' }),
+        new AIMessage({
+          content: 'Let me check the weather.',
+          tool_calls: [
+            {
+              id: 'call_123',
+              name: 'get_weather',
+              args: { location: 'NYC' },
+              type: 'tool_call',
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Sunny, 75°F',
+          tool_call_id: 'call_123',
+        }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      // Should have 2 messages: HumanMessage + converted HumanMessage
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[0].content).toBe('What is the weather?');
+      expect(result[1]).toBeInstanceOf(HumanMessage);
+
+      // Check that the converted message includes the context prefix
+      expect(result[1].content).toContain('[Previous agent context]');
+      expect(result[1].content).toContain('Let me check the weather');
+      expect(result[1].content).toContain('Sunny, 75°F');
+    });
+
+    test('should convert AI message with tool_use in content to HumanMessage', () => {
+      const messages = [
+        new HumanMessage({ content: 'Search for something' }),
+        new AIMessage({
+          content: [
+            { type: 'text', text: 'Searching...' },
+            {
+              type: 'tool_use',
+              id: 'call_456',
+              name: 'search',
+              input: { query: 'test' },
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Found results',
+          tool_call_id: 'call_456',
+        }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[1]).toBeInstanceOf(HumanMessage);
+      expect(result[1].content).toContain('[Previous agent context]');
+      expect(result[1].content).toContain('Searching...');
+      expect(result[1].content).toContain('Found results');
+    });
+
+    test('should handle multiple tool messages in sequence', () => {
+      const messages = [
+        new HumanMessage({ content: 'Do multiple things' }),
+        new AIMessage({
+          content: 'I will perform multiple actions.',
+          tool_calls: [
+            {
+              id: 'call_1',
+              name: 'action1',
+              args: { param: 'a' },
+              type: 'tool_call',
+            },
+            {
+              id: 'call_2',
+              name: 'action2',
+              args: { param: 'b' },
+              type: 'tool_call',
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Result 1',
+          tool_call_id: 'call_1',
+        }),
+        new ToolMessage({
+          content: 'Result 2',
+          tool_call_id: 'call_2',
+        }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      // Should combine all tool messages into one HumanMessage
+      expect(result).toHaveLength(2);
+      expect(result[1]).toBeInstanceOf(HumanMessage);
+      expect(result[1].content).toContain('Result 1');
+      expect(result[1].content).toContain('Result 2');
+    });
+  });
+
+  describe('messages without tool calls (should pass through)', () => {
+    test('should not modify AI message without tool calls', () => {
+      const messages = [
+        new HumanMessage({ content: 'Hello' }),
+        new AIMessage({ content: 'Hi there, how can I help?' }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[0].content).toBe('Hello');
+      expect(result[1]).toBeInstanceOf(AIMessage);
+      expect(result[1].content).toBe('Hi there, how can I help?');
+    });
+
+    test('should preserve HumanMessages and other message types', () => {
+      const messages = [
+        new HumanMessage({ content: 'Question 1' }),
+        new AIMessage({ content: 'Answer 1' }),
+        new HumanMessage({ content: 'Question 2' }),
+        new AIMessage({ content: 'Answer 2' }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      expect(result).toHaveLength(4);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[1]).toBeInstanceOf(AIMessage);
+      expect(result[2]).toBeInstanceOf(HumanMessage);
+      expect(result[3]).toBeInstanceOf(AIMessage);
+    });
+  });
+
+  describe('mixed scenarios', () => {
+    test('should handle mix of normal and tool-using messages', () => {
+      const messages = [
+        new HumanMessage({ content: 'First question' }),
+        new AIMessage({ content: 'First answer without tools' }),
+        new HumanMessage({ content: 'Second question' }),
+        new AIMessage({
+          content: 'Using a tool',
+          tool_calls: [
+            {
+              id: 'call_abc',
+              name: 'some_tool',
+              args: {},
+              type: 'tool_call',
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Tool result',
+          tool_call_id: 'call_abc',
+        }),
+        new HumanMessage({ content: 'Third question' }),
+        new AIMessage({ content: 'Third answer without tools' }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      // Original message 1: HumanMessage (preserved)
+      // Original message 2: AIMessage without tools (preserved)
+      // Original message 3: HumanMessage (preserved)
+      // Original messages 4-5: AIMessage with tool + ToolMessage (converted to 1 HumanMessage)
+      // Original message 6: HumanMessage (preserved)
+      // Original message 7: AIMessage without tools (preserved)
+      expect(result).toHaveLength(6);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[1]).toBeInstanceOf(AIMessage);
+      expect(result[2]).toBeInstanceOf(HumanMessage);
+      expect(result[3]).toBeInstanceOf(HumanMessage); // Converted
+      expect(result[4]).toBeInstanceOf(HumanMessage);
+      expect(result[5]).toBeInstanceOf(AIMessage);
+    });
+
+    test('should handle multiple tool-using sequences', () => {
+      const messages = [
+        new HumanMessage({ content: 'Do task 1' }),
+        new AIMessage({
+          content: 'Doing task 1',
+          tool_calls: [
+            {
+              id: 'call_1',
+              name: 'tool1',
+              args: {},
+              type: 'tool_call',
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Result 1',
+          tool_call_id: 'call_1',
+        }),
+        new HumanMessage({ content: 'Do task 2' }),
+        new AIMessage({
+          content: 'Doing task 2',
+          tool_calls: [
+            {
+              id: 'call_2',
+              name: 'tool2',
+              args: {},
+              type: 'tool_call',
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Result 2',
+          tool_call_id: 'call_2',
+        }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      // Each tool sequence should be converted to a HumanMessage
+      expect(result).toHaveLength(4);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[0].content).toBe('Do task 1');
+      expect(result[1]).toBeInstanceOf(HumanMessage);
+      expect(result[1].content).toContain('Doing task 1');
+      expect(result[2]).toBeInstanceOf(HumanMessage);
+      expect(result[2].content).toBe('Do task 2');
+      expect(result[3]).toBeInstanceOf(HumanMessage);
+      expect(result[3].content).toContain('Doing task 2');
+    });
+  });
+
+  describe('edge cases', () => {
+    test('should handle empty messages array', () => {
+      const messages: never[] = [];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      expect(result).toHaveLength(0);
+    });
+
+    test('should handle AI message with empty content array', () => {
+      const messages = [
+        new HumanMessage({ content: 'Hello' }),
+        new AIMessage({ content: [] }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      expect(result).toHaveLength(2);
+      expect(result[1]).toBeInstanceOf(AIMessage);
+    });
+
+    test('should work with different providers', () => {
+      const messages = [
+        new AIMessage({
+          content: 'Using tool',
+          tool_calls: [
+            {
+              id: 'call_x',
+              name: 'test',
+              args: {},
+              type: 'tool_call',
+            },
+          ],
+        }),
+        new ToolMessage({
+          content: 'Result',
+          tool_call_id: 'call_x',
+        }),
+      ];
+
+      // Test with Anthropic
+      const resultAnthropic = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+      expect(resultAnthropic).toHaveLength(1);
+      expect(resultAnthropic[0]).toBeInstanceOf(HumanMessage);
+
+      // Test with Bedrock
+      const resultBedrock = ensureThinkingBlockInMessages(
+        messages,
+        Providers.BEDROCK
+      );
+      expect(resultBedrock).toHaveLength(1);
+      expect(resultBedrock[0]).toBeInstanceOf(HumanMessage);
+    });
+
+    test('should handle tool message without preceding AI message', () => {
+      const messages = [
+        new HumanMessage({ content: 'Hello' }),
+        new ToolMessage({
+          content: 'Unexpected tool result',
+          tool_call_id: 'call_orphan',
+        }),
+      ];
+
+      const result = ensureThinkingBlockInMessages(
+        messages,
+        Providers.ANTHROPIC
+      );
+
+      // Should preserve both messages as-is since tool message has no preceding AI message with tools
+      expect(result).toHaveLength(2);
+      expect(result[0]).toBeInstanceOf(HumanMessage);
+      expect(result[1]).toBeInstanceOf(ToolMessage);
+    });
+  });
+});