dify/api/controllers/inner_api/agent/runtime_selection.py

"""Canonical bounded Run header parsing; the parsed value grants no authority."""

import base64
import json

from dify_agent.adapters.native_testing.client import AgentPhysicalDispatchPayload
from dify_agent.protocol.native_testing import (
    MAX_AGENT_CONTROL_BYTES,
    AgentRunTransportSelection,
)
from pydantic import BaseModel

from runtime_contracts.builder.execution import BuilderExecutionPolicyError


def parse_agent_run_selection(headers) -> AgentRunTransportSelection:
    return _parse_header(headers, "X-Dify-Agent-Run", AgentRunTransportSelection)


def parse_agent_config_dispatch(headers) -> AgentPhysicalDispatchPayload:
    """Decode a bounded exact operation header; it grants no permission."""
    return _parse_header(headers, "X-Dify-Agent-Dispatch", AgentPhysicalDispatchPayload)


def _parse_header[Model: BaseModel](headers, name: str, model: type[Model]) -> Model:
    values = headers.getlist(name)
    if len(values) != 1:
        raise BuilderExecutionPolicyError("native_agent_selection_invalid")
    encoded = values[0]
    try:
        if not encoded or len(encoded) > (MAX_AGENT_CONTROL_BYTES * 4 + 2) // 3:
            raise ValueError("invalid selection size")
        raw = base64.b64decode(encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True)
        if len(raw) > MAX_AGENT_CONTROL_BYTES:
            raise ValueError("invalid selection size")
        value = model.model_validate_json(raw)
        canonical = json.dumps(
            value.model_dump(mode="json"),
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=False,
            allow_nan=False,
        ).encode("utf-8")
        if raw != canonical or encoded != base64.urlsafe_b64encode(canonical).rstrip(b"=").decode("ascii"):
            raise ValueError("noncanonical selection")
        return value
    except ValueError:
        raise BuilderExecutionPolicyError("native_agent_selection_invalid") from None