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 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
    requirement_ids: list[Identifier] = Field(min_length=1, max_length=5)
    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 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 geometry_conclusion_schema(working_head: str, evidence_refs: list[str]) -> dict[str, Any]:
    """Bind a diagnostic conclusion to evidence generated for this head."""
    schema = GeometryConclusion.model_json_schema()
    properties = schema.get("properties", {})
    if isinstance(properties, dict):
        properties["working_head"] = {"const": working_head}
        evidence = properties.get("evidence_refs")
        if isinstance(evidence, dict):
            evidence["items"] = {"enum": evidence_refs}
    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