perf: narrow smart-merge and KG Firestore reads

backend/config/firestore_quick_reads.py

@@ -0,0 +1,22 @@
+# LIFECYCLE: permanent
+"""Default-on structural optimizations; independent emergency kill switches."""
+
+import os
+
+
+def smart_merge_partition_enabled() -> bool:
+    return os.getenv('FSR_SMART_MERGE_PARTITION', 'on').strip().lower() != 'off'
+
+
+def kg_citation_index_enabled() -> bool:
+    return os.getenv('FSR_KG_CITATION_INDEX', 'on').strip().lower() != 'off'
+
+
+def indexed_string_value_safe(value: object) -> bool:
+    """Only query identifiers fully represented by Firestore's UTF-8 index value."""
+    if not isinstance(value, str):
+        return False
+    try:
+        return len(value.encode('utf-8')) <= 1500
+    except UnicodeEncodeError:
+        return False

backend/database/firestore_index_registry.py

@@ -1,4 +1,4 @@
 """Repository-owned Firestore query and index requirements.
 The generated Firebase manifest follows query specs that own their builders and indexes.
 Index-only requirements stay explicit until callers migrate to query specs.
 """
@@ -1935,6 +1935,24 @@
     index_fields=(_asc('discarded'), _asc('source'), _asc('status'), _desc('created_at'), _desc('__name__')),
 )
 
+CONVERSATIONS_SMART_MERGE_PARTITION_QUERY = FirestoreQuerySpec(
+    identifier='conversations_smart_merge_device_preceding',
+    collection_group='conversations',
+    query_scope='COLLECTION',
+    filters=(
+        *CONVERSATIONS_SMART_MERGE_PRECEDING_QUERY.filters,
+        FirestoreQueryFilter('client_device_id', '==', 'client_device_id'),
+    ),
+    index_fields=(
+        _asc('discarded'),
+        _asc('source'),
+        _asc('status'),
+        _asc('client_device_id'),
+        _desc('created_at'),
+        _desc('__name__'),
+    ),
+)
+
 
 CONVERSATION_PHOTOS_NAME_RANGE_QUERY = FirestoreQuerySpec(
     identifier='conversation_photos_name_range_export',
@@ -2028,6 +2046,7 @@
     DAY3_REENGAGEMENT_RETURNED_CONVERSATIONS_QUERY,
     DAILY_SUMMARY_RECIPIENTS_QUERY,
     CONVERSATIONS_SMART_MERGE_PRECEDING_QUERY,
+    CONVERSATIONS_SMART_MERGE_PARTITION_QUERY,
 )
 
 _INDEX_ONLY_REQUIREMENT_SIGNATURES = frozenset(requirement.signature for requirement in INDEX_ONLY_REQUIREMENTS)

backend/database/knowledge_graph.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timezone
 from typing import Any, Dict, Iterable, List, Optional, Tuple, TypedDict, cast
 import uuid
+import logging
 
 from google.cloud.firestore_v1 import FieldFilter
 
@@ -10,6 +11,12 @@
 
 from ._client import db
 from .read_boundary import parse_snapshot_or_none
+from config.firestore_quick_reads import kg_citation_index_enabled, indexed_string_value_safe
+from database.firestore_read_origin import read_origin
+from google.api_core.exceptions import FailedPrecondition
+from utils.observability.fallback import record_fallback
+
+logger = logging.getLogger(__name__)
 
 users_collection = 'users'
 knowledge_nodes_collection = 'knowledge_nodes'
@@ -861,6 +868,38 @@
     _batch_delete(edges_ref)
 
 
