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

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

// Mock SessionStorage
class MockSessionStorage extends SessionStorage {
  private mockCurrentSessionId: string | null = null;
  private mockSessions: Map<string, SessionData> = new Map();
  private mockProjects: Map<string, ProjectMetadata> = new Map();
  private defaultProject: ProjectMetadata = {
    id: 'test-project-id',
    workdir: '/test/workdir',
    createdAt: new Date().toISOString(),
    isGitRepo: true,
  };

  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 getOrCreateProject(workdir: string): Promise<ProjectMetadata> {
    const existing = this.mockProjects.get(workdir);
    if (existing) return existing;

    const project: ProjectMetadata = {
      id: `project-${workdir.replace(/\//g, '-')}`,
      workdir,
      createdAt: new Date().toISOString(),
      isGitRepo: true,
    };
    this.mockProjects.set(workdir, project);
    return project;
  }

  async setCurrentSession(sessionId: string): Promise<void> {
    this.mockCurrentSessionId = sessionId;
  }

  async getCurrentSessionId(): Promise<string | null> {
    return this.mockCurrentSessionId;
  }

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

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

  async loadSession(projectId: string, sessionId: string): Promise<SessionData | null> {
    const session = this.mockSessions.get(sessionId);
    if (session && session.projectId === projectId) {
      return session;
    }
    return null;
  }

  async saveSessionMetadata(session: SessionData): Promise<void> {
    await this.saveSession(session);
  }

  async listSessionsByProject(projectId: string): Promise<SessionSummary[]> {
    return Array.from(this.mockSessions.values())
      .filter((s) => s.projectId === projectId)
      .map((s) => ({
        id: s.id,
        title: s.title || `Session ${s.id}`,
        workdir: s.workdir,
        messageCount: s.messages.length,
        createdAt: s.createdAt,
        updatedAt: s.updatedAt,
      }))
      .sort((a, b) => new Date(b.updatedAt).getTime() - new Date(a.updatedAt).getTime());
  }

  async listAllSessions(): Promise<SessionSummary[]> {
    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,
      }))
      .sort((a, b) => new Date(b.updatedAt).getTime() - new Date(a.updatedAt).getTime());
  }

  async deleteSession(projectId: string, sessionId: string): Promise<boolean> {
    const session = this.mockSessions.get(sessionId);
    if (session && session.projectId === projectId) {
      return this.mockSessions.delete(sessionId);
    }
    return false;
  }

  async cleanupOldSessions(keepCount: number = 50): Promise<number> {
    const sessions = await this.listAllSessions();
    if (sessions.length <= keepCount) return 0;
    const toDelete = sessions.slice(keepCount);
    let count = 0;
    for (const s of toDelete) {
      // Find the projectId for this session
      const session = this.mockSessions.get(s.id);
      if (session && (await this.deleteSession(session.projectId, s.id))) {
        count++;
      }
    }
    return count;
  }

  // Helper methods for testing
  _setCurrentSessionId(sessionId: string | null): void {
    this.mockCurrentSessionId = sessionId;
  }

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

  _clear(): void {
    this.mockCurrentSessionId = null;
    this.mockSessions.clear();
    this.mockProjects.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 project = await storage.getOrCreateProject('/test/workdir');

      const existingSession: SessionData = {
        id: 'existing-session',
        projectId: project.id,
        createdAt: new Date().toISOString(),
        updatedAt: new Date().toISOString(),
        workdir: '/test/workdir',
        messages: [{ role: 'user', content: 'Hello' }],
        discoveredTools: ['tool1'],
        todos: [],
      };
      storage._addSession(existingSession);
      storage._setCurrentSessionId('existing-session');

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

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

    it('不同工作目录创建新会话', async () => {
      // 首先创建一个会话
      await manager.init('/old/workdir');
      await manager.addMessage({ role: 'user', content: 'Old message' });

      // 创建新的 manager 并用不同的工作目录初始化
      const newManager = new SessionManager(storage);
      const session = await newManager.init('/new/workdir');

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

  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');
    });
  });

  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.projectId, 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.save();
    });

    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' });

      for (let i = 1; i <= 2; i++) {
        await manager.newSession(`/workdir${i}`);
        await manager.addMessage({ role: 'user', content: `Message ${i}` });
      }
    });

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

    it('删除历史会话', async () => {
      const sessions = await manager.listSessions();
      if (sessions.length > 0) {
        const toDelete = sessions[0].id;
        const result = await manager.deleteSession(toDelete);
        expect(result).toBe(true);
      }
    });

    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' });

      for (let i = 1; i <= 9; i++) {
        await manager.newSession(`/workdir${i}`);
        await manager.addMessage({ role: 'user', content: `Message ${i}` });
      }
    });

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

      expect(deleted).toBeGreaterThan(0);
      const remaining = await manager.listAllSessions();
      expect(remaining.length).toBeLessThanOrEqual(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]+$/);
  });
});