From 94cf89b5709d271ab56fdfdd2b6c8dc72e463221 Mon Sep 17 00:00:00 2001 From: Chris Bartholomew Date: Thu, 19 Mar 2026 11:38:04 -0400 Subject: [PATCH] Fix non-atomic async operation creation (#619) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix non-atomic async operation creation in _submit_async_operation Previously the method performed two separate database round-trips: 1. INSERT into async_operations with no task_payload (null) 2. submit_task → UPDATE to set task_payload A process crash or network error between steps 1 and 2 left a row with task_payload IS NULL permanently. The worker's claim query requires task_payload IS NOT NULL, so these orphaned rows could never be picked up and the queue appeared degraded indefinitely. Fix: build full_payload before the INSERT and include task_payload in the same INSERT statement, making operation creation atomic. submit_task is still called afterwards — for SyncTaskBackend it executes the task immediately (unchanged behaviour); for BrokerTaskBackend it becomes an idempotent UPDATE (payload already set) kept for symmetry. * Preserve datetime payloads in atomic async insert --- .../hindsight_api/engine/memory_engine.py | 44 ++++++++---- hindsight-api-slim/tests/test_file_retain.py | 72 +++++++++++++++++++ 2 files changed, 101 insertions(+), 15 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index eeb71e5a..32f96853 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -67,6 +67,13 @@ def fq_table(table_name: str) -> str: return f"{get_current_schema()}.{table_name}" +def _json_default(obj: Any) -> str: + """JSON serializer for types commonly carried through async task payloads.""" + if isinstance(obj, datetime): + return obj.isoformat() + raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable") + + # Tables that must be schema-qualified (for runtime validation) _PROTECTED_TABLES = frozenset( [ @@ -7407,21 +7414,10 @@ class MemoryEngine(MemoryEngineInterface): operation_id = uuid.uuid4() - # Insert operation record into database - async with acquire_with_retry(pool) as conn: - await conn.execute( - f""" - INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status) - VALUES ($1, $2, $3, $4, $5) - """, - operation_id, - bank_id, - operation_type, - json.dumps(result_metadata or {}), - "pending", - ) - - # Build and submit task payload + # Build full payload before INSERT so task_payload is included atomically. + # Previously the INSERT omitted task_payload and a separate submit_task call + # did an UPDATE — a crash between the two left a null-payload row that the + # worker's claim query (task_payload IS NOT NULL) could never pick up. full_payload = { "type": task_type, "operation_id": str(operation_id), @@ -7429,6 +7425,24 @@ class MemoryEngine(MemoryEngineInterface): **task_payload, } + # Insert operation record with task_payload in a single atomic statement + async with acquire_with_retry(pool) as conn: + await conn.execute( + f""" + INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status, task_payload) + VALUES ($1, $2, $3, $4, $5, $6::jsonb) + """, + operation_id, + bank_id, + operation_type, + json.dumps(result_metadata or {}, default=_json_default), + "pending", + json.dumps(full_payload, default=_json_default), + ) + + # For SyncTaskBackend: executes the task immediately. + # For BrokerTaskBackend: does an idempotent UPDATE (payload already set above), + # kept for symmetry and to support any future notification mechanisms. await self._task_backend.submit_task(full_payload) logger.info(f"{operation_type} task queued for bank_id={bank_id}, operation_id={operation_id}") diff --git a/hindsight-api-slim/tests/test_file_retain.py b/hindsight-api-slim/tests/test_file_retain.py index cc9488e0..d4ba80d3 100644 --- a/hindsight-api-slim/tests/test_file_retain.py +++ b/hindsight-api-slim/tests/test_file_retain.py @@ -5,6 +5,7 @@ End-to-end tests for file retain (upload, convert, retain) functionality. import asyncio import io import json +from datetime import datetime, timezone import pytest from httpx import ASGITransport, AsyncClient @@ -471,6 +472,77 @@ async def test_file_conversion_creates_separate_retain_operation(memory_no_llm_v assert len(doc["original_text"]) > 0 +@pytest.mark.asyncio +async def test_async_file_retain_serializes_datetime_timestamp(memory_no_llm_verify, sample_txt_content): + """Async file retain should accept Python datetimes in task payloads.""" + from hindsight_api.engine.parsers.base import FileParser + from hindsight_api.models import RequestContext + + bank_id = f"test_file_timestamp_bank_{datetime.now(timezone.utc).timestamp()}" + timestamp = datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc) + + context = RequestContext(internal=True) + await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context) + + class MockFile: + def __init__(self, content, filename, content_type): + self.content = content + self.filename = filename + self.content_type = content_type + + async def read(self): + return self.content + + class TimestampParser(FileParser): + async def convert(self, file_data: bytes, filename: str) -> str: + return file_data.decode("utf-8") + + def supports(self, filename: str, content_type: str | None = None) -> bool: + return filename.endswith(".txt") + + def name(self) -> str: + return "timestamp_parser" + + memory_no_llm_verify._parser_registry.register(TimestampParser()) + + mock_file = MockFile(sample_txt_content, "timestamped.txt", "text/plain") + + result = await memory_no_llm_verify.submit_async_file_retain( + bank_id=bank_id, + file_items=[ + { + "file": mock_file, + "document_id": "timestamped_doc", + "context": "timestamp test", + "metadata": {}, + "tags": [], + "timestamp": timestamp, + "parser": ["timestamp_parser"], + } + ], + document_tags=None, + request_context=context, + ) + + operation_id = result["operation_ids"][0] + pool = await memory_no_llm_verify._get_pool() + from hindsight_api.engine.memory_engine import get_current_schema + + async with pool.acquire() as conn: + row = await conn.fetchrow( + f""" + SELECT status, task_payload->>'timestamp' AS timestamp + FROM {get_current_schema()}.async_operations + WHERE operation_id = $1 + """, + operation_id, + ) + + assert row is not None + assert row["status"] == "completed" + assert row["timestamp"] == "2024-01-15T10:30:00+00:00" + + @pytest.mark.asyncio async def test_file_conversion_failure_sets_status_to_failed(memory_no_llm_verify, sample_txt_content): """Test that when file conversion fails, the operation status is set to 'failed' not 'completed'."""