fix(recall): cap entity fanout in graph expansion (#911)

* fix(recall): cap entity fanout in graph expansion to prevent slow queries

On large banks, the entity co-occurrence self-join in _expand_combined()
produces massive intermediate row counts when seeds reference high-fanout
entities (e.g. an entity with 25K+ mentions). This causes recall latency
to degrade significantly.

Changes:
- Replace unbounded entity self-join with LATERAL per-entity cap
  (graph_per_entity_limit, default 200), reducing intermediate rows
  from potentially millions to at most num_entities * 200
- Add ORDER BY unit_id DESC in LATERAL subquery for deterministic
  recency-biased sampling (rides the PK index, no extra sort)
- Add timeout fallback (graph_expansion_timeout, default 10s) that
  drops entity expansion and falls back to semantic+causal only
- Add composite index (entity_id, unit_id) on unit_entities for
  index-only scans in the LATERAL subquery
- Merge 3 unmerged migration heads into one
- Fix recall_perf.py dotenv override issue

Unlike the approach in #895, this does NOT filter out hub entities
entirely — all entities are kept but capped equally, preserving
retrieval quality for queries about frequently-mentioned entities.

Benchmarked on a 67K-unit bank (top entity = 25K mentions):
- retrieval_graph: 0.337s → 0.055s (84% faster)
- end-to-end recall: 0.912s → 0.519s (43% faster)

* fix(tests): fix broken test_combined_scoring and test_reranking_proof_count

- test_combined_scoring: replace MagicMock(spec=RetrievalResult) with real
  dataclass instances — MagicMock attributes returned nested mocks that
  failed on >= comparisons with int
- test_reranking_proof_count: remove deleted `embedding` param from
  RetrievalResult constructor, use None for occurred_start/end to get
  neutral recency (datetime.now gave recency=1.0 which boosted scores)

* refactor: rename config to link_expansion_ prefix, fix observation fanout

- Rename GRAPH_PER_ENTITY_LIMIT → LINK_EXPANSION_PER_ENTITY_LIMIT and
  GRAPH_EXPANSION_TIMEOUT → LINK_EXPANSION_TIMEOUT to follow the
  convention that these are specific to the link_expansion graph retriever
- Apply the same LATERAL per-entity cap to _expand_observations(), which
  had the same unbounded self-join through unit_entities

* style: fix formatting in config.py

hindsight-api-slim/hindsight_api/alembic/versions/h3i4j5k6l7m8_merge_heads_and_add_unit_entities_index.py

@@ -0,0 +1,42 @@
+"""Merge 3 migration heads and add unit_entities composite index
+
+Revision ID: h3i4j5k6l7m8
+Revises: a4b5c6d7e8f9, c2d3e4f5g6h7, g2h3i4j5k6l7
+Create Date: 2026-04-07
+
+Merges three unmerged migration heads into one, and adds a composite index
+(entity_id, unit_id) on unit_entities for index-only scans in the LATERAL
+entity expansion query.
+"""
+
+from collections.abc import Sequence
+
+from alembic import context, op
+
+revision: str = "h3i4j5k6l7m8"
+down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "c2d3e4f5g6h7", "g2h3i4j5k6l7")
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def _get_schema_prefix() -> str:
+    """Get schema prefix for table names (required for multi-tenant support)."""
+    schema = context.config.get_main_option("target_schema")
+    return f'"{schema}".' if schema else ""
+
+
+def upgrade() -> None:
+    schema = _get_schema_prefix()
+    # Composite index enables index-only scans for entity_id -> unit_id lookups
+    op.execute(
+        f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity_unit ON {schema}unit_entities (entity_id, unit_id)"
+    )
+    # Drop the now-redundant single-column index
+    op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity")
+
+
+def downgrade() -> None:
+    schema = _get_schema_prefix()
+    op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity_unit")
+    # Restore the single-column index
+    op.execute(f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities (entity_id)")

hindsight-api-slim/hindsight_api/config.py

@@ -262,6 +262,8 @@ ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
 ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
 ENV_RECALL_MAX_QUERY_TOKENS = "HINDSIGHT_API_RECALL_MAX_QUERY_TOKENS"
 ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
+ENV_LINK_EXPANSION_PER_ENTITY_LIMIT = "HINDSIGHT_API_LINK_EXPANSION_PER_ENTITY_LIMIT"
+ENV_LINK_EXPANSION_TIMEOUT = "HINDSIGHT_API_LINK_EXPANSION_TIMEOUT"
 
 # OpenTelemetry tracing configuration
 ENV_OTEL_TRACES_ENABLED = "HINDSIGHT_API_OTEL_TRACES_ENABLED"
@@ -475,6 +477,8 @@ DEFAULT_RECALL_MAX_CONCURRENT = 32  # Max concurrent recall operations per worke
 DEFAULT_RECALL_CONNECTION_BUDGET = 4  # Max concurrent DB connections per recall operation
 DEFAULT_RECALL_MAX_QUERY_TOKENS = 500  # Maximum tokens allowed in recall query
 DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY = 8  # Max concurrent mental model refreshes
+DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT = 200  # Max target units per entity in graph expansion
+DEFAULT_LINK_EXPANSION_TIMEOUT = 10.0  # Timeout (seconds) for entity expansion query
 
 # Retain settings
 DEFAULT_RETAIN_MAX_COMPLETION_TOKENS = 64000  # Max tokens for fact extraction LLM call
@@ -780,6 +784,8 @@ class HindsightConfig:
     recall_connection_budget: int
     recall_max_query_tokens: int
     mental_model_refresh_concurrency: int
+    link_expansion_per_entity_limit: int
+    link_expansion_timeout: float
 
     # Retain settings
     retain_max_completion_tokens: int
@@ -1286,6 +1292,10 @@ class HindsightConfig:
             mental_model_refresh_concurrency=int(
                 os.getenv(ENV_MENTAL_MODEL_REFRESH_CONCURRENCY, str(DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY))
             ),
+            link_expansion_per_entity_limit=int(
+                os.getenv(ENV_LINK_EXPANSION_PER_ENTITY_LIMIT, str(DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT))
+            ),
+            link_expansion_timeout=float(os.getenv(ENV_LINK_EXPANSION_TIMEOUT, str(DEFAULT_LINK_EXPANSION_TIMEOUT))),
             # Optimization flags
             skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
             lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",

hindsight-api-slim/hindsight_api/engine/search/link_expansion_retrieval.py

@@ -6,25 +6,30 @@ stored in memory_links:
 
 1. Entity links  — query-time self-join through unit_entities. Score = number of distinct
                    shared entities between the seed set and each candidate, computed via
-                   COUNT(DISTINCT entity_id). More accurate than precomputed entity links.
+                   COUNT(DISTINCT entity_id). Uses a LATERAL per-entity cap
+                   (graph_per_entity_limit, default 200) to prevent high-fanout entities
+                   from exploding the self-join intermediate rows.
 2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most
                     similar existing facts at insert time, similarity >= 0.7). Checked
                     in both directions since the graph is not symmetric. Score = weight.
 3. Causal links  — explicit causal chains (causes/caused_by/enables/prevents).
                    Score = weight + 1.0 (boosted as highest-quality signal).
 
