runtime/src/resynthesis/tokenizer_backend.py

"""Hash-bound Fastokens ownership for Resynthesis tokenizer I/O boundaries.

The native parent already owns one BPE vocabulary. Fastokens is the accelerated
executor for that exact tokenizer graph; it neither changes token IDs nor
replaces the trained VGE/bit-tokenizer that consumes those IDs. This module is
an explicit external I/O adapter, not a model hot path.
"""

from __future__ import annotations

import contextlib
import hashlib
import importlib
import io
import json
import sys
import threading
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType
from typing import Any

from resynthesis.config import (
    RESYNTHESIS_FASTOKENS_COMMIT,
    RESYNTHESIS_FASTOKENS_RECEIPT_PATH,
    RESYNTHESIS_FASTOKENS_RECEIPT_SHA256,
    RESYNTHESIS_FASTOKENS_ROOT,
    RESYNTHESIS_FASTOKENS_SITE,
    RESYNTHESIS_FASTOKENS_SOURCE_ARCHIVE_SHA256,
    RESYNTHESIS_FASTOKENS_VERSION,
    RESYNTHESIS_FASTOKENS_WHEEL_SHA256,
    RESYNTHESIS_NATIVE_PARENT_ROOT,
    RESYNTHESIS_PARENT_TOKENIZER_SHA256,
    RESYNTHESIS_PARENT_TOKENIZER_VOCAB_SIZE,
    RESYNTHESIS_PROJECTION_VOCAB_SIZE,
    ResynthesisConfig,
)


@dataclass(frozen=True)
class FastokensDependencyIdentity:
    """Immutable dependency identity established before tokenizer loading."""

    version: str
    commit: str
    dependency_receipt_path: str
    dependency_receipt_sha256: str
    source_archive_sha256: str
    wheel_sha256: str
    site_path: str
    verified_installed_files: int
    transformers_version: str


@dataclass(frozen=True)
class ResynthesisTokenizerIdentity:
    """Serializable proof that the parent BPE is executing through Fastokens."""

    dependency: FastokensDependencyIdentity
    tokenizer_path: str
    tokenizer_sha256: str
    tokenizer_class: str
    backend_module: str
    backend_class: str
    vocabulary_size: int
    inherited_prefix_vocabulary_size: int
    inherited_prefix_preserved: bool

    def boundary_receipt(self) -> dict[str, object]:
        """Return JSON-ready evidence at the explicit CLI/receipt boundary."""

        return {
            "schema": "nnf.resynthesis.fastokens_boundary.v1",
            "backendOwner": "Resynthesis",
            "fastokensVersion": self.dependency.version,
            "fastokensCommit": self.dependency.commit,
            "dependencyReceiptPath": self.dependency.dependency_receipt_path,
            "dependencyReceiptSha256": (
                self.dependency.dependency_receipt_sha256
            ),
            "sourceArchiveSha256": self.dependency.source_archive_sha256,
            "wheelSha256": self.dependency.wheel_sha256,
            "sitePath": self.dependency.site_path,
            "verifiedInstalledFiles": (
                self.dependency.verified_installed_files
            ),
            "transformersVersion": self.dependency.transformers_version,
            "tokenizerPath": self.tokenizer_path,
            "tokenizerSha256": self.tokenizer_sha256,
            "tokenizerClass": self.tokenizer_class,
            "backendModule": self.backend_module,
            "backendClass": self.backend_class,
            "vocabularySize": self.vocabulary_size,
            "inheritedPrefixVocabularySize": (
                self.inherited_prefix_vocabulary_size
            ),
            "inheritedPrefixPreserved": self.inherited_prefix_preserved,
            "physicalProjectionVocabularySize": (
                RESYNTHESIS_PROJECTION_VOCAB_SIZE
            ),
            # Active undecodable projection rows shrink as successor token IDs
            # consume physical coordinates. This is receipt diagnostics only.
            "activeProjectionOnlyRows": (
                max(
                    RESYNTHESIS_PROJECTION_VOCAB_SIZE
                    - self.vocabulary_size,
                    0,
                )
            ),
            # The learned bridge retains the immutable inherited suffix as its
            # transfer source. Unlike the active gap above, this width never
            # shrinks when the lexical surface appends IDs.
            "vocabularyTransferSourceRows": (
                RESYNTHESIS_PROJECTION_VOCAB_SIZE
                - min(
                    RESYNTHESIS_PROJECTION_VOCAB_SIZE,
                    self.inherited_prefix_vocabulary_size,
                )
            ),
            "appendedLexicalRowsBeyondProjection": (
                max(
                    self.vocabulary_size
                    - RESYNTHESIS_PROJECTION_VOCAB_SIZE,
                    0,
                )
            ),
            "vocabularyTransferSourceRowsMappedByModel": True,
            "physicalProjectionVocabularyIsCeiling": False,
            "vocabularyGrowthPolicy": (
                "prefix_preserving_additive_token_ids"
            ),
            "bpeIdentityPreserved": True,
            "vgeRemainsDownstreamAuthority": True,
            "fallbackAllowed": False,
        }


