video-image-production-build/tests/production/test_segment_acceptance.py

from __future__ import annotations

import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from production_build_agents.contracts.models import (
    Artifact,
    CandidateArtifact,
    GlobalDataStageDelivery,
)
from production_build_agents.contracts.document_identity import (
    exact_document_sha256,
)
from production_build_agents.contracts.production_execution_models import (
    AcceptedSegmentRef,
    ImmutableJsonRef,
    SegmentDependencyRef,
)
from production_build_agents.contracts.production_planning_models import (
    ProductionPlanningPackage,
)
from production_build_agents.run.artifacts import (
    ArtifactIntegrityError,
    content_identity_for_path,
    file_sha256,
    seal_shared_visual_anchor,
)
from production_build_agents.run.layout import (
    segment_delivery_path,
    segment_executor_candidate_path,
    segment_package_path,
    segment_validation_report_path,
    segment_validator_candidate_path,
)
from production_build_agents.run.records import write_once_json
from production_build_agents.run.segment_inputs import (
    reusable_accepted_segment_snapshot,
)
from production_build_agents.production.segment_acceptance import (
    build_accepted_segment_ref,
    segment_reuse_inputs_compatible,
)
from tests.support.production_fixtures import (
    production_plan,
    segment_bundle,
)
from tests.support.production_planner_fixtures import (
    build_production_planner_scenario,
)


