simstudioai--sim/apps/sim/executor/handlers/agent/memory.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest'
import { MEMORY } from '@/executor/constants'
import { Memory } from '@/executor/handlers/agent/memory'
import type { Message } from '@/executor/handlers/agent/types'

vi.mock('@/lib/tokenization/estimators', () => ({
  getAccurateTokenCount: vi.fn((text: string) => {
    return Math.ceil(text.length / 4)
  }),
}))

describe('Memory', () => {
  let memoryService: Memory

  beforeEach(() => {
    memoryService = new Memory()
  })

  describe('applyWindow (message-based)', () => {
    it('should keep last N messages', () => {
      const messages: Message[] = [
        { role: 'user', content: 'Message 1' },
        { role: 'assistant', content: 'Response 1' },
        { role: 'user', content: 'Message 2' },
        { role: 'assistant', content: 'Response 2' },
        { role: 'user', content: 'Message 3' },
        { role: 'assistant', content: 'Response 3' },
      ]

      const result = (memoryService as any).applyWindow(messages, 4)

      expect(result.length).toBe(4)
      expect(result[0].content).toBe('Message 2')
      expect(result[3].content).toBe('Response 3')
    })

    it('should return all messages if limit exceeds array length', () => {
      const messages: Message[] = [
        { role: 'user', content: 'Test' },
        { role: 'assistant', content: 'Response' },
      ]

      const result = (memoryService as any).applyWindow(messages, 10)
      expect(result.length).toBe(2)
    })

    it('should handle invalid window size', () => {
      const messages: Message[] = [{ role: 'user', content: 'Test' }]

      const result = (memoryService as any).applyWindow(messages, Number.NaN)
      expect(result).toEqual(messages)
    })

    it('should handle zero limit', () => {
      const messages: Message[] = [{ role: 'user', content: 'Test' }]

      const result = (memoryService as any).applyWindow(messages, 0)
      expect(result).toEqual(messages)
    })
  })

  describe('applyTokenWindow (token-based)', () => {
    it('should keep messages within token limit', () => {
      const messages: Message[] = [
        { role: 'user', content: 'Short' },
        { role: 'assistant', content: 'This is a longer response message' },
        { role: 'user', content: 'Another user message here' },
        { role: 'assistant', content: 'Final response' },
      ]

      const result = (memoryService as any).applyTokenWindow(messages, 15, 'gpt-4o')

      expect(result.length).toBeGreaterThan(0)
      expect(result.length).toBeLessThan(messages.length)
      expect(result[result.length - 1].content).toBe('Final response')
    })

    it('should include at least 1 message even if it exceeds limit', () => {
      const messages: Message[] = [
        {
          role: 'user',
          content:
            'This is a very long message that definitely exceeds our small token limit of just 5 tokens',
        },
      ]

      const result = (memoryService as any).applyTokenWindow(messages, 5, 'gpt-4o')

      expect(result.length).toBe(1)
      expect(result[0].content).toBe(messages[0].content)
    })

    it('should process messages from newest to oldest', () => {
      const messages: Message[] = [
        { role: 'user', content: 'Old message' },
        { role: 'assistant', content: 'Old response' },
        { role: 'user', content: 'New message' },
        { role: 'assistant', content: 'New response' },
      ]

      const result = (memoryService as any).applyTokenWindow(messages, 10, 'gpt-4o')

      expect(result[result.length - 1].content).toBe('New response')
    })

    it('should handle invalid token limit', () => {
      const messages: Message[] = [{ role: 'user', content: 'Test' }]

      const result = (memoryService as any).applyTokenWindow(messages, Number.NaN, 'gpt-4o')
      expect(result).toEqual(messages)
    })

    it('should handle zero or negative token limit', () => {
      const messages: Message[] = [{ role: 'user', content: 'Test' }]

      const result1 = (memoryService as any).applyTokenWindow(messages, 0, 'gpt-4o')
      expect(result1).toEqual(messages)

      const result2 = (memoryService as any).applyTokenWindow(messages, -5, 'gpt-4o')
      expect(result2).toEqual(messages)
    })

    it('should work without model specified', () => {
      const messages: Message[] = [{ role: 'user', content: 'Test message' }]

      const result = (memoryService as any).applyTokenWindow(messages, 100, undefined)
      expect(result.length).toBe(1)
    })

    it('should handle empty messages array', () => {
      const messages: Message[] = []

      const result = (memoryService as any).applyTokenWindow(messages, 100, 'gpt-4o')
      expect(result).toEqual([])
    })
  })

  describe('validateConversationId', () => {
    it('should throw error for missing conversationId', () => {
      expect(() => {
        ;(memoryService as any).validateConversationId(undefined)
      }).toThrow('Conversation ID is required')
    })

    it('should throw error for empty conversationId', () => {
      expect(() => {
        ;(memoryService as any).validateConversationId('   ')
      }).toThrow('Conversation ID is required')
    })

    it('should throw error for too long conversationId', () => {
      const longId = 'a'.repeat(MEMORY.MAX_CONVERSATION_ID_LENGTH + 1)
      expect(() => {
        ;(memoryService as any).validateConversationId(longId)
      }).toThrow('Conversation ID too long')
    })

    it('should accept valid conversationId', () => {
      expect(() => {
        ;(memoryService as any).validateConversationId('user-123')
      }).not.toThrow()
    })
  })

  describe('validateContent', () => {
    it('should throw error for content exceeding max size', () => {
      const largeContent = 'x'.repeat(MEMORY.MAX_MESSAGE_CONTENT_BYTES + 1)
      expect(() => {
        ;(memoryService as any).validateContent(largeContent)
      }).toThrow('Message content too large')
    })

    it('should accept content within limit', () => {
      const content = 'Normal sized content'
      expect(() => {
        ;(memoryService as any).validateContent(content)
      }).not.toThrow()
    })
  })

  describe('sanitizeMessageForStorage', () => {
    it('should strip file payloads but preserve tool-call fields before memory persistence', () => {
      const message: Message = {
        role: 'user',
        content: 'Analyze this file',
        executionId: 'exec-1',
        files: [
          {
            id: 'file-1',
            key: 'workspace/ws-1/example.png',
            name: 'example.png',
            url: '/api/files/serve/workspace%2Fws-1%2Fexample.png?context=workspace',
            size: 128,
            type: 'image/png',
            base64: 'iVBORw0KGgo=',
          },
        ],
        tool_calls: [{ id: 'call-1' }],
      }

      expect((memoryService as any).sanitizeMessageForStorage(message)).toEqual({
        role: 'user',
        content: 'Analyze this file',
        executionId: 'exec-1',
        tool_calls: [{ id: 'call-1' }],
      })
    })
  })

  describe('Token-based vs Message-based comparison', () => {
    it('should produce different results for same limit concept', () => {
      const messages: Message[] = [
        { role: 'user', content: 'A' },
        {
          role: 'assistant',
          content: 'This is a much longer response that takes many more tokens',
        },
        { role: 'user', content: 'B' },
      ]

      const messageResult = (memoryService as any).applyWindow(messages, 2)
      expect(messageResult.length).toBe(2)

      const tokenResult = (memoryService as any).applyTokenWindow(messages, 10, 'gpt-4o')
      expect(tokenResult.length).toBeGreaterThanOrEqual(1)
    })
  })
})