_ACTIVATION_LOCK = threading.Lock()
_DEPENDENCY_IDENTITIES: dict[
    tuple[str, str, str, str, str],
    FastokensDependencyIdentity,
] = {}


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _tokenizer_artifact_vocabulary_boundary(
    tokenizer_path: Path,
) -> dict[str, int]:
    """Read one tokenizer's exact token-to-ID map at the artifact boundary."""

    payload = json.loads(tokenizer_path.read_text(encoding="utf-8"))
    if not isinstance(payload, Mapping):
        raise RuntimeError("Resynthesis tokenizer artifact is invalid")
    model = payload.get("model")
    model_vocab = (
        model.get("vocab") if isinstance(model, Mapping) else None
    )
    added_tokens = payload.get("added_tokens")
    if not isinstance(model_vocab, Mapping) or not isinstance(
        added_tokens, Sequence
    ):
        raise RuntimeError(
            "Resynthesis tokenizer artifact has no complete vocabulary"
        )
    vocabulary: dict[str, int] = {}
    for token, token_id in model_vocab.items():
        if (
            not isinstance(token, str)
            or not isinstance(token_id, int)
            or isinstance(token_id, bool)
            or token_id < 0
        ):
            raise RuntimeError(
                "Resynthesis tokenizer model vocabulary is invalid"
            )
        vocabulary[token] = token_id
    for row in added_tokens:
        if not isinstance(row, Mapping):
            raise RuntimeError(
                "Resynthesis tokenizer added-token row is invalid"
            )
        token = row.get("content")
        token_id = row.get("id")
        if (
            not isinstance(token, str)
            or not isinstance(token_id, int)
            or isinstance(token_id, bool)
            or token_id < 0
        ):
            raise RuntimeError(
                "Resynthesis tokenizer added-token identity is invalid"
            )
        prior = vocabulary.get(token)
        if prior is not None and prior != token_id:
            raise RuntimeError(
                "Resynthesis tokenizer assigns one token multiple IDs"
            )
        vocabulary[token] = token_id
    token_ids = tuple(vocabulary.values())
    if (
        not token_ids
        or len(set(token_ids)) != len(token_ids)
        or set(token_ids) != set(range(max(token_ids) + 1))
    ):
        raise RuntimeError(
            "Resynthesis tokenizer IDs are not one contiguous vocabulary"
        )
    return vocabulary


