@illuma-ai/agents 1.0.81

package/src/specs/emergency-prune.test.ts

@@ -0,0 +1,407 @@
+// src/specs/emergency-prune.test.ts
+/**
+ * Tests for the emergency pruning feature that handles "input too long" errors
+ * by retrying with more aggressive message pruning and adding a context notice.
+ */
+import {
+  HumanMessage,
+  AIMessage,
+  SystemMessage,
+  BaseMessage,
+} from '@langchain/core/messages';
+import type * as t from '@/types';
+import { createPruneMessages } from '@/messages/prune';
+import { Providers } from '@/common';
+
+// Simple token counter for testing (1 character = 1 token)
+const createTestTokenCounter = (): t.TokenCounter => {
+  return (message: BaseMessage): number => {
+    const content = message.content as
+      | string
+      | Array<t.MessageContentComplex | string>
+      | undefined;
+    if (typeof content === 'string') {
+      return content.length;
+    }
+    if (Array.isArray(content)) {
+      return content.reduce((total, item) => {
+        if (typeof item === 'string') return total + item.length;
+        if (
+          typeof item === 'object' &&
+          'text' in item &&
+          typeof item.text === 'string'
+        ) {
+          return total + item.text.length;
+        }
+        return total;
+      }, 0);
+    }
+    return 0;
+  };
+};
+
+// Helper to create test messages
+const createTestMessages = (
+  count: number,
+  tokensPer: number
+): BaseMessage[] => {
+  const messages: BaseMessage[] = [
+    new SystemMessage('You are a helpful assistant.'),
+  ];
+
+  for (let i = 0; i < count; i++) {
+    const content = 'x'.repeat(tokensPer);
+    if (i % 2 === 0) {
+      messages.push(new HumanMessage(content));
+    } else {
+      messages.push(new AIMessage(content));
+    }
+  }
+
+  return messages;
+};
+
+// Helper to build indexTokenCountMap
+const buildIndexTokenCountMap = (
+  messages: BaseMessage[],
+  tokenCounter: t.TokenCounter
+): Record<string, number> => {
+  const map: Record<string, number> = {};
+  messages.forEach((msg, index) => {
+    map[index] = tokenCounter(msg);
+  });
+  return map;
+};
+
+/**
+ * Estimates a human-friendly description of the conversation timeframe based on message count.
+ * This mirrors the implementation in Graph.ts
+ */
+const getContextTimeframeDescription = (messageCount: number): string => {
+  if (messageCount <= 5) {
+    return 'just the last few exchanges';
+  } else if (messageCount <= 15) {
+    return 'the last several minutes';
+  } else if (messageCount <= 30) {
+    return 'roughly the past hour';
+  } else if (messageCount <= 60) {
+    return 'the past couple of hours';
+  } else if (messageCount <= 150) {
+    return 'the past few hours';
+  } else if (messageCount <= 300) {
+    return 'roughly a day\'s worth';
+  } else if (messageCount <= 700) {
+    return 'the past few days';
+  } else {
+    return 'about a week or more';
+  }
+};
+
+describe('Emergency Pruning Feature', () => {
+  const tokenCounter = createTestTokenCounter();
+
+  describe('Normal Pruning vs Emergency Pruning', () => {
+    it('should prune more aggressively with 50% reduced context', () => {
+      // Create 20 messages, each with 100 tokens = 2000 tokens total (excluding system)
+      const messages = createTestMessages(20, 100);
+      const indexTokenCountMap = buildIndexTokenCountMap(
+        messages,
+        tokenCounter
+      );
+
+      // Normal prune with 1500 token limit
+      const normalMaxTokens = 1500;
+      const normalPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: normalMaxTokens,
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context: normalContext } = normalPrune({ messages });
+
+      // Emergency prune with 50% (750 tokens)
+      const emergencyMaxTokens = Math.floor(normalMaxTokens * 0.5);
+      const emergencyPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: emergencyMaxTokens,
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context: emergencyContext } = emergencyPrune({ messages });
+
+      // Emergency should have fewer messages
+      expect(emergencyContext.length).toBeLessThan(normalContext.length);
+      console.log(
+        `Normal prune: ${normalContext.length} messages, Emergency prune: ${emergencyContext.length} messages`
+      );
+    });
+
+    it('should preserve system message and latest user message after emergency prune', () => {
+      const messages = createTestMessages(10, 200);
+      const indexTokenCountMap = buildIndexTokenCountMap(
+        messages,
+        tokenCounter
+      );
+
+      // Very aggressive prune - only 300 tokens
+      const emergencyPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: 300,
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context } = emergencyPrune({ messages });
+
+      // Should still have system message if it fits
+      if (context.length > 0) {
+        // Check that we have at least the most recent messages
+        const lastMessage = context[context.length - 1];
+        expect(lastMessage).toBeDefined();
+      }
+    });
+  });
+
+  describe('Pruning Notice Message Injection', () => {
+    it('should calculate correct number of pruned messages', () => {
+      const originalCount = 20;
+      const messages = createTestMessages(originalCount, 100);
+      const indexTokenCountMap = buildIndexTokenCountMap(
+        messages,
+        tokenCounter
+      );
+
+      const emergencyPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: 500, // Very small to force aggressive pruning
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context: reducedMessages } = emergencyPrune({ messages });
+
+      // Calculate how many were pruned (this is what we inject in the notice)
+      const prunedCount = messages.length - reducedMessages.length;
+
+      expect(prunedCount).toBeGreaterThan(0);
+      console.log(
+        `Original: ${messages.length}, After prune: ${reducedMessages.length}, Pruned: ${prunedCount}`
+      );
+    });
+
+    it('should inject personalized notice message after system message', () => {
+      const messages = createTestMessages(10, 100);
+      const indexTokenCountMap = buildIndexTokenCountMap(
+        messages,
+        tokenCounter
+      );
+
+      const emergencyPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: 800,
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context: reducedMessages } = emergencyPrune({ messages });
+
+      // Simulate the notice injection logic from Graph.ts
+      const prunedCount = messages.length - reducedMessages.length;
+      const remainingCount = reducedMessages.length;
+      const estimatedContextDescription =
+        getContextTimeframeDescription(remainingCount);
+
+      const pruneNoticeMessage = new HumanMessage({
+        content: `[CONTEXT NOTICE]
+Our conversation has grown quite long, so I've focused on ${estimatedContextDescription} of our chat (${remainingCount} recent messages). ${prunedCount} earlier messages are no longer in my immediate memory.
+
+If I seem to be missing something we discussed earlier, just give me a quick reminder and I'll pick right back up! I'm still fully engaged and ready to help with whatever you need.`,
+      });
+
+      // Insert after system message
+      const hasSystemMessage = reducedMessages[0]?.getType() === 'system';
+      const insertIndex = hasSystemMessage ? 1 : 0;
+
+      const messagesWithNotice = [
+        ...reducedMessages.slice(0, insertIndex),
+        pruneNoticeMessage,
+        ...reducedMessages.slice(insertIndex),
+      ];
+
+      // Verify notice is in correct position
+      if (hasSystemMessage) {
+        expect(messagesWithNotice[0].getType()).toBe('system');
+        expect(messagesWithNotice[1].getType()).toBe('human');
+        expect(messagesWithNotice[1].content as string).toContain(
+          '[CONTEXT NOTICE]'
+        );
+        expect(messagesWithNotice[1].content as string).toContain(
+          'recent messages'
+        );
+        expect(messagesWithNotice[1].content as string).toContain(
+          'quick reminder'
+        );
+      } else {
+        expect(messagesWithNotice[0].getType()).toBe('human');
+        expect(messagesWithNotice[0].content as string).toContain(
+          '[CONTEXT NOTICE]'
+        );
+      }
+
+      // Total messages should be reduced + 1 notice
+      expect(messagesWithNotice.length).toBe(reducedMessages.length + 1);
+
+      console.log(
+        `Notice preview:\n${(pruneNoticeMessage.content as string).substring(0, 200)}...`
+      );
+    });
+  });
+
+  describe('Context Timeframe Description', () => {
+    it('should return appropriate descriptions for different message counts', () => {
+      expect(getContextTimeframeDescription(3)).toBe(
+        'just the last few exchanges'
+      );
+      expect(getContextTimeframeDescription(10)).toBe(
+        'the last several minutes'
+      );
+      expect(getContextTimeframeDescription(25)).toBe('roughly the past hour');
+      expect(getContextTimeframeDescription(45)).toBe(
+        'the past couple of hours'
+      );
+      expect(getContextTimeframeDescription(100)).toBe('the past few hours');
+      expect(getContextTimeframeDescription(200)).toBe('roughly a day\'s worth');
+      expect(getContextTimeframeDescription(500)).toBe('the past few days');
+      expect(getContextTimeframeDescription(1000)).toBe('about a week or more');
+    });
+  });
+
+  describe('Error Detection Patterns', () => {
+    const errorPatterns = [
+      'Input is too long for the model',
+      'context length exceeded',
+      'maximum context length',
+      'ValidationException: Input is too long',
+      'prompt is too long for this model',
+      'The input is too long',
+    ];
+
+    it('should detect various "input too long" error patterns', () => {
+      const isInputTooLongError = (errorMessage: string): boolean => {
+        const lowerMessage = errorMessage.toLowerCase();
+        return (
+          lowerMessage.includes('too long') ||
+          lowerMessage.includes('input is too long') ||
+          lowerMessage.includes('context length') ||
+          lowerMessage.includes('maximum context') ||
+          lowerMessage.includes('validationexception') ||
+          lowerMessage.includes('prompt is too long')
+        );
+      };
+
+      for (const pattern of errorPatterns) {
+        expect(isInputTooLongError(pattern)).toBe(true);
+        console.log(`✓ Detected: "${pattern}"`);
+      }
+
+      // Should not match unrelated errors
+      expect(isInputTooLongError('Network timeout')).toBe(false);
+      expect(isInputTooLongError('Invalid API key')).toBe(false);
+      expect(isInputTooLongError('Rate limit exceeded')).toBe(false);
+    });
+  });
+
+  describe('Edge Cases', () => {
+    it('should handle empty messages after pruning', () => {
+      // Single very long message that exceeds the limit
+      const messages: BaseMessage[] = [
+        new SystemMessage('System prompt'),
+        new HumanMessage('x'.repeat(10000)), // Way too long
+      ];
+      const indexTokenCountMap = buildIndexTokenCountMap(
+        messages,
+        tokenCounter
+      );
+
+      const emergencyPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: 100, // Very small limit
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context } = emergencyPrune({ messages });
+
+      // Should have at least tried to keep something or be empty
+      // The key is it shouldn't throw
+      expect(Array.isArray(context)).toBe(true);
+    });
+
+    it('should work with only system message and one user message', () => {
+      const messages: BaseMessage[] = [
+        new SystemMessage('You are helpful.'),
+        new HumanMessage('Hello'),
+      ];
+      const indexTokenCountMap = buildIndexTokenCountMap(
+        messages,
+        tokenCounter
+      );
+
+      const emergencyPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: 500,
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context } = emergencyPrune({ messages });
+
+      expect(context.length).toBe(2);
+      expect(context[0].getType()).toBe('system');
+      expect(context[1].getType()).toBe('human');
+    });
+
+    it('should handle conversation without system message', () => {
+      const messages: BaseMessage[] = [
+        new HumanMessage('Hello'),
+        new AIMessage('Hi there!'),
+        new HumanMessage('How are you?'),
+      ];
+      const indexTokenCountMap = buildIndexTokenCountMap(
+        messages,
+        tokenCounter
+      );
+
+      const emergencyPrune = createPruneMessages({
+        startIndex: 0,
+        provider: Providers.BEDROCK,
+        tokenCounter,
+        maxTokens: 100,
+        thinkingEnabled: false,
+        indexTokenCountMap,
+      });
+
+      const { context } = emergencyPrune({ messages });
+
+      // Should keep the most recent messages that fit
+      expect(context.length).toBeGreaterThan(0);
+      expect(context[0].getType()).not.toBe('system');
+    });
+  });
+});