cdsl-cad/backend/app/cad_agent/application/llm_contracts.py

"""Canonical schemas for every v3 LLM state-changing command.

Providers may parse structured output first, but this module validates the raw
tool arguments a second time before a command reaches a handler.
"""

from __future__ import annotations

from copy import deepcopy
import hashlib
import json
import math
from typing import Annotated, Any, Literal, TypeVar

from jsonschema import Draft202012Validator
from pydantic import BaseModel, ConfigDict, Field, JsonValue, RootModel, ValidationError, model_validator

from app.cad_agent.domain.errors import ErrorCode, WorkflowError


class StrictDto(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True, str_strip_whitespace=True)


ShortText = Annotated[str, Field(min_length=1, max_length=360)]
Identifier = Annotated[str, Field(pattern=r"^[a-z][a-z0-9_:-]{0,95}$")]


class AcceptanceClaimInput(StrictDto):
    claim_kind: Identifier
    expected: dict[str, JsonValue] = Field(min_length=0, max_length=24)


class SpecRequirementInput(StrictDto):
    statement: Annotated[str, Field(min_length=1, max_length=1000)]
    assumptions: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    acceptance_claims: list[AcceptanceClaimInput] = Field(min_length=1, max_length=16)


class MarkdownDocument(StrictDto):
    """A frozen human-readable design artifact, never an executable payload."""
    markdown: Annotated[str, Field(min_length=1, max_length=16_000)]


class CompiledRequirementInput(StrictDto):
    """One verifier bundle for one server-parsed checklist item.

    The checklist text, ordering, source bindings, and all identifiers are
    intentionally absent: the service owns them after Markdown is frozen.
    """
    assumptions: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    acceptance_claims: list[AcceptanceClaimInput] = Field(min_length=1, max_length=16)


class CompiledRequirementsSpec(StrictDto):
    requirements: list[CompiledRequirementInput] = Field(min_length=1, max_length=64)


# Kept only so an interrupted process with an already imported old tool schema
# fails at the workflow boundary instead of failing module import. New v3.1
# tasks never expose or accept this aggregate specification.
class RequirementsSpec(StrictDto):
    outcome: Literal["ready"]
    summary: Annotated[str, Field(min_length=1, max_length=2000)]
    assumptions: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=32)
    requirements: list[SpecRequirementInput] = Field(min_length=1, max_length=32)


class RequirementsClarification(StrictDto):
    outcome: Literal["clarification"]
    source_quotes: list[Annotated[str, Field(min_length=1, max_length=500)]] = Field(min_length=2, max_length=4)
    question: Annotated[str, Field(min_length=1, max_length=500)]


class RequirementsAuthorOutput(RootModel[Annotated[RequirementsSpec | RequirementsClarification, Field(discriminator="outcome")]]):
    pass


class EmptyCommand(StrictDto):
    pass