def _validated_prefix_vocabulary_size_boundary(
    tokenizer_path: Path,
    *,
    inherited_tokenizer_path: Path | None = None,
) -> int:
    """Prove that a successor preserves every inherited token ID exactly."""

    active_path = tokenizer_path.expanduser().resolve()
    if not active_path.is_file():
        raise FileNotFoundError(
            f"Resynthesis tokenizer is absent: {active_path}"
        )
    inherited_path = (
        inherited_tokenizer_path.expanduser().resolve()
        if inherited_tokenizer_path is not None
        else Path(RESYNTHESIS_NATIVE_PARENT_ROOT).resolve()
        / "tokenizer.json"
    )
    if (
        not inherited_path.is_file()
        or _sha256(inherited_path) != RESYNTHESIS_PARENT_TOKENIZER_SHA256
    ):
        raise RuntimeError(
            "Resynthesis inherited tokenizer artifact identity differs"
        )
    inherited_vocabulary = _tokenizer_artifact_vocabulary_boundary(
        inherited_path
    )
    if len(inherited_vocabulary) != RESYNTHESIS_PARENT_TOKENIZER_VOCAB_SIZE:
        raise RuntimeError(
            "Resynthesis inherited tokenizer vocabulary identity differs"
        )
    active_vocabulary = _tokenizer_artifact_vocabulary_boundary(active_path)
    if len(active_vocabulary) < len(inherited_vocabulary) or any(
        active_vocabulary.get(token) != token_id
        for token, token_id in inherited_vocabulary.items()
    ):
        raise RuntimeError(
            "Resynthesis successor tokenizer changed its inherited prefix"
        )
    return len(active_vocabulary)


def _required_string(value: object, *, field: str) -> str:
    if not isinstance(value, str) or not value:
        raise RuntimeError(f"Fastokens receipt field {field} is absent")
    return value


def _required_mapping(value: object, *, field: str) -> Mapping[str, object]:
    if not isinstance(value, Mapping):
        raise RuntimeError(f"Fastokens receipt field {field} is malformed")
    return value


def _verified_artifact(
    record: Mapping[str, object],
    *,
    field: str,
    expected_root: Path,
    artifact_path: Path | None = None,
    expected_sha256: str | None = None,
) -> str:
    declared_path = Path(
        _required_string(record.get("path"), field=f"{field}.path")
    ).expanduser().resolve()
    try:
        declared_path.relative_to(expected_root)
    except ValueError as exc:
        raise RuntimeError(f"Fastokens {field} escaped its immutable root") from exc
    declared_sha256 = _required_string(
        record.get("sha256"),
        field=f"{field}.sha256",
    )
    if expected_sha256 is not None and declared_sha256 != expected_sha256:
        raise RuntimeError(f"Fastokens {field} identity differs from release")
    path = (
        artifact_path.expanduser().resolve()
        if artifact_path is not None
        else declared_path
    )
    declared_bytes = record.get("bytes")
    if (
        not isinstance(declared_bytes, int)
        or isinstance(declared_bytes, bool)
        or declared_bytes < 1
        or not path.is_file()
        or path.stat().st_size != declared_bytes
        or _sha256(path) != declared_sha256
    ):
        raise RuntimeError(f"Fastokens {field} bytes differ from their receipt")
    return declared_sha256


