dify/api/services/dify_builder/native_testing/agent_lineage.py

"""Persisted Run lineage before any credential, provider, file or config work.

Reads use the exact S1 private selection and current protected state owners. All
backend/object I/O stays outside their short SQL writers. This owner validates
lineage only; no invocation permit is spent and no physical capability is ready.
"""

from collections.abc import Callable
from dataclasses import dataclass
from typing import Literal

from dify_agent.client import DifyAgentClientError
from dify_agent.protocol.native_testing import (
    AgentNativeTestRef,
    AgentReadResultPayload,
    AgentRuntimeIdentity,
    AgentRunTransportSelection,
    ReadAgentRuntimeIdentity,
)
from pydantic import JsonValue

from runtime_contracts.builder.execution import BuilderExecutionPolicyError
from runtime_contracts.builder.native.agent_contracts import (
    AgentNodeAssociation,
    AgentParentNodeRef,
    AgentPreparedConfig,
)
from services.dify_builder.native_testing.agent_preparation import (
    AgentPreparationService,
    AgentRuntimeIdentityReader,
)
from services.dify_builder.native_testing.state_store import native_payload_read_scope


@dataclass(frozen=True, slots=True)
class AgentServicePrincipal:
    """Request-local marker constructed only after the Agent bearer authenticates."""


@dataclass(frozen=True, slots=True)
class NativeAgentEndpointContext:
    ref: AgentNativeTestRef
    config: AgentPreparedConfig
    association: AgentNodeAssociation | None
    physical_owner_ready: bool = False
    native_test_profile: Literal["isolated", "confirmed"] = "isolated"


@dataclass(frozen=True, slots=True)
class NativeAgentRecoveryReadContext:
    """One authenticated, selected predecessor result; never mutation authority."""

    ref: AgentNativeTestRef
    result_ref: str
    result: JsonValue


