fleet-memory/hindsight-api/hindsight_api/engine/memory_engine.py

"""
Memory Engine for Memory Banks.

This implements a sophisticated memory architecture that combines:
1. Temporal links: Memories connected by time proximity
2. Semantic links: Memories connected by meaning/similarity
3. Entity links: Memories connected by shared entities (PERSON, ORG, etc.)
4. Spreading activation: Search through the graph with activation decay
5. Dynamic weighting: Recency and frequency-based importance
"""

import asyncio
import contextvars
import json
import logging
import time
import uuid
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING, Any

import tiktoken

from ..config import get_config
from ..metrics import get_metrics_collector
from ..tracing import create_operation_span
from ..utils import mask_network_location
from .db_budget import budgeted_operation
from .operation_metadata import (
    BatchRetainChildMetadata,
    BatchRetainParentMetadata,
    ConsolidationMetadata,
    RefreshMentalModelMetadata,
    RetainMetadata,
)

# Context variable for current schema (async-safe, per-task isolation)
# Note: default is None, actual default comes from config via get_current_schema()
_current_schema: contextvars.ContextVar[str | None] = contextvars.ContextVar("current_schema", default=None)


def get_current_schema() -> str:
    """Get the current schema from context (falls back to config default)."""
    schema = _current_schema.get()
    if schema is None:
        # Fall back to configured default schema
        return get_config().database_schema
    return schema


# Initialize tiktoken encoder once at module level for efficiency
_tiktoken_encoder = tiktoken.get_encoding("cl100k_base")  # GPT-4/GPT-3.5-turbo encoding


def count_tokens(text: str) -> int:
    """Count tokens in text using tiktoken (cl100k_base encoding for GPT-4/3.5)."""
    return len(_tiktoken_encoder.encode(text))


def fq_table(table_name: str) -> str:
    """
    Get fully-qualified table name with current schema.

    Example:
        fq_table("memory_units") -> "public.memory_units"
        fq_table("memory_units") -> "tenant_xyz.memory_units" (if schema is set)
    """
    return f"{get_current_schema()}.{table_name}"


# Tables that must be schema-qualified (for runtime validation)
_PROTECTED_TABLES = frozenset(
    [
        "memory_units",
        "memory_links",
        "unit_entities",
        "entities",
        "entity_cooccurrences",
        "banks",
        "documents",
        "chunks",
        "async_operations",
        "file_storage",
    ]
)

# Enable runtime SQL validation (can be disabled in production for performance)
_VALIDATE_SQL_SCHEMAS = True


class UnqualifiedTableError(Exception):
    """Raised when SQL contains unqualified table references."""

    pass


def validate_sql_schema(sql: str) -> None:
    """
    Validate that SQL doesn't contain unqualified table references.

    This is a runtime safety check to prevent cross-tenant data access.
    Raises UnqualifiedTableError if any protected table is referenced
    without a schema prefix.

    Args:
        sql: The SQL query to validate

    Raises:
        UnqualifiedTableError: If unqualified table reference found
    """
    if not _VALIDATE_SQL_SCHEMAS:
        return

    import re

    sql_upper = sql.upper()

    for table in _PROTECTED_TABLES:
        table_upper = table.upper()

        # Pattern: SQL keyword followed by unqualified table name
        # Matches: FROM memory_units, JOIN memory_units, INTO memory_units, UPDATE memory_units
        patterns = [
            rf"FROM\s+{table_upper}(?:\s|$|,|\)|;)",
            rf"JOIN\s+{table_upper}(?:\s|$|,|\)|;)",
            rf"INTO\s+{table_upper}(?:\s|$|\()",
            rf"UPDATE\s+{table_upper}(?:\s|$)",
            rf"DELETE\s+FROM\s+{table_upper}(?:\s|$|;)",
        ]

        for pattern in patterns:
            match = re.search(pattern, sql_upper)
            if match:
                # Check if it's actually qualified (preceded by schema.)
                # Look backwards from match to see if there's a dot
                start = match.start()
                # Find the table name position in the match
                table_pos = sql_upper.find(table_upper, start)
                if table_pos > 0:
                    # Check character before table name (skip whitespace)
                    prefix = sql[:table_pos].rstrip()
                    if not prefix.endswith("."):
                        raise UnqualifiedTableError(
                            f"Unqualified table reference '{table}' in SQL. "
                            f"Use fq_table('{table}') for schema safety. "
                            f"SQL snippet: ...{sql[max(0, start - 10) : start + 50]}..."
                        )


import asyncpg
import numpy as np
from pydantic import BaseModel, Field

from .cross_encoder import CrossEncoderModel
from .embeddings import Embeddings, create_embeddings_from_env
from .interface import MemoryEngineInterface

if TYPE_CHECKING:
    from hindsight_api.extensions import OperationValidatorExtension, TenantExtension
    from hindsight_api.models import RequestContext


from enum import Enum

from ..metrics import get_metrics_collector
from ..pg0 import EmbeddedPostgres, parse_pg0_url
from .entity_resolver import EntityResolver
from .llm_wrapper import LLMConfig, requires_api_key
from .query_analyzer import QueryAnalyzer
from .reflect import run_reflect_agent
from .reflect.tools import tool_expand, tool_recall, tool_search_mental_models, tool_search_observations
from .response_models import (
    VALID_RECALL_FACT_TYPES,
    EntityObservation,
    EntityState,
    LLMCallTrace,
    MemoryFact,
    ObservationRef,
    ReflectResult,
    TokenUsage,
    ToolCallTrace,
)
from .response_models import RecallResult as RecallResultModel
from .retain import bank_utils, embedding_utils
from .retain.types import RetainContentDict
from .search import think_utils
from .search.reranking import CrossEncoderReranker
from .search.tags import TagsMatch
from .task_backend import BrokerTaskBackend, SyncTaskBackend, TaskBackend


class Budget(str, Enum):
    """Budget levels for recall/reflect operations."""

    LOW = "low"
    MID = "mid"
    HIGH = "high"


def utcnow():
    """Get current UTC time with timezone info."""
    return datetime.now(UTC)


# Logger for memory system
logger = logging.getLogger(__name__)

import tiktoken

from .db_utils import acquire_with_retry

# Cache tiktoken encoding for token budget filtering (module-level singleton)
_TIKTOKEN_ENCODING = None


def _get_tiktoken_encoding():
    """Get cached tiktoken encoding (cl100k_base for GPT-4/3.5)."""
    global _TIKTOKEN_ENCODING
    if _TIKTOKEN_ENCODING is None:
        _TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
    return _TIKTOKEN_ENCODING


