diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index f50b59cf..b47aef47 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/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})") diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index a0ae8016..5c301ba2 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/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: diff --git a/hindsight-api/hindsight_api/engine/task_backend.py b/hindsight-api/hindsight_api/engine/task_backend.py index 29420463..cfca7a81 100644 --- a/hindsight-api/hindsight_api/engine/task_backend.py +++ b/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: diff --git a/hindsight-api/hindsight_api/worker/poller.py b/hindsight-api/hindsight_api/worker/poller.py index 3cd307d5..a0e77ec1 100644 --- a/hindsight-api/hindsight_api/worker/poller.py +++ b/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")