natural-memory-nm21/memory_os_v2.py

"""Hierarchical, pageable memory components for Natural Memory v2.

The module deliberately keeps storage and neural routing separate:

* :class:`MemoryRouterV2` is the trainable sparse router.
* :class:`PagedMemoryBankV2` owns versioned pages and records.
* :class:`MemoryOSV2` coordinates writing, reading, quarantine and repair.
* :class:`KVBudgetManagerV2` models the hot-context budget.

The bank can run entirely in process memory for research, or be backed by a
checkpoint/page store by serializing ``export_payload``.  No prompt text is
assembled by this module; it returns evidence records and routing traces to
the model adapter, which can inject only the selected evidence.
"""

from __future__ import annotations

import hashlib
import json
import math
import re
import time
from dataclasses import asdict, dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, Sequence

import torch
from torch import Tensor, nn
import torch.nn.functional as F

if TYPE_CHECKING:
    from .tiered_memory_store_v2 import TieredMemoryStoreV2


STATUS_ACTIVE = "active"
STATUS_SUPERSEDED = "superseded"
STATUS_RETRACTED = "retracted"
STATUS_QUARANTINED = "quarantined"


def _now() -> int:
    return int(time.time())


def _stable_id(text: str, *, prefix: str) -> str:
    digest = hashlib.sha1(text.encode("utf-8")).hexdigest()[:16]
    return f"{prefix}_{digest}"


def _tokens(text: str) -> set[str]:
    raw = {
        token
        for token in re.findall(r"[\u4e00-\u9fff]|[A-Za-z0-9_\-]+", text.lower())
        if token not in {"我", "的", "是", "了", "请", "一下"}
    }
    # Keep the complete identifier for exact lookup, but also expose its
    # bounded components.  Real code questions usually say ``reset token``
    # while the source contains ``DEFAULT_MEMORY_RESET_TOKEN``; treating the
    # identifier as one opaque token makes that otherwise obvious evidence
    # unreachable without a dense full-corpus scan.
    subtokens = {
        part
        for token in raw
        if "_" in token or "-" in token
        for part in re.split(r"[_-]+", token)
        if len(part) >= 3
    }
    # Chinese questions and metadata are commonly written as short phrases
    # (模型名称, 外部脚本, 全量注意力) rather than whitespace-delimited words.
    # Add bounded 3-character n-grams for contiguous CJK runs so the sparse
    # index can match those phrases without turning every generic two-character
    # word (for example 候选) into a direct route.
    cjk_ngrams: set[str] = set()
    for run in re.findall(r"[\u4e00-\u9fff]+", text.lower()):
        for width in (3,):
            if len(run) < width:
                continue
            cjk_ngrams.update(run[index : index + width] for index in range(len(run) - width + 1))
    return raw | subtokens | cjk_ngrams


def _is_explicit_unknown_request(text: str) -> bool:
    """Detect a request that explicitly says the target field is unregistered.

    This is a narrow semantic fail-safe around the learned router.  It does
    not decide ordinary relevance and it never invents a value; it prevents a
    nearby, but different, personal fact from being promoted into an answer
    when the user explicitly asks about a non-existent/unregistered field.
    """

    normalized = " ".join(str(text).strip().split())
    if any(
        marker in normalized
        for marker in (
            "不存在的",
            "未登记",
            "没有登记",
            "从未记录",
            "从来没有登记",
            "没有出现在",
            "没有出现过",
        )
    ):
        return True
    # Natural user questions often express the same safety intent as a
    # conditional: do not guess when the field has no record.  Treat the
    # combination as an abstention request, without blocking ordinary uses of
    # either phrase in isolation.
    return (
        "没有记录" in normalized
        and any(marker in normalized for marker in ("不要猜", "说不知道", "明确说不知道"))
    ) or (
        "如果没有" in normalized
        and any(marker in normalized for marker in ("说不知道", "明确说不知道", "不要猜"))
    )


@dataclass
class MemoryRecordV2:
    """One versioned memory item with evidence and routing metadata."""

    record_id: str
    text: str
    key: Tensor
    summary: Tensor
    # Frozen-backbone representation used by the natural-language retriever.
    # ``key`` remains the compact page/address vector; keeping this optional
    # second view lets update/forget decisions use the same semantic space
    # after a restart without turning the address key into a full KV cache.
    semantic_key: Optional[Tensor] = None
    memory_type: str = "fact"
    entity: str = ""
    attribute: str = ""
    value: str = ""
    timestamp: int = field(default_factory=_now)
    importance: float = 0.5
    confidence: float = 0.5
    source: str = "user"
    status: str = STATUS_ACTIVE
    version: int = 0
    page_id: str = ""
    supersedes: str = ""
    related_ids: list[str] = field(default_factory=list)
    evidence: list[str] = field(default_factory=list)
    # Identifies the user turn this record was written from, as a short hash of that
    # turn's text.  One turn routinely produces several records (a sentence can carry
    # a name *and* a number, and an agent may issue one structured write per field),
    # and until now nothing linked them: an explicit "forget X" retracted only the
    # record whose key matched and left its siblings active and readable.  Empty on
    # records written before this field existed, which keeps old payloads loadable.
    origin: str = ""
    # ``slot_index`` links a V2 address to the existing hot text bank when
    # the Qwen adapter is used as a compatibility bridge.  A value of -1
    # means the record is an independent paged-memory item.
    slot_index: int = -1
    token_ids: Optional[Tensor] = None
    token_mask: Optional[Tensor] = None
    access_count: int = 0
    last_access: int = field(default_factory=_now)

    def conflict_key(self) -> str:
        if self.entity and self.attribute:
            return f"{self.entity.strip().lower()}::{self.attribute.strip().lower()}"
        return ""

    def routing_text(self) -> str:
        """Return the searchable view used by the sparse evidence router.

        The original implementation routed only on ``text``.  That loses
        structured evidence when a natural-language query uses an alias of an
        attribute or value (for example ``优先使用哪层内存`` for a record whose
        value is ``DRAM``), and it makes code records needlessly dependent on
        the exact prose used when they were written.  The fields are metadata,
        not extra prompt content, so adding them here does not increase the
        injected evidence token budget.
        """

        aliases = {
            "模型名称": "正式名字 模型名 模型叫什么 名称 名字",
            "对照模型": "原版 Qwen 基线 跑分 对照规模",
            "记忆存储优先级": "优先使用 哪层内存 容量不够 降级 DRAM RAM 显存",
            "记忆持久化方案": "第三个权重切片 数据库 SQLite 磁盘分页 持久化文件",
            "记忆读取方式": "模型自己读 模型原生 外部脚本 提示词 读取器",
            "显存安全要求": "显存原则 不要顶满 显卡 GPU 安全底线",
            "记忆质量要求": "禁止胡编乱造 写错 清理 总结 置信度",
            "研发取向": "训练 微调 应用拼装 方向 精力",
            "核心架构目标": "终极要求 模型自己拥有记忆 固定外部代码 核心",
            "记忆压缩策略": "压成一句话 高压缩率 碎片对话 浓缩",
            "KV与Slot分工": "当前上下文 长期历史 运行内存 注意力",
            "自动化边界": "模型自发完成 自己完成 外部固定参数",
            "评测重点": "大规模读取 正确率 显存 测试规模",
        }.get(self.attribute.strip(), "")
        # Operational records are grounded in real source excerpts.  Their
        # natural questions often use a concept name rather than the exact
        # sentence surrounding the excerpt, so expose a small query-facing
        # alias set in the sparse index only.  These aliases never enter the
        # prompt and therefore do not increase generation cost.
        operational_aliases = ""
        if self.attribute.strip().lower().startswith("operational:"):
            try:
                operational_index = int(self.attribute.rsplit(":", 1)[-1])
            except ValueError:
                operational_index = -1
            operational_aliases = {
                0: "重启 恢复 记忆 权重切片 memory safetensors memory shard",
                1: "百万 slot 当前 token 全部 全量注意力 top-k 稀疏读取",
                2: "写入 隔离 active quarantine 置信度 低置信度",
                3: "权威副本 热点记录 process RAM GPU cache 显存",
                4: "4B 对照 base Qwen Chunk RAG Natural Memory 实验协议",
                5: "训练 公开 cross-encoder 强 RAG 基线",
                6: "清空记忆 默认 reset token fim prefix",
                7: "语义地址 Qwen 主干 hidden state 编码",
            }.get(operational_index, "")
        return " ".join(
            value
            for value in (
                self.text,
                self.entity,
                self.attribute,
                self.value,
                aliases,
                operational_aliases,
            )
            if str(value).strip()
        )

    def lexical_score(self, query_text: str) -> float:
        query = _tokens(query_text)
        if not query:
            return 0.0
        own = _tokens(self.routing_text())
        if not own:
            return 0.0
        return len(query & own) / math.sqrt(max(1, len(query) * len(own)))

    def token_overlap_score(self, query_token_ids: Optional[Tensor]) -> float:
        """Score exact token evidence without decoding or full attention."""

        if query_token_ids is None or self.token_ids is None:
            return 0.0
        own_mask = self.token_mask
        own = self.token_ids.reshape(-1)
        if isinstance(own_mask, Tensor) and own_mask.numel() == own.numel():
            own = own[own_mask.reshape(-1).to(dtype=torch.bool)]
        query = query_token_ids.detach().reshape(-1).to(device=own.device)
        if own.numel() == 0 or query.numel() == 0:
            return 0.0
        own_unique = torch.unique(own)
        query_unique = torch.unique(query)
        shared = torch.isin(query_unique, own_unique).sum().item()
        # Query coverage is the useful direction here: a short question that
        # names an entity should strongly prefer the chunk containing that
        # entity even when the chunk itself is long.
        return float(shared) / max(1, int(query_unique.numel()))

    def touch(self) -> None:
        self.access_count += 1
        self.last_access = _now()


def memory_record_to_dict(record: MemoryRecordV2) -> dict[str, Any]:
    """Return a JSON-safe public view without exposing embedding tensors."""

    return {
        "record_id": record.record_id,
        "text": record.text,
        "memory_type": record.memory_type,
        "entity": record.entity,
        "attribute": record.attribute,
        "value": record.value,
        "timestamp": int(record.timestamp),
        "importance": float(record.importance),
        "confidence": float(record.confidence),
        "source": record.source,
        "status": record.status,
        "version": int(record.version),
        "page_id": record.page_id,
        "supersedes": record.supersedes,
        "related_ids": list(record.related_ids),
        "evidence": list(record.evidence),
        "origin": record.origin,
        "slot_index": int(record.slot_index),
        "token_count": int(record.token_ids.numel()) if isinstance(record.token_ids, Tensor) else 0,
        "key_dim": int(record.key.numel()) if isinstance(record.key, Tensor) else 0,
        "access_count": int(record.access_count),
        "last_access": int(record.last_access),
    }


@dataclass
class MemoryPageV2:
    page_id: str
    tier: str = "warm"
    capacity: int = 32
    record_ids: list[str] = field(default_factory=list)
    key: Optional[Tensor] = None
    summary: Optional[Tensor] = None
    importance: float = 0.0
    created_at: int = field(default_factory=_now)
    last_access: int = field(default_factory=_now)

    @property
    def full(self) -> bool:
        return len(self.record_ids) >= self.capacity

    def touch(self) -> None:
        self.last_access = _now()


@dataclass
class RouterDecisionV2:
    need_memory: bool
    page_ids: list[str]
    record_ids: list[str]
    page_scores: list[float]
    record_scores: list[float]
    hop_count: int
    hop_trace: list[list[str]]
    confidence: float
    stop_reason: str
    # Raw calibrated evidence diagnostics.  They are optional in spirit but
    # always populated by new decisions so an evaluation can distinguish
    # "a record was returned" from "the router had a defensible margin".
    top_score: float = 0.0
    score_margin: float = 0.0
    evidence_score: float = 0.0