+def _citation_nodes(nodes_ref: Any, memory_ids: List[str]) -> List[Any]:
+    """Use Firestore's writer-maintained reverse index, including historical rows.
+
+    Native array indexes cover every successful write and are backfilled by
+    Firestore before READY. Never infer coverage from an empty side collection.
+    Materialize all queries before mutations, so fallback cannot repeat writes.
+    """
+    if not kg_citation_index_enabled():
+        return list(nodes_ref.stream())
+    ids = set(memory_ids)
+    # A node can match every chunk of a large retraction. Keep the original
+    # single scan for bulk work rather than multiply billed reads for it.
+    if len(ids) > 30 or not all(indexed_string_value_safe(value) for value in ids):
+        return list(nodes_ref.stream())
+    try:
+        query = nodes_ref.where(filter=FieldFilter('memory_ids', 'array_contains_any', sorted(ids)))
+        return list(query.stream())
+    except FailedPrecondition as error:
+        if 'index' not in str(error).lower():
+            raise
+        record_fallback(
+            component='other',
+            from_mode='indexed_query',
+            to_mode='legacy_query',
+            reason='unavailable',
+            outcome='recovered',
+            log=logger,
+        )
+        return list(nodes_ref.stream())
+
+
+@read_origin('task:kg_citation_prune')
 def prune_memory_citations_from_kg(uid: str, memory_ids: List[str], *, db_client: Any = None) -> int:
     """Remove memory_ids from KG nodes/edges; delete entities with no remaining citations."""
     if not memory_ids:
@@ -872,7 +911,7 @@
     edges_ref = user_ref.collection(knowledge_edges_collection)
     pruned = 0
 
-    for doc in nodes_ref.stream():
+    for doc in _citation_nodes(nodes_ref, memory_ids):
         node_doc: KnowledgeNodeDoc = cast(KnowledgeNodeDoc, _typed_doc(doc))
         existing_ids = set(node_doc.get("memory_ids") or [])
         if not existing_ids.intersection(retracted):

backend/database/smart_merge.py

@@ -23,18 +23,24 @@
 
 from google.cloud import firestore
 from google.cloud.firestore_v1 import FieldFilter
+from google.api_core.exceptions import FailedPrecondition
 
 from database import conversations as conversations_db
 from database import smart_merge_audit as audit_db
 from database.quiet_check_retirement import stop_quiet_check_for_conversation
 from database._client import get_firestore_client, run_transactional
-from database.firestore_index_registry import CONVERSATIONS_SMART_MERGE_PRECEDING_QUERY
+from database.firestore_index_registry import (
+    CONVERSATIONS_SMART_MERGE_PRECEDING_QUERY,
+    CONVERSATIONS_SMART_MERGE_PARTITION_QUERY,
+)
 from database.people_stats_cache import invalidate_people_stats_cache
 from config.sync_live_dedupe import sync_live_dedupe_active_for
 from utils.conversations.smart_merge_policy import rebase_donor_segments
 from utils.sync.bridge_observation import log_bridge_observation
 from config import merge_ancestry
 from config.conversation_smart_merge import smart_merge_flatten_enabled
