AirPlan-V2/lib/air_runtime/partial_replanner.py

"""
局部重规划 — P1-21 仅重新生成受 ADR 变更影响的任务子集。
替代全量重规划，保留未受影响任务的接口约束。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from air_runtime.task_graph import TaskGraph, TaskNode, PlanDelta


@dataclass
class Interface:
    """未受影响任务暴露的公共接口约束。"""
    task_id: str
    write_set: list[str] = field(default_factory=list)
    adr_refs: list[str] = field(default_factory=list)


@dataclass
class ReplanContext:
    """受影响任务的上下文信息，供 Arc 局部重规划使用。"""
    task_id: str
    task: str
    files_dirs: str
    done_when: str
    write_set: list[str] = field(default_factory=list)
    adr_refs: list[str] = field(default_factory=list)
    status: str = ""


class PartialReplanner:
    """仅重新生成受 ADR 变更影响的任务子集。

    与 incremental_replan_mode 的区别：
    - incremental_replan_mode: 全量重建 DAG 再 diff
    - PartialReplanner: 只对受影响部分重新规划，保留稳定接口约束
    """

    def replan(self, graph: TaskGraph, invalidated_ids: list[str],
               new_adr_path: Path | None = None) -> PlanDelta:
        """局部重规划：仅生成受影响任务的替代任务。

        Args:
            graph: 当前任务图（已包含 INVALIDATED 标记）
            invalidated_ids: 被 ADR 变更级联失效的任务 ID 列表
            new_adr_path: 新 ADR 文件路径（可选，供 Arc 参考）
        """
        delta = PlanDelta()

        # 1. 收集受影响任务的上下文
        affected_context = self._collect_affected_context(graph, invalidated_ids)

        # 2. 提取未受影响任务的稳定接口
        stable_interfaces = self._extract_stable_interfaces(graph, set(invalidated_ids))

        # 3. 生成局部重规划指令文件（供 Arc 读取）
        replan_request = {
            "type": "partial-replan",
            "invalidatedTaskIds": invalidated_ids,
            "affectedContext": [ctx.__dict__ for ctx in affected_context],
            "stableInterfaces": [iface.__dict__ for iface in stable_interfaces],
            "newAdrPath": str(new_adr_path) if new_adr_path else None,
        }

        # 4. 构建增量 delta
        # removed_tasks 已在 invalidate_by_adr 中填充
        # added_tasks 留空——由 Arc 读取 replan-request.json 后生成新任务
        delta.replan_request = replan_request

        return delta

    def replan_with_constraints(self, graph: TaskGraph,
                                 invalidated_ids: list[str],
                                 frozen_interfaces: list[Interface],
                                 new_adr_path: Path | None = None) -> PlanDelta:
        """带冻结接口约束的局部重规划（3.2.17 Phase 3）。

        frozen_interfaces 来自 SAFE 已完成任务，重规划不能破坏这些接口。
        """
        delta = PlanDelta()

        affected_context = self._collect_affected_context(graph, invalidated_ids)
        stable_interfaces = self._extract_stable_interfaces(graph, set(invalidated_ids))

        all_frozen = list(stable_interfaces)
        for fi in frozen_interfaces:
            if not any(s.task_id == fi.task_id for s in all_frozen):
                all_frozen.append(fi)

        replan_request = {
            "type": "partial-replan",
            "invalidatedTaskIds": invalidated_ids,
            "affectedContext": [ctx.__dict__ for ctx in affected_context],
            "stableInterfaces": [iface.__dict__ for iface in all_frozen],
            "newAdrPath": str(new_adr_path) if new_adr_path else None,
        }
        delta.replan_request = replan_request
        return delta

    def generate_verification_tasks(self, graph: TaskGraph,
                                     boundary_task_ids: list[str]) -> PlanDelta:
        """为 BOUNDARY 已完成任务生成验证任务（3.2.17 Phase 3.4）。

        验证任务检查 BOUNDARY 任务的兼容性，确认是否与新 ADR 一致。
        """
        from air_runtime.task_graph import TaskNode, Edge
        delta = PlanDelta()
        for idx, tid in enumerate(boundary_task_ids):
            node = graph.nodes.get(tid)
            if not node or node.status != "DONE":
                continue
            verify_id = f"{tid}-VFY"
            verify_node = TaskNode(
                id=verify_id,
                status="TODO",
                task=f"[边界验证] 验证 {tid} ({node.task}) 与新 ADR 的兼容性",
                files_dirs=node.files_dirs,
                done_when="兼容性确认通过",
                write_set=list(node.write_set),
                test_required=True,
            )
            delta.added_tasks.append(verify_node)
            delta.edge_changes.added.append(Edge(
                source=tid, target=verify_id, kind="verification",
            ))
        return delta

    def _collect_affected_context(self, graph: TaskGraph,
                                   invalidated_ids: list[str]) -> list[ReplanContext]:
        """收集受影响任务的上下文。"""
        contexts = []
        for tid in invalidated_ids:
            node = graph.nodes.get(tid)
            if node:
                contexts.append(ReplanContext(
                    task_id=node.id,
                    task=node.task,
                    files_dirs=node.files_dirs,
                    done_when=node.done_when,
                    write_set=list(node.write_set),
                    adr_refs=list(node.adr_refs),
                    status=node.status,
                ))
        return contexts

    def _extract_stable_interfaces(self, graph: TaskGraph,
                                    invalidated_ids: set) -> list[Interface]:
        """提取未受影响 DONE 任务的接口约束，确保重规划不破坏依赖。"""
        interfaces = []
        for nid, node in graph.nodes.items():
            if nid not in invalidated_ids and node.status == "DONE":
                if node.write_set or node.adr_refs:
                    interfaces.append(Interface(
                        task_id=node.id,
                        write_set=list(node.write_set),
                        adr_refs=list(node.adr_refs),
                    ))
        return interfaces