class MemoryEngine(MemoryEngineInterface):
    """
    Advanced memory system using temporal and semantic linking with PostgreSQL.

    This class provides:
    - Embedding generation for semantic search
    - Entity, temporal, and semantic link creation
    - Think operations for formulating answers with opinions
    - bank profile and disposition management
    """

    def __init__(
        self,
        db_url: str | None = None,
        memory_llm_provider: str | None = None,
        memory_llm_api_key: str | None = None,
        memory_llm_model: str | None = None,
        memory_llm_base_url: str | None = None,
        # Per-operation LLM config (optional, falls back to memory_llm_* params)
        retain_llm_provider: str | None = None,
        retain_llm_api_key: str | None = None,
        retain_llm_model: str | None = None,
        retain_llm_base_url: str | None = None,
        reflect_llm_provider: str | None = None,
        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,
        pool_min_size: int | None = None,
        pool_max_size: int | None = None,
        db_command_timeout: int | None = None,
        db_acquire_timeout: int | None = None,
        task_backend: TaskBackend | None = None,
        run_migrations: bool = True,
        operation_validator: "OperationValidatorExtension | None" = None,
        tenant_extension: "TenantExtension | None" = None,
        skip_llm_verification: bool | None = None,
        lazy_reranker: bool | None = None,
    ):
        """
        Initialize the temporal + semantic memory system.

        All parameters are optional and will be read from environment variables if not provided.
        See hindsight_api.config for environment variable names and defaults.

        Args:
            db_url: PostgreSQL connection URL. Defaults to HINDSIGHT_API_DATABASE_URL env var or "pg0".
                    Also supports pg0 URLs: "pg0" or "pg0://instance-name" or "pg0://instance-name:port"
            memory_llm_provider: LLM provider. Defaults to HINDSIGHT_API_LLM_PROVIDER env var or "groq".
            memory_llm_api_key: API key for the LLM provider. Defaults to HINDSIGHT_API_LLM_API_KEY env var.
            memory_llm_model: Model name. Defaults to HINDSIGHT_API_LLM_MODEL env var.
            memory_llm_base_url: Base URL for the LLM API. Defaults based on provider.
            retain_llm_provider: LLM provider for retain operations. Falls back to memory_llm_provider.
            retain_llm_api_key: API key for retain LLM. Falls back to memory_llm_api_key.
            retain_llm_model: Model for retain operations. Falls back to memory_llm_model.
            retain_llm_base_url: Base URL for retain LLM. Falls back to memory_llm_base_url.
            reflect_llm_provider: LLM provider for reflect operations. Falls back to memory_llm_provider.
            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.
            pool_min_size: Minimum number of connections in the pool. Defaults to HINDSIGHT_API_DB_POOL_MIN_SIZE.
            pool_max_size: Maximum number of connections in the pool. Defaults to HINDSIGHT_API_DB_POOL_MAX_SIZE.
            db_command_timeout: PostgreSQL command timeout in seconds. Defaults to HINDSIGHT_API_DB_COMMAND_TIMEOUT.
            db_acquire_timeout: Connection acquisition timeout in seconds. Defaults to HINDSIGHT_API_DB_ACQUIRE_TIMEOUT.
            task_backend: Custom task backend. If not provided, uses BrokerTaskBackend for distributed processing.
            run_migrations: Whether to run database migrations during initialize(). Default: True
            operation_validator: Optional extension to validate operations before execution.
                                If provided, retain/recall/reflect operations will be validated.
            tenant_extension: Optional extension for multi-tenancy and API key authentication.
                             If provided, operations require a RequestContext for authentication.
            skip_llm_verification: Skip LLM connection verification during initialization.
                                  Defaults to HINDSIGHT_API_SKIP_LLM_VERIFICATION env var or False.
            lazy_reranker: Delay reranker initialization until first use. Useful for retain-only
                          operations that don't need the cross-encoder. Defaults to
                          HINDSIGHT_API_LAZY_RERANKER env var or False.
        """
        # Load config from environment for any missing parameters
        from ..config import get_config

        config = get_config()

        # Apply optimization flags from config if not explicitly provided
        self._skip_llm_verification = (
            skip_llm_verification if skip_llm_verification is not None else config.skip_llm_verification
        )
        self._lazy_reranker = lazy_reranker if lazy_reranker is not None else config.lazy_reranker

        # Apply defaults from config
        db_url = db_url or config.database_url
        memory_llm_provider = memory_llm_provider or config.llm_provider
        memory_llm_api_key = memory_llm_api_key or config.llm_api_key
        if not memory_llm_api_key and requires_api_key(memory_llm_provider):
            raise ValueError("LLM API key is required. Set HINDSIGHT_API_LLM_API_KEY environment variable.")
        memory_llm_model = memory_llm_model or config.llm_model
        memory_llm_base_url = memory_llm_base_url or config.get_llm_base_url() or None
        # Track pg0 instance (if used)
        self._pg0: EmbeddedPostgres | None = None

        # Initialize PostgreSQL connection URL
        # The actual URL will be set during initialize() after starting the server
        # Supports: "pg0" (default instance), "pg0://instance-name" (named instance), or regular postgresql:// URL
        self._use_pg0, self._pg0_instance_name, self._pg0_port = parse_pg0_url(db_url)
        if self._use_pg0:
            self.db_url = None
        else:
            self.db_url = db_url

        # Set default base URL if not provided
        if memory_llm_base_url is None:
            if memory_llm_provider.lower() == "groq":
                memory_llm_base_url = "https://api.groq.com/openai/v1"
            elif memory_llm_provider.lower() == "ollama":
                memory_llm_base_url = "http://localhost:11434/v1"
            else:
                memory_llm_base_url = ""

        # Connection pool (will be created in initialize())
        self._pool = None
        self._initialized = False
        self._pool_min_size = pool_min_size if pool_min_size is not None else config.db_pool_min_size
        self._pool_max_size = pool_max_size if pool_max_size is not None else config.db_pool_max_size
        self._db_command_timeout = db_command_timeout if db_command_timeout is not None else config.db_command_timeout
        self._db_acquire_timeout = db_acquire_timeout if db_acquire_timeout is not None else config.db_acquire_timeout
        self._run_migrations = run_migrations

        # Initialize entity resolver (will be created in initialize())
        self.entity_resolver = None

        # Initialize embeddings (from env vars if not provided)
        if embeddings is not None:
            self.embeddings = embeddings
        else:
            self.embeddings = create_embeddings_from_env()

        # Initialize query analyzer
        if query_analyzer is not None:
            self.query_analyzer = query_analyzer
        else:
            from .query_analyzer import DateparserQueryAnalyzer

            self.query_analyzer = DateparserQueryAnalyzer()

        # Initialize LLM configuration (default, used as fallback)
        self._llm_config = LLMConfig(
            provider=memory_llm_provider,
            api_key=memory_llm_api_key,
            base_url=memory_llm_base_url,
            model=memory_llm_model,
        )

        # Store client and model for convenience (deprecated: use _llm_config.call() instead)
        self._llm_client = self._llm_config._client
        self._llm_model = self._llm_config.model

        # Initialize per-operation LLM configs (fall back to default if not specified)
        # Retain LLM config - for fact extraction (benefits from strong structured output)
        retain_provider = retain_llm_provider or config.retain_llm_provider or memory_llm_provider
        retain_api_key = retain_llm_api_key or config.retain_llm_api_key or memory_llm_api_key
        retain_model = retain_llm_model or config.retain_llm_model or memory_llm_model
        retain_base_url = retain_llm_base_url or config.retain_llm_base_url or memory_llm_base_url
        # Apply provider-specific base URL defaults for retain
        if retain_base_url is None:
            if retain_provider.lower() == "groq":
                retain_base_url = "https://api.groq.com/openai/v1"
            elif retain_provider.lower() == "ollama":
                retain_base_url = "http://localhost:11434/v1"
            else:
                retain_base_url = ""

        self._retain_llm_config = LLMConfig(
            provider=retain_provider,
            api_key=retain_api_key,
            base_url=retain_base_url,
            model=retain_model,
        )

        # Reflect LLM config - for think/observe operations (can use lighter models)
        reflect_provider = reflect_llm_provider or config.reflect_llm_provider or memory_llm_provider
        reflect_api_key = reflect_llm_api_key or config.reflect_llm_api_key or memory_llm_api_key
        reflect_model = reflect_llm_model or config.reflect_llm_model or memory_llm_model
        reflect_base_url = reflect_llm_base_url or config.reflect_llm_base_url or memory_llm_base_url
        # Apply provider-specific base URL defaults for reflect
        if reflect_base_url is None:
            if reflect_provider.lower() == "groq":
                reflect_base_url = "https://api.groq.com/openai/v1"
            elif reflect_provider.lower() == "ollama":
                reflect_base_url = "http://localhost:11434/v1"
            else:
                reflect_base_url = ""

        self._reflect_llm_config = LLMConfig(
            provider=reflect_provider,
            api_key=reflect_api_key,
            base_url=reflect_base_url,
            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)

        # Initialize task backend
        # If no custom backend provided, use BrokerTaskBackend which stores tasks in PostgreSQL
        # The pool_getter lambda will return the pool once it's initialized
        self._task_backend = task_backend or BrokerTaskBackend(
            pool_getter=lambda: self._pool,
            schema_getter=get_current_schema,
        )

        # Backpressure mechanism: limit concurrent searches to prevent overwhelming the database
        # Configurable via HINDSIGHT_API_RECALL_MAX_CONCURRENT (default: 50)
        self._search_semaphore = asyncio.Semaphore(get_config().recall_max_concurrent)

        # Backpressure for put operations: limit concurrent puts to prevent database contention
        # Each put_batch holds a connection for the entire transaction, so we limit to 5
        # concurrent puts to avoid connection pool exhaustion and reduce write contention
        self._put_semaphore = asyncio.Semaphore(5)

        # initialize encoding eagerly to avoid delaying the first time
        _get_tiktoken_encoding()

        # Store operation validator extension (optional)
        self._operation_validator = operation_validator

        # Store tenant extension (always set, use default if none provided)
        if tenant_extension is None:
            from ..extensions.builtin.tenant import DefaultTenantExtension

            tenant_extension = DefaultTenantExtension(config={})
        self._tenant_extension = tenant_extension

    @property
    def tenant_extension(self) -> "TenantExtension | None":
        """The configured tenant extension, if any."""
        return self._tenant_extension

    async def _validate_operation(self, validation_coro) -> None:
        """
        Run validation if an operation validator is configured.

        Args:
            validation_coro: Coroutine that returns a ValidationResult

        Raises:
            OperationValidationError: If validation fails
        """
        if self._operation_validator is None:
            return

        from hindsight_api.extensions import OperationValidationError

        result = await validation_coro
        if not result.allowed:
            raise OperationValidationError(result.reason or "Operation not allowed", result.status_code)

    async def _authenticate_tenant(self, request_context: "RequestContext | None") -> str:
        """
        Authenticate tenant and set schema in context variable.

        The schema is stored in a contextvar for async-safe, per-task isolation.
        Use fq_table(table_name) to get fully-qualified table names.

        Args:
            request_context: The request context with API key. Required if tenant_extension is configured.

        Returns:
            Schema name that was set in the context.

        Raises:
            AuthenticationError: If authentication fails or request_context is missing when required.
        """
        from hindsight_api.extensions import AuthenticationError

        if request_context is None:
            raise AuthenticationError("RequestContext is required")

        # For internal/background operations (e.g., worker tasks), skip extension authentication.
        # The task was already authenticated at submission time, and execute_task sets _current_schema
        # from the task's _schema field.
        if request_context.internal:
            return _current_schema.get()

        # Authenticate through tenant extension (always set, may be default no-auth extension)
        tenant_context = await self._tenant_extension.authenticate(request_context)

        _current_schema.set(tenant_context.schema_name)
        return tenant_context.schema_name

    async def _handle_batch_retain(self, task_dict: dict[str, Any]):
        """
        Handler for batch retain tasks.

        Args:
            task_dict: Dict with 'bank_id', 'contents', 'operation_id'

        Raises:
            ValueError: If bank_id is missing
            Exception: Any exception from retain_batch_async (propagates to execute_task for retry)
        """
        bank_id = task_dict.get("bank_id")
        if not bank_id:
            raise ValueError("bank_id is required for batch retain task")
        contents = task_dict.get("contents", [])
        document_tags = task_dict.get("document_tags")
        operation_id = task_dict.get("operation_id")  # For batch API crash recovery

        logger.info(
            f"[BATCH_RETAIN_TASK] Starting background batch retain for bank_id={bank_id}, {len(contents)} items, operation_id={operation_id}"
        )

        # Restore tenant_id/api_key_id from task payload so extensions
        # (e.g., operation validators) can attribute the operation correctly.
        # internal=True to skip extension auth (worker has no API key),
        # user_initiated=True so extensions know this originated from a user request.
        from hindsight_api.models import RequestContext

        context = RequestContext(
            internal=True,
            user_initiated=True,
            tenant_id=task_dict.get("_tenant_id"),
            api_key_id=task_dict.get("_api_key_id"),
        )
        await self.retain_batch_async(
            bank_id=bank_id,
            contents=contents,
            document_tags=document_tags,
            request_context=context,
            operation_id=operation_id,
        )

        # If this retain was triggered by file conversion, update document with file metadata
        file_metadata = task_dict.get("_file_metadata")
        if file_metadata and len(contents) == 1:
            doc_id = contents[0].get("document_id")
            if doc_id:
                pool = await self._get_pool()
                async with acquire_with_retry(pool) as conn:
                    await conn.execute(
                        f"""
                        UPDATE {fq_table("documents")}
                        SET file_storage_key = $3,
                            file_original_name = $4,
                            file_content_type = $5,
                            updated_at = NOW()
                        WHERE id = $1 AND bank_id = $2
                        """,
                        doc_id,
                        bank_id,
                        file_metadata["file_storage_key"],
                        file_metadata["file_original_name"],
                        file_metadata["file_content_type"],
                    )

        logger.info(f"[BATCH_RETAIN_TASK] Completed background batch retain for bank_id={bank_id}")

    async def _handle_file_convert_retain(self, task_dict: dict[str, Any]):
        """
        Handler for file conversion tasks.

        Converts a file to markdown, then submits a separate async retain operation
        and marks this conversion as completed — all in a single transaction.
        This avoids holding a worker slot during the expensive retain pipeline.

        Args:
            task_dict: Dict with 'bank_id', 'storage_key', 'parser', etc.

        Raises:
            ValueError: If required fields are missing
            Exception: Any exception from conversion (includes filename in error)
        """
        bank_id = task_dict.get("bank_id")
        storage_key = task_dict.get("storage_key")
        document_id = task_dict.get("document_id")
        operation_id = task_dict.get("operation_id")
        filename = task_dict.get("original_filename", "unknown")

        if not all([bank_id, storage_key, document_id]):
            raise ValueError("bank_id, storage_key, and document_id are required for file_convert_retain task")

        logger.info(f"[FILE_CONVERT_RETAIN] Starting for bank_id={bank_id}, document_id={document_id}, file={filename}")

        try:
            # Retrieve file from storage
            file_data = await self._file_storage.retrieve(storage_key)

            # Convert to markdown
            parser = self._parser_registry.get_parser(
                name=task_dict.get("parser"),
                filename=filename,
                content_type=task_dict.get("content_type"),
            )
            markdown_content = await parser.convert(file_data, filename)
        except Exception as e:
            # Re-raise with filename context for better error reporting
            error_msg = f"Failed to parse file '{filename}': {str(e)}"
            logger.error(f"[FILE_CONVERT_RETAIN] {error_msg}")
            raise RuntimeError(error_msg) from e

        logger.info(
            f"[FILE_CONVERT_RETAIN] Converted file for bank_id={bank_id}, "
            f"document_id={document_id}, {len(markdown_content)} chars. Submitting retain task."
        )

        # Build retain task payload
        retain_contents = [
            {
                "content": markdown_content,
                "document_id": document_id,
                "context": task_dict.get("context"),
                "metadata": task_dict.get("metadata", {}),
                "tags": task_dict.get("tags", []),
                "timestamp": task_dict.get("timestamp"),
            }
        ]
        document_tags = task_dict.get("document_tags")

        retain_task_payload: dict[str, Any] = {"contents": retain_contents}
        if document_tags:
            retain_task_payload["document_tags"] = document_tags

        # Pass tenant/api_key context through to retain task
        if task_dict.get("_tenant_id"):
            retain_task_payload["_tenant_id"] = task_dict["_tenant_id"]
        if task_dict.get("_api_key_id"):
            retain_task_payload["_api_key_id"] = task_dict["_api_key_id"]

        # File metadata to attach after retain creates the document
        retain_task_payload["_file_metadata"] = {
            "file_storage_key": storage_key,
            "file_original_name": task_dict["original_filename"],
            "file_content_type": task_dict["content_type"],
        }

        # In one transaction: create the retain async operation AND mark this conversion as completed
        retain_operation_id = uuid.uuid4()
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            async with conn.transaction():
                # Create the retain operation record
                await conn.execute(
                    f"""
                    INSERT INTO {fq_table("async_operations")}
                    (operation_id, bank_id, operation_type, result_metadata, status)
                    VALUES ($1, $2, $3, $4, $5)
                    """,
                    retain_operation_id,
                    bank_id,
                    "retain",
                    json.dumps({}),
                    "pending",
                )

                # Mark this file_convert_retain operation as completed
                if operation_id:
                    await conn.execute(
                        f"""
                        UPDATE {fq_table("async_operations")}
                        SET status = 'completed', updated_at = NOW(), completed_at = NOW()
                        WHERE operation_id = $1
                        """,
                        uuid.UUID(operation_id),
                    )

        # Submit the retain task to the task backend (outside the transaction)
        full_retain_payload = {
            "type": "batch_retain",
            "operation_id": str(retain_operation_id),
            "bank_id": bank_id,
            **retain_task_payload,
        }
        await self._task_backend.submit_task(full_retain_payload)

        logger.info(
            f"[FILE_CONVERT_RETAIN] Completed conversion for bank_id={bank_id}, "
            f"document_id={document_id}. Retain task submitted as operation {retain_operation_id}"
        )

        # Delete file bytes from storage if configured (saves storage costs)
        from ..config import get_config

        config = get_config()
        if config.file_delete_after_retain:
            try:
                await self._file_storage.delete(storage_key)
                logger.info(f"[FILE_CONVERT_RETAIN] Deleted file bytes for {storage_key} (conversion completed)")
            except Exception as e:
                # Non-fatal - log and continue
                logger.warning(f"[FILE_CONVERT_RETAIN] Failed to delete file {storage_key}: {e}")

    async def _handle_consolidation(self, task_dict: dict[str, Any]):
        """
        Handler for consolidation tasks.

        Consolidates new memories into mental models for a bank.

        Args:
            task_dict: Dict with 'bank_id'

        Raises:
            ValueError: If bank_id is missing
            Exception: Any exception from consolidation (propagates to execute_task for retry)
        """
        bank_id = task_dict.get("bank_id")
        if not bank_id:
            raise ValueError("bank_id is required for consolidation task")

        from hindsight_api.models import RequestContext

        from .consolidation import run_consolidation_job

        # Restore tenant_id/api_key_id from task payload so downstream operations
        # (e.g., mental model refreshes) can attribute usage to the correct org.
        internal_context = RequestContext(
            internal=True,
            tenant_id=task_dict.get("_tenant_id"),
            api_key_id=task_dict.get("_api_key_id"),
        )
        result = await run_consolidation_job(
            memory_engine=self,
            bank_id=bank_id,
            request_context=internal_context,
        )

        logger.info(f"[CONSOLIDATION] bank={bank_id} completed: {result.get('memories_processed', 0)} processed")

    async def _handle_refresh_mental_model(self, task_dict: dict[str, Any]):
        """
        Handler for refresh_mental_model tasks.

        Re-runs the source query through reflect and updates the mental model content.

        Args:
            task_dict: Dict with 'bank_id', 'mental_model_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")
        mental_model_id = task_dict.get("mental_model_id")

        if not bank_id or not mental_model_id:
            raise ValueError("bank_id and mental_model_id are required for refresh_mental_model task")

        logger.info(f"[REFRESH_MENTAL_MODEL_TASK] Starting for bank_id={bank_id}, mental_model_id={mental_model_id}")

        from hindsight_api.models import RequestContext

        # Restore tenant_id/api_key_id from task payload so extensions can
        # attribute the mental_model_refresh operation to the correct org.
        internal_context = RequestContext(
            internal=True,
            tenant_id=task_dict.get("_tenant_id"),
            api_key_id=task_dict.get("_api_key_id"),
        )

        # Get the current mental model to get source_query
        mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=internal_context)
        if not mental_model:
            raise ValueError(f"Mental model {mental_model_id} not found in bank {bank_id}")

        source_query = mental_model["source_query"]

        # SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching
        # to ensure it can only access other mental models/memories with the SAME tags.
        # This prevents cross-tenant/cross-user information leakage by excluding untagged content.
        tags = mental_model.get("tags")
        tags_match = "all_strict" if tags else "any"

        # Run reflect to generate new content, excluding the mental model being refreshed
        reflect_result = await self.reflect_async(
            bank_id=bank_id,
            query=source_query,
            request_context=internal_context,
            tags=tags,
            tags_match=tags_match,
            exclude_mental_model_ids=[mental_model_id],
        )

        generated_content = reflect_result.text or "No content generated"

        # Build reflect_response payload to store
        # based_on contains MemoryFact objects for most types, but plain dicts for directives
        based_on_serialized: dict[str, list[dict[str, Any]]] = {}
        for fact_type, facts in reflect_result.based_on.items():
            serialized_facts = []
            for fact in facts:
                if isinstance(fact, dict):
                    # Plain dict (e.g., directives with id, name, content)
                    serialized_facts.append(
                        {
                            "id": str(fact["id"]),
                            "text": fact.get("text", fact.get("content", fact.get("name", ""))),
                            "type": fact_type,
                        }
                    )
                else:
                    # MemoryFact object with .id and .text attributes
                    serialized_facts.append(
                        {
                            "id": str(fact.id),
                            "text": fact.text,
                            "type": fact_type,
                        }
                    )
            based_on_serialized[fact_type] = serialized_facts

        reflect_response = {
            "text": reflect_result.text,
            "based_on": based_on_serialized,
        }

        # Update the mental model with the generated content and reflect_response
        await self.update_mental_model(
            bank_id=bank_id,
            mental_model_id=mental_model_id,
            content=generated_content,
            reflect_response=reflect_response,
            request_context=internal_context,
        )

        # Call post-operation hook if validator is configured
        if self._operation_validator:
            from hindsight_api.extensions.operation_validator import MentalModelRefreshResult

            # Count facts and mental models from based_on
            facts_used = 0
            mental_models_used = 0
            if reflect_result.based_on:
                for fact_type, facts in reflect_result.based_on.items():
                    if facts:
                        if fact_type == "mental_models":
                            mental_models_used += len(facts)
                        else:
                            facts_used += len(facts)

            # Estimate tokens
            query_tokens = len(source_query) // 4 if source_query else 0
            output_tokens = len(generated_content) // 4 if generated_content else 0
            context_tokens = 0  # refresh doesn't use additional context

            result_ctx = MentalModelRefreshResult(
                bank_id=bank_id,
                mental_model_id=mental_model_id,
                request_context=internal_context,
                query_tokens=query_tokens,
                output_tokens=output_tokens,
                context_tokens=context_tokens,
                facts_used=facts_used,
                mental_models_used=mental_models_used,
                success=True,
            )
            try:
                await self._operation_validator.on_mental_model_refresh_complete(result_ctx)
            except Exception as hook_err:
                logger.warning(f"Post-mental-model-refresh hook error (non-fatal): {hook_err}")

        logger.info(f"[REFRESH_MENTAL_MODEL_TASK] Completed for bank_id={bank_id}, mental_model_id={mental_model_id}")

    async def execute_task(self, task_dict: dict[str, Any]):
        """
        Execute a task by routing it to the appropriate handler.

        This method is called by the task backend to execute tasks.
        It receives a plain dict that can be serialized and sent over the network.

        Args:
            task_dict: Task dictionary with 'type' key and other payload data
                      Example: {'type': 'batch_retain', 'bank_id': '...', 'contents': [...]}
        """
        task_type = task_dict.get("type")
        operation_id = task_dict.get("operation_id")
        retry_count = task_dict.get("retry_count", 0)
        max_retries = 3

        # Set schema context for multi-tenant task execution
        schema = task_dict.pop("_schema", None)
        if schema:
            _current_schema.set(schema)

        # Check if operation was cancelled (only for tasks with operation_id)
        if operation_id:
            try:
                pool = await self._get_pool()
                async with acquire_with_retry(pool) as conn:
                    result = await conn.fetchrow(
                        f"SELECT operation_id FROM {fq_table('async_operations')} WHERE operation_id = $1",
                        uuid.UUID(operation_id),
                    )
                    if not result:
                        # Operation was cancelled, skip processing
                        logger.info(f"Skipping cancelled operation: {operation_id}")
                        return
            except Exception as e:
                logger.error(f"Failed to check operation status {operation_id}: {e}")
                # Continue with processing if we can't check status

        try:
            if task_type == "batch_retain":
                await self._handle_batch_retain(task_dict)
            elif task_type == "file_convert_retain":
                await self._handle_file_convert_retain(task_dict)
            elif task_type == "consolidation":
                await self._handle_consolidation(task_dict)
            elif task_type == "refresh_mental_model":
                await self._handle_refresh_mental_model(task_dict)
            else:
                logger.error(f"Unknown task type: {task_type}")
                # Don't retry unknown task types
                if operation_id:
                    await self._delete_operation_record(operation_id)
                return

            # Task succeeded - mark operation as completed
            # file_convert_retain marks itself as completed in a transaction, skip double-marking
            if operation_id and task_type != "file_convert_retain":
                await self._mark_operation_completed(operation_id)

        except Exception as e:
            # Task failed - check if we should retry
            logger.error(
                f"Task execution failed (attempt {retry_count + 1}/{max_retries + 1}): {task_type}, error: {e}"
            )
            import traceback

            error_traceback = traceback.format_exc()
            traceback.print_exc()

            # Don't retry file conversion - if conversion fails, it won't succeed on retry
            # (missing OCR, corrupted file, unsupported format, etc.)
            should_retry = retry_count < max_retries and task_type != "file_convert_retain"

            if should_retry:
                # Reschedule with incremented retry count
                task_dict["retry_count"] = retry_count + 1
                logger.info(f"Rescheduling task {task_type} (retry {retry_count + 1}/{max_retries})")
                await self._task_backend.submit_task(task_dict)
            else:
                # Max retries exceeded or non-retryable task - mark operation as failed
                reason = "non-retryable task type" if task_type == "file_convert_retain" else "max retries exceeded"
                logger.error(f"Not retrying task {task_type} ({reason}), marking as failed")
                if operation_id:
                    await self._mark_operation_failed(operation_id, str(e), error_traceback)

    async def _delete_operation_record(self, operation_id: str):
        """Helper to delete an operation record from the database."""
        try:
            pool = await self._get_pool()
            async with acquire_with_retry(pool) as conn:
                await conn.execute(
                    f"DELETE FROM {fq_table('async_operations')} WHERE operation_id = $1", uuid.UUID(operation_id)
                )
        except Exception as e:
            logger.error(f"Failed to delete async operation record {operation_id}: {e}")

    async def _mark_operation_failed(self, operation_id: str, error_message: str, error_traceback: str):
        """Helper to mark an operation as failed in the database.

        Also checks if this is a child operation and updates the parent if all siblings are done.
        Uses a single transaction to avoid race conditions when multiple children fail simultaneously.
        """
        try:
            pool = await self._get_pool()
            # Truncate error message to avoid extremely long strings
            full_error = f"{error_message}\n\nTraceback:\n{error_traceback}"
            truncated_error = full_error[:5000] if len(full_error) > 5000 else full_error

            async with acquire_with_retry(pool) as conn:
                async with conn.transaction():
                    # Mark this operation as failed
                    await conn.execute(
                        f"""
                        UPDATE {fq_table("async_operations")}
                        SET status = 'failed', error_message = $2, updated_at = NOW()
                        WHERE operation_id = $1
                        """,
                        uuid.UUID(operation_id),
                        truncated_error,
                    )
                    logger.info(f"Marked async operation as failed: {operation_id}")

                    # Check if this is a child operation and update parent if all siblings are done
                    # This happens in the same transaction after the child status is updated
                    await self._maybe_update_parent_operation(operation_id, conn)
        except Exception as e:
            logger.error(f"Failed to mark operation as failed {operation_id}: {e}")

    async def _mark_operation_completed(self, operation_id: str):
        """Helper to mark an operation as completed in the database.

        Also checks if this is a child operation and updates the parent if all siblings are done.
        Uses a single transaction to avoid race conditions when multiple children complete simultaneously.
        """
        try:
            pool = await self._get_pool()
            async with acquire_with_retry(pool) as conn:
                async with conn.transaction():
                    # Mark this operation as completed
                    await conn.execute(
                        f"""
                        UPDATE {fq_table("async_operations")}
                        SET status = 'completed', updated_at = NOW(), completed_at = NOW()
                        WHERE operation_id = $1
                        """,
                        uuid.UUID(operation_id),
                    )
                    logger.info(f"Marked async operation as completed: {operation_id}")

                    # Check if this is a child operation and update parent if all siblings are done
                    # This happens in the same transaction after the child status is updated
                    await self._maybe_update_parent_operation(operation_id, conn)
        except Exception as e:
            logger.error(f"Failed to mark operation as completed {operation_id}: {e}")

    async def _maybe_update_parent_operation(self, child_operation_id: str, conn):
        """Check if this is a child operation and update parent status if all siblings are done.

        Must be called within an active transaction that has already updated the child's status.
        Uses SELECT FOR UPDATE to lock the parent and prevent race conditions.

        Args:
            child_operation_id: The operation ID that just completed or failed
            conn: Database connection with an active transaction
        """
        try:
            # Get this operation's metadata to check if it has a parent
            row = await conn.fetchrow(
                f"""
                SELECT result_metadata, bank_id
                FROM {fq_table("async_operations")}
                WHERE operation_id = $1
                """,
                uuid.UUID(child_operation_id),
            )

            if not row:
                return

            result_metadata = json.loads(row["result_metadata"]) if row["result_metadata"] else {}
            parent_operation_id = result_metadata.get("parent_operation_id")

            if not parent_operation_id:
                # Not a child operation
                return

            bank_id = row["bank_id"]

            # Lock the parent operation to prevent concurrent updates from other children
            # Use FOR UPDATE to ensure only one child can update the parent at a time
            parent_row = await conn.fetchrow(
                f"""
                SELECT operation_id
                FROM {fq_table("async_operations")}
                WHERE operation_id = $1 AND bank_id = $2
                FOR UPDATE
                """,
                uuid.UUID(parent_operation_id),
                bank_id,
            )

            if not parent_row:
                # Parent doesn't exist (shouldn't happen)
                return

            # Get all sibling operations (including this one)
            # This query runs in the same transaction, so it sees the current child's updated status
            siblings = await conn.fetch(
                f"""
                SELECT status
                FROM {fq_table("async_operations")}
                WHERE bank_id = $1
                AND result_metadata::jsonb @> $2::jsonb
                """,
                bank_id,
                json.dumps({"parent_operation_id": parent_operation_id}),
            )

            if not siblings:
                return

            # Check if all siblings are done (completed or failed)
            all_completed = all(sib["status"] == "completed" for sib in siblings)
            any_failed = any(sib["status"] == "failed" for sib in siblings)
            all_done = all(sib["status"] in ("completed", "failed") for sib in siblings)

            if not all_done:
                # Some siblings still pending/processing
                return

            # All siblings are done - update parent status
            if any_failed:
                new_status = "failed"
                # Set parent error message to indicate child failure
                await conn.execute(
                    f"""
                    UPDATE {fq_table("async_operations")}
                    SET status = $2, error_message = $3, updated_at = NOW()
                    WHERE operation_id = $1
                    """,
                    uuid.UUID(parent_operation_id),
                    new_status,
                    "One or more sub-batches failed",
                )
            elif all_completed:
                new_status = "completed"
                await conn.execute(
                    f"""
                    UPDATE {fq_table("async_operations")}
                    SET status = $2, updated_at = NOW(), completed_at = NOW()
                    WHERE operation_id = $1
                    """,
                    uuid.UUID(parent_operation_id),
                    new_status,
                )

            logger.info(f"Updated parent operation {parent_operation_id} to status '{new_status}' (all children done)")

        except Exception as e:
            logger.error(f"Failed to update parent operation for child {child_operation_id}: {e}")
            # Re-raise to rollback the transaction
            raise

    async def initialize(self):
        """Initialize the connection pool, models, and background workers.

        Loads models (embeddings, cross-encoder) in parallel with pg0 startup
        for faster overall initialization.
        """
        if self._initialized:
            return

        # Run model loading in thread pool (CPU-bound) in parallel with pg0 startup
        loop = asyncio.get_event_loop()

        async def start_pg0():
            """Start pg0 if configured."""
            if self._use_pg0:
                kwargs = {"name": self._pg0_instance_name}
                if self._pg0_port is not None:
                    kwargs["port"] = self._pg0_port
                pg0 = EmbeddedPostgres(**kwargs)
                # Check if pg0 is already running before we start it
                was_already_running = await pg0.is_running()
                self.db_url = await pg0.ensure_running()
                # Only track pg0 (to stop later) if WE started it
                if not was_already_running:
                    self._pg0 = pg0

        async def init_embeddings():
            """Initialize embedding model."""
            # For local providers, run in thread pool to avoid blocking event loop
            if self.embeddings.provider_name == "local":
                await loop.run_in_executor(None, lambda: asyncio.run(self.embeddings.initialize()))
            else:
                await self.embeddings.initialize()

        async def init_cross_encoder():
            """Initialize cross-encoder model."""
            cross_encoder = self._cross_encoder_reranker.cross_encoder
            # For local providers, run in thread pool to avoid blocking event loop
            if cross_encoder.provider_name == "local":
                await loop.run_in_executor(None, lambda: asyncio.run(cross_encoder.initialize()))
            else:
                await cross_encoder.initialize()
            # Mark reranker as initialized
            self._cross_encoder_reranker._initialized = True

        async def init_query_analyzer():
            """Initialize query analyzer model."""
            # Query analyzer load is sync and CPU-bound
            await loop.run_in_executor(None, self.query_analyzer.load)

        async def verify_llm():
            """Verify LLM connections are working for all unique configs."""
            if not self._skip_llm_verification:
                # Verify default config
                await self._llm_config.verify_connection()
                # Verify retain config if different from default
                retain_is_different = (
                    self._retain_llm_config.provider != self._llm_config.provider
                    or self._retain_llm_config.model != self._llm_config.model
                )
                if retain_is_different:
                    await self._retain_llm_config.verify_connection()
                # Verify reflect config if different from default and retain
                reflect_is_different = (
                    self._reflect_llm_config.provider != self._llm_config.provider
                    or self._reflect_llm_config.model != self._llm_config.model
                ) and (
                    self._reflect_llm_config.provider != self._retain_llm_config.provider
                    or self._reflect_llm_config.model != self._retain_llm_config.model
                )
                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 = [
            start_pg0(),
            init_embeddings(),
            init_query_analyzer(),
        ]

        # Only init cross-encoder eagerly if not using lazy initialization
        if not self._lazy_reranker:
            init_tasks.append(init_cross_encoder())

        # Only verify LLM if not skipping
        if not self._skip_llm_verification:
            init_tasks.append(verify_llm())

        # Run pg0 and selected model initializations in parallel
        await asyncio.gather(*init_tasks)

        # Run database migrations if enabled
        if self._run_migrations:
            from ..migrations import (
                ensure_embedding_dimension,
                ensure_text_search_extension,
                ensure_vector_extension,
                run_migrations,
            )

            if not self.db_url:
                raise ValueError("Database URL is required for migrations")

            # Migrate all schemas from the tenant extension
            # The tenant extension is the single source of truth for which schemas exist
            logger.info("Running database migrations...")
            tenants = await self._tenant_extension.list_tenants()
            if tenants:
                logger.info(f"Running migrations on {len(tenants)} schema(s)...")
                for tenant in tenants:
                    schema = tenant.schema
                    if schema:
                        run_migrations(self.db_url, schema=schema)
                logger.info("Schema migrations completed")

                # Get config for vector extension setting
                config = get_config()

                # Ensure embedding column dimension matches the model's dimension
                # This is done after migrations and after embeddings.initialize()
                for tenant in tenants:
                    schema = tenant.schema
                    if schema:
                        ensure_embedding_dimension(
                            self.db_url,
                            self.embeddings.dimension,
                            schema=schema,
                            vector_extension=config.vector_extension,
                        )

                # Ensure vector indexes match the configured extension
                for tenant in tenants:
                    schema = tenant.schema
                    if schema:
                        ensure_vector_extension(self.db_url, vector_extension=config.vector_extension, schema=schema)

                # Ensure text search columns/indexes match the configured extension
                for tenant in tenants:
                    schema = tenant.schema
                    if schema:
                        ensure_text_search_extension(
                            self.db_url, text_search_extension=config.text_search_extension, schema=schema
                        )

        logger.info(f"Connecting to PostgreSQL at {mask_network_location(self.db_url)}")

        # Create connection pool
        # For read-heavy workloads with many parallel think/search operations,
        # we need a larger pool. Read operations don't need strong isolation.
        self._pool = await asyncpg.create_pool(
            self.db_url,
            min_size=self._pool_min_size,
            max_size=self._pool_max_size,
            command_timeout=self._db_command_timeout,
            statement_cache_size=0,  # Disable prepared statement cache
            timeout=self._db_acquire_timeout,  # Connection acquisition timeout (seconds)
        )

        # Initialize entity resolver with pool
        self.entity_resolver = EntityResolver(self._pool)

        # Initialize config resolver for hierarchical configuration
        from ..config_resolver import ConfigResolver

        self._config_resolver = ConfigResolver(pool=self._pool, tenant_extension=self._tenant_extension)
        logger.debug("Config resolver initialized for hierarchical configuration")

        # Initialize file storage
        from .storage import create_file_storage

        config = get_config()
        self._file_storage = create_file_storage(
            storage_type=config.file_storage_type,
            pool_getter=lambda: self._pool,
            schema=get_current_schema() if get_current_schema() != config.database_schema else None,
        )
        logger.debug(f"File storage initialized ({config.file_storage_type})")

        # Initialize parser registry
        from .parsers import FileParserRegistry, IrisParser, MarkitdownParser

        self._parser_registry = FileParserRegistry()
        try:
            self._parser_registry.register(MarkitdownParser())
            logger.debug("Registered markitdown parser")
        except ImportError:
            logger.warning("markitdown not available - file parsing disabled")
        iris_token = config.file_parser_iris_token
        iris_org_id = config.file_parser_iris_org_id
        if iris_token and iris_org_id:
            self._parser_registry.register(IrisParser(token=iris_token, org_id=iris_org_id))
            logger.debug("Registered iris parser")
        else:
            logger.debug("Iris parser not registered (VECTORIZE_TOKEN or VECTORIZE_ORG_ID not set)")

        # Set executor for task backend and initialize
        self._task_backend.set_executor(self.execute_task)
        await self._task_backend.initialize()

        self._initialized = True
        logger.info("Memory system initialized (pool and task backend started)")

    async def _get_pool(self) -> asyncpg.Pool:
        """Get the connection pool (must call initialize() first)."""
        if not self._initialized:
            await self.initialize()
        return self._pool

    async def _acquire_connection(self):
        """
        Acquire a connection from the pool with retry logic.

        Returns an async context manager that yields a connection.
        Retries on transient connection errors with exponential backoff.
        """
        pool = await self._get_pool()

        async def acquire():
            return await pool.acquire()

        return await _retry_with_backoff(acquire)

    async def health_check(self) -> dict:
        """
        Perform a health check by querying the database.

        Returns:
            dict with status and optional error message

        Note:
            Returns unhealthy until initialize() has completed successfully.
        """
        # Not healthy until fully initialized
        if not self._initialized:
            return {"status": "unhealthy", "reason": "not_initialized"}

        try:
            pool = await self._get_pool()
            async with pool.acquire() as conn:
                result = await conn.fetchval("SELECT 1")
                if result == 1:
                    return {"status": "healthy", "database": "connected"}
                else:
                    return {"status": "unhealthy", "database": "unexpected response"}
        except Exception as e:
            return {"status": "unhealthy", "database": "error", "error": str(e)}

    async def close(self):
        """Close the connection pool and shutdown background workers."""
        logger.info("close() started")

        # Shutdown task backend
        await self._task_backend.shutdown()

        # Close pool
        if self._pool is not None:
            self._pool.terminate()
            self._pool = None

        self._initialized = False

        # Stop pg0 if we started it
        if self._pg0 is not None:
            logger.info("Stopping pg0...")
            await self._pg0.stop()
            self._pg0 = None
            logger.info("pg0 stopped")

    async def wait_for_background_tasks(self):
        """
        Wait for all pending background tasks to complete.

        This is useful in tests to ensure background tasks complete before making assertions.
        """
        if hasattr(self._task_backend, "wait_for_pending_tasks"):
            await self._task_backend.wait_for_pending_tasks()

    def _format_readable_date(self, dt: datetime) -> str:
        """
        Format a datetime into a readable string for temporal matching.

        Examples:
            - June 2024
            - January 15, 2024
            - December 2023

        This helps queries like "camping in June" match facts that happened in June.

        Args:
            dt: datetime object to format

        Returns:
            Readable date string
        """
        # Format as "Month Year" for most cases
        # Could be extended to include day for very specific dates if needed
        month_name = dt.strftime("%B")  # Full month name (e.g., "June")
        year = dt.strftime("%Y")  # Year (e.g., "2024")

        # For now, use "Month Year" format
        # Could check if day is significant (not 1st or 15th) and include it
        return f"{month_name} {year}"

    async def _find_duplicate_facts_batch(
        self,
        conn,
        bank_id: str,
        texts: list[str],
        embeddings: list[list[float]],
        event_date: datetime,
        time_window_hours: int = 24,
        similarity_threshold: float = 0.95,
    ) -> list[bool]:
        """
        Check which facts are duplicates using semantic similarity + temporal window.

        For each new fact, checks if a semantically similar fact already exists
        within the time window. Uses pgvector cosine similarity for efficiency.

        Args:
            conn: Database connection
            bank_id: bank IDentifier
            texts: List of fact texts to check
            embeddings: Corresponding embeddings
            event_date: Event date for temporal filtering
            time_window_hours: Hours before/after event_date to search (default: 24)
            similarity_threshold: Minimum cosine similarity to consider duplicate (default: 0.95)

        Returns:
            List of booleans - True if fact is a duplicate (should skip), False if new
        """
        if not texts:
            return []

        # Handle edge cases where event_date is at datetime boundaries
        try:
            time_lower = event_date - timedelta(hours=time_window_hours)
        except OverflowError:
            time_lower = datetime.min
        try:
            time_upper = event_date + timedelta(hours=time_window_hours)
        except OverflowError:
            time_upper = datetime.max

        # Fetch ALL existing facts in time window ONCE (much faster than N queries)
        import time as time_mod

        fetch_start = time_mod.time()
        existing_facts = await conn.fetch(
            f"""
            SELECT id, text, embedding
            FROM {fq_table("memory_units")}
            WHERE bank_id = $1
              AND event_date BETWEEN $2 AND $3
            """,
            bank_id,
            time_lower,
            time_upper,
        )

        # If no existing facts, nothing is duplicate
        if not existing_facts:
            return [False] * len(texts)

        # Compute similarities in Python (vectorized with numpy)
        is_duplicate = []

        # Convert existing embeddings to numpy for faster computation
        embedding_arrays = []
        for row in existing_facts:
            raw_emb = row["embedding"]
            # Handle different pgvector formats
            if isinstance(raw_emb, str):
                # Parse string format: "[1.0, 2.0, ...]"
                import json

                emb = np.array(json.loads(raw_emb), dtype=np.float32)
            elif isinstance(raw_emb, (list, tuple)):
                emb = np.array(raw_emb, dtype=np.float32)
            else:
                # Try direct conversion
                emb = np.array(raw_emb, dtype=np.float32)
            embedding_arrays.append(emb)

        if not embedding_arrays:
            existing_embeddings = np.array([])
        elif len(embedding_arrays) == 1:
            # Single embedding: reshape to (1, dim)
            existing_embeddings = embedding_arrays[0].reshape(1, -1)
        else:
            # Multiple embeddings: vstack
            existing_embeddings = np.vstack(embedding_arrays)

        comp_start = time_mod.time()
        for embedding in embeddings:
            # Compute cosine similarity with all existing facts
            emb_array = np.array(embedding)
            # Cosine similarity = 1 - cosine distance
            # For normalized vectors: cosine_sim = dot product
            similarities = np.dot(existing_embeddings, emb_array)

            # Check if any existing fact is too similar
            max_similarity = np.max(similarities) if len(similarities) > 0 else 0
            is_duplicate.append(max_similarity > similarity_threshold)

        return is_duplicate

    def retain(
        self,
        bank_id: str,
        content: str,
        context: str = "",
        event_date: datetime | None = None,
        request_context: "RequestContext | None" = None,
    ) -> list[str]:
        """
        Store content as memory units (synchronous wrapper).

        This is a synchronous wrapper around retain_async() for convenience.
        For best performance, use retain_async() directly.

        Args:
            bank_id: Unique identifier for the bank
            content: Text content to store
            context: Context about when/why this memory was formed
            event_date: When the event occurred (defaults to now)
            request_context: Request context for authentication (optional, uses internal context if not provided)

        Returns:
            List of created unit IDs
        """
        # Run async version synchronously
        from hindsight_api.models import RequestContext as RC

        ctx = request_context if request_context is not None else RC()
        return asyncio.run(self.retain_async(bank_id, content, context, event_date, request_context=ctx))

    async def retain_async(
        self,
        bank_id: str,
        content: str,
        context: str = "",
        event_date: datetime | None = None,
        document_id: str | None = None,
        fact_type_override: str | None = None,
        confidence_score: float | None = None,
        *,
        request_context: "RequestContext",
    ) -> list[str]:
        """
        Store content as memory units with temporal and semantic links (ASYNC version).

        This is a convenience wrapper around retain_batch_async for a single content item.

        Args:
            bank_id: Unique identifier for the bank
            content: Text content to store
            context: Context about when/why this memory was formed
            event_date: When the event occurred (defaults to now)
            document_id: Optional document ID for tracking (always upserts if document already exists)
            fact_type_override: Override fact type ('world', 'experience')
            confidence_score: Confidence score (0.0 to 1.0)
            request_context: Request context for authentication.

        Returns:
            List of created unit IDs
        """
        # Build content dict
        content_dict: RetainContentDict = {"content": content, "context": context}
        if event_date:
            content_dict["event_date"] = event_date
        if document_id:
            content_dict["document_id"] = document_id

        # Use retain_batch_async with a single item (avoids code duplication)
        result = await self.retain_batch_async(
            bank_id=bank_id,
            contents=[content_dict],
            request_context=request_context,
            fact_type_override=fact_type_override,
            confidence_score=confidence_score,
        )

        # Return the first (and only) list of unit IDs
        return result[0] if result else []

    async def retain_batch_async(
        self,
        bank_id: str,
        contents: list[RetainContentDict],
        *,
        request_context: "RequestContext",
        document_id: str | None = None,
        fact_type_override: str | None = None,
        confidence_score: float | None = None,
        document_tags: list[str] | None = None,
        return_usage: bool = False,
        operation_id: str | None = None,
    ):
        """
        Store multiple content items as memory units in ONE batch operation.

        This is MUCH more efficient than calling retain_async multiple times:
        - Extracts facts from all contents in parallel
        - Generates ALL embeddings in ONE batch
        - Does ALL database operations in ONE transaction
        - Automatically chunks large batches to prevent timeouts

        Args:
            bank_id: Unique identifier for the bank
            contents: List of dicts with keys:
                - "content" (required): Text content to store
                - "context" (optional): Context about the memory
                - "event_date" (optional): When the event occurred
                - "document_id" (optional): Document ID for this specific content item
            document_id: **DEPRECATED** - Use "document_id" key in each content dict instead.
                        Applies the same document_id to ALL content items that don't specify their own.
            fact_type_override: Override fact type for all facts ('world', 'experience')
            confidence_score: Confidence score (0.0 to 1.0)
            return_usage: If True, returns tuple of (unit_ids, TokenUsage). Default False for backward compatibility.

        Returns:
            If return_usage=False: List of lists of unit IDs (one list per content item)
            If return_usage=True: Tuple of (unit_ids, TokenUsage)

        Example (new style - per-content document_id):
            unit_ids = await memory.retain_batch_async(
                bank_id="user123",
                contents=[
                    {"content": "Alice works at Google", "document_id": "doc1"},
                    {"content": "Bob loves Python", "document_id": "doc2"},
                    {"content": "More about Alice", "document_id": "doc1"},
                ]
            )
            # Returns: [["unit-id-1"], ["unit-id-2"], ["unit-id-3"]]

        Example (deprecated style - batch-level document_id):
            unit_ids = await memory.retain_batch_async(
                bank_id="user123",
                contents=[
                    {"content": "Alice works at Google"},
                    {"content": "Bob loves Python"},
                ],
                document_id="meeting-2024-01-15"
            )
            # Returns: [["unit-id-1"], ["unit-id-2"]]
        """
        start_time = time.time()

        if not contents:
            if return_usage:
                return [], TokenUsage()
            return []

        # Authenticate tenant and set schema in context (for fq_table())
        await self._authenticate_tenant(request_context)

        # Validate operation if validator is configured
        contents_copy = [dict(c) for c in contents]  # Convert TypedDict to regular dict for extension
        if self._operation_validator:
            from hindsight_api.extensions import RetainContext

            ctx = RetainContext(
                bank_id=bank_id,
                contents=contents_copy,
                request_context=request_context,
                document_id=document_id,
                fact_type_override=fact_type_override,
                confidence_score=confidence_score,
            )
            await self._validate_operation(self._operation_validator.validate_retain(ctx))

        # Apply batch-level document_id to contents that don't have their own (backwards compatibility)
        if document_id:
            for item in contents:
                if "document_id" not in item:
                    item["document_id"] = document_id

        # Validate no duplicate document_ids in the batch
        # Having duplicate document_ids causes race conditions in document upserts during parallel processing
        doc_ids = [item.get("document_id") for item in contents if item.get("document_id")]
        if len(doc_ids) != len(set(doc_ids)):
            from collections import Counter

            duplicates = [doc_id for doc_id, count in Counter(doc_ids).items() if count > 1]
            raise ValueError(
                f"Batch contains duplicate document_ids: {duplicates}. "
                f"Each content item in a batch must have a unique document_id to avoid race conditions."
            )

        # Auto-chunk large batches by token count to avoid timeouts and memory issues
        # Calculate total token count
        total_tokens = sum(count_tokens(item.get("content", "")) for item in contents)
        total_usage = TokenUsage()

        # Get batch size threshold from config
        config = get_config()
        tokens_per_batch = config.retain_batch_tokens

        if total_tokens > tokens_per_batch:
            # Split into smaller batches based on token count
            logger.info(
                f"Large batch detected ({total_tokens:,} tokens from {len(contents)} items). Splitting into sub-batches of ~{tokens_per_batch:,} tokens each..."
            )

            sub_batches = []
            current_batch = []
            current_batch_tokens = 0

            for item in contents:
                item_tokens = count_tokens(item.get("content", ""))

                # If adding this item would exceed the limit, start a new batch
                # (unless current batch is empty - then we must include it even if it's large)
                if current_batch and current_batch_tokens + item_tokens > tokens_per_batch:
                    sub_batches.append(current_batch)
                    current_batch = [item]
                    current_batch_tokens = item_tokens
                else:
                    current_batch.append(item)
                    current_batch_tokens += item_tokens

            # Add the last batch
            if current_batch:
                sub_batches.append(current_batch)

            logger.info(f"Split into {len(sub_batches)} sub-batches: {[len(b) for b in sub_batches]} items each")

            # Process each sub-batch
            all_results = []
            for i, sub_batch in enumerate(sub_batches, 1):
                sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch)
                logger.info(
                    f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
                )

                sub_results, sub_usage = await self._retain_batch_async_internal(
                    bank_id=bank_id,
                    contents=sub_batch,
                    request_context=request_context,
                    document_id=document_id,
                    is_first_batch=i == 1,  # Only upsert on first batch
                    fact_type_override=fact_type_override,
                    confidence_score=confidence_score,
                    document_tags=document_tags,
                    operation_id=operation_id,
                )
                all_results.extend(sub_results)
                total_usage = total_usage + sub_usage

            total_time = time.time() - start_time
            logger.info(
                f"RETAIN_BATCH_ASYNC (chunked) COMPLETE: {len(all_results)} results from {len(contents)} contents in {total_time:.3f}s"
            )
            result = all_results
        else:
            # Small batch - use internal method directly
            result, total_usage = await self._retain_batch_async_internal(
                bank_id=bank_id,
                contents=contents,
                request_context=request_context,
                document_id=document_id,
                is_first_batch=True,
                fact_type_override=fact_type_override,
                confidence_score=confidence_score,
                document_tags=document_tags,
                operation_id=operation_id,
            )

        # Call post-operation hook if validator is configured
        if self._operation_validator:
            from hindsight_api.extensions import RetainResult

            result_ctx = RetainResult(
                bank_id=bank_id,
                contents=contents_copy,
                request_context=request_context,
                document_id=document_id,
                fact_type_override=fact_type_override,
                confidence_score=confidence_score,
                unit_ids=result,
                success=True,
                error=None,
                llm_input_tokens=total_usage.input_tokens,
                llm_output_tokens=total_usage.output_tokens,
                llm_total_tokens=total_usage.total_tokens,
            )
            try:
                await self._operation_validator.on_retain_complete(result_ctx)
            except Exception as e:
                logger.warning(f"Post-retain hook error (non-fatal): {e}")

        # Trigger consolidation as a tracked async operation if enabled
        # Resolve bank-specific config to check if observations are enabled for this bank
        config = await self._config_resolver.resolve_full_config(bank_id, request_context)
        if config.enable_observations:
            try:
                await self.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
            except Exception as e:
                # Log but don't fail the retain - consolidation is non-critical
                logger.warning(f"Failed to submit consolidation task for bank {bank_id}: {e}")

        if return_usage:
            return result, total_usage
        return result

    async def _retain_batch_async_internal(
        self,
        bank_id: str,
        contents: list[RetainContentDict],
        request_context: "RequestContext",
        document_id: str | None = None,
        is_first_batch: bool = True,
        fact_type_override: str | None = None,
        confidence_score: float | None = None,
        document_tags: list[str] | None = None,
        operation_id: str | None = None,
    ) -> tuple[list[list[str]], "TokenUsage"]:
        """
        Internal method for batch processing without chunking logic.

        Assumes contents are already appropriately sized (< 50k chars).
        Called by retain_batch_async after chunking large batches.

        Uses semaphore for backpressure to limit concurrent retains.

        Args:
            bank_id: Unique identifier for the bank
            contents: List of dicts with content, context, event_date
            request_context: Request context for config resolution
            document_id: Optional document ID (always upserts if exists)
            is_first_batch: Whether this is the first batch (for chunked operations, only delete on first batch)
            fact_type_override: Override fact type for all facts
            confidence_score: Confidence score for opinions
            document_tags: Tags applied to all items in this batch

        Returns:
            Tuple of (unit ID lists, token usage for fact extraction)
        """
        # Backpressure: limit concurrent retains to prevent database contention
        async with self._put_semaphore:
            # Use the new modular orchestrator
            from .retain import orchestrator

            pool = await self._get_pool()

            # Resolve bank-specific config for this operation
            resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context)

            # Create parent span for retain operation
            with create_operation_span("retain", bank_id):
                return await orchestrator.retain_batch(
                    pool=pool,
                    embeddings_model=self.embeddings,
                    llm_config=self._retain_llm_config,
                    entity_resolver=self.entity_resolver,
                    format_date_fn=self._format_readable_date,
                    duplicate_checker_fn=self._find_duplicate_facts_batch,
                    bank_id=bank_id,
                    contents_dicts=contents,
                    document_id=document_id,
                    is_first_batch=is_first_batch,
                    fact_type_override=fact_type_override,
                    confidence_score=confidence_score,
                    document_tags=document_tags,
                    config=resolved_config,
                    operation_id=operation_id,
                    schema=request_context.tenant_id if request_context else None,
                )

    def recall(
        self,
        bank_id: str,
        query: str,
        fact_type: str,
        budget: Budget = Budget.MID,
        max_tokens: int = 4096,
        enable_trace: bool = False,
    ) -> tuple[list[dict[str, Any]], Any | None]:
        """
        Recall memories using 4-way parallel retrieval (synchronous wrapper).

        This is a synchronous wrapper around recall_async() for convenience.
        For best performance, use recall_async() directly.

        Args:
            bank_id: bank ID to recall for
            query: Recall query
            fact_type: Required filter for fact type ('world', 'experience', or 'opinion')
            budget: Budget level for graph traversal (low=100, mid=300, high=600 units)
            max_tokens: Maximum tokens to return (counts only 'text' field, default 4096)
            enable_trace: If True, returns detailed trace object

        Returns:
            Tuple of (results, trace)
        """
        # Run async version synchronously - deprecated sync method, passing None for request_context
        from hindsight_api.models import RequestContext

        return asyncio.run(
            self.recall_async(
                bank_id,
                query,
                budget=budget,
                max_tokens=max_tokens,
                enable_trace=enable_trace,
                fact_type=[fact_type],
                request_context=RequestContext(),
            )
        )

    async def recall_async(
        self,
        bank_id: str,
        query: str,
        *,
        budget: Budget | None = None,
        max_tokens: int = 4096,
        enable_trace: bool = False,
        fact_type: list[str] | None = None,
        question_date: datetime | None = None,
        include_entities: bool = False,
        max_entity_tokens: int = 500,
        include_chunks: bool = False,
        max_chunk_tokens: int = 8192,
        include_source_facts: bool = False,
        max_source_facts_tokens: int = 4096,
        request_context: "RequestContext",
        tags: list[str] | None = None,
        tags_match: TagsMatch = "any",
        _connection_budget: int | None = None,
        _quiet: bool = False,
    ) -> RecallResultModel:
        """
        Recall memories using N*4-way parallel retrieval (N fact types × 4 retrieval methods).

        This implements the core RECALL operation:
        1. Retrieval: For each fact type, run 4 parallel retrievals (semantic vector, BM25 keyword, graph activation, temporal graph)
        2. Merge: Combine using Reciprocal Rank Fusion (RRF)
        3. Rerank: Score using selected reranker (heuristic or cross-encoder)
        4. Diversify: Apply MMR for diversity
        5. Token Filter: Return results up to max_tokens budget

        Args:
            bank_id: bank ID to recall for
            query: Recall query
            fact_type: List of fact types to recall (e.g., ['world', 'experience'])
            budget: Budget level for graph traversal (low=100, mid=300, high=600 units)
            max_tokens: Maximum tokens to return (counts only 'text' field, default 4096)
                       Results are returned until token budget is reached, stopping before
                       including a fact that would exceed the limit
            enable_trace: Whether to return trace for debugging (deprecated)
            question_date: Optional date when question was asked (for temporal filtering)
            include_entities: Whether to include entity observations in the response
            max_entity_tokens: Maximum tokens for entity observations (default 500)
            include_chunks: Whether to include raw chunks in the response
            max_chunk_tokens: Maximum tokens for chunks (default 8192)
                             NOTE: Chunks are fetched independently of max_tokens filtering.
                             This means setting max_tokens=0 will return 0 facts but can still
                             return chunks from the top-scored (reranked) results.
                             Chunks are fetched in batches (estimated as (max_chunk_tokens // retain_chunk_size) * 2)
                             until the token budget is exhausted or all chunks are fetched.
                             This handles varying chunk sizes across documents.
            tags: Optional list of tags for visibility filtering (OR matching - returns
                  memories that have at least one matching tag)

        Returns:
            RecallResultModel containing:
            - results: List of MemoryFact objects (filtered by max_tokens)
            - trace: Optional trace information for debugging
            - entities: Optional dict of entity states (if include_entities=True)
            - chunks: Optional dict of chunks (if include_chunks=True, independent of max_tokens)
        """
        # Authenticate tenant and set schema in context (for fq_table())
        await self._authenticate_tenant(request_context)

        # Default to all fact types if not specified
        if fact_type is None:
            fact_type = list(VALID_RECALL_FACT_TYPES)

        # Filter out 'opinion' early (deprecated, silently ignore)
        fact_type = [ft for ft in fact_type if ft != "opinion"]

        # Validate fact types
        invalid_types = set(fact_type) - VALID_RECALL_FACT_TYPES
        if invalid_types:
            raise ValueError(
                f"Invalid fact type(s): {', '.join(sorted(invalid_types))}. "
                f"Must be one of: {', '.join(sorted(VALID_RECALL_FACT_TYPES))}"
            )
        if not fact_type:
            # All requested types were opinions - return empty result
            return RecallResultModel(results=[], entities={}, chunks={})

        # Validate operation if validator is configured
        if self._operation_validator:
            from hindsight_api.extensions import RecallContext

            ctx = RecallContext(
                bank_id=bank_id,
                query=query,
                request_context=request_context,
                budget=budget,
                max_tokens=max_tokens,
                enable_trace=enable_trace,
                fact_types=list(fact_type),
                question_date=question_date,
                include_entities=include_entities,
                max_entity_tokens=max_entity_tokens,
                include_chunks=include_chunks,
                max_chunk_tokens=max_chunk_tokens,
            )
            await self._validate_operation(self._operation_validator.validate_recall(ctx))

        # Map budget enum to thinking_budget number (default to MID if None)
        budget_mapping = {Budget.LOW: 100, Budget.MID: 300, Budget.HIGH: 1000}
        effective_budget = budget if budget is not None else Budget.MID
        thinking_budget = budget_mapping[effective_budget]

        # Log recall start with tags if present (skip if quiet mode for internal operations)
        if not _quiet:
            tags_info = f", tags={tags} ({tags_match})" if tags else ""
            logger.info(f"[RECALL {bank_id[:8]}] Starting recall for query: {query[:50]}...{tags_info}")

        # Create parent span for recall operation
        from ..tracing import get_tracer

        tracer = get_tracer()
        # Use start_as_current_span to ensure child spans are linked properly
        recall_span_context = tracer.start_as_current_span("hindsight.recall")
        recall_span = recall_span_context.__enter__()
        recall_span.set_attribute("hindsight.bank_id", bank_id)
        recall_span.set_attribute("hindsight.query", query[:100])
        recall_span.set_attribute("hindsight.fact_types", ",".join(fact_type))
        recall_span.set_attribute("hindsight.thinking_budget", thinking_budget)
        recall_span.set_attribute("hindsight.max_tokens", max_tokens)

        try:
            # Backpressure: limit concurrent recalls to prevent overwhelming the database
            result = None
            error_msg = None
            semaphore_wait_start = time.time()
            async with self._search_semaphore:
                semaphore_wait = time.time() - semaphore_wait_start
                # Retry loop for connection errors
                max_retries = 3
                for attempt in range(max_retries + 1):
                    try:
                        result = await self._search_with_retries(
                            bank_id,
                            query,
                            fact_type,
                            thinking_budget,
                            max_tokens,
                            enable_trace,
                            question_date,
                            include_entities,
                            max_entity_tokens,
                            include_chunks,
                            max_chunk_tokens,
                            request_context,
                            semaphore_wait=semaphore_wait,
                            tags=tags,
                            tags_match=tags_match,
                            connection_budget=_connection_budget,
                            quiet=_quiet,
                            include_source_facts=include_source_facts,
                            max_source_facts_tokens=max_source_facts_tokens,
                        )
                        break  # Success - exit retry loop
                    except Exception as e:
                        # Check if it's a connection error
                        is_connection_error = (
                            isinstance(e, asyncpg.TooManyConnectionsError)
                            or isinstance(e, asyncpg.CannotConnectNowError)
                            or (isinstance(e, asyncpg.PostgresError) and "connection" in str(e).lower())
                        )

                        if is_connection_error and attempt < max_retries:
                            # Wait with exponential backoff before retry
                            wait_time = 0.5 * (2**attempt)  # 0.5s, 1s, 2s
                            logger.warning(
                                f"Connection error on search attempt {attempt + 1}/{max_retries + 1}: {str(e)}. "
                                f"Retrying in {wait_time:.1f}s..."
                            )
                            await asyncio.sleep(wait_time)
                        else:
                            # Not a connection error or out of retries - call post-hook and raise
                            error_msg = str(e)
                            if self._operation_validator:
                                from hindsight_api.extensions.operation_validator import RecallResult

                                result_ctx = RecallResult(
                                    bank_id=bank_id,
                                    query=query,
                                    request_context=request_context,
                                    budget=budget,
                                    max_tokens=max_tokens,
                                    enable_trace=enable_trace,
                                    fact_types=list(fact_type),
                                    question_date=question_date,
                                    include_entities=include_entities,
                                    max_entity_tokens=max_entity_tokens,
                                    include_chunks=include_chunks,
                                    max_chunk_tokens=max_chunk_tokens,
                                    result=None,
                                    success=False,
                                    error=error_msg,
                                )
                                try:
                                    await self._operation_validator.on_recall_complete(result_ctx)
                                except Exception as hook_err:
                                    logger.warning(f"Post-recall hook error (non-fatal): {hook_err}")
                            raise
                else:
                    # Exceeded max retries
                    error_msg = "Exceeded maximum retries for search due to connection errors."
                    if self._operation_validator:
                        from hindsight_api.extensions.operation_validator import RecallResult

                        result_ctx = RecallResult(
                            bank_id=bank_id,
                            query=query,
                            request_context=request_context,
                            budget=budget,
                            max_tokens=max_tokens,
                            enable_trace=enable_trace,
                            fact_types=list(fact_type),
                            question_date=question_date,
                            include_entities=include_entities,
                            max_entity_tokens=max_entity_tokens,
                            include_chunks=include_chunks,
                            max_chunk_tokens=max_chunk_tokens,
                            result=None,
                            success=False,
                            error=error_msg,
                        )
                        try:
                            await self._operation_validator.on_recall_complete(result_ctx)
                        except Exception as hook_err:
                            logger.warning(f"Post-recall hook error (non-fatal): {hook_err}")
                    raise Exception(error_msg)

            # Call post-operation hook for success
            if self._operation_validator and result is not None:
                from hindsight_api.extensions.operation_validator import RecallResult

                result_ctx = RecallResult(
                    bank_id=bank_id,
                    query=query,
                    request_context=request_context,
                    budget=budget,
                    max_tokens=max_tokens,
                    enable_trace=enable_trace,
                    fact_types=list(fact_type),
                    question_date=question_date,
                    include_entities=include_entities,
                    max_entity_tokens=max_entity_tokens,
                    include_chunks=include_chunks,
                    max_chunk_tokens=max_chunk_tokens,
                    result=result,
                    success=True,
                    error=None,
                )
                try:
                    await self._operation_validator.on_recall_complete(result_ctx)
                except Exception as e:
                    logger.warning(f"Post-recall hook error (non-fatal): {e}")

            return result
        finally:
            recall_span_context.__exit__(None, None, None)

    async def _search_with_retries(
        self,
        bank_id: str,
        query: str,
        fact_type: list[str],
        thinking_budget: int,
        max_tokens: int,
        enable_trace: bool,
        question_date: datetime | None = None,
        include_entities: bool = False,
        max_entity_tokens: int = 500,
        include_chunks: bool = False,
        max_chunk_tokens: int = 8192,
        request_context: "RequestContext" = None,
        semaphore_wait: float = 0.0,
        tags: list[str] | None = None,
        tags_match: TagsMatch = "any",
        connection_budget: int | None = None,
        quiet: bool = False,
        include_source_facts: bool = False,
        max_source_facts_tokens: int = 4096,
    ) -> RecallResultModel:
        """
        Search implementation with modular retrieval and reranking.

        Architecture:
        1. Retrieval: 4-way parallel (semantic, keyword, graph, temporal graph)
        2. Merge: RRF to combine ranked lists
        3. Reranking: Pluggable strategy (heuristic or cross-encoder)
        4. Diversity: MMR with λ=0.5
        5. Chunks: Fetch chunks from top-scored results (BEFORE token filtering)
        6. Token Filter: Limit facts to max_tokens budget

        Args:
            bank_id: bank IDentifier
            query: Search query
            fact_type: Type of facts to search
            thinking_budget: Nodes to explore in graph traversal
            max_tokens: Maximum tokens to return (counts only 'text' field)
            enable_trace: Whether to return search trace (deprecated)
            include_entities: Whether to include entity observations
            max_entity_tokens: Maximum tokens for entity observations
            include_chunks: Whether to include raw chunks (fetched before max_tokens filtering)
            max_chunk_tokens: Maximum tokens for chunks

        Returns:
            RecallResultModel with results, trace, optional entities, and optional chunks
        """
        # Initialize tracer if requested
        from .search.tracer import SearchTracer

        tracer = (
            SearchTracer(query, thinking_budget, max_tokens, tags=tags, tags_match=tags_match) if enable_trace else None
        )
        if tracer:
            tracer.start()

        pool = await self._get_pool()
        recall_start = time.time()

        # Buffer logs for clean output in concurrent scenarios
        recall_id = f"{bank_id[:8]}-{int(time.time() * 1000) % 100000}"
        log_buffer = []
        tags_info = f", tags={tags}, tags_match={tags_match}" if tags else ""
        log_buffer.append(
            f"[RECALL {recall_id}] Query: '{query[:50]}...' (budget={thinking_budget}, max_tokens={max_tokens}{tags_info})"
        )

        # Import tracing utilities
        from ..tracing import get_tracer

        tracer_otel = get_tracer()

        try:
            # Step 1: Generate query embedding (for semantic search)
            step_start = time.time()

            embedding_span = tracer_otel.start_span("hindsight.recall_embedding")
            embedding_span.set_attribute("hindsight.bank_id", bank_id)
            embedding_span.set_attribute("hindsight.query", query[:100])

            try:
                query_embedding = embedding_utils.generate_embedding(self.embeddings, query)
                step_duration = time.time() - step_start
                log_buffer.append(f"  [1] Generate query embedding: {step_duration:.3f}s")
            finally:
                embedding_span.end()

            if tracer:
                tracer.record_query_embedding(query_embedding)
                tracer.add_phase_metric("generate_query_embedding", step_duration)

            # Step 2: Optimized parallel retrieval using batched queries
            # - Semantic + BM25 combined in 1 CTE query for ALL fact types
            # - Graph runs per fact type (complex traversal)
            # - Temporal runs per fact type (if constraint detected)
            step_start = time.time()
            query_embedding_str = str(query_embedding)

            from .search.retrieval import (
                get_default_graph_retriever,
                retrieve_all_fact_types_parallel,
            )

            # Track each retrieval start time
            retrieval_start = time.time()

            retrieval_span = tracer_otel.start_span("hindsight.recall_retrieval")
            retrieval_span.set_attribute("hindsight.bank_id", bank_id)
            retrieval_span.set_attribute("hindsight.fact_types", ",".join(fact_type))
            retrieval_span.set_attribute("hindsight.thinking_budget", thinking_budget)

            try:
                # Run optimized retrieval with connection budget
                config = get_config()
                effective_connection_budget = (
                    connection_budget if connection_budget is not None else config.recall_connection_budget
                )
                async with budgeted_operation(
                    max_connections=effective_connection_budget,
                    operation_id=f"recall-{recall_id}",
                ) as op:
                    budgeted_pool = op.wrap_pool(pool)
                    parallel_start = time.time()
                    multi_result = await retrieve_all_fact_types_parallel(
                        budgeted_pool,
                        query,
                        query_embedding_str,
                        bank_id,
                        fact_type,  # Pass all fact types at once
                        thinking_budget,
                        question_date,
                        self.query_analyzer,
                        tags=tags,
                        tags_match=tags_match,
                    )
                    parallel_duration = time.time() - parallel_start
            finally:
                retrieval_span.end()

            # Combine all results from all fact types and aggregate timings
            semantic_results = []
            bm25_results = []
            graph_results = []
            temporal_results = []
            aggregated_timings = {
                "semantic": 0.0,
                "bm25": 0.0,
                "graph": 0.0,
                "temporal": 0.0,
                "temporal_extraction": 0.0,
            }
            all_mpfp_timings = []

            detected_temporal_constraint = None
            max_conn_wait = multi_result.max_conn_wait
            for ft in fact_type:
                retrieval_result = multi_result.results_by_fact_type.get(ft)
                if not retrieval_result:
                    continue

                # Log fact types in this retrieval batch
                logger.debug(
                    f"[RECALL {recall_id}] Fact type '{ft}': semantic={len(retrieval_result.semantic)}, bm25={len(retrieval_result.bm25)}, graph={len(retrieval_result.graph)}, temporal={len(retrieval_result.temporal) if retrieval_result.temporal else 0}"
                )

                semantic_results.extend(retrieval_result.semantic)
                bm25_results.extend(retrieval_result.bm25)
                graph_results.extend(retrieval_result.graph)
                if retrieval_result.temporal:
                    temporal_results.extend(retrieval_result.temporal)
                # Track max timing for each method (since they run in parallel across fact types)
                for method, duration in retrieval_result.timings.items():
                    aggregated_timings[method] = max(aggregated_timings.get(method, 0.0), duration)
                # Capture temporal constraint (same across all fact types)
                if retrieval_result.temporal_constraint:
                    detected_temporal_constraint = retrieval_result.temporal_constraint

            # If no temporal results from any fact type, set to None
            if not temporal_results:
                temporal_results = None

            # Sort combined results by score (descending) so higher-scored results
            # get better ranks in the trace, regardless of fact type
            semantic_results.sort(key=lambda r: r.similarity if hasattr(r, "similarity") else 0, reverse=True)
            bm25_results.sort(key=lambda r: r.bm25_score if hasattr(r, "bm25_score") else 0, reverse=True)
            graph_results.sort(key=lambda r: r.activation if hasattr(r, "activation") else 0, reverse=True)
            if temporal_results:
                temporal_results.sort(
                    key=lambda r: r.combined_score if hasattr(r, "combined_score") else 0, reverse=True
                )

            retrieval_duration = time.time() - retrieval_start

            step_duration = time.time() - step_start
            total_retrievals = len(fact_type) * (4 if temporal_results else 3)
            # Format per-method timings
            timing_parts = [
                f"semantic={len(semantic_results)}({aggregated_timings['semantic']:.3f}s)",
                f"bm25={len(bm25_results)}({aggregated_timings['bm25']:.3f}s)",
                f"graph={len(graph_results)}({aggregated_timings['graph']:.3f}s)",
                f"temporal_extraction={aggregated_timings['temporal_extraction']:.3f}s",
            ]
            temporal_info = ""
            if detected_temporal_constraint:
                start_dt, end_dt = detected_temporal_constraint
                temporal_count = len(temporal_results) if temporal_results else 0
                timing_parts.append(f"temporal={temporal_count}({aggregated_timings['temporal']:.3f}s)")
                temporal_info = f" | temporal_range={start_dt.strftime('%Y-%m-%d')} to {end_dt.strftime('%Y-%m-%d')}"
            log_buffer.append(
                f"  [2] Parallel retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {parallel_duration:.3f}s{temporal_info}"
            )

            # Log graph retriever timing breakdown if available
            if all_mpfp_timings:
                retriever_name = get_default_graph_retriever().name.upper()
                mpfp_total = all_mpfp_timings[0]  # Take first fact type's timing as representative
                mpfp_parts = [
                    f"db_queries={mpfp_total.db_queries}",
                    f"edge_load={mpfp_total.edge_load_time:.3f}s",
                    f"edges={mpfp_total.edge_count}",
                    f"patterns={mpfp_total.pattern_count}",
                ]
                if mpfp_total.seeds_time > 0.01:
                    mpfp_parts.append(f"seeds={mpfp_total.seeds_time:.3f}s")
                if mpfp_total.fusion > 0.001:
                    mpfp_parts.append(f"fusion={mpfp_total.fusion:.3f}s")
                if mpfp_total.fetch > 0.001:
                    mpfp_parts.append(f"fetch={mpfp_total.fetch:.3f}s")
                log_buffer.append(f"      [{retriever_name}] {', '.join(mpfp_parts)}")
                # Log detailed hop timing for debugging slow queries
                if mpfp_total.hop_details:
                    for hd in mpfp_total.hop_details:
                        log_buffer.append(
                            f"        hop{hd['hop']}: exec={hd.get('exec_time', 0) * 1000:.0f}ms, "
                            f"uncached={hd.get('uncached_after_filter', 0)}, "
                            f"load={hd.get('load_time', 0) * 1000:.0f}ms, "
                            f"edges={hd.get('edges_loaded', 0)}"
                        )

            # Record temporal constraint in tracer if detected
            if tracer and detected_temporal_constraint:
                start_dt, end_dt = detected_temporal_constraint
                tracer.record_temporal_constraint(start_dt, end_dt)

            # Record retrieval results for tracer - per fact type
            if tracer:
                # Convert RetrievalResult to old tuple format for tracer
                def to_tuple_format(results):
                    return [(r.id, r.__dict__) for r in results]

                # Add retrieval results per fact type (to show parallel execution in UI)
                for ft_name in fact_type:
                    rr = multi_result.results_by_fact_type.get(ft_name)
                    if not rr:
                        continue

                    # Add semantic retrieval results for this fact type
                    tracer.add_retrieval_results(
                        method_name="semantic",
                        results=to_tuple_format(rr.semantic),
                        duration_seconds=rr.timings.get("semantic", 0.0),
                        score_field="similarity",
                        metadata={"limit": thinking_budget},
                        fact_type=ft_name,
                    )

                    # Add BM25 retrieval results for this fact type
                    tracer.add_retrieval_results(
                        method_name="bm25",
                        results=to_tuple_format(rr.bm25),
                        duration_seconds=rr.timings.get("bm25", 0.0),
                        score_field="bm25_score",
                        metadata={"limit": thinking_budget},
                        fact_type=ft_name,
                    )

                    # Add graph retrieval results for this fact type
                    tracer.add_retrieval_results(
                        method_name="graph",
                        results=to_tuple_format(rr.graph),
                        duration_seconds=rr.timings.get("graph", 0.0),
                        score_field="activation",
                        metadata={"budget": thinking_budget},
                        fact_type=ft_name,
                    )

                    # Add temporal retrieval results for this fact type
                    # Show temporal even with 0 results if constraint was detected
                    if rr.temporal is not None or rr.temporal_constraint is not None:
                        temporal_metadata = {"budget": thinking_budget}
                        if rr.temporal_constraint:
                            start_dt, end_dt = rr.temporal_constraint
                            temporal_metadata["constraint"] = {
                                "start": start_dt.isoformat() if start_dt else None,
                                "end": end_dt.isoformat() if end_dt else None,
                            }
                        tracer.add_retrieval_results(
                            method_name="temporal",
                            results=to_tuple_format(rr.temporal or []),
                            duration_seconds=rr.timings.get("temporal", 0.0),
                            score_field="temporal_score",
                            metadata=temporal_metadata,
                            fact_type=ft_name,
                        )

                # Record entry points (from semantic results) for legacy graph view
                for rank, retrieval in enumerate(semantic_results[:10], start=1):  # Top 10 as entry points
                    tracer.add_entry_point(retrieval.id, retrieval.text, retrieval.similarity or 0.0, rank)

                tracer.add_phase_metric(
                    "parallel_retrieval",
                    step_duration,
                    {
                        "semantic_count": len(semantic_results),
                        "bm25_count": len(bm25_results),
                        "graph_count": len(graph_results),
                        "temporal_count": len(temporal_results) if temporal_results else 0,
                    },
                )

            # Step 3: Merge with RRF
            step_start = time.time()
            from .search.fusion import reciprocal_rank_fusion

            fusion_span = tracer_otel.start_span("hindsight.recall_fusion")
            fusion_span.set_attribute("hindsight.bank_id", bank_id)
            fusion_span.set_attribute("hindsight.semantic_count", len(semantic_results))
            fusion_span.set_attribute("hindsight.bm25_count", len(bm25_results))
            fusion_span.set_attribute("hindsight.graph_count", len(graph_results))
            fusion_span.set_attribute("hindsight.temporal_count", len(temporal_results) if temporal_results else 0)

            try:
                # Merge 3 or 4 result lists depending on temporal constraint
                if temporal_results:
                    merged_candidates = reciprocal_rank_fusion(
                        [semantic_results, bm25_results, graph_results, temporal_results]
                    )
                else:
                    merged_candidates = reciprocal_rank_fusion([semantic_results, bm25_results, graph_results])

                step_duration = time.time() - step_start
                log_buffer.append(
                    f"  [3] RRF merge: {len(merged_candidates)} unique candidates in {step_duration:.3f}s"
                )
            finally:
                fusion_span.set_attribute("hindsight.merged_count", len(merged_candidates))
                fusion_span.end()

            if tracer:
                # Convert MergedCandidate to old tuple format for tracer
                tracer_merged = [
                    (mc.id, mc.retrieval.__dict__, {"rrf_score": mc.rrf_score, **mc.source_ranks})
                    for mc in merged_candidates
                ]
                tracer.add_rrf_merged(tracer_merged)
                tracer.add_phase_metric("rrf_merge", step_duration, {"candidates_merged": len(merged_candidates)})

            # Step 4: Rerank using cross-encoder (MergedCandidate -> ScoredResult)
            step_start = time.time()
            reranker_instance = self._cross_encoder_reranker

            rerank_span = tracer_otel.start_span("hindsight.recall_rerank")
            rerank_span.set_attribute("hindsight.bank_id", bank_id)
            rerank_span.set_attribute("hindsight.candidates_count", len(merged_candidates))

            scored_results: list = []
            pre_filtered_count = 0
            try:
                # Ensure reranker is initialized (for lazy initialization mode)
                await reranker_instance.ensure_initialized()

                # Pre-filter candidates to reduce reranking cost (RRF already provides good ranking)
                # This is especially important for remote rerankers with network latency
                reranker_max_candidates = get_config().reranker_max_candidates
                if len(merged_candidates) > reranker_max_candidates:
                    # Sort by RRF score and take top candidates
                    merged_candidates.sort(key=lambda mc: mc.rrf_score, reverse=True)
                    pre_filtered_count = len(merged_candidates) - reranker_max_candidates
                    merged_candidates = merged_candidates[:reranker_max_candidates]

                # Rerank using cross-encoder
                scored_results = await reranker_instance.rerank(query, merged_candidates)

                step_duration = time.time() - step_start
                pre_filter_note = f" (pre-filtered {pre_filtered_count})" if pre_filtered_count > 0 else ""
                log_buffer.append(
                    f"  [4] Reranking: {len(scored_results)} candidates scored in {step_duration:.3f}s{pre_filter_note}"
                )
            finally:
                rerank_span.set_attribute("hindsight.scored_count", len(scored_results))
                if pre_filtered_count > 0:
                    rerank_span.set_attribute("hindsight.pre_filtered_count", pre_filtered_count)
                rerank_span.end()

            # Step 4.5: Combine cross-encoder score with retrieval signals
            # This preserves retrieval work (RRF, temporal, recency) instead of pure cross-encoder ranking
            if scored_results:
                # Normalize RRF scores to [0, 1] range using min-max normalization
                rrf_scores = [sr.candidate.rrf_score for sr in scored_results]
                max_rrf = max(rrf_scores) if rrf_scores else 0.0
                min_rrf = min(rrf_scores) if rrf_scores else 0.0
                rrf_range = max_rrf - min_rrf  # Don't force to 1.0, let fallback handle it

                # Calculate recency based on occurred_start (more recent = higher score)
                now = utcnow()
                for sr in scored_results:
                    # Normalize RRF score (0-1 range, 0.5 if all same)
                    if rrf_range > 0:
                        sr.rrf_normalized = (sr.candidate.rrf_score - min_rrf) / rrf_range
                    else:
                        # All RRF scores are the same, use neutral value
                        sr.rrf_normalized = 0.5

                    # Calculate recency (decay over 365 days, minimum 0.1)
                    sr.recency = 0.5  # default for missing dates
                    if sr.retrieval.occurred_start:
                        occurred = sr.retrieval.occurred_start
                        if hasattr(occurred, "tzinfo") and occurred.tzinfo is None:
                            occurred = occurred.replace(tzinfo=UTC)
                        days_ago = (now - occurred).total_seconds() / 86400
                        sr.recency = max(0.1, 1.0 - (days_ago / 365))  # Linear decay over 1 year

                    # Get temporal proximity if available (already 0-1)
                    sr.temporal = (
                        sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5
                    )

                    # Weighted combination
                    # Cross-encoder: 60% (semantic relevance)
                    # RRF: 20% (retrieval consensus)
                    # Temporal proximity: 10% (time relevance for temporal queries)
                    # Recency: 10% (prefer recent facts)
                    sr.combined_score = (
                        0.6 * sr.cross_encoder_score_normalized
                        + 0.2 * sr.rrf_normalized
                        + 0.1 * sr.temporal
                        + 0.1 * sr.recency
                    )
                    sr.weight = sr.combined_score  # Update weight for final ranking

                # Re-sort by combined score
                scored_results.sort(key=lambda x: x.weight, reverse=True)
                log_buffer.append(
                    "  [4.6] Combined scoring: cross_encoder(0.6) + rrf(0.2) + temporal(0.1) + recency(0.1)"
                )

            # Add reranked results to tracer AFTER combined scoring (so normalized values are included)
            if tracer:
                results_dict = [sr.to_dict() for sr in scored_results]
                tracer_merged = [
                    (mc.id, mc.retrieval.__dict__, {"rrf_score": mc.rrf_score, **mc.source_ranks})
                    for mc in merged_candidates
                ]
                tracer.add_reranked(results_dict, tracer_merged)
                tracer.add_phase_metric(
                    "reranking",
                    step_duration,
                    {"reranker_type": "cross-encoder", "candidates_reranked": len(scored_results)},
                )

            # Step 5: Truncate to thinking_budget * 2 for token filtering
            rerank_limit = thinking_budget * 2
            top_scored = scored_results[:rerank_limit]
            log_buffer.append(f"  [5] Truncated to top {len(top_scored)} results")

            # Step 5.5: Fetch chunks from top-scored results (before token filtering)
            # Chunks are fetched independently of max_tokens filtering
            chunks_dict = None
            total_chunk_tokens = 0
            if include_chunks and top_scored:
                from .response_models import ChunkInfo

                # Collect chunk_ids in order of fact relevance (preserving order from top_scored)
                # Use a list to maintain order, but track seen chunks to avoid duplicates
                chunk_ids_ordered = []
                seen_chunk_ids = set()
                for sr in top_scored:
                    chunk_id = sr.retrieval.chunk_id
                    if chunk_id and chunk_id not in seen_chunk_ids:
                        chunk_ids_ordered.append(chunk_id)
                        seen_chunk_ids.add(chunk_id)

                if chunk_ids_ordered:
                    # Estimate batch size based on retain_chunk_size * 2 (rough estimate)
                    # Chunk sizes vary per document, so we fetch in batches until budget is exhausted
                    bank_config = await self._config_resolver.resolve_full_config(bank_id, request_context)
                    estimated_batch_size = max(1, (max_chunk_tokens // bank_config.retain_chunk_size) * 2)

                    chunks_dict = {}
                    encoding = _get_tiktoken_encoding()
                    chunk_offset = 0

                    # Fetch chunks in batches until we run out of budget or chunks
                    while chunk_offset < len(chunk_ids_ordered) and total_chunk_tokens < max_chunk_tokens:
                        # Get next batch of chunk IDs
                        batch_chunk_ids = chunk_ids_ordered[chunk_offset : chunk_offset + estimated_batch_size]
                        chunk_offset += estimated_batch_size

                        # Fetch chunk data from database
                        async with acquire_with_retry(pool) as conn:
                            chunks_rows = await conn.fetch(
                                f"""
                                SELECT chunk_id, chunk_text, chunk_index
                                FROM {fq_table("chunks")}
                                WHERE chunk_id = ANY($1::text[])
                                """,
                                batch_chunk_ids,
                            )

                        # Create a lookup dict for fast access (preserves order from batch_chunk_ids)
                        chunks_lookup = {row["chunk_id"]: row for row in chunks_rows}

                        # Process chunks in order, respecting token budget
                        for chunk_id in batch_chunk_ids:
                            if chunk_id not in chunks_lookup:
                                continue

                            row = chunks_lookup[chunk_id]
                            chunk_text = row["chunk_text"]
                            chunk_tokens = len(encoding.encode(chunk_text))

                            # Check if adding this chunk would exceed the limit
                            if total_chunk_tokens + chunk_tokens > max_chunk_tokens:
                                # Truncate the chunk to fit within the remaining budget
                                remaining_tokens = max_chunk_tokens - total_chunk_tokens
                                if remaining_tokens > 0:
                                    # Truncate to remaining tokens
                                    truncated_text = encoding.decode(encoding.encode(chunk_text)[:remaining_tokens])
                                    chunks_dict[chunk_id] = ChunkInfo(
                                        chunk_text=truncated_text, chunk_index=row["chunk_index"], truncated=True
                                    )
                                    total_chunk_tokens = max_chunk_tokens
                                # Budget exhausted - stop fetching more batches
                                break
                            else:
                                chunks_dict[chunk_id] = ChunkInfo(
                                    chunk_text=chunk_text, chunk_index=row["chunk_index"], truncated=False
                                )
                                total_chunk_tokens += chunk_tokens

                        # If we hit the budget limit in this batch, stop fetching more batches
                        if total_chunk_tokens >= max_chunk_tokens:
                            break

            # Step 6: Token budget filtering
            step_start = time.time()

            # Convert to dict for token filtering (backward compatibility)
            top_dicts = [sr.to_dict() for sr in top_scored]
            filtered_dicts, total_tokens = self._filter_by_token_budget(top_dicts, max_tokens)

            # Convert back to list of IDs and filter scored_results
            filtered_ids = {d["id"] for d in filtered_dicts}
            top_scored = [sr for sr in top_scored if sr.id in filtered_ids]

            step_duration = time.time() - step_start
            log_buffer.append(
                f"  [6] Token filtering: {len(top_scored)} results, {total_tokens}/{max_tokens} tokens in {step_duration:.3f}s"
            )

            if tracer:
                tracer.add_phase_metric(
                    "token_filtering",
                    step_duration,
                    {"results_selected": len(top_scored), "tokens_used": total_tokens, "max_tokens": max_tokens},
                )

            # Record visits for all retrieved nodes
            if tracer:
                for sr in scored_results:
                    tracer.visit_node(
                        node_id=sr.id,
                        text=sr.retrieval.text,
                        context=sr.retrieval.context or "",
                        event_date=sr.retrieval.occurred_start,
                        is_entry_point=(sr.id in [ep.node_id for ep in tracer.entry_points]),
                        parent_node_id=None,  # In parallel retrieval, there's no clear parent
                        link_type=None,
                        link_weight=None,
                        activation=sr.candidate.rrf_score,  # Use RRF score as activation
                        semantic_similarity=sr.retrieval.similarity or 0.0,
                        recency=sr.recency,
                        frequency=0.0,
                        final_weight=sr.weight,
                    )

            # Log fact_type distribution in results
            fact_type_counts = {}
            for sr in top_scored:
                ft = sr.retrieval.fact_type
                fact_type_counts[ft] = fact_type_counts.get(ft, 0) + 1

            fact_type_summary = ", ".join([f"{ft}={count}" for ft, count in sorted(fact_type_counts.items())])

            # Convert ScoredResult to dicts with ISO datetime strings
            top_results_dicts = []
            for sr in top_scored:
                result_dict = sr.to_dict()
                # Convert datetime objects to ISO strings for JSON serialization
                if result_dict.get("occurred_start"):
                    occurred_start = result_dict["occurred_start"]
                    result_dict["occurred_start"] = (
                        occurred_start.isoformat() if hasattr(occurred_start, "isoformat") else occurred_start
                    )
                if result_dict.get("occurred_end"):
                    occurred_end = result_dict["occurred_end"]
                    result_dict["occurred_end"] = (
                        occurred_end.isoformat() if hasattr(occurred_end, "isoformat") else occurred_end
                    )
                if result_dict.get("mentioned_at"):
                    mentioned_at = result_dict["mentioned_at"]
                    result_dict["mentioned_at"] = (
                        mentioned_at.isoformat() if hasattr(mentioned_at, "isoformat") else mentioned_at
                    )
                top_results_dicts.append(result_dict)

            # Fetch source facts for observation-type results (mirrors chunks pattern)
            source_fact_ids_by_obs: dict[str, list[str]] = {}  # obs_id -> [source_id, ...]
            source_facts_dict: dict[str, MemoryFact] | None = None
            if include_source_facts:
                observation_ids = [uuid.UUID(sr.id) for sr in top_scored if sr.retrieval.fact_type == "observation"]
                if observation_ids:
                    async with acquire_with_retry(pool) as sf_conn:
                        # Fetch source_memory_ids for all observation results
                        obs_rows = await sf_conn.fetch(
                            f"""
                            SELECT id, source_memory_ids
                            FROM {fq_table("memory_units")}
                            WHERE id = ANY($1::uuid[]) AND fact_type = 'observation'
                            """,
                            observation_ids,
                        )

                        # Collect unique source IDs in order of first appearance
                        seen_source_ids: set[str] = set()
                        source_ids_ordered: list[str] = []
                        for obs_row in obs_rows:
                            obs_id = str(obs_row["id"])
                            sids = [str(s) for s in (obs_row["source_memory_ids"] or [])]
                            source_fact_ids_by_obs[obs_id] = sids
                            for sid in sids:
                                if sid not in seen_source_ids:
                                    source_ids_ordered.append(sid)
                                    seen_source_ids.add(sid)

                        # Fetch source fact content up to token budget
                        if source_ids_ordered:
                            import uuid as uuid_module

                            source_rows = await sf_conn.fetch(
                                f"""
                                SELECT id, text, fact_type, context, occurred_start, occurred_end,
                                       mentioned_at, document_id, chunk_id, tags
                                FROM {fq_table("memory_units")}
                                WHERE id = ANY($1::uuid[])
                                """,
                                [uuid_module.UUID(sid) for sid in source_ids_ordered],
                            )
                            source_row_by_id = {str(r["id"]): r for r in source_rows}

                            encoding = _get_tiktoken_encoding()
                            source_facts_dict = {}
                            total_source_tokens = 0
                            for sid in source_ids_ordered:
                                if sid not in source_row_by_id:
                                    continue
                                r = source_row_by_id[sid]
                                fact_tokens = len(encoding.encode(r["text"]))
                                if (
                                    max_source_facts_tokens >= 0
                                    and total_source_tokens + fact_tokens > max_source_facts_tokens
                                ):
                                    break
                                source_facts_dict[sid] = MemoryFact(
                                    id=sid,
                                    text=r["text"],
                                    fact_type=r["fact_type"],
                                    context=r["context"],
                                    occurred_start=r["occurred_start"].isoformat() if r["occurred_start"] else None,
                                    occurred_end=r["occurred_end"].isoformat() if r["occurred_end"] else None,
                                    mentioned_at=r["mentioned_at"].isoformat() if r["mentioned_at"] else None,
                                    document_id=r["document_id"],
                                    chunk_id=str(r["chunk_id"]) if r["chunk_id"] else None,
                                    tags=r["tags"] or None,
                                )
                                total_source_tokens += fact_tokens

            # Get entities for each fact if include_entities is requested
            fact_entity_map = {}  # unit_id -> list of (entity_id, entity_name)
            if include_entities and top_scored:
                unit_ids = [uuid.UUID(sr.id) for sr in top_scored]
                if unit_ids:
                    async with acquire_with_retry(pool) as entity_conn:
                        entity_rows = await entity_conn.fetch(
                            f"""
                            SELECT ue.unit_id, e.id as entity_id, e.canonical_name
                            FROM {fq_table("unit_entities")} ue
                            JOIN {fq_table("entities")} e ON ue.entity_id = e.id
                            WHERE ue.unit_id = ANY($1::uuid[])
                            """,
                            unit_ids,
                        )
                        for row in entity_rows:
                            unit_id = str(row["unit_id"])
                            if unit_id not in fact_entity_map:
                                fact_entity_map[unit_id] = []
                            fact_entity_map[unit_id].append(
                                {"entity_id": str(row["entity_id"]), "canonical_name": row["canonical_name"]}
                            )

            # Convert results to MemoryFact objects
            memory_facts = []
            for result_dict in top_results_dicts:
                result_id = str(result_dict.get("id"))
                # Get entity names for this fact
                entity_names = None
                if include_entities and result_id in fact_entity_map:
                    entity_names = [e["canonical_name"] for e in fact_entity_map[result_id]]

                memory_facts.append(
                    MemoryFact(
                        id=result_id,
                        text=result_dict.get("text"),
                        fact_type=result_dict.get("fact_type", "world"),
                        entities=entity_names,
                        context=result_dict.get("context"),
                        occurred_start=result_dict.get("occurred_start"),
                        occurred_end=result_dict.get("occurred_end"),
                        mentioned_at=result_dict.get("mentioned_at"),
                        document_id=result_dict.get("document_id"),
                        chunk_id=result_dict.get("chunk_id"),
                        tags=result_dict.get("tags"),
                        source_fact_ids=source_fact_ids_by_obs.get(result_id) if include_source_facts else None,
                    )
                )

            # Fetch entity observations if requested
            entities_dict = None
            total_entity_tokens = 0
            if include_entities and fact_entity_map:
                # Collect unique entities in order of fact relevance (preserving order from top_scored)
                entities_ordered = []  # list of (entity_id, entity_name) tuples
                seen_entity_ids = set()

                for sr in top_scored:
                    unit_id = sr.id
                    if unit_id in fact_entity_map:
                        for entity in fact_entity_map[unit_id]:
                            entity_id = entity["entity_id"]
                            entity_name = entity["canonical_name"]
                            if entity_id not in seen_entity_ids:
                                entities_ordered.append((entity_id, entity_name))
                                seen_entity_ids.add(entity_id)

                # Return entities with empty observations (summaries now live in mental models)
                entities_dict = {}
                for entity_id, entity_name in entities_ordered:
                    entities_dict[entity_name] = EntityState(
                        entity_id=entity_id,
                        canonical_name=entity_name,
                        observations=[],  # Mental models provide this now
                    )

            # Finalize trace if enabled
            trace_dict = None
            if tracer:
                trace = tracer.finalize(top_results_dicts)
                trace_dict = trace.to_dict() if trace else None

            # Log final recall stats
            total_time = time.time() - recall_start
            num_chunks = len(chunks_dict) if chunks_dict else 0
            num_entities = len(entities_dict) if entities_dict else 0
            # Include wait times in log if significant
            wait_parts = []
            if semaphore_wait > 0.01:
                wait_parts.append(f"sem={semaphore_wait:.3f}s")
            if max_conn_wait > 0.01:
                wait_parts.append(f"conn={max_conn_wait:.3f}s")
            wait_info = f" | waits: {', '.join(wait_parts)}" if wait_parts else ""
            log_buffer.append(
                f"[RECALL {recall_id}] Complete: {len(top_scored)} facts ({total_tokens} tok), {num_chunks} chunks ({total_chunk_tokens} tok), {num_entities} entities ({total_entity_tokens} tok) | {fact_type_summary} | {total_time:.3f}s{wait_info}"
            )
            if not quiet:
                logger.info("\n" + "\n".join(log_buffer))

            return RecallResultModel(
                results=memory_facts,
                trace=trace_dict,
                entities=entities_dict,
                chunks=chunks_dict,
                source_facts=source_facts_dict,
            )

        except Exception as e:
            log_buffer.append(f"[RECALL {recall_id}] ERROR after {time.time() - recall_start:.3f}s: {str(e)}")
            if not quiet:
                logger.error("\n" + "\n".join(log_buffer))
            raise Exception(f"Failed to search memories: {str(e)}")

    def _filter_by_token_budget(
        self, results: list[dict[str, Any]], max_tokens: int
    ) -> tuple[list[dict[str, Any]], int]:
        """
        Filter results to fit within token budget.

        Counts tokens only for the 'text' field using tiktoken (cl100k_base encoding).
        Stops before including a fact that would exceed the budget.

        Args:
            results: List of search results
            max_tokens: Maximum tokens allowed

        Returns:
            Tuple of (filtered_results, total_tokens_used)
        """
        encoding = _get_tiktoken_encoding()

        filtered_results = []
        total_tokens = 0

        for result in results:
            text = result.get("text", "")
            text_tokens = len(encoding.encode(text))

            # Check if adding this result would exceed budget
            if total_tokens + text_tokens <= max_tokens:
                filtered_results.append(result)
                total_tokens += text_tokens
            else:
                # Stop before including a fact that would exceed limit
                break

        return filtered_results, total_tokens

    async def get_document(
        self,
        document_id: str,
        bank_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """
        Retrieve document metadata and statistics.

        Args:
            document_id: Document ID to retrieve
            bank_id: bank ID that owns the document
            request_context: Request context for authentication.

        Returns:
            Dictionary with document info or None if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            doc = await conn.fetchrow(
                f"""
                SELECT d.id, d.bank_id, d.original_text, d.content_hash,
                       d.created_at, d.updated_at, d.tags, COUNT(mu.id) as unit_count
                FROM {fq_table("documents")} d
                LEFT JOIN {fq_table("memory_units")} mu ON mu.document_id = d.id
                WHERE d.id = $1 AND d.bank_id = $2
                GROUP BY d.id, d.bank_id, d.original_text, d.content_hash, d.created_at, d.updated_at, d.tags
                """,
                document_id,
                bank_id,
            )

            if not doc:
                return None

            return {
                "id": doc["id"],
                "bank_id": doc["bank_id"],
                "original_text": doc["original_text"],
                "content_hash": doc["content_hash"],
                "memory_unit_count": doc["unit_count"],
                "created_at": doc["created_at"].isoformat() if doc["created_at"] else None,
                "updated_at": doc["updated_at"].isoformat() if doc["updated_at"] else None,
                "tags": list(doc["tags"]) if doc["tags"] else [],
            }

    async def delete_document(
        self,
        document_id: str,
        bank_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, int]:
        """
        Delete a document and all its associated memory units and links.

        Args:
            document_id: Document ID to delete
            bank_id: bank ID that owns the document
            request_context: Request context for authentication.

        Returns:
            Dictionary with counts of deleted items
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            async with conn.transaction():
                # Get memory unit IDs before deletion (for mental model invalidation)
                unit_rows = await conn.fetch(
                    f"SELECT id FROM {fq_table('memory_units')} WHERE document_id = $1", document_id
                )
                unit_ids = [str(row["id"]) for row in unit_rows]
                units_count = len(unit_ids)

                # Delete document (cascades to memory_units and all their links)
                deleted = await conn.fetchval(
                    f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id",
                    document_id,
                    bank_id,
                )

                # Invalidate deleted fact IDs from mental models
                if deleted and unit_ids:
                    await self._invalidate_facts_from_mental_models(conn, bank_id, unit_ids)

                return {"document_deleted": 1 if deleted else 0, "memory_units_deleted": units_count if deleted else 0}

    async def delete_memory_unit(
        self,
        unit_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """
        Delete a single memory unit and all its associated links.

        Due to CASCADE DELETE constraints, this will automatically delete:
        - All links from this unit (memory_links where from_unit_id = unit_id)
        - All links to this unit (memory_links where to_unit_id = unit_id)
        - All entity associations (unit_entities where unit_id = unit_id)

        Args:
            unit_id: UUID of the memory unit to delete
            request_context: Request context for authentication.

        Returns:
            Dictionary with deletion result
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            async with conn.transaction():
                # Get bank_id before deletion (for mental model invalidation)
                bank_id = await conn.fetchval(f"SELECT bank_id FROM {fq_table('memory_units')} WHERE id = $1", unit_id)

                # Delete the memory unit (cascades to links and associations)
                deleted = await conn.fetchval(
                    f"DELETE FROM {fq_table('memory_units')} WHERE id = $1 RETURNING id", unit_id
                )

                # Invalidate deleted fact ID from mental models
                if deleted and bank_id:
                    await self._invalidate_facts_from_mental_models(conn, bank_id, [str(deleted)])

                return {
                    "success": deleted is not None,
                    "unit_id": str(deleted) if deleted else None,
                    "message": "Memory unit and all its links deleted successfully"
                    if deleted
                    else "Memory unit not found",
                }

    async def delete_bank(
        self,
        bank_id: str,
        fact_type: str | None = None,
        *,
        request_context: "RequestContext",
    ) -> dict[str, int]:
        """
        Delete all data for a specific agent (multi-tenant cleanup).

        This is much more efficient than dropping all tables and allows
        multiple agents to coexist in the same database.

        Deletes (with CASCADE):
        - All memory units for this bank (optionally filtered by fact_type)
        - All entities for this bank (if deleting all memory units)
        - All associated links, unit-entity associations, and co-occurrences

        Args:
            bank_id: bank ID to delete
            fact_type: Optional fact type filter (world, experience, opinion). If provided, only deletes memories of that type.
            request_context: Request context for authentication.

        Returns:
            Dictionary with counts of deleted items
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            # Ensure connection is not in read-only mode (can happen with connection poolers)
            await conn.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
            async with conn.transaction():
                try:
                    if fact_type:
                        # Delete only memories of a specific fact type
                        units_count = await conn.fetchval(
                            f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = $2",
                            bank_id,
                            fact_type,
                        )
                        await conn.execute(
                            f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = $2",
                            bank_id,
                            fact_type,
                        )

                        # Note: We don't delete entities when fact_type is specified,
                        # as they may be referenced by other memory units
                        return {"memory_units_deleted": units_count, "entities_deleted": 0}
                    else:
                        # Delete all data for the bank
                        units_count = await conn.fetchval(
                            f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1", bank_id
                        )
                        entities_count = await conn.fetchval(
                            f"SELECT COUNT(*) FROM {fq_table('entities')} WHERE bank_id = $1", bank_id
                        )
                        documents_count = await conn.fetchval(
                            f"SELECT COUNT(*) FROM {fq_table('documents')} WHERE bank_id = $1", bank_id
                        )

                        # Delete documents (cascades to chunks)
                        await conn.execute(f"DELETE FROM {fq_table('documents')} WHERE bank_id = $1", bank_id)

                        # Delete memory units (cascades to unit_entities, memory_links)
                        await conn.execute(f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1", bank_id)

                        # Delete entities (cascades to unit_entities, entity_cooccurrences, memory_links with entity_id)
                        await conn.execute(f"DELETE FROM {fq_table('entities')} WHERE bank_id = $1", bank_id)

                        # Delete the bank profile itself
                        await conn.execute(f"DELETE FROM {fq_table('banks')} WHERE bank_id = $1", bank_id)

                        return {
                            "memory_units_deleted": units_count,
                            "entities_deleted": entities_count,
                            "documents_deleted": documents_count,
                            "bank_deleted": True,
                        }

                except Exception as e:
                    raise Exception(f"Failed to delete agent data: {str(e)}")

    async def clear_observations(
        self,
        bank_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, int]:
        """
        Clear all observations for a bank (consolidated knowledge).

        Args:
            bank_id: Bank ID to clear observations for
            request_context: Request context for authentication.

        Returns:
            Dictionary with count of deleted observations
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            async with conn.transaction():
                # Count observations before deletion
                count = await conn.fetchval(
                    f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'",
                    bank_id,
                )

                # Delete all observations
                await conn.execute(
                    f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'",
                    bank_id,
                )

                # Reset consolidation timestamp
                await conn.execute(
                    f"UPDATE {fq_table('banks')} SET last_consolidated_at = NULL WHERE bank_id = $1",
                    bank_id,
                )

                return {"deleted_count": count or 0}

    async def run_consolidation(
        self,
        bank_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, int]:
        """
        Run memory consolidation to create/update mental models.

        Args:
            bank_id: Bank ID to run consolidation for
            request_context: Request context for authentication.

        Returns:
            Dictionary with consolidation stats
        """
        await self._authenticate_tenant(request_context)

        from .consolidation import run_consolidation_job

        # Create parent span for consolidation operation
        with create_operation_span("consolidation", bank_id):
            result = await run_consolidation_job(
                memory_engine=self,
                bank_id=bank_id,
                request_context=request_context,
            )

            return {
                "processed": result.get("processed", 0),
                "created": result.get("created", 0),
                "updated": result.get("updated", 0),
                "skipped": result.get("skipped", 0),
            }

    async def get_graph_data(
        self,
        bank_id: str | None = None,
        fact_type: str | None = None,
        *,
        limit: int = 1000,
        request_context: "RequestContext",
    ):
        """
        Get graph data for visualization.

        Args:
            bank_id: Filter by bank ID
            fact_type: Filter by fact type (world, experience, opinion)
            limit: Maximum number of items to return (default: 1000)
            request_context: Request context for authentication.

        Returns:
            Dict with nodes, edges, table_rows, total_units, and limit
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            # Get memory units, optionally filtered by bank_id and fact_type
            query_conditions = []
            query_params = []
            param_count = 0

            if bank_id:
                param_count += 1
                query_conditions.append(f"bank_id = ${param_count}")
                query_params.append(bank_id)

            if fact_type:
                param_count += 1
                query_conditions.append(f"fact_type = ${param_count}")
                query_params.append(fact_type)

            where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""

            # Get total count first
            total_count_result = await conn.fetchrow(
                f"""
                SELECT COUNT(*) as total
                FROM {fq_table("memory_units")}
                {where_clause}
            """,
                *query_params,
            )
            total_count = total_count_result["total"] if total_count_result else 0

            # Get units with limit
            param_count += 1
            units = await conn.fetch(
                f"""
                SELECT id, text, event_date, context, occurred_start, occurred_end, mentioned_at, document_id, chunk_id, fact_type, tags, created_at, proof_count, source_memory_ids
                FROM {fq_table("memory_units")}
                {where_clause}
                ORDER BY mentioned_at DESC NULLS LAST, event_date DESC
                LIMIT ${param_count}
            """,
                *query_params,
                limit,
            )

            # Get links, filtering to only include links between units of the selected agent
            # Use DISTINCT ON with LEAST/GREATEST to deduplicate bidirectional links
            unit_ids = [row["id"] for row in units]
            unit_id_set = set(unit_ids)

            # Collect source memory IDs from observations
            source_memory_ids = []
            for unit in units:
                if unit["source_memory_ids"]:
                    source_memory_ids.extend(unit["source_memory_ids"])
            source_memory_ids = list(set(source_memory_ids))  # Deduplicate

            # Fetch links involving both visible units AND source memories
            all_relevant_ids = unit_ids + source_memory_ids
            if all_relevant_ids:
                links = await conn.fetch(
                    f"""
                    SELECT DISTINCT ON (LEAST(ml.from_unit_id, ml.to_unit_id), GREATEST(ml.from_unit_id, ml.to_unit_id), ml.link_type, COALESCE(ml.entity_id, '00000000-0000-0000-0000-000000000000'::uuid))
                        ml.from_unit_id,
                        ml.to_unit_id,
                        ml.link_type,
                        ml.weight,
                        e.canonical_name as entity_name
                    FROM {fq_table("memory_links")} ml
                    LEFT JOIN {fq_table("entities")} e ON ml.entity_id = e.id
                    WHERE ml.from_unit_id = ANY($1::uuid[]) OR ml.to_unit_id = ANY($1::uuid[])
                    ORDER BY LEAST(ml.from_unit_id, ml.to_unit_id), GREATEST(ml.from_unit_id, ml.to_unit_id), ml.link_type, COALESCE(ml.entity_id, '00000000-0000-0000-0000-000000000000'::uuid), ml.weight DESC
                """,
                    all_relevant_ids,
                )
            else:
                links = []

            # Copy links from source memories to observations
            # Observations inherit links from their source memories via source_memory_ids
            # Build a map from source_id to observation_ids
            source_to_observations = {}
            for unit in units:
                if unit["source_memory_ids"]:
                    for source_id in unit["source_memory_ids"]:
                        if source_id not in source_to_observations:
                            source_to_observations[source_id] = []
                        source_to_observations[source_id].append(unit["id"])

            copied_links = []
            for link in links:
                from_id = link["from_unit_id"]
                to_id = link["to_unit_id"]

                # Get observations that should inherit this link
                from_observations = source_to_observations.get(from_id, [])
                to_observations = source_to_observations.get(to_id, [])

                # If from_id is a source memory, copy links to its observations
                if from_observations:
                    for obs_id in from_observations:
                        # Only include if the target is visible
                        if to_id in unit_id_set or to_observations:
                            target = to_observations[0] if to_observations and to_id not in unit_id_set else to_id
                            if target in unit_id_set:
                                copied_links.append(
                                    {
                                        "from_unit_id": obs_id,
                                        "to_unit_id": target,
                                        "link_type": link["link_type"],
                                        "weight": link["weight"],
                                        "entity_name": link["entity_name"],
                                    }
                                )

                # If to_id is a source memory, copy links to its observations
                if to_observations and from_id in unit_id_set:
                    for obs_id in to_observations:
                        copied_links.append(
                            {
                                "from_unit_id": from_id,
                                "to_unit_id": obs_id,
                                "link_type": link["link_type"],
                                "weight": link["weight"],
                                "entity_name": link["entity_name"],
                            }
                        )

            # Keep only direct links between visible nodes
            direct_links = [
                link for link in links if link["from_unit_id"] in unit_id_set and link["to_unit_id"] in unit_id_set
            ]

            # Get entity information
            unit_entities = await conn.fetch(f"""
                SELECT ue.unit_id, e.canonical_name
                FROM {fq_table("unit_entities")} ue
                JOIN {fq_table("entities")} e ON ue.entity_id = e.id
                ORDER BY ue.unit_id
            """)

        # Build entity mapping
        entity_map = {}
        for row in unit_entities:
            unit_id = row["unit_id"]
            entity_name = row["canonical_name"]
            if unit_id not in entity_map:
                entity_map[unit_id] = []
            entity_map[unit_id].append(entity_name)

        # For observations, inherit entities from source memories
        for unit in units:
            if unit["source_memory_ids"] and unit["id"] not in entity_map:
                # Collect entities from all source memories
                source_entities = []
                for source_id in unit["source_memory_ids"]:
                    if source_id in entity_map:
                        source_entities.extend(entity_map[source_id])
                if source_entities:
                    # Deduplicate while preserving order
                    entity_map[unit["id"]] = list(dict.fromkeys(source_entities))

        # Build nodes
        nodes = []
        for row in units:
            unit_id = row["id"]
            text = row["text"]
            event_date = row["event_date"]
            context = row["context"]

            entities = entity_map.get(unit_id, [])
            entity_count = len(entities)

            # Color by entity count
            if entity_count == 0:
                color = "#e0e0e0"
            elif entity_count == 1:
                color = "#90caf9"
            else:
                color = "#42a5f5"

            nodes.append(
                {
                    "data": {
                        "id": str(unit_id),
                        "label": f"{text[:30]}..." if len(text) > 30 else text,
                        "text": text,
                        "date": event_date.isoformat() if event_date else "",
                        "context": context if context else "",
                        "entities": ", ".join(entities) if entities else "None",
                        "color": color,
                    }
                }
            )

        # Build edges (combine direct links and copied links from sources)
        edges = []
        all_links = direct_links + copied_links
        for row in all_links:
            from_id = str(row["from_unit_id"])
            to_id = str(row["to_unit_id"])
            link_type = row["link_type"]
            weight = row["weight"]
            entity_name = row.get("entity_name")

            # Color by link type
            if link_type == "temporal":
                color = "#00bcd4"
                line_style = "dashed"
            elif link_type == "semantic":
                color = "#ff69b4"
                line_style = "solid"
            elif link_type == "entity":
                color = "#ffd700"
                line_style = "solid"
            else:
                color = "#999999"
                line_style = "solid"

            edges.append(
                {
                    "data": {
                        "id": f"{from_id}-{to_id}-{link_type}",
                        "source": from_id,
                        "target": to_id,
                        "linkType": link_type,
                        "weight": weight,
                        "entityName": entity_name if entity_name else "",
                        "color": color,
                        "lineStyle": line_style,
                    }
                }
            )

        # Build table rows
        table_rows = []
        for row in units:
            unit_id = row["id"]
            entities = entity_map.get(unit_id, [])

            table_rows.append(
                {
                    "id": str(unit_id),
                    "text": row["text"],
                    "context": row["context"] if row["context"] else "N/A",
                    "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
                    "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
                    "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
                    "date": row["event_date"].strftime("%Y-%m-%d %H:%M")
                    if row["event_date"]
                    else "N/A",  # Deprecated, kept for backwards compatibility
                    "entities": ", ".join(entities) if entities else "None",
                    "document_id": row["document_id"],
                    "chunk_id": row["chunk_id"] if row["chunk_id"] else None,
                    "fact_type": row["fact_type"],
                    "tags": list(row["tags"]) if row["tags"] else [],
                    "created_at": row["created_at"].isoformat() if row["created_at"] else None,
                    "proof_count": row["proof_count"] if row["proof_count"] else None,
                }
            )

        return {"nodes": nodes, "edges": edges, "table_rows": table_rows, "total_units": total_count, "limit": limit}

    async def list_memory_units(
        self,
        bank_id: str,
        *,
        fact_type: str | None = None,
        search_query: str | None = None,
        limit: int = 100,
        offset: int = 0,
        request_context: "RequestContext",
    ):
        """
        List memory units for table view with optional full-text search.

        Args:
            bank_id: Filter by bank ID
            fact_type: Filter by fact type (world, experience, opinion)
            search_query: Full-text search query (searches text and context fields)
            limit: Maximum number of results to return
            offset: Offset for pagination
            request_context: Request context for authentication.

        Returns:
            Dict with items (list of memory units) and total count
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            # Build query conditions
            query_conditions = []
            query_params = []
            param_count = 0

            if bank_id:
                param_count += 1
                query_conditions.append(f"bank_id = ${param_count}")
                query_params.append(bank_id)

            if fact_type:
                param_count += 1
                query_conditions.append(f"fact_type = ${param_count}")
                query_params.append(fact_type)

            if search_query:
                # Full-text search on text and context fields using ILIKE
                param_count += 1
                query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})")
                query_params.append(f"%{search_query}%")

            where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""

            # Get total count
            count_query = f"""
                SELECT COUNT(*) as total
                FROM {fq_table("memory_units")}
                {where_clause}
            """
            count_result = await conn.fetchrow(count_query, *query_params)
            total = count_result["total"]

            # Get units with limit and offset
            param_count += 1
            limit_param = f"${param_count}"
            query_params.append(limit)

            param_count += 1
            offset_param = f"${param_count}"
            query_params.append(offset)

            units = await conn.fetch(
                f"""
                SELECT id, text, event_date, context, fact_type, mentioned_at, occurred_start, occurred_end, chunk_id
                FROM {fq_table("memory_units")}
                {where_clause}
                ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
                LIMIT {limit_param} OFFSET {offset_param}
            """,
                *query_params,
            )

            # Get entity information for these units
            if units:
                unit_ids = [row["id"] for row in units]
                unit_entities = await conn.fetch(
                    f"""
                    SELECT ue.unit_id, e.canonical_name
                    FROM {fq_table("unit_entities")} ue
                    JOIN {fq_table("entities")} e ON ue.entity_id = e.id
                    WHERE ue.unit_id = ANY($1::uuid[])
                    ORDER BY ue.unit_id
                """,
                    unit_ids,
                )
            else:
                unit_entities = []

            # Build entity mapping
            entity_map = {}
            for row in unit_entities:
                unit_id = row["unit_id"]
                entity_name = row["canonical_name"]
                if unit_id not in entity_map:
                    entity_map[unit_id] = []
                entity_map[unit_id].append(entity_name)

            # Build result items
            items = []
            for row in units:
                unit_id = row["id"]
                entities = entity_map.get(unit_id, [])

                items.append(
                    {
                        "id": str(unit_id),
                        "text": row["text"],
                        "context": row["context"] if row["context"] else "",
                        "date": row["event_date"].isoformat() if row["event_date"] else "",
                        "fact_type": row["fact_type"],
                        "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
                        "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
                        "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
                        "entities": ", ".join(entities) if entities else "",
                        "chunk_id": row["chunk_id"] if row["chunk_id"] else None,
                    }
                )

            return {"items": items, "total": total, "limit": limit, "offset": offset}

    async def get_memory_unit(
        self,
        bank_id: str,
        memory_id: str,
        request_context: "RequestContext",
    ):
        """
        Get a single memory unit by ID.

        Args:
            bank_id: Bank ID
            memory_id: Memory unit ID
            request_context: Request context for authentication.

        Returns:
            Dict with memory unit data or None if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            # 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, source_memory_ids
                FROM {fq_table("memory_units")}
                WHERE id = $1 AND bank_id = $2
                """,
                memory_id,
                bank_id,
            )

            if not row:
                return None

            # Get entity information
            entities_rows = await conn.fetch(
                f"""
                SELECT e.canonical_name
                FROM {fq_table("unit_entities")} ue
                JOIN {fq_table("entities")} e ON ue.entity_id = e.id
                WHERE ue.unit_id = $1
                """,
                row["id"],
            )
            entities = [r["canonical_name"] for r in entities_rows]

            result = {
                "id": str(row["id"]),
                "text": row["text"],
                "context": row["context"] if row["context"] else "",
                "date": row["event_date"].isoformat() if row["event_date"] else "",
                "type": row["fact_type"],
                "mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
                "occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
                "occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
                "entities": entities,
                "document_id": row["document_id"] if row["document_id"] else None,
                "chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None,
                "tags": row["tags"] if row["tags"] else [],
            }

            # For observations, include source_memory_ids and fetch source_memories
            if row["fact_type"] == "observation" 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,
        *,
        search_query: str | None = None,
        limit: int = 100,
        offset: int = 0,
        request_context: "RequestContext",
    ):
        """
        List documents with optional search and pagination.

        Args:
            bank_id: bank ID (required)
            search_query: Search in document ID
            limit: Maximum number of results
            offset: Offset for pagination
            request_context: Request context for authentication.

        Returns:
            Dict with items (list of documents without original_text) and total count
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            # Build query conditions
            query_conditions = []
            query_params = []
            param_count = 0

            param_count += 1
            query_conditions.append(f"bank_id = ${param_count}")
            query_params.append(bank_id)

            if search_query:
                # Search in document ID
                param_count += 1
                query_conditions.append(f"id ILIKE ${param_count}")
                query_params.append(f"%{search_query}%")

            where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""

            # Get total count
            count_query = f"""
                SELECT COUNT(*) as total
                FROM {fq_table("documents")}
                {where_clause}
            """
            count_result = await conn.fetchrow(count_query, *query_params)
            total = count_result["total"]

            # Get documents with limit and offset (without original_text for performance)
            param_count += 1
            limit_param = f"${param_count}"
            query_params.append(limit)

            param_count += 1
            offset_param = f"${param_count}"
            query_params.append(offset)

            documents = await conn.fetch(
                f"""
                SELECT
                    id,
                    bank_id,
                    content_hash,
                    created_at,
                    updated_at,
                    LENGTH(original_text) as text_length,
                    retain_params,
                    tags
                FROM {fq_table("documents")}
                {where_clause}
                ORDER BY created_at DESC
                LIMIT {limit_param} OFFSET {offset_param}
            """,
                *query_params,
            )

            # Get memory unit count for each document
            if documents:
                doc_ids = [(row["id"], row["bank_id"]) for row in documents]

                # Create placeholders for the query
                placeholders = []
                params_for_count = []
                for i, (doc_id, bank_id_val) in enumerate(doc_ids):
                    idx_doc = i * 2 + 1
                    idx_agent = i * 2 + 2
                    placeholders.append(f"(document_id = ${idx_doc} AND bank_id = ${idx_agent})")
                    params_for_count.extend([doc_id, bank_id_val])

                where_clause_count = " OR ".join(placeholders)

                unit_counts = await conn.fetch(
                    f"""
                    SELECT document_id, bank_id, COUNT(*) as unit_count
                    FROM {fq_table("memory_units")}
                    WHERE {where_clause_count}
                    GROUP BY document_id, bank_id
                """,
                    *params_for_count,
                )
            else:
                unit_counts = []

            # Build count mapping
            count_map = {(row["document_id"], row["bank_id"]): row["unit_count"] for row in unit_counts}

            # Build result items
            items = []
            for row in documents:
                doc_id = row["id"]
                bank_id_val = row["bank_id"]
                unit_count = count_map.get((doc_id, bank_id_val), 0)

                items.append(
                    {
                        "id": doc_id,
                        "bank_id": bank_id_val,
                        "content_hash": row["content_hash"],
                        "created_at": row["created_at"].isoformat() if row["created_at"] else "",
                        "updated_at": row["updated_at"].isoformat() if row["updated_at"] else "",
                        "text_length": row["text_length"] or 0,
                        "memory_unit_count": unit_count,
                        "retain_params": row["retain_params"] if row["retain_params"] else None,
                        "tags": row["tags"] if row["tags"] else [],
                    }
                )

            return {"items": items, "total": total, "limit": limit, "offset": offset}

    async def get_chunk(
        self,
        chunk_id: str,
        *,
        request_context: "RequestContext",
    ):
        """
        Get a specific chunk by its ID.

        Args:
            chunk_id: Chunk ID (format: bank_id_document_id_chunk_index)
            request_context: Request context for authentication.

        Returns:
            Dict with chunk details including chunk_text, or None if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            chunk = await conn.fetchrow(
                f"""
                SELECT
                    chunk_id,
                    document_id,
                    bank_id,
                    chunk_index,
                    chunk_text,
                    created_at
                FROM {fq_table("chunks")}
                WHERE chunk_id = $1
            """,
                chunk_id,
            )

            if not chunk:
                return None

            return {
                "chunk_id": chunk["chunk_id"],
                "document_id": chunk["document_id"],
                "bank_id": chunk["bank_id"],
                "chunk_index": chunk["chunk_index"],
                "chunk_text": chunk["chunk_text"],
                "created_at": chunk["created_at"].isoformat() if chunk["created_at"] else "",
            }

    # ==================== bank profile Methods ====================

    async def get_bank_profile(
        self,
        bank_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """
        Get bank profile (name, disposition + mission).
        Auto-creates agent with default values if not exists.

        Args:
            bank_id: bank IDentifier
            request_context: Request context for authentication.

        Returns:
            Dict with name, disposition traits, and mission
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        profile = await bank_utils.get_bank_profile(pool, bank_id)
        disposition = profile["disposition"]
        return {
            "bank_id": bank_id,
            "name": profile["name"],
            "disposition": disposition,
            "mission": profile["mission"],
        }

    async def update_bank_disposition(
        self,
        bank_id: str,
        disposition: dict[str, int],
        *,
        request_context: "RequestContext",
    ) -> None:
        """
        Update bank disposition traits.

        Args:
            bank_id: bank IDentifier
            disposition: Dict with skepticism, literalism, empathy (all 1-5)
            request_context: Request context for authentication.
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        await bank_utils.update_bank_disposition(pool, bank_id, disposition)

    async def set_bank_mission(
        self,
        bank_id: str,
        mission: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """
        Set the mission for a bank.

        Args:
            bank_id: bank IDentifier
            mission: The mission text
            request_context: Request context for authentication.

        Returns:
            Dict with bank_id and mission.
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        await bank_utils.set_bank_mission(pool, bank_id, mission)
        return {"bank_id": bank_id, "mission": mission}

    async def merge_bank_mission(
        self,
        bank_id: str,
        new_info: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """
        Merge new mission information with existing mission using LLM.
        Normalizes to first person ("I") and resolves conflicts.

        Args:
            bank_id: bank IDentifier
            new_info: New mission information to add/merge
            request_context: Request context for authentication.

        Returns:
            Dict with 'mission' (str) key
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        return await bank_utils.merge_bank_mission(pool, self._reflect_llm_config, bank_id, new_info)

    async def list_banks(
        self,
        *,
        request_context: "RequestContext",
    ) -> list[dict[str, Any]]:
        """
        List all agents in the system.

        Args:
            request_context: Request context for authentication.

        Returns:
            List of dicts with bank_id, name, disposition, mission, created_at, updated_at
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        return await bank_utils.list_banks(pool)

    # ==================== Reflect Methods ====================

    async def reflect_async(
        self,
        bank_id: str,
        query: str,
        *,
        budget: Budget | None = None,
        context: str | None = None,
        max_tokens: int = 4096,
        response_schema: dict | None = None,
        request_context: "RequestContext",
        tags: list[str] | None = None,
        tags_match: TagsMatch = "any",
        exclude_mental_model_ids: list[str] | None = None,
        _skip_span: bool = False,
    ) -> ReflectResult:
        """
        Reflect and formulate an answer using an agentic loop with tools.

        The reflect agent iteratively uses tools to:
        1. lookup: Get mental models (synthesized knowledge)
        2. recall: Search facts (semantic + temporal retrieval)
        3. learn: Create/update mental models with new insights
        4. expand: Get chunk/document context for memories

        The agent starts with empty context and must call tools to gather
        information. On the last iteration, tools are removed to force a
        final text response.

        Args:
            bank_id: bank identifier
            query: Question to answer
            budget: Budget level (currently unused, reserved for future)
            context: Additional context string to include in agent prompt
            max_tokens: Max tokens (currently unused, reserved for future)
            response_schema: Optional JSON Schema for structured output (not yet supported)
            tags: Optional tags to filter memories
            tags_match: How to match tags - "any" (OR), "all" (AND)
            exclude_mental_model_ids: Optional list of mental model IDs to exclude from search
                (used when refreshing a mental model to avoid circular reference)

        Returns:
            ReflectResult containing:
                - text: Plain text answer
                - based_on: Empty dict (agent retrieves facts dynamically)
                - structured_output: None (not yet supported for agentic reflect)
        """
        # Use cached LLM config
        if self._reflect_llm_config is None:
            raise ValueError("Memory LLM API key not set. Set HINDSIGHT_API_LLM_API_KEY environment variable.")

        # Authenticate tenant and set schema in context (for fq_table())
        await self._authenticate_tenant(request_context)

        # Validate operation if validator is configured
        if self._operation_validator:
            from hindsight_api.extensions import ReflectContext

            ctx = ReflectContext(
                bank_id=bank_id,
                query=query,
                request_context=request_context,
                budget=budget,
                context=context,
            )
            await self._validate_operation(self._operation_validator.validate_reflect(ctx))

        reflect_start = time.time()
        reflect_id = f"{bank_id[:8]}-{int(time.time() * 1000) % 100000}"
        tags_info = f", tags={tags} ({tags_match})" if tags else ""
        logger.info(f"[REFLECT {reflect_id}] Starting agentic reflect for query: {query[:50]}...{tags_info}")

        # Get bank profile for agent identity
        profile = await self.get_bank_profile(bank_id, request_context=request_context)

        # NOTE: Mental models are NOT pre-loaded to keep the initial prompt small.
        # The agent can call lookup() to list available models if needed.
        # This is critical for banks with many mental models to avoid huge prompts.

        # Compute max iterations based on budget
        config = get_config()
        base_max_iterations = config.reflect_max_iterations
        # Budget multipliers: low=0.5x, mid=1x, high=2x
        budget_multipliers = {Budget.LOW: 0.5, Budget.MID: 1.0, Budget.HIGH: 2.0}
        effective_budget = budget or Budget.LOW
        max_iterations = max(1, int(base_max_iterations * budget_multipliers.get(effective_budget, 1.0)))

        # Run agentic loop - acquire connections only when needed for DB operations
        # (not held during LLM calls which can be slow)
        pool = await self._get_pool()

        # Get bank stats for freshness info
        bank_stats = await self.get_bank_stats(bank_id, request_context=request_context)
        last_consolidated_at = bank_stats.last_consolidated_at if hasattr(bank_stats, "last_consolidated_at") else None
        pending_consolidation = bank_stats.pending_consolidation if hasattr(bank_stats, "pending_consolidation") else 0

        # Create tool callbacks that acquire connections only when needed
        from .retain import embedding_utils

        async def search_mental_models_fn(q: str, max_results: int = 5) -> dict[str, Any]:
            # Generate embedding for the query
            embeddings = await embedding_utils.generate_embeddings_batch(self.embeddings, [q])
            query_embedding = embeddings[0]
            async with pool.acquire() as conn:
                return await tool_search_mental_models(
                    conn,
                    bank_id,
                    q,
                    query_embedding,
                    max_results=max_results,
                    tags=tags,
                    tags_match=tags_match,
                    exclude_ids=exclude_mental_model_ids,
                    pending_consolidation=pending_consolidation,
                )

        async def search_observations_fn(q: str, max_tokens: int = 5000) -> dict[str, Any]:
            return await tool_search_observations(
                self,
                bank_id,
                q,
                request_context,
                max_tokens=max_tokens,
                tags=tags,
                tags_match=tags_match,
                last_consolidated_at=last_consolidated_at,
                pending_consolidation=pending_consolidation,
            )

        async def recall_fn(q: str, max_tokens: int = 4096) -> dict[str, Any]:
            return await tool_recall(
                self, bank_id, q, request_context, max_tokens=max_tokens, tags=tags, tags_match=tags_match
            )

        async def expand_fn(memory_ids: list[str], depth: str) -> dict[str, Any]:
            async with pool.acquire() as conn:
                return await tool_expand(conn, bank_id, memory_ids, depth)

        # Load directives from the dedicated directives table
        # Directives are hard rules that must be followed in all responses
        # Use isolation_mode=True to prevent tag-scoped directives from leaking into untagged operations
        # Use the same tags_match as the reflect request so directives respect the same scoping rules
        directives_raw = await self.list_directives(
            bank_id=bank_id,
            tags=tags,
            tags_match=tags_match,
            active_only=True,
            request_context=request_context,
            isolation_mode=True,
        )
        directives = directives_raw
        if directives:
            logger.info(f"[REFLECT {reflect_id}] Loaded {len(directives)} directives")

        # Check if the bank has any mental models
        async with pool.acquire() as conn:
            mental_model_count = await conn.fetchval(
                f"SELECT COUNT(*) FROM {fq_table('mental_models')} WHERE bank_id = $1",
                bank_id,
            )
        has_mental_models = mental_model_count > 0
        if has_mental_models:
            logger.info(f"[REFLECT {reflect_id}] Bank has {mental_model_count} mental models")

        # Run the agent with parent span for reflect operation (skip if called from another operation)
        if not _skip_span:
            span_context = create_operation_span("reflect", bank_id)
            span_context.__enter__()
        else:
            span_context = None

        try:
            agent_result = await run_reflect_agent(
                llm_config=self._reflect_llm_config,
                bank_id=bank_id,
                query=query,
                bank_profile=profile,
                search_mental_models_fn=search_mental_models_fn,
                search_observations_fn=search_observations_fn,
                recall_fn=recall_fn,
                expand_fn=expand_fn,
                context=context,
                max_iterations=max_iterations,
                max_tokens=max_tokens,
                response_schema=response_schema,
                directives=directives,
                has_mental_models=has_mental_models,
                budget=effective_budget,
            )

            total_time = time.time() - reflect_start
            logger.info(
                f"[REFLECT {reflect_id}] Complete: {len(agent_result.text)} chars, "
                f"{agent_result.iterations} iterations, {agent_result.tools_called} tool calls | {total_time:.3f}s"
            )

            # Convert agent tool trace to ToolCallTrace objects
            tool_trace_result = [
                ToolCallTrace(
                    tool=tc.tool,
                    reason=tc.reason,
                    input=tc.input,
                    output=tc.output,
                    duration_ms=tc.duration_ms,
                    iteration=tc.iteration,
                )
                for tc in agent_result.tool_trace
            ]

            # Convert agent LLM trace to LLMCallTrace objects
            llm_trace_result = [
                LLMCallTrace(scope=lc.scope, duration_ms=lc.duration_ms) for lc in agent_result.llm_trace
            ]

            # Extract memories from recall tool outputs - only include memories the agent actually used
            # agent_result.used_memory_ids contains validated IDs from the done action
            used_memory_ids_set = set(agent_result.used_memory_ids) if agent_result.used_memory_ids else set()
            # based_on stores facts, mental models, and directives
            # Note: directives list stores raw directive dicts (not MemoryFact), which will be converted to Directive objects
            based_on: dict[str, list[MemoryFact] | list[dict[str, Any]]] = {
                "world": [],
                "experience": [],
                "opinion": [],
                "observation": [],
                "mental-models": [],
                "directives": [],
            }
            seen_memory_ids: set[str] = set()
            for tc in agent_result.tool_trace:
                if tc.tool == "recall" and "memories" in tc.output:
                    for memory_data in tc.output["memories"]:
                        memory_id = memory_data.get("id")
                        # Only include memories that the agent declared as used (or all if none specified)
                        if memory_id and memory_id not in seen_memory_ids:
                            if used_memory_ids_set and memory_id not in used_memory_ids_set:
                                continue  # Skip memories not actually used by the agent
                            seen_memory_ids.add(memory_id)
                            fact_type = memory_data.get("type", "world")
                            if fact_type in based_on:
                                based_on[fact_type].append(
                                    MemoryFact(
                                        id=memory_id,
                                        text=memory_data.get("text", ""),
                                        fact_type=fact_type,
                                        context=None,
                                        occurred_start=memory_data.get("occurred"),
                                        occurred_end=memory_data.get("occurred"),
                                    )
                                )

            # Extract mental models from tool outputs - only include models the agent actually used
            # agent_result.used_mental_model_ids contains validated IDs from the done action
            used_model_ids_set = (
                set(agent_result.used_mental_model_ids) if agent_result.used_mental_model_ids else set()
            )
            based_on["mental-models"] = []
            seen_model_ids: set[str] = set()
            for tc in agent_result.tool_trace:
                if tc.tool == "get_mental_model":
                    # Single model lookup (with full details)
                    if tc.output.get("found") and "model" in tc.output:
                        model = tc.output["model"]
                        model_id = model.get("id")
                        if model_id and model_id not in seen_model_ids:
                            # Only include models that the agent declared as used (or all if none specified)
                            if used_model_ids_set and model_id not in used_model_ids_set:
                                continue  # Skip models not actually used by the agent
                            seen_model_ids.add(model_id)
                            # Add to based_on as MemoryFact with type "mental-models"
                            model_name = model.get("name", "")
                            model_summary = model.get("summary") or model.get("description", "")
                            based_on["mental-models"].append(
                                MemoryFact(
                                    id=model_id,
                                    text=f"{model_name}: {model_summary}",
                                    fact_type="mental-models",
                                    context=f"{model.get('type', 'concept')} ({model.get('subtype', 'structural')})",
                                    occurred_start=None,
                                    occurred_end=None,
                                )
                            )
                elif tc.tool == "search_mental_models":
                    # Search mental models - include all returned models (filtered by used_model_ids_set if specified)
                    for model in tc.output.get("mental_models", []):
                        model_id = model.get("id")
                        if model_id and model_id not in seen_model_ids:
                            # Only include models that the agent declared as used (or all if none specified)
                            if used_model_ids_set and model_id not in used_model_ids_set:
                                continue  # Skip models not actually used by the agent
                            seen_model_ids.add(model_id)
                            # Add to based_on as MemoryFact with type "mental-models"
                            model_name = model.get("name", "")
                            model_summary = model.get("summary") or model.get("description", "")
                            based_on["mental-models"].append(
                                MemoryFact(
                                    id=model_id,
                                    text=f"{model_name}: {model_summary}",
                                    fact_type="mental-models",
                                    context=f"{model.get('type', 'concept')} ({model.get('subtype', 'structural')})",
                                    occurred_start=None,
                                    occurred_end=None,
                                )
                            )
                elif tc.tool == "search_mental_models":
                    # Search mental models - include all returned mental models (filtered by used_mental_model_ids_set if specified)
                    used_mental_model_ids_set = (
                        set(agent_result.used_mental_model_ids) if agent_result.used_mental_model_ids else set()
                    )
                    for mental_model in tc.output.get("mental_models", []):
                        mental_model_id = mental_model.get("id")
                        if mental_model_id and mental_model_id not in seen_model_ids:
                            # Only include mental models that the agent declared as used (or all if none specified)
                            if used_mental_model_ids_set and mental_model_id not in used_mental_model_ids_set:
                                continue  # Skip mental models not actually used by the agent
                            seen_model_ids.add(mental_model_id)
                            # Add to based_on as MemoryFact with type "mental-models" (mental models are synthesized knowledge)
                            mental_model_name = mental_model.get("name", "")
                            mental_model_content = mental_model.get("content", "")
                            based_on["mental-models"].append(
                                MemoryFact(
                                    id=mental_model_id,
                                    text=f"{mental_model_name}: {mental_model_content}",
                                    fact_type="mental-models",
                                    context="mental model (user-curated)",
                                    occurred_start=None,
                                    occurred_end=None,
                                )
                            )
                    # List all models lookup - don't add to based_on (too verbose, just a listing)

            # Add directives to based_on["directives"]
            # Store raw directive dicts (with id, name, content) for http.py to convert to ReflectDirective
            for directive_raw in directives_raw:
                based_on["directives"].append(
                    {
                        "id": directive_raw["id"],
                        "name": directive_raw["name"],
                        "content": directive_raw["content"],
                    }
                )

            # Build directives_applied from agent result
            from hindsight_api.engine.response_models import DirectiveRef

            directives_applied_result = [
                DirectiveRef(id=d.id, name=d.name, content=d.content) 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,
                structured_output=agent_result.structured_output,
                usage=usage,
                tool_trace=tool_trace_result,
                llm_trace=llm_trace_result,
                directives_applied=directives_applied_result,
            )

            # Call post-operation hook if validator is configured
            if self._operation_validator:
                from hindsight_api.extensions.operation_validator import ReflectResultContext

                result_ctx = ReflectResultContext(
                    bank_id=bank_id,
                    query=query,
                    request_context=request_context,
                    budget=budget,
                    context=context,
                    result=result,
                    success=True,
                    error=None,
                )
                try:
                    await self._operation_validator.on_reflect_complete(result_ctx)
                except Exception as e:
                    logger.warning(f"Post-reflect hook error (non-fatal): {e}")

            return result
        finally:
            if span_context:
                span_context.__exit__(None, None, None)

    async def list_entities(
        self,
        bank_id: str,
        *,
        limit: int = 100,
        offset: int = 0,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """
        List all entities for a bank with pagination.

        Args:
            bank_id: bank IDentifier
            limit: Maximum number of entities to return
            offset: Offset for pagination
            request_context: Request context for authentication.

        Returns:
            Dict with items, total, limit, offset
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            # Get total count
            total_row = await conn.fetchrow(
                f"""
                SELECT COUNT(*) as total
                FROM {fq_table("entities")}
                WHERE bank_id = $1
                """,
                bank_id,
            )
            total = total_row["total"] if total_row else 0

            # Get paginated entities
            rows = await conn.fetch(
                f"""
                SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
                FROM {fq_table("entities")}
                WHERE bank_id = $1
                ORDER BY mention_count DESC, last_seen DESC, id ASC
                LIMIT $2 OFFSET $3
                """,
                bank_id,
                limit,
                offset,
            )

            entities = []
            for row in rows:
                # Handle metadata - may be dict, JSON string, or None
                metadata = row["metadata"]
                if metadata is None:
                    metadata = {}
                elif isinstance(metadata, str):
                    import json

                    try:
                        metadata = json.loads(metadata)
                    except json.JSONDecodeError:
                        metadata = {}

                entities.append(
                    {
                        "id": str(row["id"]),
                        "canonical_name": row["canonical_name"],
                        "mention_count": row["mention_count"],
                        "first_seen": row["first_seen"].isoformat() if row["first_seen"] else None,
                        "last_seen": row["last_seen"].isoformat() if row["last_seen"] else None,
                        "metadata": metadata,
                    }
                )
            return {
                "items": entities,
                "total": total,
                "limit": limit,
                "offset": offset,
            }

    async def list_tags(
        self,
        bank_id: str,
        *,
        pattern: str | None = None,
        limit: int = 100,
        offset: int = 0,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """
        List all unique tags for a bank with usage counts.

        Use this to discover available tags or expand wildcard patterns.
        Supports '*' as wildcard for flexible matching (case-insensitive):
        - 'user:*' matches user:alice, user:bob
        - '*-admin' matches role-admin, super-admin
        - 'env*-prod' matches env-prod, environment-prod

        Args:
            bank_id: Bank identifier
            pattern: Wildcard pattern to filter tags (use '*' as wildcard, case-insensitive)
            limit: Maximum number of tags to return
            offset: Offset for pagination
            request_context: Request context for authentication.

        Returns:
            Dict with items (list of {tag, count}), total, limit, offset
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()
        async with acquire_with_retry(pool) as conn:
            # Build pattern filter if provided (convert * to % for ILIKE)
            pattern_clause = ""
            params: list[Any] = [bank_id]
            if pattern:
                # Convert wildcard pattern: * -> % for SQL ILIKE
                sql_pattern = pattern.replace("*", "%")
                pattern_clause = "AND tag ILIKE $2"
                params.append(sql_pattern)

            # Get total count of distinct tags matching pattern
            total_row = await conn.fetchrow(
                f"""
                SELECT COUNT(DISTINCT tag) as total
                FROM {fq_table("memory_units")}, unnest(tags) AS tag
                WHERE bank_id = $1 AND tags IS NOT NULL AND tags != '{{}}'
                {pattern_clause}
                """,
                *params,
            )
            total = total_row["total"] if total_row else 0

            # Get paginated tags with counts, ordered by frequency
            limit_param = len(params) + 1
            offset_param = len(params) + 2
            params.extend([limit, offset])

            rows = await conn.fetch(
                f"""
                SELECT tag, COUNT(*) as count
                FROM {fq_table("memory_units")}, unnest(tags) AS tag
                WHERE bank_id = $1 AND tags IS NOT NULL AND tags != '{{}}'
                {pattern_clause}
                GROUP BY tag
                ORDER BY count DESC, tag ASC
                LIMIT ${limit_param} OFFSET ${offset_param}
                """,
                *params,
            )

            items = [{"tag": row["tag"], "count": row["count"]} for row in rows]

            return {
                "items": items,
                "total": total,
                "limit": limit,
                "offset": offset,
            }

    async def get_entity_state(
        self,
        bank_id: str,
        entity_id: str,
        entity_name: str,
        *,
        limit: int = 10,
        request_context: "RequestContext",
    ) -> EntityState:
        """
        Get the current state of an entity.

        NOTE: Entity observations/summaries have been moved to mental models.
        This method returns an entity with empty observations.

        Args:
            bank_id: bank IDentifier
            entity_id: Entity UUID
            entity_name: Canonical name of the entity
            limit: Maximum number of observations to include (kept for backwards compat)
            request_context: Request context for authentication.

        Returns:
            EntityState with empty observations (summaries now in mental models)
        """
        await self._authenticate_tenant(request_context)
        return EntityState(entity_id=entity_id, canonical_name=entity_name, observations=[])

    # =========================================================================
    # Statistics & Operations (for HTTP API layer)
    # =========================================================================

    async def get_bank_stats(
        self,
        bank_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Get statistics about memory nodes and links for a bank."""
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            # Get node counts by fact_type
            node_stats = await conn.fetch(
                f"""
                SELECT fact_type, COUNT(*) as count
                FROM {fq_table("memory_units")}
                WHERE bank_id = $1
                GROUP BY fact_type
                """,
                bank_id,
            )

            # Get link counts by link_type
            link_stats = await conn.fetch(
                f"""
                SELECT ml.link_type, COUNT(*) as count
                FROM {fq_table("memory_links")} ml
                JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
                WHERE mu.bank_id = $1
                GROUP BY ml.link_type
                """,
                bank_id,
            )

            # Get link counts by fact_type (from nodes)
            link_fact_type_stats = await conn.fetch(
                f"""
                SELECT mu.fact_type, COUNT(*) as count
                FROM {fq_table("memory_links")} ml
                JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
                WHERE mu.bank_id = $1
                GROUP BY mu.fact_type
                """,
                bank_id,
            )

            # Get link counts by fact_type AND link_type
            link_breakdown_stats = await conn.fetch(
                f"""
                SELECT mu.fact_type, ml.link_type, COUNT(*) as count
                FROM {fq_table("memory_links")} ml
                JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
                WHERE mu.bank_id = $1
                GROUP BY mu.fact_type, ml.link_type
                """,
                bank_id,
            )

            # Get pending and failed operations counts
            ops_stats = await conn.fetch(
                f"""
                SELECT status, COUNT(*) as count
                FROM {fq_table("async_operations")}
                WHERE bank_id = $1
                GROUP BY status
                """,
                bank_id,
            )

            return {
                "bank_id": bank_id,
                "node_counts": {row["fact_type"]: row["count"] for row in node_stats},
                "link_counts": {row["link_type"]: row["count"] for row in link_stats},
                "link_counts_by_fact_type": {row["fact_type"]: row["count"] for row in link_fact_type_stats},
                "link_breakdown": [
                    {"fact_type": row["fact_type"], "link_type": row["link_type"], "count": row["count"]}
                    for row in link_breakdown_stats
                ],
                "operations": {row["status"]: row["count"] for row in ops_stats},
            }

    async def get_entity(
        self,
        bank_id: str,
        entity_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """Get entity details including metadata and observations."""
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            entity_row = await conn.fetchrow(
                f"""
                SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
                FROM {fq_table("entities")}
                WHERE bank_id = $1 AND id = $2
                """,
                bank_id,
                uuid.UUID(entity_id),
            )

        if not entity_row:
            return None

        return {
            "id": str(entity_row["id"]),
            "canonical_name": entity_row["canonical_name"],
            "mention_count": entity_row["mention_count"],
            "first_seen": entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
            "last_seen": entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
            "metadata": entity_row["metadata"] or {},
            "observations": [],
        }

    def _parse_observations(self, observations_raw: list):
        """Parse raw observation dicts into typed Observation models.

        Returns list of Observation models with computed trend/evidence_span/evidence_count.
        """
        from .reflect.observations import Observation, ObservationEvidence

        observations: list[Observation] = []
        for obs in observations_raw:
            if not isinstance(obs, dict):
                continue

            try:
                parsed = Observation(
                    title=obs.get("title", ""),
                    content=obs.get("content", ""),
                    evidence=[
                        ObservationEvidence(
                            memory_id=ev.get("memory_id", ""),
                            quote=ev.get("quote", ""),
                            relevance=ev.get("relevance", ""),
                            timestamp=ev.get("timestamp"),
                        )
                        for ev in obs.get("evidence", [])
                        if isinstance(ev, dict)
                    ],
                    created_at=obs.get("created_at"),
                )
                observations.append(parsed)
            except Exception as e:
                logger.warning(f"Failed to parse observation: {e}")
                continue

        return observations

    async def _count_memories_since(
        self,
        bank_id: str,
        since_timestamp: str | None,
        pool=None,
    ) -> int:
        """
        Count memories created after a given timestamp.

        Args:
            bank_id: Bank identifier
            since_timestamp: ISO timestamp string. If None, returns total count.
            pool: Optional database pool (uses default if not provided)

        Returns:
            Number of memories created since the timestamp
        """
        if pool is None:
            pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            if since_timestamp:
                # Parse the timestamp
                from datetime import datetime

                try:
                    ts = datetime.fromisoformat(since_timestamp.replace("Z", "+00:00"))
                except ValueError:
                    # Invalid timestamp, return total count
                    ts = None

                if ts:
                    count = await conn.fetchval(
                        f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND created_at > $2",
                        bank_id,
                        ts,
                    )
                    return count or 0

            # No timestamp or invalid, return total count
            count = await conn.fetchval(
                f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1",
                bank_id,
            )
            return count or 0

    async def _invalidate_facts_from_mental_models(
        self,
        conn,
        bank_id: str,
        fact_ids: list[str],
    ) -> int:
        """
        Remove fact IDs from observation source_memory_ids when memories are deleted.

        Observations are stored in memory_units with fact_type='observation'
        and have a source_memory_ids column (UUID[]) tracking their source memories.

        Args:
            conn: Database connection
            bank_id: Bank identifier
            fact_ids: List of fact IDs to remove from observations

        Returns:
            Number of observations updated
        """
        if not fact_ids:
            return 0

        # Convert string IDs to UUIDs for the array comparison
        import uuid as uuid_module

        fact_uuids = [uuid_module.UUID(fid) for fid in fact_ids]

        # Update observations (memory_units with fact_type='observation')
        # by removing the deleted fact IDs from source_memory_ids
        # Use array subtraction: source_memory_ids - deleted_ids
        result = await conn.execute(
            f"""
            UPDATE {fq_table("memory_units")}
            SET source_memory_ids = (
                SELECT COALESCE(array_agg(elem), ARRAY[]::uuid[])
                FROM unnest(source_memory_ids) AS elem
                WHERE elem != ALL($2::uuid[])
            ),
                updated_at = NOW()
            WHERE bank_id = $1
            AND fact_type = 'observation'
            AND source_memory_ids && $2::uuid[]
            """,
            bank_id,
            fact_uuids,
        )

        # Parse the result to get number of updated rows
        updated_count = int(result.split()[-1]) if result and "UPDATE" in result else 0
        if updated_count > 0:
            logger.info(
                f"[OBSERVATIONS] Invalidated {len(fact_ids)} fact IDs from {updated_count} observations in bank {bank_id}"
            )
        return updated_count

    # =========================================================================
    # MENTAL MODELS (CONSOLIDATED) - Read-only access to auto-consolidated mental models
    # =========================================================================

    async def list_mental_models_consolidated(
        self,
        bank_id: str,
        *,
        tags: list[str] | None = None,
        tags_match: str = "any",
        limit: int = 100,
        offset: int = 0,
        request_context: "RequestContext",
    ) -> list[dict[str, Any]]:
        """List auto-consolidated observations for a bank.

        Observations are stored in memory_units with fact_type='observation'.
        They are automatically created and updated by the consolidation engine.

        Args:
            bank_id: Bank identifier
            tags: Optional tags to filter by
            tags_match: How to match tags - 'any', 'all', or 'exact'
            limit: Maximum number of results
            offset: Offset for pagination
            request_context: Request context for authentication

        Returns:
            List of observation dicts
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            # Build tag filter
            tag_filter = ""
            params: list[Any] = [bank_id, limit, offset]
            if tags:
                if tags_match == "all":
                    tag_filter = " AND tags @> $4::varchar[]"
                elif tags_match == "exact":
                    tag_filter = " AND tags = $4::varchar[]"
                else:  # any
                    tag_filter = " AND tags && $4::varchar[]"
                params.append(tags)

            rows = await conn.fetch(
                f"""
                SELECT id, bank_id, text, proof_count, history, tags, source_memory_ids, created_at, updated_at
                FROM {fq_table("memory_units")}
                WHERE bank_id = $1 AND fact_type = 'observation' {tag_filter}
                ORDER BY updated_at DESC NULLS LAST
                LIMIT $2 OFFSET $3
                """,
                *params,
            )

            return [self._row_to_observation_consolidated(row) for row in rows]

    async def get_observation_consolidated(
        self,
        bank_id: str,
        observation_id: str,
        *,
        include_source_memories: bool = True,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """Get a single observation by ID.

        Args:
            bank_id: Bank identifier
            observation_id: Observation ID
            include_source_memories: Whether to include full source memory details
            request_context: Request context for authentication

        Returns:
            Observation dict or None if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            row = await conn.fetchrow(
                f"""
                SELECT id, bank_id, text, proof_count, history, tags, source_memory_ids, created_at, updated_at
                FROM {fq_table("memory_units")}
                WHERE bank_id = $1 AND id = $2 AND fact_type = 'observation'
                """,
                bank_id,
                observation_id,
            )

            if not row:
                return None

            result = self._row_to_observation_consolidated(row)

            # Fetch source memories if requested and source_memory_ids exist
            if include_source_memories and result.get("source_memory_ids"):
                source_ids = [uuid.UUID(sid) if isinstance(sid, str) else sid for sid in result["source_memory_ids"]]
                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

    def _row_to_observation_consolidated(self, row: Any) -> dict[str, Any]:
        """Convert a database row to an observation dict."""
        import json

        history = row["history"]
        if isinstance(history, str):
            history = json.loads(history)
        elif history is None:
            history = []

        # Convert source_memory_ids to strings
        source_memory_ids = row.get("source_memory_ids") or []
        source_memory_ids = [str(sid) for sid in source_memory_ids]

        return {
            "id": str(row["id"]),
            "bank_id": row["bank_id"],
            "text": row["text"],
            "proof_count": row["proof_count"] or 1,
            "history": history,
            "tags": row["tags"] or [],
            "source_memory_ids": source_memory_ids,
            "source_memories": [],  # Populated separately when fetching full details
            "created_at": row["created_at"].isoformat() if row["created_at"] else None,
            "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
        }

    # =========================================================================
    # MENTAL MODELS CRUD
    # =========================================================================

    async def list_mental_models(
        self,
        bank_id: str,
        *,
        tags: list[str] | None = None,
        tags_match: str = "any",
        limit: int = 100,
        offset: int = 0,
        request_context: "RequestContext",
    ) -> list[dict[str, Any]]:
        """List pinned mental models for a bank.

        Args:
            bank_id: Bank identifier
            tags: Optional tags to filter by
            tags_match: How to match tags - 'any', 'all', or 'exact'
            limit: Maximum number of results
            offset: Offset for pagination
            request_context: Request context for authentication

        Returns:
            List of pinned mental model dicts
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            # Build tag filter
            tag_filter = ""
            params: list[Any] = [bank_id, limit, offset]
            if tags:
                if tags_match == "all":
                    tag_filter = " AND tags @> $4::varchar[]"
                elif tags_match == "exact":
                    tag_filter = " AND tags = $4::varchar[]"
                else:  # any
                    tag_filter = " AND tags && $4::varchar[]"
                params.append(tags)

            rows = await conn.fetch(
                f"""
                SELECT id, bank_id, name, source_query, content, tags,
                       last_refreshed_at, created_at, reflect_response,
                       max_tokens, trigger
                FROM {fq_table("mental_models")}
                WHERE bank_id = $1 {tag_filter}
                ORDER BY last_refreshed_at DESC
                LIMIT $2 OFFSET $3
                """,
                *params,
            )

            return [self._row_to_mental_model(row) for row in rows]

    async def get_mental_model(
        self,
        bank_id: str,
        mental_model_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """Get a single pinned mental model by ID.

        Args:
            bank_id: Bank identifier
            mental_model_id: Pinned mental model UUID
            request_context: Request context for authentication

        Returns:
            Pinned mental model dict or None if not found
        """
        await self._authenticate_tenant(request_context)

        # Pre-operation validation (credit check / usage metering)
        if self._operation_validator:
            from hindsight_api.extensions.operation_validator import MentalModelGetContext

            ctx = MentalModelGetContext(
                bank_id=bank_id,
                mental_model_id=mental_model_id,
                request_context=request_context,
            )
            await self._validate_operation(self._operation_validator.validate_mental_model_get(ctx))

        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            row = await conn.fetchrow(
                f"""
                SELECT id, bank_id, name, source_query, content, tags,
                       last_refreshed_at, created_at, reflect_response,
                       max_tokens, trigger
                FROM {fq_table("mental_models")}
                WHERE bank_id = $1 AND id = $2
                """,
                bank_id,
                mental_model_id,
            )

            result = self._row_to_mental_model(row) if row else None

        # Post-operation hook (usage recording)
        if result and self._operation_validator:
            from hindsight_api.extensions.operation_validator import MentalModelGetResult

            content = result.get("content", "")
            output_tokens = len(content) // 4 if content else 0

            result_ctx = MentalModelGetResult(
                bank_id=bank_id,
                mental_model_id=mental_model_id,
                request_context=request_context,
                output_tokens=output_tokens,
                success=True,
            )
            try:
                await self._operation_validator.on_mental_model_get_complete(result_ctx)
            except Exception as hook_err:
                logger.warning(f"Post-mental-model-get hook error (non-fatal): {hook_err}")

        return result

    async def create_mental_model(
        self,
        bank_id: str,
        name: str,
        source_query: str,
        content: str,
        *,
        mental_model_id: str | None = None,
        tags: list[str] | None = None,
        max_tokens: int | None = None,
        trigger: dict[str, Any] | None = None,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Create a new pinned mental model.

        Args:
            bank_id: Bank identifier
            name: Human-readable name for the mental model
            source_query: The query that generated this mental model
            content: The synthesized content
            mental_model_id: Optional UUID for the mental model (auto-generated if not provided)
            tags: Optional tags for scoped visibility
            max_tokens: Token limit for content generation during refresh
            trigger: Trigger settings (e.g., refresh_after_consolidation)
            request_context: Request context for authentication

        Returns:
            The created pinned mental model dict
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        # Generate embedding for the content
        embedding_text = f"{name} {content}"
        embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [embedding_text])
        # Convert embedding to string for asyncpg vector type
        embedding_str = str(embedding[0]) if embedding else None

        async with acquire_with_retry(pool) as conn:
            if mental_model_id:
                row = await conn.fetchrow(
                    f"""
                    INSERT INTO {fq_table("mental_models")}
                    (id, bank_id, name, source_query, content, embedding, tags, max_tokens, trigger)
                    VALUES ($1, $2, $3, $4, $5, $6, $7, COALESCE($8, 2048), COALESCE($9, '{{"refresh_after_consolidation": false}}'::jsonb))
                    RETURNING id, bank_id, name, source_query, content, tags,
                              last_refreshed_at, created_at, reflect_response,
                              max_tokens, trigger
                    """,
                    mental_model_id,
                    bank_id,
                    name,
                    source_query,
                    content,
                    embedding_str,
                    tags or [],
                    max_tokens,
                    json.dumps(trigger) if trigger else None,
                )
            else:
                row = await conn.fetchrow(
                    f"""
                    INSERT INTO {fq_table("mental_models")}
                    (bank_id, name, source_query, content, embedding, tags, max_tokens, trigger)
                    VALUES ($1, $2, $3, $4, $5, $6, COALESCE($7, 2048), COALESCE($8, '{{"refresh_after_consolidation": false}}'::jsonb))
                    RETURNING id, bank_id, name, source_query, content, tags,
                              last_refreshed_at, created_at, reflect_response,
                              max_tokens, trigger
                    """,
                    bank_id,
                    name,
                    source_query,
                    content,
                    embedding_str,
                    tags or [],
                    max_tokens,
                    json.dumps(trigger) if trigger else None,
                )

        logger.info(f"[MENTAL_MODELS] Created pinned mental model '{name}' for bank {bank_id}")
        return self._row_to_mental_model(row)

    async def refresh_mental_model(
        self,
        bank_id: str,
        mental_model_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """Refresh a pinned mental model by re-running its source query.

        This method:
        1. Gets the pinned mental model
        2. Runs the source_query through reflect
        3. Updates the content with the new synthesis
        4. Updates last_refreshed_at

        Args:
            bank_id: Bank identifier
            mental_model_id: Pinned mental model UUID
            request_context: Request context for authentication

        Returns:
            Updated pinned mental model dict or None if not found
        """
        await self._authenticate_tenant(request_context)

        # Get the current mental model
        mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=request_context)
        if not mental_model:
            return None

        # Create parent span for mental model refresh operation
        with create_operation_span("mental_model_refresh", bank_id):
            # SECURITY: If the mental model has tags, pass them to reflect with "all_strict" matching
            # to ensure it can only access other mental models/memories with the SAME tags.
            # This prevents cross-tenant/cross-user information leakage by excluding untagged content.
            tags = mental_model.get("tags")
            tags_match = "all_strict" if tags else "any"

            # Run reflect with the source query, excluding the mental model being refreshed
            # Skip creating a nested "hindsight.reflect" span since we already have "hindsight.mental_model_refresh"
            reflect_result = await self.reflect_async(
                bank_id=bank_id,
                query=mental_model["source_query"],
                request_context=request_context,
                tags=tags,
                tags_match=tags_match,
                exclude_mental_model_ids=[mental_model_id],
                _skip_span=True,
            )

            # Build reflect_response payload to store
            # based_on contains MemoryFact objects for most types, but plain dicts for directives
            based_on_serialized_payload: dict[str, list[dict[str, Any]]] = {}
            for fact_type, facts in reflect_result.based_on.items():
                serialized_facts = []
                for fact in facts:
                    if isinstance(fact, dict):
                        # Plain dict (e.g., directives with id, name, content)
                        serialized_facts.append(
                            {
                                "id": str(fact["id"]),
                                "text": fact.get("text", fact.get("content", fact.get("name", ""))),
                                "type": fact_type,
                                "context": fact.get("context", None),
                            }
                        )
                    else:
                        # MemoryFact object with .id, .text, .context attributes
                        serialized_facts.append(
                            {
                                "id": str(fact.id),
                                "text": fact.text,
                                "type": fact_type,
                                "context": fact.context,
                            }
                        )
                based_on_serialized_payload[fact_type] = serialized_facts

            reflect_response_payload = {
                "text": reflect_result.text,
                "based_on": based_on_serialized_payload,
                "mental_models": [],  # Mental models are included in based_on["mental-models"]
            }

            # Update the mental model with new content and reflect_response
            return await self.update_mental_model(
                bank_id,
                mental_model_id,
                content=reflect_result.text,
                reflect_response=reflect_response_payload,
                request_context=request_context,
            )

    async def update_mental_model(
        self,
        bank_id: str,
        mental_model_id: str,
        *,
        name: str | None = None,
        content: str | None = None,
        source_query: str | None = None,
        max_tokens: int | None = None,
        tags: list[str] | None = None,
        trigger: dict[str, Any] | None = None,
        reflect_response: dict[str, Any] | None = None,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """Update a pinned mental model.

        Args:
            bank_id: Bank identifier
            mental_model_id: Pinned mental model UUID
            name: New name (if changing)
            content: New content (if changing)
            source_query: New source query (if changing)
            max_tokens: New max tokens (if changing)
            tags: New tags (if changing)
            trigger: New trigger settings (if changing)
            reflect_response: Full reflect API response payload (if changing)
            request_context: Request context for authentication

        Returns:
            Updated pinned mental model dict or None if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            # Build dynamic update
            updates = []
            params: list[Any] = [bank_id, mental_model_id]
            param_idx = 3

            if name is not None:
                updates.append(f"name = ${param_idx}")
                params.append(name)
                param_idx += 1

            if content is not None:
                updates.append(f"content = ${param_idx}")
                params.append(content)
                param_idx += 1
                updates.append("last_refreshed_at = NOW()")
                # Also update embedding (convert to string for asyncpg vector type)
                embedding_text = f"{name or ''} {content}"
                embedding = await embedding_utils.generate_embeddings_batch(self.embeddings, [embedding_text])
                if embedding:
                    updates.append(f"embedding = ${param_idx}")
                    params.append(str(embedding[0]))
                    param_idx += 1

            if reflect_response is not None:
                updates.append(f"reflect_response = ${param_idx}")
                params.append(json.dumps(reflect_response))
                param_idx += 1

            if source_query is not None:
                updates.append(f"source_query = ${param_idx}")
                params.append(source_query)
                param_idx += 1

            if max_tokens is not None:
                updates.append(f"max_tokens = ${param_idx}")
                params.append(max_tokens)
                param_idx += 1

            if tags is not None:
                updates.append(f"tags = ${param_idx}")
                params.append(tags)
                param_idx += 1

            if trigger is not None:
                updates.append(f"trigger = ${param_idx}")
                params.append(json.dumps(trigger))
                param_idx += 1

            if not updates:
                return None

            query = f"""
                UPDATE {fq_table("mental_models")}
                SET {", ".join(updates)}
                WHERE bank_id = $1 AND id = $2
                RETURNING id, bank_id, name, source_query, content, tags,
                          last_refreshed_at, created_at, reflect_response,
                          max_tokens, trigger
            """

            row = await conn.fetchrow(query, *params)

            return self._row_to_mental_model(row) if row else None

    async def delete_mental_model(
        self,
        bank_id: str,
        mental_model_id: str,
        *,
        request_context: "RequestContext",
    ) -> bool:
        """Delete a pinned mental model.

        Args:
            bank_id: Bank identifier
            mental_model_id: Pinned mental model UUID
            request_context: Request context for authentication

        Returns:
            True if deleted, False if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            result = await conn.execute(
                f"DELETE FROM {fq_table('mental_models')} WHERE bank_id = $1 AND id = $2",
                bank_id,
                mental_model_id,
            )

        return result == "DELETE 1"

    def _row_to_mental_model(self, row) -> dict[str, Any]:
        """Convert a database row to a mental model dict."""
        reflect_response = row.get("reflect_response")
        # Parse JSON string to dict if needed (asyncpg may return JSONB as string)
        if isinstance(reflect_response, str):
            try:
                reflect_response = json.loads(reflect_response)
            except json.JSONDecodeError:
                reflect_response = None
        trigger = row.get("trigger")
        if isinstance(trigger, str):
            try:
                trigger = json.loads(trigger)
            except json.JSONDecodeError:
                trigger = None
        return {
            "id": str(row["id"]),
            "bank_id": row["bank_id"],
            "name": row["name"],
            "source_query": row["source_query"],
            "content": row["content"],
            "tags": row["tags"] or [],
            "max_tokens": row.get("max_tokens"),
            "trigger": trigger,
            "last_refreshed_at": row["last_refreshed_at"].isoformat() if row["last_refreshed_at"] else None,
            "created_at": row["created_at"].isoformat() if row["created_at"] else None,
            "reflect_response": reflect_response,
        }

    # =========================================================================
    # Directives - Hard rules injected into prompts
    # =========================================================================

    async def list_directives(
        self,
        bank_id: str,
        *,
        tags: list[str] | None = None,
        tags_match: str = "any",
        active_only: bool = True,
        limit: int = 100,
        offset: int = 0,
        request_context: "RequestContext",
        isolation_mode: bool = False,
    ) -> list[dict[str, Any]]:
        """List directives for a bank.

        Args:
            bank_id: Bank identifier
            tags: Optional tags to filter by
            tags_match: How to match tags - 'any', 'all', or 'exact'
            active_only: Only return active directives (default True)
            limit: Maximum number of results
            offset: Offset for pagination
            request_context: Request context for authentication
            isolation_mode: When True and tags=None, only return directives with no tags.
                This prevents tag-scoped directives from leaking into untagged operations.
                Default False (normal API behavior - returns all directives when tags=None)

        Returns:
            List of directive dicts
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            # Build filters
            from .search.tags import build_tags_where_clause

            filters = ["bank_id = $1"]
            params: list[Any] = [bank_id]
            param_idx = 2

            if active_only:
                filters.append("is_active = TRUE")

            # Apply tags filter for directives:
            # Directives have special scoping rules:
            #   - Untagged directives (tags=[] or null) always apply regardless of reflect tags
            #   - Tagged directives only apply when the reflect operation includes matching tags
            #   - If tags=None and isolation_mode=True: only untagged directives (no leakage)
            #   - If tags=None and isolation_mode=False: all directives (normal API behavior)
            if tags:
                tags_clause, tags_params, param_idx = build_tags_where_clause(
                    tags=tags, param_offset=param_idx, table_alias="", match=tags_match
                )
                if tags_clause:
                    # Always include untagged directives; tagged ones must match the reflect tags
                    scoped_clause = tags_clause.replace("AND ", "", 1)
                    filters.append(f"((tags IS NULL OR tags = '{{}}') OR ({scoped_clause}))")
                    params.extend(tags_params)
            elif isolation_mode:
                # Isolation mode: only include directives with empty/null tags
                # This ensures tag-scoped directives don't apply to untagged operations
                filters.append("(tags IS NULL OR tags = '{}')")

            params.extend([limit, offset])

            rows = await conn.fetch(
                f"""
                SELECT id, bank_id, name, content, priority, is_active, tags, created_at, updated_at
                FROM {fq_table("directives")}
                WHERE {" AND ".join(filters)}
                ORDER BY priority DESC, created_at DESC
                LIMIT ${param_idx} OFFSET ${param_idx + 1}
                """,
                *params,
            )

            return [self._row_to_directive(row) for row in rows]

    async def get_directive(
        self,
        bank_id: str,
        directive_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """Get a single directive by ID.

        Args:
            bank_id: Bank identifier
            directive_id: Directive UUID
            request_context: Request context for authentication

        Returns:
            Directive dict or None if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            row = await conn.fetchrow(
                f"""
                SELECT id, bank_id, name, content, priority, is_active, tags, created_at, updated_at
                FROM {fq_table("directives")}
                WHERE bank_id = $1 AND id = $2
                """,
                bank_id,
                directive_id,
            )

            return self._row_to_directive(row) if row else None

    async def create_directive(
        self,
        bank_id: str,
        name: str,
        content: str,
        *,
        priority: int = 0,
        is_active: bool = True,
        tags: list[str] | None = None,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Create a new directive.

        Args:
            bank_id: Bank identifier
            name: Human-readable name for the directive
            content: The directive text to inject into prompts
            priority: Higher priority directives are injected first (default 0)
            is_active: Whether this directive is active (default True)
            tags: Optional tags for filtering
            request_context: Request context for authentication

        Returns:
            The created directive dict
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            row = await conn.fetchrow(
                f"""
                INSERT INTO {fq_table("directives")}
                (bank_id, name, content, priority, is_active, tags)
                VALUES ($1, $2, $3, $4, $5, $6)
                RETURNING id, bank_id, name, content, priority, is_active, tags, created_at, updated_at
                """,
                bank_id,
                name,
                content,
                priority,
                is_active,
                tags or [],
            )

        logger.info(f"[DIRECTIVES] Created directive '{name}' for bank {bank_id}")
        return self._row_to_directive(row)

    async def update_directive(
        self,
        bank_id: str,
        directive_id: str,
        *,
        name: str | None = None,
        content: str | None = None,
        priority: int | None = None,
        is_active: bool | None = None,
        tags: list[str] | None = None,
        request_context: "RequestContext",
    ) -> dict[str, Any] | None:
        """Update a directive.

        Args:
            bank_id: Bank identifier
            directive_id: Directive UUID
            name: New name (optional)
            content: New content (optional)
            priority: New priority (optional)
            is_active: New active status (optional)
            tags: New tags (optional)
            request_context: Request context for authentication

        Returns:
            Updated directive dict or None if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        # Build update query dynamically
        updates = ["updated_at = now()"]
        params: list[Any] = []
        param_idx = 1

        if name is not None:
            updates.append(f"name = ${param_idx}")
            params.append(name)
            param_idx += 1

        if content is not None:
            updates.append(f"content = ${param_idx}")
            params.append(content)
            param_idx += 1

        if priority is not None:
            updates.append(f"priority = ${param_idx}")
            params.append(priority)
            param_idx += 1

        if is_active is not None:
            updates.append(f"is_active = ${param_idx}")
            params.append(is_active)
            param_idx += 1

        if tags is not None:
            updates.append(f"tags = ${param_idx}")
            params.append(tags)
            param_idx += 1

        params.extend([bank_id, directive_id])

        async with acquire_with_retry(pool) as conn:
            row = await conn.fetchrow(
                f"""
                UPDATE {fq_table("directives")}
                SET {", ".join(updates)}
                WHERE bank_id = ${param_idx} AND id = ${param_idx + 1}
                RETURNING id, bank_id, name, content, priority, is_active, tags, created_at, updated_at
                """,
                *params,
            )

            return self._row_to_directive(row) if row else None

    async def delete_directive(
        self,
        bank_id: str,
        directive_id: str,
        *,
        request_context: "RequestContext",
    ) -> bool:
        """Delete a directive.

        Args:
            bank_id: Bank identifier
            directive_id: Directive UUID
            request_context: Request context for authentication

        Returns:
            True if deleted, False if not found
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            result = await conn.execute(
                f"DELETE FROM {fq_table('directives')} WHERE bank_id = $1 AND id = $2",
                bank_id,
                directive_id,
            )

        return result == "DELETE 1"

    def _row_to_directive(self, row) -> dict[str, Any]:
        """Convert a database row to a directive dict."""
        return {
            "id": str(row["id"]),
            "bank_id": row["bank_id"],
            "name": row["name"],
            "content": row["content"],
            "priority": row["priority"],
            "is_active": row["is_active"],
            "tags": row["tags"] or [],
            "created_at": row["created_at"].isoformat() if row["created_at"] else None,
            "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
        }

    async def list_operations(
        self,
        bank_id: str,
        *,
        status: str | None = None,
        limit: int = 20,
        offset: int = 0,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """List async operations for a bank with optional filtering and pagination.

        Args:
            bank_id: Bank identifier
            status: Optional status filter (pending, completed, failed)
            limit: Maximum number of operations to return (default 20)
            offset: Number of operations to skip (default 0)
            request_context: Request context for authentication

        Returns:
            Dict with total count and list of operations, sorted by most recent first
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            # Build WHERE clause
            where_conditions = ["bank_id = $1"]
            params: list[Any] = [bank_id]

            if status:
                # Map API status to DB statuses (pending includes processing)
                if status == "pending":
                    where_conditions.append("status IN ('pending', 'processing')")
                else:
                    where_conditions.append(f"status = ${len(params) + 1}")
                    params.append(status)

            where_clause = " AND ".join(where_conditions)

            # Get total count (with filter)
            total_row = await conn.fetchrow(
                f"SELECT COUNT(*) as total FROM {fq_table('async_operations')} WHERE {where_clause}",
                *params,
            )
            total = total_row["total"] if total_row else 0

            # Get operations with pagination (include result_metadata to check for parent operations)
            operations = await conn.fetch(
                f"""
                SELECT operation_id, operation_type, created_at, status, error_message, result_metadata
                FROM {fq_table("async_operations")}
                WHERE {where_clause}
                ORDER BY created_at DESC
                LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}
                """,
                *params,
                limit,
                offset,
            )

            # Build operation list using status from database
            # Parent operations have their status updated when all children complete/fail
            operation_list = []
            for row in operations:
                # Map DB status to API status (pending includes processing)
                db_status = row["status"]
                api_status = "pending" if db_status in ("pending", "processing") else db_status

                operation_list.append(
                    {
                        "id": str(row["operation_id"]),
                        "task_type": row["operation_type"],
                        "items_count": 0,
                        "document_id": None,
                        "created_at": row["created_at"].isoformat(),
                        "status": api_status,
                        "error_message": row["error_message"],
                    }
                )

            return {
                "total": total,
                "operations": operation_list,
            }

    async def get_operation_status(
        self,
        bank_id: str,
        operation_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Get the status of a specific async operation.

        For parent operations, the status is automatically updated in the database when all children complete/fail.

        Returns:
            - status: "pending", "completed", or "failed" (from database)
            - updated_at: last update timestamp
            - completed_at: completion timestamp (if completed)
            - child_operations: (for parent operations) list of child operation statuses
        """
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        op_uuid = uuid.UUID(operation_id)

        async with acquire_with_retry(pool) as conn:
            row = await conn.fetchrow(
                f"""
                SELECT operation_id, operation_type, created_at, updated_at, completed_at, status, error_message, result_metadata
                FROM {fq_table("async_operations")}
                WHERE operation_id = $1 AND bank_id = $2
                """,
                op_uuid,
                bank_id,
            )

            if row:
                # Check if this is a parent operation
                result_metadata = json.loads(row["result_metadata"]) if row["result_metadata"] else {}
                is_parent = result_metadata.get("is_parent", False)

                # Use status from database (parent status is updated when all children complete/fail)
                db_status = row["status"]
                api_status = "pending" if db_status in ("pending", "processing") else db_status

                # For parent operations, include child operations list
                if is_parent:
                    # Query child operations
                    child_rows = await conn.fetch(
                        f"""
                        SELECT operation_id, status, error_message, result_metadata
                        FROM {fq_table("async_operations")}
                        WHERE bank_id = $1
                        AND result_metadata::jsonb @> $2::jsonb
                        ORDER BY (result_metadata->>'sub_batch_index')::int
                        """,
                        bank_id,
                        json.dumps({"parent_operation_id": operation_id}),
                    )

                    # Build child operations list and check if parent status needs updating
                    child_statuses = []
                    all_done = True
                    any_failed = False
                    all_completed = True

                    for child_row in child_rows:
                        child_metadata = (
                            json.loads(child_row["result_metadata"]) if child_row["result_metadata"] else {}
                        )
                        child_status = child_row["status"]

                        child_statuses.append(
                            {
                                "operation_id": str(child_row["operation_id"]),
                                "status": child_status,
                                "sub_batch_index": child_metadata.get("sub_batch_index"),
                                "items_count": child_metadata.get("items_count"),
                                "error_message": child_row["error_message"],
                            }
                        )

                        if child_status not in ("completed", "failed"):
                            all_done = False
                        if child_status == "failed":
                            any_failed = True
                        if child_status != "completed":
                            all_completed = False

                    # Self-healing: if parent status is out of sync with children, update it
                    if all_done and api_status == "pending":
                        correct_status = "failed" if any_failed else "completed"
                        logger.warning(
                            f"Parent operation {operation_id} status out of sync (DB: pending, should be: {correct_status}). Fixing."
                        )
                        await conn.execute(
                            f"""
                            UPDATE {fq_table("async_operations")}
                            SET status = $2, updated_at = NOW(), completed_at = NOW()
                            WHERE operation_id = $1
                            """,
                            op_uuid,
                            correct_status,
                        )
                        api_status = correct_status

                    return {
                        "operation_id": operation_id,
                        "status": api_status,
                        "operation_type": row["operation_type"],
                        "created_at": row["created_at"].isoformat() if row["created_at"] else None,
                        "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
                        "completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
                        "error_message": row["error_message"],
                        "result_metadata": result_metadata,
                        "child_operations": child_statuses,
                    }
                else:
                    # Regular operation (not a parent)
                    return {
                        "operation_id": operation_id,
                        "status": api_status,
                        "operation_type": row["operation_type"],
                        "created_at": row["created_at"].isoformat() if row["created_at"] else None,
                        "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
                        "completed_at": row["completed_at"].isoformat() if row["completed_at"] else None,
                        "error_message": row["error_message"],
                        "result_metadata": result_metadata,
                    }
            else:
                # Operation not found
                return {
                    "operation_id": operation_id,
                    "status": "not_found",
                    "operation_type": None,
                    "created_at": None,
                    "updated_at": None,
                    "completed_at": None,
                    "error_message": None,
                }

    async def cancel_operation(
        self,
        bank_id: str,
        operation_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Cancel a pending async operation."""
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        op_uuid = uuid.UUID(operation_id)

        async with acquire_with_retry(pool) as conn:
            # Check if operation exists and belongs to this memory bank
            result = await conn.fetchrow(
                f"SELECT bank_id FROM {fq_table('async_operations')} WHERE operation_id = $1 AND bank_id = $2",
                op_uuid,
                bank_id,
            )

            if not result:
                raise ValueError(f"Operation {operation_id} not found for bank {bank_id}")

            # Delete the operation
            await conn.execute(f"DELETE FROM {fq_table('async_operations')} WHERE operation_id = $1", op_uuid)

            return {
                "success": True,
                "message": f"Operation {operation_id} cancelled",
                "operation_id": operation_id,
                "bank_id": bank_id,
            }

    async def update_bank(
        self,
        bank_id: str,
        *,
        name: str | None = None,
        mission: str | None = None,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Update bank name and/or mission."""
        await self._authenticate_tenant(request_context)
        pool = await self._get_pool()

        async with acquire_with_retry(pool) as conn:
            if name is not None:
                await conn.execute(
                    f"""
                    UPDATE {fq_table("banks")}
                    SET name = $2, updated_at = NOW()
                    WHERE bank_id = $1
                    """,
                    bank_id,
                    name,
                )

            if mission is not None:
                await conn.execute(
                    f"""
                    UPDATE {fq_table("banks")}
                    SET mission = $2, updated_at = NOW()
                    WHERE bank_id = $1
                    """,
                    bank_id,
                    mission,
                )

        # Return updated profile
        return await self.get_bank_profile(bank_id, request_context=request_context)

    async def _submit_async_operation(
        self,
        bank_id: str,
        operation_type: str,
        task_type: str,
        task_payload: dict[str, Any],
        *,
        result_metadata: dict[str, Any] | None = None,
        dedupe_by_bank: bool = False,
    ) -> dict[str, Any]:
        """Generic helper to submit an async operation.

        Args:
            bank_id: Bank identifier
            operation_type: Operation type for the async_operations record (e.g., 'consolidation', 'retain')
            task_type: Task type for the task payload (e.g., 'consolidation', 'batch_retain')
            task_payload: Additional task payload fields (operation_id and bank_id are added automatically)
            result_metadata: Optional metadata to store with the operation record
            dedupe_by_bank: If True, skip creating a new task if one is already pending for this bank+operation_type

        Returns:
            Dict with operation_id and optionally deduplicated=True if an existing task was found
        """
        import json

        pool = await self._get_pool()

        # Check for existing pending task if deduplication is enabled
        # Note: We only check 'pending', not 'processing', because a processing task
        # uses a watermark from when it started - new memories added after that point
        # would need another consolidation run to be processed.
        if dedupe_by_bank:
            async with acquire_with_retry(pool) as conn:
                existing = await conn.fetchrow(
                    f"""
                    SELECT operation_id FROM {fq_table("async_operations")}
                    WHERE bank_id = $1 AND operation_type = $2 AND status = 'pending'
                    LIMIT 1
                    """,
                    bank_id,
                    operation_type,
                )
                if existing:
                    logger.debug(
                        f"{operation_type} task already pending for bank_id={bank_id}, "
                        f"skipping duplicate (existing operation_id={existing['operation_id']})"
                    )
                    return {
                        "operation_id": str(existing["operation_id"]),
                        "deduplicated": True,
                    }

        operation_id = uuid.uuid4()

        # Insert operation record into database
        async with acquire_with_retry(pool) as conn:
            await conn.execute(
                f"""
                INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status)
                VALUES ($1, $2, $3, $4, $5)
                """,
                operation_id,
                bank_id,
                operation_type,
                json.dumps(result_metadata or {}),
                "pending",
            )

        # Build and submit task payload
        full_payload = {
            "type": task_type,
            "operation_id": str(operation_id),
            "bank_id": bank_id,
            **task_payload,
        }

        await self._task_backend.submit_task(full_payload)

        logger.info(f"{operation_type} task queued for bank_id={bank_id}, operation_id={operation_id}")

        return {
            "operation_id": str(operation_id),
        }

    async def submit_async_retain(
        self,
        bank_id: str,
        contents: list[dict[str, Any]],
        *,
        request_context: "RequestContext",
        document_tags: list[str] | None = None,
    ) -> dict[str, Any]:
        """Submit a batch retain operation to run asynchronously.

        For large batches (exceeding retain_batch_chars threshold), automatically splits
        into smaller sub-batches and creates a parent operation that tracks all children.
        """
        await self._authenticate_tenant(request_context)

        # Validate no duplicate document_ids in the batch
        # Having duplicate document_ids causes race conditions in document upserts during parallel processing
        doc_ids = [item.get("document_id") for item in contents if item.get("document_id")]
        if len(doc_ids) != len(set(doc_ids)):
            from collections import Counter

            duplicates = [doc_id for doc_id, count in Counter(doc_ids).items() if count > 1]
            raise ValueError(
                f"Batch contains duplicate document_ids: {duplicates}. "
                f"Each content item in a batch must have a unique document_id to avoid race conditions."
            )

        # Calculate total token count and determine if we need to split
        total_tokens = sum(count_tokens(item.get("content", "")) for item in contents)
        config = get_config()
        tokens_per_batch = config.retain_batch_tokens

        # Split into sub-batches based on token count
        sub_batches = []
        current_batch = []
        current_batch_tokens = 0

        for item in contents:
            item_tokens = count_tokens(item.get("content", ""))

            # If adding this item would exceed the limit, start a new batch
            # (unless current batch is empty - then we must include it even if it's large)
            if current_batch and current_batch_tokens + item_tokens > tokens_per_batch:
                sub_batches.append(current_batch)
                current_batch = [item]
                current_batch_tokens = item_tokens
            else:
                current_batch.append(item)
                current_batch_tokens += item_tokens

        # Add the last batch
        if current_batch:
            sub_batches.append(current_batch)

        # Log splitting info if we actually split
        if len(sub_batches) > 1:
            logger.info(
                f"Large async retain batch ({total_tokens:,} tokens from {len(contents)} items). "
                f"Split into {len(sub_batches)} sub-batches: {[len(b) for b in sub_batches]} items each"
            )

        # Always create parent operation (even for single batch - simpler, more reliable code path)
        import uuid

        parent_operation_id = uuid.uuid4()
        pool = await self._get_pool()

        # Create typed metadata for parent operation
        parent_metadata = BatchRetainParentMetadata(
            items_count=len(contents),
            total_tokens=total_tokens,
            num_sub_batches=len(sub_batches),
        )

        async with acquire_with_retry(pool) as conn:
            await conn.execute(
                f"""
                INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status)
                VALUES ($1, $2, $3, $4, $5)
                """,
                parent_operation_id,
                bank_id,
                "batch_retain",
                json.dumps(parent_metadata.to_dict()),
                "pending",  # Will be updated by status aggregation
            )

        logger.info(f"Created parent operation {parent_operation_id} for {len(sub_batches)} sub-batch(es)")

        # Submit child operations for each sub-batch
        for i, sub_batch in enumerate(sub_batches, 1):
            if len(sub_batches) > 1:
                sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch)
                logger.info(
                    f"Submitting sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
                )

            task_payload: dict[str, Any] = {"contents": sub_batch}
            if document_tags:
                task_payload["document_tags"] = document_tags
            # Pass tenant_id and api_key_id through task payload
            if request_context.tenant_id:
                task_payload["_tenant_id"] = request_context.tenant_id
            if request_context.api_key_id:
                task_payload["_api_key_id"] = request_context.api_key_id

            # Create typed metadata for child operation
            child_metadata = BatchRetainChildMetadata(
                items_count=len(sub_batch),
                parent_operation_id=str(parent_operation_id),
                sub_batch_index=i,
                total_sub_batches=len(sub_batches),
            )

            # Create child operation with reference to parent
            await self._submit_async_operation(
                bank_id=bank_id,
                operation_type="retain",
                task_type="batch_retain",
                task_payload=task_payload,
                result_metadata=child_metadata.to_dict(),
                dedupe_by_bank=False,
            )

        return {
            "operation_id": str(parent_operation_id),
            "items_count": len(contents),
        }

    async def submit_async_file_retain(
        self,
        bank_id: str,
        file_items: list[dict[str, Any]],
        parser: str,
        document_tags: list[str] | None,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """
        Submit batch file conversion + retain operation.

        Each file is converted to markdown and then retained as a memory.
        Files are stored in object storage and conversion happens asynchronously.

        Args:
            bank_id: Bank ID
            file_items: List of file items, each containing:
                - file: UploadFile object (FastAPI)
                - document_id: Document ID
                - context: Optional context
                - metadata: Optional metadata dict
                - tags: Optional tags list
                - timestamp: Optional timestamp
            parser: Parser name (e.g., "markitdown")
            document_tags: Tags applied to all documents
            request_context: Request context for authentication

        Returns:
            dict with operation_id and files_count
        """
        await self._authenticate_tenant(request_context)

        config = get_config()

        # Validate file count
        if len(file_items) > config.file_conversion_max_batch_size:
            raise ValueError(f"Too many files. Maximum {config.file_conversion_max_batch_size} files per request.")

        # Read all files and validate total batch size
        files_data = []
        total_batch_size = 0

        for item in file_items:
            file = item["file"]
            file_data = await file.read()
            total_batch_size += len(file_data)
            files_data.append((item, file, file_data))

        # Validate total batch size
        if total_batch_size > config.file_conversion_max_batch_size_bytes:
            total_mb = total_batch_size / (1024 * 1024)
            raise ValueError(
                f"Total batch size ({total_mb:.1f}MB) exceeds maximum of {config.file_conversion_max_batch_size_mb}MB"
            )

        # Submit individual operation for each file
        operation_ids = []
        for item, file, file_data in files_data:
            # Generate storage key
            storage_key = f"banks/{bank_id}/files/{item['document_id']}/{file.filename}"

            # Store file in object storage
            await self._file_storage.store(
                file_data=file_data,
                key=storage_key,
                metadata={
                    "content_type": file.content_type or "application/octet-stream",
                    "original_filename": file.filename,
                    "bank_id": bank_id,
                    "document_id": item["document_id"],
                },
            )

            # Create individual operation and submit task
            task_payload: dict[str, Any] = {
                "document_id": item["document_id"],
                "storage_key": storage_key,
                "original_filename": file.filename,
                "content_type": file.content_type or "application/octet-stream",
                "parser": parser,
                "context": item.get("context"),
                "metadata": item.get("metadata", {}),
                "tags": item.get("tags", []),
                "document_tags": document_tags or [],
                "timestamp": item.get("timestamp"),
            }

            # Pass tenant_id and api_key_id through task payload
            if request_context.tenant_id:
                task_payload["_tenant_id"] = request_context.tenant_id
            if request_context.api_key_id:
                task_payload["_api_key_id"] = request_context.api_key_id

            result = await self._submit_async_operation(
                bank_id=bank_id,
                operation_type="file_convert_retain",
                task_type="file_convert_retain",
                task_payload=task_payload,
                result_metadata={
                    "original_filename": file.filename,
                },
                dedupe_by_bank=False,
            )
            operation_ids.append(result["operation_id"])

        return {
            "operation_ids": operation_ids,
            "files_count": len(file_items),
        }

    async def submit_async_consolidation(
        self,
        bank_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Submit a consolidation operation to run asynchronously.

        Deduplicates by bank_id - if there's already a pending consolidation for this bank,
        returns the existing operation_id instead of creating a new one.

        Args:
            bank_id: Bank identifier
            request_context: Request context for authentication

        Returns:
            Dict with operation_id
        """
        await self._authenticate_tenant(request_context)

        # Pass tenant_id and api_key_id through task payload so the worker
        # can provide request context to extension hooks (e.g., usage metering
        # for mental model refreshes triggered by consolidation).
        task_payload: dict[str, Any] = {}
        if request_context.tenant_id:
            task_payload["_tenant_id"] = request_context.tenant_id
        if request_context.api_key_id:
            task_payload["_api_key_id"] = request_context.api_key_id

        return await self._submit_async_operation(
            bank_id=bank_id,
            operation_type="consolidation",
            task_type="consolidation",
            task_payload=task_payload,
            dedupe_by_bank=True,
        )

    async def submit_async_refresh_mental_model(
        self,
        bank_id: str,
        mental_model_id: str,
        *,
        request_context: "RequestContext",
    ) -> dict[str, Any]:
        """Submit an async mental model refresh operation.

        This schedules a background task to re-run the source query and update the content.

        Args:
            bank_id: Bank identifier
            mental_model_id: Mental model UUID to refresh
            request_context: Request context for authentication

        Returns:
            Dict with operation_id
        """
        await self._authenticate_tenant(request_context)

        # Pre-operation validation (credit check)
        if self._operation_validator:
            from hindsight_api.extensions.operation_validator import MentalModelRefreshContext

            ctx = MentalModelRefreshContext(
                bank_id=bank_id,
                mental_model_id=mental_model_id,
                request_context=request_context,
            )
            await self._validate_operation(self._operation_validator.validate_mental_model_refresh(ctx))

        # Verify mental model exists
        mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=request_context)
        if not mental_model:
            raise ValueError(f"Mental model {mental_model_id} not found in bank {bank_id}")

        # Pass tenant_id and api_key_id through task payload so the worker
        # can provide request context to extension hooks.
        task_payload: dict[str, Any] = {
            "mental_model_id": mental_model_id,
        }
        if request_context.tenant_id:
            task_payload["_tenant_id"] = request_context.tenant_id
        if request_context.api_key_id:
            task_payload["_api_key_id"] = request_context.api_key_id

        return await self._submit_async_operation(
            bank_id=bank_id,
            operation_type="refresh_mental_model",
            task_type="refresh_mental_model",
            task_payload=task_payload,
            result_metadata={"mental_model_id": mental_model_id, "name": mental_model["name"]},
            dedupe_by_bank=False,
        )