ai-terminal-assistant/packages/core/tests/unit/session/manager.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { SessionManager } from '../../../src/session/manager.js';
import { SessionStorage } from '../../../src/session/storage.js';
import type { SessionData, Todo } from '../../../src/session/types.js';
import type { ModelMessage } from 'ai';

// Mock SessionStorage
class MockSessionStorage extends SessionStorage {
  private mockCurrentSession: SessionData | null = null;
  private mockSessions: Map<string, SessionData> = new Map();

  constructor() {
    super('/tmp/test-sessions');
  }

  async ensureDir(): Promise<void> {
    // no-op for testing
  }

  generateSessionId(): string {
    return `test-${Date.now()}-${Math.random().toString(36).substring(2, 6)}`;
  }

  async saveCurrentSession(session: SessionData): Promise<void> {
    this.mockCurrentSession = { ...session, updatedAt: new Date().toISOString() };
  }

  async loadCurrentSession(): Promise<SessionData | null> {
    return this.mockCurrentSession;
  }

  async archiveCurrentSession(): Promise<void> {
    if (this.mockCurrentSession) {
      this.mockSessions.set(this.mockCurrentSession.id, { ...this.mockCurrentSession });
      this.mockCurrentSession = null;
    }
  }

  async clearCurrentSession(): Promise<void> {
    this.mockCurrentSession = null;
  }

  async listSessions(): Promise<{ id: string; title: string; workdir: string; messageCount: number; createdAt: string; updatedAt: string }[]> {
    return Array.from(this.mockSessions.values()).map((s) => ({
      id: s.id,
      title: s.title || `Session ${s.id}`,
      workdir: s.workdir,
      messageCount: s.messages.length,
      createdAt: s.createdAt,
      updatedAt: s.updatedAt,
    }));
  }

  async loadSession(sessionId: string): Promise<SessionData | null> {
    return this.mockSessions.get(sessionId) || null;
  }

  async saveSession(session: SessionData): Promise<void> {
    this.mockSessions.set(session.id, { ...session, updatedAt: new Date().toISOString() });
  }

  async deleteSession(sessionId: string): Promise<boolean> {
    return this.mockSessions.delete(sessionId);
  }

  async cleanupOldSessions(keepCount: number = 50): Promise<number> {
    const sessions = await this.listSessions();
    if (sessions.length <= keepCount) return 0;
    const toDelete = sessions.slice(keepCount);
    let count = 0;
    for (const s of toDelete) {
      if (await this.deleteSession(s.id)) count++;
    }
    return count;
  }

  // Helper methods for testing
  _setCurrentSession(session: SessionData | null): void {
    this.mockCurrentSession = session;
  }

  _addSession(session: SessionData): void {
    this.mockSessions.set(session.id, session);
  }

  _clear(): void {
    this.mockCurrentSession = null;
    this.mockSessions.clear();
  }
}

