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}"
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}")

View file

@ -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'."""