def _persist_pass_bundle(
    root: Path,
    *,
    include_unselected_actual_input: bool = False,
    accepted_dependency: AcceptedSegmentRef | None = None,
    segment_index: int = 1,
    run_id: str = "source-run",
):
    root.mkdir(parents=True, exist_ok=True)
    (
        package,
        candidate,
        delivery,
        validator_candidate,
        report,
        _,
    ) = segment_bundle(segment_index, run_id=run_id)
    run_dir = root / run_id
    source_path = root / "source.png"
    source_path.write_bytes(b"source")
    old_source_uri = package.artifacts[0].uri
    source_artifact = package.artifacts[0].model_copy(
        update={
            "uri": str(source_path.resolve()),
            **content_identity_for_path(source_path),
        }
    )
    old_anchor = package.shared_visual_anchor
    anchor_path = run_dir / "production_inputs" / "shared_visual_anchor.png"
    anchor_path.parent.mkdir(parents=True, exist_ok=True)
    anchor_path.write_bytes(b"shared anchor")
    shared_anchor = seal_shared_visual_anchor(anchor_path)
    package = package.model_copy(
        update={
            "artifacts": [source_artifact],
            "shared_visual_anchor": shared_anchor,
        }
    )
    global_artifacts = [source_artifact]
    optional_artifact = None
    if include_unselected_actual_input:
        optional_path = root / "optional-audio.wav"
        optional_path.write_bytes(b"optional-audio")
        optional_artifact = Artifact(
            artifact_id="Task2-v1-artifact-1",
            artifact_type="audio",
            uri=str(optional_path.resolve()),
            source_uri=None,
            description="Planner 未选择但 Executor 实际使用的素材",
            **content_identity_for_path(optional_path),
        )
        global_artifacts.append(optional_artifact)
    global_delivery = build_production_planner_scenario(
        run_id="source-global"
    ).global_data_delivery.model_copy(
        update={"active_artifacts": global_artifacts}
    )
    global_delivery_path = root / "global_data_stage_delivery.json"
    global_delivery_path.write_text(
        global_delivery.model_dump_json(),
        encoding="utf-8",
    )
    package = package.model_copy(
        update={
            "global_data_delivery_ref": ImmutableJsonRef(
                uri=str(global_delivery_path.resolve()),
                document_sha256=exact_document_sha256(
                    global_delivery_path.read_bytes()
                ),
            )
        }
    )
    if accepted_dependency is not None:
        package = package.model_copy(
            update={
                "dependencies": [
                    SegmentDependencyRef(
                        segment_id=accepted_dependency.segment_id,
                        segment_delivery_ref=(
                            accepted_dependency.segment_delivery_ref
                        ),
                        validation_report_ref=(
                            accepted_dependency.validation_report_ref
                        ),
                    )
                ]
            }
        )

    old_to_new = {
        old_source_uri: str(source_path.resolve()),
        old_anchor.uri: shared_anchor.uri,
    }
    formal_artifacts = []
    for artifact in delivery.artifacts:
        suffix = Path(artifact.uri).suffix
        path = run_dir / "formal_artifacts" / f"{artifact.artifact_id}{suffix}"
        path.parent.mkdir(parents=True, exist_ok=True)
        if suffix == ".ass":
            path.write_text("subtitle", encoding="utf-8")
        else:
            path.write_bytes(f"artifact:{artifact.artifact_id}".encode())
        old_to_new[artifact.uri] = str(path.resolve())
        formal_artifacts.append(
            artifact.model_copy(
                update={
                    "uri": str(path.resolve()),
                    **content_identity_for_path(path),
                }
            )
        )

    candidate = candidate.model_copy(
        update={
            "artifacts": [
                CandidateArtifact(
                    artifact_type=item.artifact_type,
                    uri=old_to_new[item.uri],
                    source_uri=item.source_uri,
                    description=item.description,
                )
                for item in candidate.artifacts
            ],
            "lineage_claims": [
                claim.model_copy(
                    update={
                        "artifact_uri": old_to_new[claim.artifact_uri],
                        "input_artifact_uris": [
                            old_to_new.get(uri, uri)
                            for uri in claim.input_artifact_uris
                        ]
                        + (
                            [optional_artifact.uri]
                            if (
                                optional_artifact is not None
                                and index == 0
                            )
                            else []
                        )
                        + (
                            [accepted_dependency.primary_artifact.uri]
                            if accepted_dependency is not None and index == 0
                            else []
                        ),
                    }
                )
                for index, claim in enumerate(candidate.lineage_claims)
            ],
            "shot_outputs": [
                claim.model_copy(
                    update={
                        "artifact_uri": old_to_new[claim.artifact_uri]
                    }
                )
                for claim in candidate.shot_outputs
            ],
            "narration_artifact_uri": old_to_new[
                candidate.narration_artifact_uri
            ],
            "subtitle_artifact_uri": old_to_new[
                candidate.subtitle_artifact_uri
            ],
            "primary_artifact_uri": old_to_new[
                candidate.primary_artifact_uri
            ],
        }
    )
    package_path = write_once_json(
        segment_package_path(
            run_dir,
            package.segment_id,
            package.plan_version,
        ),
        package.model_dump(mode="json"),
    )
    candidate_path = write_once_json(
        segment_executor_candidate_path(
            run_dir,
            package.segment_id,
            package.plan_version,
        ),
        candidate.model_dump(mode="json"),
    )
    delivery = delivery.model_copy(
        update={
            "segment_package_uri": str(package_path.resolve()),
            "segment_package_sha256": file_sha256(package_path),
            "candidate_uri": str(candidate_path.resolve()),
            "artifacts": formal_artifacts,
            "tool_calls": [
                call.model_copy(
                    update={
                        "arguments": (
                            {
                                **call.arguments,
                                "dynamic_segment_input": (
                                    accepted_dependency.primary_artifact.uri
                                ),
                            }
                            if accepted_dependency is not None and index == 0
                            else call.arguments
                        ),
                        "output_refs": [
                            old_to_new.get(uri, uri)
                            for uri in call.output_refs
                        ]
                    }
                )
                for index, call in enumerate(delivery.tool_calls)
            ],
            "artifact_graph": delivery.artifact_graph.model_copy(
                update={
                    "lineages": [
                        lineage.model_copy(
                            update={
                                "input_artifact_ids": [
                                    (
                                        shared_anchor.anchor_id
                                        if artifact_id
                                        == old_anchor.anchor_id
                                        else artifact_id
                                    )
                                    for artifact_id
                                    in lineage.input_artifact_ids
                                ]
                                + (
                                    [optional_artifact.artifact_id]
                                    if (
                                        optional_artifact is not None
                                        and index == 0
                                    )
                                    else []
                                )
                                + (
                                    [
                                        accepted_dependency.primary_artifact.artifact_id
                                    ]
                                    if (
                                        accepted_dependency is not None
                                        and index == 0
                                    )
                                    else []
                                )
                            }
                        )
                        for index, lineage in enumerate(
                            delivery.artifact_graph.lineages
                        )
                    ]
                }
            ),
        }
    )
    write_once_json(
        segment_delivery_path(
            run_dir,
            package.segment_id,
            package.plan_version,
        ),
        delivery.model_dump(mode="json"),
    )
    write_once_json(
        segment_validator_candidate_path(
            run_dir,
            package.segment_id,
            package.plan_version,
        ),
        validator_candidate.model_dump(mode="json"),
    )
    write_once_json(
        segment_validation_report_path(
            run_dir,
            package.segment_id,
            package.plan_version,
        ),
        report.model_dump(mode="json"),
    )
    return package, run_dir