+from config.firestore_quick_reads import smart_merge_partition_enabled, indexed_string_value_safe
+from database.firestore_read_origin import read_origin
 from utils.conversations.size_limit import (
     CONVERSATION_ABSORPTION_BYTE_CAP,
     projected_document_bytes,
@@ -119,6 +125,7 @@
     discarded: bool = False,
     transaction: Any = None,
     firestore_client: Any = None,
+    client_device_id: Optional[str] = None,
 ) -> list[dict[str, Any]]:
     """Newest-first rows of one source and discard flag created before ``created_before``.
 
@@ -142,13 +149,59 @@
     query = query.order_by('created_at', direction=firestore.Query.DESCENDING).select(list(_PRECEDING_FIELDS))
     if limit is not None:
         query = query.limit(limit)
-    stream = query.stream(transaction=transaction) if transaction is not None else query.stream()
-    rows = []
-    for snapshot in stream:
-        data = snapshot.to_dict() or {}
-        data['id'] = snapshot.id
-        rows.append(data)
-    return rows
+
+    def collect(selected):
+        stream = selected.stream(transaction=transaction) if transaction is not None else selected.stream()
+        rows = []
+        for snapshot in stream:
+            data = snapshot.to_dict() or {}
+            data['id'] = snapshot.id
+            rows.append(data)
+        return rows
+
+    # A filter before a finite limit changes the global six-row window. Keep
+    # that lookup unchanged; the expensive discarded-barrier scan is unbounded.
+    # Missing device fields are not matched by equality-to-null in Firestore.
+    if (
+        limit is None
+        and isinstance(client_device_id, str)
+        and client_device_id
+        and smart_merge_partition_enabled()
+        and indexed_string_value_safe(client_device_id)
+    ):
+        with read_origin('task:smart_merge_partition_lookup'):
+            try:
+                partition_query = (
+                    CONVERSATIONS_SMART_MERGE_PARTITION_QUERY.build(
+                        _collection(client, uid),
+                        {
+                            'discarded': discarded,
+                            'source': source,
+                            'statuses': list(_ALL_STATUSES),
+                            'created_before': created_before,
+                            'client_device_id': client_device_id,
+                        },
+                        field_filter_factory=FieldFilter,
+                    )
+                    .order_by('created_at', direction=firestore.Query.DESCENDING)
+                    .select(list(_PRECEDING_FIELDS))
+                )
+                return collect(partition_query)
+            except FailedPrecondition as error:
+                if 'index' not in str(error).lower():
+                    raise
+                # The lead reconciles the index before shipping, but a rolling
+                # deploy must remain correct while the index is still building.
+                record_fallback(
+                    component='other',
+                    from_mode='indexed_query',
+                    to_mode='legacy_query',
+                    reason='unavailable',
+                    outcome='recovered',
+                    log=logger,
+                )
+                return collect(query)
+    return collect(query)
 
 
 def _epoch(value: Any) -> Optional[float]:
@@ -219,6 +272,7 @@
         discarded=True,
         transaction=transaction,
         firestore_client=firestore_client,
+        client_device_id=donor.get('client_device_id'),
     ):
         if row.get('client_device_id') != donor.get('client_device_id'):
             continue
@@ -259,7 +313,7 @@
     return run_transactional(client, write)
 
 
 def absorb_conversation(
     uid: str,
     survivor_id: str,
     donor_id: str,

backend/docs/feature-flag-registry.md

@@ -433,7 +433,9 @@ and an explicit empty literal renders as `''`.
 | `EPISODES_ENABLED` | Default on unless set to a non-on value; live episodes overlay (per-moment gpt-5-nano update, provisional items, /v1/episodes capability). Off reports enabled=false and apps keep today's per-conversation UI | backend | env | open | — | — | — | — | keep | keep | — | dazheng |
 | `FAIR_USE_ENABLED` | Enforce fair-use metering; set false to stop enforcement | backend | env | closed | — | true (backend-listen (chart)) | true (backend-listen (chart)) | — | keep | keep | — | unowned |
 | `FAIR_USE_KILL_SWITCH` | Emergency stop for fair-use enforcement | backend | env | inverted | — | false (backend-listen (chart)) | false (backend-listen (chart)) | — | keep | keep | — | unowned |
+| `FSR_KG_CITATION_INDEX` | Default on; query native reverse citations for KG node pruning; off restores the node scan | backend | env | open | — | — | — | — | keep | keep | — | dazheng |
 | `FSR_MAINTENANCE_COORDINATION_ENABLED` | Default on in prod independently of read-lever rollback; explicit false restores per-replica maintenance | backend | env | open | — | — | — | — | keep | keep | — | dazheng |
+| `FSR_SMART_MERGE_PARTITION` | Default on; partition the unbounded smart-merge discarded-barrier scan by explicit device; off restores the original scan | backend | env | open | — | — | — | — | keep | keep | — | dazheng |
 | `JOBSTORE_SHADOW_WRITES` | Default on unless 0/false/off/no; off stops the A2 Postgres shadow mirror (and its parity reads) while BATCH_JOBSTORE_MODE=shadow | backend | env | open | — | — | — | — | keep | keep | — | dazheng |
 | `LISTEN_FINALIZATION_BYOK_ABANDONMENT_ENABLED` | Abandon finalization after BYOK failure | backend | env | open | true | true (backend-listen (chart), gke/backend-listen) | true (backend-listen (chart), gke/backend-listen) | — | keep | keep | — | unowned |
 | `LISTEN_FINALIZATION_DURABLE_ATTEMPT_CAP_ENABLED` | Cap durable listen worker attempts while retaining transcript-only conversations | backend | env | open | — | — | true | — | keep | keep | — | dazheng |
@@ -843,7 +845,9 @@ their code default (`fail` tells you which way a missing value resolves).
 - `FIRESTORE_CACHE_ENABLED` — Enable Firestore response cache (fail: closed)
 - `FIRST_SUMMARY_DEADLINE_SECONDS` — First required summary dispatch budget; default 120 seconds, bounded 60 to 120 (fail: closed)
 - `FIRST_SUMMARY_UID_ALLOWLIST` — Optional summary canary owner allowlist; empty admits only the positive percentage cohort when global gate is on (fail: closed)
+- `FSR_KG_CITATION_INDEX` — Default on; query native reverse citations for KG node pruning; off restores the node scan (fail: open)
 - `FSR_MAINTENANCE_COORDINATION_ENABLED` — Default on in prod independently of read-lever rollback; explicit false restores per-replica maintenance (fail: open)
+- `FSR_SMART_MERGE_PARTITION` — Default on; partition the unbounded smart-merge discarded-barrier scan by explicit device; off restores the original scan (fail: open)
 - `HELD_CONTENT_COHORT` — Fallback tokens uid:<uid>,pct:<0-100> or explicit all, used only after rollout_cohorts/held-content is confirmed absent; empty, unset, or malformed admits nobody (fail: closed)
 - `HELD_CONTENT_ENABLED` — Default-off at-cap choice API; additionally requires X-Omi-Held-Content: 1. Capture custody and recovery must ship before enablement. (fail: closed)
 - `JOBSTORE_SHADOW_WRITES` — Default on unless 0/false/off/no; off stops the A2 Postgres shadow mirror (and its parity reads) while BATCH_JOBSTORE_MODE=shadow (fail: open)

backend/docs/runbooks/firestore-quick-reads.md

@@ -0,0 +1,68 @@
+# Firestore structural query reads
+
+PR #23442 now retains two default-on structural optimizations. The public catalog
+and request-cohort cache changes were dropped during the 2026-10-27 main merge:
+main already has public-list single-flight/Redis reuse and deadline-bounded shared
+cohort observations. This PR adds no request or cross-request cache lifetime.
+
+| Kill switch (`off` restores the prior path) | Scope | Origin |
+| --- | --- | --- |
+| `FSR_SMART_MERGE_PARTITION` | Unbounded discarded-barrier scan for explicit devices | `task:smart_merge_partition_lookup` |
+| `FSR_KG_CITATION_INDEX` | Reverse-citation node lookup during KG pruning | `task:kg_citation_prune` |
+
+## Index and policy gates
+
+The lead must reconcile the generated manifest and verify
+`conversations_smart_merge_device_preceding` READY before serving this source.
+Use the existing schema approval/readiness path; this lane performs no applies,
+access changes, production writes or deployment. A missing/building-index
+`FailedPrecondition` falls back to the old query before returning rows; other
+errors propagate. Both attempts retain the caller's transaction.
+
+The finite predecessor lookup keeps its global six-row window. Missing or empty
+device IDs retain the old scan so absent fields remain visible. Discarded barriers
+stay unbounded within the same device, including deleted barriers. Only values
+fully represented by the index are eligible: device and citation identifiers
+over 1,500 UTF-8 bytes retain the original scan. Firestore documents
+[index-value truncation and inconsistent query results](https://docs.cloud.google.com/firestore/quotas)
+above that limit; this is optimization eligibility, not a new input restriction.
+Main's donor
+rebasing, source/lineage checks, open-loop extraction and retirement paths remain.
+
+Firestore's native `memory_ids` array index is the reverse-citation writer: it
+updates with existing node writes and covers historical data before serving.
+There is no separate side collection or backfill. One query for at most 30 distinct memory IDs is materialized before mutations.
+Larger retractions keep the original single scan, so overlapping citation chunks
+cannot amplify reads. A missing index falls back before any write. Global surviving-node and orphan-edge scans remain: the
+operation must still remove unrelated pre-existing dangling edges. No canonical
+promotion, deletion/hold authority, transcription, speaker or Free-plan change.
+
+## Measurement and economics
+
+The 10-26 fleet caller pools are **16.749388M reads/day** for
+`smart_merge:find_preceding_conversations` and **8.915616M/day** for
+`knowledge_graph:prune_memory_citations_from_kg`. At $0.43–0.47/M and 30 days:
+
+- Smart merge's entire caller pool is $216–236/month. The removable fraction is
+  only explicit-device unbounded barriers multiplied by the off-device rows they
+  currently download. The saved aggregate does not identify either fraction.
+  **$0–236/month gross ceiling; no defensible point forecast.** The synthetic
+  121-to-1 lookup is mechanism evidence, not the fleet traffic mix.
+- KG removes only the first of two full node scans; the edge scan remains. Its
+  full caller pool is $115–126/month, but this change's structural ceiling is
+  **$0–63/month**, before empty-query floors and the unknown fraction of bulk retractions. The 206-to-106
+  synthetic fixture corresponds to $56–61/month only if that graph shape holds
+  fleet-wide; it is not a measured forecast.
+
+The retained bundle's **$0–299/month** is an upper bound, before added composite
+index storage and fallback overhead. Existing writes maintain the native indexes;
+no new application document writes are introduced. New verified/realized savings
+are **$0**. Earlier 22.5M/day predictions are retired, not booked.
+
+Compare old and new callers together: smart merge's inner collector is attributed
+to `smart_merge:collect` (including the unchanged finite query), KG's new lookup to
+`knowledge_graph:_citation_nodes`, and remaining sweeps to the old prune caller.
+Track both new origins, legacy callers, fallback counts, query/RPC latency,
+merge/barrier correctness and graph-cleanup equivalence. After lead serving,
+compare matched windows against the 10-24/25 billing baseline and confirm net
+reads/writes/storage at D+1 in the billing export. A renamed caller is not saving.

backend/tests/support/firestore_query_driver_registry.py

@@ -1,4 +1,4 @@
 """Explicit driver registry for every statically discovered ``database.*`` querying function.
 
 Each entry is either a ``DriverEntry`` (driven across the full product of its
 shape-affecting parameter domains), a ``CoveredByEntry`` (query shapes are
@@ -1561,6 +1561,13 @@
         'database.knowledge_graph.prune_memory_citations_from_kg', base={'uid': UID, 'memory_ids': ['mem-1', 'mem-2']}
     )
 )