def _verified_dependency(
    cfg: ResynthesisConfig,
) -> FastokensDependencyIdentity:
    """Verify one canonical dependency at its configured release coordinates."""

    root = Path(RESYNTHESIS_FASTOKENS_ROOT).resolve()
    canonical_site = Path(RESYNTHESIS_FASTOKENS_SITE).resolve()
    site = Path(cfg.fastokens_site_path).expanduser().resolve()
    receipt_path = Path(cfg.fastokens_receipt_path).expanduser().resolve()
    wheel_path = Path(cfg.fastokens_wheel_path).expanduser().resolve()
    if not receipt_path.is_file():
        raise FileNotFoundError(
            f"Resynthesis Fastokens dependency receipt is absent: {receipt_path}"
        )
    receipt_sha256 = _sha256(receipt_path)
    if (
        cfg.fastokens_receipt_sha256 != RESYNTHESIS_FASTOKENS_RECEIPT_SHA256
        or receipt_sha256 != cfg.fastokens_receipt_sha256
        or cfg.fastokens_wheel_sha256 != RESYNTHESIS_FASTOKENS_WHEEL_SHA256
        or cfg.fastokens_source_archive_sha256
        != RESYNTHESIS_FASTOKENS_SOURCE_ARCHIVE_SHA256
    ):
        raise RuntimeError("Resynthesis Fastokens dependency receipt has drifted")
    loaded = json.loads(receipt_path.read_text(encoding="utf-8"))
    receipt = _required_mapping(loaded, field="root")
    if (
        receipt.get("schema") != "nnf.resynthesis.fastokens_dependency.v1"
        or receipt.get("name") != "fastokens"
        or receipt.get("version") != RESYNTHESIS_FASTOKENS_VERSION
        or receipt.get("upstreamCommit") != RESYNTHESIS_FASTOKENS_COMMIT
        or receipt.get("buildPassed") is not True
    ):
        raise RuntimeError("Resynthesis Fastokens dependency identity differs")
    installed_site = Path(
        _required_string(receipt.get("installedSite"), field="installedSite")
    ).resolve()
    if installed_site != canonical_site or not site.is_dir():
        raise RuntimeError("Resynthesis Fastokens installed site differs")

    source_record = _required_mapping(
        receipt.get("sourceArchive"),
        field="sourceArchive",
    )
    source_path = Path(
        _required_string(
            source_record.get("path"),
            field="sourceArchive.path",
        )
    ).expanduser().resolve()
    source_sha256 = _required_string(
        source_record.get("sha256"),
        field="sourceArchive.sha256",
    )
    if (
        source_path != root / "source.tar.gz"
        or source_sha256 != cfg.fastokens_source_archive_sha256
    ):
        raise RuntimeError("Fastokens source archive identity differs from release")
    if receipt_path == Path(RESYNTHESIS_FASTOKENS_RECEIPT_PATH).resolve():
        source_sha256 = _verified_artifact(
            source_record,
            field="source archive",
            expected_root=root,
            expected_sha256=cfg.fastokens_source_archive_sha256,
        )
    wheel_sha256 = _verified_artifact(
        _required_mapping(receipt.get("wheel"), field="wheel"),
        field="wheel",
        expected_root=root,
        artifact_path=wheel_path,
        expected_sha256=cfg.fastokens_wheel_sha256,
    )

    installed_files = receipt.get("installedFiles")
    if (
        not isinstance(installed_files, Sequence)
        or isinstance(installed_files, (str, bytes))
        or not installed_files
    ):
        raise RuntimeError("Fastokens installed-file receipt is malformed")
    observed_paths: set[str] = set()
    for index, raw_record in enumerate(installed_files):
        record = _required_mapping(raw_record, field=f"installedFiles[{index}]")
        relative = _required_string(
            record.get("path"),
            field=f"installedFiles[{index}].path",
        )
        if relative in observed_paths:
            raise RuntimeError("Fastokens installed-file receipt has duplicates")
        observed_paths.add(relative)
        candidate = (site / relative).resolve()
        try:
            candidate.relative_to(site)
        except ValueError as exc:
            raise RuntimeError("Fastokens installed file escaped its site") from exc
        expected_sha256 = _required_string(
            record.get("sha256"),
            field=f"installedFiles[{index}].sha256",
        )
        if not candidate.is_file() or _sha256(candidate) != expected_sha256:
            raise RuntimeError(
                f"Fastokens installed file differs from its receipt: {relative}"
            )
    required_runtime_files = {
        "fastokens/__init__.py",
        "fastokens/_compat.py",
        "fastokens/_native.abi3.so",
        "fastokens-0.2.1.dist-info/METADATA",
    }
    if not required_runtime_files.issubset(observed_paths):
        raise RuntimeError("Fastokens receipt omits required runtime files")

    site_string = str(site)
    if site_string not in sys.path:
        sys.path.insert(0, site_string)
    importlib.invalidate_caches()
    module = importlib.import_module("fastokens")
    _assert_module_owned_by_site(module, site)
    patch = getattr(module, "patch_transformers", None)
    if not callable(patch):
        raise RuntimeError("Fastokens patch_transformers is unavailable")
    patch_output = io.StringIO()
    with contextlib.redirect_stdout(patch_output):
        patch()
    if getattr(module, "_patched", None) is not True:
        raise RuntimeError("Fastokens did not retain Transformers patch authority")

    transformers = importlib.import_module("transformers")
    transformers_version = getattr(transformers, "__version__", None)
    if not isinstance(transformers_version, str) or not transformers_version:
        raise RuntimeError("Transformers version identity is unavailable")
    return FastokensDependencyIdentity(
        version=RESYNTHESIS_FASTOKENS_VERSION,
        commit=RESYNTHESIS_FASTOKENS_COMMIT,
        dependency_receipt_path=str(receipt_path),
        dependency_receipt_sha256=receipt_sha256,
        source_archive_sha256=source_sha256,
        wheel_sha256=wheel_sha256,
        site_path=str(site),
        verified_installed_files=len(installed_files),
        transformers_version=transformers_version,
    )