class MemoryRouterV2(nn.Module):
    """Trainable multi-head router for sparse page and record retrieval.

    The router does not generate the answer.  It scores candidate memory
    units, predicts whether another hop is useful, and exposes per-head scores
    for diagnostics.  The projection weights are trained separately from the
    frozen Qwen backbone using hard-negative episodes.
    """

    def __init__(
        self,
        hidden_size: int,
        *,
        router_dim: int = 128,
        num_heads: int = 8,
        max_hops: int = 3,
    ) -> None:
        super().__init__()
        if router_dim % num_heads != 0:
            raise ValueError("router_dim must be divisible by num_heads")
        if max_hops < 1:
            raise ValueError("max_hops must be positive")
        self.hidden_size = hidden_size
        self.router_dim = router_dim
        self.num_heads = num_heads
        self.max_hops = max_hops
        self.head_dim = router_dim // num_heads
        self.query_projection = nn.Linear(hidden_size, router_dim, bias=False)
        self.key_projection = nn.Linear(hidden_size, router_dim, bias=False)
        self.pair_scorer = nn.Sequential(
            nn.Linear(router_dim * 3, router_dim),
            nn.SiLU(),
            nn.Linear(router_dim, 1),
        )
        policy_hidden = max(32, min(256, router_dim * 2))
        self.need_memory = nn.Sequential(
            nn.Linear(hidden_size, policy_hidden),
            nn.SiLU(),
            nn.Linear(policy_hidden, 1),
        )
        self.hop_controller = nn.Sequential(
            nn.Linear(hidden_size, policy_hidden),
            nn.SiLU(),
            nn.Linear(policy_hidden, max_hops + 1),
        )
        self.head_gate = nn.Linear(hidden_size, num_heads)

    def _reshape(self, value: Tensor) -> Tensor:
        return value.view(*value.shape[:-1], self.num_heads, self.head_dim)

    def encode_query(self, query: Tensor) -> Tensor:
        """Project a model hidden state into the compact address space."""

        if query.shape[-1] != self.hidden_size:
            raise ValueError(f"query last dimension must be {self.hidden_size}")
        return F.normalize(self.query_projection(query), dim=-1)

    def encode_key(self, key: Tensor) -> Tensor:
        """Project model keys into the compact address space."""

        if key.shape[-1] != self.hidden_size:
            raise ValueError(f"key last dimension must be {self.hidden_size}")
        return F.normalize(self.key_projection(key), dim=-1)

    def projected_scores(self, query: Tensor, projected_candidates: Tensor) -> tuple[Tensor, Tensor]:
        """Score compact candidates without storing full hidden states.

        This is the storage-saving path: the bank keeps only the projected
        candidate keys, while the current query is projected on demand.
        """

        if query.ndim != 2 or query.shape[-1] != self.hidden_size:
            raise ValueError("query must have shape [B, hidden_size]")
        if projected_candidates.ndim == 2:
            projected_candidates = projected_candidates.unsqueeze(0).expand(query.shape[0], -1, -1)
        if projected_candidates.ndim != 3 or projected_candidates.shape[-1] != self.router_dim:
            raise ValueError("projected_candidates must have shape [B,N,router_dim]")
        q = self._reshape(self.encode_query(query))
        k = self._reshape(F.normalize(projected_candidates, dim=-1))
        head_scores = torch.einsum("bhc,bnhc->bnh", q, k)
        gates = torch.softmax(self.head_gate(query), dim=-1)[:, None, :]
        cosine_score = (head_scores * gates).sum(dim=-1)
        q_expanded = q.reshape(query.shape[0], 1, -1).expand(-1, projected_candidates.shape[1], -1)
        k_flat = k.reshape(query.shape[0], projected_candidates.shape[1], -1)
        pair_input = torch.cat((q_expanded, k_flat, q_expanded - k_flat), dim=-1)
        learned_score = self.pair_scorer(pair_input).squeeze(-1)
        return cosine_score + learned_score, head_scores

    def pair_scores(self, query: Tensor, candidates: Tensor) -> tuple[Tensor, Tensor]:
        """Return aggregate and per-head candidate scores.

        ``query`` is ``[B,H]`` and ``candidates`` is ``[B,N,H]`` or ``[N,H]``.
        """

        if query.ndim != 2:
            raise ValueError("query must have shape [B,H]")
        if candidates.ndim == 2:
            candidates = candidates.unsqueeze(0).expand(query.shape[0], -1, -1)
        if candidates.ndim != 3 or candidates.shape[0] != query.shape[0]:
            raise ValueError("candidates must have shape [B,N,H] or [N,H]")
        return self.projected_scores(query, self.encode_key(candidates))

    def forward(self, query: Tensor, candidates: Tensor) -> dict[str, Tensor]:
        scores, head_scores = self.pair_scores(query, candidates)
        return {
            "scores": scores,
            "head_scores": head_scores,
            "need_memory_logits": self.need_memory(query).squeeze(-1),
            "hop_logits": self.hop_controller(query),
        }


