refactor(lean): share proof-state validation (#761)

Author: morluto

src/jacobian/lean_frontend/_state_validation.py

@@ -0,0 +1,93 @@
+"""Validation of immutable proof states against a caller-selected Lean profile."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Protocol
+
+from pydantic import ValidationError
+
+from jacobian.capability_service import CapabilityInvocationError
+from jacobian.contracts.capabilities import CapabilityDiagnostic
+from jacobian.contracts.lean import LeanEnvironment
+from jacobian.contracts.lean_exploration import LeanProofStateArtifact
+from jacobian.lean_frontend.artifacts import (
+    _environment_imports,
+    _source_digest,
+    _state_digest_payload,
+)
+from jacobian.references import LeanCheckerInstallation
+from jacobian.storage.errors import StorageError
+from jacobian.storage.repository import ArtifactRepository
+
+
+class _StoredProofStateResources(Protocol):
+    @property
+    def store(self) -> ArtifactRepository: ...
+
+    @property
+    def semantics_uri(self) -> str: ...
+
+    @property
+    def state_schema_uri(self) -> str: ...
+
+    @property
+    def installations(
+        self,
+    ) -> Mapping[LeanEnvironment, LeanCheckerInstallation]: ...
+
+
+def _load_validated_proof_state(
+    resources: _StoredProofStateResources,
+    state_uri: str,
+    *,
+    expected_environment: LeanEnvironment,
+    expected_environment_digest: str,
+    invalid_state_hint: str,
+) -> LeanProofStateArtifact:
+    """Load one state bound to the profile selected by the current request."""
+
+    try:
+        stored = resources.store.get(state_uri)
+        if (
+            stored.manifest.schema_uri != resources.state_schema_uri
+            or stored.manifest.semantics_uri != resources.semantics_uri
+        ):
+            raise ValueError("artifact is not a Lean proof state")
+        state = LeanProofStateArtifact.model_validate(stored.payload)
+    except (StorageError, ValidationError, ValueError) as exc:
+        raise CapabilityInvocationError(
+            CapabilityDiagnostic(
+                code="INVALID_LEAN_PROOF_STATE",
+                stage="state_loading",
+                message="The supplied state artifact is unavailable or invalid.",
+                hint=invalid_state_hint,
+            )
+        ) from exc
+
+    installation = resources.installations[expected_environment]
+    if (
+        state.environment is not expected_environment
+        or state.environment_digest != expected_environment_digest
+        or state.imports != _environment_imports(expected_environment)
+        or state.lean_version != installation.lean_version
+        or state.lean_commit != installation.lean_commit
+        or state.mathlib_commit != installation.mathlib_commit
+        or state.source_digest != _source_digest(state.statement, state.tactic_prefix)
+        or state.state_digest != _state_digest_payload(state)
+    ):
+        raise CapabilityInvocationError(
+            CapabilityDiagnostic(
+                code="STALE_LEAN_PROOF_STATE",
+                stage="state_validation",
+                message=(
+                    "The proof state no longer matches its source or the "
+                    "current pinned Lean environment."
+                ),
+                hint="Recreate the proof state under the current environment.",
+            )
+        )
+    return state
+
+
+__all__ = ["_StoredProofStateResources", "_load_validated_proof_state"]

src/jacobian/lean_frontend/metavariable_fields.py

@@ -33,8 +33,6 @@
     CapabilityResult,
     CapabilityScope,
 )
-from jacobian.contracts.lean import LeanEnvironment
-from jacobian.contracts.lean_exploration import LeanProofStateArtifact
 from jacobian.contracts.lean_metavariable_fields import (
     LeanElaborationContext,
     LeanMetavariableFieldsArtifact,
@@ -43,20 +41,18 @@
     LeanStructuredMetavariable,
 )
 from jacobian.contracts.results import Execution, ExecutionStatus
+from jacobian.lean_frontend._state_validation import _load_validated_proof_state
 from jacobian.lean_frontend.artifacts import (
     _environment_digest,
-    _environment_imports,
     _proof_state_command,
     _source_digest,
-    _state_digest_payload,
 )
 from jacobian.lean_frontend.exploration import (
     _Resources,
     _runtime_ms,
     _validate_source_parts,
 )
 from jacobian.lean_frontend.repl import _response_errors
-from jacobian.storage.errors import StorageError
 
 
 class LeanMetavariableFieldsAdapter:
@@ -113,10 +109,14 @@ def invoke(self, request: CapabilityRequest) -> CapabilityResult:
             validated.environment,
             installation,
         )
-        bound_state = self._load_bound_state(
+        bound_state = _load_validated_proof_state(
+            self.resources,
             validated.state_uri,
             expected_environment=validated.environment,
             expected_environment_digest=environment_digest,
+            invalid_state_hint=(
+                "Use a state URI returned by a proof-state capability."
+            ),
         )
         if bound_state.completed:
             raise CapabilityInvocationError(
@@ -326,56 +326,6 @@ def invoke(self, request: CapabilityRequest) -> CapabilityResult:
             artifact_uris=(validated.state_uri, artifact.artifact_uri),
         )
 
-    def _load_bound_state(
-        self,
-        state_uri: str,
-        *,
-        expected_environment: LeanEnvironment,
-        expected_environment_digest: str,
-    ) -> LeanProofStateArtifact:
-        try:
-            stored = self.resources.store.get(state_uri)
-            if (
-                stored.manifest.schema_uri != self.resources.state_schema_uri
-                or stored.manifest.semantics_uri != self.resources.semantics_uri
-            ):
-                raise ValueError("artifact is not a Lean proof state")
-            state = LeanProofStateArtifact.model_validate(stored.payload)
-        except (StorageError, ValidationError, ValueError) as exc:
-            raise CapabilityInvocationError(
-                CapabilityDiagnostic(
-                    code="INVALID_LEAN_PROOF_STATE",
-                    stage="state_loading",
-                    message="The supplied state artifact is unavailable or invalid.",
-                    hint="Use a state URI returned by a proof-state capability.",
-                )
-            ) from exc
-        installation = self.resources.installations[expected_environment]
-        expected_imports = _environment_imports(expected_environment)
-        if (
-            state.environment is not expected_environment
-            or state.environment_digest != expected_environment_digest
-            or state.imports != expected_imports
-            or state.lean_version != installation.lean_version
-            or state.lean_commit != installation.lean_commit
-            or state.mathlib_commit != installation.mathlib_commit
-            or state.source_digest
-            != _source_digest(state.statement, state.tactic_prefix)
-            or state.state_digest != _state_digest_payload(state)
-        ):
-            raise CapabilityInvocationError(
-                CapabilityDiagnostic(
-                    code="STALE_LEAN_PROOF_STATE",
-                    stage="state_validation",
-                    message=(
-                        "The proof state no longer matches its source or the "
-                        "current pinned Lean environment."
-                    ),
-                    hint="Recreate the proof state under the current environment.",
-                )
-            )
-        return state
-
 
 def install_lean_metavariable_fields_capability(
     resources: _Resources,

src/jacobian/lean_frontend/proof_state.py

@@ -23,22 +23,19 @@
     CapabilityResult,
     CapabilityScope,
 )
-from jacobian.contracts.lean import LeanEnvironment
 from jacobian.contracts.lean_exploration import (
-    LeanProofStateArtifact,
     LeanProofStateOutput,
     LeanProofStateRequest,
     LeanProofStateTransitionArtifact,
     LeanProofSuccessorState,
     LeanTypedGoal,
 )
 from jacobian.contracts.results import Execution, ExecutionStatus
+from jacobian.lean_frontend._state_validation import _load_validated_proof_state
 from jacobian.lean_frontend.artifacts import (
     _environment_digest,
-    _environment_imports,
     _proof_state_command,
     _source_digest,
-    _state_digest_payload,
     _state_payload,
 )
 from jacobian.lean_frontend.exploration import (
@@ -49,7 +46,6 @@
     _validate_source_parts,
 )
 from jacobian.lean_frontend.repl import _response_errors
-from jacobian.storage.errors import StorageError
 
 
 class LeanProofStateAdapter:
@@ -132,10 +128,12 @@ def invoke(self, request: CapabilityRequest) -> CapabilityResult:
             proof_prefix = validated.proof_prefix
             bound_state = None
         else:
-            bound_state = self._load_bound_state(
+            bound_state = _load_validated_proof_state(
+                self.resources,
                 validated.state_uri,
                 expected_environment=validated.environment,
                 expected_environment_digest=environment_digest,
+                invalid_state_hint="Use a state URI returned by this capability.",
             )
             if bound_state.completed:
                 raise CapabilityInvocationError(
@@ -402,55 +400,5 @@ def invoke(self, request: CapabilityRequest) -> CapabilityResult:
             artifact_uris=artifact_uris,
         )
 
-    def _load_bound_state(
-        self,
-        state_uri: str,
-        *,
-        expected_environment: LeanEnvironment,
-        expected_environment_digest: str,
-    ) -> LeanProofStateArtifact:
-        try:
-            stored = self.resources.store.get(state_uri)
-            if (
-                stored.manifest.schema_uri != self.resources.state_schema_uri
-                or stored.manifest.semantics_uri != self.resources.semantics_uri
-            ):
-                raise ValueError("artifact is not a Lean proof state")
-            state = LeanProofStateArtifact.model_validate(stored.payload)
-        except (StorageError, ValidationError, ValueError) as exc:
-            raise CapabilityInvocationError(
-                CapabilityDiagnostic(
-                    code="INVALID_LEAN_PROOF_STATE",
-                    stage="state_loading",
-                    message="The supplied state artifact is unavailable or invalid.",
-                    hint="Use a state URI returned by this capability.",
-                )
-            ) from exc
-        installation = self.resources.installations[expected_environment]
-        expected_imports = _environment_imports(expected_environment)
-        if (
-            state.environment is not expected_environment
-            or state.environment_digest != expected_environment_digest
-            or state.imports != expected_imports
-            or state.lean_version != installation.lean_version
-            or state.lean_commit != installation.lean_commit
-            or state.mathlib_commit != installation.mathlib_commit
-            or state.source_digest
-            != _source_digest(state.statement, state.tactic_prefix)
-            or state.state_digest != _state_digest_payload(state)
-        ):
-            raise CapabilityInvocationError(
-                CapabilityDiagnostic(
-                    code="STALE_LEAN_PROOF_STATE",
-                    stage="state_validation",
-                    message=(
-                        "The proof state no longer matches its source or the "
-                        "current pinned Lean environment."
-                    ),
-                    hint="Recreate the proof state under the current environment.",
-                )
-            )
-        return state
-
 
 __all__ = ["LeanProofStateAdapter"]