Fix non-atomic async operation creation (#619)

* 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

Author: cdbartholomew

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,28 +7414,35 @@ async def _submit_async_operation(
 
         operation_id = uuid.uuid4()
 
-        # Insert operation record into database
+        # 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),
+            "bank_id": bank_id,
+            **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)
-                VALUES ($1, $2, $3, $4, $5)
+                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 {}),
+                json.dumps(result_metadata or {}, default=_json_default),
                 "pending",
+                json.dumps(full_payload, default=_json_default),
             )
 
-        # Build and submit task payload
-        full_payload = {
-            "type": task_type,
-            "operation_id": str(operation_id),
-            "bank_id": bank_id,
-            **task_payload,
-        }
-
+        # 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}")

hindsight-api-slim/tests/test_file_retain.py

@@ -5,6 +5,7 @@
 import asyncio
 import io
 import json
+from datetime import datetime, timezone
 
 import pytest
 from httpx import ASGITransport, AsyncClient
@@ -471,6 +472,77 @@ async def read(self):
     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'."""