auto-coder 1.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

autocoder/common/file_checkpoint/manager.py

@@ -9,6 +9,7 @@ import uuid
 import logging
 import difflib
 import time
+import hashlib
 from typing import Dict, List, Optional, Tuple, Any
 from datetime import datetime
 
@@ -22,6 +23,23 @@ from autocoder.common.file_checkpoint.conversation_checkpoint import Conversatio
 logger = logging.getLogger(__name__)
 
 
+def _generate_checkpoint_id(input_string: str) -> str:
+    """
+    根据输入字符串生成安全的checkpoint_id
+    
+    Args:
+        input_string: 输入字符串（通常是change_group_id或文件路径）
+        
+    Returns:
+        str: MD5哈希生成的checkpoint_id
+    """
+    if not input_string:
+        return str(uuid.uuid4())
+    
+    # 使用MD5生成固定长度的哈希值
+    return hashlib.md5(input_string.encode('utf-8')).hexdigest()
+
+
 class FileChangeManager:
     """文件变更管理器，提供高层次的API接口"""
     
@@ -48,15 +66,14 @@ class FileChangeManager:
             parent_dir = os.path.dirname(store_dir)
             conversation_store_dir = os.path.join(parent_dir, "conversation_checkpoints")
         
-        try:            
-            self.conversation_store = ConversationCheckpointStore(conversation_store_dir, max_history)
-            logger.info(f"对话检查点存储初始化成功: {conversation_store_dir}")
-        except ImportError as e:
-            logger.warning(f"对话检查点存储初始化失败: {str(e)}")
-            self.conversation_store = None
+        # 让异常直接抛出，确保 conversation_store 始终可用
+        self.conversation_store = ConversationCheckpointStore(conversation_store_dir, max_history)
+        logger.info(f"对话检查点存储初始化成功: {conversation_store_dir}")
     
     def apply_changes_with_conversation(self, changes: Dict[str, FileChange], 
-                                   conversations: List[Dict[str, Any]],
+                                   conversation_id: str,
+                                   first_message_id: Optional[str] = None,
+                                   last_message_id: Optional[str] = None,
                                    change_group_id: Optional[str] = None,
                                    metadata: Optional[Dict[str, Any]] = None) -> ApplyResult:
         """
@@ -64,31 +81,55 @@ class FileChangeManager:
         
         Args:
             changes: 文件变更字典，格式为 {file_path: FileChange}
-            conversations: 当前对话历史
+            conversation_id: 对话会话ID
+            first_message_id: 对话消息区间起始ID
+            last_message_id: 对话消息区间结束ID
             change_group_id: 变更组ID，用于将相关变更归为一组
             metadata: 元数据，可包含额外信息
             
         Returns:
             ApplyResult: 应用结果对象
         """
+        # 记录原始的change_group_id是否为None
+        original_group_id_is_none = change_group_id is None
+        
         # 应用文件变更
         result = self.apply_changes(changes, change_group_id)
         
-        if result.success and self.conversation_store is not None:
+        if result.success:
             try:                
+                # 创建并保存对话检查点 
+                # 生成安全的checkpoint_id（使用MD5哈希）
+                if not original_group_id_is_none:
+                    # 用户提供了group_id，使用MD5哈希
+                    checkpoint_id = _generate_checkpoint_id(change_group_id)
+                else:
+                    # 用户没有提供group_id，我们需要获取实际使用的group_id
+                    if result.change_ids:
+                        # 获取第一个变更的记录来找到实际的group_id
+                        first_change = self.change_store.get_change(result.change_ids[0])
+                        if first_change and first_change.group_id:
+                            checkpoint_id = _generate_checkpoint_id(first_change.group_id)
+                        else:
+                            checkpoint_id = result.change_ids[0]
+                    else:
+                        checkpoint_id = str(uuid.uuid4())
                 
-                # 创建并保存对话检查点
-                checkpoint_id = change_group_id or result.change_ids[0] if result.change_ids else str(uuid.uuid4())
                 checkpoint = ConversationCheckpoint(
                     checkpoint_id=checkpoint_id,
                     timestamp=time.time(),
-                    conversations=conversations,
+                    conversation_id=conversation_id,
+                    first_message_id=first_message_id,
+                    last_message_id=last_message_id,
                     metadata=metadata
                 )
