fix(memory): preserve trusted state head across ledger writes

backend/database/memory_apply_store.py

@@ -26,10 +26,7 @@
 )
 from models.memory_operations import MemoryOperation
 from models.product_memory import MemoryItemStatus, MemoryItem
-from utils.memory.v3.account_generation_source import (
-    V3_TRUSTED_ACCOUNT_GENERATION_SCHEMA_VERSION,
-    V3_TRUSTED_ACCOUNT_GENERATION_SOURCE,
-)
+from models.memory_state_head import trusted_memory_state_head_fields
 
 
 class MemoryFirestoreApplyError(Exception):
@@ -452,16 +449,15 @@ def _write_apply_result(
 
 
 def _memory_state_head_from_control(control_state: MemoryControlState) -> MemoryApplyDoc:
-    state_head: MemoryApplyDoc = {
-        "schema_version": V3_TRUSTED_ACCOUNT_GENERATION_SCHEMA_VERSION,
-        "uid": control_state.uid,
-        "source": V3_TRUSTED_ACCOUNT_GENERATION_SOURCE,
-        "account_generation": control_state.account_generation,
-        "head_commit_id": control_state.head_commit_id,
-        "commit_sequence": control_state.commit_sequence,
-        "updated_at": control_state.updated_at,
-    }
-    return state_head
+    trusted_fields = trusted_memory_state_head_fields(
+        uid=control_state.uid,
+        account_generation=control_state.account_generation,
+        head_commit_id=control_state.head_commit_id,
+        commit_sequence=control_state.commit_sequence,
+    )
+    if trusted_fields is None:  # MemoryControlState has already validated this input.
+        raise MemoryFirestoreApplyError("invalid memory state-head control fields")
+    return cast(MemoryApplyDoc, {**trusted_fields, "updated_at": control_state.updated_at})
 
 
 def _required_model(*, ref: Any, transaction: Any, model: type[M], label: str) -> M:

backend/database/memory_ledger.py

@@ -10,6 +10,10 @@
 from database import projection_repair
 from database.firestore_transaction_retry import run_with_transaction_contention_retry
 from models.memories import confidence_fields_for_evidence
+from models.memory_state_head import (
+    trusted_memory_state_head_fields_from_control,
+    trusted_memory_state_head_fields_from_state,
+)
 from ._client import db
 
 T = TypeVar("T")
@@ -40,6 +44,7 @@ def _typed_doc(doc: Any) -> Dict[str, Any]:
 users_collection = 'users'
 memory_state_collection = 'memory_state'
 memory_state_document = 'head'
+memory_apply_control_document = 'apply_control'
 memory_commits_collection = 'memory_commits'
 
 
@@ -50,6 +55,39 @@ def __init__(self, expected_parent: Optional[str], current_head: Optional[str]):
         self.current_head = current_head
 
 
+def _state_head_write_payload(
+    *,
+    transaction: Any,
+    user_ref: Any,
+    uid: str,
+    state: Dict[str, Any],
+    commit: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Keep canonical trusted fields when the legacy ledger advances its own head.
+
+    The legacy ledger's ``current_head_commit_id`` is independent from the
+    canonical apply stream's ``head_commit_id``.  Older ledger writes replaced
+    the shared document and erased the canonical fields.  If an old write has
+    already clobbered them, repair only from the transactional canonical control
+    record; do not fabricate a trusted generation.
+    """
+    trusted_fields = trusted_memory_state_head_fields_from_state(state, uid=uid)
+    if trusted_fields is None:
+        control_ref = user_ref.collection(memory_state_collection).document(memory_apply_control_document)
+        control_snapshot = control_ref.get(transaction=transaction)
+        control = _typed_doc(control_snapshot) if control_snapshot.exists else {}
+        trusted_fields = trusted_memory_state_head_fields_from_control(control, uid=uid)
+
+    payload = {
+        'current_head_commit_id': commit['commit_id'],
+        'projection_version': 1,
+        'updated_at': commit['commit_time'],
+    }
+    if trusted_fields is not None:
+        payload.update(trusted_fields)
+    return payload
+
+
 def mutation(mutation_type: str, **payload: Any) -> Dict[str, Any]:
     return {'type': mutation_type, **payload}
 
@@ -288,11 +326,13 @@ def _append_commit_transaction(
     transaction.set(commit_ref, commit)
     transaction.set(
         state_ref,
-        {
-            'current_head_commit_id': commit['commit_id'],
-            'projection_version': 1,
-            'updated_at': commit['commit_time'],
-        },
+        _state_head_write_payload(
+            transaction=transaction,
+            user_ref=user_ref,
+            uid=uid,
+            state=state,
+            commit=commit,
+        ),
     )
     return {'commit': commit, 'applied': True}
 
@@ -334,11 +374,13 @@ def _append_commit_with_builder_transaction(
     transaction.set(commit_ref, commit)
     transaction.set(
         state_ref,
-        {
-            'current_head_commit_id': commit['commit_id'],
-            'projection_version': 1,
-            'updated_at': commit['commit_time'],
-        },
+        _state_head_write_payload(
+            transaction=transaction,
+            user_ref=user_ref,
+            uid=uid,
+            state=state,
+            commit=commit,
+        ),
     )
     return {'commit': commit, 'applied': True}
 

backend/models/memory_state_head.py

@@ -0,0 +1,67 @@
+"""Shared schema for the trusted canonical fields on ``memory_state/head``.
+
+The state-head document also carries legacy-ledger metadata.  Writers must
+therefore preserve these trusted fields rather than replacing the whole
+document with only their own metadata.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Mapping
+
+MEMORY_STATE_HEAD_SCHEMA_VERSION = 1
+MEMORY_STATE_HEAD_SOURCE = "memory_state_head"
+
+
+def trusted_memory_state_head_fields(
+    *,
+    uid: str,
+    account_generation: Any,
+    head_commit_id: Any,
+    commit_sequence: Any,
+) -> dict[str, Any] | None:
+    """Return the trusted fields only when they satisfy the shared contract."""
+    if not uid:
+        return None
+    if isinstance(account_generation, bool) or not isinstance(account_generation, int) or account_generation < 0:
+        return None
+    if not isinstance(head_commit_id, str) or not head_commit_id:
+        return None
+    if isinstance(commit_sequence, bool) or not isinstance(commit_sequence, int) or commit_sequence < 0:
+        return None
+    return {
+        "schema_version": MEMORY_STATE_HEAD_SCHEMA_VERSION,
+        "uid": uid,
+        "source": MEMORY_STATE_HEAD_SOURCE,
+        "account_generation": account_generation,
+        "head_commit_id": head_commit_id,
+        "commit_sequence": commit_sequence,
+    }
+
+
+def trusted_memory_state_head_fields_from_state(state: Mapping[str, Any], *, uid: str) -> dict[str, Any] | None:
+    """Validate and extract trusted fields from a state-head document."""
+    if (
+        state.get("schema_version") != MEMORY_STATE_HEAD_SCHEMA_VERSION
+        or state.get("uid") != uid
+        or state.get("source") != MEMORY_STATE_HEAD_SOURCE
+    ):
+        return None
+    return trusted_memory_state_head_fields(
+        uid=uid,
+        account_generation=state.get("account_generation"),
+        head_commit_id=state.get("head_commit_id"),
+        commit_sequence=state.get("commit_sequence"),
+    )
+
+
+def trusted_memory_state_head_fields_from_control(control: Mapping[str, Any], *, uid: str) -> dict[str, Any] | None:
+    """Build trusted fields from the canonical apply-control record for repair."""
+    if control.get("uid") != uid:
+        return None
+    return trusted_memory_state_head_fields(
+        uid=uid,
+        account_generation=control.get("account_generation"),
+        head_commit_id=control.get("head_commit_id"),
+        commit_sequence=control.get("commit_sequence"),
+    )

backend/tests/unit/test_memory_ledger.py

@@ -1,4 +1,4 @@
 """Memory ledger unit tests.
 
 ``database.memory_ledger`` binds ``db`` from ``database._client`` and imports
 ``transactional`` from ``google.cloud.firestore_v1`` at import time, so a fake
@@ -20,6 +20,8 @@
 
 from database import firestore_transaction_retry
 from testing.import_isolation import load_module_fresh, stub_modules
+from tests.unit.fake_firestore import FakeFirestore
+from utils.memory.v3.account_generation_source import read_memory_v3_trusted_account_generation
 
 _BACKEND = Path(__file__).resolve().parents[2]
 
@@ -77,6 +79,56 @@
     }
 
 
+class _LedgerSnapshot:
+    def __init__(self, data):
+        self._data = data
+        self.exists = data is not None
+
+    def to_dict(self):
+        return self._data
+
+
+class _LedgerDocument:
+    def __init__(self, db, path):
+        self._db = db
+        self.path = path
+
+    def collection(self, name):
+        return _LedgerCollection(self._db, f"{self.path}/{name}")
+
+    def get(self, transaction=None):
+        return _LedgerSnapshot(self._db.docs.get(self.path))
+
+
+class _LedgerCollection:
+    def __init__(self, db, path):
+        self._db = db
+        self.path = path
+
+    def document(self, document_id):
+        return _LedgerDocument(self._db, f"{self.path}/{document_id}")
+
+
+class _LedgerDb:
+    def __init__(self, docs):
+        self.docs = dict(docs)
+
+    def collection(self, name):
+        return _LedgerCollection(self, name)
+
+
+class _LedgerTransaction:
+    def __init__(self):
+        self.sets = []
+
+    def set(self, ref, payload):
+        self.sets.append((ref.path, payload))
+
+
+def _ledger_state_write(transaction):
+    return next(payload for path, payload in transaction.sets if path.endswith("/memory_state/head"))
+
+
 def test_fold_commits_replays_head_and_valid_time():
     january = datetime(2026, 1, 15, tzinfo=timezone.utc)
     february = datetime(2026, 2, 1, tzinfo=timezone.utc)
@@ -262,6 +314,84 @@
         raise AssertionError("Expected same-parent sibling append to fail")
 
 
+def test_ledger_append_repairs_clobbered_trusted_state_head_from_canonical_apply_control():
+    uid = "u1"
+    database = _LedgerDb(
+        {
+            f"users/{uid}/memory_state/head": {
+                "current_head_commit_id": "legacy-head",
+                "projection_version": 1,
+            },
+            f"users/{uid}/memory_state/apply_control": {
+                "uid": uid,
+                "account_generation": 7,
+                "head_commit_id": "canonical-head",
+                "commit_sequence": 11,
+            },
+        }
+    )
+    transaction = _LedgerTransaction()
+
+    result = memory_ledger._append_commit_transaction(
+        transaction,
+        database,
+        uid,
+        "legacy-head",
+        [memory_ledger.add_fact(_fact("m1", "Lives in NYC"))],
+        None,
+        datetime(2026, 7, 14, tzinfo=timezone.utc),
+        None,
+        False,
+    )
+
+    state_head = _ledger_state_write(transaction)
+    assert state_head["current_head_commit_id"] == result["commit"]["commit_id"]
+    assert state_head["head_commit_id"] == "canonical-head"
+    assert state_head["commit_sequence"] == 11
+
+    trusted = read_memory_v3_trusted_account_generation(
+        uid=uid,
+        db_client=FakeFirestore({f"users/{uid}/memory_state/head": state_head}),
+    )
+    assert trusted.read_error_reason is None
+    assert trusted.account_generation == 7
+
+
+def test_ledger_builder_append_preserves_existing_trusted_state_head_without_control_fallback():
+    uid = "u1"
+    database = _LedgerDb(
+        {
+            f"users/{uid}/memory_state/head": {
+                "current_head_commit_id": "legacy-head",
+                "projection_version": 1,
+                "schema_version": 1,
+                "uid": uid,
+                "source": "memory_state_head",
+                "account_generation": 7,
+                "head_commit_id": "canonical-head",
+                "commit_sequence": 11,
+            },
+        }
+    )
+    transaction = _LedgerTransaction()
+
+    result = memory_ledger._append_commit_with_builder_transaction(
+        transaction,
+        database,
+        uid,
+        "legacy-head",
+        lambda _transaction: {"mutations": [memory_ledger.add_fact(_fact("m1", "Lives in NYC"))]},
+        None,
+        datetime(2026, 7, 14, tzinfo=timezone.utc),
+        False,
+    )
+
+    state_head = _ledger_state_write(transaction)
+    assert state_head["current_head_commit_id"] == result["commit"]["commit_id"]
+    assert state_head["head_commit_id"] == "canonical-head"
+    assert state_head["commit_sequence"] == 11
+
+
 def test_typed_transactional_isolates_sdk_wrapper_state_per_concurrent_call(monkeypatch):
     barrier = Barrier(2)
     created_wrappers = []

backend/utils/memory/v3/account_generation_source.py

@@ -10,9 +10,10 @@
 from typing import Any, cast
 
 from database.memory_collections import MemoryCollections
+from models.memory_state_head import MEMORY_STATE_HEAD_SCHEMA_VERSION, MEMORY_STATE_HEAD_SOURCE
 
-V3_TRUSTED_ACCOUNT_GENERATION_SCHEMA_VERSION = 1
-V3_TRUSTED_ACCOUNT_GENERATION_SOURCE = 'memory_state_head'
+V3_TRUSTED_ACCOUNT_GENERATION_SCHEMA_VERSION = MEMORY_STATE_HEAD_SCHEMA_VERSION
+V3_TRUSTED_ACCOUNT_GENERATION_SOURCE = MEMORY_STATE_HEAD_SOURCE
 
 
 class V3AccountGenerationFailureReason(str, Enum):