class PagedMemoryBankV2:
    """Versioned memory pages with sparse routing and multi-hop expansion."""

    def __init__(
        self,
        hidden_size: int,
        *,
        page_capacity: int = 32,
        max_pages: int = 32768,
        hot_pages: int = 8,
        top_k_pages: int = 4,
        top_k_records: int = 8,
        max_hops: int = 3,
        router: Optional[MemoryRouterV2] = None,
        key_dim: Optional[int] = None,
        coarse_index_bits: int = 20,
        tier_store: Optional["TieredMemoryStoreV2"] = None,
        max_resident_pages: int = 256,
        runtime_device: Optional[torch.device] = None,
        gpu_cache_records: int = 256,
        gpu_cache_tokens: int = 131072,
        gpu_cache_reserve_mb: int = 2048,
        gpu_cache_adaptive: bool = True,
        record_scorer: Optional[Callable[[Tensor, Tensor], Tensor]] = None,
    ) -> None:
        if page_capacity < 1 or max_pages < 1:
            raise ValueError("page_capacity and max_pages must be positive")
        self.hidden_size = hidden_size
        self.page_capacity = page_capacity
        self.max_pages = max_pages
        self.hot_pages = hot_pages
        self.top_k_pages = top_k_pages
        self.top_k_records = top_k_records
        self.max_hops = max_hops
        self.router = router
        self.key_dim = int(key_dim or (router.router_dim if router is not None else hidden_size))
        if not 4 <= coarse_index_bits <= 20:
            raise ValueError("coarse_index_bits must be between 4 and 20")
        self.coarse_index_bits = int(coarse_index_bits)
        if max_resident_pages < hot_pages:
            raise ValueError("max_resident_pages must be >= hot_pages")
        self.tier_store = tier_store
        self.max_resident_pages = int(max_resident_pages)
        if gpu_cache_records < 0 or gpu_cache_tokens < 0:
            raise ValueError("gpu cache limits must be non-negative")
        if gpu_cache_reserve_mb < 0:
            raise ValueError("gpu_cache_reserve_mb must be non-negative")
        self.gpu_cache_records = int(gpu_cache_records)
        self.gpu_cache_tokens = int(gpu_cache_tokens)
        self.gpu_cache_reserve_mb = int(gpu_cache_reserve_mb)
        self.gpu_cache_reserve_bytes = self.gpu_cache_reserve_mb * 1024 * 1024
        self.gpu_cache_adaptive = bool(gpu_cache_adaptive)
        # Optional learned exact reranker.  Page routing remains bounded and
        # is still performed by MemoryRouterV2; this callback only scores the
        # small candidate set inside selected pages.
        self.record_scorer = record_scorer
        # Additive prior weights used by ``_record_scores``.  They are exposed because they
        # dominate the learned term -- measured, the neural score is 12-17% of the summed
        # total while these priors are 83-88% -- so they, not the reranker, decide which
        # records get injected.  Overridable per bank from the model config.
        self.prior_weights = {
            "lexical": 0.25,
            "token_overlap": 0.45,
            "rare_lexical": 1.25,
            "shape_bonus": 0.35,
            "structured_bonus": 0.15,
        }
        self._gpu_device: Optional[torch.device] = None
        self._gpu_record_cache: dict[str, dict[str, Tensor]] = {}
        self._gpu_cache_order: list[str] = []
        self._gpu_cache_tokens_used = 0
        self._gpu_cache_bytes_used = 0
        self._gpu_cache_hits = 0
        self._gpu_cache_misses = 0
        self._gpu_cache_fallbacks = 0
        self._gpu_cache_alloc_failures = 0
        self._gpu_cache_last_free_bytes: Optional[int] = None
        generator = torch.Generator(device="cpu").manual_seed(1729 + self.key_dim)
        self.coarse_planes = F.normalize(
            torch.randn(self.coarse_index_bits, self.key_dim, generator=generator), dim=-1
        )
        self.records: dict[str, MemoryRecordV2] = {}
        self.pages: dict[str, MemoryPageV2] = {}
        # Only pages with free capacity participate in write placement.  The
        # bounded order avoids scanning every page at million-record scale.
        self._open_page_ids: set[str] = set()
        self._open_page_order: list[str] = []
        self._coarse_buckets: dict[int, set[str]] = {}
        self._page_signatures: dict[str, set[int]] = {}
        # Sparse lexical page addresses recover exact pages for rare query
        # terms without turning the reader into a full-page scan.  The neural
        # router still scores the bounded candidate set afterwards.
        self._lexical_buckets: dict[str, set[str]] = {}
        self._page_terms: dict[str, set[str]] = {}
        # Evidence-bearing records get a second sparse inverted index.  This
        # prevents a page centroid from hiding a useful record among raw
        # background chunks without making the reader scan the whole bank.
        self._lexical_record_buckets: dict[str, set[str]] = {}
        self._record_lexical_terms: dict[str, set[str]] = {}
        self._record_page_ids: dict[str, str] = {}
        # Structured address indexes are separate from the semantic router.
        # They let distinctive identifiers (user ids, filenames, symbols) reach
        # the exact record even when a page centroid or LSH bucket is noisy.
        self._address_page_buckets: dict[str, set[str]] = {}
        self._address_record_buckets: dict[str, set[str]] = {}
        self._value_record_buckets: dict[str, set[str]] = {}
        self._entity_record_buckets: dict[str, set[str]] = {}
        self._page_address_terms: dict[str, set[str]] = {}
        self._page_entities: dict[str, set[str]] = {}
        self._record_address_terms: dict[str, set[str]] = {}
        self._record_entities: dict[str, str] = {}
        self._text_index: dict[str, str] = {}
        self.active_by_conflict: dict[str, str] = {}
        self.quarantine: dict[str, MemoryRecordV2] = {}
        self._counter = 0
        self._last_coarse_candidates = 0
        self._suspend_refresh = 0
        if self.tier_store is not None:
            self._restore_page_headers_from_store()
            for item in self.tier_store.load_quarantine():
                record = MemoryRecordV2(**item)
                self.quarantine[record.record_id] = record
            self.active_by_conflict.update(
                self.tier_store.active_conflicts(active_status=STATUS_ACTIVE)
            )
        if runtime_device is not None:
            self.configure_gpu_cache(
                runtime_device,
                max_records=self.gpu_cache_records,
                max_tokens=self.gpu_cache_tokens,
                reserve_mb=self.gpu_cache_reserve_mb,
                adaptive=self.gpu_cache_adaptive,
            )

    def configure_gpu_cache(
        self,
        device: torch.device | str,
        *,
        max_records: Optional[int] = None,
        max_tokens: Optional[int] = None,
        reserve_mb: Optional[int] = None,
        adaptive: Optional[bool] = None,
    ) -> None:
        """Configure a bounded VRAM cache for recently used memory records.

        The canonical copy remains in process RAM and is what gets exported
        into the embedded third safetensors shard.  This cache contains only
        routing keys and the small token payloads needed by recent reads; it
        is deliberately not a second full memory store.
        """

        target = torch.device(device)
        if max_records is not None:
            if int(max_records) < 0:
                raise ValueError("max_records must be non-negative")
            self.gpu_cache_records = int(max_records)
        if max_tokens is not None:
            if int(max_tokens) < 0:
                raise ValueError("max_tokens must be non-negative")
            self.gpu_cache_tokens = int(max_tokens)
        if reserve_mb is not None:
            if int(reserve_mb) < 0:
                raise ValueError("reserve_mb must be non-negative")
            self.gpu_cache_reserve_mb = int(reserve_mb)
            self.gpu_cache_reserve_bytes = self.gpu_cache_reserve_mb * 1024 * 1024
        if adaptive is not None:
            self.gpu_cache_adaptive = bool(adaptive)
        self._gpu_device = target if target.type == "cuda" else None
        self._clear_gpu_cache()

    def _clear_gpu_cache(self) -> None:
        self._gpu_record_cache.clear()
        self._gpu_cache_order.clear()
        self._gpu_cache_tokens_used = 0
        self._gpu_cache_bytes_used = 0

    @staticmethod
    def _tensor_bytes(value: Optional[Tensor]) -> int:
        if not isinstance(value, Tensor):
            return 0
        return int(value.numel()) * int(value.element_size())

    def _record_cache_bytes(self, record: MemoryRecordV2, *, include_tokens: bool) -> int:
        total = self._tensor_bytes(record.key) + self._tensor_bytes(record.summary)
        if include_tokens:
            total += self._tensor_bytes(record.token_ids)
            total += self._tensor_bytes(record.token_mask)
        return total

    def _has_gpu_headroom(self, requested_bytes: int) -> bool:
        """Keep the cache below a VRAM safety line while the model is running."""

        if self._gpu_device is None or not self.gpu_cache_adaptive:
            return True
        try:
            free_bytes, _ = torch.cuda.mem_get_info(self._gpu_device)
        except (RuntimeError, AssertionError, TypeError):
            return True
        self._gpu_cache_last_free_bytes = int(free_bytes)
        return int(free_bytes) - int(requested_bytes) >= self.gpu_cache_reserve_bytes

    def _evict_oldest_gpu_record(self) -> None:
        if not self._gpu_cache_order:
            return
        evicted_id = self._gpu_cache_order.pop(0)
        evicted = self._gpu_record_cache.pop(evicted_id, {})
        evicted_tokens = evicted.get("token_ids")
        if isinstance(evicted_tokens, Tensor):
            self._gpu_cache_tokens_used -= int(evicted_tokens.numel())
        self._gpu_cache_bytes_used -= sum(self._tensor_bytes(value) for value in evicted.values())
        self._gpu_cache_bytes_used = max(0, self._gpu_cache_bytes_used)

    def _promote_record(self, record: MemoryRecordV2) -> None:
        """Promote one hot record to VRAM without changing its RAM copy."""

        if (
            self._gpu_device is None
            or self.gpu_cache_records <= 0
            or record.record_id in self._gpu_record_cache
        ):
            if record.record_id in self._gpu_record_cache:
                self._gpu_cache_hits += 1
                self._gpu_cache_order.remove(record.record_id)
                self._gpu_cache_order.append(record.record_id)
            return
        self._gpu_cache_misses += 1
        source_token_count = int(record.token_ids.numel()) if isinstance(record.token_ids, Tensor) else 0
        token_count = source_token_count
        if self.gpu_cache_tokens == 0 or token_count > self.gpu_cache_tokens:
            token_count = 0
        include_tokens = token_count > 0 and isinstance(record.token_ids, Tensor)
        requested_bytes = self._record_cache_bytes(record, include_tokens=include_tokens)
        while self._gpu_cache_order and (
            len(self._gpu_cache_order) >= self.gpu_cache_records
            or self._gpu_cache_tokens_used + token_count > self.gpu_cache_tokens
            or not self._has_gpu_headroom(requested_bytes)
        ):
            self._evict_oldest_gpu_record()
        if not self._has_gpu_headroom(requested_bytes):
            # The authoritative RAM copy remains available to the router.
            # A hot-cache miss must never make inference fail just because the
            # GPU is busy with model weights or KV activations.
            self._gpu_cache_fallbacks += 1
            return
        try:
            cached: dict[str, Tensor] = {
                "key": record.key.detach().to(self._gpu_device, non_blocking=True),
                "summary": record.summary.detach().to(self._gpu_device, non_blocking=True),
            }
            if include_tokens:
                cached["token_ids"] = record.token_ids.detach().to(self._gpu_device, non_blocking=True)
                if isinstance(record.token_mask, Tensor):
                    cached["token_mask"] = record.token_mask.detach().to(self._gpu_device, non_blocking=True)
        except RuntimeError as exc:
            if "out of memory" not in str(exc).lower():
                raise
            self._gpu_cache_alloc_failures += 1
            self._gpu_cache_fallbacks += 1
            return
        if include_tokens:
            self._gpu_cache_tokens_used += token_count
        cache_bytes = sum(self._tensor_bytes(value) for value in cached.values())
        self._gpu_record_cache[record.record_id] = cached
        self._gpu_cache_order.append(record.record_id)
        self._gpu_cache_bytes_used += cache_bytes

    def _record_key(self, record: MemoryRecordV2) -> Tensor:
        self._promote_record(record)
        cached = self._gpu_record_cache.get(record.record_id)
        if cached is not None:
            return cached["key"]
        return record.key

    def gpu_record_payload(self, record: MemoryRecordV2) -> tuple[Tensor, Optional[Tensor]]:
        """Return a hot record's token payload, preferring its VRAM copy."""

        self._promote_record(record)
        cached = self._gpu_record_cache.get(record.record_id)
        if cached is not None and "token_ids" in cached:
            return cached["token_ids"], cached.get("token_mask")
        return record.token_ids, record.token_mask

    def promote_records(self, records: Iterable[MemoryRecordV2]) -> None:
        for record in records:
            self._promote_record(record)

    def has_token_evidence(
        self,
        query_key: Tensor,
        query_token_ids: Optional[Tensor],
        query_text: str = "",
        *,
        min_shared_tokens: int = 3,
    ) -> bool:
        """Check bounded exact evidence before the learned abstention gate."""

        if query_token_ids is None:
            return False
        # Attribute aliases and sparse operational terms are intentionally
        # indexed outside the serialized token payload.  They are still
        # trustworthy routing evidence: otherwise the learned need-memory
        # gate can abstain before the exact lexical address gets a chance to
        # reach the record reranker.
        lexical_pages, lexical_records = self._lexical_evidence_hits(query_text)
        address_pages, address_records = self._address_hits(query_text)
        if lexical_pages or lexical_records or address_pages or address_records:
            return True
        query_unique = torch.unique(query_token_ids.detach().reshape(-1).cpu())
        if query_unique.numel() < min_shared_tokens:
            return False
        for page_id in self._candidate_page_ids(query_key, query_text):
            page = self._hydrate_page(page_id)
            if page is None:
                continue
            for record_id in page.record_ids:
                record = self.records.get(record_id)
                if record is None or record.status != STATUS_ACTIVE or record.token_ids is None:
                    continue
                own = record.token_ids.reshape(-1)
                if isinstance(record.token_mask, Tensor) and record.token_mask.numel() == own.numel():
                    own = own[record.token_mask.reshape(-1).to(dtype=torch.bool)]
                shared = torch.isin(query_unique, torch.unique(own)).sum().item()
                if int(shared) >= int(min_shared_tokens):
                    return True
        return False

    def _restore_page_headers_from_store(self) -> None:
        """Restore page addresses without materializing every cold record."""

        if self.tier_store is None:
            return
        headers = self.tier_store.page_headers()
        for item in headers:
            page = MemoryPageV2(**item)
            self.pages[page.page_id] = page
            self._counter = max(self._counter, int(page.page_id.rsplit("_", 1)[-1]))
            if not page.full:
                self._open_page_ids.add(page.page_id)
                self._open_page_order.append(page.page_id)
            self._index_page(page, persist=False)

    def _rebuild_coarse_index(self, coarse_index_bits: Optional[int] = None) -> None:
        """Recreate LSH planes and page buckets after a config upgrade."""

        if coarse_index_bits is not None:
            if not 4 <= int(coarse_index_bits) <= 20:
                raise ValueError("coarse_index_bits must be between 4 and 20")
            self.coarse_index_bits = int(coarse_index_bits)
        generator = torch.Generator(device="cpu").manual_seed(1729 + self.key_dim)
        self.coarse_planes = F.normalize(
            torch.randn(self.coarse_index_bits, self.key_dim, generator=generator), dim=-1
        )
        self._coarse_buckets.clear()
        self._page_signatures.clear()
        self._lexical_buckets.clear()
        self._page_terms.clear()
        self._lexical_record_buckets.clear()
        self._record_lexical_terms.clear()
        self._record_page_ids.clear()
        self._value_record_buckets.clear()
        for page in self.pages.values():
            self._index_page(page, persist=False)
        if self.tier_store is not None:
            self.tier_store.clear_coarse_buckets()
            page_by_record = {
                record_id: page.page_id
                for page in self.pages.values()
                for record_id in page.record_ids
            }
            for record_id, key in self.tier_store.record_keys():
                page_id = page_by_record.get(record_id)
                if page_id is None:
                    continue
                signature = self._coarse_signature(key)
                self._page_signatures.setdefault(page_id, set()).add(signature)
            for page_id, signatures in self._page_signatures.items():
                self.tier_store.replace_page_buckets(page_id, signatures)

    def _encode_key(self, key: Tensor) -> Tensor:
        key = key.detach().float().reshape(-1)
        if key.numel() == self.key_dim:
            return F.normalize(key, dim=0)
        if self.router is not None and key.numel() == self.hidden_size:
            router_device = next(self.router.parameters()).device
            return self.router.encode_key(key.to(router_device).unsqueeze(0))[0].detach().float().cpu()
        raise ValueError(
            f"memory key must contain {self.hidden_size} model values or {self.key_dim} address values"
        )

    def _score_candidates(self, query_key: Tensor, candidate_keys: Tensor) -> Tensor:
        if self.router is not None and query_key.numel() == self.hidden_size:
            router_device = next(self.router.parameters()).device
            scores, _ = self.router.projected_scores(
                query_key.to(router_device).reshape(1, -1),
                candidate_keys.to(router_device).reshape(1, -1, self.key_dim),
            )
            # Training uses logits for cross-entropy; storage uses a bounded
            # relevance value so page thresholds remain stable across router
            # checkpoints and do not discard valid negative logits.
            return torch.sigmoid(scores[0])
        query = self._encode_key(query_key)
        candidates = F.normalize(candidate_keys.float().to(query.device), dim=-1)
        return torch.matmul(candidates, query)

    def _coarse_signature(self, key: Tensor) -> int:
        address = self._encode_key(key).cpu()
        bits = (torch.mv(self.coarse_planes, address) >= 0).to(torch.int64)
        signature = 0
        for index, bit in enumerate(bits.tolist()):
            signature |= int(bit) << index
        return signature

    def _remove_page_from_index(self, page: MemoryPageV2) -> None:
        signatures = self._page_signatures.pop(page.page_id, set())
        for signature in signatures:
            bucket = self._coarse_buckets.get(signature)
            if bucket is not None:
                bucket.discard(page.page_id)
                if not bucket:
                    self._coarse_buckets.pop(signature, None)
        terms = self._page_terms.pop(page.page_id, set())
        for term in terms:
            bucket = self._lexical_buckets.get(term)
            if bucket is not None:
                bucket.discard(page.page_id)
                if not bucket:
                    self._lexical_buckets.pop(term, None)
        for term in self._page_address_terms.pop(page.page_id, set()):
            bucket = self._address_page_buckets.get(term)
            if bucket is not None:
                bucket.discard(page.page_id)
                if not bucket:
                    self._address_page_buckets.pop(term, None)
        for entity in self._page_entities.pop(page.page_id, set()):
            bucket = self._entity_record_buckets.get(entity)
            if bucket is not None:
                bucket.difference_update(page.record_ids)
                if not bucket:
                    self._entity_record_buckets.pop(entity, None)
        for record_id in page.record_ids:
            for term in self._record_lexical_terms.pop(record_id, set()):
                bucket = self._lexical_record_buckets.get(term)
                if bucket is not None:
                    bucket.discard(record_id)
                    if not bucket:
                        self._lexical_record_buckets.pop(term, None)
            self._record_page_ids.pop(record_id, None)
            for term in self._record_address_terms.pop(record_id, set()):
                bucket = self._address_record_buckets.get(term)
                if bucket is not None:
                    bucket.discard(record_id)
                    if not bucket:
                        self._address_record_buckets.pop(term, None)
            record = self.records.get(record_id)
            if record is not None:
                for term in _tokens(record.value):
                    bucket = self._value_record_buckets.get(term)
                    if bucket is not None:
                        bucket.discard(record_id)
                        if not bucket:
                            self._value_record_buckets.pop(term, None)
            self._record_entities.pop(record_id, None)

    def _index_page(self, page: MemoryPageV2, *, persist: bool = True) -> None:
        # Make re-indexing idempotent when a cold page is hydrated after
        # restart and its record text becomes available.
        self._remove_page_from_index(page)
        signatures: set[int] = set()
        if page.key is not None:
            signatures.add(self._coarse_signature(page.key))
        # Index record addresses as well as the page centroid.  A centroid
        # can blur multiple topics, while a record address is exact and still
        # costs only one small LSH bucket insertion per record.
        for record_id in page.record_ids:
            record = self.records.get(record_id)
            if record is not None:
                signatures.add(self._coarse_signature(record.key))
        self._page_signatures[page.page_id] = signatures
        for signature in signatures:
            self._coarse_buckets.setdefault(signature, set()).add(page.page_id)
        terms: set[str] = set()
        for record_id in page.record_ids:
            record = self.records.get(record_id)
            if record is not None and record.status == STATUS_ACTIVE:
                record_terms = _tokens(record.routing_text())
                terms.update(record_terms)
                self._record_page_ids[record_id] = page.page_id
                # Random/background chunks may be present only to model
                # corpus noise.  Do not let them occupy the direct lexical
                # route when they have no semantic key or token payload.
                if record.semantic_key is not None or record.token_ids is not None:
                    self._record_lexical_terms[record_id] = record_terms
                    for term in record_terms:
                        self._lexical_record_buckets.setdefault(term, set()).add(record_id)
        self._page_terms[page.page_id] = terms
        for term in terms:
            self._lexical_buckets.setdefault(term, set()).add(page.page_id)
        address_terms: set[str] = set()
        entities: set[str] = set()
        for record_id in page.record_ids:
            record = self.records.get(record_id)
            if record is None or record.status != STATUS_ACTIVE:
                continue
            entity = record.entity.strip().lower()
            if entity:
                entities.add(entity)
                self._record_entities[record_id] = entity
                self._entity_record_buckets.setdefault(entity, set()).add(record_id)
            record_address_terms = {
                term
                for field in (record.entity, record.attribute)
                for term in _tokens(str(field))
                if self._is_distinctive_address_term(term)
                or (
                    record.attribute.strip().lower().startswith(("symbol:", "class:", "def:"))
                    and self._is_distinctive_symbol_term(term)
                )
            }
            self._record_address_terms[record_id] = record_address_terms
            address_terms.update(record_address_terms)
            for term in record_address_terms:
                self._address_record_buckets.setdefault(term, set()).add(record_id)
            for term in _tokens(record.value):
                self._value_record_buckets.setdefault(term, set()).add(record_id)
        self._page_entities[page.page_id] = entities
        self._page_address_terms[page.page_id] = address_terms
        for term in address_terms:
            self._address_page_buckets.setdefault(term, set()).add(page.page_id)
        if self.tier_store is not None and persist:
            self.tier_store.upsert_page(page)
            self.tier_store.replace_page_buckets(page.page_id, signatures)

    def _hydrate_records(self, record_ids: Iterable[str]) -> None:
        """Load only the records needed by the current candidate pages."""

        if self.tier_store is None:
            return
        missing = [record_id for record_id in record_ids if record_id not in self.records]
        if not missing:
            return
        touched_pages: set[str] = set()
        for item in self.tier_store.load_records(missing):
            self.records[item["record_id"]] = MemoryRecordV2(**item)
            text = item.get("text", "").strip().lower()
            if text and item.get("status") == STATUS_ACTIVE:
                self._text_index[text] = item["record_id"]
            page_id = str(item.get("page_id", ""))
            if page_id:
                touched_pages.add(page_id)
        # Rebuild each touched page once, after all of its requested records
        # are resident.  Re-indexing inside the load loop could temporarily
        # erase address terms belonging to still-hydrating fragments.
        for page_id in touched_pages:
            page = self.pages.get(page_id)
            if page is not None:
                self._index_page(page, persist=False)

    def _hydrate_page(self, page_id: str) -> Optional[MemoryPageV2]:
        page = self.pages.get(page_id)
        if page is None or self.tier_store is None:
            return page
        self._hydrate_records(page.record_ids)
        return page

    def _evict_cold_records(self) -> None:
        """Keep page headers resident and unload non-hot record payloads."""

        if self.tier_store is None or self.max_resident_pages <= 0:
            return
        ranked = sorted(
            self.pages.values(),
            key=lambda page: (page.tier == "hot", page.last_access, page.importance),
            reverse=True,
        )
        keep = {page.page_id for page in ranked[: self.max_resident_pages]}
        for page in self.pages.values():
            if page.page_id in keep:
                continue
            for record_id in page.record_ids:
                self.records.pop(record_id, None)

    def _store_record(self, record: MemoryRecordV2) -> None:
        if self.tier_store is not None:
            self.tier_store.upsert_record(record)

    def _candidate_page_ids(self, query_key: Tensor, query_text: str = "") -> list[str]:
        """Use locality-sensitive coarse buckets before exact page scoring."""

        page_ids = list(self.pages)
        if len(page_ids) <= 128:
            self._last_coarse_candidates = len(page_ids)
            return page_ids
        signature = self._coarse_signature(query_key)
        probes = [signature]
        # Probe Hamming distance 1 and 2.  This is bounded (<= 79 buckets for
        # 12 bits) and avoids a query-time scan over every page.
        for bit in range(self.coarse_index_bits):
            probes.append(signature ^ (1 << bit))
        for left in range(self.coarse_index_bits):
            for right in range(left + 1, self.coarse_index_bits):
                probes.append(signature ^ (1 << left) ^ (1 << right))
        query_terms = _tokens(query_text)
        lexical_pages: set[str] = set()
        for term in query_terms:
            bucket = self._lexical_buckets.get(term)
            # Ignore common terms so this remains a sparse address lookup.
            if bucket is not None and len(bucket) <= 128:
                lexical_pages.update(bucket)
            record_bucket = self._lexical_record_buckets.get(term)
            if record_bucket is not None and len(record_bucket) <= 128:
                lexical_pages.update(
                    self._record_page_ids[record_id]
                    for record_id in record_bucket
                    if record_id in self._record_page_ids
                )
        # Structured address terms are allowed to bypass the ordinary lexical
        # bucket-size cutoff.  This is still sparse because the address index
        # only stores distinctive ids, filenames and symbols.
        address_pages, _ = self._address_hits(query_text)
        lexical_pages.update(address_pages)
        if self.tier_store is not None:
            hot_ids = [
                page.page_id for page in self.pages.values() if page.tier == "hot"
            ]
            selected = self.tier_store.candidate_page_ids(
                probes,
                hot_page_ids=hot_ids,
                limit=max(128, self.top_k_pages * 64),
            )
            selected_set = set(selected)
            selected_set.update(lexical_pages)
            limit = max(128, self.top_k_pages * 64)
            if len(selected_set) > limit:
                selected_set = set(selected[:limit]) | lexical_pages
            self._last_coarse_candidates = len(selected_set)
            return list(selected_set)
        selected: set[str] = set()
        for bucket in probes:
            selected.update(self._coarse_buckets.get(bucket, ()))
        selected.update(lexical_pages)
        # Hot pages are always eligible.  They are few by construction and
        # protect high-value memories when a page centroid is still moving.
        selected.update(
            page.page_id for page in self.pages.values() if page.tier == "hot"
        )
        if not selected:
            # Cold-start safety: a bounded sample is preferable to silently
            # claiming that no memory exists.
            selected.update(page_ids[: min(128, len(page_ids))])
        self._last_coarse_candidates = len(selected)
        return list(selected)

    def _rare_lexical_address_score(self, query_text: str, terms: set[str]) -> float:
        """Return a bounded address prior for rare query terms.

        The learned router is intentionally not the only source of truth at
        million-token scale.  A page containing an exact, rare identifier
        such as a project code is an address hit, even when the page centroid
        is diluted by hundreds of unrelated chunks.  This prior is sparse and
        bounded: common language terms are ignored, and it never scans pages.
        """

        if not query_text or not terms:
            return 0.0
        hits = 0
        for term in _tokens(query_text):
            if term not in terms:
                continue
            # Evidence-bearing records have their own sparse lexical index.
            # Use it before the page-level bucket: a large page can contain a
            # rare, exact target even when the page itself is lexically broad.
            record_bucket = self._lexical_record_buckets.get(term)
            page_bucket = self._lexical_buckets.get(term)
            if record_bucket is not None and len(record_bucket) <= 32:
                hits += 1
            elif page_bucket is not None and len(page_bucket) <= 8:
                hits += 1
        return min(2.0, float(hits))

    @staticmethod
    def _is_distinctive_address_term(term: str) -> bool:
        """Keep structured address lookup sparse at very large scale."""

        normalized = str(term).strip().lower()
        return len(normalized) >= 8 or any(
            char.isdigit() or char in "_-./\\" for char in normalized
        )

    @staticmethod
    def _is_distinctive_symbol_term(term: str) -> bool:
        """Keep ordinary Python symbol names addressable without filenames."""

        normalized = str(term).strip().lower()
        return len(normalized) >= 3 and all("a" <= char <= "z" or char == "_" for char in normalized)

    @staticmethod
    def _is_distinctive_value_term(term: str) -> bool:
        """Recognize compact values such as ``DRAM`` without indexing prose."""

        normalized = str(term).strip().lower()
        if len(normalized) < 4:
            return False
        if any(char.isdigit() or char in "_-./\\" for char in normalized):
            return True
        return all("a" <= char <= "z" for char in normalized)

    def _address_hits(self, query_text: str) -> tuple[set[str], set[str]]:
        """Return bounded page/record hits for distinctive address terms."""

        page_ids: set[str] = set()
        record_ids: set[str] = set()
        for term in _tokens(query_text):
            is_address_term = self._is_distinctive_address_term(term)
            is_value_term = self._is_distinctive_value_term(term)
            is_symbol_term = self._is_distinctive_symbol_term(term)
            if not is_address_term and not is_value_term and not is_symbol_term:
                continue
            if is_address_term or is_symbol_term:
                pages = self._address_page_buckets.get(term)
                records = self._address_record_buckets.get(term)
                # A common source word should not fan out into a full scan.
                # Codes, filenames and symbols normally have small buckets.
                if pages is not None and len(pages) <= 256:
                    page_ids.update(pages)
                if records is not None and len(records) <= 256:
                    record_ids.update(records)
            if is_value_term:
                value_records = self._value_record_buckets.get(term)
                if value_records is not None and len(value_records) <= 128:
                    record_ids.update(value_records)
                    page_ids.update(
                        self._record_page_ids[record_id]
                        for record_id in value_records
                        if record_id in self._record_page_ids
                    )
        return page_ids, record_ids

    def _lexical_evidence_hits(self, query_text: str) -> tuple[set[str], set[str]]:
        """Find sparse exact evidence hits for operational/fragment queries.

        A filename or symbol can use the structured address index, but a
        question such as ``what is the reset token`` has no filename to
        address.  The evidence index handles that case without scanning the
        bank: it unions only small inverted-list buckets and requires either
        two shared terms or one distinctive ASCII/identifier term.
        """

        matched_terms: dict[str, set[str]] = {}
        for term in _tokens(query_text):
            bucket = self._lexical_record_buckets.get(term)
            if bucket is None or len(bucket) > 64:
                continue
            for record_id in bucket:
                matched_terms.setdefault(record_id, set()).add(term)
        if not matched_terms:
            return set(), set()
        record_ids: set[str] = set()
        for record_id, terms in matched_terms.items():
            ascii_terms = [
                term
                for term in terms
                if all("a" <= char <= "z" or char == "_" for char in term)
            ]
            distinctive = any(
                self._is_distinctive_address_term(term)
                or self._is_distinctive_symbol_term(term)
                or self._is_distinctive_value_term(term)
                for term in ascii_terms
            )
            # ``候选`` is two single Chinese characters and is not a sparse
            # address.  Pure CJK evidence therefore needs a longer overlap;
            # one or more ASCII identifier terms may trigger immediately.
            if distinctive or (ascii_terms and len(terms) >= 2) or len(terms) >= 3:
                record_ids.add(record_id)
        page_ids = {
            self._record_page_ids[record_id]
            for record_id in record_ids
            if record_id in self._record_page_ids
        }
        return page_ids, record_ids

    def has_explicit_address(self, query_text: str) -> bool:
        """Return whether a query contains an indexed sparse address.

        The model-side reader uses this as a safe fast-path signal. It does
        not infer relevance from arbitrary lexical overlap; only the existing
        bounded address index can activate it.
        """

        page_ids, record_ids = self._address_hits(query_text)
        return bool(page_ids or record_ids)

    def _record_matches_explicit_address(
        self,
        record: MemoryRecordV2,
        query_text: str,
    ) -> bool:
        """Match entity plus an attribute alias such as ``symbol:parse_args``."""

        query_lower = str(query_text).strip().lower()
        query_terms = _tokens(query_lower)
        entity = record.entity.strip().lower()
        if not entity:
            return False
        entity_match = entity in query_lower
        # Numeric identifiers are not globally unique: ``训练用户00006`` and
        # ``评估用户00006`` must not collide merely because both contain
        # ``00006``.  Token fallback is reserved for non-natural-language
        # addresses such as filenames and symbols.
        if not entity_match and not any("\u4e00" <= char <= "\u9fff" for char in entity):
            entity_match = any(
                term in query_terms
                for term in _tokens(entity)
                if self._is_distinctive_address_term(term)
            )
        attribute = record.attribute.strip().lower()
        symbol_attribute = attribute.startswith(("symbol:", "class:", "def:"))
        attribute_match = bool(attribute and attribute in query_lower) or any(
            term in query_terms
            for term in _tokens(attribute)
            if self._is_distinctive_address_term(term)
            or (symbol_attribute and self._is_distinctive_symbol_term(term))
        )
        value_match = any(
            term in query_terms
            for term in _tokens(record.value)
            if self._is_distinctive_value_term(term)
        )
        # A concrete source file is a stronger address than a shared symbol
        # such as ``main`` or ``run``.  Do not admit same-named definitions
        # from unrelated files into the evidence prefix.
        explicit_file = bool(
            re.search(r"(?:[a-z0-9_.\\/-]+)\.(?:py|md|json|toml|yaml|yml)", query_lower)
        )
        if explicit_file and symbol_attribute and not entity_match:
            return False
        # Code navigation questions commonly provide only a symbol such as
        # ``main`` or ``_path`` and omit the filename.  The symbol itself is
        # still a valid sparse address; keep the bounded candidate set so the
        # evidence layer can expose ambiguity instead of dropping all hits.
        if not entity_match and not attribute_match and not value_match:
            return False
        # Attribute-only addressing is safe for code symbols because
        # ``symbol:<name>`` is itself the requested identifier.  Natural
        # attributes such as ``档案代号`` are not unique across entities, so
        # they must still require the entity or a distinctive value match.
        if not entity_match and (value_match or (attribute_match and symbol_attribute)):
            return True
        if not attribute:
            return entity_match or value_match
        return entity_match and (attribute_match or value_match or attribute in query_lower)

    def _new_page(self, *, tier: str = "warm") -> MemoryPageV2:
        if len(self.pages) >= self.max_pages:
            self.consolidate(max_pages=max(1, self.max_pages - 1))
        if len(self.pages) >= self.max_pages:
            raise RuntimeError(
                "memory page capacity exhausted; consolidate or increase max_pages"
            )
        self._counter += 1
        page = MemoryPageV2(
            page_id=f"page_{self._counter:08d}",
            tier=tier,
            capacity=self.page_capacity,
        )
        self.pages[page.page_id] = page
        self._open_page_ids.add(page.page_id)
        self._open_page_order.append(page.page_id)
        return page

    def _target_page(self, key: Tensor, *, importance: float) -> MemoryPageV2:
        open_ids = [
            page_id
            for page_id in self._open_page_order
            if page_id in self._open_page_ids
            and page_id in self.pages
            and not self.pages[page_id].full
        ]
        # If many pages remain partially filled, compare only a bounded recent
        # window plus hot pages. Write placement must stay sublinear in the
        # number of pages.
        if len(open_ids) > 128:
            recent = open_ids[-128:]
            hot = [
                page.page_id
                for page in self.pages.values()
                if page.tier == "hot" and not page.full
            ]
            open_ids = list(dict.fromkeys(recent + hot))
        candidates = [self.pages[page_id] for page_id in open_ids]
        if not candidates:
            return self._new_page(tier="hot" if importance >= 0.8 else "warm")
        key = self._encode_key(key)
        best_page: Optional[MemoryPageV2] = None
        best_score = -float("inf")
        for page in candidates:
            if page.key is None:
                score = -0.1 * len(page.record_ids)
            else:
                score = float(torch.dot(key.to(page.key.device), F.normalize(page.key.float(), dim=0)))
            score += 0.02 * page.importance
            if score > best_score:
                best_score = score
                best_page = page
        assert best_page is not None
        return best_page

    def _update_page(self, page: MemoryPageV2, record: MemoryRecordV2) -> None:
        self._remove_page_from_index(page)
        page.record_ids.append(record.record_id)
        record.page_id = page.page_id
        weight = 1.0 / max(1, len(page.record_ids))
        if page.key is None:
            page.key = record.key.detach().float().clone()
            page.summary = record.summary.detach().float().clone()
        else:
            page.key = (1.0 - weight) * page.key + weight * record.key.detach().float()
            if page.summary is None:
                page.summary = record.summary.detach().float().clone()
            else:
                page.summary = (1.0 - weight) * page.summary + weight * record.summary.detach().float()
        page.importance = max(page.importance, float(record.importance))
        page.touch()
        if page.full:
            self._open_page_ids.discard(page.page_id)
        self._index_page(page)

    def _find_duplicate(self, record: MemoryRecordV2) -> Optional[MemoryRecordV2]:
        text_key = record.text.strip().lower()
        if text_key:
            record_id = self._text_index.get(text_key)
            duplicate = self.records.get(record_id) if record_id else None
            if duplicate is None and self.tier_store is not None:
                item = self.tier_store.find_by_text(record.text, active_status=STATUS_ACTIVE)
                if item is not None:
                    self._hydrate_records([item["record_id"]])
                    duplicate = self.records.get(item["record_id"])
            if duplicate is not None and duplicate.status == STATUS_ACTIVE:
                return duplicate
        conflict_key = record.conflict_key()
        if conflict_key:
            old_id = self.active_by_conflict.get(conflict_key)
            if old_id and old_id in self.records:
                old = self.records[old_id]
                if old.status == STATUS_ACTIVE and old.value.strip().lower() == record.value.strip().lower():
                    return old
            if self.tier_store is not None:
                item = self.tier_store.find_by_conflict(
                    record.entity,
                    record.attribute,
                    active_status=STATUS_ACTIVE,
                )
                if item is not None:
                    self._hydrate_records([item["record_id"]])
                    old = self.records.get(item["record_id"])
                    if old is not None and old.value.strip().lower() == record.value.strip().lower():
                        return old
        return None

    def write(
        self,
        *,
        text: str,
        key: Tensor,
        summary: Optional[Tensor] = None,
        semantic_key: Optional[Tensor] = None,
        memory_type: str = "fact",
        entity: str = "",
        attribute: str = "",
        value: str = "",
        importance: float = 0.5,
        confidence: float = 0.5,
        source: str = "user",
        evidence: Optional[Sequence[str]] = None,
        related_ids: Optional[Sequence[str]] = None,
        slot_index: int = -1,
        token_ids: Optional[Tensor] = None,
        token_mask: Optional[Tensor] = None,
        trusted: bool = True,
        force: bool = False,
        origin: str = "",
    ) -> tuple[MemoryRecordV2, str]:
        """Insert or version a record, returning ``(record, action)``."""

        importance = float(max(0.0, min(1.0, importance)))
        confidence = float(max(0.0, min(1.0, confidence)))
        if summary is None:
            summary = key
        record = MemoryRecordV2(
            record_id=_stable_id(f"{text}|{_now()}|{len(self.records)}", prefix="mem"),
            text=text,
            key=self._encode_key(key).cpu().clone(),
            summary=self._encode_key(summary).cpu().clone(),
            semantic_key=(
                semantic_key.detach().float().reshape(-1).cpu().clone()
                if semantic_key is not None
                else None
            ),
            memory_type=memory_type,
            entity=entity,
            attribute=attribute,
            value=value,
            importance=importance,
            confidence=confidence,
            source=source,
            evidence=list(evidence or [text]),
            related_ids=list(related_ids or []),
            slot_index=int(slot_index),
            token_ids=(token_ids.detach().cpu().long().reshape(-1).clone() if token_ids is not None else None),
            token_mask=(token_mask.detach().cpu().bool().reshape(-1).clone() if token_mask is not None else None),
            origin=str(origin or ""),
        )
        if not trusted and not force:
            record.status = STATUS_QUARANTINED
            self.quarantine[record.record_id] = record
            if self.tier_store is not None:
                self.tier_store.upsert_quarantine(record)
            return record, "quarantined"
        duplicate = self._find_duplicate(record)
        if duplicate is not None:
            duplicate.confidence = max(duplicate.confidence, record.confidence)
            duplicate.importance = max(duplicate.importance, record.importance)
            duplicate.evidence.extend(item for item in record.evidence if item not in duplicate.evidence)
            duplicate.touch()
            return duplicate, "duplicate"
        if record.slot_index >= 0:
            for old in self.records.values():
                if old.slot_index == record.slot_index and old.status == STATUS_ACTIVE:
                    old.status = STATUS_SUPERSEDED
                    record.version = max(record.version, old.version + 1)
                    record.supersedes = old.record_id
                    self._store_record(old)
        conflict_key = record.conflict_key()
        if conflict_key and conflict_key in self.active_by_conflict:
            old_id = self.active_by_conflict[conflict_key]
            old = self.records.get(old_id)
            if old is not None and old.status == STATUS_ACTIVE:
                old.status = STATUS_SUPERSEDED
                record.version = old.version + 1
                record.supersedes = old.record_id
                self._store_record(old)
        elif conflict_key and self.tier_store is not None:
            item = self.tier_store.find_by_conflict(
                record.entity,
                record.attribute,
                active_status=STATUS_ACTIVE,
            )
            old = None
            if item is not None:
                self._hydrate_records([item["record_id"]])
                old = self.records.get(item["record_id"])
            if old is not None:
                old.status = STATUS_SUPERSEDED
                record.version = old.version + 1
                record.supersedes = old.record_id
                self.active_by_conflict[conflict_key] = old.record_id
                self._store_record(old)
        self._absorb_unidentified_same_turn(record)
        page = self._target_page(record.key, importance=record.importance)
        self.records[record.record_id] = record
        if record.text.strip():
            self._text_index[record.text.strip().lower()] = record.record_id
        self._update_page(page, record)
        self._store_record(record)
        if conflict_key:
            self.active_by_conflict[conflict_key] = record.record_id
        if self._suspend_refresh == 0:
            self._refresh_tiers()
        return record, "updated" if record.supersedes else "inserted"

    def write_batch(self, records: Iterable[dict[str, Any]]) -> list[tuple[MemoryRecordV2, str]]:
        """Write many records with one durable transaction and one tier pass."""

        items = list(records)
        if not items:
            return []
        self._suspend_refresh += 1
        try:
            if self.tier_store is not None:
                with self.tier_store.transaction():
                    output = [self.write(**item) for item in items]
            else:
                output = [self.write(**item) for item in items]
        finally:
            self._suspend_refresh -= 1
        if self._suspend_refresh == 0:
            self._refresh_tiers()
        return output

    def _refresh_tiers(self) -> None:
        ranked = sorted(
            self.pages.values(),
            key=lambda page: (page.importance, page.last_access, len(page.record_ids)),
            reverse=True,
        )
        hot_ids = {page.page_id for page in ranked[: self.hot_pages]}
        resident_ids = {
            page.page_id for page in ranked[: self.max_resident_pages]
        }
        for page in self.pages.values():
            old_tier = page.tier
            if page.page_id in hot_ids:
                page.tier = "hot"
            elif self.tier_store is not None and page.page_id not in resident_ids:
                page.tier = "cold"
            else:
                page.tier = "warm"
            if self.tier_store is not None and old_tier != page.tier:
                self.tier_store.set_page_tier(page.page_id, page.tier)
        if self.tier_store is None and self._gpu_device is not None:
            hot_records: list[MemoryRecordV2] = []
            for page in sorted(
                (item for item in self.pages.values() if item.tier == "hot"),
                key=lambda item: (item.last_access, item.importance),
                reverse=True,
            ):
                hot_records.extend(
                    self.records[record_id]
                    for record_id in page.record_ids
                    if record_id in self.records
                    and self.records[record_id].status == STATUS_ACTIVE
                )
            self.promote_records(hot_records)
        self._evict_cold_records()

    def _page_scores(
        self,
        query_key: Tensor,
        query_text: str = "",
        query_token_ids: Optional[Tensor] = None,
    ) -> list[tuple[MemoryPageV2, float]]:
        scores = []
        candidate_page_ids = self._candidate_page_ids(query_key, query_text)
        for page_id in candidate_page_ids:
            page = self._hydrate_page(page_id)
            if page is None:
                continue
            if page.key is None:
                continue
            score = float(self._score_candidates(query_key, page.key.unsqueeze(0))[0].item())
            active_records = [
                self.records[record_id]
                for record_id in page.record_ids
                if record_id in self.records and self.records[record_id].status == STATUS_ACTIVE
            ]
            if active_records:
                record_keys = torch.stack(
                    [self._record_key(record).to(device=query_key.device) for record in active_records],
                    dim=0,
                )
                # Exact record-level reranking prevents a mixed-topic page
                # centroid from hiding a highly relevant slot.
                score = max(score, float(self._score_candidates(query_key, record_keys).max().item()))
                # Prefer pages that contain usable evidence over pages made
                # entirely of address/background vectors.  The raw records
                # remain in RAM; this only affects bounded admission order.
                if any(record.token_ids is not None for record in active_records):
                    score += 0.20
            score += 0.05 * page.importance
            age_seconds = max(0, _now() - page.last_access)
            score += 0.02 / (1.0 + age_seconds / 3600.0)
            if query_text:
                page_text = " ".join(
                    self.records[item].routing_text()
                    for item in page.record_ids
                    if item in self.records
                )
                score += 0.65 * max(
                    (self.records[item].lexical_score(query_text) for item in page.record_ids if item in self.records),
                    default=0.0,
                )
                # The record-level index is deliberately limited to usable
                # evidence records. This makes an exact operational phrase
                # beat a large README/background page without scanning all
                # records at query time.
                record_lexical = max(
                    (
                        self.records[item].lexical_score(query_text)
                        for item in page.record_ids
                        if item in self.records
                        and self.records[item].status == STATUS_ACTIVE
                        and (
                            self.records[item].semantic_key is not None
                            or self.records[item].token_ids is not None
                        )
                    ),
                    default=0.0,
                )
                score += 0.45 * record_lexical
                score += 1.25 * self._rare_lexical_address_score(
                    query_text,
                    self._page_terms.get(page.page_id, set()),
                )
                # A query that explicitly names a sufficiently distinctive
                # entity must not be routed by broad semantic similarity
                # alone.  This is an address-layer guard for near-duplicate
                # memories (for example, two users with the same attribute),
                # not an answer generator: it only promotes the page that
                # already contains the named entity and optional attribute.
                query_lower = query_text.strip().lower()
                exact_entity = False
                exact_attribute = False
                for record in active_records:
                    entity = record.entity.strip().lower()
                    if len(entity) < 4 or entity not in query_lower:
                        continue
                    exact_entity = True
                    attribute = record.attribute.strip().lower()
                    if attribute and attribute in query_lower:
                        exact_attribute = True
                        break
                if exact_entity:
                    score += 2.50
                if exact_attribute:
                    score += 0.75
            if query_token_ids is not None:
                score += 0.45 * max(
                    (
                        self.records[item].token_overlap_score(query_token_ids)
                        for item in page.record_ids
                        if item in self.records and self.records[item].status == STATUS_ACTIVE
                    ),
                    default=0.0,
                )
            scores.append((page, score))
        return sorted(scores, key=lambda item: item[1], reverse=True)

    def _record_scores(
        self,
        page: MemoryPageV2,
        query_key: Tensor,
        query_text: str,
        query_token_ids: Optional[Tensor] = None,
        *,
        allow_superseded: bool = False,
    ) -> list[tuple[MemoryRecordV2, float]]:
        output = []
        self._hydrate_page(page.page_id)
        live_records: list[MemoryRecordV2] = []
        for record_id in page.record_ids:
            record = self.records.get(record_id)
            if record is None:
                continue
            if record.status != STATUS_ACTIVE and not allow_superseded:
                continue
            live_records.append(record)
        if not live_records:
            return []
        candidate_keys = torch.stack(
            [self._record_key(record).to(device=query_key.device) for record in live_records],
            dim=0,
        )
        semantic_keys = [record.semantic_key for record in live_records]
        if self.record_scorer is not None and any(isinstance(value, Tensor) for value in semantic_keys):
            # Pages may contain learned target records beside raw background
            # chunks.  Score both groups without allowing a random filler key
            # to contaminate the learned reranker.
            routed_scores = self._score_candidates(query_key, candidate_keys).to(
                device=query_key.device, dtype=torch.float32
            )
            semantic_indices = [
                index for index, value in enumerate(semantic_keys) if isinstance(value, Tensor)
            ]
            semantic_candidates = torch.stack(
                [semantic_keys[index] for index in semantic_indices], dim=0
            ).to(device=query_key.device)
            semantic_scores = self.record_scorer(query_key, semantic_candidates).reshape(-1)
            for position, score in zip(semantic_indices, semantic_scores):
                routed_scores[position] = score.to(dtype=torch.float32)
        else:
            routed_scores = self._score_candidates(query_key, candidate_keys)
        for record, routed_score in zip(live_records, routed_scores):
            score = float(routed_score.item())
            # Prior weights are configurable because they dominate the learned term:
            # measured on this scorer, the neural score is only 12-17% of the summed
            # total while these priors are 83-88%, so changing the learned ranker alone
            # measurably cannot change which records are injected (verified end-to-end).
            score += self.prior_weights["lexical"] * record.lexical_score(query_text)
            score += self.prior_weights["token_overlap"] * record.token_overlap_score(query_token_ids)
            score += 0.10 * record.confidence + 0.08 * record.importance
            score += self.prior_weights["rare_lexical"] * self._rare_lexical_address_score(
                query_text,
                _tokens(record.routing_text()),
            )
            # An item that can actually be injected into the model is more
            # useful than a centroid-only background chunk.  This is a small
            # evidence-quality prior, not a replacement for relevance.
            if record.token_ids is not None:
                score += self.prior_weights["shape_bonus"]
            if record.entity or record.attribute or record.value:
                score += self.prior_weights["structured_bonus"]
            entity = record.entity.strip().lower()
            query_lower = query_text.strip().lower()
            if len(entity) >= 4 and entity in query_lower:
                # Keep the same identity prior at record level so the exact
                # entity remains first after page admission and reranking.
                score += 2.50
                attribute = record.attribute.strip().lower()
                if attribute and attribute in query_lower:
                    score += 0.75
            output.append((record, score))
        return sorted(output, key=lambda item: item[1], reverse=True)

    def query(
        self,
        *,
        query_key: Tensor,
        query_text: str = "",
        query_token_ids: Optional[Tensor] = None,
        top_k_pages: Optional[int] = None,
        top_k_records: Optional[int] = None,
        max_hops: Optional[int] = None,
        min_score: float = -1.0,
    ) -> tuple[list[MemoryRecordV2], RouterDecisionV2]:
        """Sparse multi-hop query over pages and active records."""

        page_limit = max(1, top_k_pages or self.top_k_pages)
        record_limit = top_k_records or self.top_k_records
        hop_limit = min(max_hops or self.max_hops, self.max_hops)
        address_page_ids, address_record_ids = self._address_hits(query_text)
        lexical_page_ids, lexical_record_ids = self._lexical_evidence_hits(query_text)
        content_record_ids = lexical_record_ids.difference(address_record_ids)
        routed_page_ids = address_page_ids | lexical_page_ids
        if routed_page_ids:
            # Explicit addresses are already a high-confidence routing signal.
            # Avoid paying the neural page rerank cost for unrelated candidates;
            # exact records below still receive the final bounded score.
            ranked_pages = []
            for page_id in sorted(routed_page_ids):
                page = self._hydrate_page(page_id)
                if page is None:
                    continue
                page_score = 3.0 + 0.05 * page.importance
                if page_id in lexical_page_ids:
                    page_score += 0.5
                ranked_pages.append((page, page_score))
        else:
            ranked_pages = self._page_scores(query_key, query_text, query_token_ids)
        selected_pages = [item for item in ranked_pages[:page_limit] if item[1] >= min_score]
        # Preserve exact address hits even when the learned page score ranks a
        # mixed-topic page above the target page.  Only the small indexed set of
        # distinctive address pages is admitted; this does not scan the bank.
        ranked_by_page = {page.page_id: (page, score) for page, score in ranked_pages}
        selected_page_ids = {page.page_id for page, _ in selected_pages}
        if routed_page_ids:
            best_selected_score = max((score for _, score in selected_pages), default=0.0)
            for page_id in sorted(routed_page_ids):
                if page_id in selected_page_ids:
                    continue
                ranked = ranked_by_page.get(page_id)
                if ranked is None:
                    page = self._hydrate_page(page_id)
                    if page is None:
                        continue
                    ranked = (page, best_selected_score)
                page, score = ranked
                selected_pages.append((page, max(float(score), best_selected_score) + 3.0))
                selected_page_ids.add(page_id)
        current_page_ids = [page.page_id for page, _ in selected_pages]
        selected_records: list[tuple[MemoryRecordV2, float]] = []
        hop_trace: list[list[str]] = []
        visited: set[str] = set()
        # Direct address records are inserted before semantic candidates.  The
        # neural router still scores ordinary candidates, but a distinctive
        # filename/symbol or user id must never be lost to a centroid collision.
        direct_records: list[tuple[MemoryRecordV2, float]] = []
        direct_record_ids = address_record_ids | content_record_ids
        if direct_record_ids:
            self._hydrate_records(direct_record_ids)
            for record_id in sorted(direct_record_ids):
                record = self.records.get(record_id)
                if record is None or record.status != STATUS_ACTIVE:
                    continue
                if record_id in address_record_ids and not self._record_matches_explicit_address(record, query_text):
                    continue
                score = (10.0 if record_id in address_record_ids else 8.0) + 0.25 * record.lexical_score(query_text)
                if record.token_ids is not None:
                    score += 0.35
                if record.entity or record.attribute or record.value:
                    score += 0.15
                direct_records.append((record, score))
            direct_records.sort(key=lambda item: item[1], reverse=True)
            # If a query names an ambiguous symbol/value, return a bounded
            # evidence set rather than silently choosing one occurrence.  The
            # caller still receives only a small Top-K set for ordinary
            # queries; fan-out is capped by the number of sparse address hits.
            direct_limit = max(record_limit, min(8, len(direct_records)))
            selected_records.extend(direct_records[:direct_limit])
            for record, _ in selected_records:
                if record.record_id not in visited:
                    visited.add(record.record_id)
                    record.touch()
                    self._store_record(record)
        # An exact address is already the final routing decision. Do not pay
        # for semantic page/record reranking or unrelated multi-hop expansion
        # when the query names a distinctive entity, filename, symbol, or user
        # id. Ambiguous queries still use the complete learned path below.
        skip_semantic_expansion = bool(direct_record_ids and direct_records)
        for hop in range(0 if skip_semantic_expansion else hop_limit):
            hop_records: list[str] = []
            candidate_pages = [self.pages[item] for item in current_page_ids if item in self.pages]
            candidates = []
            for page in candidate_pages:
                candidates.extend(self._record_scores(page, query_key, query_text, query_token_ids))
            candidates.sort(key=lambda item: item[1], reverse=True)
            for record, score in candidates:
                if record.record_id in visited:
                    continue
                visited.add(record.record_id)
                selected_records.append((record, score))
                record.touch()
                self._store_record(record)
                hop_records.append(record.record_id)
                if len(selected_records) >= record_limit:
                    break
            hop_trace.append(hop_records)
            if len(selected_records) >= record_limit or not hop_records:
                break
            related_pages: list[str] = []
            for record_id in hop_records:
                record = self.records[record_id]
                self._hydrate_records(record.related_ids)
                for related_id in record.related_ids:
                    related = self.records.get(related_id)
                    if related is not None and related.page_id not in related_pages:
                        related_pages.append(related.page_id)
            if not related_pages:
                break
            current_page_ids = related_pages[:page_limit]
        selected_records.sort(key=lambda item: item[1], reverse=True)
        # When the query names one distinctive entity and its attribute, a
        # semantically similar second record is a liability: it can make the
        # language model blend two people or two versions into one answer.
        # Treat the explicit entity/attribute pair as an address constraint
        # and retain only records at that address.  Broad queries without a
        # sufficiently distinctive entity keep the normal top-k behavior.
        exact_address_records = [
            item
            for item in selected_records
            if self._record_matches_explicit_address(item[0], query_text)
        ]
        if exact_address_records:
            selected_records = exact_address_records
        # Ambiguous code symbols need more than one evidence card: returning
        # only the first occurrence makes a symbol such as ``run`` look like
        # a unique answer even when it exists in several files.  Keep this
        # expansion local to symbol-address queries so ordinary personal
        # memory reads retain the normal small Top-K prompt.
        symbol_address_query = any(
            record.attribute.strip().lower().startswith(("symbol:", "class:", "def:"))
            for record, _ in direct_records
        )
        output_limit = min(
            len(selected_records),
            max(record_limit, 8) if symbol_address_query else record_limit,
        )
        records = [record for record, _ in selected_records[:output_limit]]
        # Keep only the bounded hot working set on VRAM.  The authoritative
        # copies remain in RAM and are exported to the embedded weight shard.
        self.promote_records(records)
        page_scores = [score for _, score in selected_pages]
        record_scores = [score for _, score in selected_records[:record_limit]]
        top_score = float(record_scores[0]) if record_scores else 0.0
        second_score = float(record_scores[1]) if len(record_scores) > 1 else 0.0
        score_margin = max(0.0, top_score - second_score) if record_scores else 0.0
        confidence = max(0.0, min(1.0, top_score))
        decision = RouterDecisionV2(
            need_memory=bool(records),
            page_ids=[page.page_id for page, _ in selected_pages],
            record_ids=[record.record_id for record in records],
            page_scores=page_scores,
            record_scores=record_scores,
            hop_count=len(hop_trace),
            hop_trace=hop_trace,
            confidence=confidence,
            stop_reason="evidence_found" if records else "no_relevant_evidence",
            top_score=top_score,
            score_margin=score_margin,
            evidence_score=confidence,
        )
        for page, _ in selected_pages:
            page.touch()
        self._refresh_tiers()
        return records, decision

    def correct(
        self,
        *,
        text: str,
        key: Tensor,
        entity: str,
        attribute: str,
        value: str,
        evidence: Optional[Sequence[str]] = None,
        confidence: float = 1.0,
    ) -> tuple[MemoryRecordV2, str]:
        """Write an explicit correction as a new version."""

        return self.write(
            text=text,
            key=key,
            entity=entity,
            attribute=attribute,
            value=value,
            evidence=evidence,
            confidence=confidence,
            importance=1.0,
            source="user_correction",
            trusted=True,
            force=True,
        )

    def approve(self, record_id: str) -> MemoryRecordV2:
        record = self.quarantine.pop(record_id, None)
        if record is None:
            raise KeyError(f"unknown quarantined record: {record_id}")
        if self.tier_store is not None:
            self.tier_store.delete_quarantine(record_id)
        record.status = STATUS_ACTIVE
        page = self._target_page(record.key, importance=record.importance)
        self.records[record.record_id] = record
        self._update_page(page, record)
        self._store_record(record)
        conflict_key = record.conflict_key()
        if conflict_key:
            old_id = self.active_by_conflict.get(conflict_key)
            if old_id and old_id not in self.records:
                self._hydrate_records([old_id])
            old = self.records.get(old_id) if old_id else None
            if old is not None and old.status == STATUS_ACTIVE:
                old.status = STATUS_SUPERSEDED
                record.version = old.version + 1
                record.supersedes = old.record_id
                self._store_record(old)
            self.active_by_conflict[conflict_key] = record.record_id
        self._refresh_tiers()
        return record

    def retract(self, record_id: str) -> None:
        record = self.records.get(record_id)
        if record is None and self.tier_store is not None:
            self._hydrate_records([record_id])
            record = self.records.get(record_id)
        if record is None:
            raise KeyError(record_id)
        record.status = STATUS_RETRACTED
        self._remove_gpu_record(record_id)
        self._store_record(record)
        conflict_key = record.conflict_key()
        if conflict_key and self.active_by_conflict.get(conflict_key) == record_id:
            del self.active_by_conflict[conflict_key]
        self._refresh_tiers()

    def _absorb_unidentified_same_turn(self, record: MemoryRecordV2) -> list[str]:
        """Supersede the same turn's identity-less record once a structured one lands.

        The automatic layer writes one record per user turn and runs *before* the model
        generates, so it cannot know what that turn will turn out to be about.  When the
        sentence is not the ``<attribute> 是 <value>`` shape its metadata inference
        understands (「我叫Wpy」, 「我住在上海」), the automatic record ends up with no
        entity and no attribute -- hence no conflict key, hence a correction can never
        retire it, and a stale value stays retrievable beside the corrected one.  Measured:
        two contradictory identity records both stayed ``active`` at version 0.

        The structured record for that turn does carry a key.  Absorbing the turn's
        unidentified record into it -- superseded rather than retracted, since it is the
        same fact in a worse shape -- leaves exactly one active record per turn and removes
        the uncorrectable leftover.

        Only records sharing this record's ``origin`` are touched, and only active ones, so
        this cannot reach another turn's facts or an already retired one.
        """

        origin = str(record.origin or "").strip()
        if not origin or not (record.entity and record.attribute):
            return []
        absorbed: list[str] = []
        for sibling in list(self.records.values()):
            if sibling.record_id == record.record_id or sibling.status != STATUS_ACTIVE:
                continue
            if sibling.origin != origin or (sibling.entity and sibling.attribute):
                continue
            sibling.status = STATUS_SUPERSEDED
            self._store_record(sibling)
            absorbed.append(sibling.record_id)
        return absorbed

    def retract_origin(self, origin: str) -> list[str]:
        """Retract every active record written from the same user turn.

        Retraction has always been *key*-scoped: a forget request retracts the record
        whose key matched.  One turn usually writes more than one record, though -- a
        sentence like "my emergency contact is Wang, extension 7781" produces a contact
        record and an extension record -- so forgetting "the emergency contact
        information" used to leave the extension active and readable.  A measured
        end-to-end run showed the model then answering with that leftover value, i.e.
        the memory layer leaking a fact the user had just revoked.

        Grouping by origin makes the unit of forgetting the unit of telling, which is
        what the user meant.  Records with no origin (written before this field
        existed) are never matched, so this cannot retract anything unrelated.
        """

        origin = str(origin or "").strip()
        if not origin:
            return []
        retracted: list[str] = []
        for record in list(self.records.values()):
            if record.status != STATUS_ACTIVE or record.origin != origin:
                continue
            self.retract(record.record_id)
            retracted.append(record.record_id)
        return retracted

    def consolidate(self, *, max_pages: Optional[int] = None) -> int:
        """Move low-value pages toward compact summaries without deleting evidence."""

        target = max_pages or self.max_pages
        if len(self.pages) <= target:
            self._refresh_tiers()
            return 0
        ranked = sorted(self.pages.values(), key=lambda page: (page.importance, page.last_access))
        merged = 0
        for source in ranked:
            if len(self.pages) <= target:
                break
            if source.tier == "hot" or not source.record_ids:
                continue
            destinations = [
                page
                for page in self.pages.values()
                if page.page_id != source.page_id and not page.full
            ]
            # Consolidation may only use an existing page.  Creating a new
            # page while trying to enforce a page limit would recurse forever
            # when every page is already full.
            if not destinations:
                continue
            source_key = source.key if source.key is not None else torch.zeros(self.key_dim)
            source_key = self._encode_key(source_key)
            destination = max(
                destinations,
                key=lambda page: (
                    float(torch.dot(source_key, F.normalize(page.key.float(), dim=0)))
                    if page.key is not None
                    else -0.1 * len(page.record_ids)
                ),
            )
            for record_id in list(source.record_ids):
                if record_id not in destination.record_ids and not destination.full:
                    destination.record_ids.append(record_id)
                    self.records[record_id].page_id = destination.page_id
            if destination.full:
                self._open_page_ids.discard(destination.page_id)
            else:
                self._open_page_ids.add(destination.page_id)
            self._remove_page_from_index(destination)
            self._index_page(destination)
            if destination.key is None:
                destination.key = source.key
            if destination.summary is None:
                destination.summary = source.summary
            destination.importance = max(destination.importance, source.importance)
            self._remove_page_from_index(source)
            del self.pages[source.page_id]
            self._open_page_ids.discard(source.page_id)
            merged += 1
        self._refresh_tiers()
        return merged

    def stats(self) -> dict[str, Any]:
        active = [record for record in self.records.values() if record.status == STATUS_ACTIVE]
        stored = self.tier_store.count() if self.tier_store is not None else {}
        record_count = int(stored.get("records", len(self.records)))
        active_count = int(stored.get("status_active", len(active)))
        superseded_count = int(
            stored.get(
                "status_superseded",
                sum(record.status == STATUS_SUPERSEDED for record in self.records.values()),
            )
        )
        retracted_count = int(
            stored.get(
                "status_retracted",
                sum(record.status == STATUS_RETRACTED for record in self.records.values()),
            )
        )
        quarantined_count = int(stored.get("quarantined", len(self.quarantine)))
        return {
            "pages": len(self.pages),
            "max_pages": self.max_pages,
            "page_capacity": self.page_capacity,
            "capacity_records": self.max_pages * self.page_capacity,
            "open_pages": len(self._open_page_ids),
            "records": record_count,
            "active_records": active_count,
            "superseded_records": superseded_count,
            "retracted_records": retracted_count,
            "quarantined_records": quarantined_count,
            "hot_pages": int(stored.get("hot_pages", sum(page.tier == "hot" for page in self.pages.values()))),
            "warm_pages": int(stored.get("warm_pages", sum(page.tier == "warm" for page in self.pages.values()))),
            "cold_pages": int(stored.get("cold_pages", sum(page.tier == "cold" for page in self.pages.values()))),
            "active_conflict_keys": len(self.active_by_conflict),
            "coarse_index_buckets": (
                self.tier_store.coarse_bucket_count()
                if self.tier_store is not None
                else len(self._coarse_buckets)
            ),
            "last_coarse_candidates": self._last_coarse_candidates,
            "storage_mode": "tiered" if self.tier_store is not None else "embedded",
            "resident_records": len(self.records),
            "resident_pages": len(self.pages),
            "gpu_cache_records": len(self._gpu_record_cache),
            "gpu_cache_tokens": self._gpu_cache_tokens_used,
            "gpu_cache_bytes": self._gpu_cache_bytes_used,
            "gpu_cache_hits": self._gpu_cache_hits,
            "gpu_cache_misses": self._gpu_cache_misses,
            "gpu_cache_fallbacks": self._gpu_cache_fallbacks,
            "gpu_cache_alloc_failures": self._gpu_cache_alloc_failures,
            "gpu_cache_device": str(self._gpu_device) if self._gpu_device is not None else "none",
            "gpu_cache_limit_records": self.gpu_cache_records,
            "gpu_cache_limit_tokens": self.gpu_cache_tokens,
            "gpu_cache_reserve_mb": self.gpu_cache_reserve_mb,
            "gpu_cache_adaptive": self.gpu_cache_adaptive,
            "gpu_cache_last_free_mb": (
                round(self._gpu_cache_last_free_bytes / (1024 * 1024), 2)
                if self._gpu_cache_last_free_bytes is not None
                else None
            ),
        }

    def _resolve_record(self, record_id: str) -> MemoryRecordV2:
        record = self.records.get(str(record_id))
        if record is None and self.tier_store is not None:
            self._hydrate_records([str(record_id)])
            record = self.records.get(str(record_id))
        if record is None:
            record = self.quarantine.get(str(record_id))
        if record is None:
            raise KeyError(str(record_id))
        return record

    def list_records(
        self,
        *,
        query_text: str = "",
        status: str = STATUS_ACTIVE,
        limit: int = 100,
        offset: int = 0,
    ) -> list[MemoryRecordV2]:
        """List records for the local management API without exposing vectors."""

        if limit < 1 or limit > 10000:
            raise ValueError("limit must be between 1 and 10000")
        if offset < 0:
            raise ValueError("offset must be non-negative")
        allowed = {STATUS_ACTIVE, STATUS_SUPERSEDED, STATUS_RETRACTED, STATUS_QUARANTINED, "all"}
        if status not in allowed:
            raise ValueError(f"status must be one of: {', '.join(sorted(allowed))}")
        candidates = list(self.records.values()) + list(self.quarantine.values())
        if status != "all":
            candidates = [item for item in candidates if item.status == status]
        if query_text.strip():
            candidates = [
                item for item in candidates
                if item.lexical_score(query_text) > 0.0 or query_text.strip().lower() in item.text.lower()
            ]
            candidates.sort(
                key=lambda item: (item.lexical_score(query_text), item.last_access, item.timestamp),
                reverse=True,
            )
        else:
            candidates.sort(key=lambda item: (item.last_access, item.timestamp), reverse=True)
        return candidates[offset : offset + limit]

    def edit_record(
        self,
        record_id: str,
        *,
        text: Optional[str] = None,
        key: Optional[Tensor] = None,
        summary: Optional[Tensor] = None,
        memory_type: Optional[str] = None,
        entity: Optional[str] = None,
        attribute: Optional[str] = None,
        value: Optional[str] = None,
        importance: Optional[float] = None,
        confidence: Optional[float] = None,
        evidence: Optional[Sequence[str]] = None,
        token_ids: Optional[Tensor] = None,
        token_mask: Optional[Tensor] = None,
        source: str = "user_edit",
    ) -> MemoryRecordV2:
        """Create a new version and preserve the old record as superseded."""

        old = self._resolve_record(record_id)
        if old.status != STATUS_ACTIVE:
            raise ValueError(f"only active records can be edited: {record_id}")
        old.status = STATUS_SUPERSEDED
        self._remove_gpu_record(old.record_id)
        self._store_record(old)
        old_conflict = old.conflict_key()
        if old_conflict and self.active_by_conflict.get(old_conflict) == old.record_id:
            del self.active_by_conflict[old_conflict]
        next_text = old.text if text is None else str(text)
        next_entity = old.entity if entity is None else str(entity)
        next_attribute = old.attribute if attribute is None else str(attribute)
        next_value = old.value if value is None else str(value)
        next_importance = old.importance if importance is None else float(importance)
        next_confidence = old.confidence if confidence is None else float(confidence)
        next_key = old.key if key is None else key
        next_summary = old.summary if summary is None else summary
        now = _now()
        new_record = MemoryRecordV2(
            record_id=_stable_id(f"{old.record_id}|edit|{now}|{len(self.records)}", prefix="mem"),
            text=next_text,
            key=self._encode_key(next_key).cpu().clone(),
            summary=self._encode_key(next_summary).cpu().clone(),
            memory_type=old.memory_type if memory_type is None else str(memory_type),
            entity=next_entity,
            attribute=next_attribute,
            value=next_value,
            timestamp=now,
            importance=max(0.0, min(1.0, next_importance)),
            confidence=max(0.0, min(1.0, next_confidence)),
            source=str(source),
            status=STATUS_ACTIVE,
            version=old.version + 1,
            supersedes=old.record_id,
            related_ids=list(old.related_ids),
            evidence=list(old.evidence if evidence is None else evidence),
            # A new version of a fact came from whatever turn produced the fact.
            origin=old.origin,
            slot_index=old.slot_index,
            token_ids=(token_ids.detach().cpu().long().reshape(-1).clone() if token_ids is not None else old.token_ids),
            token_mask=(token_mask.detach().cpu().bool().reshape(-1).clone() if token_mask is not None else old.token_mask),
        )
        self.records[new_record.record_id] = new_record
        if new_record.text.strip():
            self._text_index[new_record.text.strip().lower()] = new_record.record_id
        page = self._target_page(new_record.key, importance=new_record.importance)
        self._update_page(page, new_record)
        self._store_record(new_record)
        new_conflict = new_record.conflict_key()
        if new_conflict:
            self.active_by_conflict[new_conflict] = new_record.record_id
        self._refresh_tiers()
        return new_record

    def _remove_gpu_record(self, record_id: str) -> None:
        cached = self._gpu_record_cache.pop(record_id, None)
        if cached is None:
            return
        if record_id in self._gpu_cache_order:
            self._gpu_cache_order.remove(record_id)
        token_ids = cached.get("token_ids")
        if isinstance(token_ids, Tensor):
            self._gpu_cache_tokens_used -= int(token_ids.numel())
        self._gpu_cache_bytes_used -= sum(self._tensor_bytes(value) for value in cached.values())
        self._gpu_cache_tokens_used = max(0, self._gpu_cache_tokens_used)
        self._gpu_cache_bytes_used = max(0, self._gpu_cache_bytes_used)

    def audit(self) -> dict[str, Any]:
        """Check page membership, conflict indexes and relationship references."""

        issues: list[str] = []
        page_membership: set[str] = set()
        for page in self.pages.values():
            if len(page.record_ids) > page.capacity:
                issues.append(f"page_over_capacity:{page.page_id}")
            for record_id in page.record_ids:
                page_membership.add(record_id)
                record = self.records.get(record_id)
                if record is None:
                    issues.append(f"missing_record:{record_id}")
                elif record.page_id != page.page_id:
                    issues.append(f"page_pointer_mismatch:{record_id}")
        for record in self.records.values():
            if record.status == STATUS_ACTIVE and record.record_id not in page_membership:
                issues.append(f"active_record_not_indexed:{record.record_id}")
            for related_id in record.related_ids:
                if related_id not in self.records and related_id not in self.quarantine:
                    issues.append(f"dangling_related_id:{record.record_id}->{related_id}")
        for conflict_key, record_id in self.active_by_conflict.items():
            record = self.records.get(record_id)
            if record is None or record.status != STATUS_ACTIVE or record.conflict_key() != conflict_key:
                issues.append(f"invalid_conflict_index:{conflict_key}")
        return {
            "healthy": not issues,
            "issues": issues[:100],
            "issue_count": len(issues),
            "stats": self.stats(),
        }

    def flush_storage(self) -> None:
        """Flush the durable tier without serializing cold pages into RAM."""

        if self.tier_store is not None:
            self.tier_store.flush()

    def close_storage(self) -> None:
        """Close the durable page database so Windows can remove or rotate it."""

        if self.tier_store is not None:
            self.tier_store.close()

    def export_payload(self) -> dict[str, Any]:
        """Export tensors plus JSON-safe metadata for checkpoint backends."""

        if self.tier_store is not None:
            for page in self.pages.values():
                self._hydrate_page(page.page_id)
        records = []
        for record in self.records.values():
            item = asdict(record)
            item["key"] = record.key.detach().cpu()
            item["summary"] = record.summary.detach().cpu()
            records.append(item)
        pages = []
        for page in self.pages.values():
            item = asdict(page)
            item["key"] = page.key.detach().cpu() if page.key is not None else None
            item["summary"] = page.summary.detach().cpu() if page.summary is not None else None
            pages.append(item)
        quarantine = []
        for record in self.quarantine.values():
            item = asdict(record)
            item["key"] = record.key.detach().cpu()
            item["summary"] = record.summary.detach().cpu()
            quarantine.append(item)
        return {
            "format_version": 2,
            "hidden_size": self.hidden_size,
            "page_capacity": self.page_capacity,
            "max_pages": self.max_pages,
            "hot_pages": self.hot_pages,
            "top_k_pages": self.top_k_pages,
            "top_k_records": self.top_k_records,
            "max_hops": self.max_hops,
            "key_dim": self.key_dim,
            "coarse_index_bits": self.coarse_index_bits,
            "record_count": len(self.records),
            "counter": self._counter,
            "records": records,
            "pages": pages,
            "quarantine": quarantine,
            "active_by_conflict": dict(self.active_by_conflict),
        }

    @classmethod
    def from_payload(
        cls,
        payload: dict[str, Any],
        *,
        router: Optional[MemoryRouterV2] = None,
        tier_store: Optional["TieredMemoryStoreV2"] = None,
        max_resident_pages: int = 256,
        runtime_device: Optional[torch.device] = None,
        gpu_cache_records: int = 256,
        gpu_cache_tokens: int = 131072,
        gpu_cache_reserve_mb: int = 2048,
        gpu_cache_adaptive: bool = True,
        record_scorer: Optional[Callable[[Tensor, Tensor], Tensor]] = None,
    ) -> "PagedMemoryBankV2":
        bank = cls(
            int(payload["hidden_size"]),
            page_capacity=int(payload.get("page_capacity", 32)),
            max_pages=int(payload.get("max_pages", 32768)),
            hot_pages=int(payload.get("hot_pages", 8)),
            top_k_pages=int(payload.get("top_k_pages", 4)),
            top_k_records=int(payload.get("top_k_records", 8)),
            max_hops=int(payload.get("max_hops", 3)),
            router=router,
            key_dim=int(payload.get("key_dim", router.router_dim if router is not None else payload["hidden_size"])),
            coarse_index_bits=int(payload.get("coarse_index_bits", 20)),
            tier_store=tier_store,
            max_resident_pages=max_resident_pages,
            runtime_device=runtime_device,
            gpu_cache_records=gpu_cache_records,
            gpu_cache_tokens=gpu_cache_tokens,
            gpu_cache_reserve_mb=gpu_cache_reserve_mb,
            gpu_cache_adaptive=gpu_cache_adaptive,
            record_scorer=record_scorer,
        )
        bank._counter = int(payload.get("counter", 0))
        for item in payload.get("records", []):
            item = dict(item)
            item["key"] = torch.as_tensor(item["key"]).float()
            item["summary"] = torch.as_tensor(item["summary"]).float()
            if item.get("semantic_key") is not None:
                item["semantic_key"] = torch.as_tensor(item["semantic_key"]).float()
            if item.get("token_ids") is not None:
                item["token_ids"] = torch.as_tensor(item["token_ids"]).long()
            if item.get("token_mask") is not None:
                item["token_mask"] = torch.as_tensor(item["token_mask"]).bool()
            bank.records[item["record_id"]] = MemoryRecordV2(**item)
            if item.get("text", "").strip() and item.get("status") == STATUS_ACTIVE:
                bank._text_index[item["text"].strip().lower()] = item["record_id"]
        for item in payload.get("pages", []):
            item = dict(item)
            if item.get("key") is not None:
                item["key"] = torch.as_tensor(item["key"]).float()
            if item.get("summary") is not None:
                item["summary"] = torch.as_tensor(item["summary"]).float()
            if item.get("semantic_key") is not None:
                item["semantic_key"] = torch.as_tensor(item["semantic_key"]).float()
            bank.pages[item["page_id"]] = MemoryPageV2(**item)
        bank._open_page_ids = {
            page.page_id for page in bank.pages.values() if not page.full
        }
        bank._open_page_order = [
            page.page_id
            for page in sorted(
                bank.pages.values(), key=lambda item: (item.created_at, item.page_id)
            )
            if page.page_id in bank._open_page_ids
        ]
        for page in bank.pages.values():
            bank._index_page(page)
        for item in payload.get("quarantine", []):
            item = dict(item)
            item["key"] = torch.as_tensor(item["key"]).float()
            item["summary"] = torch.as_tensor(item["summary"]).float()
            if item.get("token_ids") is not None:
                item["token_ids"] = torch.as_tensor(item["token_ids"]).long()
            if item.get("token_mask") is not None:
                item["token_mask"] = torch.as_tensor(item["token_mask"]).bool()
            bank.quarantine[item["record_id"]] = MemoryRecordV2(**item)
        bank.active_by_conflict = dict(payload.get("active_by_conflict", {}))
        bank._refresh_tiers()
        return bank


