@illuma-ai/agents 1.0.90 → 1.0.94

package/src/graphs/contextManagement.test.ts

@@ -0,0 +1,625 @@
+import {
+  HumanMessage,
+  AIMessage,
+  SystemMessage,
+  ToolMessage,
+  BaseMessage,
+} from '@langchain/core/messages';
+import type { TokenCounter } from '@/types/run';
+import { createEmergencySummary, summarize } from '@/messages/summarize';
+import { getContextUtilization } from '@/messages/prune';
+
+/**
+ * Simple token counter that approximates 1 token per 4 characters.
+ * Mirrors the pattern used in summarize.test.ts.
+ */
+const simpleTokenCounter: TokenCounter = (msg: BaseMessage): number => {
+  const content =
+    typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
+  return Math.ceil(content.length / 4);
+};
+
+// ============================================================================
+// 1. getContextUtilization
+// ============================================================================
+
+describe('getContextUtilization', () => {
+  it('returns correct percentage for given token counts', () => {
+    const indexTokenCountMap: Record<string, number | undefined> = {
+      '0': 500,
+      '1': 300,
+      '2': 200,
+    };
+    const instructionTokens = 100;
+    const maxContextTokens = 10000;
+
+    // Total: 500 + 300 + 200 + 100 = 1100 => 11%
+    const result = getContextUtilization(
+      indexTokenCountMap,
+      instructionTokens,
+      maxContextTokens
+    );
+    expect(result).toBeCloseTo(11, 1);
+  });
+
+  it('returns 0 when maxContextTokens is 0', () => {
+    const result = getContextUtilization({ '0': 100 }, 50, 0);
+    expect(result).toBe(0);
+  });
+
+  it('returns 0 when maxContextTokens is negative', () => {
+    const result = getContextUtilization({ '0': 100 }, 50, -1);
+    expect(result).toBe(0);
+  });
+
+  it('handles empty indexTokenCountMap', () => {
+    const result = getContextUtilization({}, 200, 1000);
+    // Only instruction tokens: 200/1000 = 20%
+    expect(result).toBeCloseTo(20, 1);
+  });
+
+  it('handles undefined values in indexTokenCountMap', () => {
+    const indexTokenCountMap: Record<string, number | undefined> = {
+      '0': 100,
+      '1': undefined,
+      '2': 300,
+    };
+    // Total: 100 + 0 + 300 + 50 = 450 => 45%
+    const result = getContextUtilization(indexTokenCountMap, 50, 1000);
+    expect(result).toBeCloseTo(45, 1);
+  });
+
+  it('returns over 100 when tokens exceed max', () => {
+    const result = getContextUtilization({ '0': 5000 }, 1000, 4000);
+    // Total: 5000 + 1000 = 6000 => 150%
+    expect(result).toBeCloseTo(150, 1);
+  });
+});
+
+// ============================================================================
+// 2. Emergency compaction trigger (>95% utilization)
+// ============================================================================
+
+describe('Emergency compaction trigger (>95% utilization)', () => {
+  it('createEmergencySummary produces valid emergency summary from messages', () => {
+    const messages = [
+      new SystemMessage('System instructions'),
+      new HumanMessage('Please analyze the data'),
+      new AIMessage('I will analyze the data now'),
+      new ToolMessage({
+        content: 'Data results',
+        tool_call_id: 'tc1',
+        name: 'search',
+      }),
+      new HumanMessage('What about the second part?'),
+      new AIMessage('Processing second part'),
+    ];
+
+    const summary = createEmergencySummary(messages);
+
+    expect(summary).toContain('[Emergency Context Summary]');
+    expect(summary).toContain('Original request: Please analyze the data');
+    expect(summary).toContain('Last response: Processing second part');
+    expect(summary).toContain('Tools used: search');
+    expect(summary).toContain('Messages compacted: 6');
+  });
+
+  it('emergency summary can be wrapped in a SystemMessage for injection', () => {
+    const messages = [new HumanMessage('Hello'), new AIMessage('Hi there')];
+
+    const summary = createEmergencySummary(messages);
+    const summaryMsg = new SystemMessage(
+      `[Context Compacted — Emergency]\n${summary}\n\n[CONTEXT COMPACTION NOTICE]\nPrevious conversation was compacted to stay within context limits.`
+    );
+
+    expect(summaryMsg.getType()).toBe('system');
+    expect(summaryMsg.content).toContain('[Context Compacted — Emergency]');
+    expect(summaryMsg.content).toContain('[Emergency Context Summary]');
+    expect(summaryMsg.content).toContain('CONTEXT COMPACTION NOTICE');
+  });
+
+  it('keeps system message + summary + last N messages after emergency compaction', () => {
+    const finalMessages = [
+      new SystemMessage('System prompt'),
+      new HumanMessage('Q1'),
+      new AIMessage('A1'),
+      new HumanMessage('Q2'),
+      new AIMessage('A2'),
+      new HumanMessage('Q3'),
+      new AIMessage('A3'),
+      new HumanMessage('Q4'),
+    ];
+
+    // Simulate the emergency compaction logic from Graph.ts
+    const emergencySummary = createEmergencySummary(finalMessages);
+    const summaryMsg = new SystemMessage(
+      `[Context Compacted — Emergency]\n${emergencySummary}`
+    );
+    const systemIdx = finalMessages[0]?.getType() === 'system' ? 1 : 0;
+    const keepCount = Math.min(4, finalMessages.length - systemIdx);
+    const compacted = [
+      ...finalMessages.slice(0, systemIdx),
+      summaryMsg,
+      ...finalMessages.slice(finalMessages.length - keepCount),
+    ];
+
+    // System + summary + last 4 messages
+    expect(compacted).toHaveLength(6);
+    expect(compacted[0].getType()).toBe('system'); // Original system
+    expect(compacted[1].getType()).toBe('system'); // Summary
+    expect(compacted[1].content as string).toContain(
+      '[Context Compacted — Emergency]'
+    );
+    expect(compacted[compacted.length - 1].getType()).toBe('human'); // Last message (Q4)
+  });
+});
+
+// ============================================================================
+// 3. Proactive compaction trigger (>85% utilization)
+// ============================================================================
+
+describe('Proactive compaction trigger (>85% utilization)', () => {
+  it('summarize is called with correct messages when utilization is between 85-95%', async () => {
+    const messages = [
+      new SystemMessage('System prompt'),
+      new HumanMessage('First question'),
+      new AIMessage('First answer'),
+      new HumanMessage('Second question'),
+      new AIMessage('Second answer'),
+      new HumanMessage('Third question'),
+      new AIMessage('Third answer'),
+      new HumanMessage('Latest question'),
+    ];
+
+    // Messages to compact: skip system (idx 0) and last 3
+    const systemIdx = messages[0]?.getType() === 'system' ? 1 : 0;
+    const messagesToCompact = messages.slice(systemIdx, messages.length - 3);
+
+    expect(messagesToCompact).toHaveLength(4); // Q1, A1, Q2, A2
+    expect(messagesToCompact[0].getType()).toBe('human');
+    expect(messagesToCompact[messagesToCompact.length - 1].getType()).toBe(
+      'ai'
+    );
+
+    const mockCallback = jest
+      .fn()
+      .mockResolvedValue('## Summary\nCompacted context from 4 messages');
+
+    const result = await summarize(messagesToCompact, mockCallback, {
+      tokenCounter: simpleTokenCounter,
+      summaryBudget: 5000,
+    });
+
+    expect(result.tier).toBe('full');
+    expect(result.summary).toContain('Compacted context');
+    expect(result.messagesCompacted).toBe(4);
+    expect(mockCallback).toHaveBeenCalledTimes(1);
+  });
+
+  it('proactive compaction preserves system + summary + last 3 messages', async () => {
+    const messages = [
+      new SystemMessage('System prompt'),
+      new HumanMessage('Q1'),
+      new AIMessage('A1'),
+      new HumanMessage('Q2'),
+      new AIMessage('A2'),
+      new HumanMessage('Q3'),
+      new AIMessage('A3'),
+      new HumanMessage('Latest'),
+    ];
+
+    const mockCallback = jest
+      .fn()
+      .mockResolvedValue('Summary of compacted messages');
+
+    const systemIdx = messages[0]?.getType() === 'system' ? 1 : 0;
+    const messagesToCompact = messages.slice(systemIdx, messages.length - 3);
+
+    const summaryResult = await summarize(messagesToCompact, mockCallback);
+
+    const summaryMsg = new SystemMessage(
+      `[Context Compacted — ${summaryResult.tier}]\n${summaryResult.summary}`
+    );
+    const compacted = [
+      ...messages.slice(0, systemIdx),
+      summaryMsg,
+      ...messages.slice(messages.length - 3),
+    ];
+
+    // System + summary + last 3
+    expect(compacted).toHaveLength(5);
+    expect(compacted[0].getType()).toBe('system');
+    expect(compacted[1].getType()).toBe('system');
+    expect(compacted[1].content as string).toContain(
+      '[Context Compacted — full]'
+    );
+    expect(compacted[4].getType()).toBe('human'); // Latest
+  });
+});
+
+// ============================================================================
+// 4. TokenCounter fallback
+// ============================================================================
+
+describe('TokenCounter fallback', () => {
+  /**
+   * When tokenCounter is null, Graph.ts creates a fallback estimator
+   * using the rule: 4 characters = 1 token.
+   */
+  const fallbackCounter = (msg: BaseMessage): number => {
+    const content =
+      typeof msg.content === 'string'
+        ? msg.content
+        : JSON.stringify(msg.content);
+    return Math.ceil(content.length / 4);
+  };
+
+  it('estimates tokens as content length / 4', () => {
+    const msg = new HumanMessage('Hello world!'); // 12 chars
+    expect(fallbackCounter(msg)).toBe(3); // ceil(12/4) = 3
+  });
+
+  it('handles empty content', () => {
+    const msg = new HumanMessage('');
+    expect(fallbackCounter(msg)).toBe(0); // ceil(0/4) = 0
+  });
+
+  it('handles single character', () => {
+    const msg = new HumanMessage('a'); // 1 char
+    expect(fallbackCounter(msg)).toBe(1); // ceil(1/4) = 1
+  });
+
+  it('handles exactly 4 characters', () => {
+    const msg = new HumanMessage('abcd'); // 4 chars
+    expect(fallbackCounter(msg)).toBe(1); // ceil(4/4) = 1
+  });
+
+  it('handles 5 characters (rounds up)', () => {
+    const msg = new HumanMessage('abcde'); // 5 chars
+    expect(fallbackCounter(msg)).toBe(2); // ceil(5/4) = 2
+  });
+
+  it('handles complex (array) content by stringifying', () => {
+    const msg = new HumanMessage({
+      content: [{ type: 'text', text: 'Hello' }],
+    });
+    const stringified = JSON.stringify(msg.content);
+    expect(fallbackCounter(msg)).toBe(Math.ceil(stringified.length / 4));
+  });
+
+  it('handles long content', () => {
+    const longText = 'x'.repeat(1000); // 1000 chars
+    const msg = new HumanMessage(longText);
+    expect(fallbackCounter(msg)).toBe(250); // ceil(1000/4) = 250
+  });
+});
+
+// ============================================================================
+// 5. isInputTooLongError detection patterns
+// ============================================================================
+
+describe('isInputTooLongError detection', () => {
+  /**
+   * Mirrors the error detection logic in Graph.ts (lines ~1838-1844).
+   */
+  const isInputTooLongError = (errorMessage: string): boolean => {
+    const lower = errorMessage.toLowerCase();
+    return (
+      lower.includes('too long') ||
+      lower.includes('input is too long') ||
+      lower.includes('context length') ||
+      lower.includes('maximum context') ||
+      lower.includes('validationexception') ||
+      lower.includes('prompt is too long')
+    );
+  };
+
+  const errorPatterns = [
+    'input is too long',
+    'context length exceeded',
+    'maximum context length',
+    'validationexception: input is too long',
+    'prompt is too long',
+    'Your input is too long. Please reduce',
+    "This model's maximum context length is 128000 tokens",
+    'Request too long for model',
+  ];
+
+  for (const pattern of errorPatterns) {
+    it(`detects "${pattern}"`, () => {
+      expect(isInputTooLongError(pattern)).toBe(true);
+    });
+  }
+
+  it('detects mixed case error messages', () => {
+    expect(isInputTooLongError('INPUT IS TOO LONG')).toBe(true);
+    expect(isInputTooLongError('Context Length Exceeded')).toBe(true);
+    expect(isInputTooLongError('ValidationException: Prompt Is Too Long')).toBe(
+      true
+    );
+  });
+
+  it('does not detect unrelated errors', () => {
+    expect(isInputTooLongError('authentication failed')).toBe(false);
+    expect(isInputTooLongError('rate limit exceeded')).toBe(false);
+    expect(isInputTooLongError('network timeout')).toBe(false);
+    expect(isInputTooLongError('internal server error')).toBe(false);
+    expect(isInputTooLongError('invalid API key')).toBe(false);
+  });
+
+  it('does not match partial unrelated words', () => {
+    // "belong" contains "long" but NOT "too long"
+    expect(isInputTooLongError('this does not belong here')).toBe(false);
+  });
+});
+
+// ============================================================================
+// 6. Summarize-then-replace recovery flow
+// ============================================================================
+
+describe('summarize-then-replace recovery', () => {
+  it('produces compacted messages with summary and last 3 messages', async () => {
+    const messages = [
+      new SystemMessage('System prompt'),
+      new HumanMessage('First question'),
+      new AIMessage('First answer'),
+      new HumanMessage('Second question'),
+      new AIMessage('Second answer'),
+      new HumanMessage('Third question'),
+      new AIMessage('Third answer'),
+      new HumanMessage('Latest question'),
+    ];
+
+    const mockCallback = jest
+      .fn()
+      .mockResolvedValue('## Summary\nCompacted context');
+    const systemIdx = messages[0]?.getType() === 'system' ? 1 : 0;
+    const messagesToCompact = messages.slice(systemIdx, messages.length - 3);
+
+    const result = await summarize(messagesToCompact, mockCallback);
+
+    expect(result.tier).toBe('full');
+    expect(result.summary).toContain('Compacted context');
+    expect(result.messagesCompacted).toBe(4);
+
+    // Verify the recovery message structure
+    const summaryMsg = new SystemMessage(
+      `[Context Compacted — ${result.tier}]\n${result.summary}`
+    );
+    const recoveryMessages = [
+      ...messages.slice(0, systemIdx), // System message
+      summaryMsg, // Summary
+      ...messages.slice(messages.length - 3), // Last 3
+    ];
+
+    expect(recoveryMessages).toHaveLength(5);
+    expect(recoveryMessages[0].getType()).toBe('system'); // Original system
+    expect(recoveryMessages[1].getType()).toBe('system'); // Summary
+    expect(recoveryMessages[2].getType()).toBe('human'); // Third question (Q3)
+    expect(recoveryMessages[3].getType()).toBe('ai'); // Third answer (A3)
+    expect(recoveryMessages[4].getType()).toBe('human'); // Latest question
+  });
+
+  it('falls back to emergency summary when callback fails', async () => {
+    const messages = [
+      new SystemMessage('System prompt'),
+      new HumanMessage('Q1'),
+      new AIMessage('A1'),
+      new HumanMessage('Q2'),
+      new AIMessage('A2'),
+      new HumanMessage('Latest'),
+    ];
+
+    const mockCallback = jest.fn().mockRejectedValue(new Error('LLM down'));
+    const systemIdx = messages[0]?.getType() === 'system' ? 1 : 0;
+    const messagesToCompact = messages.slice(systemIdx, messages.length - 3);
+
+    const result = await summarize(messagesToCompact, mockCallback);
+
+    expect(result.tier).toBe('emergency');
+    expect(result.summary).toContain('[Emergency Context Summary]');
+    expect(result.messagesCompacted).toBe(2);
+  });
+
+  it('recovery messages are fewer than original messages', async () => {
+    const messages = [
+      new SystemMessage('System'),
+      new HumanMessage('Q1'),
+      new AIMessage('A1'),
+      new HumanMessage('Q2'),
+      new AIMessage('A2'),
+      new HumanMessage('Q3'),
+      new AIMessage('A3'),
+      new HumanMessage('Q4'),
+      new AIMessage('A4'),
+      new HumanMessage('Q5'),
+    ];
+
+    const mockCallback = jest.fn().mockResolvedValue('Short summary');
+    const systemIdx = messages[0]?.getType() === 'system' ? 1 : 0;
+    const messagesToCompact = messages.slice(systemIdx, messages.length - 3);
+
+    const result = await summarize(messagesToCompact, mockCallback);
+
+    const summaryMsg = new SystemMessage(
+      `[Context Compacted — ${result.tier}]\n${result.summary}`
+    );
+    const recoveryMessages = [
+      ...messages.slice(0, systemIdx),
+      summaryMsg,
+      ...messages.slice(messages.length - 3),
+    ];
+
+    // Original: 10, Recovery: system + summary + last 3 = 5
+    expect(recoveryMessages.length).toBeLessThan(messages.length);
+    expect(recoveryMessages).toHaveLength(5);
+  });
+});
+
+// ============================================================================
+// 7. StructuredOutputTruncatedError recovery
+// ============================================================================
+
+describe('StructuredOutputTruncatedError recovery', () => {
+  it('compaction reduces context for retry', async () => {
+    const messages = [
+      new SystemMessage('Instructions'),
+      new HumanMessage('Generate JSON'),
+      new AIMessage('Partial response that was very long and got truncated'),
+      new HumanMessage('More context to process'),
+    ];
+
+    const mockCallback = jest
+      .fn()
+      .mockResolvedValue('Compacted instructions and context');
+
+    // Compact middle messages (between system and last 2)
+    const systemIdx = messages[0]?.getType() === 'system' ? 1 : 0;
+    const messagesToCompact = messages.slice(systemIdx, messages.length - 2);
+
+    const result = await summarize(messagesToCompact, mockCallback, {
+      tokenCounter: simpleTokenCounter,
+    });
+
+    expect(result.tier).toBe('full');
+    expect(result.summary).toBe('Compacted instructions and context');
+
+    // After compaction: system + summary + last 2
+    const summaryMsg = new SystemMessage(
+      `[Context Compacted]\n${result.summary}`
+    );
+    const compacted = [messages[0], summaryMsg, ...messages.slice(-2)];
+
+    expect(compacted).toHaveLength(4);
+    expect(compacted.length).toBeLessThanOrEqual(messages.length);
+    expect(compacted[0].getType()).toBe('system');
+    expect(compacted[1].getType()).toBe('system');
+    expect(compacted[3].getType()).toBe('human');
+  });
+
+  it('emergency tier handles structured output recovery when LLM is unavailable', async () => {
+    const messages = [
+      new SystemMessage('Instructions'),
+      new HumanMessage('Generate complex JSON'),
+      new AIMessage('Very long partial response'),
+      new HumanMessage('Additional context'),
+      new AIMessage('Another partial'),
+      new HumanMessage('Final request'),
+    ];
+
+    const mockCallback = jest
+      .fn()
+      .mockRejectedValue(new Error('model overloaded'));
+
+    const systemIdx = messages[0]?.getType() === 'system' ? 1 : 0;
+    const messagesToCompact = messages.slice(systemIdx, messages.length - 2);
+
+    const result = await summarize(messagesToCompact, mockCallback);
+
+    expect(result.tier).toBe('emergency');
+    expect(result.summary).toContain('[Emergency Context Summary]');
+
+    // Even emergency recovery should produce a valid compacted set
+    const summaryMsg = new SystemMessage(
+      `[Context Compacted — emergency]\n${result.summary}`
+    );
+    const compacted = [messages[0], summaryMsg, ...messages.slice(-2)];
+
+    expect(compacted).toHaveLength(4);
+    expect(compacted.length).toBeLessThan(messages.length);
+  });
+});
+
+// ============================================================================
+// 8. ON_CONTEXT_PRESSURE event shape
+// ============================================================================
+
+describe('ON_CONTEXT_PRESSURE event', () => {
+  /**
+   * Tests the payload shape that Graph.ts emits via safeDispatchCustomEvent
+   * for the ON_CONTEXT_PRESSURE event. The threshold is determined by:
+   *   >95% => 'emergency'
+   *   >85% => 'critical'
+   *   >70% => 'warning'
+   */
+  /* eslint-disable @typescript-eslint/explicit-function-return-type */
+  const buildPressurePayload = (
+    utilizationPct: number,
+    maxTokens: number,
+    agentId: string
+  ) => ({
+    utilizationPct,
+    totalTokens: Math.round((maxTokens * utilizationPct) / 100),
+    maxTokens,
+    agentId,
+    threshold:
+      utilizationPct > 95
+        ? ('emergency' as const)
+        : utilizationPct > 85
+          ? ('critical' as const)
+          : ('warning' as const),
+  });
+  /* eslint-enable @typescript-eslint/explicit-function-return-type */
+
+  it('has correct shape for warning threshold (70-85%)', () => {
+    const payload = buildPressurePayload(75, 100000, 'test-agent');
+
+    expect(payload.threshold).toBe('warning');
+    expect(payload.utilizationPct).toBe(75);
+    expect(payload.totalTokens).toBe(75000);
+    expect(payload.maxTokens).toBe(100000);
+    expect(payload.agentId).toBe('test-agent');
+  });
+
+  it('has correct shape for critical threshold (85-95%)', () => {
+    const payload = buildPressurePayload(90, 100000, 'test-agent');
+
+    expect(payload.threshold).toBe('critical');
+    expect(payload.utilizationPct).toBe(90);
+    expect(payload.totalTokens).toBe(90000);
+  });
+
+  it('has correct shape for emergency threshold (>95%)', () => {
+    const payload = buildPressurePayload(97, 100000, 'test-agent');
+
+    expect(payload.threshold).toBe('emergency');
+    expect(payload.utilizationPct).toBe(97);
+    expect(payload.totalTokens).toBe(97000);
+  });
+
+  it('boundary: exactly 85% is warning, not critical', () => {
+    const payload = buildPressurePayload(85, 100000, 'agent-1');
+    expect(payload.threshold).toBe('warning');
+  });
+
+  it('boundary: 85.1% is critical', () => {
+    const payload = buildPressurePayload(85.1, 100000, 'agent-1');
+    expect(payload.threshold).toBe('critical');
+  });
+
+  it('boundary: exactly 95% is critical, not emergency', () => {
+    const payload = buildPressurePayload(95, 100000, 'agent-1');
+    expect(payload.threshold).toBe('critical');
+  });
+
+  it('boundary: 95.1% is emergency', () => {
+    const payload = buildPressurePayload(95.1, 100000, 'agent-1');
+    expect(payload.threshold).toBe('emergency');
+  });
+
+  it('totalTokens is derived from utilization and maxTokens', () => {
+    const payload = buildPressurePayload(80, 200000, 'agent-2');
+    expect(payload.totalTokens).toBe(160000);
+  });
+
+  it('event is only emitted when utilization > 70%', () => {
+    // Mirrors the guard in Graph.ts: `if (utilization > 70) { ... emit ... }`
+    const utilization = 65;
+    const shouldEmit = utilization > 70;
+    expect(shouldEmit).toBe(false);
+
+    const utilization2 = 71;
+    const shouldEmit2 = utilization2 > 70;
+    expect(shouldEmit2).toBe(true);
+  });
+});