diff --git a/NeuralComposeEEG/src/neuralcompose_eeg/structured_state.py b/NeuralComposeEEG/src/neuralcompose_eeg/structured_state.py new file mode 100644 index 0000000..6552c28 --- /dev/null +++ b/NeuralComposeEEG/src/neuralcompose_eeg/structured_state.py @@ -0,0 +1,363 @@ +"""Serialize offline encoder probabilities into a replayable shadow-state contract. + +This bridge is deliberately downstream of an offline encoder artifact. It +stores no waveform values, dialogue content, task labels, or policy actions. +The resulting records are suitable only for deterministic replay while the +encoder program remains ``pipeline_only``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +from pathlib import Path +from typing import Any, Iterable + +import numpy as np + +from .contracts import ContractError +from .dataset import CanonicalDataset, load_dataset +from .provenance import sha256_json + + +STATE_SCHEMA = "nc-eeg-encoder-state-v0" +MANIFEST_SCHEMA = "nc-eeg-encoder-state-manifest-v0" +STATE_BRIDGE_VERSION = "nc-eeg-structured-state-bridge-v0" +_REQUIRED_GLOBAL_STATUS = { + "status": "insufficient_evidence", + "science_status": "pipeline_only", + "shadow_only": True, + "live_control": False, + "promotion_status": "not_eligible", +} +_ENCODER_PROVENANCE_KEYS = frozenset( + { + "model_id", + "model_revision", + "initialization", + "checkpoint_sha256", + "extractor_sha256", + "source_kind", + } +) +_HEX_DIGITS = "0123456789abcdef" +_FORBIDDEN_SERIALIZED_TERMS = ( + "dialogue", + "dialectic", + "prompt", + "speech", + "text", + "policy", + "action", + "eeg_values", + "waveform", + "label", + "target", +) + + +def _canonical_json(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + + +def _publish_text(path: Path, text: str) -> None: + """Write via same-directory rename so an interrupted run cannot publish a + truncated artifact over a previously valid one.""" + partial = path.with_name(f".{path.name}.partial") + try: + partial.write_text(text, encoding="utf-8") + os.replace(partial, path) + finally: + partial.unlink(missing_ok=True) + + +def _require_sha256(value: Any, field: str) -> str: + if not isinstance(value, str) or len(value) != 64 or value.strip(_HEX_DIGITS): + raise ContractError(f"{field} must be a lowercase SHA-256 hex digest") + return value + + +def _state_id(dataset_sha256: str, raw_window_sha256: str) -> str: + return hashlib.sha256(f"{dataset_sha256}:{raw_window_sha256}".encode("utf-8")).hexdigest() + + +def _settled(path: Path, text: str, description: str) -> bool: + """Report whether `path` already holds exactly `text`, refusing any other + existing content. Called for both artifacts before either is written, so a + refusal leaves the published pair untouched.""" + try: + existing = path.read_bytes() + except FileNotFoundError: + return False + except OSError as exc: + raise ContractError(f"cannot read existing {description}: {exc}") from exc + if existing != text.encode("utf-8"): + raise ContractError(f"refusing to overwrite existing {description} that differs from this run") + return True + + +def _sha256_lines(records: Iterable[dict[str, Any]]) -> str: + digest = hashlib.sha256() + for record in records: + digest.update(_canonical_json(record).encode("utf-8")) + digest.update(b"\n") + return digest.hexdigest() + + +def _entropy(probabilities: np.ndarray) -> float: + nonzero = probabilities[probabilities > 0] + return float(-np.sum(nonzero * np.log(nonzero))) + + +def _validate_encoder_provenance(value: dict[str, Any]) -> dict[str, str]: + if not isinstance(value, dict): + raise ContractError("encoder provenance must be an object") + unexpected = set(value) - _ENCODER_PROVENANCE_KEYS + if unexpected: + raise ContractError(f"encoder provenance has unsupported fields: {sorted(unexpected)}") + required = {"model_id", "model_revision", "source_kind"} + missing = required - set(value) + if missing: + raise ContractError(f"encoder provenance lacks fields: {sorted(missing)}") + normalized: dict[str, str] = {} + for key, item in value.items(): + if not isinstance(item, str) or not item: + raise ContractError(f"encoder provenance {key} must be a nonempty string") + normalized[key] = item + return normalized + + +def _validate_probabilities(probabilities: np.ndarray, expected_rows: int, class_count: int) -> np.ndarray: + result = np.asarray(probabilities, dtype=np.float64) + if result.ndim != 2 or result.shape != (expected_rows, class_count): + raise ContractError("probabilities must have [canonical_window, observable_class] shape") + if not np.all(np.isfinite(result)) or np.any(result < 0) or np.any(result > 1): + raise ContractError("probabilities must be finite values in [0, 1]") + if not np.allclose(result.sum(axis=1), 1.0, atol=1e-6): + raise ContractError("probabilities must sum to one for every canonical window") + return result + + +def _require_no_forbidden_serialized_terms(value: Any) -> None: + serialized = _canonical_json(value).casefold() + matched = [term for term in _FORBIDDEN_SERIALIZED_TERMS if term in serialized] + if matched: + raise ContractError(f"structured state has forbidden fields or values: {matched}") + + +def build_shadow_state_records( + dataset: CanonicalDataset, + probabilities: np.ndarray, + *, + encoder_provenance: dict[str, Any], +) -> list[dict[str, Any]]: + """Create deterministic, read-only states from canonical offline outputs.""" + provenance = _validate_encoder_provenance(encoder_provenance) + values = _validate_probabilities(probabilities, len(dataset.labels), len(dataset.label_order)) + records: list[dict[str, Any]] = [] + dataset_sha256 = dataset.artifact_sha256() + for row, probability in enumerate(values): + top_index = int(np.argmax(probability)) + observable_probabilities = { + label: float(probability[index]) + for index, label in enumerate(dataset.label_order) + } + state = { + "schema_version": STATE_SCHEMA, + **_REQUIRED_GLOBAL_STATUS, + "state_id": _state_id(dataset_sha256, str(dataset.raw_window_hashes[row])), + "source": { + "dataset_sha256": dataset_sha256, + "source_manifest_sha256": dataset.source_manifest_sha256, + "raw_window_sha256": str(dataset.raw_window_hashes[row]), + "session_id": str(dataset.session_ids[row]), + "start_sample": int(dataset.start_samples[row]), + }, + "encoder": provenance, + "observable_state": { + "probabilities": observable_probabilities, + "argmax_observable": dataset.label_order[top_index], + "confidence": float(probability[top_index]), + "entropy_nats": _entropy(probability), + }, + } + _require_no_forbidden_serialized_terms(state) + records.append(state) + return records + + +def write_shadow_state_artifacts( + dataset: CanonicalDataset, + probabilities: np.ndarray, + *, + encoder_provenance: dict[str, Any], + states_output: Path, + manifest_output: Path, +) -> dict[str, Any]: + """Persist canonical JSONL records and a hash-bound replay manifest. + + The pair is the unit of publication, not either file. Both payloads are + built and validated before anything is written, and the manifest is written + last as the commit marker — a reader that finds no manifest knows the + records were never committed. + + Republication is idempotent, never destructive: byte-identical inputs are a + no-op, and any differing existing file is refused with both files left + untouched. An interrupted first publication (records present, manifest + absent) can be completed, but only when the existing records match exactly. + """ + records = build_shadow_state_records(dataset, probabilities, encoder_provenance=encoder_provenance) + manifest = { + "schema_version": MANIFEST_SCHEMA, + "state_schema_version": STATE_SCHEMA, + "bridge_version": STATE_BRIDGE_VERSION, + **_REQUIRED_GLOBAL_STATUS, + "record_count": len(records), + "records_sha256": _sha256_lines(records), + "dataset_sha256": dataset.artifact_sha256(), + "source_manifest_sha256": dataset.source_manifest_sha256, + "encoder": _validate_encoder_provenance(encoder_provenance), + "input_probability_sha256": sha256_json(np.asarray(probabilities, dtype=np.float64).tolist()), + } + _require_no_forbidden_serialized_terms(manifest) + if states_output.resolve() == manifest_output.resolve(): + raise ContractError("state and manifest outputs must be distinct paths") + + states_text = "".join(f"{_canonical_json(record)}\n" for record in records) + manifest_text = json.dumps(manifest, indent=2, sort_keys=True) + "\n" + states_settled = _settled(states_output, states_text, "state records") + manifest_settled = _settled(manifest_output, manifest_text, "state manifest") + if states_settled and manifest_settled: + return manifest + + states_output.parent.mkdir(parents=True, exist_ok=True) + manifest_output.parent.mkdir(parents=True, exist_ok=True) + if not states_settled: + _publish_text(states_output, states_text) + _publish_text(manifest_output, manifest_text) + return manifest + + +def load_shadow_state_records(states_path: Path, manifest_path: Path) -> list[dict[str, Any]]: + """Read records for deterministic replay after validating their complete contract.""" + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ContractError(f"cannot read state manifest: {exc}") from exc + if manifest.get("schema_version") != MANIFEST_SCHEMA: + raise ContractError("state manifest has the wrong schema") + for key, expected in _REQUIRED_GLOBAL_STATUS.items(): + if manifest.get(key) != expected: + raise ContractError(f"state manifest {key} violates shadow-only status") + try: + records = [json.loads(line) for line in states_path.read_text(encoding="utf-8").splitlines() if line] + except (OSError, json.JSONDecodeError) as exc: + raise ContractError(f"cannot read state records: {exc}") from exc + if len(records) != manifest.get("record_count"): + raise ContractError("state record count does not match manifest") + if _sha256_lines(records) != manifest.get("records_sha256"): + raise ContractError("state records do not match manifest hash") + # The hash above only proves the records are the ones this manifest was + # written for. It says nothing about whether they *agree* with it, so a + # rehashed file could still claim a different encoder or dataset than the + # manifest it ships with. Bind the two together explicitly. + manifest_encoder = _validate_encoder_provenance(manifest.get("encoder")) + manifest_dataset_sha256 = _require_sha256(manifest.get("dataset_sha256"), "manifest dataset_sha256") + manifest_source_sha256 = _require_sha256( + manifest.get("source_manifest_sha256"), "manifest source_manifest_sha256" + ) + seen_ids: set[str] = set() + for record in records: + if record.get("schema_version") != STATE_SCHEMA: + raise ContractError("state record has the wrong schema") + for key, expected in _REQUIRED_GLOBAL_STATUS.items(): + if record.get(key) != expected: + raise ContractError(f"state record {key} violates shadow-only status") + source = record.get("source") + observable = record.get("observable_state") + if not isinstance(source, dict) or not isinstance(observable, dict): + raise ContractError("state record needs source and observable_state objects") + if record.get("encoder") != manifest_encoder: + raise ContractError("state record encoder does not match the manifest") + if source.get("dataset_sha256") != manifest_dataset_sha256: + raise ContractError("state record dataset_sha256 does not match the manifest") + if source.get("source_manifest_sha256") != manifest_source_sha256: + raise ContractError("state record source_manifest_sha256 does not match the manifest") + raw_window_sha256 = _require_sha256(source.get("raw_window_sha256"), "raw_window_sha256") + state_id = record.get("state_id") + if state_id in seen_ids: + raise ContractError("state ids must be unique SHA-256 values") + if _require_sha256(state_id, "state_id") != _state_id(manifest_dataset_sha256, raw_window_sha256): + raise ContractError("state_id is not derived from its dataset and raw window") + seen_ids.add(state_id) + probabilities = observable.get("probabilities") + if not isinstance(probabilities, dict) or not probabilities: + raise ContractError("state record needs observable probabilities") + values = np.asarray(list(probabilities.values()), dtype=np.float64) + _validate_probabilities(values.reshape(1, -1), 1, len(values)) + confidence = observable.get("confidence") + entropy = observable.get("entropy_nats") + if not isinstance(confidence, (int, float)) or not math.isfinite(float(confidence)): + raise ContractError("state confidence must be finite") + if not isinstance(entropy, (int, float)) or not math.isfinite(float(entropy)): + raise ContractError("state entropy must be finite") + if observable.get("argmax_observable") not in probabilities: + raise ContractError("state argmax must be one of its observable probabilities") + selected_probability = float(probabilities[observable["argmax_observable"]]) + maximum_probability = float(values.max()) + if not math.isclose(selected_probability, maximum_probability, rel_tol=0.0, abs_tol=1e-12): + raise ContractError("state argmax does not select a maximum observable probability") + if not math.isclose(float(confidence), maximum_probability, rel_tol=0.0, abs_tol=1e-12): + raise ContractError("state confidence does not match its maximum observable probability") + if not math.isclose(float(entropy), _entropy(values), rel_tol=0.0, abs_tol=1e-12): + raise ContractError("state entropy does not match its observable probabilities") + _require_no_forbidden_serialized_terms(record) + return records + + +def _load_prediction_archive(path: Path, dataset: CanonicalDataset) -> np.ndarray: + try: + with np.load(path, allow_pickle=False) as archive: + hashes = archive["raw_window_hashes"].astype("U64") + probabilities = archive["probabilities"].astype(np.float64) + except (OSError, KeyError, ValueError) as exc: + raise ContractError(f"invalid offline prediction archive: {exc}") from exc + if len(set(hashes.tolist())) != len(hashes) or set(hashes.tolist()) != set(dataset.raw_window_hashes.tolist()): + raise ContractError("prediction archive must cover each canonical window exactly once") + mapped = {window_hash: row for window_hash, row in zip(hashes.tolist(), probabilities, strict=True)} + return np.vstack([mapped[window_hash] for window_hash in dataset.raw_window_hashes.tolist()]) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", required=True, type=Path) + parser.add_argument("--probabilities", required=True, type=Path) + parser.add_argument("--encoder-provenance", required=True, type=Path) + parser.add_argument("--states-output", required=True, type=Path) + parser.add_argument("--manifest-output", required=True, type=Path) + parser.add_argument("--verify", action="store_true") + args = parser.parse_args(argv) + dataset = load_dataset(args.dataset) + try: + encoder_provenance = json.loads(args.encoder_provenance.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise SystemExit(f"cannot read encoder provenance: {exc}") from exc + probabilities = _load_prediction_archive(args.probabilities, dataset) + manifest = write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=encoder_provenance, + states_output=args.states_output, + manifest_output=args.manifest_output, + ) + if args.verify: + load_shadow_state_records(args.states_output, args.manifest_output) + print(f"wrote {args.states_output} and {args.manifest_output}; status={manifest['status']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/NeuralComposeEEG/tests/test_structured_state.py b/NeuralComposeEEG/tests/test_structured_state.py new file mode 100644 index 0000000..d6d884b --- /dev/null +++ b/NeuralComposeEEG/tests/test_structured_state.py @@ -0,0 +1,407 @@ +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from neuralcompose_eeg.contracts import ContractError +from neuralcompose_eeg.dataset import build_canonical_dataset +from neuralcompose_eeg.structured_state import ( + MANIFEST_SCHEMA, + STATE_SCHEMA, + _sha256_lines, + build_shadow_state_records, + load_shadow_state_records, + write_shadow_state_artifacts, +) + +from .test_pipeline import _manifest, _preprocessing_path + + +class StructuredStateTests(unittest.TestCase): + def _dataset_and_probabilities(self, root: Path): + dataset, _ = build_canonical_dataset(_manifest(root), _preprocessing_path()) + values = np.zeros((len(dataset.labels), len(dataset.label_order)), dtype=np.float64) + values[:, 0] = 0.7 + values[:, 1] = 0.2 + values[:, 2] = 0.1 + return dataset, values + + @staticmethod + def _encoder_provenance() -> dict[str, str]: + return { + "model_id": "fixture-shadow-encoder", + "model_revision": "fixture-v0", + "source_kind": "synthetic_contract_fixture", + } + + def test_round_trip_is_deterministic_and_shadow_only(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + states = root / "shadow" / "states.jsonl" + manifest = root / "shadow" / "manifest.json" + first = write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + first_state_bytes = states.read_bytes() + first_manifest_bytes = manifest.read_bytes() + second = write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + replay = load_shadow_state_records(states, manifest) + self.assertEqual(first, second) + self.assertEqual(first_state_bytes, states.read_bytes()) + self.assertEqual(first_manifest_bytes, manifest.read_bytes()) + self.assertEqual(first["schema_version"], MANIFEST_SCHEMA) + self.assertEqual(first["status"], "insufficient_evidence") + self.assertEqual(first["science_status"], "pipeline_only") + self.assertFalse(first["live_control"]) + self.assertEqual(len(replay), len(dataset.labels)) + self.assertTrue(all(record["schema_version"] == STATE_SCHEMA for record in replay)) + serialized = states.read_text().casefold() + for forbidden in ("dialogue", "dialectic", "prompt", "speech", "policy", "action", "waveform", "target"): + self.assertNotIn(forbidden, serialized) + + def test_replay_rejects_tampered_record(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + states = root / "shadow" / "states.jsonl" + manifest = root / "shadow" / "manifest.json" + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + rows = states.read_text().splitlines() + row = json.loads(rows[0]) + row["observable_state"]["confidence"] = 0.01 + rows[0] = json.dumps(row, sort_keys=True, separators=(",", ":")) + states.write_text("\n".join(rows) + "\n") + with self.assertRaisesRegex(ContractError, "manifest hash"): + load_shadow_state_records(states, manifest) + + def test_replay_rejects_self_inconsistent_confidence_even_with_rehashed_manifest(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + states = root / "shadow" / "states.jsonl" + manifest = root / "shadow" / "manifest.json" + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + records = [json.loads(line) for line in states.read_text().splitlines()] + records[0]["observable_state"]["confidence"] = 0.1 + states.write_text("".join(json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" for record in records)) + manifest_value = json.loads(manifest.read_text()) + manifest_value["records_sha256"] = _sha256_lines(records) + manifest.write_text(json.dumps(manifest_value, sort_keys=True) + "\n") + with self.assertRaisesRegex(ContractError, "confidence does not match"): + load_shadow_state_records(states, manifest) + + def test_prohibits_dialogue_or_policy_provenance(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + with self.assertRaisesRegex(ContractError, "unsupported fields"): + build_shadow_state_records( + dataset, + probabilities, + encoder_provenance={**self._encoder_provenance(), "dialogue_source": "forbidden"}, + ) + + def test_rejects_invalid_probability_rows(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + probabilities[0, 0] = 0.6 + with self.assertRaisesRegex(ContractError, "sum to one"): + build_shadow_state_records( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + ) + + def test_rejects_probability_shape_mismatch(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + for malformed in (probabilities[:, :-1], probabilities[0]): + with self.assertRaisesRegex(ContractError, "shape"): + build_shadow_state_records( + dataset, + malformed, + encoder_provenance=self._encoder_provenance(), + ) + + def test_rejects_nonfinite_probabilities(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for bad in (np.nan, np.inf): + dataset, probabilities = self._dataset_and_probabilities(root) + probabilities[0, 0] = bad + with self.assertRaisesRegex(ContractError, r"finite values in \[0, 1\]"): + build_shadow_state_records( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + ) + + def test_records_dataset_and_encoder_provenance(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + records = build_shadow_state_records( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + ) + for record in records: + self.assertEqual(record["encoder"], self._encoder_provenance()) + self.assertEqual(record["source"]["dataset_sha256"], dataset.artifact_sha256()) + self.assertEqual(record["source"]["source_manifest_sha256"], dataset.source_manifest_sha256) + self.assertEqual(len(record["source"]["raw_window_sha256"]), 64) + self.assertEqual(len({record["state_id"] for record in records}), len(records)) + + def _rehash(self, states: Path, manifest: Path, records: list[dict]) -> None: + """Republish tampered records with a manifest hash that matches them, so + the reader cannot fall back on the hash check to reject the content.""" + states.write_text( + "".join(json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" for record in records) + ) + manifest_value = json.loads(manifest.read_text()) + manifest_value["records_sha256"] = _sha256_lines(records) + manifest.write_text(json.dumps(manifest_value, sort_keys=True) + "\n") + + def _published(self, root: Path) -> tuple[Path, Path, list[dict]]: + dataset, probabilities = self._dataset_and_probabilities(root) + states = root / "shadow" / "states.jsonl" + manifest = root / "shadow" / "manifest.json" + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + return states, manifest, [json.loads(line) for line in states.read_text().splitlines()] + + def test_replay_rejects_manifest_granting_live_control_or_promotion(self) -> None: + for key, granted in (("live_control", True), ("promotion_status", "eligible")): + with self.subTest(key=key), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + states, manifest, _ = self._published(root) + manifest_value = json.loads(manifest.read_text()) + manifest_value[key] = granted + manifest.write_text(json.dumps(manifest_value, sort_keys=True) + "\n") + with self.assertRaisesRegex(ContractError, f"{key} violates shadow-only status"): + load_shadow_state_records(states, manifest) + + def test_replay_rejects_record_granting_live_control(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + states, manifest, records = self._published(root) + records[0]["live_control"] = True + self._rehash(states, manifest, records) + with self.assertRaisesRegex(ContractError, "live_control violates shadow-only status"): + load_shadow_state_records(states, manifest) + + def test_replay_rejects_argmax_that_is_not_the_maximum(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + states, manifest, records = self._published(root) + observable = records[0]["observable_state"] + lowest = min(observable["probabilities"], key=observable["probabilities"].get) + observable["argmax_observable"] = lowest + observable["confidence"] = observable["probabilities"][lowest] + self._rehash(states, manifest, records) + with self.assertRaisesRegex(ContractError, "argmax does not select a maximum"): + load_shadow_state_records(states, manifest) + + def test_publication_leaves_no_partial_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + published = root / "published" + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=published / "states.jsonl", + manifest_output=published / "manifest.json", + ) + self.assertEqual( + sorted(path.name for path in published.iterdir()), + ["manifest.json", "states.jsonl"], + ) + + def test_rejected_republication_preserves_the_previous_artifact(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + states = root / "shadow" / "states.jsonl" + manifest = root / "shadow" / "manifest.json" + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + previous = states.read_bytes() + probabilities[0, 0] = np.nan + with self.assertRaises(ContractError): + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + self.assertEqual(states.read_bytes(), previous) + self.assertEqual([path.name for path in states.parent.iterdir() if ".partial" in path.name], []) + load_shadow_state_records(states, manifest) + + # ── record/manifest provenance agreement ───────────────────────────── + # + # records_sha256 only proves the records are the ones the manifest was + # written for. Each case below rehashes the manifest so that check passes, + # leaving the agreement checks as the only thing standing between a + # mismatched artifact and a successful replay. + + def test_replay_rejects_record_encoder_mismatch_with_manifest(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + states, manifest, records = self._published(root) + records[0]["encoder"] = {**self._encoder_provenance(), "model_id": "substituted-encoder"} + self._rehash(states, manifest, records) + with self.assertRaisesRegex(ContractError, "encoder does not match the manifest"): + load_shadow_state_records(states, manifest) + + def test_replay_rejects_record_source_mismatch_with_manifest(self) -> None: + for field in ("dataset_sha256", "source_manifest_sha256"): + with self.subTest(field=field), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + states, manifest, records = self._published(root) + records[0]["source"][field] = "0" * 64 + self._rehash(states, manifest, records) + with self.assertRaisesRegex(ContractError, f"{field} does not match the manifest"): + load_shadow_state_records(states, manifest) + + def test_replay_rejects_invalid_or_unrecomputed_state_id(self) -> None: + cases = { + "not_hexadecimal": ("z" * 64, "state_id must be a lowercase SHA-256"), + "wrong_length": ("abc123", "state_id must be a lowercase SHA-256"), + "valid_hex_but_not_derived": ("a" * 64, "not derived from its dataset and raw window"), + } + for name, (state_id, expected) in cases.items(): + with self.subTest(case=name), tempfile.TemporaryDirectory() as directory: + root = Path(directory) + states, manifest, records = self._published(root) + records[0]["state_id"] = state_id + self._rehash(states, manifest, records) + with self.assertRaisesRegex(ContractError, expected): + load_shadow_state_records(states, manifest) + + def test_replay_rejects_raw_window_that_is_not_a_digest(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + states, manifest, records = self._published(root) + records[0]["source"]["raw_window_sha256"] = "not-a-digest" + self._rehash(states, manifest, records) + with self.assertRaisesRegex(ContractError, "raw_window_sha256 must be a lowercase SHA-256"): + load_shadow_state_records(states, manifest) + + # ── the pair is the unit of publication ────────────────────────────── + + def test_differing_republication_is_refused_and_preserves_pair(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + states = root / "shadow" / "states.jsonl" + manifest = root / "shadow" / "manifest.json" + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + published = (states.read_bytes(), manifest.read_bytes()) + # A different but entirely valid run over the same output paths. + probabilities[:, 0], probabilities[:, 2] = 0.1, 0.7 + with self.assertRaisesRegex(ContractError, "refusing to overwrite"): + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + self.assertEqual((states.read_bytes(), manifest.read_bytes()), published) + self.assertEqual([path.name for path in states.parent.iterdir() if ".partial" in path.name], []) + load_shadow_state_records(states, manifest) + + def test_matching_states_only_publication_can_complete_manifest(self) -> None: + """An interrupted first publication — records committed, manifest never + written — is completable, because the records on disk match exactly.""" + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + states = root / "shadow" / "states.jsonl" + manifest = root / "shadow" / "manifest.json" + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + committed = states.read_bytes() + manifest.unlink() + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=states, + manifest_output=manifest, + ) + self.assertEqual(states.read_bytes(), committed) + self.assertEqual(len(load_shadow_state_records(states, manifest)), len(dataset.labels)) + + def test_states_and_manifest_outputs_must_be_distinct(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + dataset, probabilities = self._dataset_and_probabilities(root) + shared = root / "artifact.json" + with self.assertRaisesRegex(ContractError, "must be distinct paths"): + write_shadow_state_artifacts( + dataset, + probabilities, + encoder_provenance=self._encoder_provenance(), + states_output=shared, + manifest_output=shared, + ) + self.assertFalse(shared.exists()) + + +if __name__ == "__main__": + unittest.main()