@dataclass
class KVBudgetManagerV2:
    """Token budget and overflow policy for the hot working context."""

    max_tokens: int = 32768
    hard_max_tokens: int = 131072
    compaction_trigger: float = 0.90
    keep_recent_tokens: int = 8192

    def __post_init__(self) -> None:
        if not 1 <= self.max_tokens <= self.hard_max_tokens:
            raise ValueError("max_tokens must be inside [1, hard_max_tokens]")
        if not 0.5 <= self.compaction_trigger < 1.0:
            raise ValueError("compaction_trigger must be in [0.5, 1)")
        self.keep_recent_tokens = min(self.keep_recent_tokens, self.max_tokens)

    @property
    def trigger_tokens(self) -> int:
        return max(1, int(self.max_tokens * self.compaction_trigger))

    def needs_compaction(self, token_count: int) -> bool:
        return int(token_count) >= self.trigger_tokens

    def overflow(self, token_count: int) -> int:
        return max(0, int(token_count) - self.max_tokens)

    def retention_plan(self, token_count: int) -> dict[str, int | bool]:
        overflow = self.overflow(token_count)
        return {
            "token_count": int(token_count),
            "max_tokens": self.max_tokens,
            "needs_compaction": self.needs_compaction(token_count),
            "overflow_tokens": overflow,
            "preserve_recent_tokens": self.keep_recent_tokens,
            "tokens_for_memory": max(0, overflow + max(0, int(token_count * 0.1))),
        }