class SegmentAcceptanceTest(unittest.TestCase):
    def _accepted(self, package, run_dir):
        return build_accepted_segment_ref(
            package,
            run_dir=run_dir,
            package_record_path=segment_package_path(
                run_dir,
                package.segment_id,
                package.plan_version,
            ),
        )

    def test_complete_current_pass_bundle_can_be_accepted(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            package, run_dir = _persist_pass_bundle(Path(temp_dir))
            accepted = build_accepted_segment_ref(
                package,
                run_dir=run_dir,
                package_record_path=segment_package_path(
                    run_dir,
                    package.segment_id,
                    package.plan_version,
                ),
            )
            self.assertEqual(accepted.segment_id, package.segment_id)
            self.assertEqual(
                accepted.primary_artifact.segment_id,
                package.segment_id,
            )
            self.assertEqual(
                accepted.subtitle_artifact.segment_id,
                package.segment_id,
            )

    def test_tampered_subtitle_is_rejected_before_acceptance(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            package, run_dir = _persist_pass_bundle(Path(temp_dir))
            delivery_path = segment_delivery_path(
                run_dir,
                package.segment_id,
                package.plan_version,
            )
            from production_build_agents.contracts.production_models import (
                SegmentDelivery,
            )

            delivery = SegmentDelivery.model_validate_json(
                delivery_path.read_text(encoding="utf-8")
            )
            subtitle = next(
                item
                for item in delivery.artifacts
                if item.artifact_id == delivery.subtitle_artifact_id
            )
            Path(subtitle.uri).write_text("tampered", encoding="utf-8")
            with self.assertRaises(ArtifactIntegrityError):
                build_accepted_segment_ref(
                    package,
                    run_dir=run_dir,
                    package_record_path=segment_package_path(
                        run_dir,
                        package.segment_id,
                        package.plan_version,
                    ),
                )

    def test_optional_global_pool_growth_does_not_break_reuse(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            package, run_dir = _persist_pass_bundle(root)
            accepted = self._accepted(package, run_dir)
            delivery_path = Path(package.global_data_delivery_ref.uri)
            global_delivery = GlobalDataStageDelivery.model_validate_json(
                delivery_path.read_bytes()
            )
            optional_path = root / "new-optional.wav"
            optional_path.write_bytes(b"new optional")
            optional = Artifact(
                artifact_id="Task2-v1-artifact-1",
                artifact_type="audio",
                uri=str(optional_path.resolve()),
                source_uri=None,
                description="新增但未使用的可选素材",
                **content_identity_for_path(optional_path),
            )
            current_delivery = global_delivery.model_copy(
                update={
                    "active_artifacts": [
                        *global_delivery.active_artifacts,
                        optional,
                    ]
                }
            )
            current_delivery_path = root / "current-global-delivery.json"
            current_delivery_path.write_text(
                current_delivery.model_dump_json(),
                encoding="utf-8",
            )
            current = package.model_copy(
                update={
                    "plan_version": 2,
                    "global_data_delivery_ref": ImmutableJsonRef(
                        uri=str(current_delivery_path.resolve()),
                        document_sha256=exact_document_sha256(
                            current_delivery_path.read_bytes()
                        ),
                    ),
                }
            )
            self.assertTrue(
                segment_reuse_inputs_compatible(current, accepted)
            )

    def test_authorized_business_input_change_forces_execute(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            package, run_dir = _persist_pass_bundle(root)
            accepted = self._accepted(package, run_dir)
            current = package.model_copy(
                update={
                    "plan_version": 2,
                    "shots": [
                        shot.model_copy(
                            update={"objective": "授权修订后的镜头目标"}
                        )
                        for shot in package.shots
                    ],
                }
            )

            self.assertFalse(
                segment_reuse_inputs_compatible(current, accepted)
            )

    def test_unused_accepted_segment_pool_growth_does_not_break_reuse(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            package, run_dir = _persist_pass_bundle(root / "current")
            accepted = self._accepted(package, run_dir)
            dependency_package, dependency_run_dir = _persist_pass_bundle(
                root / "dependency",
                segment_index=2,
                run_id="dependency-run",
            )
            dependency_accepted = self._accepted(
                dependency_package,
                dependency_run_dir,
            )
            current = package.model_copy(
                update={
                    "plan_version": 2,
                    "dependencies": [
                        SegmentDependencyRef(
                            segment_id=dependency_accepted.segment_id,
                            segment_delivery_ref=(
                                dependency_accepted.segment_delivery_ref
                            ),
                            validation_report_ref=(
                                dependency_accepted.validation_report_ref
                            ),
                        )
                    ],
                }
            )
            self.assertTrue(
                segment_reuse_inputs_compatible(current, accepted)
            )

    def test_independent_dag_replan_invalidates_dynamic_cross_segment_input(
        self,
    ) -> None:
        """无显式 DAG 边时，也不能把已变化的旧上游泄漏给消费方。"""

        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            producer_package, producer_run_dir = _persist_pass_bundle(
                root / "producer",
                segment_index=1,
                run_id="dynamic-run",
            )
            producer_accepted = self._accepted(
                producer_package,
                producer_run_dir,
            )
            consumer_package, consumer_run_dir = _persist_pass_bundle(
                root / "consumer",
                accepted_dependency=producer_accepted,
                segment_index=2,
                run_id="dynamic-run",
            )
            consumer_accepted = self._accepted(
                consumer_package,
                consumer_run_dir,
            )

            previous = production_plan(
                [producer_package, consumer_package]
            )
            independent_segments = [
                previous.segments[0].model_copy(
                    update={"depends_on": [], "priority": 20}
                ),
                previous.segments[1].model_copy(
                    update={"depends_on": [], "priority": 10}
                ),
            ]
            previous = previous.model_copy(
                update={"segments": independent_segments}
            )
            changed_producer = independent_segments[0].model_copy(
                update={
                    "shots": [
                        independent_segments[0].shots[0].model_copy(
                            update={"objective": "Replan 后的新镜头目标"}
                        )
                    ]
                }
            )
            current = previous.model_copy(
                update={
                    "plan_version": 2,
                    "segments": [
                        changed_producer,
                        independent_segments[1],
                    ],
                }
            )
            planning = build_production_planner_scenario(
                run_id="dynamic-run"
            ).planning_package.model_copy(
                update={
                    "mode": "REPLAN",
                    "target_plan_version": 2,
                    "previous_plan_ref": (
                        producer_accepted.segment_package_ref
                    ),
                    "failure_report_ref": (
                        producer_accepted.validation_report_ref
                    ),
                    "accepted_segments": [
                        producer_accepted,
                        consumer_accepted,
                    ],
                    "authorized_revision_segment_ids": ["Segment1"],
                }
            )

            current_producer = producer_package.model_copy(
                update={
                    "plan_version": 2,
                    "shots": changed_producer.shots,
                }
            )
            current_consumer = consumer_package.model_copy(
                update={
                    "plan_version": 2,
                    "dependencies": [],
                }
            )
            scenario = build_production_planner_scenario(
                run_id="dynamic-run"
            )
            packages = {
                "Segment1": current_producer,
                "Segment2": current_consumer,
            }
            with patch(
                "production_build_agents.run.segment_inputs."
                "materialize_segment_package",
                side_effect=lambda *_args, segment_id, **_kwargs: (
                    packages[segment_id]
                ),
            ):
                snapshot = reusable_accepted_segment_snapshot(
                    scenario.brief,
                    scenario.global_data_delivery,
                    scenario.production_input_catalog,
                    planning,
                    current,
                    {},
                    planning_package_ref=planning.production_brief_ref,
                )
            self.assertEqual(list(snapshot), ["Segment2"])

            self.assertFalse(
                segment_reuse_inputs_compatible(
                    current_consumer,
                    consumer_accepted,
                )
            )

    def test_missing_actually_used_optional_artifact_forces_execute(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            package, run_dir = _persist_pass_bundle(
                root,
                include_unselected_actual_input=True,
            )
            accepted = self._accepted(package, run_dir)
            source_delivery = GlobalDataStageDelivery.model_validate_json(
                Path(package.global_data_delivery_ref.uri).read_bytes()
            )
            reduced = source_delivery.model_copy(
                update={
                    "active_artifacts": [
                        source_delivery.active_artifacts[0]
                    ]
                }
            )
            reduced_path = root / "reduced-global-delivery.json"
            reduced_path.write_text(
                reduced.model_dump_json(),
                encoding="utf-8",
            )
            current = package.model_copy(
                update={
                    "plan_version": 2,
                    "global_data_delivery_ref": ImmutableJsonRef(
                        uri=str(reduced_path.resolve()),
                        document_sha256=exact_document_sha256(
                            reduced_path.read_bytes()
                        ),
                    ),
                }
            )
            self.assertFalse(
                segment_reuse_inputs_compatible(current, accepted)
            )
            changed_path = root / "changed-optional-audio.wav"
            changed_path.write_bytes(b"changed optional audio")
            changed_optional = source_delivery.active_artifacts[1].model_copy(
                update={
                    "uri": str(changed_path.resolve()),
                    **content_identity_for_path(changed_path),
                }
            )
            changed_delivery = source_delivery.model_copy(
                update={
                    "active_artifacts": [
                        source_delivery.active_artifacts[0],
                        changed_optional,
                    ]
                }
            )
            changed_delivery_path = root / "changed-global-delivery.json"
            changed_delivery_path.write_text(
                changed_delivery.model_dump_json(),
                encoding="utf-8",
            )
            changed_current = package.model_copy(
                update={
                    "plan_version": 2,
                    "global_data_delivery_ref": ImmutableJsonRef(
                        uri=str(changed_delivery_path.resolve()),
                        document_sha256=exact_document_sha256(
                            changed_delivery_path.read_bytes()
                        ),
                    ),
                }
            )
            self.assertFalse(
                segment_reuse_inputs_compatible(
                    changed_current,
                    accepted,
                )
            )

    def test_reuse_compatibility_fails_closed_on_source_bundle_tamper(
        self,
    ) -> None:
        for target in ("artifact", "report"):
            with self.subTest(target=target):
                with tempfile.TemporaryDirectory() as temp_dir:
                    root = Path(temp_dir)
                    package, run_dir = _persist_pass_bundle(root)
                    accepted = self._accepted(package, run_dir)
                    current = package.model_copy(
                        update={"plan_version": 2}
                    )
                    if target == "artifact":
                        from production_build_agents.contracts.production_models import (
                            SegmentDelivery,
                        )

                        delivery = SegmentDelivery.model_validate_json(
                            Path(
                                accepted.segment_delivery_ref.uri
                            ).read_bytes()
                        )
                        artifact = next(
                            item
                            for item in delivery.artifacts
                            if item.artifact_id
                            == delivery.subtitle_artifact_id
                        )
                        Path(artifact.uri).write_text(
                            "tampered",
                            encoding="utf-8",
                        )
                    else:
                        Path(accepted.validation_report_ref.uri).write_text(
                            "{}",
                            encoding="utf-8",
                        )
                    with self.assertRaises(ValueError):
                        segment_reuse_inputs_compatible(
                            current,
                            accepted,
                        )


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