def _assert_module_owned_by_site(module: ModuleType, site: Path) -> None:
    module_path_raw = getattr(module, "__file__", None)
    if not isinstance(module_path_raw, str):
        raise RuntimeError("Fastokens module has no filesystem identity")
    module_path = Path(module_path_raw).resolve()
    try:
        module_path.relative_to(site)
    except ValueError as exc:
        raise RuntimeError("Fastokens was imported outside the pinned site") from exc


def activate_resynthesis_fastokens(
    cfg: ResynthesisConfig | None = None,
) -> FastokensDependencyIdentity:
    """Verify and activate the one pinned tokenizer backend for this process."""

    active_cfg = cfg or ResynthesisConfig()
    location_key = (
        str(Path(active_cfg.fastokens_receipt_path).expanduser().resolve()),
        active_cfg.fastokens_receipt_sha256,
        str(Path(active_cfg.fastokens_site_path).expanduser().resolve()),
        str(Path(active_cfg.fastokens_wheel_path).expanduser().resolve()),
        active_cfg.fastokens_wheel_sha256,
    )
    with _ACTIVATION_LOCK:
        identity = _DEPENDENCY_IDENTITIES.get(location_key)
        if identity is None:
            identity = _verified_dependency(active_cfg)
            _DEPENDENCY_IDENTITIES[location_key] = identity
        return identity


def _resolve_configured_tokenizer_path(cfg: ResynthesisConfig) -> Path:
    tokenizer_path = Path(cfg.tokenizer_path).expanduser()
    if not tokenizer_path.is_absolute():
        tokenizer_path = Path(cfg.base_model_dir).expanduser() / tokenizer_path
    return tokenizer_path.resolve()


def load_resynthesis_tokenizer(
    model_dir: str | Path = RESYNTHESIS_NATIVE_PARENT_ROOT,
    *,
    cfg: ResynthesisConfig | None = None,
    local_files_only: bool = True,
    trust_remote_code: bool = True,
) -> Any:
    """Load a prefix-preserving Resynthesis BPE through Fastokens.

    The inherited tokenizer is the immutable ID prefix. A successor may append
    lexical rows, but it may not renumber, replace, or remove any inherited
    token. The inherited physical projection width is therefore geometry, not
    a vocabulary ceiling; model-owned transfer tensors learn appended rows.
    """

    resolved_model_dir = Path(model_dir).expanduser().resolve()
    tokenizer_path = resolved_model_dir / "tokenizer.json"
    if not tokenizer_path.is_file():
        raise FileNotFoundError(f"Resynthesis tokenizer is absent: {tokenizer_path}")
    tokenizer_sha256 = _sha256(tokenizer_path)
    if (
        resolved_model_dir == Path(RESYNTHESIS_NATIVE_PARENT_ROOT).resolve()
        and tokenizer_sha256 != RESYNTHESIS_PARENT_TOKENIZER_SHA256
    ):
        raise RuntimeError("Resynthesis native-parent tokenizer has drifted")
    inherited_tokenizer_path = (
        _resolve_configured_tokenizer_path(cfg)
        if cfg is not None
        else Path(RESYNTHESIS_NATIVE_PARENT_ROOT).resolve()
        / "tokenizer.json"
    )
    _validated_prefix_vocabulary_size_boundary(
        tokenizer_path,
        inherited_tokenizer_path=inherited_tokenizer_path,
    )

    dependency = activate_resynthesis_fastokens(cfg)
    transformers = importlib.import_module("transformers")
    auto_tokenizer = getattr(transformers, "AutoTokenizer", None)
    if auto_tokenizer is None:
        raise RuntimeError("Transformers AutoTokenizer is unavailable")
    tokenizer = auto_tokenizer.from_pretrained(
        str(resolved_model_dir),
        local_files_only=local_files_only,
        trust_remote_code=trust_remote_code,
    )
    identity = tokenizer_identity(
        tokenizer,
        dependency=dependency,
        tokenizer_path=tokenizer_path,
        inherited_tokenizer_path=inherited_tokenizer_path,
    )
    if identity.tokenizer_sha256 != tokenizer_sha256:
        raise RuntimeError("Fastokens tokenizer identity did not bind parent bytes")
    setattr(tokenizer, "_resynthesis_tokenizer_identity", identity)
    return tokenizer