-                self.conversation_store.save_checkpoint(checkpoint)
+                
+                saved_id = self.conversation_store.save_checkpoint(checkpoint)
                 logger.info(f"已保存对话检查点: {checkpoint_id}")
             except Exception as e:
                 logger.error(f"保存对话检查点失败: {str(e)}")
+        else:
+            logger.warning(f"文件变更应用失败，跳过对话检查点保存")
         
         return result
     
@@ -265,16 +306,17 @@ class FileChangeManager:
         
         # 获取关联的对话检查点
         checkpoint = None
-        checkpoint_id = change_record.group_id or change_id
-        if self.conversation_store is not None:
-            try:
-                checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
-                if checkpoint:
-                    logger.info(f"找到关联的对话检查点: {checkpoint_id}")
-                else:
-                    logger.info(f"未找到关联的对话检查点: {checkpoint_id}")
-            except Exception as e:
-                logger.error(f"获取对话检查点失败: {str(e)}")
+        # 生成安全的checkpoint_id
+        checkpoint_id = _generate_checkpoint_id(change_record.group_id) if change_record.group_id else change_id
+        
+        try:
+            checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
+            if checkpoint:
+                logger.info(f"找到关联的对话检查点: {checkpoint_id}")
+            else:
+                logger.debug(f"未找到关联的对话检查点: {checkpoint_id}")
+        except Exception as e:
+            logger.error(f"获取对话检查点失败: {str(e)}")
         
         # 撤销文件变更
         undo_result = self.undo_change(change_id)
@@ -359,15 +401,16 @@ class FileChangeManager:
         
         # 获取关联的对话检查点
         checkpoint = None
-        if self.conversation_store is not None:
-            try:
-                checkpoint = self.conversation_store.get_checkpoint(group_id)
-                if checkpoint:
-                    logger.info(f"找到关联的对话检查点: {group_id}")
-                else:
-                    logger.info(f"未找到关联的对话检查点: {group_id}")
-            except Exception as e:
-                logger.error(f"获取对话检查点失败: {str(e)}")
+        checkpoint_id = _generate_checkpoint_id(group_id) if group_id else group_id
+        
+        try:
+            checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
+            if checkpoint:
+                logger.info(f"找到关联的对话检查点: {checkpoint_id}")
+            else:
+                logger.debug(f"未找到关联的对话检查点: {checkpoint_id}")
+        except Exception as e:
+            logger.error(f"获取对话检查点失败: {str(e)}")
         
         # 撤销文件变更
         undo_result = self.undo_change_group(group_id)
@@ -422,16 +465,16 @@ class FileChangeManager:
         
         # 获取关联的对话检查点
         checkpoint = None
-        checkpoint_id = target_change.group_id or version_id
-        if self.conversation_store is not None:
-            try:
-                checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
-                if checkpoint:
-                    logger.info(f"找到关联的对话检查点: {checkpoint_id}")
-                else:
-                    logger.info(f"未找到关联的对话检查点: {checkpoint_id}")
-            except Exception as e:
-                logger.error(f"获取对话检查点失败: {str(e)}")
+        checkpoint_id = _generate_checkpoint_id(target_change.group_id) if target_change.group_id else version_id
+        
+        try:
+            checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
+            if checkpoint:
+                logger.info(f"找到关联的对话检查点: {checkpoint_id}")
+            else:
+                logger.debug(f"未找到关联的对话检查点: {checkpoint_id}")
+        except Exception as e:
+            logger.error(f"获取对话检查点失败: {str(e)}")
         
         # 撤销文件变更
         undo_result = self.undo_to_version(version_id)
@@ -551,12 +594,11 @@ class FileChangeManager:
             has_conversation = False
             
             # 检查是否有对话检查点