@dataclass
class NativeAgentEndpointLineage:
    backend: AgentRuntimeIdentityReader
    preparation_for_ref: Callable[..., AgentPreparationService]

    @native_payload_read_scope()
    def read_recovery_result(self, *, principal, selection, payload: AgentReadResultPayload):
        """Read one selected predecessor object after exact service/runtime proof.

        Only native-result-read calls this port. Ordinary selections return None;
        an existing recovery must prove the full old tuple and A's SQL predicate.
        No identity response alone grants object access or mutation authority.
        """
        if not isinstance(principal, AgentServicePrincipal) or not isinstance(selection, AgentRunTransportSelection):
            raise BuilderExecutionPolicyError("native_endpoint_principal_unavailable")
        if selection.parent_node is None:
            return None
        parent = AgentParentNodeRef.model_validate(selection.parent_node.model_dump(mode="json"))
        if payload.ref.model_dump(mode="json") != parent.ref.model_dump(mode="json"):
            raise BuilderExecutionPolicyError("native_endpoint_owner_mismatch")
        store = self.preparation_for_ref(parent.ref).state_store
        relation = store.load_selected_recovery(ref=parent.ref, parent_node=parent)
        if relation is None:
            return None
        if (selection.agent_run_id, selection.request_digest) == (
            relation.agent_run_id,
            relation.successor.config_digest,
        ):
            return None
        if (selection.agent_run_id, selection.request_digest, selection.agent_generation) != (
            relation.agent_run_id,
            relation.predecessor.config_digest,
            relation.expected_agent_generation,
        ):
            raise BuilderExecutionPolicyError("native_agent_identity_mismatch")
        try:
            actual = AgentRuntimeIdentity.model_validate(
                self.backend.read_run_identity_sync(
                    selection.agent_run_id,
                    ReadAgentRuntimeIdentity(
                        parent_node=selection.parent_node, request_digest=selection.request_digest
                    ),
                ).model_dump(mode="json")
            )
        except (DifyAgentClientError, ValueError, OSError):
            raise BuilderExecutionPolicyError("native_agent_identity_unavailable") from None
        if (
            actual.run_id,
            actual.parent_node,
            actual.request_digest,
            actual.agent_generation,
            actual.execution_mode,
            actual.native_test_profile,
        ) != (
            selection.agent_run_id,
            selection.parent_node,
            selection.request_digest,
            selection.agent_generation,
            "builder_test",
            "confirmed",
        ):
            raise BuilderExecutionPolicyError("native_agent_identity_mismatch")
        result = store.read_recovery_predecessor_result(
            ref=parent.ref, parent_node=parent, selection=relation, result_ref=payload.result_ref
        )
        return NativeAgentRecoveryReadContext(payload.ref, payload.result_ref, result)

    @native_payload_read_scope()
    def resolve(
        self,
        *,
        principal: AgentServicePrincipal,
        selection: AgentRunTransportSelection,
        boundary: str,
        caller: dict,
    ) -> NativeAgentEndpointContext | None:
        if not isinstance(principal, AgentServicePrincipal) or not isinstance(selection, AgentRunTransportSelection):
            raise BuilderExecutionPolicyError("native_endpoint_principal_unavailable")
        if caller.get("agent_run_id", selection.agent_run_id) != selection.agent_run_id:
            raise BuilderExecutionPolicyError("native_endpoint_owner_mismatch")
        # Capture protected ownership before identity I/O so an owner/config
        # mutation during that read cannot become the new comparison baseline.
        store = None
        parent = None
        context = None
        config = None
        association = None
        control = boundary in {"prepare-run", "prepared-config"}
        settlement = boundary == "native-settlement"
        if selection.parent_node is not None:
            parent = AgentParentNodeRef.model_validate(selection.parent_node.model_dump(mode="json"))
            store = self.preparation_for_ref(parent.ref).state_store
            if settlement:
                context, config, association = store.load_agent_settlement_identity(ref=parent.ref, parent_node=parent)
            else:
                context = store.load_context(ref=parent.ref)
                config = store.load_prepared_config(ref=parent.ref, parent_node=parent, config_ref=None)
                association = store.discover_agent_run(ref=parent.ref, parent_node=parent)
            if config.config_digest != selection.request_digest or (association is None and not control):
                raise BuilderExecutionPolicyError("native_agent_identity_mismatch")
        try:
            readback = self.backend.read_run_identity_sync(
                selection.agent_run_id,
                ReadAgentRuntimeIdentity(
                    parent_node=selection.parent_node,
                    request_digest=selection.request_digest,
                ),
            )
            actual = AgentRuntimeIdentity.model_validate(
                readback.model_dump(mode="json") if isinstance(readback, AgentRuntimeIdentity) else readback
            )
        except (
            DifyAgentClientError,
            ValueError,
            OSError,
        ):
            raise BuilderExecutionPolicyError("native_agent_identity_unavailable") from None
        if (
            actual.run_id,
            actual.parent_node,
            actual.request_digest,
            actual.agent_generation,
        ) != (
            selection.agent_run_id,
            selection.parent_node,
            selection.request_digest,
            selection.agent_generation,
        ):
            raise BuilderExecutionPolicyError("native_agent_identity_mismatch")
        if actual.execution_mode == "ordinary" and actual.parent_node is None:
            return None
        if (
            actual.execution_mode != "builder_test"
            or store is None
            or parent is None
            or config is None
            or context is None
        ):
            raise BuilderExecutionPolicyError("native_agent_identity_mismatch")
        # Each S1 read checks current resources/RBAC after its own object I/O.
        # Repeat across the complete private identity/object boundary as well.
        if settlement:
            current_context, current_config, current_association = store.load_agent_settlement_identity(
                ref=parent.ref, parent_node=parent
            )
        else:
            current_config = store.load_prepared_config(
                ref=parent.ref, parent_node=parent, config_ref=config.config_ref
            )
            current_association = store.discover_agent_run(ref=parent.ref, parent_node=parent)
            current_context = store.load_context(ref=parent.ref)
        if current_config != config or current_association != association or current_context != context:
            raise BuilderExecutionPolicyError("native_agent_owner_mismatch")
        if association is not None and (
            association.parent_node != parent
            or association.agent_run_id != actual.run_id
            or association.config_ref != config.config_ref
            or association.config_digest != actual.request_digest
            or association.workflow_native_run_id != config.workflow_native_run_id
            or association.workspace != config.workspace
            or config.workflow_native_run_id != context.native_run_id
        ):
            raise BuilderExecutionPolicyError("native_agent_owner_mismatch")
        for key, expected in (
            ("tenant_id", context.tenant_id),
            ("app_id", context.app_id),
            ("user_id", context.actor_id),
            ("user_from", "account"),
        ):
            if key in caller and caller[key] != expected:
                raise BuilderExecutionPolicyError("native_endpoint_owner_mismatch")
        # Creation and an unbound config read belong to initial generation zero.
        # A resumed bound Run may re-read its immutable config before acquiring
        # the workspace; exact current identity and association were proved above.
        if control and actual.agent_generation != 0 and (boundary != "prepared-config" or association is None):
            raise BuilderExecutionPolicyError("native_agent_identity_mismatch")
        return NativeAgentEndpointContext(
            ref=AgentNativeTestRef.model_validate(parent.ref.model_dump(mode="json")),
            config=config,
            association=association,
            native_test_profile=actual.native_test_profile,
        )