fix: add missing list_tenants to test mocks and update CLI for async refresh

hindsight-api/hindsight_api/alembic/versions/s4n5o6p7q8r9_add_consolidated_at_to_memory_units.py

@@ -0,0 +1,53 @@
+"""Add consolidated_at column to memory_units for incremental consolidation tracking.
+
+This allows consolidation to track progress at the memory level rather than
+using a bank-level watermark. If consolidation crashes, already-processed
+memories won't be reprocessed.
+
+Revision ID: s4n5o6p7q8r9
+Revises: r3m4n5o6p7q8
+Create Date: 2025-01-22
+"""
+
+from collections.abc import Sequence
+
+from alembic import context, op
+
+revision: str = "s4n5o6p7q8r9"
+down_revision: str | Sequence[str] | None = "r3m4n5o6p7q8"
+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()
+
+    # Add consolidated_at column to memory_units
+    op.execute(
+        f"""
+        ALTER TABLE {schema}memory_units
+        ADD COLUMN IF NOT EXISTS consolidated_at TIMESTAMPTZ DEFAULT NULL
+        """
+    )
+
+    # Create index for efficient querying of unconsolidated memories
+    op.execute(
+        f"""
+        CREATE INDEX IF NOT EXISTS idx_memory_units_unconsolidated
+        ON {schema}memory_units (bank_id, created_at)
+        WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')
+        """
+    )
+
+
+def downgrade() -> None:
+    schema = _get_schema_prefix()
+
+    op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_unconsolidated")
+    op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS consolidated_at")

hindsight-api/hindsight_api/api/http.py

@@ -1191,11 +1191,8 @@ class OperationResponse(BaseModel):
 class ConsolidationResponse(BaseModel):
     """Response model for consolidation trigger endpoint."""
 
-    status: str = Field(description="Status of the consolidation (completed or queued)")
-    processed: int = Field(description="Number of memories processed")
-    created: int = Field(description="Number of mental models created")
-    updated: int = Field(description="Number of mental models updated")
-    message: str = Field(description="Human-readable summary")
+    operation_id: str = Field(description="ID of the async consolidation operation")
+    deduplicated: bool = Field(default=False, description="True if an existing pending task was reused")
 
 
 class OperationsListResponse(BaseModel):
@@ -2058,42 +2055,19 @@ def _register_routes(app: FastAPI):
                 )
                 total_documents = doc_count_result["count"] if doc_count_result else 0
 
-                # Get consolidation stats
-                bank_row = await conn.fetchrow(
+                # Get consolidation stats from memory-level tracking
+                consolidation_stats = await conn.fetchrow(
                     f"""
-                    SELECT last_consolidated_at
-                    FROM {fq_table("banks")}
+                    SELECT
+                        MAX(consolidated_at) as last_consolidated_at,
+                        COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) as pending
+                    FROM {fq_table("memory_units")}
                     WHERE bank_id = $1
                     """,
                     bank_id,
                 )
-                last_consolidated_at = bank_row["last_consolidated_at"] if bank_row else None
-
-                # Count memories pending consolidation (created after last_consolidated_at)
-                if last_consolidated_at:
-                    pending_consolidation_result = await conn.fetchrow(
-                        f"""
-                        SELECT COUNT(*) as count
-                        FROM {fq_table("memory_units")}
-                        WHERE bank_id = $1
-                          AND created_at > $2
-                          AND fact_type IN ('experience', 'world')
-                        """,
-                        bank_id,
-                        last_consolidated_at,
-                    )
-                else:
-                    # If never consolidated, count all experience/world memories
-                    pending_consolidation_result = await conn.fetchrow(
-                        f"""
-                        SELECT COUNT(*) as count
-                        FROM {fq_table("memory_units")}
-                        WHERE bank_id = $1
-                          AND fact_type IN ('experience', 'world')
-                        """,
-                        bank_id,
-                    )
-                pending_consolidation = pending_consolidation_result["count"] if pending_consolidation_result else 0
+                last_consolidated_at = consolidation_stats["last_consolidated_at"] if consolidation_stats else None
+                pending_consolidation = consolidation_stats["pending"] if consolidation_stats else 0
 
                 # Count total mental models
                 mental_model_count_result = await conn.fetchrow(
@@ -2354,9 +2328,9 @@ def _register_routes(app: FastAPI):
 
     @app.post(
         "/v1/default/banks/{bank_id}/reflections/{reflection_id}/refresh",
-        response_model=ReflectionResponse,
+        response_model=AsyncOperationSubmitResponse,
         summary="Refresh reflection",
-        description="Re-run the source query through reflect and update the content.",
+        description="Submit an async task to re-run the source query through reflect and update the content.",
         operation_id="refresh_reflection",
         tags=["Reflections"],
     )
@@ -2365,16 +2339,16 @@ def _register_routes(app: FastAPI):
         reflection_id: str,
         request_context: RequestContext = Depends(get_request_context),
     ):
-        """Refresh a reflection by re-running its source query."""
+        """Refresh a reflection by re-running its source query (async)."""
         try:
-            reflection = await app.state.memory.refresh_reflection(
+            result = await app.state.memory.submit_async_refresh_reflection(
                 bank_id=bank_id,
                 reflection_id=reflection_id,
                 request_context=request_context,
             )
-            if reflection is None:
-                raise HTTPException(status_code=404, detail=f"Reflection '{reflection_id}' not found")
-            return ReflectionResponse(**reflection)
+            return AsyncOperationSubmitResponse(operation_id=result["operation_id"], status="queued")
+        except ValueError as e:
+            raise HTTPException(status_code=404, detail=str(e))
         except (AuthenticationError, HTTPException):
             raise
         except Exception as e:
@@ -3206,18 +3180,12 @@ def _register_routes(app: FastAPI):
         tags=["Banks"],
     )
     async def api_trigger_consolidation(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
-        """Trigger consolidation for a bank."""
+        """Trigger consolidation for a bank (async)."""
         try:
-            result = await app.state.memory.run_consolidation(bank_id, request_context=request_context)
-            processed = result.get("processed", 0)
-            created = result.get("created", 0)
-            updated = result.get("updated", 0)
+            result = await app.state.memory.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
             return ConsolidationResponse(
-                status="completed",
-                processed=processed,
-                created=created,
-                updated=updated,
-                message=f"Consolidation completed: {processed} memories processed, {created} mental models created, {updated} updated",
+                operation_id=result["operation_id"],
+                deduplicated=result.get("deduplicated", False),
             )
         except (AuthenticationError, HTTPException):
             raise

hindsight-api/hindsight_api/config.py

@@ -39,6 +39,11 @@ ENV_REFLECT_LLM_API_KEY = "HINDSIGHT_API_REFLECT_LLM_API_KEY"
 ENV_REFLECT_LLM_MODEL = "HINDSIGHT_API_REFLECT_LLM_MODEL"
 ENV_REFLECT_LLM_BASE_URL = "HINDSIGHT_API_REFLECT_LLM_BASE_URL"
 
+ENV_CONSOLIDATION_LLM_PROVIDER = "HINDSIGHT_API_CONSOLIDATION_LLM_PROVIDER"
+ENV_CONSOLIDATION_LLM_API_KEY = "HINDSIGHT_API_CONSOLIDATION_LLM_API_KEY"
+ENV_CONSOLIDATION_LLM_MODEL = "HINDSIGHT_API_CONSOLIDATION_LLM_MODEL"
+ENV_CONSOLIDATION_LLM_BASE_URL = "HINDSIGHT_API_CONSOLIDATION_LLM_BASE_URL"
+
 ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
 ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
 ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
@@ -293,6 +298,11 @@ class HindsightConfig:
     reflect_llm_model: str | None
     reflect_llm_base_url: str | None
 
+    consolidation_llm_provider: str | None
+    consolidation_llm_api_key: str | None
+    consolidation_llm_model: str | None
+    consolidation_llm_base_url: str | None
+
     # Embeddings
     embeddings_provider: str
     embeddings_local_model: str
@@ -385,6 +395,10 @@ class HindsightConfig:
             reflect_llm_api_key=os.getenv(ENV_REFLECT_LLM_API_KEY) or None,
             reflect_llm_model=os.getenv(ENV_REFLECT_LLM_MODEL) or None,
             reflect_llm_base_url=os.getenv(ENV_REFLECT_LLM_BASE_URL) or None,
+            consolidation_llm_provider=os.getenv(ENV_CONSOLIDATION_LLM_PROVIDER) or None,
+            consolidation_llm_api_key=os.getenv(ENV_CONSOLIDATION_LLM_API_KEY) or None,
+            consolidation_llm_model=os.getenv(ENV_CONSOLIDATION_LLM_MODEL) or None,
+            consolidation_llm_base_url=os.getenv(ENV_CONSOLIDATION_LLM_BASE_URL) or None,
             # Embeddings
             embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER),
             embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL),
@@ -531,6 +545,10 @@ class HindsightConfig:
             reflect_provider = self.reflect_llm_provider or self.llm_provider
             reflect_model = self.reflect_llm_model or self.llm_model
             logger.info(f"LLM (reflect): provider={reflect_provider}, model={reflect_model}")
+        if self.consolidation_llm_provider or self.consolidation_llm_model:
+            consolidation_provider = self.consolidation_llm_provider or self.llm_provider
+            consolidation_model = self.consolidation_llm_model or self.llm_model
+            logger.info(f"LLM (consolidation): provider={consolidation_provider}, model={consolidation_model}")
         logger.info(f"Embeddings: provider={self.embeddings_provider}")
         logger.info(f"Reranker: provider={self.reranker_provider}")
         logger.info(f"Graph retriever: {self.graph_retriever}")

hindsight-api/hindsight_api/engine/consolidation/consolidator.py