class NextAction(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    intent: ShortText
    # The server binds this list from every frozen checklist target. It is not
    # author input, so a five-item UI-era limit must not reject a valid task.
    requirement_ids: list[Identifier] = Field(min_length=1, max_length=64)
    atomic_id: Identifier
    expected_change: ShortText

    @model_validator(mode="after")
    def _requirement_ids_are_unique(self) -> "NextAction":
        if len(self.requirement_ids) != len(set(self.requirement_ids)):
            raise ValueError("requirement_ids must not contain duplicates")
        return self


class StatelessNextAction(StrictDto):
    intent: ShortText
    operation: Identifier
    expected_change: ShortText


class TopologyRequest(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    kind: Literal["face", "edge", "vertex", "plane", "axis", "body"] | None = None
    limit: int = Field(default=16, ge=1, le=64)


class GeometryConclusion(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    evidence_refs: list[Identifier] = Field(min_length=1, max_length=16)
    root_cause: Annotated[str, Field(min_length=1, max_length=360)]
    decision: Literal["return_to_action_selection", "rollback"]
    corrective_intent: str | None = Field(default=None, min_length=1, max_length=360)

    @model_validator(mode="after")
    def _evidence_refs_are_unique(self) -> "GeometryConclusion":
        if len(self.evidence_refs) != len(set(self.evidence_refs)):
            raise ValueError("evidence_refs must not contain duplicates")
        return self


class RollbackCheckpoint(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$")]
    checkpoint_token: Identifier
    reason: Annotated[str, Field(min_length=1, max_length=360)]


class StatelessTopologyRequest(StrictDto):
    kind: Literal["face", "edge", "vertex", "plane", "axis", "body"] | None = None
    limit: int = Field(default=16, ge=1, le=64)


class StatelessGeometryConclusion(StrictDto):
    root_cause: Annotated[str, Field(min_length=1, max_length=360)]
    decision: Literal["return_to_action_selection", "rollback"]
    corrective_intent: str | None = Field(default=None, min_length=1, max_length=360)


class StatelessRollbackCheckpoint(StrictDto):
    checkpoint_token: Identifier
    reason: Annotated[str, Field(min_length=1, max_length=360)]


class ClaimCoverage(StrictDto):
    claim_id: Identifier
    status: Literal["pass", "pending", "fail", "not_applicable"]
    evidence_refs: list[Identifier] = Field(default_factory=list, max_length=16)

    @model_validator(mode="after")
    def _evidence_refs_are_unique(self) -> "ClaimCoverage":
        if len(self.evidence_refs) != len(set(self.evidence_refs)):
            raise ValueError("evidence_refs must not contain duplicates")
        return self


class CandidateReview(StrictDto):
    candidate_id: Identifier = Field(description="Server-issued candidate ID from the review facts.")
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$", description="Current server-issued working head from the review facts.")]
    verdict: Literal["accept", "reject"] = Field(description="Required independent decision. Set accept only when the supplied candidate evidence supports every covered claim; otherwise set reject.")
    claim_coverage: list[ClaimCoverage] = Field(min_length=1, max_length=128, description="Required coverage decision for every claim ID in the supplied candidate facts.")
    evidence: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    issues: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)


class StatelessCandidateReview(StrictDto):
    verdict: Literal["accept", "reject"]
    evidence: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    issues: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)


class VisualClaimDecision(StrictDto):
    status: Literal["pass", "fail"]
    evidence: Annotated[str, Field(min_length=1, max_length=360)]


class StatelessFinalReview(StrictDto):
    verdict: Literal["pass", "repair"]
    visual_claims: list[VisualClaimDecision] = Field(default_factory=list, max_length=128)
    evidence: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    issues: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)


class ImageMeasurement(StrictDto):
    name: Annotated[str, Field(min_length=1, max_length=160)]
    value: float | None = None
    unit: Literal["mm", "degree", "count", "unknown"] = "unknown"
    evidence: Annotated[str, Field(min_length=1, max_length=360)]
    confidence: float = Field(ge=0, le=1)


class ImageObservation(StrictDto):
    summary: Annotated[str, Field(min_length=1, max_length=2000)]
    visible_features: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=64)
    measurements: list[ImageMeasurement] = Field(default_factory=list, max_length=128)
    view_directions: list[Annotated[str, Field(min_length=1, max_length=120)]] = Field(default_factory=list, max_length=16)
    uncertainties: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=64)
    assumptions: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=64)


class FinalReview(StrictDto):
    working_head: Annotated[str, Field(pattern=r"^[a-z0-9_:-]{5,192}$", description="Current server-issued working head from the final review facts.")]
    verdict: Literal["pass", "repair"] = Field(description="Required independent final decision. Set pass only when the supplied evidence supports every claim; otherwise set repair.")
    claim_coverage: list[ClaimCoverage] = Field(min_length=1, max_length=128, description="Required coverage decision for every claim ID in the supplied final-review facts.")
    evidence: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)
    issues: list[Annotated[str, Field(min_length=1, max_length=360)]] = Field(default_factory=list, max_length=16)


