state_manager.py

"""
状态管理器 — 从 Orchestrator 中提取的状态持久化与恢复逻辑。

负责:
    - 任务运行时状态的保存与恢复（logs/orchestrator-state.json）
    - 崩溃恢复
    - 阶段状态跟踪（StageStatus）
    - 构建最终结果
"""

from __future__ import annotations

import json
import logging
import os
import tempfile
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# 状态数据类
# ---------------------------------------------------------------------------


@dataclass
class StageStatus:
    """单个阶段的执行状态。

    Attributes:
        status: 阶段状态（"pending" / "in_progress" / "completed" / "skipped"）。
        generator_retry_count: Generator 退回重试次数。
        planner_retry_count: Planner 退回重试次数。
        spec_version: 当前 spec 版本号。
        last_generator_output: Generator 上次输出（用于重试时附加实现层）。
        last_review_report: 上次审查报告（用于重试时附加反馈层）。
        history_summaries: 历史审查摘要列表。
    """

    status: str = "pending"  # pending / in_progress / completed / skipped
    generator_retry_count: int = 0
    planner_retry_count: int = 0
    spec_version: int = 1
    last_generator_output: str = ""
    last_review_report: str = ""
    history_summaries: list[str] = field(default_factory=list)
    generator_retry_reset_used: bool = False  # 每阶段最多重置 generator_retry_count 一次

    def to_dict(self) -> dict[str, Any]:
        """序列化为字典。"""
        return {
            "status": self.status,
            "generator_retry_count": self.generator_retry_count,
            "planner_retry_count": self.planner_retry_count,
            "spec_version": self.spec_version,
            "last_generator_output": self.last_generator_output,
            "last_review_report": self.last_review_report,
            "history_summaries": self.history_summaries,
            "generator_retry_reset_used": self.generator_retry_reset_used,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> StageStatus:
        """从字典反序列化。"""
        return cls(
            status=data.get("status", "pending"),
            generator_retry_count=data.get("generator_retry_count", 0),
            planner_retry_count=data.get("planner_retry_count", 0),
            spec_version=data.get("spec_version", 1),
            last_generator_output=data.get("last_generator_output", ""),
            last_review_report=data.get("last_review_report", ""),
            history_summaries=data.get("history_summaries", []),
            generator_retry_reset_used=data.get("generator_retry_reset_used", False),
        )


# ---------------------------------------------------------------------------
# StateManager
# ---------------------------------------------------------------------------


class StateManager:
    """状态管理器，负责任务运行时状态的持久化与恢复。

    从 Orchestrator 中提取的状态管理相关方法，保持原有逻辑不变。

    Args:
        project_root: 项目根目录路径。
        config: 配置字典。
    """

    def __init__(
        self,
        project_root: str,
        config: dict[str, Any],
    ) -> None:
        self._project_root: str = project_root
        self._config: dict[str, Any] = config

        # 从配置中读取参数
        self._token_budget: int = self._config.get("task_token_budget", 500000)
        self._max_history_summaries: int = self._config.get(
            "max_history_summaries", 3
        )

        # 状态文件路径
        logs_dir = os.path.join(project_root, "logs")
        os.makedirs(logs_dir, exist_ok=True)
        self._state_file_path: str = os.path.join(
            logs_dir, "orchestrator-state.json"
        )

        # 运行时状态
        self._task_id: str = ""
        self._current_stage: int = 0
        self._total_stages: int = 0
        self._stage_statuses: dict[str, StageStatus] = {}
        self._token_used: int = 0
        self._safe_mode: bool = False
        self._safe_mode_reason: str = ""
        self._safe_mode_entered_at: str = ""
        self._consecutive_failures: int = 0
        self._consecutive_successes: int = 0
        self._last_decision: str = ""
        self._created_at: str = ""
        self._updated_at: str = ""
        self._state_dirty: bool = False

        # 尝试从状态文件恢复
        self._recover_state()

        logger.info(
            "StateManager 初始化完成: project_root=%s, token_budget=%d, "
            "safe_mode=%s",
            project_root,
            self._token_budget,
            self._safe_mode,
        )

    # ------------------------------------------------------------------
    # 属性（property）
    # ------------------------------------------------------------------

    @property
    def task_id(self) -> str:
        return self._task_id

    @task_id.setter
    def task_id(self, value: str) -> None:
        self._task_id = value

    @property
    def current_stage(self) -> int:
        return self._current_stage

    @current_stage.setter
    def current_stage(self, value: int) -> None:
        self._current_stage = value

    @property
    def total_stages(self) -> int:
        return self._total_stages

    @total_stages.setter
    def total_stages(self, value: int) -> None:
        self._total_stages = value

    @property
    def stage_statuses(self) -> dict[str, StageStatus]:
        return self._stage_statuses

    @property
    def last_decision(self) -> str:
        return self._last_decision

    @last_decision.setter
    def last_decision(self, value: str) -> None:
        self._last_decision = value

    @property
    def safe_mode(self) -> bool:
        return self._safe_mode

    @safe_mode.setter
    def safe_mode(self, value: bool) -> None:
        self._safe_mode = value

    @property
    def safe_mode_reason(self) -> str:
        return self._safe_mode_reason

    @safe_mode_reason.setter
    def safe_mode_reason(self, value: str) -> None:
        self._safe_mode_reason = value

    @property
    def safe_mode_entered_at(self) -> str:
        return self._safe_mode_entered_at

    @safe_mode_entered_at.setter
    def safe_mode_entered_at(self, value: str) -> None:
        self._safe_mode_entered_at = value

    @property
    def token_used(self) -> int:
        return self._token_used

    @token_used.setter
    def token_used(self, value: int) -> None:
        self._token_used = value

    @property
    def token_budget(self) -> int:
        return self._token_budget

    @property
    def consecutive_failures(self) -> int:
        return self._consecutive_failures

    @consecutive_failures.setter
    def consecutive_failures(self, value: int) -> None:
        self._consecutive_failures = value

    @property
    def consecutive_successes(self) -> int:
        return self._consecutive_successes

    @consecutive_successes.setter
    def consecutive_successes(self, value: int) -> None:
        self._consecutive_successes = value

    @property
    def created_at(self) -> str:
        return self._created_at

    @created_at.setter
    def created_at(self, value: str) -> None:
        self._created_at = value

    @property
    def updated_at(self) -> str:
        return self._updated_at

    @property
    def state_dirty(self) -> bool:
        return self._state_dirty

    @state_dirty.setter
    def state_dirty(self, value: bool) -> None:
        self._state_dirty = value

    # ------------------------------------------------------------------
    # 状态恢复与持久化
    # ------------------------------------------------------------------

    def _recover_state(self) -> bool:
        """从状态文件恢复，返回是否成功恢复。

        恢复成功后，所有运行时状态变量将被还原为上次保存的值。

        Returns:
            True 表示成功恢复，False 表示无状态文件或恢复失败。
        """
        if not os.path.isfile(self._state_file_path):
            logger.info("未找到状态文件，将作为新任务启动")
            return False

        try:
            with open(self._state_file_path, "r", encoding="utf-8") as f:
                state: dict[str, Any] = json.load(f)

            self._task_id = state.get("task_id", "")
            self._current_stage = state.get("current_stage", 0)
            self._token_used = state.get("token_usage", {}).get("used", 0)
            self._safe_mode = state.get("safe_mode", False)
            self._safe_mode_reason = state.get("safe_mode_reason", "")
            self._safe_mode_entered_at = state.get("safe_mode_entered_at", "")
            self._consecutive_failures = state.get("consecutive_failures", 0)
            self._consecutive_successes = state.get("consecutive_successes", 0)
            self._last_decision = state.get("last_decision", "")
            self._created_at = state.get("created_at", "")
            self._updated_at = state.get("updated_at", "")

            # 恢复阶段状态
            stage_status_data = state.get("stage_status", {})
            for key, data in stage_status_data.items():
                self._stage_statuses[key] = StageStatus.from_dict(data)

            # 恢复总阶段数（从已完成的阶段中推断）
            completed_stages = [
                k for k, v in self._stage_statuses.items()
                if v.status == "completed"
            ]
            if completed_stages:
                max_stage_num = max(
                    int(k.split("-")[1]) for k in completed_stages
                )
                self._total_stages = max(self._total_stages, max_stage_num)

            logger.info(
                "状态恢复成功: task_id=%s, current_stage=%d, "
                "token_used=%d, safe_mode=%s",
                self._task_id,
                self._current_stage,
                self._token_used,
                self._safe_mode,
            )
            return True

        except (json.JSONDecodeError, KeyError, ValueError) as e:
            logger.error("状态文件恢复失败: %s，将作为新任务启动", e)
            return False

    def _save_state(self) -> None:
        """标记状态为脏（延迟写入）。

        仅设置脏标记，不立即写磁盘。实际写入由 flush_state() 触发，
        或在关键节点（阶段完成、任务结束等）由 flush_state() 强制执行。
        """
        self._state_dirty = True

    def flush_state(self) -> None:
        """强制将状态写入磁盘。

        设置脏标记并执行实际的磁盘写入。用于关键节点如阶段完成、
        任务结束、安全模式进入、人类介入决策后等场景。
        """
        if not self._state_dirty:
            return

        self._updated_at = self._now_iso()
        self._state_dirty = False

        state: dict[str, Any] = {
            "task_id": self._task_id,
            "current_stage": self._current_stage,
            "stage_status": {
                key: value.to_dict()
                for key, value in self._stage_statuses.items()
            },
            "token_usage": {
                "used": self._token_used,
                "budget": self._token_budget,
            },
            "safe_mode": self._safe_mode,
            "safe_mode_reason": self._safe_mode_reason,
            "safe_mode_entered_at": self._safe_mode_entered_at,
            "consecutive_failures": self._consecutive_failures,
            "consecutive_successes": self._consecutive_successes,
            "last_decision": self._last_decision,
            "created_at": self._created_at,
            "updated_at": self._updated_at,
            # T31 — 持续学习：记录各阶段历史摘要数量
            "process_docs": {
                key: {
                    "history_summaries_count": len(value.history_summaries),
                    "max_history_summaries": self._max_history_summaries,
                }
                for key, value in self._stage_statuses.items()
            },
        }

        try:
            # 原子写入：先写临时文件，再重命名
            fd, tmp_path = tempfile.mkstemp(
                suffix=".tmp",
                dir=os.path.dirname(self._state_file_path),
            )
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                json.dump(state, f, ensure_ascii=False, indent=2)
            os.replace(tmp_path, self._state_file_path)

            logger.debug("状态文件已保存: %s", self._state_file_path)

        except OSError as e:
            logger.error("状态文件保存失败: %s", e)

    def flush(self) -> None:
        """flush_state 的简写别名，供 Orchestrator 调用。"""
        self.flush_state()

    def get_stage_status(self, stage_key: str) -> StageStatus:
        """获取或创建阶段状态（公开接口）。"""
        return self._get_stage_status(stage_key)

    def build_result(self, status: str, detail: dict[str, Any]) -> dict[str, Any]:
        """构建最终返回结果（公开接口）。"""
        return self._build_result(status, detail)

    def mark_dirty(self) -> None:
        """标记状态为脏（公开接口）。"""
        self._state_dirty = True

    def save_state(self) -> None:
        """保存状态到磁盘（公开接口）。"""
        self._save_state()

    # ------------------------------------------------------------------
    # 阶段状态
    # ------------------------------------------------------------------

    def _get_stage_status(self, stage_key: str) -> StageStatus:
        """获取或创建阶段状态。

        Args:
            stage_key: 阶段标识（如 "stage-1"）。

        Returns:
            对应的 StageStatus 实例。
        """
        if stage_key not in self._stage_statuses:
            self._stage_statuses[stage_key] = StageStatus()
        return self._stage_statuses[stage_key]

    # ------------------------------------------------------------------
    # 结果构建
    # ------------------------------------------------------------------

    def _build_result(
        self,
        status: str,
        detail: dict[str, Any],
    ) -> dict[str, Any]:
        """构建最终返回结果。

        Args:
            status: 最终状态（"completed" / "aborted" / "safe_mode"）。
            detail: 附加详情。

        Returns:
            结构化的结果字典。
        """
        completed_stages = [
            key for key, value in self._stage_statuses.items()
            if value.status == "completed"
        ]

        result: dict[str, Any] = {
            "task_id": self._task_id,
            "status": status,
            "total_stages": self._total_stages,
            "completed_stages": completed_stages,
            "token_usage": {
                "used": self._token_used,
                "budget": self._token_budget,
            },
            "safe_mode": self._safe_mode,
            "last_decision": self._last_decision,
        }
        result.update(detail)

        logger.info(
            "任务结束: task_id=%s, status=%s, completed_stages=%d/%d, "
            "token_used=%d",
            self._task_id,
            status,
            len(completed_stages),
            self._total_stages,
            self._token_used,
        )

        return result

    # ------------------------------------------------------------------
    # 工具方法
    # ------------------------------------------------------------------

    @staticmethod
    def _now_iso() -> str:
        """获取当前时间的 ISO 8601 格式字符串（UTC）。"""
        return datetime.now(timezone.utc).isoformat()

    @staticmethod
    def now_iso() -> str:
        """获取当前时间的 ISO 8601 格式字符串（UTC）— 公开接口。"""
        return StateManager._now_iso()