@@ -93,12 +93,14 @@ async def run_consolidation_job(
         logger.debug(f"Consolidation disabled for bank {bank_id}")
         return {"status": "disabled", "bank_id": bank_id}
 
-    async with memory_engine._pool.acquire() as conn:
-        # Get bank profile and last_consolidated_at
+    pool = memory_engine._pool
+
+    # Get bank profile
+    async with pool.acquire() as conn:
         t0 = time.time()
         bank_row = await conn.fetchrow(
             f"""
-            SELECT bank_id, name, mission, last_consolidated_at
+            SELECT bank_id, name, mission
             FROM {fq_table("banks")}
             WHERE bank_id = $1
             """,
@@ -110,31 +112,51 @@ async def run_consolidation_job(
             return {"status": "bank_not_found", "bank_id": bank_id}
 
         mission = bank_row["mission"] or "General memory consolidation"
-        last_consolidated_at = bank_row["last_consolidated_at"]
         perf.record_timing("fetch_bank", time.time() - t0)
 
-        # Fetch memories created after last_consolidated_at (exclude mental_model type)
-        t0 = time.time()
-        if last_consolidated_at:
+        # Count total unconsolidated memories for progress logging
+        total_count = await conn.fetchval(
+            f"""
+            SELECT COUNT(*)
+            FROM {fq_table("memory_units")}
+            WHERE bank_id = $1
+              AND consolidated_at IS NULL
+              AND fact_type IN ('experience', 'world')
+            """,
+            bank_id,
+        )
+
+    if total_count == 0:
+        logger.debug(f"No new memories to consolidate for bank {bank_id}")
+        return {"status": "no_new_memories", "bank_id": bank_id, "memories_processed": 0}
+
+    logger.info(f"[CONSOLIDATION] bank={bank_id} total_unconsolidated={total_count}")
+    perf.log(f"[1] Found {total_count} pending memories to consolidate")
+
+    # Process each memory with individual commits for crash recovery
+    stats = {
+        "memories_processed": 0,
+        "mental_models_created": 0,
+        "mental_models_updated": 0,
+        "mental_models_merged": 0,
+        "actions_executed": 0,
+        "skipped": 0,
+    }
+
+    batch_num = 0
+    while True:
+        batch_num += 1
+        batch_start = time.time()
+
+        # Fetch next batch of unconsolidated memories
+        async with pool.acquire() as conn:
+            t0 = time.time()
             memories = await conn.fetch(
                 f"""
-                SELECT id, text, fact_type, occurred_start, event_date, tags
-                FROM {fq_table("memory_units")}
-                WHERE bank_id = $1 AND created_at > $2
-                  AND fact_type IN ('experience', 'world')
-                ORDER BY created_at ASC
-                LIMIT $3
-                """,
-                bank_id,
-                last_consolidated_at,
-                max_memories_per_batch,
-            )
-        else:
-            memories = await conn.fetch(
-                f"""
-                SELECT id, text, fact_type, occurred_start, event_date, tags
+                SELECT id, text, fact_type, occurred_start, event_date, tags, mentioned_at
                 FROM {fq_table("memory_units")}
                 WHERE bank_id = $1
+                  AND consolidated_at IS NULL
                   AND fact_type IN ('experience', 'world')
                 ORDER BY created_at ASC
                 LIMIT $2
@@ -142,45 +164,16 @@ async def run_consolidation_job(
                 bank_id,
                 max_memories_per_batch,
             )
-        perf.record_timing("fetch_memories", time.time() - t0)
+            perf.record_timing("fetch_memories", time.time() - t0)
 
         if not memories:
-            logger.debug(f"No new memories to consolidate for bank {bank_id}")
-            # Update timestamp anyway to prevent reprocessing
-            await _update_last_consolidated_at(conn, bank_id)
-            return {"status": "no_new_memories", "bank_id": bank_id, "memories_processed": 0}
+            break  # No more unconsolidated memories
 
-        logger.info(
-            f"[CONSOLIDATION] bank={bank_id} memories={len(memories)} "
-            f"batch_size={max_memories_per_batch} since={last_consolidated_at or 'beginning'}"
-        )
-        perf.log(f"[1] Found {len(memories)} pending memories to consolidate")
+        for memory in memories:
+            mem_start = time.time()
 
-        # Process each memory sequentially
-        # Important: We process ALL pending memories before updating the watermark
-        # to avoid losing memories when many have the same timestamp
-        stats = {
-            "memories_processed": 0,
-            "mental_models_created": 0,
-            "mental_models_updated": 0,
-            "mental_models_merged": 0,
-            "actions_executed": 0,  # Total actions (can be > memories_processed due to multiple actions per fact)
-            "skipped": 0,
-        }
-
-        # Track processed memory IDs to avoid reprocessing
-        processed_ids: set[uuid.UUID] = set()
-        batch_num = 0
-
-        while memories:
-            batch_num += 1
-            batch_start = time.time()
-
-            for memory in memories:
-                if memory["id"] in processed_ids:
-                    continue
-
-                mem_start = time.time()
+            # Process the memory (uses its own connection internally)
+            async with pool.acquire() as conn:
                 result = await _process_memory(
                     conn=conn,
                     memory_engine=memory_engine,
@@ -190,118 +183,81 @@ async def run_consolidation_job(
                     request_context=request_context,
                     perf=perf,
                 )
-                mem_time = time.time() - mem_start
-                perf.record_timing("process_memory_total", mem_time)
 
-                processed_ids.add(memory["id"])
-                stats["memories_processed"] += 1
-
-                action = result.get("action")
-                if action == "created":
-                    stats["mental_models_created"] += 1
-                    stats["actions_executed"] += 1
-                elif action == "updated":
-                    stats["mental_models_updated"] += 1
-                    stats["actions_executed"] += 1
-                elif action == "merged":
-                    stats["mental_models_merged"] += 1
-                    stats["actions_executed"] += 1
-                elif action == "multiple":
-                    # Multiple actions from one fact (tag routing)
-                    stats["mental_models_created"] += result.get("created", 0)
-                    stats["mental_models_updated"] += result.get("updated", 0)
-                    stats["mental_models_merged"] += result.get("merged", 0)
-                    stats["actions_executed"] += result.get("total_actions", 0)
-                elif action == "skipped":
-                    stats["skipped"] += 1
-
-            batch_time = time.time() - batch_start
-            perf.log(
-                f"[2] Batch {batch_num}: {len(memories)} memories in {batch_time:.3f}s "
-                f"(avg {batch_time / len(memories):.3f}s/memory)"
-            )
-
-            # Fetch next batch of memories (excluding already processed)
-            t0 = time.time()
-            if last_consolidated_at:
-                memories = await conn.fetch(
+                # Mark memory as consolidated (committed immediately)
+                await conn.execute(
                     f"""
-                    SELECT id, text, fact_type, occurred_start, event_date, tags
-                    FROM {fq_table("memory_units")}
-                    WHERE bank_id = $1 AND created_at > $2
-                      AND fact_type IN ('experience', 'world')
-                      AND id != ALL($4)
-                    ORDER BY created_at ASC
-                    LIMIT $3
+                    UPDATE {fq_table("memory_units")}
+                    SET consolidated_at = NOW()
+                    WHERE id = $1
                     """,
-                    bank_id,
-                    last_consolidated_at,
-                    max_memories_per_batch,
-                    list(processed_ids),
+                    memory["id"],
                 )
-            else:
-                memories = await conn.fetch(
-                    f"""
-                    SELECT id, text, fact_type, occurred_start, event_date, tags
-                    FROM {fq_table("memory_units")}
-                    WHERE bank_id = $1
-                      AND fact_type IN ('experience', 'world')
-                      AND id != ALL($3)
-                    ORDER BY created_at ASC
-                    LIMIT $2
-                    """,
-                    bank_id,
-                    max_memories_per_batch,
-                    list(processed_ids),
+
+            mem_time = time.time() - mem_start
+            perf.record_timing("process_memory_total", mem_time)
+
+            stats["memories_processed"] += 1
+
+            action = result.get("action")
+            if action == "created":
+                stats["mental_models_created"] += 1
+                stats["actions_executed"] += 1
+            elif action == "updated":
+                stats["mental_models_updated"] += 1
+                stats["actions_executed"] += 1
+            elif action == "merged":
+                stats["mental_models_merged"] += 1
+                stats["actions_executed"] += 1
+            elif action == "multiple":
+                stats["mental_models_created"] += result.get("created", 0)
+                stats["mental_models_updated"] += result.get("updated", 0)
+                stats["mental_models_merged"] += result.get("merged", 0)
+                stats["actions_executed"] += result.get("total_actions", 0)
+            elif action == "skipped":
+                stats["skipped"] += 1
+
+            # Log progress periodically
+            if stats["memories_processed"] % 10 == 0:
+                logger.info(
+                    f"[CONSOLIDATION] bank={bank_id} progress: "
+                    f"{stats['memories_processed']}/{total_count} memories processed"
                 )
-            perf.record_timing("fetch_memories", time.time() - t0)
 
-        # Update last_consolidated_at only after ALL memories are processed
-        t0 = time.time()
-        await _update_last_consolidated_at(conn, bank_id)
-        perf.record_timing("update_watermark", time.time() - t0)
-
-        # Build summary
+        batch_time = time.time() - batch_start
         perf.log(
-            f"[3] Results: {stats['memories_processed']} memories → "
-            f"{stats['actions_executed']} actions "
-            f"({stats['mental_models_created']} created, "
-            f"{stats['mental_models_updated']} updated, "
-            f"{stats['mental_models_merged']} merged, "
-            f"{stats['skipped']} skipped)"
+            f"[2] Batch {batch_num}: {len(memories)} memories in {batch_time:.3f}s "
+            f"(avg {batch_time / len(memories):.3f}s/memory)"
         )
 
-        # Add timing breakdown
-        timing_parts = []
-        if "recall" in perf.timings:
-            timing_parts.append(f"recall={perf.timings['recall']:.3f}s")
-        if "llm" in perf.timings:
-            timing_parts.append(f"llm={perf.timings['llm']:.3f}s")
-        if "embedding" in perf.timings:
-            timing_parts.append(f"embedding={perf.timings['embedding']:.3f}s")
-        if "db_write" in perf.timings:
-            timing_parts.append(f"db_write={perf.timings['db_write']:.3f}s")
-
-        if timing_parts:
-            perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
-
-        perf.flush()
-
-        return {"status": "completed", "bank_id": bank_id, **stats}
-
-
-async def _update_last_consolidated_at(conn: "Connection", bank_id: str) -> None:
-    """Update the bank's last_consolidated_at timestamp."""
-    await conn.execute(
-        f"""
-        UPDATE {fq_table("banks")}
-        SET last_consolidated_at = $1
-        WHERE bank_id = $2
-        """,
-        datetime.now(timezone.utc),
-        bank_id,
+    # Build summary
+    perf.log(
+        f"[3] Results: {stats['memories_processed']} memories -> "
+        f"{stats['actions_executed']} actions "
+        f"({stats['mental_models_created']} created, "
+        f"{stats['mental_models_updated']} updated, "
+        f"{stats['mental_models_merged']} merged, "
+        f"{stats['skipped']} skipped)"
     )
 
+    # Add timing breakdown
+    timing_parts = []
+    if "recall" in perf.timings:
+        timing_parts.append(f"recall={perf.timings['recall']:.3f}s")
+    if "llm" in perf.timings:
+        timing_parts.append(f"llm={perf.timings['llm']:.3f}s")
+    if "embedding" in perf.timings:
+        timing_parts.append(f"embedding={perf.timings['embedding']:.3f}s")
+    if "db_write" in perf.timings:
+        timing_parts.append(f"db_write={perf.timings['db_write']:.3f}s")
+
+    if timing_parts:
+        perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
+
+    perf.flush()
+
+    return {"status": "completed", "bank_id": bank_id, **stats}
+
 
 async def _process_memory(
     conn: "Connection",
@@ -372,6 +328,7 @@ async def _process_memory(
                 memory_id=memory_id,
                 action=action,
                 mental_models=related_mental_models,
+                source_mentioned_at=memory.get("mentioned_at"),
                 perf=perf,
             )
             results.append(result)
@@ -384,6 +341,7 @@ async def _process_memory(
                 action=action,
                 event_date=memory.get("event_date"),
                 occurred_start=memory.get("occurred_start"),
+                mentioned_at=memory.get("mentioned_at"),
                 perf=perf,
             )
             results.append(result)
@@ -416,12 +374,14 @@ async def _execute_update_action(
     memory_id: uuid.UUID,
     action: dict[str, Any],
     mental_models: list[dict[str, Any]],
+    source_mentioned_at: datetime | None = None,
     perf: ConsolidationPerfLog | None = None,
 ) -> dict[str, Any]:
     """
     Execute an update action on an existing mental model.
 
-    Updates the mental model text, adds to history, and increments proof_count.
+    Updates the mental model text, adds to history, increments proof_count,
+    and updates mentioned_at if the new source memory has a more recent date.
     """
     learning_id = action.get("learning_id")
     new_text = action.get("text")
@@ -458,6 +418,7 @@ async def _execute_update_action(
         perf.record_timing("embedding", time.time() - t0)
 
     # Update the mental model
+    # Update mentioned_at if source memory has a more recent date
     t0 = time.time()
     await conn.execute(
         f"""
@@ -467,7 +428,8 @@ async def _execute_update_action(
             history = $3,
             source_memory_ids = $4,
             proof_count = $5,
-            updated_at = now()
+            updated_at = now(),
+            mentioned_at = GREATEST(mentioned_at, COALESCE($7, mentioned_at))
         WHERE id = $6
         """,
         new_text,
@@ -476,6 +438,7 @@ async def _execute_update_action(
         source_ids,
         len(source_ids),
         uuid.UUID(learning_id),
+        source_mentioned_at,
     )
 
     # Create links from memory to mental model
@@ -496,6 +459,7 @@ async def _execute_create_action(
     action: dict[str, Any],
     event_date: datetime | None = None,
     occurred_start: datetime | None = None,
+    mentioned_at: datetime | None = None,
     perf: ConsolidationPerfLog | None = None,
 ) -> dict[str, Any]:
     """
@@ -520,6 +484,7 @@ async def _execute_create_action(
         tags=tags,
         event_date=event_date,
         occurred_start=occurred_start,
+        mentioned_at=mentioned_at,
         perf=perf,
     )
 
@@ -751,7 +716,7 @@ Focus on DURABLE knowledge that serves this mission, not ephemeral state.
     ]
 
     try:
-        result = await memory_engine._llm_config.call(
+        result = await memory_engine._consolidation_llm_config.call(
             messages=messages,
             skip_validation=True,  # Raw JSON response
             scope="consolidation",
@@ -790,6 +755,7 @@ async def _create_mental_model_directly(
     tags: list[str] | None = None,
     event_date: datetime | None = None,
     occurred_start: datetime | None = None,
+    mentioned_at: datetime | None = None,
     perf: ConsolidationPerfLog | None = None,
 ) -> dict[str, Any]:
     """
@@ -809,6 +775,7 @@ async def _create_mental_model_directly(
     now = datetime.now(timezone.utc)
     mm_event_date = event_date or now
     mm_occurred_start = occurred_start or now
+    mm_mentioned_at = mentioned_at or now
     mm_tags = tags or []
 
     t0 = time.time()
@@ -817,9 +784,9 @@ async def _create_mental_model_directly(
         f"""
         INSERT INTO {fq_table("memory_units")} (
             id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
-            tags, event_date, occurred_start
+            tags, event_date, occurred_start, mentioned_at
         )
-        VALUES ($1, $2, $3, 'mental_model', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8)
+        VALUES ($1, $2, $3, 'mental_model', $4::vector, 1, $5, '[]'::jsonb, $6, $7, $8, $9)
         RETURNING id
         """,
         mental_model_id,
@@ -830,6 +797,7 @@ async def _create_mental_model_directly(
         mm_tags,
         mm_event_date,
         mm_occurred_start,
+        mm_mentioned_at,
     )
 
     # Create links between memory and mental model (includes entity links, memory_links)

hindsight-api/hindsight_api/engine/llm_wrapper.py

@@ -647,7 +647,13 @@ class LLMProvider:
                         success=True,
                     )
 
-                    return LLMToolCallResult(content=content, tool_calls=tool_calls, finish_reason=finish_reason)
+                    return LLMToolCallResult(
+                        content=content,
+                        tool_calls=tool_calls,
+                        finish_reason=finish_reason,
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                    )
 
                 except APIConnectionError as e:
                     last_exception = e
@@ -797,6 +803,10 @@ class LLMProvider:
                 content = "".join(content_parts) if content_parts else None
                 finish_reason = "tool_calls" if tool_calls else "stop"
 
+                # Extract token usage
+                input_tokens = response.usage.input_tokens or 0
+                output_tokens = response.usage.output_tokens or 0
+
                 # Record metrics
                 metrics = get_metrics_collector()
                 metrics.record_llm_call(
@@ -804,12 +814,18 @@ class LLMProvider:
                     model=self.model,
                     scope=scope,
                     duration=time.time() - start_time,
-                    input_tokens=response.usage.input_tokens or 0,
-                    output_tokens=response.usage.output_tokens or 0,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
                     success=True,
                 )
 
-                return LLMToolCallResult(content=content, tool_calls=tool_calls, finish_reason=finish_reason)
+                return LLMToolCallResult(
+                    content=content,
+                    tool_calls=tool_calls,
+                    finish_reason=finish_reason,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                )
 
             except (APIConnectionError, APIStatusError) as e:
                 if isinstance(e, APIStatusError) and e.status_code in (401, 403):
@@ -930,7 +946,13 @@ class LLMProvider:
                     success=True,
                 )
 
-                return LLMToolCallResult(content=content, tool_calls=tool_calls, finish_reason=finish_reason)
+                return LLMToolCallResult(
+                    content=content,
+                    tool_calls=tool_calls,
+                    finish_reason=finish_reason,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                )
 
             except genai_errors.APIError as e:
                 if e.code in (401, 403):

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -222,6 +222,10 @@ class MemoryEngine(MemoryEngineInterface):
         reflect_llm_api_key: str | None = None,
         reflect_llm_model: str | None = None,
         reflect_llm_base_url: str | None = None,
+        consolidation_llm_provider: str | None = None,
+        consolidation_llm_api_key: str | None = None,
+        consolidation_llm_model: str | None = None,
+        consolidation_llm_base_url: str | None = None,
         embeddings: Embeddings | None = None,
         cross_encoder: CrossEncoderModel | None = None,
         query_analyzer: QueryAnalyzer | None = None,
@@ -257,6 +261,10 @@ class MemoryEngine(MemoryEngineInterface):
             reflect_llm_api_key: API key for reflect LLM. Falls back to memory_llm_api_key.
             reflect_llm_model: Model for reflect operations. Falls back to memory_llm_model.
             reflect_llm_base_url: Base URL for reflect LLM. Falls back to memory_llm_base_url.
+            consolidation_llm_provider: LLM provider for consolidation operations. Falls back to memory_llm_provider.
+            consolidation_llm_api_key: API key for consolidation LLM. Falls back to memory_llm_api_key.
+            consolidation_llm_model: Model for consolidation operations. Falls back to memory_llm_model.
+            consolidation_llm_base_url: Base URL for consolidation LLM. Falls back to memory_llm_base_url.
             embeddings: Embeddings implementation. If not provided, created from env vars.
             cross_encoder: Cross-encoder model. If not provided, created from env vars.
             query_analyzer: Query analyzer implementation. If not provided, uses DateparserQueryAnalyzer.
@@ -398,6 +406,27 @@ class MemoryEngine(MemoryEngineInterface):
             model=reflect_model,
         )
 
+        # Consolidation LLM config - for mental model consolidation (can use efficient models)
+        consolidation_provider = consolidation_llm_provider or config.consolidation_llm_provider or memory_llm_provider
+        consolidation_api_key = consolidation_llm_api_key or config.consolidation_llm_api_key or memory_llm_api_key
+        consolidation_model = consolidation_llm_model or config.consolidation_llm_model or memory_llm_model
+        consolidation_base_url = consolidation_llm_base_url or config.consolidation_llm_base_url or memory_llm_base_url
+        # Apply provider-specific base URL defaults for consolidation
+        if consolidation_base_url is None:
+            if consolidation_provider.lower() == "groq":
+                consolidation_base_url = "https://api.groq.com/openai/v1"
+            elif consolidation_provider.lower() == "ollama":
+                consolidation_base_url = "http://localhost:11434/v1"
+            else:
+                consolidation_base_url = ""
+
+        self._consolidation_llm_config = LLMConfig(
+            provider=consolidation_provider,
+            api_key=consolidation_api_key,
+            base_url=consolidation_base_url,
+            model=consolidation_model,
+        )
+
         # Initialize cross-encoder reranker (cached for performance)
         self._cross_encoder_reranker = CrossEncoderReranker(cross_encoder=cross_encoder)
 
@@ -584,7 +613,10 @@ class MemoryEngine(MemoryEngineInterface):
                 ]
                 for fact_type, facts in reflect_result.based_on.items()
             },
-            "mental_models": [],  # Mental models are included in based_on["mental-models"]
+            # Extract mental models from based_on["mental-models"] for easy UI access
+            "mental_models": [
+                {"id": str(fact.id), "text": fact.text} for fact in reflect_result.based_on.get("mental-models", [])
+            ],
         }
 
         # Update the reflection with the generated content and reflect_response
@@ -598,6 +630,79 @@ class MemoryEngine(MemoryEngineInterface):
 
         logger.info(f"[CREATE_REFLECTION_TASK] Completed for bank_id={bank_id}, reflection_id={reflection_id}")
 
+    async def _handle_refresh_reflection(self, task_dict: dict[str, Any]):
+        """
+        Handler for refresh_reflection tasks.
+
+        Re-runs the source query through reflect and updates the reflection content.
+
+        Args:
+            task_dict: Dict with 'bank_id', 'reflection_id', 'operation_id'
+
+        Raises:
+            ValueError: If required fields are missing
+            Exception: Any exception from reflect/update (propagates to execute_task for retry)
+        """
+        bank_id = task_dict.get("bank_id")
+        reflection_id = task_dict.get("reflection_id")
+
+        if not bank_id or not reflection_id:
+            raise ValueError("bank_id and reflection_id are required for refresh_reflection task")
+
+        logger.info(f"[REFRESH_REFLECTION_TASK] Starting for bank_id={bank_id}, reflection_id={reflection_id}")
+
+        from hindsight_api.models import RequestContext
+
+        internal_context = RequestContext()
+
+        # Get the current reflection to get source_query
+        reflection = await self.get_reflection(bank_id, reflection_id, request_context=internal_context)
+        if not reflection:
+            raise ValueError(f"Reflection {reflection_id} not found in bank {bank_id}")
+
+        source_query = reflection["source_query"]
+
+        # Run reflect to generate new content, excluding the reflection being refreshed
+        reflect_result = await self.reflect_async(
+            bank_id=bank_id,
+            query=source_query,
+            request_context=internal_context,
+            exclude_reflection_ids=[reflection_id],
+        )
+
+        generated_content = reflect_result.text or "No content generated"
+
+        # Build reflect_response payload to store
+        reflect_response = {
+            "text": reflect_result.text,
+            "based_on": {
+                fact_type: [
+                    {
+                        "id": str(fact.id),
+                        "text": fact.text,
+                        "type": fact_type,
+                    }
+                    for fact in facts
+                ]
+                for fact_type, facts in reflect_result.based_on.items()
+            },
+            # Extract mental models from based_on["mental-models"] for easy UI access
+            "mental_models": [
+                {"id": str(fact.id), "text": fact.text} for fact in reflect_result.based_on.get("mental-models", [])
+            ],
+        }
+
+        # Update the reflection with the generated content and reflect_response
+        await self.update_reflection(
+            bank_id=bank_id,
+            reflection_id=reflection_id,
+            content=generated_content,
+            reflect_response=reflect_response,
+            request_context=internal_context,
+        )
+
+        logger.info(f"[REFRESH_REFLECTION_TASK] Completed for bank_id={bank_id}, reflection_id={reflection_id}")
+
     async def execute_task(self, task_dict: dict[str, Any]):
         """
         Execute a task by routing it to the appropriate handler.
@@ -638,6 +743,8 @@ class MemoryEngine(MemoryEngineInterface):
                 await self._handle_consolidation(task_dict)
             elif task_type == "create_reflection":
                 await self._handle_create_reflection(task_dict)
+            elif task_type == "refresh_reflection":
+                await self._handle_refresh_reflection(task_dict)
             else:
                 logger.error(f"Unknown task type: {task_type}")
                 # Don't retry unknown task types
@@ -792,6 +899,23 @@ class MemoryEngine(MemoryEngineInterface):
                 )
                 if reflect_is_different:
                     await self._reflect_llm_config.verify_connection()
+                # Verify consolidation config if different from all others
+                consolidation_is_different = (
+                    (
+                        self._consolidation_llm_config.provider != self._llm_config.provider
+                        or self._consolidation_llm_config.model != self._llm_config.model
+                    )
+                    and (
+                        self._consolidation_llm_config.provider != self._retain_llm_config.provider
+                        or self._consolidation_llm_config.model != self._retain_llm_config.model
+                    )
+                    and (
+                        self._consolidation_llm_config.provider != self._reflect_llm_config.provider
+                        or self._consolidation_llm_config.model != self._reflect_llm_config.model
+                    )
+                )
+                if consolidation_is_different:
+                    await self._consolidation_llm_config.verify_connection()
 
         # Build list of initialization tasks
         init_tasks = [
@@ -2995,11 +3119,11 @@ class MemoryEngine(MemoryEngineInterface):
         await self._authenticate_tenant(request_context)
         pool = await self._get_pool()
         async with acquire_with_retry(pool) as conn:
-            # Get the memory unit
+            # Get the memory unit (include source_memory_ids for mental models)
             row = await conn.fetchrow(
                 f"""
                 SELECT id, text, context, event_date, occurred_start, occurred_end,
-                       mentioned_at, fact_type, document_id, chunk_id, tags
+                       mentioned_at, fact_type, document_id, chunk_id, tags, source_memory_ids
                 FROM {fq_table("memory_units")}
                 WHERE id = $1 AND bank_id = $2
                 """,
@@ -3022,7 +3146,7 @@ class MemoryEngine(MemoryEngineInterface):
             )
             entities = [r["canonical_name"] for r in entities_rows]
 
-            return {
+            result = {
                 "id": str(row["id"]),
                 "text": row["text"],
                 "context": row["context"] if row["context"] else "",
@@ -3037,6 +3161,35 @@ class MemoryEngine(MemoryEngineInterface):
                 "tags": row["tags"] if row["tags"] else [],
             }
 
+            # For mental models, include source_memory_ids and fetch source_memories
+            if row["fact_type"] == "mental_model" and row["source_memory_ids"]:
+                source_ids = row["source_memory_ids"]
+                result["source_memory_ids"] = [str(sid) for sid in source_ids]
+
+                # Fetch source memories
+                source_rows = await conn.fetch(
+                    f"""
+                    SELECT id, text, fact_type, context, occurred_start, mentioned_at
+                    FROM {fq_table("memory_units")}
+                    WHERE id = ANY($1::uuid[])
+                    ORDER BY mentioned_at DESC NULLS LAST
+                    """,
+                    source_ids,
+                )
+                result["source_memories"] = [
+                    {
+                        "id": str(r["id"]),
+                        "text": r["text"],
+                        "type": r["fact_type"],
+                        "context": r["context"],
+                        "occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None,
+                        "mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None,
+                    }
+                    for r in source_rows
+                ]
+
+            return result
+
     async def list_documents(
         self,
         bank_id: str,
@@ -3671,13 +3824,22 @@ class MemoryEngine(MemoryEngineInterface):
             DirectiveRef(id=d.id, name=d.name, rules=d.rules) for d in agent_result.directives_applied
         ]
 
+        # Convert agent usage to TokenUsage format
+        from hindsight_api.engine.response_models import TokenUsage
+
+        usage = TokenUsage(
+            input_tokens=agent_result.usage.input_tokens,
+            output_tokens=agent_result.usage.output_tokens,
+            total_tokens=agent_result.usage.total_tokens,
+        )
+
         # Return response (compatible with existing API)
         result = ReflectResult(
             text=agent_result.text,
             based_on=based_on,
             new_opinions=[],  # Learnings stored as mental models
             structured_output=agent_result.structured_output,
-            usage=None,  # Token tracking not yet implemented for agentic loop
+            usage=usage,
             tool_trace=tool_trace_result,
             llm_trace=llm_trace_result,
             directives_applied=directives_applied_result,
@@ -5650,3 +5812,40 @@ class MemoryEngine(MemoryEngineInterface):
             result_metadata={"reflection_id": reflection_id, "name": name, "source_query": source_query},
             dedupe_by_bank=False,
         )
+
+    async def submit_async_refresh_reflection(
+        self,
+        bank_id: str,
+        reflection_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """Submit an async reflection refresh operation.
+
+        This schedules a background task to re-run the source query and update the content.
+
+        Args:
+            bank_id: Bank identifier
+            reflection_id: Reflection UUID to refresh
+            request_context: Request context for authentication
+
+        Returns:
+            Dict with operation_id
+        """
+        await self._authenticate_tenant(request_context)
+
+        # Verify reflection exists
+        reflection = await self.get_reflection(bank_id, reflection_id, request_context=request_context)
+        if not reflection:
+            raise ValueError(f"Reflection {reflection_id} not found in bank {bank_id}")
+
+        return await self._submit_async_operation(
+            bank_id=bank_id,
+            operation_type="refresh_reflection",
+            task_type="refresh_reflection",
+            task_payload={
+                "reflection_id": reflection_id,
+            },
+            result_metadata={"reflection_id": reflection_id, "name": reflection["name"]},
+            dedupe_by_bank=False,
+        )

hindsight-api/hindsight_api/engine/reflect/agent.py

@@ -10,10 +10,11 @@ Uses hierarchical retrieval:
 import asyncio
 import json
 import logging
+import re
 import time
 from typing import TYPE_CHECKING, Any, Awaitable, Callable
 
-from .models import DirectiveInfo, LLMCall, ReflectAgentResult, ToolCall
+from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
 from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools
 from .tools_schema import get_reflect_tools
 
@@ -77,12 +78,27 @@ def _is_done_tool(name: str) -> bool:
     return _normalize_tool_name(name) == "done"
 
 
+# Pattern to match done() call as text - handles done({...}) with nested JSON
+_DONE_CALL_PATTERN = re.compile(r"done\s*\(\s*\{.*$", re.DOTALL)
+
+
+def _clean_answer_text(text: str) -> str:
+    """Clean up answer text by removing any done() tool call syntax.
+
+    Some LLMs output the done() call as text instead of a proper tool call.
+    This strips out patterns like: done({"answer": "...", ...})
+    """
+    # Remove done() call pattern from the end of the text
+    cleaned = _DONE_CALL_PATTERN.sub("", text).strip()
+    return cleaned if cleaned else text
+
+
 async def _generate_structured_output(
     answer: str,
     response_schema: dict,
     llm_config: "LLMProvider",
     reflect_id: str,
-) -> dict[str, Any] | None:
+) -> tuple[dict[str, Any] | None, int, int]:
     """Generate structured output from an answer using the provided JSON schema.
 
     Args:
@@ -92,7 +108,8 @@ async def _generate_structured_output(
         reflect_id: Reflect ID for logging
 
     Returns:
-        Structured output dict if successful, None otherwise
+        Tuple of (structured_output, input_tokens, output_tokens).
+        structured_output is None if generation fails.
     """
     try:
         from typing import Any as TypingAny
@@ -149,7 +166,7 @@ Return ONLY a valid JSON object that matches this exact schema. Pay special atte
 
 Do not include any explanation, only the JSON object."""
 
-        structured_result = await llm_config.call(
+        structured_result, usage = await llm_config.call(
             messages=[
                 {
                     "role": "system",
@@ -160,6 +177,7 @@ Do not include any explanation, only the JSON object."""
             response_format=DynamicModel,
             scope="reflect_structured",
             skip_validation=True,  # We'll handle the dict ourselves
+            return_usage=True,
         )
 
         # Convert to dict
@@ -172,11 +190,11 @@ Do not include any explanation, only the JSON object."""
             structured_output = json.loads(str(structured_result))
 
         logger.info(f"[REFLECT {reflect_id}] Generated structured output with {len(structured_output)} fields")
-        return structured_output
+        return structured_output, usage.input_tokens, usage.output_tokens
 
     except Exception as e:
         logger.warning(f"[REFLECT {reflect_id}] Failed to generate structured output: {e}")
-        return None
+        return None, 0, 0
 
 
 async def run_reflect_agent(
@@ -246,13 +264,32 @@ async def run_reflect_agent(
     llm_trace: list[dict[str, Any]] = []
     context_history: list[dict[str, Any]] = []  # For final prompt fallback
 
+    # Token usage tracking - accumulate across all LLM calls
+    total_input_tokens = 0
+    total_output_tokens = 0
+
     # Track available IDs for validation (prevents hallucinated citations)
     available_memory_ids: set[str] = set()
     available_reflection_ids: set[str] = set()
     available_mental_model_ids: set[str] = set()
 
     def _get_llm_trace() -> list[LLMCall]:
-        return [LLMCall(scope=c["scope"], duration_ms=c["duration_ms"]) for c in llm_trace]
+        return [
+            LLMCall(
+                scope=c["scope"],
+                duration_ms=c["duration_ms"],
+                input_tokens=c.get("input_tokens", 0),
+                output_tokens=c.get("output_tokens", 0),
+            )
+            for c in llm_trace
+        ]
+
+    def _get_usage() -> TokenUsageSummary:
+        return TokenUsageSummary(
+            input_tokens=total_input_tokens,
+            output_tokens=total_output_tokens,
+            total_tokens=total_input_tokens + total_output_tokens,
+        )
 
     def _log_completion(answer: str, iterations: int, forced: bool = False):
         elapsed_ms = int((time.time() - start_time) * 1000)
@@ -286,21 +323,36 @@ async def run_reflect_agent(
             # Force text response on last iteration - no tools
             prompt = build_final_prompt(query, context_history, bank_profile, context)
             llm_start = time.time()
-            response = await llm_config.call(
+            response, usage = await llm_config.call(
                 messages=[
                     {"role": "system", "content": FINAL_SYSTEM_PROMPT},
                     {"role": "user", "content": prompt},
                 ],
                 scope="reflect_agent_final",
                 max_completion_tokens=max_tokens,
+                return_usage=True,
             )
-            llm_trace.append({"scope": "final", "duration_ms": int((time.time() - llm_start) * 1000)})
-            answer = response.strip()
+            llm_duration = int((time.time() - llm_start) * 1000)
+            total_input_tokens += usage.input_tokens
+            total_output_tokens += usage.output_tokens
+            llm_trace.append(
+                {
+                    "scope": "final",
+                    "duration_ms": llm_duration,
+                    "input_tokens": usage.input_tokens,
+                    "output_tokens": usage.output_tokens,
+                }
+            )
+            answer = _clean_answer_text(response.strip())
 
             # Generate structured output if schema provided
             structured_output = None
             if response_schema and answer:
-                structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
+                structured_output, struct_in, struct_out = await _generate_structured_output(
+                    answer, response_schema, llm_config, reflect_id
+                )
+                total_input_tokens += struct_in
+                total_output_tokens += struct_out
 
             _log_completion(answer, iteration + 1, forced=True)
             return ReflectAgentResult(
@@ -310,6 +362,7 @@ async def run_reflect_agent(
                 tools_called=total_tools_called,
                 tool_trace=tool_trace,
                 llm_trace=_get_llm_trace(),
+                usage=_get_usage(),
                 directives_applied=directives_applied,
             )
 
@@ -324,7 +377,16 @@ async def run_reflect_agent(
                 tool_choice="required" if iteration == 0 else "auto",  # Force tool use on first iteration
             )
             llm_duration = int((time.time() - llm_start) * 1000)
-            llm_trace.append({"scope": f"agent_{iteration + 1}", "duration_ms": llm_duration})
+            total_input_tokens += result.input_tokens
+            total_output_tokens += result.output_tokens
+            llm_trace.append(
+                {
+                    "scope": f"agent_{iteration + 1}",
+                    "duration_ms": llm_duration,
+                    "input_tokens": result.input_tokens,
+                    "output_tokens": result.output_tokens,
+                }
+            )
 
         except Exception as e:
             err_duration = int((time.time() - llm_start) * 1000)
@@ -338,21 +400,36 @@ async def run_reflect_agent(
                 continue
             prompt = build_final_prompt(query, context_history, bank_profile, context)
             llm_start = time.time()
-            response = await llm_config.call(
+            response, usage = await llm_config.call(
                 messages=[
                     {"role": "system", "content": FINAL_SYSTEM_PROMPT},
                     {"role": "user", "content": prompt},
                 ],
                 scope="reflect_agent_final",
                 max_completion_tokens=max_tokens,
+                return_usage=True,
             )
-            llm_trace.append({"scope": "final", "duration_ms": int((time.time() - llm_start) * 1000)})
-            answer = response.strip()
+            llm_duration = int((time.time() - llm_start) * 1000)
+            total_input_tokens += usage.input_tokens
+            total_output_tokens += usage.output_tokens
+            llm_trace.append(
+                {
+                    "scope": "final",
+                    "duration_ms": llm_duration,
+                    "input_tokens": usage.input_tokens,
+                    "output_tokens": usage.output_tokens,
+                }
+            )
+            answer = _clean_answer_text(response.strip())
 
             # Generate structured output if schema provided
             structured_output = None
             if response_schema and answer:
-                structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
+                structured_output, struct_in, struct_out = await _generate_structured_output(
+                    answer, response_schema, llm_config, reflect_id
+                )
+                total_input_tokens += struct_in
+                total_output_tokens += struct_out
 
             _log_completion(answer, iteration + 1, forced=True)
             return ReflectAgentResult(
@@ -362,20 +439,23 @@ async def run_reflect_agent(
                 tools_called=total_tools_called,
                 tool_trace=tool_trace,
                 llm_trace=_get_llm_trace(),
+                usage=_get_usage(),
                 directives_applied=directives_applied,
             )
 
         # No tool calls - LLM wants to respond with text
         if not result.tool_calls:
             if result.content:
-                answer = result.content.strip()
+                answer = _clean_answer_text(result.content.strip())
 
                 # Generate structured output if schema provided
                 structured_output = None
                 if response_schema and answer:
-                    structured_output = await _generate_structured_output(
+                    structured_output, struct_in, struct_out = await _generate_structured_output(
                         answer, response_schema, llm_config, reflect_id
                     )
+                    total_input_tokens += struct_in
+                    total_output_tokens += struct_out
 
                 _log_completion(answer, iteration + 1)
                 return ReflectAgentResult(
@@ -385,26 +465,42 @@ async def run_reflect_agent(
                     tools_called=total_tools_called,
                     tool_trace=tool_trace,
                     llm_trace=_get_llm_trace(),
+                    usage=_get_usage(),
                     directives_applied=directives_applied,
                 )
             # Empty response, force final
             prompt = build_final_prompt(query, context_history, bank_profile, context)
             llm_start = time.time()
-            response = await llm_config.call(
+            response, usage = await llm_config.call(
                 messages=[
                     {"role": "system", "content": FINAL_SYSTEM_PROMPT},
                     {"role": "user", "content": prompt},
                 ],
                 scope="reflect_agent_final",
                 max_completion_tokens=max_tokens,
+                return_usage=True,
             )
-            llm_trace.append({"scope": "final", "duration_ms": int((time.time() - llm_start) * 1000)})
-            answer = response.strip()
+            llm_duration = int((time.time() - llm_start) * 1000)
+            total_input_tokens += usage.input_tokens
+            total_output_tokens += usage.output_tokens
+            llm_trace.append(
+                {
+                    "scope": "final",
+                    "duration_ms": llm_duration,
+                    "input_tokens": usage.input_tokens,
+                    "output_tokens": usage.output_tokens,
+                }
+            )
+            answer = _clean_answer_text(response.strip())
 
             # Generate structured output if schema provided
             structured_output = None
             if response_schema and answer:
-                structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
+                structured_output, struct_in, struct_out = await _generate_structured_output(
+                    answer, response_schema, llm_config, reflect_id
+                )
+                total_input_tokens += struct_in
+                total_output_tokens += struct_out
 
             _log_completion(answer, iteration + 1, forced=True)
             return ReflectAgentResult(
@@ -414,6 +510,7 @@ async def run_reflect_agent(
                 tools_called=total_tools_called,
                 tool_trace=tool_trace,
                 llm_trace=_get_llm_trace(),
+                usage=_get_usage(),
                 directives_applied=directives_applied,
             )
 
@@ -455,6 +552,7 @@ async def run_reflect_agent(
                 total_tools_called,
                 tool_trace,
                 _get_llm_trace(),
+                _get_usage(),
                 _log_completion,
                 reflect_id,
                 directives_applied=directives_applied,
@@ -575,6 +673,7 @@ async def run_reflect_agent(
         tools_called=total_tools_called,
         tool_trace=tool_trace,
         llm_trace=_get_llm_trace(),
+        usage=_get_usage(),
         directives_applied=directives_applied,
     )
 
@@ -600,6 +699,7 @@ async def _process_done_tool(
     total_tools_called: int,
     tool_trace: list[ToolCall],
     llm_trace: list[LLMCall],
+    usage: TokenUsageSummary,
     log_completion: Callable,
     reflect_id: str,
     directives_applied: list[DirectiveInfo],
@@ -620,8 +720,17 @@ async def _process_done_tool(
 
     # Generate structured output if schema provided
     structured_output = None
+    final_usage = usage
     if response_schema and llm_config and answer:
-        structured_output = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
+        structured_output, struct_in, struct_out = await _generate_structured_output(
+            answer, response_schema, llm_config, reflect_id
+        )
+        # Add structured output tokens to usage
+        final_usage = TokenUsageSummary(
+            input_tokens=usage.input_tokens + struct_in,
+            output_tokens=usage.output_tokens + struct_out,
+            total_tokens=usage.total_tokens + struct_in + struct_out,
+        )
 
     log_completion(answer, iterations)
     return ReflectAgentResult(
@@ -631,6 +740,7 @@ async def _process_done_tool(
         tools_called=total_tools_called,
         tool_trace=tool_trace,
         llm_trace=llm_trace,
+        usage=final_usage,
         used_memory_ids=used_memory_ids,
         used_reflection_ids=used_reflection_ids,
         used_mental_model_ids=used_mental_model_ids,

hindsight-api/hindsight_api/engine/reflect/models.py

@@ -85,6 +85,8 @@ class LLMCall(BaseModel):
 
     scope: str = Field(description="Call scope: agent_1, agent_2, final, etc.")
     duration_ms: int = Field(description="Execution time in milliseconds")
+    input_tokens: int = Field(default=0, description="Input tokens used")
+    output_tokens: int = Field(default=0, description="Output tokens used")
 
 
 class DirectiveInfo(BaseModel):
@@ -95,6 +97,14 @@ class DirectiveInfo(BaseModel):
     rules: list[str] = Field(default_factory=list, description="Directive rules/observations that were applied")
 
 
+class TokenUsageSummary(BaseModel):
+    """Total token usage across all LLM calls."""
+
+    input_tokens: int = Field(default=0, description="Total input tokens used")
+    output_tokens: int = Field(default=0, description="Total output tokens used")
+    total_tokens: int = Field(default=0, description="Total tokens (input + output)")
+
+
 class ReflectAgentResult(BaseModel):
     """Result from the reflect agent."""
 
@@ -106,6 +116,9 @@ class ReflectAgentResult(BaseModel):
     tools_called: int = Field(default=0, description="Total number of tool calls made")
     tool_trace: list[ToolCall] = Field(default_factory=list, description="Trace of all tool calls made")
     llm_trace: list[LLMCall] = Field(default_factory=list, description="Trace of all LLM calls made")
+    usage: TokenUsageSummary = Field(
+        default_factory=TokenUsageSummary, description="Total token usage across all LLM calls"
+    )
     used_memory_ids: list[str] = Field(default_factory=list, description="Validated memory IDs actually used in answer")
     used_reflection_ids: list[str] = Field(
         default_factory=list, description="Validated reflection IDs actually used in answer"

hindsight-api/hindsight_api/engine/response_models.py

@@ -28,6 +28,8 @@ class LLMToolCallResult(BaseModel):
     content: str | None = Field(default=None, description="Text content if any")
     tool_calls: list[LLMToolCall] = Field(default_factory=list, description="Tool calls requested by the LLM")
     finish_reason: str | None = Field(default=None, description="Reason the LLM stopped: 'stop', 'tool_calls', etc.")
+    input_tokens: int = Field(default=0, description="Input tokens used in this call")
+    output_tokens: int = Field(default=0, description="Output tokens used in this call")
 
 
 class ToolCallTrace(BaseModel):

hindsight-api/hindsight_api/extensions/__init__.py

@@ -21,6 +21,10 @@ from hindsight_api.extensions.context import DefaultExtensionContext, ExtensionC
 from hindsight_api.extensions.http import HttpExtension
 from hindsight_api.extensions.loader import load_extension
 from hindsight_api.extensions.operation_validator import (
+    # Consolidation operation
+    ConsolidateContext,
+    ConsolidateResult,
+    # Core operations
     OperationValidationError,
     OperationValidatorExtension,
     RecallContext,
@@ -33,6 +37,7 @@ from hindsight_api.extensions.operation_validator import (
 )
 from hindsight_api.extensions.tenant import (
     AuthenticationError,
+    Tenant,
     TenantContext,
     TenantExtension,
 )
@@ -47,7 +52,7 @@ __all__ = [
     "DefaultExtensionContext",
     # HTTP Extension
     "HttpExtension",
-    # Operation Validator
+    # Operation Validator - Core
     "OperationValidationError",
     "OperationValidatorExtension",
     "RecallContext",
@@ -57,10 +62,14 @@ __all__ = [
     "RetainContext",
     "RetainResult",
     "ValidationResult",
+    # Operation Validator - Consolidation
+    "ConsolidateContext",
+    "ConsolidateResult",
     # Tenant/Auth
     "ApiKeyTenantExtension",
     "AuthenticationError",
     "RequestContext",
+    "Tenant",
     "TenantContext",
     "TenantExtension",
 ]

hindsight-api/hindsight_api/extensions/builtin/tenant.py

@@ -1,6 +1,6 @@
 """Built-in tenant extension implementations."""
 
-from hindsight_api.extensions.tenant import AuthenticationError, TenantContext, TenantExtension
+from hindsight_api.extensions.tenant import AuthenticationError, Tenant, TenantContext, TenantExtension
 from hindsight_api.models import RequestContext
 
 
@@ -31,3 +31,7 @@ class ApiKeyTenantExtension(TenantExtension):
         if context.api_key != self.expected_api_key:
             raise AuthenticationError("Invalid API key")
         return TenantContext(schema_name="public")
+
+    async def list_tenants(self) -> list[Tenant]:
+        """Return public schema for single-tenant setup."""
+        return [Tenant(schema="public")]

hindsight-api/hindsight_api/extensions/operation_validator.py

@@ -1,4 +1,4 @@
-"""Operation Validator Extension for validating retain/recall/reflect operations."""
+"""Operation Validator Extension for validating retain/recall/reflect/consolidate operations."""
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
@@ -97,6 +97,19 @@ class ReflectContext:
     context: str | None = None
 
 
+# =============================================================================
+# Consolidation Pre-operation Context
+# =============================================================================
+
+
+@dataclass
+class ConsolidateContext:
+    """Context for a consolidation operation validation (pre-operation)."""
+
+    bank_id: str
+    request_context: "RequestContext"
+
+
 # =============================================================================
 # Post-operation Contexts (includes results)
 # =============================================================================
@@ -164,9 +177,28 @@ class ReflectResultContext:
     error: str | None = None
 
 
+# =============================================================================
+# Consolidation Post-operation Context
+# =============================================================================
+
+
+@dataclass
+class ConsolidateResult:
+    """Result context for post-consolidation hook."""
+
+    bank_id: str
+    request_context: "RequestContext"
+    # Result
+    processed: int = 0
+    created: int = 0
+    updated: int = 0
+    success: bool = True
+    error: str | None = None
+
+
 class OperationValidatorExtension(Extension, ABC):
     """
-    Validates and hooks into retain/recall/reflect operations.
+    Validates and hooks into retain/recall/reflect/consolidate operations.
 
     This extension allows implementing custom logic such as:
     - Rate limiting (pre-operation)
@@ -185,9 +217,13 @@ class OperationValidatorExtension(Extension, ABC):
         -> config = {"max_requests": "100"}
 
     Hook execution order:
-        1. validate_retain/validate_recall/validate_reflect (pre-operation)
+        1. validate_* (pre-operation)
         2. [operation executes]
-        3. on_retain_complete/on_recall_complete/on_reflect_complete (post-operation)
+        3. on_*_complete (post-operation)
+
+    Supported operations:
+        - retain, recall, reflect (core memory operations)
+        - consolidate (mental models consolidation)
     """
 
     # =========================================================================
@@ -325,3 +361,44 @@ class OperationValidatorExtension(Extension, ABC):
                 - error: Error message (if failed)
         """
         pass
+
+    # =========================================================================
+    # Consolidation - Pre-operation validation hook (optional - override to implement)
+    # =========================================================================
+
+    async def validate_consolidate(self, ctx: ConsolidateContext) -> ValidationResult:
+        """
+        Validate a consolidation operation before execution.
+
+        Override to implement custom validation logic for consolidation.
+
+        Args:
+            ctx: Context containing:
+                - bank_id: Bank identifier
+                - request_context: Request context with auth info
+
+        Returns:
+            ValidationResult indicating whether the operation is allowed.
+        """
+        return ValidationResult.accept()
+
+    # =========================================================================
+    # Consolidation - Post-operation hook (optional - override to implement)
+    # =========================================================================
+
+    async def on_consolidate_complete(self, result: ConsolidateResult) -> None:
+        """
+        Called after a consolidation operation completes (success or failure).
+
+        Override to implement post-operation logic such as usage tracking or audit logging.
+
+        Args:
+            result: Result context containing:
+                - bank_id: Bank identifier
+                - processed: Number of memories processed
+                - created: Number of mental models created
+                - updated: Number of mental models updated
+                - success: Whether the operation succeeded
+                - error: Error message (if failed)
+        """
+        pass

hindsight-api/hindsight_api/extensions/tenant.py

@@ -28,6 +28,18 @@ class TenantContext:
     schema_name: str
 
 
+@dataclass
+class Tenant:
+    """
+    Represents a tenant for worker discovery.
+
+    Used by list_tenants() to return tenant information including
+    the PostgreSQL schema name for database operations.
+    """
+
+    schema: str
+
+
 class TenantExtension(Extension, ABC):
     """
     Extension for multi-tenancy and API key authentication.
@@ -61,3 +73,17 @@ class TenantExtension(Extension, ABC):
             AuthenticationError: If authentication fails.
         """
         ...
+
+    @abstractmethod
+    async def list_tenants(self) -> list[Tenant]:
+        """
+        List all tenants that should be processed by workers.
+
+        This method is used by the worker to discover all tenants that need
+        task polling. Workers will poll for pending tasks in each tenant's schema.
+
+        Returns:
+            List of Tenant objects containing schema information.
+            For single-tenant setups, return [Tenant(schema="public")].
+        """
+        ...

hindsight-api/hindsight_api/main.py

@@ -184,6 +184,10 @@ def main():
             reflect_llm_api_key=config.reflect_llm_api_key,
             reflect_llm_model=config.reflect_llm_model,
             reflect_llm_base_url=config.reflect_llm_base_url,
+            consolidation_llm_provider=config.consolidation_llm_provider,
+            consolidation_llm_api_key=config.consolidation_llm_api_key,
+            consolidation_llm_model=config.consolidation_llm_model,
+            consolidation_llm_base_url=config.consolidation_llm_base_url,
             embeddings_provider=config.embeddings_provider,
             embeddings_local_model=config.embeddings_local_model,
             embeddings_tei_url=config.embeddings_tei_url,

hindsight-api/hindsight_api/worker/main.py

@@ -181,6 +181,8 @@ def main():
         nonlocal memory, poller
         import uvicorn
 
+        from ..extensions import TenantExtension, load_extension
+
         # Initialize MemoryEngine
         # Workers use SyncTaskBackend because they execute tasks directly,
         # they don't need to store tasks (they poll from DB)
@@ -193,7 +195,15 @@ def main():
 
         print(f"Database connected: {config.database_url}")
 
-        # Create and start the poller
+        # Load tenant extension for dynamic schema discovery
+        tenant_extension = load_extension("TENANT", TenantExtension)
+
+        if tenant_extension:
+            print("Tenant extension loaded - schemas will be discovered dynamically on each poll")
+        else:
+            print("No tenant extension configured, using public schema only")
+
+        # Create a single poller that handles all schemas dynamically
         poller = WorkerPoller(
             pool=memory._pool,
             worker_id=args.worker_id,
@@ -201,6 +211,7 @@ def main():
             poll_interval_ms=args.poll_interval,
             batch_size=args.batch_size,
             max_retries=args.max_retries,
+            tenant_extension=tenant_extension,
         )
 
         # Create the HTTP app for metrics/health

hindsight-api/hindsight_api/worker/poller.py

@@ -11,11 +11,14 @@ import logging
 import time
 import traceback
 from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     import asyncpg
 
+    from hindsight_api.extensions.tenant import TenantExtension
+
 logger = logging.getLogger(__name__)
 
 # Progress logging interval in seconds
@@ -29,12 +32,23 @@ def fq_table(table: str, schema: str | None = None) -> str:
     return table
 
 
+@dataclass
+class ClaimedTask:
+    """A task claimed from the database with its schema context."""
+
+    operation_id: str
+    task_dict: dict[str, Any]
+    schema: str | None
+
+
 class WorkerPoller:
     """
     Polls PostgreSQL for pending tasks and executes them.
 
     Uses FOR UPDATE SKIP LOCKED for safe distributed claiming,
     allowing multiple workers to process tasks without conflicts.
+
+    Supports dynamic multi-tenant discovery via tenant_extension.
     """
 
     def __init__(
@@ -46,6 +60,7 @@ class WorkerPoller:
         batch_size: int = 10,
         max_retries: int = 3,
         schema: str | None = None,
+        tenant_extension: "TenantExtension | None" = None,
     ):
         """
         Initialize the worker poller.
@@ -57,7 +72,9 @@ class WorkerPoller:
             poll_interval_ms: Interval between polls when no tasks found (milliseconds)
             batch_size: Maximum number of tasks to claim per poll cycle
             max_retries: Maximum retry attempts before marking task as failed
-            schema: Database schema for multi-tenant support (optional)
+            schema: Database schema for single-tenant support (ignored if tenant_extension is set)
+            tenant_extension: Extension for dynamic multi-tenant discovery. If set, list_tenants()
+                            is called on each poll cycle to discover schemas dynamically.
         """
         self._pool = pool
         self._worker_id = worker_id
@@ -66,27 +83,56 @@ class WorkerPoller:
         self._batch_size = batch_size
         self._max_retries = max_retries
         self._schema = schema
+        self._tenant_extension = tenant_extension
         self._shutdown = asyncio.Event()
         self._current_tasks: set[asyncio.Task] = set()
         self._in_flight_count = 0
         self._in_flight_lock = asyncio.Lock()
         self._last_progress_log = 0.0
         self._tasks_completed_since_log = 0
-        self._active_banks: set[str] = set()
+        # Track active tasks locally: operation_id -> (op_type, bank_id, schema)
+        self._active_tasks: dict[str, tuple[str, str, str | None]] = {}
 
-    async def claim_batch(self) -> list[tuple[str, dict[str, Any]]]:
+    async def _get_schemas(self) -> list[str | None]:
+        """Get list of schemas to poll. Returns [None] for public schema."""
+        if self._tenant_extension is not None:
+            tenants = await self._tenant_extension.list_tenants()
+            # Convert "public" to None for SQL compatibility, keep others as-is
+            return [t.schema if t.schema != "public" else None for t in tenants]
+        # Single schema mode
+        return [self._schema]
+
+    async def claim_batch(self) -> list[ClaimedTask]:
         """
-        Claim up to batch_size pending tasks atomically.
+        Claim up to batch_size pending tasks atomically across all tenant schemas.
 
         Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers.
 
         For consolidation tasks specifically, skips pending tasks if there's already
         a processing consolidation for the same bank (to avoid duplicate work).
 
+        If tenant_extension is configured, dynamically discovers schemas on each call.
+
         Returns:
-            List of tuples (operation_id, task_dict)
+            List of ClaimedTask objects containing operation_id, task_dict, and schema
         """
-        table = fq_table("async_operations", self._schema)
+        schemas = await self._get_schemas()
+        all_tasks: list[ClaimedTask] = []
+        remaining_batch = self._batch_size
+
+        for schema in schemas:
+            if remaining_batch <= 0:
+                break
+
+            tasks = await self._claim_batch_for_schema(schema, remaining_batch)
+            all_tasks.extend(tasks)
+            remaining_batch -= len(tasks)
+
+        return all_tasks
+
+    async def _claim_batch_for_schema(self, schema: str | None, limit: int) -> list[ClaimedTask]:
+        """Claim tasks from a specific schema."""
+        table = fq_table("async_operations", schema)
 
         async with self._pool.acquire() as conn:
             async with conn.transaction():
@@ -113,7 +159,7 @@ class WorkerPoller:
                     LIMIT $1
                     FOR UPDATE SKIP LOCKED
                     """,
-                    self._batch_size,
+                    limit,
                 )
 
                 if not rows:
@@ -131,12 +177,19 @@ class WorkerPoller:
                     operation_ids,
                 )
 
-                # Parse and return task payloads
-                return [(str(row["operation_id"]), json.loads(row["task_payload"])) for row in rows]
+                # Parse and return task payloads with schema context
+                return [
+                    ClaimedTask(
+                        operation_id=str(row["operation_id"]),
+                        task_dict=json.loads(row["task_payload"]),
+                        schema=schema,
+                    )
+                    for row in rows
+                ]
 
-    async def _mark_completed(self, operation_id: str):
+    async def _mark_completed(self, operation_id: str, schema: str | None):
         """Mark a task as completed."""
-        table = fq_table("async_operations", self._schema)
+        table = fq_table("async_operations", schema)
         await self._pool.execute(
             f"""
             UPDATE {table}
@@ -146,9 +199,9 @@ class WorkerPoller:
             operation_id,
         )
 
-    async def _mark_failed(self, operation_id: str, error_message: str):
+    async def _mark_failed(self, operation_id: str, error_message: str, schema: str | None):
         """Mark a task as failed with error message."""
-        table = fq_table("async_operations", self._schema)
+        table = fq_table("async_operations", schema)
         # Truncate error message if too long (max 5000 chars in schema)
         error_message = error_message[:5000] if len(error_message) > 5000 else error_message
         await self._pool.execute(
@@ -161,9 +214,9 @@ class WorkerPoller:
             error_message,
         )
 
-    async def _retry_or_fail(self, operation_id: str, error_message: str):
+    async def _retry_or_fail(self, operation_id: str, error_message: str, schema: str | None):
         """Increment retry count or mark as failed if max retries exceeded."""
-        table = fq_table("async_operations", self._schema)
+        table = fq_table("async_operations", schema)
 
         # Get current retry count
         row = await self._pool.fetchrow(
@@ -180,7 +233,7 @@ class WorkerPoller:
         if retry_count >= self._max_retries:
             # Max retries exceeded, mark as failed
             await self._mark_failed(
-                operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}"
+                operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}", schema
             )
             logger.error(f"Task {operation_id} failed after {retry_count} retries")
         else:
@@ -196,20 +249,29 @@ class WorkerPoller:
             )
             logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})")
 
-    async def execute_task(self, operation_id: str, task_dict: dict[str, Any]):
+    async def execute_task(self, task: ClaimedTask):
         """Execute a single task and update its status."""
-        task_type = task_dict.get("type", "unknown")
-        bank_id = task_dict.get("bank_id", "unknown")
+        task_type = task.task_dict.get("type", "unknown")
+        bank_id = task.task_dict.get("bank_id", "unknown")
+
+        # Track this task as active
+        async with self._in_flight_lock:
+            self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema)
 
         try:
-            logger.debug(f"Executing task {operation_id} (type={task_type}, bank={bank_id})")
-            await self._executor(task_dict)
-            await self._mark_completed(operation_id)
-            logger.debug(f"Task {operation_id} completed successfully")
+            schema_info = f", schema={task.schema}" if task.schema else ""
+            logger.debug(f"Executing task {task.operation_id} (type={task_type}, bank={bank_id}{schema_info})")
+            await self._executor(task.task_dict)
+            await self._mark_completed(task.operation_id, task.schema)
+            logger.debug(f"Task {task.operation_id} completed successfully")
         except Exception as e:
             error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
-            logger.error(f"Task {operation_id} failed: {e}")
-            await self._retry_or_fail(operation_id, error_msg)
+            logger.error(f"Task {task.operation_id} failed: {e}")
+            await self._retry_or_fail(task.operation_id, error_msg, task.schema)
+        finally:
+            # Remove from active tasks
+            async with self._in_flight_lock:
+                self._active_tasks.pop(task.operation_id, None)
 
     async def recover_own_tasks(self) -> int:
         """
@@ -219,25 +281,33 @@ class WorkerPoller:
         On startup, we reset any tasks stuck in 'processing' for this worker_id
         back to 'pending' so they can be picked up again.
 
+        If tenant_extension is configured, recovers across all tenant schemas.
+
         Returns:
             Number of tasks recovered
         """
-        table = fq_table("async_operations", self._schema)
+        schemas = await self._get_schemas()
+        total_count = 0
 
-        result = await self._pool.execute(
-            f"""
-            UPDATE {table}
-            SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
-            WHERE status = 'processing' AND worker_id = $1
-            """,
-            self._worker_id,
-        )
+        for schema in schemas:
+            table = fq_table("async_operations", schema)
 
-        # Parse "UPDATE N" to get count
-        count = int(result.split()[-1]) if result else 0
-        if count > 0:
-            logger.info(f"Worker {self._worker_id} recovered {count} stale tasks from previous run")
-        return count
+            result = await self._pool.execute(
+                f"""
+                UPDATE {table}
+                SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
+                WHERE status = 'processing' AND worker_id = $1
+                """,
+                self._worker_id,
+            )
+
+            # Parse "UPDATE N" to get count
+            count = int(result.split()[-1]) if result else 0
+            total_count += count
+
+        if total_count > 0:
+            logger.info(f"Worker {self._worker_id} recovered {total_count} stale tasks from previous run")
+        return total_count
 
     async def run(self):
         """
@@ -245,6 +315,8 @@ class WorkerPoller:
 
         Continuously polls for pending tasks, claims them, and executes them
         until shutdown is signaled.
+
+        If tenant_extension is configured, dynamically discovers schemas on each poll.
         """
         # Recover any tasks from a previous crash before starting
         await self.recover_own_tasks()
@@ -253,17 +325,22 @@ class WorkerPoller:
 
         while not self._shutdown.is_set():
             try:
-                # Claim a batch of tasks
+                # Claim a batch of tasks (across all tenant schemas if configured)
                 tasks = await self.claim_batch()
 
                 if tasks:
                     # Log batch info
-                    task_types = {}
-                    for _, task_dict in tasks:
-                        t = task_dict.get("type", "unknown")
+                    task_types: dict[str, int] = {}
+                    schemas_seen: set[str | None] = set()
+                    for task in tasks:
+                        t = task.task_dict.get("type", "unknown")
                         task_types[t] = task_types.get(t, 0) + 1
+                        schemas_seen.add(task.schema)
                     types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items())
-                    logger.info(f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str}")
+                    schemas_str = ", ".join(s or "public" for s in schemas_seen)
+                    logger.info(
+                        f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str} (schemas: {schemas_str})"
+                    )
 
                     # Track in-flight tasks
                     async with self._in_flight_lock:
@@ -272,7 +349,7 @@ class WorkerPoller:
                     # Execute tasks concurrently
                     try:
                         await asyncio.gather(
-                            *[self.execute_task(op_id, task_dict) for op_id, task_dict in tasks],
+                            *[self.execute_task(task) for task in tasks],
                             return_exceptions=True,
                         )
                     finally:
@@ -336,58 +413,60 @@ class WorkerPoller:
         self._last_progress_log = now
 
         try:
-            table = fq_table("async_operations", self._schema)
-            async with self._pool.acquire() as conn:
-                # Get global stats by status
-                stats = await conn.fetch(
-                    f"""
-                    SELECT status, COUNT(*) as count
-                    FROM {table}
-                    WHERE created_at > now() - interval '24 hours'
-                    GROUP BY status
-                    """
-                )
-
-                # Get currently processing tasks grouped by type and bank
-                processing = await conn.fetch(
-                    f"""
-                    SELECT operation_type, bank_id, COUNT(*) as count
-                    FROM {table}
-                    WHERE status = 'processing'
-                    GROUP BY operation_type, bank_id
-                    """
-                )
-
-            # Build stats dict
-            status_counts = {row["status"]: row["count"] for row in stats}
-            pending = status_counts.get("pending", 0)
-            processing_count = status_counts.get("processing", 0)
-            completed = status_counts.get("completed", 0)
-            failed = status_counts.get("failed", 0)
-
-            # Build processing breakdown
-            processing_info = []
-            banks_working = set()
-            for row in processing:
-                op_type = row["operation_type"]
-                bank_id = row["bank_id"]
-                count = row["count"]
-                banks_working.add(bank_id)
-                processing_info.append(f"{op_type}:{bank_id}({count})")
-
-            # Format log
+            # Get local active tasks (this worker only)
             async with self._in_flight_lock:
                 in_flight = self._in_flight_count
+                active_tasks = dict(self._active_tasks)  # Copy to avoid holding lock
 
+            # Build local processing breakdown grouped by (op_type, bank_id)
+            task_groups: dict[tuple[str, str], int] = {}
+            for op_type, bank_id, _ in active_tasks.values():
+                key = (op_type, bank_id)
+                task_groups[key] = task_groups.get(key, 0) + 1
+
+            processing_info = [f"{op}:{bank}({cnt})" for (op, bank), cnt in task_groups.items()]
             processing_str = ", ".join(processing_info[:10]) if processing_info else "none"
             if len(processing_info) > 10:
                 processing_str += f" +{len(processing_info) - 10} more"
 
+            # Get global stats from DB across all schemas
+            schemas = await self._get_schemas()
+            global_pending = 0
+            all_worker_counts: dict[str, int] = {}
+
+            async with self._pool.acquire() as conn:
+                for schema in schemas:
+                    table = fq_table("async_operations", schema)
+
+                    row = await conn.fetchrow(f"SELECT COUNT(*) as count FROM {table} WHERE status = 'pending'")
+                    global_pending += row["count"] if row else 0
+
+                    # Get processing breakdown by worker
+                    worker_rows = await conn.fetch(
+                        f"""
+                        SELECT worker_id, COUNT(*) as count
+                        FROM {table}
+                        WHERE status = 'processing'
+                        GROUP BY worker_id
+                        """
+                    )
+                    for wr in worker_rows:
+                        wid = wr["worker_id"] or "unknown"
+                        all_worker_counts[wid] = all_worker_counts.get(wid, 0) + wr["count"]
+
+            # Format other workers' processing counts
+            other_workers = []
+            for wid, cnt in all_worker_counts.items():
+                if wid != self._worker_id:
+                    other_workers.append(f"{wid}:{cnt}")
+            others_str = ", ".join(other_workers) if other_workers else "none"
+
+            schemas_str = ", ".join(s or "public" for s in schemas)
             logger.info(
                 f"[WORKER_STATS] worker={self._worker_id} in_flight={in_flight} | "
-                f"global: pending={pending} processing={processing_count} "
-                f"completed_24h={completed} failed_24h={failed} | "
-                f"active: {processing_str}"
+                f"global: pending={global_pending} (schemas: {schemas_str}) | "
+                f"others: {others_str} | "
+                f"my_active: {processing_str}"
             )
 
         except Exception as e:

hindsight-api/tests/test_custom_embedding_dimension.py

@@ -9,18 +9,18 @@ Includes tests for:
 
 import asyncio
 import os
-import pytest
 from datetime import datetime
+
+import pytest
 from sqlalchemy import create_engine, text
 
 from hindsight_api import MemoryEngine, RequestContext
-from hindsight_api.engine.embeddings import LocalSTEmbeddings, OpenAIEmbeddings, CohereEmbeddings
-from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder, CohereCrossEncoder
+from hindsight_api.engine.cross_encoder import CohereCrossEncoder, LocalSTCrossEncoder
+from hindsight_api.engine.embeddings import CohereEmbeddings, LocalSTEmbeddings, OpenAIEmbeddings
 from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
 from hindsight_api.engine.task_backend import SyncTaskBackend
-from hindsight_api.extensions import TenantExtension, TenantContext
-from hindsight_api.migrations import run_migrations, ensure_embedding_dimension
-
+from hindsight_api.extensions import TenantContext, TenantExtension
+from hindsight_api.migrations import ensure_embedding_dimension, run_migrations
 
 # =============================================================================
 # Shared Utilities
@@ -36,6 +36,11 @@ class SchemaTenantExtension(TenantExtension):
     async def authenticate(self, request_context: RequestContext) -> TenantContext:
         return TenantContext(schema_name=self.schema_name)
 
+    async def list_tenants(self) -> list:
+        from hindsight_api.extensions.tenant import Tenant
+
+        return [Tenant(schema=self.schema_name)]
+
 
 def get_test_schema(prefix: str, worker_id: str) -> str:
     """Get unique schema name per xdist worker."""

hindsight-api/tests/test_extensions.py

@@ -24,6 +24,9 @@ from hindsight_api.extensions import (
     TenantExtension,
     ValidationResult,
     load_extension,
+    # Consolidation operation
+    ConsolidateContext,
+    ConsolidateResult,
 )
 
 
@@ -128,14 +131,18 @@ class TrackingValidator(OperationValidatorExtension):
 
     def __init__(self, config: dict):
         super().__init__(config)
-        # Pre-hook tracking
+        # Pre-hook tracking - Core operations
         self.pre_retain_calls: list[RetainContext] = []
         self.pre_recall_calls: list[RecallContext] = []
         self.pre_reflect_calls: list[ReflectContext] = []
-        # Post-hook tracking
+        # Post-hook tracking - Core operations
         self.post_retain_calls: list[RetainResult] = []
         self.post_recall_calls: list[RecallResult] = []
         self.post_reflect_calls: list[ReflectResultContext] = []
+        # Pre-hook tracking - Consolidation
+        self.pre_consolidate_calls: list[ConsolidateContext] = []
+        # Post-hook tracking - Consolidation
+        self.post_consolidate_calls: list[ConsolidateResult] = []
 
     async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
         self.pre_retain_calls.append(ctx)
@@ -158,6 +165,14 @@ class TrackingValidator(OperationValidatorExtension):
     async def on_reflect_complete(self, result: ReflectResultContext) -> None:
         self.post_reflect_calls.append(result)
 
+    # Consolidation hooks
+    async def validate_consolidate(self, ctx: ConsolidateContext) -> ValidationResult:
+        self.pre_consolidate_calls.append(ctx)
+        return ValidationResult.accept()
+
+    async def on_consolidate_complete(self, result: ConsolidateResult) -> None:
+        self.post_consolidate_calls.append(result)
+
 
 class TestMemoryEngineValidation:
     """Tests for validation integration with MemoryEngine.

hindsight-api/tests/test_http_api_integration.py

@@ -969,24 +969,22 @@ async def test_reflect_returns_token_usage(api_client):
     assert "text" in result
     assert len(result["text"]) > 0
 
-    # Verify usage field exists (may be None for agentic reflect which makes multiple LLM calls)
+    # Verify usage field exists and is populated (agentic reflect aggregates all LLM calls)
     assert "usage" in result, "Response should include 'usage' field"
     usage = result["usage"]
 
-    # Usage is optional - agentic reflect doesn't aggregate multiple LLM call usages
-    if usage is not None:
-        assert "input_tokens" in usage, "Usage should have 'input_tokens'"
-        assert "output_tokens" in usage, "Usage should have 'output_tokens'"
-        assert "total_tokens" in usage, "Usage should have 'total_tokens'"
+    # Usage must be present - agentic reflect now aggregates token usage from all LLM calls
+    assert usage is not None, "Usage should not be None - reflect aggregates all LLM call usages"
+    assert "input_tokens" in usage, "Usage should have 'input_tokens'"
+    assert "output_tokens" in usage, "Usage should have 'output_tokens'"
+    assert "total_tokens" in usage, "Usage should have 'total_tokens'"
 
-        # Verify token counts are valid
-        assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}"
-        assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
-        assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
+    # Verify token counts are valid
+    assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}"
+    assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
+    assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
 
-        print(f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
-    else:
-        print("Reflect usage is None (expected for agentic reflect)")
+    print(f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
 
 
 @pytest.mark.asyncio

hindsight-api/tests/test_main_module.py

@@ -355,14 +355,14 @@ class TestMainModuleExtensionLoading:
 
 # Mock extensions for testing
 from hindsight_api.extensions import (
-    TenantExtension,
-    TenantContext,
-    RequestContext,
     OperationValidatorExtension,
-    ValidationResult,
-    RetainContext,
     RecallContext,
     ReflectContext,
+    RequestContext,
+    RetainContext,
+    TenantContext,
+    TenantExtension,
+    ValidationResult,
 )
 
 
@@ -376,6 +376,11 @@ class MockTenantExtension(TenantExtension):
     async def authenticate(self, request_context: RequestContext) -> TenantContext:
         return TenantContext(schema_name="public")
 
+    async def list_tenants(self) -> list:
+        from hindsight_api.extensions.tenant import Tenant
+
+        return [Tenant(schema="public")]
+
     def set_context(self, context) -> None:
         self._context_set = True
 

hindsight-api/tests/test_reflect_agent.py

@@ -13,9 +13,52 @@ from unittest.mock import AsyncMock, MagicMock, patch
 from hindsight_api.engine.reflect.agent import (
     _normalize_tool_name,
     _is_done_tool,
+    _clean_answer_text,
     run_reflect_agent,
 )
-from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult
+from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
+
+
+class TestCleanAnswerText:
+    """Test cleanup of answer text that includes done() tool call syntax."""
+
+    def test_clean_text_with_done_call(self):
+        """Text ending with done() call should have it stripped."""
+        text = '''The team's OKRs focus on performance.done({"answer":"The team's OKRs","memory_ids":[]})'''
+        cleaned = _clean_answer_text(text)
+        assert cleaned == "The team's OKRs focus on performance."
+        assert "done(" not in cleaned
+
+    def test_clean_text_with_done_call_and_whitespace(self):
+        """done() call with whitespace should be stripped."""
+        text = '''Answer text here. done( {"answer": "short", "memory_ids": []} )'''
+        cleaned = _clean_answer_text(text)
+        assert cleaned == "Answer text here."
+
+    def test_clean_text_without_done_call(self):
+        """Text without done() call should be unchanged."""
+        text = "This is a normal answer without any tool calls."
+        cleaned = _clean_answer_text(text)
+        assert cleaned == text
+
+    def test_clean_text_with_done_word_in_content(self):
+        """The word 'done' in regular text should not be stripped."""
+        text = "The task is done and completed successfully."
+        cleaned = _clean_answer_text(text)
+        assert cleaned == text
+
+    def test_clean_empty_text(self):
+        """Empty text should return empty."""
+        assert _clean_answer_text("") == ""
+
+    def test_clean_text_multiline_done(self):
+        """done() call spanning multiple lines should be stripped."""
+        text = '''Summary of findings.done({
+            "answer": "Summary",
+            "memory_ids": ["id1", "id2"]
+        })'''
+        cleaned = _clean_answer_text(text)
+        assert cleaned == "Summary of findings."
 
 
 class TestToolNameNormalization:
@@ -70,8 +113,10 @@ class TestReflectAgentMocked:
         """Create a mock LLM provider."""
         llm = MagicMock()
         llm.call_with_tools = AsyncMock()
-        # Also mock call() for final iteration fallback
-        llm.call = AsyncMock(return_value="Fallback answer from final iteration")
+        # Also mock call() for final iteration fallback - returns (response, usage) tuple
+        llm.call = AsyncMock(
+            return_value=("Fallback answer from final iteration", TokenUsage(input_tokens=100, output_tokens=50, total_tokens=150))
+        )
         return llm
 
     @pytest.fixture

hindsight-api/tests/test_schema_isolation.py

@@ -11,8 +11,8 @@ import uuid
 import pytest
 import pytest_asyncio
 
-from hindsight_api.extensions import RequestContext, TenantContext, TenantExtension
 from hindsight_api.engine.memory_engine import _current_schema, fq_table
+from hindsight_api.extensions import RequestContext, TenantContext, TenantExtension
 from hindsight_api.migrations import run_migrations
 
 
@@ -52,6 +52,11 @@ class MultiSchemaTestTenantExtension(TenantExtension):
 
         raise AuthenticationError(f"Unknown API key: {context.api_key}")
 
+    async def list_tenants(self) -> list:
+        from hindsight_api.extensions.tenant import Tenant
+
+        return [Tenant(schema=schema) for schema in self.valid_schemas]
+
 
 async def drop_schema(conn, schema_name: str) -> None:
     """Drop a schema and all its contents."""

hindsight-api/tests/test_server_module.py

@@ -249,14 +249,14 @@ class TestServerModuleExtensionLoading:
 
 # Mock extensions for testing
 from hindsight_api.extensions import (
-    TenantExtension,
-    TenantContext,
-    RequestContext,
     OperationValidatorExtension,
-    ValidationResult,
-    RetainContext,
     RecallContext,
     ReflectContext,
+    RequestContext,
+    RetainContext,
+    TenantContext,
+    TenantExtension,
+    ValidationResult,
 )
 
 
@@ -270,6 +270,11 @@ class MockTenantExtension(TenantExtension):
     async def authenticate(self, request_context: RequestContext) -> TenantContext:
         return TenantContext(schema_name="public")
 
+    async def list_tenants(self) -> list:
+        from hindsight_api.extensions.tenant import Tenant
+
+        return [Tenant(schema="public")]
+
     def set_context(self, context) -> None:
         self._context_set = True
 

hindsight-api/tests/test_worker.py

@@ -162,6 +162,11 @@ class TestWorkerPoller:
         claimed = await poller.claim_batch()
         assert len(claimed) == 3
 
+        # ClaimedTask objects have operation_id, task_dict, schema attributes
+        for task in claimed:
+            assert task.operation_id is not None
+            assert task.task_dict is not None
+
         # Verify tasks are marked as processing with worker_id
         rows = await pool.fetch(
             "SELECT status, worker_id FROM async_operations WHERE bank_id = $1",
@@ -206,6 +211,7 @@ class TestWorkerPoller:
     async def test_execute_task_marks_completed(self, pool, clean_operations):
         """Test that successful task execution marks task as completed."""
         from hindsight_api.worker import WorkerPoller
+        from hindsight_api.worker.poller import ClaimedTask
 
         # Create a pending task
         bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
@@ -234,7 +240,8 @@ class TestWorkerPoller:
 
         # Execute the task
         task_dict = json.loads(payload)
-        await poller.execute_task(str(op_id), task_dict)
+        claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
+        await poller.execute_task(claimed_task)
 
         assert len(executed) == 1
 
@@ -250,6 +257,7 @@ class TestWorkerPoller:
     async def test_execute_task_retries_on_failure(self, pool, clean_operations):
         """Test that failed task execution triggers retry mechanism."""
         from hindsight_api.worker import WorkerPoller
+        from hindsight_api.worker.poller import ClaimedTask
 
         # Create a pending task with retry_count=0
         bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
@@ -277,7 +285,8 @@ class TestWorkerPoller:
 
         # Execute (should fail and retry)
         task_dict = json.loads(payload)
-        await poller.execute_task(str(op_id), task_dict)
+        claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
+        await poller.execute_task(claimed_task)
 
         # Verify task is back to pending with incremented retry_count
         row = await pool.fetchrow(
@@ -292,6 +301,7 @@ class TestWorkerPoller:
     async def test_execute_task_fails_after_max_retries(self, pool, clean_operations):
         """Test that task is marked failed after exceeding max retries."""
         from hindsight_api.worker import WorkerPoller
+        from hindsight_api.worker.poller import ClaimedTask
 
         # Create a task that has already used all retries
         bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
@@ -319,7 +329,8 @@ class TestWorkerPoller:
 
         # Execute (should fail permanently)
         task_dict = json.loads(payload)
-        await poller.execute_task(str(op_id), task_dict)
+        claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
+        await poller.execute_task(claimed_task)
 
         # Verify task is marked as failed
         row = await pool.fetchrow(
@@ -384,9 +395,8 @@ class TestWorkerPoller:
 
         # Should only claim the consolidation for the other bank
         assert len(claimed) == 1
-        claimed_op_id, claimed_payload = claimed[0]
-        assert claimed_op_id == str(other_op_id)
-        assert claimed_payload["bank_id"] == other_bank_id
+        assert claimed[0].operation_id == str(other_op_id)
+        assert claimed[0].task_dict["bank_id"] == other_bank_id
 
         # Verify the pending consolidation for first bank is still pending
         row = await pool.fetchrow(
@@ -437,8 +447,7 @@ class TestWorkerPoller:
 
         # Should claim the retain task (non-consolidation tasks are unaffected)
         assert len(claimed) == 1
-        claimed_op_id, _ = claimed[0]
-        assert claimed_op_id == str(retain_op_id)
+        assert claimed[0].operation_id == str(retain_op_id)
 
 
 class TestWorkerRecovery:
@@ -601,7 +610,7 @@ class TestConcurrentWorkers:
                 batch_size=5,  # Each worker tries to claim 5
             )
             claimed = await poller.claim_batch()
-            workers_claimed[worker_id] = [op_id for op_id, _ in claimed]
+            workers_claimed[worker_id] = [task.operation_id for task in claimed]
 
         # Run all workers concurrently
         await asyncio.gather(
@@ -825,3 +834,186 @@ class TestSyncTaskBackend:
 
         # Should not raise, error is logged
         await backend.submit_task({"type": "test"})
+
+
+class TestDynamicTenantDiscovery:
+    """Tests for dynamic tenant discovery via TenantExtension."""
+
+    @pytest.mark.asyncio
+    async def test_poller_discovers_tenants_dynamically(self, pool, clean_operations):
+        """Test that poller calls list_tenants() on each poll cycle."""
+        from hindsight_api.extensions.tenant import Tenant, TenantExtension
+        from hindsight_api.worker import WorkerPoller
+
+        # Create a mock tenant extension that tracks calls
+        class MockTenantExtension(TenantExtension):
+            def __init__(self):
+                self.list_tenants_calls = 0
+                self.tenants_to_return: list[Tenant] = [Tenant(schema="public")]
+
+            async def authenticate(self, context):
+                raise NotImplementedError("Not used in this test")
+
+            async def list_tenants(self) -> list[Tenant]:
+                self.list_tenants_calls += 1
+                return self.tenants_to_return
+
+        mock_extension = MockTenantExtension()
+
+        # Create pending tasks in public schema
+        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
+        for i in range(2):
+            op_id = uuid.uuid4()
+            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
+            await pool.execute(
+                """
+                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
+                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
+                """,
+                op_id,
+                bank_id,
+                payload,
+            )
+
+        poller = WorkerPoller(
+            pool=pool,
+            worker_id="test-worker-1",
+            executor=lambda x: None,
+            batch_size=10,
+            tenant_extension=mock_extension,
+        )
+
+        # First claim_batch should call list_tenants
+        claimed1 = await poller.claim_batch()
+        assert mock_extension.list_tenants_calls == 1
+        assert len(claimed1) == 2
+
+        # Add more tasks
+        for i in range(2):
+            op_id = uuid.uuid4()
+            payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id})
+            await pool.execute(
+                """
+                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
+                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
+                """,
+                op_id,
+                bank_id,
+                payload,
+            )
+
+        # Second claim_batch should call list_tenants again
+        claimed2 = await poller.claim_batch()
+        assert mock_extension.list_tenants_calls == 2
+        assert len(claimed2) == 2
+
+    @pytest.mark.asyncio
+    async def test_poller_picks_up_new_tenants_without_restart(self, pool, clean_operations):
+        """Test that new tenants are discovered on subsequent poll cycles."""
+        from hindsight_api.extensions.tenant import Tenant, TenantExtension
+        from hindsight_api.worker import WorkerPoller
+
+        class DynamicTenantExtension(TenantExtension):
+            def __init__(self):
+                # Start with just public
+                self.tenants: list[Tenant] = [Tenant(schema="public")]
+                self.list_tenants_calls = 0
+
+            async def authenticate(self, context):
+                raise NotImplementedError("Not used in this test")
+
+            async def list_tenants(self) -> list[Tenant]:
+                self.list_tenants_calls += 1
+                return self.tenants
+
+        dynamic_extension = DynamicTenantExtension()
+
+        # Create a task in public schema
+        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
+        op_id = uuid.uuid4()
+        payload = json.dumps({"type": "test_task", "bank_id": bank_id})
+        await pool.execute(
+            """
+            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
+            VALUES ($1, $2, 'test', 'pending', $3::jsonb)
+            """,
+            op_id,
+            bank_id,
+            payload,
+        )
+
+        poller = WorkerPoller(
+            pool=pool,
+            worker_id="test-worker-1",
+            executor=lambda x: None,
+            batch_size=10,
+            tenant_extension=dynamic_extension,
+        )
+
+        # First poll - only public schema
+        claimed1 = await poller.claim_batch()
+        assert len(claimed1) == 1
+        assert claimed1[0].schema is None  # public is represented as None
+        assert dynamic_extension.list_tenants_calls == 1
+
+        # Simulate tenant list changing (but we won't add a non-existent schema)
+        # In real world, the schema would be created before list_tenants returns it
+        # Here we just verify that list_tenants is called again
+
+        # Add another task to public
+        op_id2 = uuid.uuid4()
+        payload2 = json.dumps({"type": "test_task", "bank_id": bank_id})
+        await pool.execute(
+            """
+            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
+            VALUES ($1, $2, 'test', 'pending', $3::jsonb)
+            """,
+            op_id2,
+            bank_id,
+            payload2,
+        )
+
+        # Second poll - list_tenants should be called again
+        claimed2 = await poller.claim_batch()
+        assert len(claimed2) == 1
+        assert dynamic_extension.list_tenants_calls == 2  # Called again on second poll
+
+        # Third poll with no tasks - still calls list_tenants
+        claimed3 = await poller.claim_batch()
+        assert len(claimed3) == 0
+        assert dynamic_extension.list_tenants_calls == 3  # Called again even with no tasks
+
+    @pytest.mark.asyncio
+    async def test_poller_without_tenant_extension_uses_public(self, pool, clean_operations):
+        """Test that poller uses public schema when no tenant extension is configured."""
+        from hindsight_api.worker import WorkerPoller
+
+        # Create pending tasks
+        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
+        for i in range(3):
+            op_id = uuid.uuid4()
+            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
+            await pool.execute(
+                """
+                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
+                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
+                """,
+                op_id,
+                bank_id,
+                payload,
+            )
+
+        # No tenant_extension provided
+        poller = WorkerPoller(
+            pool=pool,
+            worker_id="test-worker-1",
+            executor=lambda x: None,
+            batch_size=10,
+        )
+
+        claimed = await poller.claim_batch()
+        assert len(claimed) == 3
+
+        # All tasks should have schema=None (public)
+        for task in claimed:
+            assert task.schema is None

hindsight-cli/src/api.rs

@@ -485,7 +485,7 @@ impl ApiClient {
         })
     }
 
-    pub fn refresh_reflection(&self, bank_id: &str, reflection_id: &str, _verbose: bool) -> Result<types::ReflectionResponse> {
+    pub fn refresh_reflection(&self, bank_id: &str, reflection_id: &str, _verbose: bool) -> Result<types::AsyncOperationSubmitResponse> {
         self.runtime.block_on(async {
             let response = self.client.refresh_reflection(bank_id, reflection_id, None).await?;
             Ok(response.into_inner())

hindsight-cli/src/commands/reflection.rs

@@ -233,7 +233,7 @@ pub fn refresh(
     output_format: OutputFormat,
 ) -> Result<()> {
     let spinner = if output_format == OutputFormat::Pretty {
-        Some(ui::create_spinner("Refreshing reflection..."))
+        Some(ui::create_spinner("Submitting reflection refresh..."))
     } else {
         None
     };
@@ -245,13 +245,17 @@ pub fn refresh(
     }
 
     match response {
-        Ok(reflection) => {
+        Ok(operation) => {
             if output_format == OutputFormat::Pretty {
-                ui::print_success(&format!("Reflection '{}' refreshed successfully", reflection_id));
+                ui::print_success(&format!(
+                    "Reflection refresh submitted. Operation ID: {}",
+                    operation.operation_id
+                ));
+                println!("  {} {}", ui::dim("Status:"), operation.status);
                 println!();
-                print_reflection_detail(&reflection);
+                println!("{}", ui::dim("Use 'hindsight operations get' to check the operation status."));
             } else {
-                output::print_output(&reflection, output_format)?;
+                output::print_output(&operation, output_format)?;
             }
             Ok(())
         }

hindsight-clients/python/.openapi-generator/FILES

@@ -14,6 +14,7 @@ hindsight_client_api/configuration.py
 hindsight_client_api/exceptions.py
 hindsight_client_api/models/__init__.py
 hindsight_client_api/models/add_background_request.py
+hindsight_client_api/models/async_operation_submit_response.py
 hindsight_client_api/models/background_response.py
 hindsight_client_api/models/bank_list_item.py
 hindsight_client_api/models/bank_list_response.py

hindsight-clients/python/hindsight_client_api/__init__.py

@@ -39,6 +39,7 @@ from hindsight_client_api.exceptions import ApiException
 
 # import models into sdk package
 from hindsight_client_api.models.add_background_request import AddBackgroundRequest
+from hindsight_client_api.models.async_operation_submit_response import AsyncOperationSubmitResponse
 from hindsight_client_api.models.background_response import BackgroundResponse
 from hindsight_client_api.models.bank_list_item import BankListItem
 from hindsight_client_api.models.bank_list_response import BankListResponse

hindsight-clients/python/hindsight_client_api/api/reflections_api.py

@@ -19,6 +19,7 @@ from typing_extensions import Annotated
 from pydantic import Field, StrictStr, field_validator
 from typing import Any, List, Optional
 from typing_extensions import Annotated
+from hindsight_client_api.models.async_operation_submit_response import AsyncOperationSubmitResponse
 from hindsight_client_api.models.create_reflection_request import CreateReflectionRequest
 from hindsight_client_api.models.create_reflection_response import CreateReflectionResponse
 from hindsight_client_api.models.reflection_list_response import ReflectionListResponse
@@ -1300,10 +1301,10 @@ class ReflectionsApi:
         _content_type: Optional[StrictStr] = None,
         _headers: Optional[Dict[StrictStr, Any]] = None,
         _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
-    ) -> ReflectionResponse:
+    ) -> AsyncOperationSubmitResponse:
         """Refresh reflection
 
-        Re-run the source query through reflect and update the content.
+        Submit an async task to re-run the source query through reflect and update the content.
 
         :param bank_id: (required)
         :type bank_id: str
@@ -1344,7 +1345,7 @@ class ReflectionsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '200': "ReflectionResponse",
+            '200': "AsyncOperationSubmitResponse",
             '422': "HTTPValidationError",
         }
         response_data = await self.api_client.call_api(
@@ -1376,10 +1377,10 @@ class ReflectionsApi:
         _content_type: Optional[StrictStr] = None,
         _headers: Optional[Dict[StrictStr, Any]] = None,
         _host_index: Annotated[StrictInt, Field(ge=0, le=0)] = 0,
-    ) -> ApiResponse[ReflectionResponse]:
+    ) -> ApiResponse[AsyncOperationSubmitResponse]:
         """Refresh reflection
 
-        Re-run the source query through reflect and update the content.
+        Submit an async task to re-run the source query through reflect and update the content.
 
         :param bank_id: (required)
         :type bank_id: str
@@ -1420,7 +1421,7 @@ class ReflectionsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '200': "ReflectionResponse",
+            '200': "AsyncOperationSubmitResponse",
             '422': "HTTPValidationError",
         }
         response_data = await self.api_client.call_api(
@@ -1455,7 +1456,7 @@ class ReflectionsApi:
     ) -> RESTResponseType:
         """Refresh reflection
 
-        Re-run the source query through reflect and update the content.
+        Submit an async task to re-run the source query through reflect and update the content.
 
         :param bank_id: (required)
         :type bank_id: str
@@ -1496,7 +1497,7 @@ class ReflectionsApi:
         )
 
         _response_types_map: Dict[str, Optional[str]] = {
-            '200': "ReflectionResponse",
+            '200': "AsyncOperationSubmitResponse",
             '422': "HTTPValidationError",
         }
         response_data = await self.api_client.call_api(

hindsight-clients/python/hindsight_client_api/models/__init__.py

@@ -15,6 +15,7 @@
 
 # import models into model package
 from hindsight_client_api.models.add_background_request import AddBackgroundRequest
+from hindsight_client_api.models.async_operation_submit_response import AsyncOperationSubmitResponse
 from hindsight_client_api.models.background_response import BackgroundResponse
 from hindsight_client_api.models.bank_list_item import BankListItem
 from hindsight_client_api.models.bank_list_response import BankListResponse

hindsight-clients/python/hindsight_client_api/models/async_operation_submit_response.py

@@ -0,0 +1,89 @@
+# coding: utf-8
+
+"""
+    Hindsight HTTP API
+
+    HTTP API for Hindsight
+
+    The version of the OpenAPI document: 0.1.0
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+from __future__ import annotations
+import pprint
+import re  # noqa: F401
+import json
+
+from pydantic import BaseModel, ConfigDict, StrictStr
+from typing import Any, ClassVar, Dict, List
+from typing import Optional, Set
+from typing_extensions import Self
+
+class AsyncOperationSubmitResponse(BaseModel):
+    """
+    Response model for submitting an async operation.
+    """ # noqa: E501
+    operation_id: StrictStr
+    status: StrictStr
+    __properties: ClassVar[List[str]] = ["operation_id", "status"]
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+        validate_assignment=True,
+        protected_namespaces=(),
+    )
+
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.model_dump(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Optional[Self]:
+        """Create an instance of AsyncOperationSubmitResponse from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Return the dictionary representation of the model using alias.
+
+        This has the following differences from calling pydantic's
+        `self.model_dump(by_alias=True)`:
+
+        * `None` is only added to the output dict for nullable fields that
+          were set at model initialization. Other fields with value `None`
+          are ignored.
+        """
+        excluded_fields: Set[str] = set([
+        ])
+
+        _dict = self.model_dump(
+            by_alias=True,
+            exclude=excluded_fields,
+            exclude_none=True,
+        )
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
+        """Create an instance of AsyncOperationSubmitResponse from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return cls.model_validate(obj)
+
+        _obj = cls.model_validate({
+            "operation_id": obj.get("operation_id"),
+            "status": obj.get("status")
+        })
+        return _obj
+
+

hindsight-clients/python/hindsight_client_api/models/consolidation_response.py

@@ -17,8 +17,8 @@ import pprint
 import re  # noqa: F401
 import json
 
-from pydantic import BaseModel, ConfigDict, Field, StrictInt, StrictStr
-from typing import Any, ClassVar, Dict, List
+from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictStr
+from typing import Any, ClassVar, Dict, List, Optional
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -26,12 +26,9 @@ class ConsolidationResponse(BaseModel):
     """
     Response model for consolidation trigger endpoint.
     """ # noqa: E501
-    status: StrictStr = Field(description="Status of the consolidation (completed or queued)")
-    processed: StrictInt = Field(description="Number of memories processed")
-    created: StrictInt = Field(description="Number of mental models created")
-    updated: StrictInt = Field(description="Number of mental models updated")
-    message: StrictStr = Field(description="Human-readable summary")
-    __properties: ClassVar[List[str]] = ["status", "processed", "created", "updated", "message"]
+    operation_id: StrictStr = Field(description="ID of the async consolidation operation")
+    deduplicated: Optional[StrictBool] = Field(default=False, description="True if an existing pending task was reused")
+    __properties: ClassVar[List[str]] = ["operation_id", "deduplicated"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -84,11 +81,8 @@ class ConsolidationResponse(BaseModel):
             return cls.model_validate(obj)
 
         _obj = cls.model_validate({
-            "status": obj.get("status"),
-            "processed": obj.get("processed"),
-            "created": obj.get("created"),
-            "updated": obj.get("updated"),
-            "message": obj.get("message")
+            "operation_id": obj.get("operation_id"),
+            "deduplicated": obj.get("deduplicated") if obj.get("deduplicated") is not None else False
         })
         return _obj
 

hindsight-clients/typescript/generated/sdk.gen.ts

@@ -455,7 +455,7 @@ export const updateReflection = <ThrowOnError extends boolean = false>(
 /**
  * Refresh reflection
  *
- * Re-run the source query through reflect and update the content.
+ * Submit an async task to re-run the source query through reflect and update the content.
  */
 export const refreshReflection = <ThrowOnError extends boolean = false>(
   options: Options<RefreshReflectionData, ThrowOnError>,

hindsight-clients/typescript/generated/types.gen.ts

@@ -24,6 +24,22 @@ export type AddBackgroundRequest = {
   update_disposition?: boolean;
 };
 
+/**
+ * AsyncOperationSubmitResponse
+ *
+ * Response model for submitting an async operation.
+ */
+export type AsyncOperationSubmitResponse = {
+  /**
+   * Operation Id
+   */
+  operation_id: string;
+  /**
+   * Status
+   */
+  status: string;
+};
+
 /**
  * BackgroundResponse
  *
@@ -295,35 +311,17 @@ export type ChunkResponse = {
  */
 export type ConsolidationResponse = {
   /**
-   * Status
+   * Operation Id
    *
-   * Status of the consolidation (completed or queued)
+   * ID of the async consolidation operation
    */
-  status: string;
+  operation_id: string;
   /**
-   * Processed
+   * Deduplicated
    *
-   * Number of memories processed
+   * True if an existing pending task was reused
    */
-  processed: number;
-  /**
-   * Created
-   *
-   * Number of mental models created
-   */
-  created: number;
-  /**
-   * Updated
-   *
-   * Number of mental models updated
-   */
-  updated: number;
-  /**
-   * Message
-   *
-   * Human-readable summary
-   */
-  message: string;
+  deduplicated?: boolean;
 };
 
 /**
@@ -2486,7 +2484,7 @@ export type RefreshReflectionResponses = {
   /**
    * Successful Response
    */
-  200: ReflectionResponse;
+  200: AsyncOperationSubmitResponse;
 };
 
 export type RefreshReflectionResponse =

hindsight-control-plane/src/app/banks/[bankId]/page.tsx

@@ -10,6 +10,7 @@ import { ThinkView } from "@/components/think-view";
 import { SearchDebugView } from "@/components/search-debug-view";
 import { BankProfileView } from "@/components/bank-profile-view";
 import { ReflectionsView } from "@/components/reflections-view";
+import { useFeatures } from "@/lib/features-context";
 
 type NavItem = "recall" | "reflect" | "data" | "documents" | "entities" | "profile";
 type DataSubTab = "world" | "experience" | "models" | "reflections";
@@ -18,10 +19,12 @@ export default function BankPage() {
   const params = useParams();
   const router = useRouter();
   const searchParams = useSearchParams();
+  const { features } = useFeatures();
 
   const bankId = params.bankId as string;
   const view = (searchParams.get("view") || "profile") as NavItem;
   const subTab = (searchParams.get("subTab") || "world") as DataSubTab;
+  const mentalModelsEnabled = features?.mental_models ?? false;
 
   const handleTabChange = (tab: NavItem) => {
     router.push(`/banks/${bankId}?view=${tab}`);
@@ -120,6 +123,11 @@ export default function BankPage() {
                       }`}
                     >
                       Mental Models
+                      {!mentalModelsEnabled && (
+                        <span className="ml-2 text-xs px-1.5 py-0.5 rounded bg-muted text-muted-foreground">
+                          Off
+                        </span>
+                      )}
                       {subTab === "models" && (
                         <div className="absolute bottom-0 left-0 right-0 h-0.5 bg-primary" />
                       )}
@@ -143,7 +151,40 @@ export default function BankPage() {
                 <div>
                   {subTab === "world" && <DataView key="world" factType="world" />}
                   {subTab === "experience" && <DataView key="experience" factType="experience" />}
-                  {subTab === "models" && <DataView key="models" factType="mental_model" />}
+                  {subTab === "models" &&
+                    (mentalModelsEnabled ? (
+                      <DataView key="models" factType="mental_model" />
+                    ) : (
+                      <div className="flex flex-col items-center justify-center py-16 text-center">
+                        <div className="text-muted-foreground mb-2">
+                          <svg
+                            xmlns="http://www.w3.org/2000/svg"
+                            width="48"
+                            height="48"
+                            viewBox="0 0 24 24"
+                            fill="none"
+                            stroke="currentColor"
+                            strokeWidth="1.5"
+                            strokeLinecap="round"
+                            strokeLinejoin="round"
+                          >
+                            <path d="M12 2a10 10 0 1 0 10 10A10 10 0 0 0 12 2Z" />
+                            <path d="M12 8v4" />
+                            <path d="M12 16h.01" />
+                          </svg>
+                        </div>
+                        <h3 className="text-lg font-semibold text-foreground mb-1">
+                          Mental Models Not Enabled
+                        </h3>
+                        <p className="text-sm text-muted-foreground max-w-md">
+                          Mental models consolidation is disabled on this server. Set{" "}
+                          <code className="px-1 py-0.5 bg-muted rounded text-xs">
+                            HINDSIGHT_API_ENABLE_MENTAL_MODELS=true
+                          </code>{" "}
+                          to enable.
+                        </p>
+                      </div>
+                    ))}
                   {subTab === "reflections" && <ReflectionsView key="reflections" />}
                 </div>
               </div>

hindsight-control-plane/src/app/layout.tsx

@@ -1,6 +1,7 @@
 import type { Metadata } from "next";
 import "./globals.css";
 import { BankProvider } from "@/lib/bank-context";
+import { FeaturesProvider } from "@/lib/features-context";
 import { ThemeProvider } from "@/lib/theme-context";
 
 export const metadata: Metadata = {
@@ -20,7 +21,9 @@ export default function RootLayout({
     <html lang="en" suppressHydrationWarning>
       <body className="bg-background text-foreground">
         <ThemeProvider>
-          <BankProvider>{children}</BankProvider>
+          <FeaturesProvider>
+            <BankProvider>{children}</BankProvider>
+          </FeaturesProvider>
         </ThemeProvider>
       </body>
     </html>

hindsight-control-plane/src/components/bank-profile-view.tsx

@@ -5,6 +5,7 @@ import ReactMarkdown from "react-markdown";
 import { useRouter } from "next/navigation";
 import { client } from "@/lib/api";
 import { useBank } from "@/lib/bank-context";
+import { useFeatures } from "@/lib/features-context";
 import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { Textarea } from "@/components/ui/textarea";
@@ -212,6 +213,8 @@ function DispositionEditor({
 export function BankProfileView() {
   const router = useRouter();
   const { currentBank, setCurrentBank, loadBanks } = useBank();
+  const { features } = useFeatures();
+  const mentalModelsEnabled = features?.mental_models ?? false;
   const [profile, setProfile] = useState<BankProfile | null>(null);
   const [stats, setStats] = useState<BankStats | null>(null);
   const [operations, setOperations] = useState<Operation[]>([]);
@@ -391,12 +394,10 @@ export function BankProfileView() {
 
     setIsConsolidating(true);
     try {
-      const result = await client.triggerConsolidation(currentBank);
+      await client.triggerConsolidation(currentBank);
+      // Reload to show the new operation in the list
       await loadData();
-      alert(
-        result.message ||
-          `Consolidation completed: ${result.created} created, ${result.updated} updated`
-      );
+      await loadOperations();
     } catch (error) {
       console.error("Error triggering consolidation:", error);
       alert("Error triggering consolidation: " + (error as Error).message);
@@ -534,20 +535,32 @@ export function BankProfileView() {
                   Edit Profile
                 </DropdownMenuItem>
                 <DropdownMenuSeparator />
-                <DropdownMenuItem onClick={handleTriggerConsolidation} disabled={isConsolidating}>
+                <DropdownMenuItem
+                  onClick={handleTriggerConsolidation}
+                  disabled={isConsolidating || !mentalModelsEnabled}
+                  title={!mentalModelsEnabled ? "Mental models feature is not enabled" : undefined}
+                >
                   {isConsolidating ? (
                     <Loader2 className="w-4 h-4 mr-2 animate-spin" />
                   ) : (
                     <Brain className="w-4 h-4 mr-2" />
                   )}
                   {isConsolidating ? "Consolidating..." : "Run Consolidation"}
+                  {!mentalModelsEnabled && (
+                    <span className="ml-auto text-xs text-muted-foreground">Off</span>
+                  )}
                 </DropdownMenuItem>
                 <DropdownMenuItem
                   onClick={() => setShowClearMentalModelsDialog(true)}
+                  disabled={!mentalModelsEnabled}
                   className="text-amber-600 dark:text-amber-400 focus:text-amber-700 dark:focus:text-amber-300"
+                  title={!mentalModelsEnabled ? "Mental models feature is not enabled" : undefined}
                 >
                   <Trash2 className="w-4 h-4 mr-2" />
                   Clear Mental Models
+                  {!mentalModelsEnabled && (
+                    <span className="ml-auto text-xs text-muted-foreground">Off</span>
+                  )}
                 </DropdownMenuItem>
                 <DropdownMenuSeparator />
                 <DropdownMenuItem
@@ -649,12 +662,28 @@ export function BankProfileView() {
               {stats.nodes_by_fact_type?.experience || 0}
             </p>
           </div>
-          <div className="bg-amber-500/10 border border-amber-500/20 rounded-xl p-4 text-center">
-            <p className="text-xs text-amber-600 dark:text-amber-400 font-semibold uppercase tracking-wide">
+          <div
+            className={`rounded-xl p-4 text-center ${
+              mentalModelsEnabled
+                ? "bg-amber-500/10 border border-amber-500/20"
+                : "bg-muted/50 border border-muted"
+            }`}
+            title={!mentalModelsEnabled ? "Mental models feature is not enabled" : undefined}
+          >
+            <p
+              className={`text-xs font-semibold uppercase tracking-wide ${
+                mentalModelsEnabled ? "text-amber-600 dark:text-amber-400" : "text-muted-foreground"
+              }`}
+            >
               Mental Models
+              {!mentalModelsEnabled && <span className="ml-1 normal-case">(Off)</span>}
             </p>
-            <p className="text-2xl font-bold text-amber-600 dark:text-amber-400 mt-1">
-              {stats.total_mental_models || 0}
+            <p
+              className={`text-2xl font-bold mt-1 ${
+                mentalModelsEnabled ? "text-amber-600 dark:text-amber-400" : "text-muted-foreground"
+              }`}
+            >
+              {mentalModelsEnabled ? stats.total_mental_models || 0 : "—"}
             </p>
           </div>
           <div className="bg-rose-500/10 border border-rose-500/20 rounded-xl p-4 text-center">

hindsight-control-plane/src/components/data-view.tsx

@@ -616,10 +616,24 @@ export function DataView({ factType }: DataViewProps) {
                             <Table className="table-fixed">
                               <TableHeader>
                                 <TableRow className="bg-muted/50">
-                                  <TableHead className="w-[45%]">Memory</TableHead>
-                                  <TableHead className="w-[20%]">Entities</TableHead>
-                                  <TableHead className="w-[15%]">Occurred</TableHead>
-                                  <TableHead className="w-[15%]">Mentioned</TableHead>
+                                  <TableHead
+                                    className={factType === "mental_model" ? "w-[55%]" : "w-[45%]"}
+                                  >
+                                    {factType === "mental_model" ? "Mental Model" : "Memory"}
+                                  </TableHead>
+                                  {factType === "mental_model" ? (
+                                    <>
+                                      <TableHead className="w-[10%]">Sources</TableHead>
+                                      <TableHead className="w-[15%]">Created</TableHead>
+                                      <TableHead className="w-[15%]">Mentioned</TableHead>
+                                    </>
+                                  ) : (
+                                    <>
+                                      <TableHead className="w-[20%]">Entities</TableHead>
+                                      <TableHead className="w-[15%]">Occurred</TableHead>
+                                      <TableHead className="w-[15%]">Mentioned</TableHead>
+                                    </>
+                                  )}
                                   <TableHead className="w-[5%]"></TableHead>
                                 </TableRow>
                               </TableHeader>
@@ -637,6 +651,12 @@ export function DataView({ factType }: DataViewProps) {
                                         day: "numeric",
                                       })
                                     : null;
+                                  const createdDisplay = row.created_at
+                                    ? new Date(row.created_at).toLocaleDateString("en-US", {
+                                        month: "short",
+                                        day: "numeric",
+                                      })
+                                    : null;
 
                                   return (
                                     <TableRow
@@ -656,40 +676,62 @@ export function DataView({ factType }: DataViewProps) {
                                           </div>
                                         )}
                                       </TableCell>
-                                      <TableCell className="py-2">
-                                        {row.entities ? (
-                                          <div className="flex gap-1 flex-wrap">
-                                            {row.entities
-                                              .split(", ")
-                                              .slice(0, 2)
-                                              .map((entity: string, i: number) => (
-                                                <span
-                                                  key={i}
-                                                  className="text-[10px] px-1.5 py-0.5 rounded-full bg-primary/10 text-primary font-medium"
-                                                >
-                                                  {entity}
-                                                </span>
-                                              ))}
-                                            {row.entities.split(", ").length > 2 && (
-                                              <span className="text-[10px] text-muted-foreground">
-                                                +{row.entities.split(", ").length - 2}
+                                      {factType === "mental_model" ? (
+                                        <>
+                                          <TableCell className="text-xs py-2 text-foreground text-center">
+                                            {row.proof_count || 1}
+                                          </TableCell>
+                                          <TableCell className="text-xs py-2 text-foreground">
+                                            {createdDisplay || (
+                                              <span className="text-muted-foreground">-</span>
+                                            )}
+                                          </TableCell>
+                                          <TableCell className="text-xs py-2 text-foreground">
+                                            {mentionedDisplay || (
+                                              <span className="text-muted-foreground">-</span>
+                                            )}
+                                          </TableCell>
+                                        </>
+                                      ) : (
+                                        <>
+                                          <TableCell className="py-2">
+                                            {row.entities ? (
+                                              <div className="flex gap-1 flex-wrap">
+                                                {row.entities
+                                                  .split(", ")
+                                                  .slice(0, 2)
+                                                  .map((entity: string, i: number) => (
+                                                    <span
+                                                      key={i}
+                                                      className="text-[10px] px-1.5 py-0.5 rounded-full bg-primary/10 text-primary font-medium"
+                                                    >
+                                                      {entity}
+                                                    </span>
+                                                  ))}
+                                                {row.entities.split(", ").length > 2 && (
+                                                  <span className="text-[10px] text-muted-foreground">
+                                                    +{row.entities.split(", ").length - 2}
+                                                  </span>
+                                                )}
+                                              </div>
+                                            ) : (
+                                              <span className="text-xs text-muted-foreground">
+                                                -
                                               </span>
                                             )}
-                                          </div>
-                                        ) : (
-                                          <span className="text-xs text-muted-foreground">-</span>
-                                        )}
-                                      </TableCell>
-                                      <TableCell className="text-xs py-2 text-foreground">
-                                        {occurredDisplay || (
-                                          <span className="text-muted-foreground">-</span>
-                                        )}
-                                      </TableCell>
-                                      <TableCell className="text-xs py-2 text-foreground">
-                                        {mentionedDisplay || (
-                                          <span className="text-muted-foreground">-</span>
-                                        )}
-                                      </TableCell>
+                                          </TableCell>
+                                          <TableCell className="text-xs py-2 text-foreground">
+                                            {occurredDisplay || (
+                                              <span className="text-muted-foreground">-</span>
+                                            )}
+                                          </TableCell>
+                                          <TableCell className="text-xs py-2 text-foreground">
+                                            {mentionedDisplay || (
+                                              <span className="text-muted-foreground">-</span>
+                                            )}
+                                          </TableCell>
+                                        </>
+                                      )}
                                       <TableCell className="py-2">
                                         <Button
                                           onClick={(e) => {

hindsight-control-plane/src/components/memory-detail-panel.tsx

@@ -38,19 +38,11 @@ export function MemoryDetailPanel({
       return;
     }
 
-    const isMentalModel = memory?.fact_type === "mental_model" || memory?.type === "mental_model";
-
     setLoading(true);
 
-    const fetchPromise = isMentalModel
-      ? client.getMentalModel(bankId, memoryId).then((data) => ({
-          ...data,
-          fact_type: "mental_model",
-          type: "mental_model",
-        }))
-      : client.getMemory(memoryId, bankId);
-
-    fetchPromise
+    // Use getMemory for all memory types - it now returns source_memories for mental models
+    client
+      .getMemory(memoryId, bankId)
       .then((data) => {
         setFullMemory(data);
       })
@@ -66,6 +58,8 @@ export function MemoryDetailPanel({
 
   // Use full memory data if available, otherwise fall back to the partial data passed in
   const displayMemory = fullMemory || memory;
+  const isMentalModel =
+    displayMemory?.fact_type === "mental_model" || displayMemory?.type === "mental_model";
 
   const copyToClipboard = async (text: string) => {
     try {
@@ -133,8 +127,8 @@ export function MemoryDetailPanel({
                 </div>
               </div>
 
-              {/* Context */}
-              {displayMemory.context && (
+              {/* Context (not shown for mental models) */}
+              {displayMemory.context && !isMentalModel && (
                 <div className="p-4 bg-muted/50 rounded-lg">
                   <div className="text-xs font-bold text-muted-foreground uppercase mb-2">
                     Context
@@ -520,6 +514,49 @@ export function MemoryDetailPanel({
                 )}
               </div>
             )}
+
+            {/* Source Memories (for mental models) */}
+            {displayMemory.source_memories && displayMemory.source_memories.length > 0 && (
+              <div className={`${compact ? "p-2" : "p-3"} bg-muted rounded-lg`}>
+                <div className={`${labelSize} font-bold text-muted-foreground uppercase mb-2`}>
+                  Source Memories ({displayMemory.source_memories.length})
+                </div>
+                <div className="space-y-2">
+                  {displayMemory.source_memories.map((source: any, i: number) => (
+                    <div
+                      key={source.id || i}
+                      className="p-2 bg-background/50 rounded border border-border/50"
+                    >
+                      <div className="flex items-start justify-between gap-2 mb-1">
+                        <span
+                          className={`px-1.5 py-0.5 rounded text-[10px] flex-shrink-0 ${
+                            source.type === "experience"
+                              ? "bg-green-500/10 text-green-600"
+                              : "bg-blue-500/10 text-blue-600"
+                          }`}
+                        >
+                          {source.type}
+                        </span>
+                        <Button
+                          variant="outline"
+                          size="sm"
+                          className="h-5 text-[10px] px-2"
+                          onClick={() => setSourceMemoryModalId(source.id)}
+                        >
+                          View
+                        </Button>
+                      </div>
+                      <p className={`${textSize} mb-1`}>{source.text}</p>
+                      {source.context && (
+                        <p className="text-[10px] text-muted-foreground italic">
+                          Context: {source.context}
+                        </p>
+                      )}
+                    </div>
+                  ))}
+                </div>
+              </div>
+            )}
           </div>
         )}
       </div>
@@ -528,6 +565,12 @@ export function MemoryDetailPanel({
       {modalType && modalId && (
         <DocumentChunkModal type={modalType} id={modalId} onClose={closeModal} />
       )}
+
+      {/* Source Memory Modal */}
+      <MemoryDetailModal
+        memoryId={sourceMemoryModalId}
+        onClose={() => setSourceMemoryModalId(null)}
+      />
     </>
   );
 }

hindsight-control-plane/src/components/reflections-view.tsx

@@ -541,13 +541,46 @@ function ReflectionDetailPanel({
     if (!currentBank) return;
 
     setRefreshing(true);
+    const originalRefreshedAt = reflection.last_refreshed_at;
+
     try {
-      const updated = await client.refreshReflection(currentBank, reflection.id);
-      onRefreshed(updated);
+      // Submit the refresh task
+      await client.refreshReflection(currentBank, reflection.id);
+
+      // Poll until last_refreshed_at changes
+      const pollInterval = 1000; // 1 second
+      const maxAttempts = 120; // 2 minutes max
+      let attempts = 0;
+
+      const poll = async (): Promise<void> => {
+        attempts++;
+        try {
+          const updated = await client.getReflection(currentBank, reflection.id);
+          if (updated.last_refreshed_at !== originalRefreshedAt) {
+            // Refresh complete
+            onRefreshed(updated);
+            setRefreshing(false);
+            return;
+          }
+          if (attempts >= maxAttempts) {
+            // Timeout
+            setRefreshing(false);
+            alert("Refresh is taking longer than expected. Check the operations list for status.");
+            return;
+          }
+          // Continue polling
+          setTimeout(poll, pollInterval);
+        } catch (error) {
+          console.error("Error polling reflection:", error);
+          setRefreshing(false);
+        }
+      };
+
+      // Start polling after a short delay
+      setTimeout(poll, pollInterval);
     } catch (error) {
       console.error("Error refreshing reflection:", error);
       alert("Error refreshing: " + (error as Error).message);
-    } finally {
       setRefreshing(false);
     }
   };

hindsight-control-plane/src/lib/api.ts

@@ -261,11 +261,8 @@ export class ControlPlaneClient {
    */
   async triggerConsolidation(bankId: string) {
     return this.fetchApi<{
-      status: string;
-      processed: number;
-      created: number;
-      updated: number;
-      message: string;
+      operation_id: string;
+      deduplicated: boolean;
     }>(`/api/banks/${bankId}/consolidate`, {
       method: "POST",
     });
@@ -658,23 +655,11 @@ export class ControlPlaneClient {
   }
 
   /**
-   * Refresh a reflection (re-run source query)
+   * Refresh a reflection (re-run source query) - async operation
    */
   async refreshReflection(bankId: string, reflectionId: string) {
     return this.fetchApi<{
-      id: string;
-      bank_id: string;
-      name: string;
-      source_query: string;
-      content: string;
-      tags: string[];
-      last_refreshed_at: string;
-      created_at: string;
-      reflect_response?: {
-        text: string;
-        based_on: Record<string, Array<{ id: string; text: string; type: string }>>;
-        mental_models?: Array<{ id: string; text: string }>;
-      };
+      operation_id: string;
     }>(`/api/banks/${bankId}/reflections/${reflectionId}/refresh`, {
       method: "POST",
     });

hindsight-control-plane/src/lib/features-context.tsx

@@ -0,0 +1,63 @@
+"use client";
+
+import React, { createContext, useContext, useState, useEffect } from "react";
+import { client } from "./api";
+
+interface Features {
+  mental_models: boolean;
+  mcp: boolean;
+  worker: boolean;
+}
+
+interface FeaturesContextType {
+  features: Features | null;
+  loading: boolean;
+  error: string | null;
+}
+
+const defaultFeatures: Features = {
+  mental_models: false,
+  mcp: false,
+  worker: false,
+};
+
+const FeaturesContext = createContext<FeaturesContextType | undefined>(undefined);
+
+export function FeaturesProvider({ children }: { children: React.ReactNode }) {
+  const [features, setFeatures] = useState<Features | null>(null);
+  const [loading, setLoading] = useState(true);
+  const [error, setError] = useState<string | null>(null);
+
+  useEffect(() => {
+    const loadFeatures = async () => {
+      try {
+        const response = await client.getVersion();
+        setFeatures(response.features);
+        setError(null);
+      } catch (err) {
+        console.error("Error loading features:", err);
+        setError("Failed to load feature flags");
+        // Use defaults on error
+        setFeatures(defaultFeatures);
+      } finally {
+        setLoading(false);
+      }
+    };
+
+    loadFeatures();
+  }, []);
+
+  return (
+    <FeaturesContext.Provider value={{ features, loading, error }}>
+      {children}
+    </FeaturesContext.Provider>
+  );
+}
+
+export function useFeatures() {
+  const context = useContext(FeaturesContext);
+  if (context === undefined) {
+    throw new Error("useFeatures must be used within a FeaturesProvider");
+  }
+  return context;
+}

hindsight-docs/static/openapi.json

@@ -1205,7 +1205,7 @@
           "Reflections"
         ],
         "summary": "Refresh reflection",
-        "description": "Re-run the source query through reflect and update the content.",
+        "description": "Submit an async task to re-run the source query through reflect and update the content.",
         "operationId": "refresh_reflection",
         "parameters": [
           {
@@ -1249,7 +1249,7 @@
             "content": {
               "application/json": {
                 "schema": {
-                  "$ref": "#/components/schemas/ReflectionResponse"
+                  "$ref": "#/components/schemas/AsyncOperationSubmitResponse"
                 }
               }
             }
@@ -2980,6 +2980,29 @@
           "update_disposition": true
         }
       },
+      "AsyncOperationSubmitResponse": {
+        "properties": {
+          "operation_id": {
+            "type": "string",
+            "title": "Operation Id"
+          },
+          "status": {
+            "type": "string",
+            "title": "Status"
+          }
+        },
+        "type": "object",
+        "required": [
+          "operation_id",
+          "status"
+        ],
+        "title": "AsyncOperationSubmitResponse",
+        "description": "Response model for submitting an async operation.",
+        "example": {
+          "operation_id": "550e8400-e29b-41d4-a716-446655440000",
+          "status": "queued"
+        }
+      },
       "BackgroundResponse": {
         "properties": {
           "mission": {
@@ -3427,39 +3450,21 @@
       },
       "ConsolidationResponse": {
         "properties": {
-          "status": {
+          "operation_id": {
             "type": "string",
-            "title": "Status",
-            "description": "Status of the consolidation (completed or queued)"
+            "title": "Operation Id",
+            "description": "ID of the async consolidation operation"
           },
-          "processed": {
-            "type": "integer",
-            "title": "Processed",
-            "description": "Number of memories processed"
-          },
-          "created": {
-            "type": "integer",
-            "title": "Created",
-            "description": "Number of mental models created"
-          },
-          "updated": {
-            "type": "integer",
-            "title": "Updated",
-            "description": "Number of mental models updated"
-          },
-          "message": {
-            "type": "string",
-            "title": "Message",
-            "description": "Human-readable summary"
+          "deduplicated": {
+            "type": "boolean",
+            "title": "Deduplicated",
+            "description": "True if an existing pending task was reused",
+            "default": false
           }
         },
         "type": "object",
         "required": [
-          "status",
-          "processed",
-          "created",
-          "updated",
-          "message"
+          "operation_id"
         ],
         "title": "ConsolidationResponse",
         "description": "Response model for consolidation trigger endpoint."