video-image-production-build/tests/agents/test_planner.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from langchain_core.messages import AIMessage
from pydantic import ValidationError

from production_build_agents.preprocess.brief_io import load_production_brief
from production_build_agents.capabilities import (
    GLOBAL_DATA_SKILL_IDS,
    global_data_skill_catalog,
)
from production_build_agents.contracts.models import (
    ArtifactExpectation,
    ArtifactRejection,
    GlobalDataPlan,
    SourceAsset,
)
from production_build_agents.contracts.identifiers import source_asset_id_for
from production_build_agents.agents.planner.agent import (
    PlanContractValidationError,
    PlannerOutputError,
    _failed_task_context,
    _planner_correction_content,
    _planner_failure_message,
    materialize_task_package,
    ready_tasks,
    run_planner_agent,
    run_replan_agent,
    validate_plan_runtime,
    validate_replan,
)
from tests.support.planner_fixtures import (
    GLOBAL_DATA_AUDIT_PATHS,
    build_global_data_plan,
    build_planned_task,
    build_planner_model,
    build_task_package,
)
from tests.support.fake_models import ToolAwareFakeChatModel


SOURCE_BRIEF_PATH = (
    Path(__file__).parents[1] / "fixtures" / "minimal_brief.json"
)
_BRIEF_FIXTURE_DIR = tempfile.TemporaryDirectory(
    prefix="production-build-planner-brief-"
)
BRIEF_PATH = Path(_BRIEF_FIXTURE_DIR.name) / "minimal_brief.json"
_brief_payload = json.loads(SOURCE_BRIEF_PATH.read_text(encoding="utf-8"))
_brief_payload["schema_version"] = "0.3"
BRIEF_PATH.write_text(
    json.dumps(_brief_payload, ensure_ascii=False, indent=2),
    encoding="utf-8",
)
PLAN_PATH = Path(_BRIEF_FIXTURE_DIR.name) / "plan.json"