describe('SessionManager - 会话管理器', () => {
  let storage: MockSessionStorage;
  let manager: SessionManager;

  beforeEach(() => {
    storage = new MockSessionStorage();
    manager = new SessionManager(storage);
    vi.useFakeTimers();
  });

  afterEach(() => {
    manager.stopAutoSave();
    vi.useRealTimers();
  });

  describe('init - 初始化', () => {
    it('无现有会话时创建新会话', async () => {
      const session = await manager.init('/test/workdir');

      expect(session).toBeDefined();
      expect(session.id).toBeDefined();
      expect(session.workdir).toBe('/test/workdir');
      expect(session.messages).toHaveLength(0);
      expect(session.discoveredTools).toHaveLength(0);
      expect(session.todos).toHaveLength(0);
    });

    it('同一工作目录恢复现有会话', async () => {
      const existingSession: SessionData = {
        id: 'existing-session',
        createdAt: new Date().toISOString(),
        updatedAt: new Date().toISOString(),
        workdir: '/test/workdir',
        messages: [{ role: 'user', content: 'Hello' }],
        discoveredTools: ['tool1'],
        todos: [],
      };
      storage._setCurrentSession(existingSession);

      const session = await manager.init('/test/workdir');

      expect(session.id).toBe('existing-session');
      expect(session.messages).toHaveLength(1);
    });

    it('不同工作目录创建新会话并归档旧会话', async () => {
      const existingSession: SessionData = {
        id: 'old-session',
        createdAt: new Date().toISOString(),
        updatedAt: new Date().toISOString(),
        workdir: '/old/workdir',
        messages: [{ role: 'user', content: 'Old message' }],
        discoveredTools: [],
        todos: [],
      };
      storage._setCurrentSession(existingSession);

      const session = await manager.init('/new/workdir');

      expect(session.id).not.toBe('old-session');
      expect(session.workdir).toBe('/new/workdir');
      // 旧会话应该被归档
      const sessions = await storage.listSessions();
      expect(sessions.some((s) => s.id === 'old-session')).toBe(true);
    });
  });

  describe('消息管理', () => {
    beforeEach(async () => {
      await manager.init('/test/workdir');
    });

    it('添加单条消息', async () => {
      const message: ModelMessage = { role: 'user', content: 'Test message' };
      await manager.addMessage(message);

      const messages = manager.getMessages();
      expect(messages).toHaveLength(1);
      expect(messages[0].content).toBe('Test message');
    });

    it('批量设置消息', async () => {
      const messages: ModelMessage[] = [
        { role: 'user', content: 'Message 1' },
        { role: 'assistant', content: 'Response 1' },
        { role: 'user', content: 'Message 2' },
      ];
      await manager.setMessages(messages);

      expect(manager.getMessages()).toHaveLength(3);
    });

    it('无当前会话时添加消息不报错', async () => {
      const newManager = new SessionManager(new MockSessionStorage());
      // 不调用 init，直接添加消息
      await newManager.addMessage({ role: 'user', content: 'Test' });
      expect(newManager.getMessages()).toHaveLength(0);
    });
  });

  describe('工具发现管理', () => {
    beforeEach(async () => {
      await manager.init('/test/workdir');
    });

    it('设置已发现的工具', async () => {
      await manager.setDiscoveredTools(['tool1', 'tool2', 'tool3']);

      expect(manager.getDiscoveredTools()).toEqual(['tool1', 'tool2', 'tool3']);
    });

    it('更新已发现的工具', async () => {
      await manager.setDiscoveredTools(['tool1']);
      await manager.setDiscoveredTools(['tool1', 'tool2']);

      expect(manager.getDiscoveredTools()).toEqual(['tool1', 'tool2']);
    });
  });

  describe('待办事项管理', () => {
    beforeEach(async () => {
      await manager.init('/test/workdir');
    });

    it('设置待办事项', async () => {
      const todos: Todo[] = [
        {
          id: '1',
          content: 'Task 1',
          status: 'pending',
          createdAt: new Date().toISOString(),
          updatedAt: new Date().toISOString(),
        },
        {
          id: '2',
          content: 'Task 2',
          status: 'in_progress',
          createdAt: new Date().toISOString(),
          updatedAt: new Date().toISOString(),
        },
      ];
      await manager.setTodos(todos);

      expect(manager.getTodos()).toHaveLength(2);
    });

    it('无当前会话时返回空数组', () => {
      const newManager = new SessionManager(new MockSessionStorage());
      expect(newManager.getTodos()).toEqual([]);
    });
  });

  describe('newSession - 创建新会话', () => {
    beforeEach(async () => {
      await manager.init('/test/workdir');
      await manager.addMessage({ role: 'user', content: 'Test' });
    });

    it('创建新会话并归档旧会话', async () => {
      const oldSessionId = manager.getSessionId();
      const newSession = await manager.newSession();

      expect(newSession.id).not.toBe(oldSessionId);
      expect(newSession.messages).toHaveLength(0);
    });

    it('使用指定工作目录创建新会话', async () => {
      const newSession = await manager.newSession('/new/workdir');

      expect(newSession.workdir).toBe('/new/workdir');
    });

    it('空消息会话不归档', async () => {
      const emptySession = await manager.newSession('/empty/workdir');
      const anotherSession = await manager.newSession('/another/workdir');

      // 空会话不应该被归档
      const sessions = await manager.listSessions();
      expect(sessions.every((s) => s.id !== emptySession.id)).toBe(true);
    });
  });

  describe('子会话管理', () => {
    beforeEach(async () => {
      await manager.init('/test/workdir');
    });

    it('创建子会话', () => {
      const parentId = manager.getSessionId()!;
      const childSession = manager.createChildSession(parentId, 'explore', 'Search task');

      expect(childSession.parentId).toBe(parentId);
      expect(childSession.agentName).toBe('explore');
      expect(childSession.title).toBe('Search task');
      expect(childSession.workdir).toBe('/test/workdir');
    });

    it('子会话使用默认标题', () => {
      const parentId = manager.getSessionId()!;
      const childSession = manager.createChildSession(parentId, 'code-reviewer');

      expect(childSession.title).toBe('子任务 (@code-reviewer)');
    });

    it('保存子会话', async () => {
      const parentId = manager.getSessionId()!;
      const childSession = manager.createChildSession(parentId, 'explore');
      childSession.messages.push({ role: 'user', content: 'Explore task' });

      await manager.saveChildSession(childSession);

      const saved = await storage.loadSession(childSession.id);
      expect(saved).not.toBeNull();
      expect(saved?.parentId).toBe(parentId);
    });
  });

  describe('会话恢复', () => {
    let archivedSessionId: string;

    beforeEach(async () => {
      // 清理之前的状态
      storage._clear();
      // 创建并归档一个会话
      await manager.init('/test/workdir');
      await manager.addMessage({ role: 'user', content: 'Archived message' });
      archivedSessionId = manager.getSessionId()!;
      await manager.newSession('/another/workdir');
    });

    it('恢复历史会话', async () => {
      const restored = await manager.restoreSession(archivedSessionId);

      expect(restored).not.toBeNull();
      expect(restored?.id).toBe(archivedSessionId);
      expect(manager.getMessages()).toHaveLength(1);
    });

    it('恢复不存在的会话返回 null', async () => {
      const result = await manager.restoreSession('non-existent-id');

      expect(result).toBeNull();
    });
  });

  describe('会话列表和删除', () => {
    beforeEach(async () => {
      // 清理之前的状态
      storage._clear();
      // 创建第一个会话
      await manager.init('/workdir0');
      await manager.addMessage({ role: 'user', content: 'Message 0' });
      // 创建后续会话（使用 newSession 避免 init 的额外归档）
      for (let i = 1; i <= 2; i++) {
        await manager.newSession(`/workdir${i}`);
        await manager.addMessage({ role: 'user', content: `Message ${i}` });
      }
      // 最后归档当前会话
      await manager.newSession('/final');
    });

    it('列出历史会话', async () => {
      const sessions = await manager.listSessions();
      expect(sessions.length).toBe(3);
    });

    it('删除历史会话', async () => {
      const sessions = await manager.listSessions();
      const toDelete = sessions[0].id;

      const result = await manager.deleteSession(toDelete);

      expect(result).toBe(true);
      const remaining = await manager.listSessions();
      expect(remaining.length).toBe(2);
    });

    it('删除不存在的会话返回 false', async () => {
      const result = await manager.deleteSession('non-existent');

      expect(result).toBe(false);
    });
  });

  describe('getSessionId', () => {
    it('返回当前会话 ID', async () => {
      await manager.init('/test/workdir');
      const id = manager.getSessionId();

      expect(id).toBeDefined();
      expect(typeof id).toBe('string');
    });

    it('无会话时返回 undefined', () => {
      const newManager = new SessionManager(new MockSessionStorage());
      expect(newManager.getSessionId()).toBeUndefined();
    });
  });

  describe('自动保存', () => {
    it('自动保存每 30 秒执行', async () => {
      await manager.init('/test/workdir');
      const saveSpy = vi.spyOn(manager, 'save');

      // 前进 30 秒
      await vi.advanceTimersByTimeAsync(30000);

      expect(saveSpy).toHaveBeenCalled();
    });

    it('stopAutoSave 停止自动保存', async () => {
      await manager.init('/test/workdir');
      const saveSpy = vi.spyOn(manager, 'save');

      manager.stopAutoSave();
      await vi.advanceTimersByTimeAsync(60000);

      // 只有 init 时调用了一次
      expect(saveSpy).toHaveBeenCalledTimes(0);
    });
  });

  describe('close - 关闭管理器', () => {
    it('关闭时保存并停止自动保存', async () => {
      await manager.init('/test/workdir');
      const saveSpy = vi.spyOn(manager, 'save');
      const stopSpy = vi.spyOn(manager, 'stopAutoSave');

      await manager.close();

      expect(saveSpy).toHaveBeenCalled();
      expect(stopSpy).toHaveBeenCalled();
    });
  });

  describe('cleanup - 清理旧会话', () => {
    beforeEach(async () => {
      // 清理之前的状态
      storage._clear();
      // 创建第一个会话
      await manager.init('/workdir0');
      await manager.addMessage({ role: 'user', content: 'Message 0' });
      // 创建 9 个后续会话（使用 newSession 避免 init 的额外归档）
      for (let i = 1; i <= 9; i++) {
        await manager.newSession(`/workdir${i}`);
        await manager.addMessage({ role: 'user', content: `Message ${i}` });
      }
      // 最后归档当前会话
      await manager.newSession('/final');
    });

    it('清理保留指定数量的会话', async () => {
      const deleted = await manager.cleanup(5);

      expect(deleted).toBe(5);
      const remaining = await manager.listSessions();
      expect(remaining.length).toBe(5);
    });

    it('会话数量不足时不清理', async () => {
      const deleted = await manager.cleanup(20);

      expect(deleted).toBe(0);
    });
  });
});

describe('SessionStorage - 会话存储', () => {
  it('generateSessionId 生成唯一 ID', () => {
    const storage = new SessionStorage('/tmp/test');
    const id1 = storage.generateSessionId();
    const id2 = storage.generateSessionId();

    expect(id1).not.toBe(id2);
    expect(id1).toMatch(/^\d{4}-\d{2}-\d{2}_[a-z0-9]+$/);
  });
});