def tokenizer_identity(
    tokenizer: Any,
    *,
    dependency: FastokensDependencyIdentity | None = None,
    tokenizer_path: str | Path | None = None,
    inherited_tokenizer_path: str | Path | None = None,
) -> ResynthesisTokenizerIdentity:
    """Verify a Fastokens tokenizer and its immutable inherited ID prefix."""

    active_dependency = dependency or activate_resynthesis_fastokens()
    backend = getattr(tokenizer, "backend_tokenizer", None)
    backend_module = type(backend).__module__
    backend_class = type(backend).__name__
    if backend_module != "fastokens._compat" or backend_class != "_TokenizerShim":
        raise RuntimeError("Resynthesis tokenizer did not load through Fastokens")
    vocabulary_size = len(tokenizer)
    if tokenizer_path is None:
        loaded_root = getattr(tokenizer, "name_or_path", None)
        loaded_path = (
            Path(loaded_root).expanduser().resolve() / "tokenizer.json"
            if isinstance(loaded_root, str) and loaded_root
            else Path(RESYNTHESIS_NATIVE_PARENT_ROOT).resolve()
            / "tokenizer.json"
        )
        active_tokenizer_path = (
            loaded_path
            if loaded_path.is_file()
            else Path(RESYNTHESIS_NATIVE_PARENT_ROOT).resolve()
            / "tokenizer.json"
        )
    else:
        active_tokenizer_path = Path(tokenizer_path).expanduser().resolve()
    artifact_vocabulary_size = (
        _validated_prefix_vocabulary_size_boundary(
            active_tokenizer_path,
            inherited_tokenizer_path=(
                Path(inherited_tokenizer_path)
                if inherited_tokenizer_path is not None
                else None
            ),
        )
    )
    if vocabulary_size != artifact_vocabulary_size:
        raise RuntimeError(
            "Resynthesis Fastokens vocabulary differs from its artifact"
        )
    tokenizer_sha256 = _sha256(active_tokenizer_path)
    if (
        active_tokenizer_path
        == Path(RESYNTHESIS_NATIVE_PARENT_ROOT).resolve() / "tokenizer.json"
        and tokenizer_sha256 != RESYNTHESIS_PARENT_TOKENIZER_SHA256
    ):
        raise RuntimeError("Resynthesis parent tokenizer identity changed")
    return ResynthesisTokenizerIdentity(
        dependency=active_dependency,
        tokenizer_path=str(active_tokenizer_path),
        tokenizer_sha256=tokenizer_sha256,
        tokenizer_class=type(tokenizer).__name__,
        backend_module=backend_module,
        backend_class=backend_class,
        vocabulary_size=vocabulary_size,
        inherited_prefix_vocabulary_size=(
            RESYNTHESIS_PARENT_TOKENIZER_VOCAB_SIZE
        ),
        inherited_prefix_preserved=True,
    )


def tokenizer_boundary_receipt(tokenizer: Any) -> dict[str, object]:
    """Serialize verified backend evidence at a durable receipt boundary."""

    attached = getattr(tokenizer, "_resynthesis_tokenizer_identity", None)
    if isinstance(attached, ResynthesisTokenizerIdentity):
        identity = attached
    else:
        identity = tokenizer_identity(tokenizer)
    return identity.boundary_receipt()