class MemoryOSV2:
    """Model-facing coordinator for a hierarchical memory bank."""

    def __init__(
        self,
        hidden_size: int,
        *,
        router: Optional[MemoryRouterV2] = None,
        bank: Optional[PagedMemoryBankV2] = None,
        kv_budget: Optional[KVBudgetManagerV2] = None,
        read_threshold: float = 0.65,
        write_threshold: float = 0.50,
        runtime_device: Optional[torch.device] = None,
        gpu_cache_records: int = 256,
        gpu_cache_tokens: int = 131072,
        gpu_cache_reserve_mb: int = 2048,
        gpu_cache_adaptive: bool = True,
        record_scorer: Optional[Callable[[Tensor, Tensor], Tensor]] = None,
        min_read_margin: float = 0.0,
        require_evidence: bool = False,
        attribute_coverage: Optional[Callable[[Tensor], tuple[str, float]]] = None,
    ) -> None:
        self.router = router or MemoryRouterV2(hidden_size)
        self.bank = bank or PagedMemoryBankV2(
            hidden_size,
            router=self.router,
            runtime_device=runtime_device,
            gpu_cache_records=gpu_cache_records,
            gpu_cache_tokens=gpu_cache_tokens,
            gpu_cache_reserve_mb=gpu_cache_reserve_mb,
            gpu_cache_adaptive=gpu_cache_adaptive,
            record_scorer=record_scorer,
        )
        if record_scorer is not None:
            self.bank.record_scorer = record_scorer
        # Optional coverage gate.  ``attribute_coverage(query_key)`` returns the attribute
        # the question asks about (or "" when the head says none), and the runtime refuses
        # when the bank does not hold it.  This exists because score geometry cannot detect
        # an unheld attribute: measured over nine score features, the best fitted head
        # reaches AUC 0.61 and three of four scorers sit at ~0.5, while same-shape
        # candidates look equally plausible for any question.
        self.attribute_coverage = attribute_coverage
        self.kv_budget = kv_budget or KVBudgetManagerV2()
        self.read_threshold = read_threshold
        self.write_threshold = write_threshold
        self.min_read_margin = float(min_read_margin)
        self.require_evidence = bool(require_evidence)

    def write(self, **kwargs: Any) -> tuple[MemoryRecordV2, str]:
        importance = float(kwargs.get("importance", 0.5))
        confidence = float(kwargs.get("confidence", 0.5))
        force = bool(kwargs.get("force", False))
        trusted = bool(kwargs.get("trusted", True))
        if not force and (importance < self.write_threshold or confidence < 0.25):
            kwargs["trusted"] = False
        return self.bank.write(**kwargs)

    def write_batch(self, records: Iterable[dict[str, Any]]) -> list[tuple[MemoryRecordV2, str]]:
        """Apply the V2 write gate and commit a batch atomically when possible."""

        prepared: list[dict[str, Any]] = []
        for item in records:
            kwargs = dict(item)
            importance = float(kwargs.get("importance", 0.5))
            confidence = float(kwargs.get("confidence", 0.5))
            force = bool(kwargs.get("force", False))
            if not force and (importance < self.write_threshold or confidence < 0.25):
                kwargs["trusted"] = False
            prepared.append(kwargs)
        return self.bank.write_batch(prepared)

    def read(
        self,
        *,
        query_key: Tensor,
        query_text: str = "",
        query_token_ids: Optional[Tensor] = None,
        top_k_pages: Optional[int] = None,
        top_k_records: Optional[int] = None,
        max_hops: Optional[int] = None,
    ) -> tuple[list[MemoryRecordV2], RouterDecisionV2]:
        if _is_explicit_unknown_request(query_text):
            return [], RouterDecisionV2(
                need_memory=False,
                page_ids=[],
                record_ids=[],
                page_scores=[],
                record_scores=[],
                hop_count=0,
                hop_trace=[],
                confidence=0.0,
                stop_reason="explicit_unknown_request",
            )
        if self.attribute_coverage is not None and query_key.numel() > 0:
            verdict = self.attribute_coverage(query_key)
            # ``None`` means the gate declines to judge (its vocabulary does not cover this
            # bank); the runtime then behaves exactly as if the gate were absent.
            if verdict is not None:
                attribute, coverage_probability = verdict
                # Match the predicted attribute against every entity in the bank.  The
                # previous check hard-coded the entity name (`f"user::{attribute}"`), while
                # the head's own `coverage()` predicate derives presence from
                # ``key.split("::", 1)[1]`` and therefore treats the entity as arbitrary.
                # Whenever a bank stores records under any other entity (explicit writes,
                # agent-authored memories, imported users) the gate then engaged and
                # refused *every* question, including attributes it actually held.
                attribute_key = attribute.strip().lower() if attribute else ""
                covered = bool(attribute_key) and any(
                    key.rsplit("::", 1)[-1].strip().lower() == attribute_key
                    for key in self.bank.active_by_conflict
                )
                if not covered:
                    return [], RouterDecisionV2(
                        need_memory=False,
                        page_ids=[],
                        record_ids=[],
                        page_scores=[],
                        record_scores=[],
                        hop_count=0,
                        hop_trace=[],
                        confidence=float(coverage_probability),
                        stop_reason=("attribute_not_covered" if attribute_key
                                     else "no_attribute_recognised"),
                    )
        if self.router is not None and query_key.numel() == self.router.hidden_size:
            router_device = next(self.router.parameters()).device
            need_probability = float(
                torch.sigmoid(self.router.need_memory(query_key.to(router_device).reshape(1, -1)))[0, 0].item()
            )
            token_evidence = self.bank.has_token_evidence(query_key, query_token_ids, query_text)
            if need_probability < self.read_threshold and not token_evidence:
                return [], RouterDecisionV2(
                    need_memory=False,
                    page_ids=[],
                    record_ids=[],
                    page_scores=[],
                    record_scores=[],
                    hop_count=0,
                    hop_trace=[],
                    confidence=need_probability,
                    stop_reason="router_abstained",
                )
        records, decision = self.bank.query(
            query_key=query_key,
            query_text=query_text,
            query_token_ids=query_token_ids,
            top_k_pages=top_k_pages,
            top_k_records=top_k_records,
            max_hops=max_hops,
            min_score=-1.0,
        )
        token_evidence = self.bank.has_token_evidence(query_key, query_token_ids, query_text)
        if (
            self.require_evidence
            and records
            and not token_evidence
            and decision.score_margin < self.min_read_margin
        ):
            decision.need_memory = False
            decision.record_ids = []
            decision.stop_reason = "insufficient_evidence_margin"
            return [], decision
        if decision.confidence < self.read_threshold and not token_evidence:
            decision.need_memory = False
            decision.record_ids = []
            decision.stop_reason = "below_read_threshold"
            return [], decision
        if token_evidence and decision.confidence < self.read_threshold:
            decision.stop_reason = "token_evidence_override"
        return records, decision

    def correct(self, **kwargs: Any) -> tuple[MemoryRecordV2, str]:
        return self.bank.correct(**kwargs)

    def approve(self, record_id: str) -> MemoryRecordV2:
        return self.bank.approve(record_id)

    def retract(self, record_id: str) -> None:
        self.bank.retract(record_id)

    def stats(self) -> dict[str, Any]:
        output = self.bank.stats()
        output["kv_budget"] = {
            "max_tokens": self.kv_budget.max_tokens,
            "trigger_tokens": self.kv_budget.trigger_tokens,
        }
        return output

    def list_records(self, **kwargs: Any) -> list[MemoryRecordV2]:
        return self.bank.list_records(**kwargs)

    def edit_record(self, record_id: str, **kwargs: Any) -> MemoryRecordV2:
        return self.bank.edit_record(record_id, **kwargs)

    def retract_record(self, record_id: str) -> None:
        self.bank.retract(record_id)

    def retract_origin(self, origin: str) -> list[str]:
        """Retract every active record written from the same user turn."""
        return self.bank.retract_origin(origin)

    def audit(self) -> dict[str, Any]:
        return self.bank.audit()

    def flush_storage(self) -> None:
        self.bank.flush_storage()

    def close_storage(self) -> None:
        self.bank.close_storage()

    def export_payload(self) -> dict[str, Any]:
        payload = self.bank.export_payload()
        payload["read_threshold"] = self.read_threshold
        payload["write_threshold"] = self.write_threshold
        payload["min_read_margin"] = self.min_read_margin
        payload["require_evidence"] = self.require_evidence
        payload["kv_budget"] = asdict(self.kv_budget)
        return payload

    @classmethod
    def from_payload(
        cls,
        payload: dict[str, Any],
        *,
        router: Optional[MemoryRouterV2] = None,
        tier_store: Optional["TieredMemoryStoreV2"] = None,
        max_resident_pages: int = 256,
        runtime_device: Optional[torch.device] = None,
        gpu_cache_records: int = 256,
        gpu_cache_tokens: int = 131072,
        gpu_cache_reserve_mb: int = 2048,
        gpu_cache_adaptive: bool = True,
        record_scorer: Optional[Callable[[Tensor, Tensor], Tensor]] = None,
        min_read_margin: float = 0.0,
        require_evidence: bool = False,
    ) -> "MemoryOSV2":
        bank = PagedMemoryBankV2.from_payload(
            payload,
            router=router,
            tier_store=tier_store,
            max_resident_pages=max_resident_pages,
            runtime_device=runtime_device,
            gpu_cache_records=gpu_cache_records,
            gpu_cache_tokens=gpu_cache_tokens,
            gpu_cache_reserve_mb=gpu_cache_reserve_mb,
            gpu_cache_adaptive=gpu_cache_adaptive,
            record_scorer=record_scorer,
        )
        budget = KVBudgetManagerV2(**payload.get("kv_budget", {}))
        return cls(
            bank.hidden_size,
            router=router,
            bank=bank,
            kv_budget=budget,
            read_threshold=float(payload.get("read_threshold", 0.65)),
            write_threshold=float(payload.get("write_threshold", 0.5)),
            record_scorer=record_scorer,
            min_read_margin=float(payload.get("min_read_margin", min_read_margin)),
            require_evidence=bool(payload.get("require_evidence", require_evidence)),
        )


def router_training_loss(
    router: MemoryRouterV2,
    query: Tensor,
    candidates: Tensor,
    positive_index: Tensor,
    *,
    need_memory_label: Optional[Tensor] = None,
    hop_label: Optional[Tensor] = None,
) -> dict[str, Tensor]:
    """Compute supervised router losses with hard negatives."""

    output = router(query, candidates)
    losses: dict[str, Tensor] = {}
    losses["candidate"] = F.cross_entropy(output["scores"], positive_index)
    if need_memory_label is not None:
        losses["need_memory"] = F.binary_cross_entropy_with_logits(
            output["need_memory_logits"], need_memory_label.float()
        )
    if hop_label is not None:
        losses["hop"] = F.cross_entropy(output["hop_logits"], hop_label)
    losses["total"] = sum(losses.values())
    return losses


__all__ = [
    "KVBudgetManagerV2",
    "MemoryOSV2",
    "MemoryPageV2",
    "MemoryRecordV2",
    "MemoryRouterV2",
    "PagedMemoryBankV2",
    "RouterDecisionV2",
    "router_training_loss",
    "STATUS_ACTIVE",
    "STATUS_QUARANTINED",
    "STATUS_RETRACTED",
    "STATUS_SUPERSEDED",
]