-            if self.conversation_store is not None:
-                try:
-                    checkpoint = self.conversation_store.get_checkpoint(group_id)
-                    has_conversation = checkpoint is not None
-                except Exception as e:
-                    logger.error(f"获取对话检查点失败: {str(e)}")
+            try:
+                checkpoint = self.conversation_store.get_checkpoint(group_id)
+                has_conversation = checkpoint is not None
+            except Exception as e:
+                logger.error(f"获取对话检查点失败: {str(e)}")
             
             checkpoints.append({
                 "id": group_id,
@@ -567,6 +609,180 @@ class FileChangeManager:
         
         return checkpoints
     
+    def get_checkpointed_message_ids(self, conversation_id: str) -> List[str]:
+        """
+        获取所有有checkpoint的message_id（只返回last_message_id）
+        
+        Args:
+            conversation_id: 对话ID
+            
+        Returns:
+            List[str]: 有checkpoint关联的last_message_id列表
+        """
+        checkpointed_message_ids = []
+        
+        try:
+            # 获取所有变更组
+            change_groups = self.get_change_groups(limit=1000)  # 获取更多的记录
+            
+            # 遍历每个变更组，查找对应的checkpoint
+            for group_id, _, _ in change_groups:
+                # 生成安全的checkpoint_id
+                checkpoint_id = _generate_checkpoint_id(group_id) if group_id else group_id
+                
+                try:
+                    checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
+                    if checkpoint and checkpoint.conversation_id == conversation_id:
+                        # 只收集last_message_id
+                        if checkpoint.last_message_id and checkpoint.last_message_id not in checkpointed_message_ids:
+                            checkpointed_message_ids.append(checkpoint.last_message_id)
+                except Exception as e:
+                    logger.debug(f"获取checkpoint {checkpoint_id} 失败: {str(e)}")
+                    continue
+            
+            # 也检查没有group_id的单独变更
+            latest_changes = self.change_store.get_latest_changes(limit=100)
+            for change in latest_changes:
+                if not change.group_id:  # 单独的变更
+                    checkpoint_id = change.change_id
+                    try:
+                        checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
+                        if checkpoint and checkpoint.conversation_id == conversation_id:
+                            if checkpoint.last_message_id and checkpoint.last_message_id not in checkpointed_message_ids:
+                                checkpointed_message_ids.append(checkpoint.last_message_id)
+                    except Exception:
+                        continue
+            
+            logger.info(f"找到 {len(checkpointed_message_ids)} 个有checkpoint的last_message_id")
+            return checkpointed_message_ids
+            
+        except Exception as e:
+            logger.error(f"获取checkpointed message_ids失败: {str(e)}")
+            return checkpointed_message_ids
+    
+    def rollback_to_message(self, message_id: str, conversation_id: Optional[str] = None) -> Tuple[UndoResult, Optional[ConversationCheckpoint]]:
+        """
+        回滚到指定message_id所在的checkpoint
+        
+        Args:
+            message_id: 目标消息ID
+            conversation_id: 可选的对话ID，用于过滤checkpoint
+            
+        Returns:
+            Tuple[UndoResult, Optional[ConversationCheckpoint]]: 回滚结果和找到的对话检查点
+        """
+        try:
+            # 查找包含该message_id的checkpoint
+            found_checkpoint = None
+            found_group_id = None
+            found_change_id = None
+            
+            # 先检查有group_id的变更
+            change_groups = self.get_change_groups(limit=1000)
+            for group_id, _, _ in change_groups:
+                checkpoint_id = _generate_checkpoint_id(group_id) if group_id else group_id
+                
+                try:
+                    checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
+                    if checkpoint:
+                        # 检查message_id是否在这个checkpoint的范围内
+                        is_in_range = False
+                        
+                        # 检查是否匹配conversation_id（如果提供了）
+                        if conversation_id and checkpoint.conversation_id != conversation_id:
+                            continue
+                            
+                        if checkpoint.first_message_id == message_id or checkpoint.last_message_id == message_id:
+                            is_in_range = True
+                        elif checkpoint.first_message_id and checkpoint.last_message_id:
+                            # 简化版本：如果message_id在checkpoint的范围内，我们假设它匹配
+                            # 注意：这里无法准确判断message_id是否在范围内，除非有完整的消息列表
+                            # 但如果first或last匹配，我们已经在上面处理了
+                            pass
+                        
+                        if is_in_range:
+                            found_checkpoint = checkpoint
+                            found_group_id = group_id
+                            logger.info(f"找到包含message_id {message_id} 的checkpoint (group: {group_id})")
+                            break
+                            
+                except Exception as e:
+                    logger.debug(f"检查checkpoint {checkpoint_id} 时出错: {str(e)}")
+                    continue
+            
+            # 如果没找到，检查单独的变更
+            if not found_checkpoint:
+                latest_changes = self.change_store.get_latest_changes(limit=100)
+                for change in latest_changes:
+                    if not change.group_id:
+                        checkpoint_id = change.change_id
+                        try:
+                            checkpoint = self.conversation_store.get_checkpoint(checkpoint_id)
+                            if checkpoint:
+                                if (checkpoint.first_message_id == message_id or 
+                                    checkpoint.last_message_id == message_id):
+                                    found_checkpoint = checkpoint
+                                    found_change_id = change.change_id
+                                    logger.info(f"找到包含message_id {message_id} 的checkpoint (change: {change.change_id})")
+                                    break
+                        except Exception:
+                            continue
+            
+            if not found_checkpoint:
+                logger.warning(f"未找到包含message_id {message_id} 的checkpoint")
+                return UndoResult(success=False, errors={"general": f"未找到包含message_id {message_id} 的checkpoint"}), None
+            
+            # 执行回滚
+            logger.info(f"开始回滚到message_id {message_id} 对应的checkpoint")
+            
+            # 获取当前所有的变更记录
+            all_changes = self.change_store.get_latest_changes(limit=1000)
+            
+            # 找到目标checkpoint对应的变更记录
+            target_timestamp = None
+            if found_group_id:
+                # 找到该组最后一个变更的时间戳
+                group_changes = self.change_store.get_changes_by_group(found_group_id)
+                if group_changes:
+                    target_timestamp = max(change.timestamp for change in group_changes)
+            elif found_change_id:
+                # 找到该变更的时间戳
+                for change in all_changes:
+                    if change.change_id == found_change_id:
+                        target_timestamp = change.timestamp
+                        break
+            
+            if not target_timestamp:
+                return UndoResult(success=False, errors={"general": "无法确定回滚目标时间点"}), found_checkpoint
+            
+            # 撤销所有在目标时间戳之后的变更
+            result = UndoResult(success=True)
+            changes_to_undo = []
+            
+            for change in all_changes:
+                if change.timestamp > target_timestamp:
+                    changes_to_undo.append(change)
+            
+            # 按时间戳降序排序，确保按照相反的顺序撤销
+            changes_to_undo.sort(key=lambda x: x.timestamp, reverse=True)
+            
+            logger.info(f"需要撤销 {len(changes_to_undo)} 个变更")
+            
+            # 逐个撤销变更
+            for change in changes_to_undo:
+                change_result = self.undo_change(change.change_id)
+                result.success = result.success and change_result.success
+                result.restored_files.extend(change_result.restored_files)
+                result.errors.update(change_result.errors)
+            
+            logger.info(f"回滚到message_id {message_id} 完成，恢复了 {len(result.restored_files)} 个文件")
+            
+            return result, found_checkpoint
+            
+        except Exception as e:
+            logger.error(f"回滚到message_id {message_id} 失败: {str(e)}")
+            return UndoResult(success=False, errors={"general": str(e)}), None
+    
     def get_diff_text(self, old_content: str, new_content: str) -> str:
         """
         获取两个文本内容的差异文本

autocoder/common/file_checkpoint/test_backup.py

@@ -40,24 +40,7 @@ class TestFileBackupManager:
         assert manager.backup_dir == temp_backup_dir
         assert os.path.exists(temp_backup_dir)
     
-    def test_init_with_default_dir(self, monkeypatch):
-        """测试使用默认目录初始化"""
-        # 创建一个临时主目录
-        temp_home = tempfile.mkdtemp()
-        try:
-            # 模拟用户主目录
-            monkeypatch.setattr(os.path, 'expanduser', lambda path: temp_home)
-            
-            manager = FileBackupManager()
-            
-            expected_dir = os.path.join(temp_home, ".autocoder", "backups")
-            assert manager.backup_dir == expected_dir
-            assert os.path.exists(expected_dir)
-        finally:
-            # 清理临时目录
-            if os.path.exists(temp_home):
-                shutil.rmtree(temp_home)
-    
+
     def test_backup_file(self, temp_backup_dir, sample_file):
         """测试备份文件功能"""
         manager = FileBackupManager(backup_dir=temp_backup_dir)