-All three signals are bounded at retain time, so no LATERAL fan-out caps are needed
-at query time. Each expansion is a simple aggregation over a small result set.
+Entity expansion is bounded by graph_per_entity_limit (LATERAL cap per entity).
+A timeout fallback (graph_expansion_timeout) drops entity expansion entirely if the
+query still exceeds the budget.
 
 For non-observation fact types the three expansions are issued as a single CTE query
 (one roundtrip, one connection) with a `source` discriminator column so the Python
 merge step can apply per-signal score transformations.
 """
 
+import asyncio
 import logging
 import math
 import time
 
+from ...config import get_config
 from ..db_utils import acquire_with_retry
 from ..memory_engine import fq_table
 from .graph_retrieval import GraphRetriever
@@ -262,35 +267,48 @@ class LinkExpansionRetriever(GraphRetriever):
                     idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
                     → replaces costly BitmapAnd of two separate scans
         """
+        config = get_config()
         ml = fq_table("memory_links")
         mu = fq_table("memory_units")
         ue = fq_table("unit_entities")
 
+        per_entity_limit = config.link_expansion_per_entity_limit
+
+        # Entity CTE with LATERAL fanout cap.
+        # Every seed entity (including high-frequency ones) is kept, but each
+        # entity's expansion is capped to per_entity_limit target units.  The
+        # LATERAL subquery orders by unit_id DESC so the most recently inserted
+        # units are preferred (a recency proxy that is free — it rides the PK
+        # index with no extra sort).
         entity_cte = f"""
+            seed_entities AS (
+                SELECT DISTINCT ue.entity_id
+                FROM {ue} ue
+                WHERE ue.unit_id = ANY($1::uuid[])
+            ),
             entity_expanded AS (
-                -- Entity co-occurrence via unit_entities self-join.
-                -- Finds units sharing entities with seeds at query time — more accurate
-                -- than precomputed entity links (no stale 50-neighbor cap).
-                -- Score = COUNT(DISTINCT shared entities), mapped to [0,1] via tanh.
                 SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
                        mu.occurred_end, mu.mentioned_at,
                        mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
-                       COUNT(DISTINCT ue_seed.entity_id)::float AS score,
+                       COUNT(DISTINCT se.entity_id)::float AS score,
                        'entity'::text AS source
-                FROM {ue} ue_seed
-                JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id
-                JOIN {mu} mu ON mu.id = ue_target.unit_id
-                WHERE ue_seed.unit_id = ANY($1::uuid[])
-                  AND ue_target.unit_id != ALL($1::uuid[])
-                  AND mu.fact_type = $2
+                FROM seed_entities se
+                CROSS JOIN LATERAL (
+                    SELECT ue_target.unit_id
+                    FROM {ue} ue_target
+                    WHERE ue_target.entity_id = se.entity_id
+                      AND ue_target.unit_id != ALL($1::uuid[])
+                    ORDER BY ue_target.unit_id DESC
+                    LIMIT {per_entity_limit}
+                ) t
+                JOIN {mu} mu ON mu.id = t.unit_id
+                WHERE mu.fact_type = $2
                 GROUP BY mu.id
                 ORDER BY score DESC
                 LIMIT $3
             )"""
 
