ai-terminal-assistant/packages/core/src/checkpoint/checkpoint-store.ts

/**
 * 检查点存储
 * 负责检查点的 CRUD 操作
 */

import { nanoid } from 'nanoid';
import type { ShadowGit } from './shadow-git.js';
import type { CheckpointLock } from './lock.js';
import type { CommitMessageGenerator } from './commit-message.js';
import type {
  CheckpointMetadata,
  CheckpointConfig,
  CheckpointTrigger,
} from './types.js';

/**
 * 检查点提交消息前缀
 */
const CHECKPOINT_PREFIX = 'checkpoint:';

/**
 * 创建检查点选项
 */
export interface CreateCheckpointOptions {
  name?: string;
  description?: string;
  trigger?: CheckpointTrigger;
  toolCall?: { tool: string; params: Record<string, unknown> };
  messageId?: string;
  sessionId?: string;
  turnIndex?: number;
}

/**
 * 检查点存储
 */
export class CheckpointStore {
  private shadowGit: ShadowGit;
  private lock: CheckpointLock;
  private commitMessageGenerator: CommitMessageGenerator;
  private config: CheckpointConfig;
  private index: Map<string, CheckpointMetadata> = new Map();

  constructor(
    shadowGit: ShadowGit,
    lock: CheckpointLock,
    commitMessageGenerator: CommitMessageGenerator,
    config: CheckpointConfig
  ) {
    this.shadowGit = shadowGit;
    this.lock = lock;
    this.commitMessageGenerator = commitMessageGenerator;
    this.config = config;
  }

  /**
   * 从 Git 历史加载检查点索引
   */
  async loadIndex(): Promise<void> {
    try {
      const commits = await this.shadowGit.getCommits(this.config.maxCheckpoints);

      for (const commit of commits) {
        if (commit.message.startsWith(CHECKPOINT_PREFIX)) {
          try {
            const jsonStr = commit.message.slice(CHECKPOINT_PREFIX.length);
            const metadata = JSON.parse(jsonStr) as CheckpointMetadata;
            metadata.commitHash = commit.hash;
            this.index.set(metadata.id, metadata);
          } catch {
            // 解析失败，跳过
          }
        }
      }
    } catch {
      // 仓库可能是空的
    }
  }

  /**
   * 创建检查点
   */
  async create(
    options: CreateCheckpointOptions,
    currentSessionId?: string | null,
    onCheckpointCreated?: (id: string) => void
  ): Promise<CheckpointMetadata> {
    return this.lock.withLock(async () => {
      return this.createInternal(options, currentSessionId, onCheckpointCreated);
    });
  }

  /**
   * 创建内部检查点（不使用锁，供内部调用）
   */
  async createInternal(
    options: CreateCheckpointOptions,
    currentSessionId?: string | null,
    onCheckpointCreated?: (id: string) => void
  ): Promise<CheckpointMetadata> {
    const id = nanoid(10);
    const timestamp = Date.now();
    const trigger = options.trigger || 'manual';

    // 创建元数据
    const metadata: CheckpointMetadata = {
      id,
      name: options.name,
      description: options.description,
      timestamp,
      trigger,
      toolCall: options.toolCall,
      commitHash: '',
      filesChanged: 0,
      messageId: options.messageId,
      sessionId: options.sessionId || currentSessionId || undefined,
      turnIndex: options.turnIndex,
    };

    // 获取变更文件数
    let filesChanged: Array<{ path: string; type: string }> = [];
    try {
      const diff = await this.shadowGit.getWorkingDirDiff();
      metadata.filesChanged = diff.files.length;
      filesChanged = diff.files;
    } catch {
      // 忽略
    }

    // 生成智能提交消息
    const humanReadableMessage = this.commitMessageGenerator.generateMessage(
      trigger,
      options.toolCall,
      filesChanged as Array<{ path: string; type: 'added' | 'modified' | 'deleted' | 'renamed' }>
    );

    // 创建 commit
    const commitMessage = CHECKPOINT_PREFIX + JSON.stringify({
      ...metadata,
      _readableMessage: humanReadableMessage,
    });
    const commitHash = await this.shadowGit.createCommit(commitMessage);
    metadata.commitHash = commitHash;

    // 更新索引
    this.index.set(id, metadata);

    // 通知回调
    onCheckpointCreated?.(id);

    return metadata;
  }

  /**
   * 列出所有检查点
   */
  list(): CheckpointMetadata[] {
    return Array.from(this.index.values()).sort(
      (a, b) => b.timestamp - a.timestamp
    );
  }

  /**
   * 获取检查点
   */
  get(idOrHash: string): CheckpointMetadata | null {
    // 先按 ID 查找
    if (this.index.has(idOrHash)) {
      return this.index.get(idOrHash)!;
    }

    // 再按 commit hash 查找
    for (const checkpoint of this.index.values()) {
      if (checkpoint.commitHash.startsWith(idOrHash)) {
        return checkpoint;
      }
    }

    return null;
  }

  /**
   * 获取最新检查点
   */
  getLatest(): CheckpointMetadata | null {
    const checkpoints = this.list();
    return checkpoints[0] || null;
  }

  /**
   * 删除检查点
   */
  delete(checkpointId: string): CheckpointMetadata | null {
    const checkpoint = this.index.get(checkpointId);
    if (!checkpoint) {
      return null;
    }

    this.index.delete(checkpointId);
    return checkpoint;
  }

  /**
   * 判断是否应该为指定工具创建检查点
   */
  shouldCreateForTool(tool: string): boolean {
    if (!this.config.enabled) return false;

    const { autoCheckpoint } = this.config;

    switch (tool) {
      case 'write_file':
        return autoCheckpoint.beforeWrite;
      case 'edit_file':
        return autoCheckpoint.beforeEdit;
      case 'delete_file':
        return autoCheckpoint.beforeDelete;
      case 'move_file':
      case 'copy_file':
        return autoCheckpoint.beforeMove;
      case 'bash':
        return autoCheckpoint.beforeBash;
      default:
        return false;
    }
  }

  /**
   * 生成检查点描述
   */
  generateDescription(tool: string, params: Record<string, unknown>): string {
    switch (tool) {
      case 'write_file':
        return `Write file: ${params.file_path || params.path}`;
      case 'edit_file':
        return `Edit file: ${params.file_path || params.path}`;
      case 'delete_file':
        return `Delete file: ${params.file_path || params.path}`;
      case 'move_file':
        return `Move: ${params.source} -> ${params.destination}`;
      case 'copy_file':
        return `Copy: ${params.source} -> ${params.destination}`;
      case 'bash':
        return `Bash: ${String(params.command).slice(0, 50)}`;
      default:
        return `Tool: ${tool}`;
    }
  }

  /**
   * 获取检查点数量
   */
  get size(): number {
    return this.index.size;
  }

  /**
   * 获取配置
   */
  getConfig(): CheckpointConfig {
    return this.config;
  }
}