class PlannerDagTest(unittest.TestCase):
    def test_planner_core_is_generic_and_global_rules_live_in_skill(self) -> None:
        root = (
            Path(__file__).parents[2]
            / "production_build_agents"
            / "agents"
            / "planner"
        )
        core = (root / "prompt.md").read_text(encoding="utf-8")
        skill = (
            root / "skills" / "global-data-planning" / "SKILL.md"
        ).read_text(encoding="utf-8")
        self.assertNotIn("Global Data 完成", core)
        self.assertNotIn("reference-inspection", core)
        self.assertIn("Global Data 完成", skill)
        self.assertIn("reference-inspection", skill)

    def test_skill_catalog_contains_planner_facing_summary(self) -> None:
        catalog = global_data_skill_catalog()
        self.assertEqual(
            [item["skill_id"] for item in catalog],
            list(GLOBAL_DATA_SKILL_IDS),
        )
        self.assertTrue(all(item["summary"] for item in catalog))
        self.assertTrue(
            all("deliverable_types" in item for item in catalog)
        )
        self.assertTrue(
            all("allowed_tools" not in item for item in catalog)
        )
        self.assertTrue(
            all("allowed_remote_tool_ids" not in item for item in catalog)
        )

    def test_planner_can_create_one_free_form_task(self) -> None:
        expected = build_global_data_plan()
        plan = run_planner_agent(
            str(BRIEF_PATH),
            plan_version=1,
            run_id="Run-planner-one",
            model=build_planner_model(expected),
        )
        self.assertEqual(plan, expected)
        self.assertEqual(len(plan.tasks), 1)
        self.assertEqual(plan.tasks[0].skill_id, "structured-analysis")

    def test_planner_can_read_brief_on_demand_before_final_plan(self) -> None:
        expected = build_global_data_plan()
        model = ToolAwareFakeChatModel(
            responses=[
                AIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "read_production_brief",
                            "args": {
                                "source_paths": GLOBAL_DATA_AUDIT_PATHS
                            },
                            "id": "read-1",
                            "type": "tool_call",
                        }
                    ],
                ),
                AIMessage(content=expected.model_dump_json()),
            ]
        )
        plan = run_planner_agent(
            str(BRIEF_PATH),
            plan_version=1,
            run_id="Run-planner-read",
            model=model,
        )
        self.assertIn("$.核心制作点[0]", plan.tasks[0].source_paths)

    def test_planner_only_requires_shared_global_data_reads(self) -> None:
        plan = build_global_data_plan()
        model = build_planner_model(
            plan,
            read_paths=GLOBAL_DATA_AUDIT_PATHS,
        )

        actual = run_planner_agent(
            str(BRIEF_PATH),
            plan_version=1,
            run_id="Run-planner-shared-audit",
            model=model,
        )

        self.assertEqual(actual, plan)

    def test_multiple_tasks_and_dependencies_are_preserved(self) -> None:
        tasks = [
            build_planned_task(task_id="Task1", priority=20),
            build_planned_task(
                task_id="Task2",
                objective="生成视觉参考图",
                skill_id="image-production",
                deliverable_type="image",
                priority=10,
            ),
            build_planned_task(
                task_id="Task3",
                objective="整理前两项资料",
                depends_on=["Task1", "Task2"],
            ),
        ]
        plan = run_planner_agent(
            str(BRIEF_PATH),
            plan_version=1,
            run_id="Run-planner-many",
            model=build_planner_model(build_global_data_plan(tasks=tasks)),
        )
        self.assertEqual(ready_tasks(plan)[0].task_id, "Task2")

    def test_cycle_is_rejected_by_contract(self) -> None:
        payload = {
            "plan_id": "GlobalDataPlan",
            "plan_version": 1,
            "goal": "准备资料",
            "revision_summary": "循环图应被拒绝",
            "stage_requirements": [
                item.model_dump(mode="json")
                for item in build_global_data_plan().stage_requirements
            ],
            "tasks": [
                build_planned_task(
                    task_id="Task1",
                    depends_on=["Task2"],
                ).model_dump(mode="json"),
                build_planned_task(
                    task_id="Task2",
                    depends_on=["Task1"],
                ).model_dump(mode="json"),
            ],
        }
        cyclic = GlobalDataPlan.model_validate(payload)
        _, brief = load_production_brief(BRIEF_PATH)
        with self.assertRaisesRegex(PlannerOutputError, "循环依赖"):
            validate_plan_runtime(
                cyclic,
                brief=brief,
                plan_version=1,
                max_tasks=8,
            )

    def test_unknown_skill_is_rejected(self) -> None:
        payload = build_global_data_plan().model_dump(mode="json")
        payload["tasks"][0]["skill_id"] = "unknown-skill"
        model = ToolAwareFakeChatModel(
            responses=[AIMessage(content=json.dumps(payload, ensure_ascii=False))]
        )
        with self.assertRaisesRegex(PlannerOutputError, "未注册"):
            run_planner_agent(
                str(BRIEF_PATH),
                plan_version=1,
                run_id="Run-planner-unknown-skill",
                model=model,
            )

    def test_unknown_source_path_is_rejected(self) -> None:
        payload = build_global_data_plan().model_dump(mode="json")
        payload["tasks"][0]["source_paths"] = ["$.不存在"]
        payload["stage_requirements"][0]["source_paths"] = ["$.不存在"]
        model = ToolAwareFakeChatModel(
            responses=[AIMessage(content=json.dumps(payload, ensure_ascii=False))]
        )
        with self.assertRaisesRegex(PlannerOutputError, "不存在的来源路径"):
            run_planner_agent(
                str(BRIEF_PATH),
                plan_version=1,
                run_id="Run-planner-unknown-source",
                model=model,
            )

    def test_all_global_audit_paths_require_critical_disposition(
        self,
    ) -> None:
        plan = build_global_data_plan()
        reduced_task = plan.tasks[0].model_copy(
            update={"source_paths": ["$.核心制作点[0]"]}
        )
        reduced_requirement = plan.stage_requirements[0].model_copy(
            update={"source_paths": ["$.核心制作点[0]"]}
        )
        reduced = plan.model_copy(
            update={
                "tasks": [reduced_task],
                "stage_requirements": [reduced_requirement],
            }
        )
        _, brief = load_production_brief(BRIEF_PATH)
        with self.assertRaisesRegex(
            PlannerOutputError,
            "missing_critical_audit_path",
        ):
            validate_plan_runtime(
                reduced,
                brief=brief,
                plan_version=1,
                max_tasks=8,
            )

    def test_source_asset_issue_keeps_exact_identity_for_correction(
        self,
    ) -> None:
        _, brief = load_production_brief(BRIEF_PATH)
        source_uri = "https://example.test/reference/skeleton.mp4"
        source_asset = SourceAsset(
            source_asset_id=source_asset_id_for("video", source_uri),
            artifact_type="video",
            source_uri=source_uri,
            source_paths=["$.核心制作点[0]"],
        )
        brief = brief.model_copy(
            update={"source_assets": [source_asset]}
        )
        plan = build_global_data_plan()

        with self.assertRaises(
            PlanContractValidationError
        ) as caught:
            validate_plan_runtime(
                plan,
                brief=brief,
                plan_version=1,
                max_tasks=8,
            )

        missing = next(
            item
            for item in caught.exception.structured_errors
            if item["code"] == "undispositioned_source_asset"
        )
        self.assertEqual(
            missing,
            {
                "code": "undispositioned_source_asset",
                "message": (
                    "Brief 中的 SourceAsset 未被 critical "
                    "source_identity Expectation 处置"
                ),
                "source_asset_id": source_asset.source_asset_id,
                "artifact_type": "video",
                "source_uri": source_uri,
                "source_paths": ["$.核心制作点[0]"],
            },
        )

        correction = _planner_correction_content(
            caught.exception.structured_errors,
            plan_version=1,
        )
        self.assertIn(source_asset.source_asset_id, correction)
        self.assertIn('"artifact_type": "video"', correction)
        self.assertIn(
            '"source_paths": ["$.核心制作点[0]"]',
            correction,
        )
        self.assertIn(
            "对于 undispositioned_source_asset",
            correction,
        )

    def test_final_planner_failure_separates_attempts(self) -> None:
        first = [
            {
                "code": "source_asset_count_mismatch",
                "expectation_id": "Requirement3-Expectation1",
            }
        ]
        final = [
            {
                "code": "undispositioned_source_asset",
                "source_asset_id": "SourceAsset-" + "a" * 64,
            }
        ]

        message = _planner_failure_message([first, final])

        self.assertIn("final_issues=", message)
        self.assertIn("previous_attempt_issues=", message)
        final_section, previous_section = message.split(
            "；previous_attempt_issues=",
            maxsplit=1,
        )
        self.assertIn("undispositioned_source_asset", final_section)
        self.assertNotIn("source_asset_count_mismatch", final_section)
        self.assertIn("source_asset_count_mismatch", previous_section)

    def test_scope_issue_exposes_both_path_sets_for_correction(
        self,
    ) -> None:
        _, brief = load_production_brief(BRIEF_PATH)
        source_uri = "https://example.test/reference/curtain.png"
        source_asset = SourceAsset(
            source_asset_id=source_asset_id_for("image", source_uri),
            artifact_type="image",
            source_uri=source_uri,
            source_paths=["$.核心制作点[1]"],
        )
        brief = brief.model_copy(update={"source_assets": [source_asset]})
        plan = build_global_data_plan()
        requirement = plan.stage_requirements[0]
        expectation = requirement.artifact_expectations[0].model_copy(
            update={
                "artifact_type": source_asset.artifact_type,
                "minimum_count": 1,
                "verification_capabilities": [
                    "source_identity",
                    "technical_integrity",
                ],
                "source_asset_ids": [source_asset.source_asset_id],
            }
        )
        requirement = requirement.model_copy(
            update={
                "source_paths": ["$.核心制作点[0]"],
                "artifact_expectations": [expectation],
            }
        )
        task = plan.tasks[0].model_copy(
            update={
                "source_paths": ["$.核心制作点[0]"],
                "skill_id": "reference-inspection",
                "deliverable_type": "reference_collection",
            }
        )
        plan = plan.model_copy(
            update={
                "stage_requirements": [requirement],
                "tasks": [task],
            }
        )

        with self.assertRaises(
            PlanContractValidationError
        ) as caught:
            validate_plan_runtime(
                plan,
                brief=brief,
                plan_version=1,
                max_tasks=8,
            )

        scope_issue = next(
            item
            for item in caught.exception.structured_errors
            if item["code"]
            == "source_asset_outside_requirement_scope"
        )
        self.assertEqual(
            scope_issue["requirement_source_paths"],
            ["$.核心制作点[0]"],
        )
        self.assertEqual(
            scope_issue["source_asset_source_paths"],
            ["$.核心制作点[1]"],
        )
        correction = _planner_correction_content(
            [scope_issue],
            plan_version=1,
        )
        self.assertIn(
            "source_asset_outside_requirement_scope",
            correction,
        )
        self.assertIn(
            "requirement_source_paths",
            correction,
        )
        self.assertIn(
            "source_asset_source_paths",
            correction,
        )

    def test_structured_analysis_cannot_cover_required_image(self) -> None:
        plan = build_global_data_plan()
        payload = plan.model_dump(mode="json")
        payload["stage_requirements"][0]["artifact_expectations"] = [
            {
                "expectation_id": "Requirement1-Expectation1",
                "artifact_type": "image",
                "minimum_count": 1,
                "usage_scope": "人物一致性",
                "verification_capabilities": ["visual_content"],
            }
        ]
        incompatible = GlobalDataPlan.model_validate(payload)
        _, brief = load_production_brief(BRIEF_PATH)
        with self.assertRaisesRegex(
            PlannerOutputError,
            "无法交付",
        ):
            validate_plan_runtime(
                incompatible,
                brief=brief,
                plan_version=1,
                max_tasks=8,
            )

    def test_audio_reference_plan_repairs_subjective_acceptance(
        self,
    ) -> None:
        task = build_planned_task(
            skill_id="reference-inspection",
            deliverable_type="reference_collection",
        )
        invalid = build_global_data_plan(tasks=[task])
        valid = build_global_data_plan(tasks=[task])
        invalid = invalid.model_copy(
            update={
                "stage_requirements": [
                    invalid.stage_requirements[0].model_copy(
                        update={
                            "description": (
                                "确认音频为治愈系无人声器乐并可正常使用"
                            ),
                            "artifact_expectations": [
                                ArtifactExpectation(
                                    expectation_id=(
                                        "Requirement1-Expectation1"
                                    ),
                                    artifact_type="audio",
                                    minimum_count=1,
                                    usage_scope="全片人声基准",
                                    verification_capabilities=[
                                        "semantic_content",
                                    ],
                                )
                            ]
                        }
                    )
                ]
            }
        )
        valid_requirement = invalid.stage_requirements[0].model_copy(
            update={
                "description": (
                    "采纳制作表指定音频，并确认媒体类型和技术属性可用"
                ),
                "artifact_expectations": [
                    invalid.stage_requirements[
                        0
                    ].artifact_expectations[0].model_copy(
                        update={
                    "verification_capabilities": [
                        "technical_integrity",
                            ]
                        }
                    )
                ],
            }
        )
        valid = valid.model_copy(
            update={
                "stage_requirements": [valid_requirement],
            }
        )

        plan = run_planner_agent(
            str(BRIEF_PATH),
            plan_version=1,
            run_id="Run-planner-audio-repair",
            model=build_planner_model(invalid, valid),
        )

        self.assertEqual(plan, valid)

    def test_requirement_ids_must_be_unique_and_continuous(self) -> None:
        payload = build_global_data_plan().model_dump(mode="json")
        requirement = dict(payload["stage_requirements"][0])
        requirement["requirement_id"] = "Requirement3"
        payload["stage_requirements"].append(requirement)
        with self.assertRaisesRegex(ValidationError, "连续编号"):
            GlobalDataPlan.model_validate(payload)

    def test_initial_plan_cannot_replace_artifact(self) -> None:
        plan = build_global_data_plan()
        altered = plan.model_copy(
            update={
                "tasks": [
                    plan.tasks[0].model_copy(
                        update={
                            "replaces_artifact_ids": [
                                "Task1-v1-artifact-1"
                            ]
                        }
                    )
                ]
            }
        )
        _, brief = load_production_brief(BRIEF_PATH)
        with self.assertRaisesRegex(PlannerOutputError, "初始"):
            validate_plan_runtime(
                altered,
                brief=brief,
                plan_version=1,
                max_tasks=8,
            )

    def test_task_budget_is_enforced_without_fixing_business_categories(self) -> None:
        plan = build_global_data_plan(
            tasks=[
                build_planned_task(task_id=f"Task{index}")
                for index in range(1, 4)
            ]
        )
        with self.assertRaisesRegex(PlannerOutputError, "超过预算"):
            run_planner_agent(
                str(BRIEF_PATH),
                plan_version=1,
                run_id="Run-planner-budget",
                max_tasks=2,
                model=build_planner_model(plan),
            )

    def test_ready_node_materializes_complete_task_package(self) -> None:
        plan = build_global_data_plan()
        task = materialize_task_package(
            plan,
            ready_tasks(plan)[0],
            run_id="Run-materialize",
            production_brief_path=BRIEF_PATH,
            plan_path=PLAN_PATH,
        )
        self.assertEqual(task.schema_version, "0.3")
        self.assertEqual(task.task_id, "Task1")
        self.assertEqual(task.run_id, "Run-materialize")
        self.assertEqual(task.plan_id, plan.plan_id)
        self.assertEqual(task.plan_version, plan.plan_version)
        self.assertEqual(
            task.production_brief_uri,
            str(BRIEF_PATH.resolve()),
        )
        self.assertEqual(task.plan_uri, str(PLAN_PATH.resolve()))
        self.assertEqual(task.dependency_deliveries, [])
        self.assertEqual(
            set(task.model_dump()),
            {
                "schema_version",
                "run_id",
                "plan_id",
                "plan_version",
                "task_id",
                "production_brief_uri",
                "plan_uri",
                "dependency_deliveries",
            },
        )

    def test_dependency_uses_upstream_delivery_manifest(self) -> None:
        planned = build_planned_task(
            task_id="Task2",
            depends_on=["Task1"],
        )
        manifest = "/tmp/Task1.executor-delivery.json"
        plan = build_global_data_plan(
            tasks=[
                build_planned_task(task_id="Task1"),
                planned,
            ]
        )
        task = materialize_task_package(
            plan,
            planned,
            run_id="Run-dependency",
            production_brief_path=BRIEF_PATH,
            plan_path=PLAN_PATH,
            dependency_artifacts={"Task1": manifest},
        )
        self.assertEqual(len(task.dependency_deliveries), 1)
        dependency = task.dependency_deliveries[0]
        self.assertEqual(dependency.task_id, "Task1")
        self.assertEqual(dependency.delivery_uri, manifest)

    def test_replan_requires_next_version_and_immutable_passed_tasks(self) -> None:
        task1 = build_planned_task(task_id="Task1")
        previous = build_global_data_plan(tasks=[task1], plan_version=1)
        valid = build_global_data_plan(tasks=[task1], plan_version=2)
        validate_replan(previous, valid, passed_task_ids={"Task1"})

        changed = build_global_data_plan(
            tasks=[
                task1.model_copy(update={"objective": "修改已经通过的任务"})
            ],
            plan_version=2,
        )
        with self.assertRaisesRegex(PlannerOutputError, "不得修改已 PASS"):
            validate_replan(
                previous,
                changed,
                passed_task_ids={"Task1"},
            )

        skipped = build_global_data_plan(tasks=[task1], plan_version=3)
        with self.assertRaisesRegex(PlannerOutputError, "只增加一个版本"):
            validate_replan(
                previous,
                skipped,
                passed_task_ids=set(),
            )

    def test_replan_contract_error_is_corrected_in_same_planner_run(
        self,
    ) -> None:
        task1 = build_planned_task(task_id="Task1")
        previous = build_global_data_plan(
            tasks=[task1],
            plan_version=1,
        )
        invalid = build_global_data_plan(
            tasks=[
                task1.model_copy(
                    update={"objective": "错误修改已经通过的任务"}
                )
            ],
            plan_version=2,
        )
        corrected = build_global_data_plan(
            tasks=[task1],
            plan_version=2,
        )

        actual = run_replan_agent(
            str(BRIEF_PATH),
            previous,
            run_id="Run-replan-contract-correction",
            failure_scope="global_data_stage",
            failed_task=None,
            validation_report={"verdict": "FAIL"},
            passed_task_ids={"Task1"},
            remaining_replans=4,
            model=build_planner_model(invalid, corrected),
        )

        self.assertEqual(actual, corrected)

    def test_failed_task_context_is_resolved_from_previous_plan(
        self,
    ) -> None:
        previous = build_global_data_plan(plan_version=1)
        package = build_task_package(
            BRIEF_PATH,
            plan_version=1,
        )

        context = _failed_task_context(previous, package)

        self.assertIsNotNone(context)
        assert context is not None
        self.assertEqual(
            context["planned_task"],
            previous.tasks[0].model_dump(mode="json"),
        )
        self.assertEqual(
            context["selected_expectations"][0]["expectation"],
            previous.stage_requirements[
                0
            ].artifact_expectations[0].model_dump(mode="json"),
        )
        self.assertNotIn(
            "artifact_expectations",
            context["selected_expectations"][0]["requirement"],
        )

    def test_replan_cannot_weaken_critical_acceptance(self) -> None:
        task = build_planned_task()
        previous = build_global_data_plan(
            tasks=[task],
            plan_version=1,
        )
        weakened_requirement = previous.stage_requirements[0].model_copy(
            update={
                "artifact_expectations": [
                    previous.stage_requirements[
                        0
                    ].artifact_expectations[0].model_copy(
                        update={
                            "verification_capabilities": [
                                "source_identity"
                            ]
                        }
                    )
                ]
            }
        )
        new_plan = previous.model_copy(
            update={
                "plan_version": 2,
                "stage_requirements": [weakened_requirement],
            }
        )

        with self.assertRaisesRegex(
            PlannerOutputError,
            "不得弱化",
        ):
            validate_replan(
                previous,
                new_plan,
                passed_task_ids=set(),
            )

    def test_replan_can_strengthen_unpassed_critical_requirement(
        self,
    ) -> None:
        previous = build_global_data_plan(plan_version=1)
        expectation = previous.stage_requirements[
            0
        ].artifact_expectations[0]
        strengthened_requirement = previous.stage_requirements[0].model_copy(
            update={
                "artifact_expectations": [
                    expectation.model_copy(
                        update={"minimum_count": 2}
                    )
                ]
            }
        )
        strengthened = previous.model_copy(
            update={
                "plan_version": 2,
                "stage_requirements": [strengthened_requirement],
            }
        )

        validate_replan(
            previous,
            strengthened,
            passed_task_ids=set(),
        )

    def test_requirement_can_repeat_same_artifact_type_by_identity(self) -> None:
        payload = build_global_data_plan().model_dump(mode="json")
        second = dict(
            payload["stage_requirements"][0]["artifact_expectations"][0]
        )
        second["expectation_id"] = "Requirement1-Expectation2"
        second["usage_scope"] = "另一独立用途"
        payload["stage_requirements"][0][
            "artifact_expectations"
        ].append(second)
        payload["tasks"][0]["expectation_ids"].append(
            "Requirement1-Expectation2"
        )
        plan = GlobalDataPlan.model_validate(payload)
        self.assertEqual(
            [item.expectation_id for item in plan.stage_requirements[0].artifact_expectations],
            ["Requirement1-Expectation1", "Requirement1-Expectation2"],
        )

    def test_replan_must_keep_all_required_artifact_replacements(
        self,
    ) -> None:
        previous = build_global_data_plan()
        previous = previous.model_copy(
            update={
                "tasks": [
                    previous.tasks[0].model_copy(
                        update={
                            "replaces_artifact_ids": [
                                "Task1-v1-artifact-1"
                            ]
                        }
                    )
                ]
            }
        )
        current = previous.model_copy(
            update={
                "plan_version": 2,
                "tasks": [
                    previous.tasks[0].model_copy(
                        update={"replaces_artifact_ids": []}
                    )
                ],
            }
        )

        with self.assertRaisesRegex(
            PlannerOutputError,
            "必须保留历史替换",
        ):
            validate_replan(
                previous,
                current,
                passed_task_ids=[],
                artifact_rejections=[
                    ArtifactRejection(
                        artifact_id="Task9-v1-artifact-1",
                        reason_code="semantic_rejection",
                        reason="Stage 语义验收未通过",
                    )
                ],
            )


if __name__ == "__main__":
    unittest.main()