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
This commit is contained in:
Chris Bartholomew 2026-03-19 11:38:04 -04:00 committed by GitHub
parent 439424559e
commit 94cf89b570
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 101 additions and 15 deletions

View file

@ -67,6 +67,13 @@ def fq_table(table_name: str) -> str:
return f"{get_current_schema()}.{table_name}" 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) # Tables that must be schema-qualified (for runtime validation)
_PROTECTED_TABLES = frozenset( _PROTECTED_TABLES = frozenset(
[ [
@ -7407,21 +7414,10 @@ class MemoryEngine(MemoryEngineInterface):
operation_id = uuid.uuid4() operation_id = uuid.uuid4()
# Insert operation record into database # Build full payload before INSERT so task_payload is included atomically.
async with acquire_with_retry(pool) as conn: # Previously the INSERT omitted task_payload and a separate submit_task call
await conn.execute( # did an UPDATE — a crash between the two left a null-payload row that the
f""" # worker's claim query (task_payload IS NOT NULL) could never pick up.
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
full_payload = { full_payload = {
"type": task_type, "type": task_type,
"operation_id": str(operation_id), "operation_id": str(operation_id),
@ -7429,6 +7425,24 @@ class MemoryEngine(MemoryEngineInterface):
**task_payload, **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) await self._task_backend.submit_task(full_payload)
logger.info(f"{operation_type} task queued for bank_id={bank_id}, operation_id={operation_id}") logger.info(f"{operation_type} task queued for bank_id={bank_id}, operation_id={operation_id}")

View file

@ -5,6 +5,7 @@ End-to-end tests for file retain (upload, convert, retain) functionality.
import asyncio import asyncio
import io import io
import json import json
from datetime import datetime, timezone
import pytest import pytest
from httpx import ASGITransport, AsyncClient 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 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 @pytest.mark.asyncio
async def test_file_conversion_failure_sets_status_to_failed(memory_no_llm_verify, sample_txt_content): 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'.""" """Test that when file conversion fails, the operation status is set to 'failed' not 'completed'."""