-        all_rows = await conn.fetch(
-            f"""
-            WITH {entity_cte},
+        semantic_causal_cte = f"""
             semantic_expanded AS (
                 -- Semantic kNN: both outgoing (seeds → their kNN at insert time) and
                 -- incoming (facts inserted after seeds that found seeds as kNN).
@@ -350,18 +368,37 @@ class LinkExpansionRetriever(GraphRetriever):
                   AND mu.fact_type = $2
                 ORDER BY mu.id, ml.weight DESC
                 LIMIT $3
-            )
+            )"""
+
+        full_query = f"""
+            WITH {entity_cte},
+            {semantic_causal_cte}
             SELECT * FROM entity_expanded
             UNION ALL
             SELECT * FROM semantic_expanded
             UNION ALL
             SELECT * FROM causal_expanded
-            """,
-            seed_ids,
-            fact_type,
-            budget,
-            self.causal_weight_threshold,
-        )
+            """
+
+        params = [seed_ids, fact_type, budget, self.causal_weight_threshold]
+
+        try:
+            all_rows = await asyncio.wait_for(
+                conn.fetch(full_query, *params),
+                timeout=config.link_expansion_timeout,
+            )
+        except asyncio.TimeoutError:
+            logger.warning(
+                f"[LinkExpansion] Entity expansion timed out after {config.link_expansion_timeout}s "
+                f"for fact_type={fact_type}, falling back to semantic+causal only"
+            )
+            fallback_query = f"""
+                WITH {semantic_causal_cte}
+                SELECT * FROM semantic_expanded
+                UNION ALL
+                SELECT * FROM causal_expanded
+                """
+            all_rows = await conn.fetch(fallback_query, *params)
 
         entity_rows = [r for r in all_rows if r["source"] == "entity"]
         semantic_rows = [r for r in all_rows if r["source"] == "semantic"]
