obsidian_ollama/tests/conversation-state.test.ts

import { ConversationStateManager } from '../src/conversation-state';
import type { OllamaMessage } from '../src/types';

describe('ConversationStateManager', () => {
  let manager: ConversationStateManager;

  beforeEach(() => {
    manager = new ConversationStateManager(
      'You are an assistant that can help answer questions using the contents of a vault'
    );
  });

  describe('constructor', () => {
    it('should initialize with default system message in long-term context', () => {
      const customManager = new ConversationStateManager();
      const longTerm = customManager.getLongTermContext();
      expect(longTerm).toHaveLength(1);
      expect(longTerm[0].role).toBe('system');
      expect(longTerm[0].content).toContain('immediately emit the tool_call');
    });

    it('should initialize with empty short-term and medium-term contexts', () => {
      expect(manager.getShortTermContext()).toEqual([]);
      expect(manager.getMediumTermContext()).toEqual([]);
    });
  });

  describe('updateShortTermContext', () => {
    it('should add messages to short-term context', () => {
      const message: OllamaMessage = { role: 'user', content: 'Hello' };
      manager.updateShortTermContext(message);
      expect(manager.getShortTermContext()).toContainEqual(message);
    });

    it('should enforce maxShortTermTurns limit of 10', () => {
      for (let i = 0; i < 15; i++) {
        manager.updateShortTermContext({ role: 'user', content: `msg-${i}` });
      }
      const context = manager.getShortTermContext();
      expect(context).toHaveLength(10);
      expect(context[0].content).toBe('msg-5');
      expect(context[9].content).toBe('msg-14');
    });
  });

  describe('updateMediumTermContext', () => {
    it('should add messages to medium-term context', () => {
      const message: OllamaMessage = { role: 'system', content: 'KB result' };
      manager.updateMediumTermContext(message);
      expect(manager.getMediumTermContext()).toContainEqual(message);
    });

    it('should enforce maxMediumTermMessages limit of 20', () => {
      for (let i = 0; i < 25; i++) {
        manager.updateMediumTermContext({ role: 'system', content: `msg-${i}` });
      }
      const context = manager.getMediumTermContext();
      expect(context).toHaveLength(20);
      expect(context[0].content).toBe('msg-5');
      expect(context[19].content).toBe('msg-24');
    });
  });

  describe('setPersona', () => {
    it('should replace default system message with custom persona', () => {
      manager.setPersona('You are a coding expert.');
      const longTerm = manager.getLongTermContext();
      expect(longTerm).toHaveLength(1);
      expect(longTerm[0].content).toBe('You are a coding expert.');
    });

    it('should allow multiple persona updates', () => {
      manager.setPersona('First persona.');
      manager.setPersona('Second persona.');
      const longTerm = manager.getLongTermContext();
      expect(longTerm).toHaveLength(1);
      expect(longTerm[0].content).toBe('Second persona.');
    });
  });

  describe('getConversationContext', () => {
    it('should return all three context layers', () => {
      manager.updateShortTermContext({ role: 'user', content: 'Hi' });
      manager.updateMediumTermContext({ role: 'system', content: 'KB' });

      const context = manager.getConversationContext('test');
      expect(context.shortTermContext).toHaveLength(1);
      expect(context.mediumTermContext).toHaveLength(1);
      expect(context.longTermContext).toHaveLength(1);
    });
  });

  describe('getCompleteMessages', () => {
    it('should return messages in correct order: long, medium, short, current', () => {
      manager.updateShortTermContext({ role: 'user', content: 'Short' });
      manager.updateShortTermContext({ role: 'assistant', content: 'Short reply' });
      manager.updateMediumTermContext({ role: 'system', content: 'Medium' });

      const messages = manager.getCompleteMessages('Current');

      // Long-term comes first
      expect(messages[0].role).toBe('system');
      expect(messages[0].content).toContain('contents of a vault');

      // Medium-term follows
      expect(messages[1].content).toBe('Medium');

      // Short-term follows
      expect(messages[2].content).toBe('Short');
      expect(messages[3].content).toBe('Short reply');

      // Current user message last
      expect(messages[messages.length - 1].content).toBe('Current');
      expect(messages[messages.length - 1].role).toBe('user');
    });
  });

  describe('clear', () => {
    it('should reset short-term and medium-term contexts', () => {
      manager.updateShortTermContext({ role: 'user', content: 'msg' });
      manager.updateMediumTermContext({ role: 'system', content: 'msg' });
      manager.clear();
      expect(manager.getShortTermContext()).toEqual([]);
      expect(manager.getMediumTermContext()).toEqual([]);
    });

    it('should restore default system message in long-term context', () => {
      manager.setPersona('Custom persona');
      manager.clear();
      const longTerm = manager.getLongTermContext();
      expect(longTerm[0].content).toContain('immediately emit the tool_call');
    });

    it('should accept a custom system prompt when clearing', () => {
      manager.setPersona('Custom persona');
      manager.clear('Custom system prompt');
      const longTerm = manager.getLongTermContext();
      expect(longTerm[0].content).toBe('Custom system prompt');
    });
  });

  describe('setMediumTermContextFromQuery', () => {
    it('should clear previous medium-term context and add query result', () => {
      manager.updateMediumTermContext({ role: 'system', content: 'Old' });
      manager.setMediumTermContextFromQuery('New KB result');
      const medium = manager.getMediumTermContext();
      expect(medium).toHaveLength(1);
      expect(medium[0].content).toContain('New KB result');
      expect(medium[0].content).toContain('Knowledge base results for current query');
    });

    it('should ignore whitespace-only query results', () => {
      manager.setMediumTermContextFromQuery('   ');
      expect(manager.getMediumTermContext()).toEqual([]);
    });

    it('should ignore empty query results', () => {
      manager.setMediumTermContextFromQuery('');
      expect(manager.getMediumTermContext()).toEqual([]);
    });
  });

  describe('immutability', () => {
    it('should return copies from getters to prevent external mutation', () => {
      manager.updateShortTermContext({ role: 'user', content: 'test' });
      const shortCopy = manager.getShortTermContext();
      shortCopy.push({ role: 'assistant', content: 'injected' });
      expect(manager.getShortTermContext()).not.toContainEqual({
        role: 'assistant',
        content: 'injected',
      });
    });

    it('should return independent copies on repeated calls', () => {
      const copy1 = manager.getShortTermContext();
      const copy2 = manager.getShortTermContext();
      expect(copy1).not.toBe(copy2);
    });
  });

  describe('edge cases', () => {
    it('should handle empty user message in getCompleteMessages', () => {
      const messages = manager.getCompleteMessages('');
      expect(messages[messages.length - 1].content).toBe('');
      expect(messages[messages.length - 1].role).toBe('user');
    });

    it('should handle messages with tool_calls', () => {
      const message: OllamaMessage = {
        role: 'assistant',
        content: '',
        tool_calls: [
          {
            id: 'call_1',
            type: 'function',
            function: { name: 'test', arguments: '{}' },
          },
        ],
      };
      manager.updateShortTermContext(message);
      expect(manager.getShortTermContext()).toContainEqual(message);
    });

    it('should preserve long-term context across clear and restore', () => {
      manager.clear();
      const longTerm = manager.getLongTermContext();
      expect(longTerm).toHaveLength(1);
      expect(longTerm[0].role).toBe('system');
      expect(longTerm[0].content).toContain('immediately emit the tool_call');
    });
  });
});