+_add(
+    CoveredByEntry(
+        'database.knowledge_graph._citation_nodes',
+        covered_by=('database.knowledge_graph.prune_memory_citations_from_kg',),
+        reason='The prune driver exercises the native reverse-citation query before global orphan cleanup.',
+    )
+)
 _add(
     DriverEntry(
         'database.knowledge_graph.upsert_knowledge_node', base={'uid': UID, 'node_data': {'label': 'shape-node'}}
@@ -2245,7 +2252,11 @@
     DriverEntry(
         'database.smart_merge.find_preceding_conversations',
         base={'uid': UID, 'source': 'omi', 'created_before': T0},
-        domains={'discarded': [False, True], 'transaction': [None, ref_transaction()]},
+        domains={
+            'discarded': [False, True],
+            'transaction': [None, ref_transaction()],
+            'client_device_id': [None, 'pendant-1'],
+        },
         neutrals={'limit': _LIMIT_OPT},
     )
 )

backend/tests/unit/test_conversation_smart_merge.py

@@ -1,4 +1,4 @@
 """Smart merge end to end over the real transaction code and a strict in-memory Firestore.
 
 Jev, conversation processing, retraction/audio seams and the search index are
 faked at their module seams; the absorb/claim/complete transactions, the
@@ -131,7 +131,16 @@
         return row
 
     def preceding(
-        self, uid, *, source, created_before, limit=None, discarded=False, transaction=None, firestore_client=None
+        self,
+        uid,
+        *,
+        source,
+        created_before,
+        limit=None,
+        discarded=False,
+        transaction=None,
+        firestore_client=None,
+        client_device_id=None,
     ):
         if transaction is not None:
             assert_read_allowed = getattr(transaction, '_assert_read_allowed', None)
@@ -148,6 +157,7 @@
             and value.get('source') == source
             and value.get('status') in smart_merge_db._ALL_STATUSES
             and value['created_at'] < created_before
+            and (not client_device_id or limit is not None or value.get('client_device_id') == client_device_id)
         ]
         rows.sort(key=lambda row: row['created_at'], reverse=True)
         projected = [{k: v for k, v in row.items() if k != 'transcript_segments'} for row in rows]