@@ -401,17 +438,31 @@ class LinkExpansionRetriever(GraphRetriever):
                 f"{len(source_ids_found)} source_memory_ids found"
             )
 
+        config = get_config()
         ue = fq_table("unit_entities")
+        per_entity_limit = config.link_expansion_per_entity_limit
 
         connected_sources_cte = f"""
-            connected_sources AS (
-                -- Find sources sharing entities with seed observation sources
-                -- via unit_entities self-join (query-time, no precomputed links needed).
-                SELECT DISTINCT ue_target.unit_id AS source_id
+            source_entities AS (
+                SELECT DISTINCT ue_seed.entity_id
                 FROM seed_sources ss
                 JOIN {ue} ue_seed ON ue_seed.unit_id = ss.source_id
-                JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id
-                WHERE ue_target.unit_id != ss.source_id
+            ),
+            connected_sources AS (
+                -- Find sources sharing entities with seed observation sources
+                -- via LATERAL-capped self-join (prevents hub entity fanout).
+                SELECT DISTINCT t.unit_id AS source_id
+                FROM source_entities se
+                CROSS JOIN LATERAL (
+                    SELECT ue_target.unit_id
+                    FROM {ue} ue_target
+                    WHERE ue_target.entity_id = se.entity_id
+                    ORDER BY ue_target.unit_id DESC
+                    LIMIT {per_entity_limit}
+                ) t
+                WHERE NOT EXISTS (
+                    SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
+                )
             )"""
 
         entity_rows = await conn.fetch(

hindsight-api-slim/tests/test_combined_scoring.py

@@ -7,7 +7,6 @@ relevance score, independent of the cross-encoder model's score calibration.
 """
 
 from datetime import datetime, timedelta, timezone
-from unittest.mock import MagicMock
 
 import pytest
 
@@ -23,13 +22,18 @@ def _make_result(
     occurred_start: datetime | None = None,
     temporal_proximity: float | None = None,
 ) -> ScoredResult:
-    retrieval = MagicMock(spec=RetrievalResult)
-    retrieval.occurred_start = occurred_start
-    retrieval.temporal_proximity = temporal_proximity
+    retrieval = RetrievalResult(
+        id="test",
+        text="test",
+        fact_type="world",
+        occurred_start=occurred_start,
+        temporal_proximity=temporal_proximity,
+    )
 
-    candidate = MagicMock(spec=MergedCandidate)
-    candidate.retrieval = retrieval
-    candidate.rrf_score = 0.05
+    candidate = MergedCandidate(
+        retrieval=retrieval,
+        rrf_score=0.05,
+    )
 
     return ScoredResult(
         candidate=candidate,

hindsight-api-slim/tests/test_graph_entity_fanout_cap.py

@@ -0,0 +1,211 @@
+"""
+Tests for LATERAL entity fanout cap in graph expansion.
+
+Verifies that the per-entity LIMIT in _expand_combined prevents high-fanout
+entities from exploding the self-join, while still returning entity-based
+graph results.
+"""
+
+import asyncio
+from datetime import datetime, timezone
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_high_fanout_entity_returns_results(memory, request_context):
+    """
+    A high-fanout entity (appearing in many facts) should still produce
+    graph retrieval results — the LATERAL cap limits rows per entity but
+    does not drop the entity entirely.
+    """
+    bank_id = f"test_fanout_cap_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Create many facts sharing one common entity ("Acme Corp") plus
+        # a few with a unique entity so we can query for the unique one
+        # and verify graph expansion finds siblings via "Acme Corp".
+        contents = [
+            # Target: unique entity "Zara" shares "Acme Corp" with the rest
+            {
+                "content": "Zara joined Acme Corp as a senior engineer last month",
+                "context": "hr update",
+                "entities": [{"text": "Zara"}, {"text": "Acme Corp"}],
+            },
+        ]
+        # Add many facts that all share "Acme Corp" — creates a high-fanout entity
+        for i in range(60):
+            contents.append(
+                {
+                    "content": f"Employee {i} completed onboarding at Acme Corp in department {i % 5}",
+                    "context": "hr update",
+                    "entities": [{"text": f"Employee {i}"}, {"text": "Acme Corp"}],
+                }
+            )
+
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=contents,
+            request_context=request_context,
+        )
+
+        from hindsight_api.engine.memory_engine import Budget
+
+        # Query for "Zara" — semantic search finds Zara's fact as a seed,
+        # then graph expansion should find other Acme Corp facts via the
+        # shared entity, even though "Acme Corp" has 60+ mentions.
+        result = await memory.recall_async(
+            bank_id=bank_id,
+            query="Zara",
+            budget=Budget.HIGH,
+            max_tokens=4096,
+            enable_trace=True,
+            request_context=request_context,
+            _quiet=True,
+        )
+
+        assert result.results is not None
+        assert len(result.results) > 0
+
+        # Verify graph retrieval ran and found results
+        retrieval_results = result.trace.get("retrieval_results", [])
+        graph_results = [r for r in retrieval_results if r.get("method_name") == "graph"]
+        assert len(graph_results) > 0, "Graph retrieval should have run"
+
+        # At least one graph result should contain Acme Corp content
+        # (found via shared entity, not just semantic similarity)
+        all_texts = [r.text for r in result.results]
+        acme_found = any("Acme Corp" in t for t in all_texts)
+        assert acme_found, "Should find Acme Corp facts via entity graph expansion"
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_entity_expansion_timeout_fallback(memory, request_context):
+    """
+    When graph_expansion_timeout is set very low, entity expansion should
+    time out gracefully and fall back to semantic+causal links only,
+    rather than failing the entire recall.
+    """
+    bank_id = f"test_timeout_fallback_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=[
+                {
+                    "content": "Alice works on the backend API at TechCorp",
+                    "context": "team info",
+                    "entities": [{"text": "Alice"}, {"text": "TechCorp"}],
+                },
+                {
+                    "content": "Bob maintains the frontend at TechCorp",
+                    "context": "team info",
+                    "entities": [{"text": "Bob"}, {"text": "TechCorp"}],
+                },
+            ],
+            request_context=request_context,
+        )
+
+        from hindsight_api.config import _get_raw_config
+        from hindsight_api.engine.memory_engine import Budget
+
+        config = _get_raw_config()
+        original_timeout = config.link_expansion_timeout
+
+        try:
+            # Set an impossibly low timeout to force the fallback path
+            config.link_expansion_timeout = 0.0001
+
+            result = await memory.recall_async(
+                bank_id=bank_id,
+                query="Alice",
+                budget=Budget.MID,
+                max_tokens=2048,
+                enable_trace=True,
+                request_context=request_context,
+                _quiet=True,
+            )
+
+            # Recall should succeed even when entity expansion times out
+            assert result.results is not None
+            assert len(result.results) > 0
+
+            # Alice should still be found via semantic search
+            result_texts = [r.text for r in result.results]
+            alice_found = any("Alice" in t for t in result_texts)
+            assert alice_found, "Should find Alice via semantic search despite graph timeout"
+        finally:
+            config.link_expansion_timeout = original_timeout
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_per_entity_limit_caps_expansion(memory, request_context):
+    """
+    With graph_per_entity_limit set to a small value, entity expansion should
+    still work but return fewer results from high-fanout entities.
+    """
+    bank_id = f"test_per_entity_limit_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Create facts with a shared entity
+        contents = [
+            {
+                "content": "Lead engineer Dana oversees the Widgets project at MegaCorp",
+                "context": "project info",
+                "entities": [{"text": "Dana"}, {"text": "MegaCorp"}],
+            },
+        ]
+        for i in range(30):
+            contents.append(
+                {
+                    "content": f"MegaCorp hired contractor {i} for the Q4 push",
+                    "context": "hiring info",
+                    "entities": [{"text": f"Contractor {i}"}, {"text": "MegaCorp"}],
+                }
+            )
+
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=contents,
+            request_context=request_context,
+        )
+
+        from hindsight_api.config import _get_raw_config
+        from hindsight_api.engine.memory_engine import Budget
+
+        config = _get_raw_config()
+        original_limit = config.link_expansion_per_entity_limit
+
+        try:
+            # Set a very small per-entity limit
+            config.link_expansion_per_entity_limit = 5
+
+            result = await memory.recall_async(
+                bank_id=bank_id,
+                query="Dana",
+                budget=Budget.HIGH,
+                max_tokens=4096,
+                enable_trace=True,
+                request_context=request_context,
+                _quiet=True,
+            )
+
+            # Recall should succeed with the cap
+            assert result.results is not None
+            assert len(result.results) > 0
+
+            # Graph retrieval should have run
+            retrieval_results = result.trace.get("retrieval_results", [])
+            graph_results = [r for r in retrieval_results if r.get("method_name") == "graph"]
+            assert len(graph_results) > 0, "Graph retrieval should have run"
+        finally:
+            config.link_expansion_per_entity_limit = original_limit
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)

hindsight-api-slim/tests/test_reranking_proof_count.py

@@ -14,24 +14,19 @@ UTC = timezone.utc
 def create_mock_scored_result(proof_count: int | None = None, ce_score: float = 0.8) -> ScoredResult:
     """Helper to create a minimal ScoredResult suitable for scoring tests."""
     retrieval = RetrievalResult(
-        id=uuid4(),
+        id=str(uuid4()),
         text="Test mock fact",
         fact_type="observation" if proof_count is not None else "world",
-        document_id=uuid4(),
-        chunk_id=uuid4(),
-        embedding=[0.1]*384,
-        similarity=0.9,
+        document_id=str(uuid4()),
+        chunk_id=str(uuid4()),
         proof_count=proof_count,
-        # Default neutral dates for testing so only proof_count changes score
-        occurred_start=datetime.now(UTC),
-        occurred_end=datetime.now(UTC)
+        # Use None for neutral recency so only proof_count changes score
+        occurred_start=None,
+        occurred_end=None
     )
     candidate = MergedCandidate(
-        id=retrieval.id,
         retrieval=retrieval,
-        semantic_rank=1,
-        bm25_rank=1,
-        rrf_score=0.1
+        rrf_score=0.1,
     )
     return ScoredResult(
         candidate=candidate,

hindsight-dev/benchmarks/perf/recall_perf.py

@@ -27,6 +27,10 @@ import statistics
 import time
 from typing import Any
 
+# Capture DB URL early before hindsight_api imports trigger dotenv override
+# (config.py uses load_dotenv(override=True) which stomps env vars)
+_EARLY_DB_URL = os.environ.get("HINDSIGHT_API_DATABASE_URL")
+
 from rich.console import Console
 from rich.progress import BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.table import Table
@@ -546,7 +550,7 @@ def _build_engine(*, disable_observations: bool = False) -> "Any":
     """Create a MemoryEngine using mock LLM and DB from env."""
     from hindsight_api import MemoryEngine
 
-    db_url = os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0")
+    db_url = _EARLY_DB_URL or os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0")
     if disable_observations:
         os.environ["HINDSIGHT_API_ENABLE_OBSERVATIONS"] = "false"
     engine = MemoryEngine(