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 _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