fix: cancel in-flight async ops when bank is deleted (#545)
* fix: cancel async ops on bank delete via CASCADE FK + heartbeat checkpoints - Add migration e5f6g7h8i9j0: FK ON DELETE CASCADE from async_operations and webhooks to banks, so deleting a bank auto-removes all its ops/webhooks - Add _check_op_alive() helper: returns False if op row was deleted (cascade) - Add consolidation checkpoint: after each LLM batch commit, abort early if op was deleted mid-run (returns status='cancelled') - Add retain checkpoint: between sub-batches, abort early if op was deleted - _mark_operation_completed/failed/completed_and_fire_webhook: gracefully handle missing row (UPDATE 0) with log instead of silent error - Thread operation_id into run_consolidation_job() for checkpoint access - Fix y0t1u2v3w4x5 and a1b2c3d4e5f6 migrations: add IF NOT EXISTS to prevent failure on idempotent re-runs - Add 10 tests covering cascade delete, _check_op_alive, graceful mark methods, consolidation checkpoint, and retain checkpoint * refactor: use RETURNING + fetchrow instead of execute + string comparison * fix: add bank upsert before async_operations FK inserts and update tests - memory_engine.py: upsert bank in submit_async_retain before async_operations INSERT - http.py: upsert bank in api_create_webhook before webhooks INSERT - test_worker.py, test_async_batch_retain.py, test_webhooks.py: add _ensure_bank helper calls before direct async_operations/webhooks inserts to satisfy FK constraint * fix: mock bank_utils.get_bank_profile in unit test with mocked pool
This commit is contained in:
Nicolò Boschi
GitHub
parent
220851e6f4
commit
0560f6260d
11 changed files with 515 additions and 13 deletions
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -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)
|
||||
""")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
311
hindsight-api/tests/test_op_cancellation.py
Normal file
311
hindsight-api/tests/test_op_cancellation.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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(
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
Loading…
Reference in a new issue