fleet-memory/hindsight-api/hindsight_api/extensions/context.py
Chris Bartholomew ce45d301ce
Add operation validator extension support with proper HTTP error handling (#86)
* Load operation validator extension in main entry point

Enable the operation validator extension to be loaded from environment
configuration and passed to MemoryEngine, allowing pre/post operation
hooks for usage metering, rate limiting, and audit logging.

* Fix reflect background task authentication and add internal flag

- Pass API key to background opinion storage task for proper auth
- Add internal flag to RequestContext for tracking internal operations
- Background opinion storage now authenticates correctly with tenant

* Add api_key_id to RequestContext for usage tracking

- Add api_key_id field to RequestContext to track which API key was used
- Enables per-API-key usage analytics in the metering system

* Fix HTTP error handling for authentication and validation errors

- Add status_code parameter to ValidationResult and OperationValidationError
- Convert OperationValidationError to HTTPException with proper status codes
- Fix authentication errors to return 401 instead of raising internal errors
- Re-raise HTTPException in exception handlers to prevent swallowing errors

* Fix AuthenticationError handling in memory engine

- Raise AuthenticationError from memory_engine._authenticate_tenant instead
  of HTTPException so unit tests pass
- Add AuthenticationError handling in HTTP layer to convert to 401 responses
- Fixes failing TestMemoryEngineTenantAuth tests

* Add global exception handler for AuthenticationError

Returns proper 401 status code for all authentication failures
across all endpoints, not just the ones with explicit handlers.

* Simplify exception handling: use global AuthenticationError handler

- Remove redundant individual exception handlers
- Add 'except AuthenticationError: raise' before generic Exception handlers
  to let global handler process auth errors uniformly

* Refactor background tasks to use tenant_id instead of api_key

This makes the core more generic - it passes tenant_id (which is
extension-agnostic) rather than api_key (which is cloud-specific).

- Add tenant_id field to RequestContext
- Pass tenant_id instead of api_key to background tasks
- Extensions can check internal=True with tenant_id to bypass normal auth

* Fix exception propagation: include HTTPException in re-raise

After cleanup of redundant exception handlers, 404 errors were
returning 500 because HTTPException was caught by the generic
except Exception handler. Fixed by combining AuthenticationError
and HTTPException in the re-raise pattern.
2026-01-01 20:19:52 -05:00

117 lines
3.8 KiB
Python

"""Extension context providing a controlled API for extensions to interact with the system."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from hindsight_api.engine.interface import MemoryEngineInterface
class ExtensionContext(ABC):
"""
Abstract context providing a controlled API for extensions.
Extensions receive this context instead of direct access to internal
components like MemoryEngine or database connections. This provides:
- A stable API that won't break when internals change
- Security by limiting what extensions can access
- Clear documentation of what extensions can do
Built-in implementation:
hindsight_api.extensions.builtin.context.DefaultExtensionContext
Example usage in an extension:
class MyTenantExtension(TenantExtension):
async def on_startup(self) -> None:
# Run migrations for a new tenant schema
await self.context.run_migration("tenant_acme")
class MyHttpExtension(HttpExtension):
def get_router(self, memory):
# Use memory engine for custom endpoints
engine = self.context.get_memory_engine()
...
"""
@abstractmethod
async def run_migration(self, schema: str) -> None:
"""
Run database migrations for a specific schema.
This creates the schema if it doesn't exist and runs all pending
migrations. Uses advisory locks to coordinate between distributed workers.
Args:
schema: PostgreSQL schema name (e.g., "tenant_acme").
The schema will be created if it doesn't exist.
Raises:
RuntimeError: If migrations fail to complete.
Example:
# Provision a new tenant schema
await context.run_migration("tenant_acme")
"""
...
@abstractmethod
def get_memory_engine(self) -> "MemoryEngineInterface":
"""
Get the memory engine interface.
Returns the MemoryEngineInterface for performing memory operations
like retain, recall, reflect, and entity/document management.
Returns:
MemoryEngineInterface instance.
Example:
engine = context.get_memory_engine()
result = await engine.recall_async(bank_id, query)
"""
...
class DefaultExtensionContext(ExtensionContext):
"""
Default implementation of ExtensionContext.
Uses the system's database URL and migration infrastructure.
"""
def __init__(
self,
database_url: str,
memory_engine: "MemoryEngineInterface | None" = None,
):
"""
Initialize the context.
Args:
database_url: SQLAlchemy database URL for migrations.
memory_engine: Optional MemoryEngine instance for memory operations.
"""
self._database_url = database_url
self._memory_engine = memory_engine
async def run_migration(self, schema: str) -> None:
"""Run migrations for a specific schema."""
from hindsight_api.migrations import run_migrations
# Prefer getting URL from memory engine (handles pg0 case where URL is set after init)
db_url = self._database_url
if self._memory_engine is not None:
engine_url = getattr(self._memory_engine, "db_url", None)
if engine_url:
db_url = engine_url
run_migrations(db_url, schema=schema)
def get_memory_engine(self) -> "MemoryEngineInterface":
"""Get the memory engine interface."""
if self._memory_engine is None:
raise RuntimeError(
"Memory engine not configured in ExtensionContext. "
"Ensure the context was created with a memory_engine parameter."
)
return self._memory_engine