def requirements_spec_schema(claim_one_of: dict[str, Any]) -> dict[str, Any]:
    schema = RequirementsAuthorOutput.model_json_schema()
    requirement = schema.get("$defs", {}).get("SpecRequirementInput")
    if isinstance(requirement, dict):
        claims = requirement.get("properties", {}).get("acceptance_claims")
        if isinstance(claims, dict):
            claims["items"] = deepcopy(claim_one_of)
    return schema


def compiled_requirements_schema(claim_one_of: dict[str, Any], target_count: int) -> dict[str, Any]:
    schema = CompiledRequirementsSpec.model_json_schema()
    definitions = schema.get("$defs", {})
    requirement = definitions.get("CompiledRequirementInput") if isinstance(definitions, dict) else None
    if isinstance(requirement, dict):
        claims = requirement.get("properties", {}).get("acceptance_claims")
        if isinstance(claims, dict):
            claims["items"] = deepcopy(claim_one_of)
    requirements = schema.get("properties", {}).get("requirements")
    if isinstance(requirements, dict):
        requirements["minItems"] = target_count
        requirements["maxItems"] = target_count
    return schema


def sanitize_compiled_requirements_arguments(raw_arguments_json: str) -> str | WorkflowError:
    """Drop harmless compiler chatter before strict requirements validation.

    ``compile_requirements_spec`` is a compiler stage: the service only needs
    the ordered verifier bundles for the frozen checklist items. Real models
    sometimes add explanatory fields such as a top-level ``assumptions`` or
    per-item ``statement`` even when the dynamic tool schema forbids them. Those
    fields are not executable and are not part of the frozen contract, so they
    should not abort a task before modeling starts.

    The verifier ``expected`` payload is intentionally not sanitized here. It
    remains governed by the registry's strict per-claim schema because those
    values drive deterministic validation.
    """
    value = canonical_json_object(raw_arguments_json)
    if isinstance(value, WorkflowError):
        return value
    requirements = value.get("requirements")
    sanitized: dict[str, Any] = {}
    if isinstance(requirements, list):
        sanitized_requirements: list[Any] = []
        for requirement in requirements:
            if not isinstance(requirement, dict):
                sanitized_requirements.append(requirement)
                continue
            item: dict[str, Any] = {}
            if "assumptions" in requirement:
                item["assumptions"] = requirement["assumptions"]
            if "acceptance_claims" in requirement:
                claims = requirement["acceptance_claims"]
                if isinstance(claims, list):
                    item["acceptance_claims"] = [
                        {key: claim[key] for key in ("claim_kind", "expected") if isinstance(claim, dict) and key in claim}
                        if isinstance(claim, dict) else claim
                        for claim in claims
                    ]
                else:
                    item["acceptance_claims"] = claims
            sanitized_requirements.append(item)
        sanitized["requirements"] = sanitized_requirements
    else:
        sanitized["requirements"] = requirements
    return json.dumps(sanitized, ensure_ascii=False, separators=(",", ":"))


def stateless_next_action_schema(atomic_ids: list[str]) -> dict[str, Any]:
    schema = StatelessNextAction.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["operation"] = {"enum": atomic_ids}
    return schema


def stateless_final_review_schema(visual_claim_count: int) -> dict[str, Any]:
    schema = StatelessFinalReview.model_json_schema()
    properties = schema.get("properties", {})
    visual = properties.get("visual_claims") if isinstance(properties, dict) else None
    if isinstance(visual, dict):
        visual["minItems"] = visual_claim_count
        visual["maxItems"] = visual_claim_count
    return schema


def stateless_rollback_checkpoint_schema(checkpoint_tokens: list[str]) -> dict[str, Any]:
    schema = StatelessRollbackCheckpoint.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["checkpoint_token"] = {"enum": checkpoint_tokens}
    return schema


def topology_request_schema(working_head: str) -> dict[str, Any]:
    schema = TopologyRequest.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
    return schema


def rollback_checkpoint_schema(working_head: str, checkpoint_tokens: list[str]) -> dict[str, Any]:
    """Bind a rollback request to immutable checkpoints in the active lineage."""
    schema = RollbackCheckpoint.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
        properties["checkpoint_token"] = {"enum": checkpoint_tokens}
    return schema


