fix: multi-tenant schema context for worker task execution (#208)

Background tasks (async retain, consolidation, reflections) fail in
multi-tenant deployments because the worker executes tasks without
setting the tenant schema context. This causes two failures:

1. The cancellation check in execute_task queries public.async_operations
   instead of the tenant's schema, finds no row, and skips the task as
   "cancelled" — even though it wasn't.

2. Even if that were fixed, _authenticate_tenant would throw
   AuthenticationError because background tasks have no API key.

Changes:
- Poller passes task.schema into task_dict so execute_task can set it
- execute_task sets _current_schema before the cancellation check
- Task handlers use RequestContext(internal=True) to signal background ops
- _authenticate_tenant skips extension auth for internal requests when
  schema is already set
- BrokerTaskBackend uses schema_getter for dynamic schema resolution
  when submitting tasks and waiting for results
- Pass tenant_extension to WorkerPoller in create_app

hindsight-api/hindsight_api/api/http.py

@@ -1373,6 +1373,7 @@ def create_app(
                 poll_interval_ms=config.worker_poll_interval_ms,
                 batch_size=config.worker_batch_size,
                 max_retries=config.worker_max_retries,
+                tenant_extension=getattr(memory, "_tenant_extension", None),
             )
             poller_task = asyncio.create_task(poller.run())
             logging.info(f"Worker poller started (worker_id={worker_id})")

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -433,7 +433,10 @@ class MemoryEngine(MemoryEngineInterface):
         # 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)
+        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)
@@ -497,6 +500,13 @@ class MemoryEngine(MemoryEngineInterface):
         if request_context is None:
             raise AuthenticationError("RequestContext is required when tenant extension is configured")
 
+        # For internal/background operations (e.g., worker tasks), skip extension authentication
+        # if the schema has already been set by execute_task via the _schema field.
+        if request_context.internal:
+            current = _current_schema.get()
+            if current and current != "public":
+                return current
+
         # Let AuthenticationError propagate - HTTP layer will convert to 401
         tenant_context = await self._tenant_extension.authenticate(request_context)
 
@@ -523,10 +533,10 @@ class MemoryEngine(MemoryEngineInterface):
             f"[BATCH_RETAIN_TASK] Starting background batch retain for bank_id={bank_id}, {len(contents)} items"
         )
 
-        # Use internal request context for background tasks
+        # Use internal request context for background tasks (skips tenant auth when schema is pre-set)
         from hindsight_api.models import RequestContext
 
-        internal_context = RequestContext()
+        internal_context = RequestContext(internal=True)
         await self.retain_batch_async(bank_id=bank_id, contents=contents, request_context=internal_context)
 
         logger.info(f"[BATCH_RETAIN_TASK] Completed background batch retain for bank_id={bank_id}")
@@ -552,7 +562,7 @@ class MemoryEngine(MemoryEngineInterface):
 
         from .consolidation import run_consolidation_job
 
-        internal_context = RequestContext()
+        internal_context = RequestContext(internal=True)
         result = await run_consolidation_job(
             memory_engine=self,
             bank_id=bank_id,
@@ -587,7 +597,7 @@ class MemoryEngine(MemoryEngineInterface):
 
         from hindsight_api.models import RequestContext
 
-        internal_context = RequestContext()
+        internal_context = RequestContext(internal=True)
 
         # Run reflect to generate content
         reflect_result = await self.reflect_async(
@@ -649,7 +659,7 @@ class MemoryEngine(MemoryEngineInterface):
 
         from hindsight_api.models import RequestContext
 
-        internal_context = RequestContext()
+        internal_context = RequestContext(internal=True)
 
         # 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)
@@ -711,6 +721,11 @@ class MemoryEngine(MemoryEngineInterface):
         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:

hindsight-api/hindsight_api/engine/task_backend.py

@@ -144,17 +144,21 @@ class BrokerTaskBackend(TaskBackend):
         self,
         pool_getter: Callable[[], "asyncpg.Pool"],
         schema: str | None = None,
+        schema_getter: Callable[[], str | None] | None = None,
     ):
         """
         Initialize the broker task backend.
 
         Args:
             pool_getter: Callable that returns the asyncpg connection pool
-            schema: Database schema for multi-tenant support (optional)
+            schema: Database schema for multi-tenant support (optional, static)
+            schema_getter: Callable that returns current schema dynamically (optional).
+                          If set, takes precedence over static schema for submit_task.
         """
         super().__init__()
         self._pool_getter = pool_getter
         self._schema = schema
+        self._schema_getter = schema_getter
 
     async def initialize(self):
         """Initialize the backend."""
@@ -180,7 +184,8 @@ class BrokerTaskBackend(TaskBackend):
         bank_id = task_dict.get("bank_id")
         payload_json = json.dumps(task_dict)
 
-        table = fq_table("async_operations", self._schema)
+        schema = self._schema_getter() if self._schema_getter else self._schema
+        table = fq_table("async_operations", schema)
 
         if operation_id:
             # Update existing operation with task payload
@@ -231,7 +236,8 @@ class BrokerTaskBackend(TaskBackend):
         import asyncio
 
         pool = self._pool_getter()
-        table = fq_table("async_operations", self._schema)
+        schema = self._schema_getter() if self._schema_getter else self._schema
+        table = fq_table("async_operations", schema)
 
         start_time = asyncio.get_event_loop().time()
         while asyncio.get_event_loop().time() - start_time < timeout:

hindsight-api/hindsight_api/worker/poller.py

@@ -261,6 +261,9 @@ class WorkerPoller:
         try:
             schema_info = f", schema={task.schema}" if task.schema else ""
             logger.debug(f"Executing task {task.operation_id} (type={task_type}, bank={bank_id}{schema_info})")
+            # Pass schema to executor so it can set the correct context
+            if task.schema:
+                task.task_dict["_schema"] = task.schema
             await self._executor(task.task_dict)
             await self._mark_completed(task.operation_id, task.schema)
             logger.debug(f"Task {task.operation_id} completed successfully")