diff --git a/hindsight-api/hindsight_api/alembic/versions/a1b2c3d4e5f6_add_file_storage_table.py b/hindsight-api/hindsight_api/alembic/versions/a1b2c3d4e5f6_add_file_storage_table.py index 4187081d..58fb7750 100644 --- a/hindsight-api/hindsight_api/alembic/versions/a1b2c3d4e5f6_add_file_storage_table.py +++ b/hindsight-api/hindsight_api/alembic/versions/a1b2c3d4e5f6_add_file_storage_table.py @@ -34,7 +34,7 @@ def upgrade() -> None: # Create file_storage table (minimal: just key + data) op.execute( f""" - CREATE TABLE {schema}file_storage ( + CREATE TABLE IF NOT EXISTS {schema}file_storage ( storage_key TEXT PRIMARY KEY, data BYTEA NOT NULL ) diff --git a/hindsight-api/hindsight_api/alembic/versions/e5f6g7h8i9j0_cascade_delete_ops_on_bank_delete.py b/hindsight-api/hindsight_api/alembic/versions/e5f6g7h8i9j0_cascade_delete_ops_on_bank_delete.py new file mode 100644 index 00000000..6f941dc1 --- /dev/null +++ b/hindsight-api/hindsight_api/alembic/versions/e5f6g7h8i9j0_cascade_delete_ops_on_bank_delete.py @@ -0,0 +1,73 @@ +"""Add CASCADE DELETE FK from async_operations and webhooks to banks. + +When a bank is deleted, all its async_operations and webhooks rows are +automatically deleted by the database. This ensures that any in-flight +worker tasks detect the deletion via _check_op_alive() and abort early. + +Revision ID: e5f6g7h8i9j0 +Revises: d4e5f6g7h8i9 +Create Date: 2026-03-11 +""" + +from collections.abc import Sequence + +from alembic import context, op + +revision: str = "e5f6g7h8i9j0" +down_revision: str | Sequence[str] | None = "d4e5f6g7h8i9" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _get_schema_prefix() -> str: + schema = context.config.get_main_option("target_schema") + return f'"{schema}".' if schema else "" + + +def upgrade() -> None: + schema = _get_schema_prefix() + + # Remove orphaned async_operations rows whose bank no longer exists + # (can happen because there was no FK before this migration). + op.execute( + f""" + DELETE FROM {schema}async_operations + WHERE bank_id IS NOT NULL + AND bank_id NOT IN (SELECT bank_id FROM {schema}banks) + """ + ) + + # Remove orphaned webhooks rows whose bank no longer exists. + op.execute( + f""" + DELETE FROM {schema}webhooks + WHERE bank_id IS NOT NULL + AND bank_id NOT IN (SELECT bank_id FROM {schema}banks) + """ + ) + + # Add FK with ON DELETE CASCADE so that deleting a bank automatically + # cleans up all its pending/processing operations and webhook configs. + op.execute( + f""" + ALTER TABLE {schema}async_operations + ADD CONSTRAINT fk_async_operations_bank_id + FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id) + ON DELETE CASCADE + """ + ) + + op.execute( + f""" + ALTER TABLE {schema}webhooks + ADD CONSTRAINT fk_webhooks_bank_id + FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id) + ON DELETE CASCADE + """ + ) + + +def downgrade() -> None: + schema = _get_schema_prefix() + op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS fk_async_operations_bank_id") + op.execute(f"ALTER TABLE {schema}webhooks DROP CONSTRAINT IF EXISTS fk_webhooks_bank_id") diff --git a/hindsight-api/hindsight_api/alembic/versions/y0t1u2v3w4x5_add_result_metadata_gin_index.py b/hindsight-api/hindsight_api/alembic/versions/y0t1u2v3w4x5_add_result_metadata_gin_index.py index 11a143d2..d3b49daa 100644 --- a/hindsight-api/hindsight_api/alembic/versions/y0t1u2v3w4x5_add_result_metadata_gin_index.py +++ b/hindsight-api/hindsight_api/alembic/versions/y0t1u2v3w4x5_add_result_metadata_gin_index.py @@ -35,7 +35,7 @@ def upgrade() -> None: # Add GIN index for JSONB containment queries (@> operator) op.execute(f""" - CREATE INDEX idx_async_operations_result_metadata + CREATE INDEX IF NOT EXISTS idx_async_operations_result_metadata ON {schema}async_operations USING gin(result_metadata) """) diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 0fc8d245..401f88d7 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -4091,6 +4091,10 @@ def _register_routes(app: FastAPI): try: pool = await app.state.memory._get_pool() from hindsight_api.engine.memory_engine import fq_table + from hindsight_api.engine.retain import bank_utils + + # Ensure the bank row exists before inserting into webhooks (FK constraint). + await bank_utils.get_bank_profile(pool, bank_id) webhook_id = uuid.uuid4() now = datetime.utcnow().isoformat() + "Z" diff --git a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py index 80aafd42..578172d2 100644 --- a/hindsight-api/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api/hindsight_api/engine/consolidation/consolidator.py @@ -161,6 +161,7 @@ async def run_consolidation_job( memory_engine: "MemoryEngine", bank_id: str, request_context: "RequestContext", + operation_id: str | None = None, ) -> dict[str, Any]: """ Run consolidation job for a bank. @@ -386,6 +387,13 @@ async def run_consolidation_job( [(m["id"],) for m in llm_batch], ) + # Checkpoint: abort if the operation (and thus the bank) was deleted mid-run. + if operation_id and not await memory_engine._check_op_alive(operation_id): + logger.info( + f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early" + ) + return {"status": "cancelled", "bank_id": bank_id, **stats} + for result in results: stats["memories_processed"] += 1 action = result.get("action") diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 8b186143..a5e0b778 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -820,6 +820,7 @@ class MemoryEngine(MemoryEngineInterface): memory_engine=self, bank_id=bank_id, request_context=internal_context, + operation_id=task_dict.get("operation_id"), ) logger.info(f"[CONSOLIDATION] bank={bank_id} completed: {result.get('memories_processed', 0)} processed") @@ -1249,6 +1250,24 @@ class MemoryEngine(MemoryEngineInterface): except Exception as e: logger.error(f"Failed to delete async operation record {operation_id}: {e}") + async def _check_op_alive(self, operation_id: str) -> bool: + """Return False if the operation row no longer exists (e.g. bank was deleted via CASCADE). + + Long-running operations should call this at natural checkpoints (e.g. after each + committed batch) to detect bank deletion early and abort cleanly. + """ + try: + pool = await self._get_pool() + async with acquire_with_retry(pool) as conn: + row = await conn.fetchrow( + f"SELECT operation_id FROM {fq_table('async_operations')} WHERE operation_id = $1", + uuid.UUID(operation_id), + ) + return row is not None + except Exception as e: + logger.error(f"Failed to check operation liveness {operation_id}: {e}") + return True # Assume alive on DB error to avoid false-positive aborts + 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. @@ -1264,15 +1283,19 @@ class MemoryEngine(MemoryEngineInterface): async with acquire_with_retry(pool) as conn: async with conn.transaction(): # Mark this operation as failed - await conn.execute( + row = await conn.fetchrow( f""" UPDATE {fq_table("async_operations")} SET status = 'failed', error_message = $2, updated_at = NOW() WHERE operation_id = $1 + RETURNING operation_id """, uuid.UUID(operation_id), truncated_error, ) + if row is None: + logger.info(f"Operation {operation_id} no longer exists (bank deleted), skipping mark-failed") + return 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 @@ -1292,14 +1315,20 @@ class MemoryEngine(MemoryEngineInterface): async with acquire_with_retry(pool) as conn: async with conn.transaction(): # Mark this operation as completed - await conn.execute( + row = await conn.fetchrow( f""" UPDATE {fq_table("async_operations")} SET status = 'completed', updated_at = NOW(), completed_at = NOW() WHERE operation_id = $1 + RETURNING operation_id """, uuid.UUID(operation_id), ) + if row is None: + logger.info( + f"Operation {operation_id} no longer exists (bank deleted), skipping mark-completed" + ) + return 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 @@ -1329,14 +1358,20 @@ class MemoryEngine(MemoryEngineInterface): pool = await self._get_pool() async with acquire_with_retry(pool) as conn: async with conn.transaction(): - await conn.execute( + row = await conn.fetchrow( f""" UPDATE {fq_table("async_operations")} SET status = 'completed', updated_at = NOW(), completed_at = NOW() WHERE operation_id = $1 + RETURNING operation_id """, uuid.UUID(operation_id), ) + if row is None: + logger.info( + f"Operation {operation_id} no longer exists (bank deleted), skipping mark-completed" + ) + return logger.info(f"Marked async operation as completed: {operation_id}") await self._maybe_update_parent_operation(operation_id, conn) @@ -2057,6 +2092,15 @@ class MemoryEngine(MemoryEngineInterface): # Process each sub-batch all_results = [] for i, sub_batch in enumerate(sub_batches, 1): + # Checkpoint: abort if the operation was deleted (bank was deleted) between sub-batches. + if operation_id and not await self._check_op_alive(operation_id): + logger.info( + f"[BATCH_RETAIN] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping after {i - 1}/{len(sub_batches)} sub-batches" + ) + if return_usage: + return all_results, total_usage + return all_results + 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" @@ -7402,6 +7446,10 @@ class MemoryEngine(MemoryEngineInterface): parent_operation_id = uuid.uuid4() pool = await self._get_pool() + # Ensure the bank row exists before inserting async_operations (which now has a FK). + # Banks are created lazily on first retain, but the FK requires the row to exist first. + await bank_utils.get_bank_profile(pool, bank_id) + # Create typed metadata for parent operation parent_metadata = BatchRetainParentMetadata( items_count=len(contents), diff --git a/hindsight-api/tests/test_async_batch_retain.py b/hindsight-api/tests/test_async_batch_retain.py index d82d502d..625c5f06 100644 --- a/hindsight-api/tests/test_async_batch_retain.py +++ b/hindsight-api/tests/test_async_batch_retain.py @@ -9,6 +9,15 @@ import pytest from hindsight_api.extensions import RequestContext +async def _ensure_bank(pool, bank_id: str) -> None: + """Upsert a minimal bank row so FK on async_operations passes.""" + await pool.execute( + "INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING", + bank_id, + bank_id, + ) + + @pytest.mark.asyncio async def test_duplicate_document_ids_rejected_async(memory, request_context): """Test that async retain rejects batches with duplicate document_ids.""" @@ -156,6 +165,7 @@ async def test_parent_operation_status_aggregation_pending(memory, request_conte """Test that parent operation shows 'pending' when children are pending.""" bank_id = "test_parent_pending" pool = await memory._get_pool() + await _ensure_bank(pool, bank_id) # Manually create a parent operation parent_id = uuid.uuid4() @@ -230,6 +240,7 @@ async def test_parent_operation_status_aggregation_failed(memory, request_contex """Test that parent operation shows 'failed' when any child fails.""" bank_id = "test_parent_failed" pool = await memory._get_pool() + await _ensure_bank(pool, bank_id) # Manually create a parent operation parent_id = uuid.uuid4() @@ -309,6 +320,7 @@ async def test_parent_operation_status_aggregation_completed(memory, request_con """Test that parent operation shows 'completed' when all children are completed.""" bank_id = "test_parent_completed" pool = await memory._get_pool() + await _ensure_bank(pool, bank_id) # Manually create a parent operation parent_id = uuid.uuid4() diff --git a/hindsight-api/tests/test_async_retain_tags.py b/hindsight-api/tests/test_async_retain_tags.py index 83c6e2bc..ed24a4ab 100644 --- a/hindsight-api/tests/test_async_retain_tags.py +++ b/hindsight-api/tests/test_async_retain_tags.py @@ -1,6 +1,6 @@ """Unit tests for async retain tag propagation.""" -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -34,13 +34,14 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload(): contents = [{"content": "Async retain payload test."}] document_tags = ["scope:tools", "user:alice"] - result = await MemoryEngine.submit_async_retain( - engine, - bank_id="bank-1", - contents=contents, - document_tags=document_tags, - request_context=request_context, - ) + with patch("hindsight_api.engine.memory_engine.bank_utils.get_bank_profile", new_callable=AsyncMock): + result = await MemoryEngine.submit_async_retain( + engine, + bank_id="bank-1", + contents=contents, + document_tags=document_tags, + request_context=request_context, + ) # Check result structure assert "operation_id" in result diff --git a/hindsight-api/tests/test_op_cancellation.py b/hindsight-api/tests/test_op_cancellation.py new file mode 100644 index 00000000..3b3039ac --- /dev/null +++ b/hindsight-api/tests/test_op_cancellation.py @@ -0,0 +1,311 @@ +"""Tests for operation cancellation when a bank is deleted. + +Covers: +- CASCADE DELETE: deleting a bank removes async_operations and webhooks rows +- _check_op_alive: returns True when op exists, False when deleted +- _mark_operation_completed / _mark_operation_failed: graceful no-op when row is gone +- Consolidation checkpoint: stops early after a batch commit if op was deleted +- Retain checkpoint: stops between sub-batches if op was deleted +""" + +import uuid +from unittest.mock import AsyncMock, patch + +import pytest +import pytest_asyncio + +from hindsight_api.engine.memory_engine import MemoryEngine + + +pytestmark = pytest.mark.xdist_group("op_cancellation_tests") + +_BANK_PREFIX = "test-op-cancel" + + +@pytest_asyncio.fixture +async def pool(pg0_db_url): + import asyncpg + from hindsight_api.pg0 import resolve_database_url + + resolved_url = await resolve_database_url(pg0_db_url) + p = await asyncpg.create_pool(resolved_url, min_size=1, max_size=5, command_timeout=30) + yield p + await p.close() + + +@pytest_asyncio.fixture(autouse=True) +async def cleanup(pool): + """Remove test rows before and after each test.""" + await pool.execute(f"DELETE FROM banks WHERE bank_id LIKE '{_BANK_PREFIX}%'") + yield + await pool.execute(f"DELETE FROM banks WHERE bank_id LIKE '{_BANK_PREFIX}%'") + + +async def _insert_bank(pool, bank_id: str): + await pool.execute( + "INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING", + bank_id, + bank_id, + ) + + +async def _insert_op(pool, bank_id: str, op_id: uuid.UUID | None = None) -> uuid.UUID: + op_id = op_id or uuid.uuid4() + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status) + VALUES ($1, $2, 'consolidation', 'processing') + """, + op_id, + bank_id, + ) + return op_id + + +# --------------------------------------------------------------------------- +# CASCADE DELETE tests +# --------------------------------------------------------------------------- + + +class TestCascadeDeleteOnBankDeletion: + @pytest.mark.asyncio + async def test_bank_deletion_cascades_to_async_operations(self, pool): + bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}" + await _insert_bank(pool, bank_id) + op_id = await _insert_op(pool, bank_id) + + # Verify op exists + row = await pool.fetchrow("SELECT operation_id FROM async_operations WHERE operation_id = $1", op_id) + assert row is not None + + # Delete the bank — should cascade to async_operations + await pool.execute("DELETE FROM banks WHERE bank_id = $1", bank_id) + + row = await pool.fetchrow("SELECT operation_id FROM async_operations WHERE operation_id = $1", op_id) + assert row is None, "async_operations row should be deleted by CASCADE" + + @pytest.mark.asyncio + async def test_bank_deletion_cascades_to_webhooks(self, pool): + bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}" + await _insert_bank(pool, bank_id) + webhook_id = uuid.uuid4() + await pool.execute( + """ + INSERT INTO webhooks (id, bank_id, url, event_types) + VALUES ($1, $2, 'https://example.com/hook', '{}') + """, + webhook_id, + bank_id, + ) + + row = await pool.fetchrow("SELECT id FROM webhooks WHERE id = $1", webhook_id) + assert row is not None + + await pool.execute("DELETE FROM banks WHERE bank_id = $1", bank_id) + + row = await pool.fetchrow("SELECT id FROM webhooks WHERE id = $1", webhook_id) + assert row is None, "webhooks row should be deleted by CASCADE" + + +# --------------------------------------------------------------------------- +# _check_op_alive tests +# --------------------------------------------------------------------------- + + +class TestCheckOpAlive: + @pytest.mark.asyncio + async def test_returns_true_when_op_exists(self, memory: MemoryEngine, request_context): + bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}" + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + op_id = uuid.uuid4() + async with memory._pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status) + VALUES ($1, $2, 'consolidation', 'processing') + """, + op_id, + bank_id, + ) + + assert await memory._check_op_alive(str(op_id)) is True + + @pytest.mark.asyncio + async def test_returns_false_when_op_deleted(self, memory: MemoryEngine, request_context): + bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}" + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + op_id = uuid.uuid4() + async with memory._pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status) + VALUES ($1, $2, 'consolidation', 'processing') + """, + op_id, + bank_id, + ) + await conn.execute("DELETE FROM async_operations WHERE operation_id = $1", op_id) + + assert await memory._check_op_alive(str(op_id)) is False + + @pytest.mark.asyncio + async def test_returns_false_after_bank_cascade_delete(self, memory: MemoryEngine, request_context): + bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}" + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + op_id = uuid.uuid4() + async with memory._pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status) + VALUES ($1, $2, 'consolidation', 'processing') + """, + op_id, + bank_id, + ) + + # Delete the bank — cascades to the op row + await memory.delete_bank(bank_id=bank_id, request_context=request_context) + + assert await memory._check_op_alive(str(op_id)) is False + + +# --------------------------------------------------------------------------- +# _mark_operation_completed / _mark_operation_failed graceful no-op +# --------------------------------------------------------------------------- + + +class TestMarkOperationGracefulOnMissingRow: + @pytest.mark.asyncio + async def test_mark_completed_does_not_raise_when_row_missing(self, memory: MemoryEngine): + # Row never existed — should log and return cleanly + missing_id = str(uuid.uuid4()) + await memory._mark_operation_completed(missing_id) # no exception + + @pytest.mark.asyncio + async def test_mark_failed_does_not_raise_when_row_missing(self, memory: MemoryEngine): + missing_id = str(uuid.uuid4()) + await memory._mark_operation_failed(missing_id, "some error", "traceback here") # no exception + + @pytest.mark.asyncio + async def test_mark_completed_and_fire_webhook_does_not_raise_when_row_missing( + self, memory: MemoryEngine + ): + missing_id = str(uuid.uuid4()) + await memory._mark_operation_completed_and_fire_webhook( + operation_id=missing_id, + bank_id="nonexistent-bank", + status="completed", + result=None, + ) # no exception + + +# --------------------------------------------------------------------------- +# Consolidation checkpoint +# --------------------------------------------------------------------------- + + +class TestConsolidationCheckpoint: + @pytest.mark.asyncio + async def test_consolidation_stops_early_when_op_cancelled(self, memory: MemoryEngine, request_context): + """Consolidation returns 'cancelled' status after the first batch if _check_op_alive is False.""" + from hindsight_api.config import _get_raw_config + from hindsight_api.engine.consolidation.consolidator import run_consolidation_job + + config = _get_raw_config() + original = config.enable_observations + config.enable_observations = True + + try: + bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}" + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + # Insert a few unconsolidated memories directly so we control the batch without LLM + async with memory._pool.acquire() as conn: + for i in range(3): + await conn.execute( + """ + INSERT INTO memory_units + (id, bank_id, text, fact_type, created_at, updated_at) + VALUES (gen_random_uuid(), $1, $2, 'experience', NOW(), NOW()) + """, + bank_id, + f"Test memory {i} for cancellation test", + ) + + op_id = str(uuid.uuid4()) + call_count = 0 + + async def _fake_check(operation_id: str) -> bool: + nonlocal call_count + call_count += 1 + # Return False on the very first checkpoint call + return False + + with patch.object(memory, "_check_op_alive", side_effect=_fake_check): + result = await run_consolidation_job( + memory_engine=memory, + bank_id=bank_id, + request_context=request_context, + operation_id=op_id, + ) + + assert result["status"] == "cancelled" + assert call_count >= 1 + finally: + config.enable_observations = original + + +# --------------------------------------------------------------------------- +# Retain checkpoint +# --------------------------------------------------------------------------- + + +class TestRetainCheckpoint: + @pytest.mark.asyncio + async def test_retain_stops_between_sub_batches_when_cancelled( + self, memory: MemoryEngine, request_context + ): + """retain_batch_async returns partial results if _check_op_alive is False between sub-batches.""" + from hindsight_api.config import _get_raw_config + + bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}" + await memory.get_bank_profile(bank_id=bank_id, request_context=request_context) + + # Force sub-batch splitting by temporarily lowering the token threshold + config = _get_raw_config() + original_tokens = config.retain_batch_tokens + # Set threshold very low so each item becomes its own sub-batch + config.retain_batch_tokens = 1 + + try: + op_id = str(uuid.uuid4()) + check_calls = 0 + + async def _fake_check(operation_id: str) -> bool: + nonlocal check_calls + check_calls += 1 + # Cancel after the first sub-batch completes + return check_calls <= 1 + + contents = [ + {"content": f"Memory item {i} about something interesting."} for i in range(4) + ] + + with patch.object(memory, "_check_op_alive", side_effect=_fake_check): + result = await memory.retain_batch_async( + bank_id=bank_id, + contents=contents, + request_context=request_context, + operation_id=op_id, + ) + + # Should have stopped early: fewer results than total items + assert len(result) < len(contents), ( + f"Expected early stop but got {len(result)}/{len(contents)} results" + ) + assert check_calls >= 1 + finally: + config.retain_batch_tokens = original_tokens diff --git a/hindsight-api/tests/test_webhooks.py b/hindsight-api/tests/test_webhooks.py index 707c2e63..874fcbc0 100644 --- a/hindsight-api/tests/test_webhooks.py +++ b/hindsight-api/tests/test_webhooks.py @@ -135,6 +135,15 @@ async def webhook_manager(memory: MemoryEngine) -> WebhookManager: return WebhookManager(pool=memory._pool, global_webhooks=[]) +async def _ensure_bank(pool, bank_id: str) -> None: + """Upsert a minimal bank row so FK constraints on async_operations/webhooks pass.""" + await pool.execute( + "INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING", + bank_id, + bank_id, + ) + + class TestFireEvent: """Integration tests for WebhookManager.fire_event().""" @@ -147,6 +156,7 @@ class TestFireEvent: webhook_id = uuid.uuid4() async with memory._pool.acquire() as conn: + await _ensure_bank(memory._pool, bank_id) await conn.execute( """ INSERT INTO webhooks (id, bank_id, url, secret, event_types, enabled, created_at, updated_at) @@ -195,6 +205,7 @@ class TestFireEvent: ): """fire_event() also queues delivery tasks for global webhooks (not stored in DB).""" bank_id = f"wh-global-{uuid.uuid4().hex[:8]}" + await _ensure_bank(memory._pool, bank_id) global_webhook = WebhookConfig( id="", # No DB row bank_id=None, @@ -245,6 +256,7 @@ class TestFireEvent: bank_id = f"wh-mismatch-{uuid.uuid4().hex[:8]}" webhook_id = uuid.uuid4() + await _ensure_bank(memory._pool, bank_id) async with memory._pool.acquire() as conn: await conn.execute( """ @@ -339,6 +351,7 @@ class TestHandleWebhookDelivery: operation_id = str(uuid.uuid4()) bank_id = f"wh-exec-{uuid.uuid4().hex[:8]}" + await _ensure_bank(memory._pool, bank_id) # Insert a real async_operations row so _mark_operation_completed has something to update async with memory._pool.acquire() as conn: await conn.execute( @@ -715,6 +728,7 @@ class TestRetainCompletedWebhook: bank_id = f"wh-retain-{uuid.uuid4().hex[:8]}" webhook_id = uuid.uuid4() + await _ensure_bank(memory._pool, bank_id) async with memory._pool.acquire() as conn: await conn.execute( """ diff --git a/hindsight-api/tests/test_worker.py b/hindsight-api/tests/test_worker.py index 510b092b..64bafdde 100644 --- a/hindsight-api/tests/test_worker.py +++ b/hindsight-api/tests/test_worker.py @@ -20,6 +20,15 @@ import pytest_asyncio from hindsight_api.engine.task_backend import BrokerTaskBackend, SyncTaskBackend +async def _ensure_bank(pool, bank_id: str) -> None: + """Upsert a minimal bank row so FK on async_operations passes.""" + await pool.execute( + "INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING", + bank_id, + bank_id, + ) + + # Use loadgroup to ensure these tests run in the same worker # since they share database state pytestmark = pytest.mark.xdist_group("worker_tests") @@ -64,6 +73,7 @@ class TestBrokerTaskBackend: # Create an operation record first operation_id = uuid.uuid4() bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) await pool.execute( """ @@ -104,6 +114,7 @@ class TestBrokerTaskBackend: await backend.initialize() bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) task_dict = { "type": "access_count_update", "bank_id": bank_id, @@ -133,6 +144,7 @@ class TestWorkerPoller: # Create some pending tasks bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) for i in range(3): op_id = uuid.uuid4() payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) @@ -182,6 +194,7 @@ class TestWorkerPoller: # Create 10 pending tasks bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) for i in range(10): op_id = uuid.uuid4() payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) @@ -219,6 +232,7 @@ class TestWorkerPoller: bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" op_id = uuid.uuid4() payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id}) + await _ensure_bank(pool, bank_id) await pool.execute( """ INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id) @@ -284,6 +298,7 @@ class TestWorkerPoller: bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" op_id = uuid.uuid4() payload = json.dumps({"type": "consolidation", "operation_id": str(op_id), "bank_id": bank_id}) + await _ensure_bank(pool, bank_id) await pool.execute( """ INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at) @@ -341,6 +356,7 @@ class TestWorkerPoller: bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" op_id = uuid.uuid4() payload = json.dumps({"type": "consolidation", "operation_id": str(op_id), "bank_id": bank_id}) + await _ensure_bank(pool, bank_id) await pool.execute( """ INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at, retry_count) @@ -394,6 +410,7 @@ class TestWorkerPoller: bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" op_id = uuid.uuid4() payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id}) + await _ensure_bank(pool, bank_id) await pool.execute( """ INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id) @@ -452,6 +469,7 @@ class TestWorkerPoller: from hindsight_api.worker import WorkerPoller bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) # Create a processing consolidation for bank processing_op_id = uuid.uuid4() @@ -479,6 +497,7 @@ class TestWorkerPoller: # Create a pending consolidation for different bank (should be claimed) other_bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, other_bank_id) other_op_id = uuid.uuid4() await pool.execute( """ @@ -517,6 +536,7 @@ class TestWorkerPoller: from hindsight_api.worker import WorkerPoller bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) # Create a processing consolidation for bank await pool.execute( @@ -564,6 +584,7 @@ class TestWorkerRecovery: # Create tasks that were being processed by this worker (simulating a crash) bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) worker_id = "crashed-worker" task_ids = [] @@ -608,6 +629,7 @@ class TestWorkerRecovery: from hindsight_api.worker import WorkerPoller bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) # Create tasks for worker-1 (the one that will recover) for i in range(2): @@ -688,6 +710,7 @@ class TestConcurrentWorkers: # Create 10 pending tasks bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) task_ids = [] for i in range(10): op_id = uuid.uuid4() @@ -749,6 +772,7 @@ class TestConcurrentWorkers: # Create tasks - some pending, some already processing bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) # Create 3 pending tasks for i in range(3): @@ -804,6 +828,7 @@ class TestWorkerDecommission: """Test that decommissioning a worker releases all its processing tasks.""" # Create tasks being processed by a worker bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) worker_id = "worker-to-decommission" for i in range(5): @@ -847,6 +872,7 @@ class TestWorkerDecommission: async def test_decommission_does_not_affect_other_workers(self, pool, clean_operations): """Test that decommissioning one worker doesn't affect another worker's tasks.""" bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) # Create tasks for worker-1 for i in range(3): @@ -964,6 +990,7 @@ class TestDynamicTenantDiscovery: # Create pending tasks in public schema bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) for i in range(2): op_id = uuid.uuid4() payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) @@ -1031,6 +1058,7 @@ class TestDynamicTenantDiscovery: # Create a task in public schema bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) op_id = uuid.uuid4() payload = json.dumps({"type": "test_task", "bank_id": bank_id}) await pool.execute( @@ -1090,6 +1118,7 @@ class TestDynamicTenantDiscovery: # Create pending tasks bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) for i in range(3): op_id = uuid.uuid4() payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id}) @@ -1229,6 +1258,7 @@ async def test_worker_fire_and_forget_nonblocking(pool, clean_operations): ) bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) # Submit initial 2 tasks task_ids = [] @@ -1331,6 +1361,7 @@ async def test_worker_slot_limits_enforced(pool, clean_operations): # Submit 10 tasks bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) for i in range(10): op_id = uuid.uuid4() payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id})