def _bind_claim_coverage_ids(schema: dict[str, Any], claim_ids: list[str]) -> None:
    definitions = schema.get("$defs", {})
    coverage = definitions.get("ClaimCoverage") if isinstance(definitions, dict) else None
    if not isinstance(coverage, dict):
        return
    properties = coverage.get("properties", {})
    if isinstance(properties, dict):
        properties["claim_id"] = {"enum": claim_ids}


Dto = TypeVar("Dto", bound=StrictDto)


def raw_arguments_hash(raw_arguments_json: str) -> str:
    return hashlib.sha256(raw_arguments_json.encode("utf-8")).hexdigest()


def json_depth(value: Any, current: int = 0) -> int:
    if isinstance(value, dict):
        return max([current, *(json_depth(item, current + 1) for item in value.values())])
    if isinstance(value, list):
        return max([current, *(json_depth(item, current + 1) for item in value)])
    return current


def canonical_json_object(raw_arguments_json: str, *, max_bytes: int = 48_000, max_depth: int = 16) -> dict[str, Any] | WorkflowError:
    """Bound and parse raw arguments before any schema-specific validation."""
    if len(raw_arguments_json.encode("utf-8")) > max_bytes:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments exceed the byte limit.")
    try:
        value = json.loads(raw_arguments_json)
    except json.JSONDecodeError as error:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments are not valid JSON.", field_errors=({"path": "/", "message": error.msg},))
    if not isinstance(value, dict):
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments must be a JSON object.")
    if json_depth(value) > max_depth:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments exceed the nesting-depth limit.")
    if _contains_non_finite_number(value):
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments must not contain NaN or infinite numbers.")
    return value


def _contains_non_finite_number(value: Any) -> bool:
    if isinstance(value, float):
        return not math.isfinite(value)
    if isinstance(value, dict):
        return any(_contains_non_finite_number(item) for item in value.values())
    if isinstance(value, list):
        return any(_contains_non_finite_number(item) for item in value)
    return False


def canonical_validate_schema(raw_arguments_json: str, schema: dict[str, Any]) -> WorkflowError | None:
    """Revalidate raw arguments against the current dynamic JSON Schema."""
    value = canonical_json_object(raw_arguments_json)
    if isinstance(value, WorkflowError):
        return value
    errors = [
        {"path": "/" + "/".join(str(part) for part in error.absolute_path), "message": error.message}
        for error in sorted(Draft202012Validator(schema).iter_errors(value), key=lambda item: (list(item.absolute_path), item.message))
    ]
    if errors:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments do not match the active dynamic schema.", field_errors=tuple(errors))
    return None


def canonical_validate(raw_arguments_json: str, model: type[Dto], *, max_bytes: int = 48_000, max_depth: int = 16) -> Dto | WorkflowError:
    """Parse raw tool arguments once and return field-level DTO failures safely."""
    value = canonical_json_object(raw_arguments_json, max_bytes=max_bytes, max_depth=max_depth)
    if isinstance(value, WorkflowError):
        return value
    try:
        return model.model_validate(value)
    except ValidationError as error:
        fields = tuple({"path": "/" + "/".join(str(part) for part in issue["loc"]), "message": issue["msg"]} for issue in error.errors())
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Tool arguments do not match the canonical schema.", field_errors=fields)


def validate_one_tool_call(tool_calls: list[dict[str, Any]], allowed_name: str) -> tuple[str, str] | WorkflowError:
    """Require exactly one known tool call; no provider parser is trusted."""
    if len(tool_calls) != 1:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "Exactly one tool call is required.")
    function = tool_calls[0].get("function") if isinstance(tool_calls[0], dict) else None
    name = str(function.get("name") or "") if isinstance(function, dict) else ""
    raw = str(function.get("arguments") or "") if isinstance(function, dict) else ""
    if name != allowed_name:
        return WorkflowError(ErrorCode.AUTHOR_FORMAT_INVALID, "The returned tool is not allowed in this workflow state.", details={"expected_tool": allowed_name, "actual_tool": name})
    return name, raw