ai-terminal-assistant/packages/core/src/checkpoint/session-tracker.ts

/**
 * 会话级检查点跟踪
 * 参考 Aider 的 aider_commit_hashes 和 commit_before_message 机制
 */

import { nanoid } from 'nanoid';
import type { CheckpointManager } from './manager.js';
import type {
  SessionState,
  SessionStats,
  CheckpointMetadata,
  RollbackResult,
} from './types.js';

/**
 * 会话跟踪器
 * 跟踪当前会话的所有检查点和文件修改
 */
export class SessionTracker {
  private currentSession: SessionState | null = null;
  private checkpointManager: CheckpointManager;

  constructor(checkpointManager: CheckpointManager) {
    this.checkpointManager = checkpointManager;
  }

  /**
   * 开始新会话
   */
  async startSession(): Promise<string> {
    // 创建会话开始检查点
    const startCp = await this.checkpointManager.createCheckpoint({
      trigger: 'session_start',
      description: 'Session start',
    });

    const sessionId = nanoid(10);

    this.currentSession = {
      id: sessionId,
      startTime: Date.now(),
      startCheckpoint: startCp.id,
      checkpoints: [startCp.id],
      modifiedFiles: [],
    };

    return sessionId;
  }

  /**
   * 结束当前会话
   */
  async endSession(): Promise<void> {
    if (!this.currentSession) return;

    // 创建会话结束检查点
    await this.checkpointManager.createCheckpoint({
      trigger: 'session_end',
      description: 'Session end',
    });

    this.currentSession = null;
  }

  /**
   * 获取当前会话 ID
   */
  getCurrentSessionId(): string | null {
    return this.currentSession?.id ?? null;
  }

  /**
   * 检查是否有活跃会话
   */
  hasActiveSession(): boolean {
    return this.currentSession !== null;
  }

  /**
   * 记录检查点到当前会话
   */
  recordCheckpoint(checkpointId: string): void {
    if (this.currentSession) {
      this.currentSession.checkpoints.push(checkpointId);
    }
  }

  /**
   * 记录文件修改
   */
  recordFileChange(filePath: string): void {
    if (this.currentSession) {
      if (!this.currentSession.modifiedFiles.includes(filePath)) {
        this.currentSession.modifiedFiles.push(filePath);
      }
    }
  }

  /**
   * 记录多个文件修改
   */
  recordFileChanges(filePaths: string[]): void {
    for (const filePath of filePaths) {
      this.recordFileChange(filePath);
    }
  }

  /**
   * 获取会话中修改的所有文件
   */
  getModifiedFiles(): string[] {
    return this.currentSession?.modifiedFiles ?? [];
  }

  /**
   * 获取会话中的所有检查点
   */
  async getSessionCheckpoints(): Promise<CheckpointMetadata[]> {
    if (!this.currentSession) {
      return [];
    }

    const checkpoints: CheckpointMetadata[] = [];
    for (const id of this.currentSession.checkpoints) {
      const cp = await this.checkpointManager.getCheckpoint(id);
      if (cp) {
        checkpoints.push(cp);
      }
    }

    return checkpoints.sort((a, b) => a.timestamp - b.timestamp);
  }

  /**
   * 撤销当前会话的所有修改
   */
  async undoSession(): Promise<RollbackResult> {
    if (!this.currentSession) {
      throw new Error('No active session');
    }

    return this.checkpointManager.rollback({
      target: this.currentSession.startCheckpoint,
      skipSafetyCheck: true, // 会话级撤销跳过安全检查
    });
  }

  /**
   * 回滚到会话中的指定检查点
   */
  async rollbackToCheckpoint(checkpointId: string): Promise<RollbackResult> {
    if (!this.currentSession) {
      throw new Error('No active session');
    }

    // 验证检查点属于当前会话
    if (!this.currentSession.checkpoints.includes(checkpointId)) {
      throw new Error('Checkpoint does not belong to current session');
    }

    return this.checkpointManager.rollback({
      target: checkpointId,
    });
  }

  /**
   * 获取会话统计信息
   */
  getSessionStats(): SessionStats {
    if (!this.currentSession) {
      throw new Error('No active session');
    }

    return {
      duration: Date.now() - this.currentSession.startTime,
      checkpointCount: this.currentSession.checkpoints.length,
      modifiedFilesCount: this.currentSession.modifiedFiles.length,
      modifiedFiles: [...this.currentSession.modifiedFiles],
    };
  }

  /**
   * 获取当前会话状态
   */
  getSessionState(): SessionState | null {
    if (!this.currentSession) {
      return null;
    }

    return {
      ...this.currentSession,
      modifiedFiles: [...this.currentSession.modifiedFiles],
      checkpoints: [...this.currentSession.checkpoints],
    };
  }

  /**
   * 获取会话开始时的检查点
   */
  getStartCheckpoint(): string | null {
    return this.currentSession?.startCheckpoint ?? null;
  }

  /**
   * 获取会话持续时间（毫秒）
   */
  getSessionDuration(): number {
    if (!this.currentSession) {
      return 0;
    }
    return Date.now() - this.currentSession.startTime;
  }
}

/**
 * 创建会话跟踪器实例
 */
export function createSessionTracker(
  checkpointManager: CheckpointManager
): SessionTracker {
  return new SessionTracker(checkpointManager);
}