Compare commits

...
Author SHA1 Message Date
Nicolò Boschi 7519ecac95 fixes 2026-01-30 10:57:30 +01:00
Nicolò Boschi ab0f8cec33 fix: deadlock in worker polling 2026-01-30 10:34:24 +01:00
Nicolò Boschi 19676d69a7 fix: deadlock in worker polling 2026-01-30 10:29:30 +01:00
10 changed files with 451 additions and 133 deletions
+2 -1
View File
@@ -1404,9 +1404,10 @@ def create_app(
worker_id=worker_id,
executor=memory.execute_task,
poll_interval_ms=config.worker_poll_interval_ms,
batch_size=config.worker_batch_size,
max_retries=config.worker_max_retries,
tenant_extension=getattr(memory, "_tenant_extension", None),
max_slots=config.worker_max_slots,
consolidation_max_slots=config.worker_consolidation_max_slots,
)
poller_task = asyncio.create_task(poller.run())
logging.info(f"Worker poller started (worker_id={worker_id})")
+10 -4
View File
@@ -143,8 +143,9 @@ ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID"
ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
ENV_WORKER_BATCH_SIZE = "HINDSIGHT_API_WORKER_BATCH_SIZE"
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS"
# Reflect agent settings
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
@@ -229,8 +230,9 @@ DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
DEFAULT_WORKER_ID = None # Will use hostname if not specified
DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
DEFAULT_WORKER_BATCH_SIZE = 10 # Tasks to claim per poll cycle
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker
# Reflect agent settings
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
@@ -419,8 +421,9 @@ class HindsightConfig:
worker_id: str | None
worker_poll_interval_ms: int
worker_max_retries: int
worker_batch_size: int
worker_http_port: int
worker_max_slots: int
worker_consolidation_max_slots: int
# Reflect agent settings
reflect_max_iterations: int
@@ -582,8 +585,11 @@ class HindsightConfig:
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))),
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
worker_batch_size=int(os.getenv(ENV_WORKER_BATCH_SIZE, str(DEFAULT_WORKER_BATCH_SIZE))),
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))),
worker_consolidation_max_slots=int(
os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS))
),
# Reflect agent settings
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
)
+2 -1
View File
@@ -253,8 +253,9 @@ def main():
worker_id=config.worker_id,
worker_poll_interval_ms=config.worker_poll_interval_ms,
worker_max_retries=config.worker_max_retries,
worker_batch_size=config.worker_batch_size,
worker_http_port=config.worker_http_port,
worker_max_slots=config.worker_max_slots,
worker_consolidation_max_slots=config.worker_consolidation_max_slots,
reflect_max_iterations=config.reflect_max_iterations,
mental_model_refresh_concurrency=config.mental_model_refresh_concurrency,
)
+4 -8
View File
@@ -124,12 +124,6 @@ def main():
default=config.worker_poll_interval_ms,
help=f"Poll interval in milliseconds (default: {config.worker_poll_interval_ms}, env: HINDSIGHT_API_WORKER_POLL_INTERVAL_MS)",
)
parser.add_argument(
"--batch-size",
type=int,
default=config.worker_batch_size,
help=f"Tasks to claim per poll (default: {config.worker_batch_size}, env: HINDSIGHT_API_WORKER_BATCH_SIZE)",
)
parser.add_argument(
"--max-retries",
type=int,
@@ -168,8 +162,9 @@ def main():
print(f"Starting Hindsight Worker: {args.worker_id}")
print(f" Poll interval: {args.poll_interval}ms")
print(f" Batch size: {args.batch_size}")
print(f" Max retries: {args.max_retries}")
print(f" Max slots: {config.worker_max_slots}")
print(f" Consolidation max slots: {config.worker_consolidation_max_slots}")
print(f" HTTP server: {args.http_host}:{args.http_port}")
print()
@@ -213,9 +208,10 @@ def main():
worker_id=args.worker_id,
executor=memory.execute_task,
poll_interval_ms=args.poll_interval,
batch_size=args.batch_size,
max_retries=args.max_retries,
tenant_extension=tenant_extension,
max_slots=config.worker_max_slots,
consolidation_max_slots=config.worker_consolidation_max_slots,
)
# Create the HTTP app for metrics/health
+205 -87
View File
@@ -57,10 +57,11 @@ class WorkerPoller:
worker_id: str,
executor: Callable[[dict[str, Any]], Awaitable[None]],
poll_interval_ms: int = 500,
batch_size: int = 10,
max_retries: int = 3,
schema: str | None = None,
tenant_extension: "TenantExtension | None" = None,
max_slots: int = 10,
consolidation_max_slots: int = 2,
):
"""
Initialize the worker poller.
@@ -70,28 +71,32 @@ class WorkerPoller:
worker_id: Unique identifier for this worker
executor: Async function to execute tasks (typically MemoryEngine.execute_task)
poll_interval_ms: Interval between polls when no tasks found (milliseconds)
batch_size: Maximum number of tasks to claim per poll cycle
max_retries: Maximum retry attempts before marking task as failed
schema: Database schema for single-tenant support (ignored if tenant_extension is set)
tenant_extension: Extension for dynamic multi-tenant discovery. If set, list_tenants()
is called on each poll cycle to discover schemas dynamically.
max_slots: Maximum concurrent tasks per worker
consolidation_max_slots: Maximum concurrent consolidation tasks per worker
"""
self._pool = pool
self._worker_id = worker_id
self._executor = executor
self._poll_interval_ms = poll_interval_ms
self._batch_size = batch_size
self._max_retries = max_retries
self._schema = schema
self._tenant_extension = tenant_extension
self._max_slots = max_slots
self._consolidation_max_slots = consolidation_max_slots
self._shutdown = asyncio.Event()
self._current_tasks: set[asyncio.Task] = set()
self._in_flight_count = 0
self._in_flight_lock = asyncio.Lock()
self._last_progress_log = 0.0
self._tasks_completed_since_log = 0
# Track active tasks locally: operation_id -> (op_type, bank_id, schema)
self._active_tasks: dict[str, tuple[str, str, str | None]] = {}
# Track active tasks locally: operation_id -> (op_type, bank_id, schema, asyncio.Task)
self._active_tasks: dict[str, tuple[str, str, str | None, asyncio.Task]] = {}
# Track in-flight tasks by operation type
self._in_flight_by_type: dict[str, int] = {}
async def _get_schemas(self) -> list[str | None]:
"""Get list of schemas to poll. Returns [None] for public schema."""
@@ -102,67 +107,114 @@ class WorkerPoller:
# Single schema mode
return [self._schema]
async def _get_available_slots(self) -> tuple[int, int]:
"""
Calculate available slots for claiming tasks.
Returns:
(total_available, consolidation_available) tuple
"""
async with self._in_flight_lock:
total_in_flight = self._in_flight_count
consolidation_in_flight = self._in_flight_by_type.get("consolidation", 0)
total_available = max(0, self._max_slots - total_in_flight)
consolidation_available = max(0, self._consolidation_max_slots - consolidation_in_flight)
return total_available, consolidation_available
async def wait_for_active_tasks(self, timeout: float = 10.0) -> bool:
"""
Wait for all active background tasks to complete (test helper).
This is a test-only utility that allows tests to synchronize with
fire-and-forget background tasks without using sleep().
Args:
timeout: Maximum time to wait in seconds
Returns:
True if all tasks completed, False if timeout was reached
"""
start_time = asyncio.get_event_loop().time()
while True:
async with self._in_flight_lock:
if self._in_flight_count == 0:
return True
elapsed = asyncio.get_event_loop().time() - start_time
if elapsed >= timeout:
return False
# Short sleep to avoid busy-waiting
await asyncio.sleep(0.01)
async def claim_batch(self) -> list[ClaimedTask]:
"""
Claim up to batch_size pending tasks atomically across all tenant schemas.
Claim pending tasks atomically across all tenant schemas,
respecting slot limits (total and consolidation).
Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers.
For consolidation tasks specifically, skips pending tasks if there's already
a processing consolidation for the same bank (to avoid duplicate work).
If tenant_extension is configured, dynamically discovers schemas on each call.
Returns:
List of ClaimedTask objects containing operation_id, task_dict, and schema
"""
# Calculate available slots
total_available, consolidation_available = await self._get_available_slots()
if total_available <= 0:
return []
schemas = await self._get_schemas()
all_tasks: list[ClaimedTask] = []
remaining_batch = self._batch_size
remaining_total = total_available
remaining_consolidation = consolidation_available
for schema in schemas:
if remaining_batch <= 0:
if remaining_total <= 0:
break
tasks = await self._claim_batch_for_schema(schema, remaining_batch)
tasks = await self._claim_batch_for_schema(schema, remaining_total, remaining_consolidation)
# Update remaining slots based on what was claimed
for task in tasks:
op_type = task.task_dict.get("operation_type", "unknown")
if op_type == "consolidation":
remaining_consolidation -= 1
all_tasks.extend(tasks)
remaining_batch -= len(tasks)
remaining_total -= len(tasks)
return all_tasks
async def _claim_batch_for_schema(self, schema: str | None, limit: int) -> list[ClaimedTask]:
"""Claim tasks from a specific schema."""
async def _claim_batch_for_schema(
self, schema: str | None, limit: int, consolidation_limit: int
) -> list[ClaimedTask]:
"""Claim tasks from a specific schema respecting slot limits."""
try:
return await self._claim_batch_for_schema_inner(schema, limit)
return await self._claim_batch_for_schema_inner(schema, limit, consolidation_limit)
except Exception as e:
logger.warning(f"Worker {self._worker_id} failed to claim tasks for schema {schema or 'public'}: {e}")
return []
async def _claim_batch_for_schema_inner(self, schema: str | None, limit: int) -> list[ClaimedTask]:
"""Inner implementation for claiming tasks from a specific schema."""
async def _claim_batch_for_schema_inner(
self, schema: str | None, limit: int, consolidation_limit: int
) -> list[ClaimedTask]:
"""Inner implementation for claiming tasks from a specific schema with slot limits."""
table = fq_table("async_operations", schema)
async with self._pool.acquire() as conn:
async with conn.transaction():
# Select and lock pending tasks
# For consolidation: skip if same bank already has one processing
rows = await conn.fetch(
# Strategy: Claim non-consolidation tasks first, then consolidation up to limit
# 1. Claim non-consolidation tasks (up to limit)
non_consolidation_rows = await conn.fetch(
f"""
SELECT operation_id, task_payload
FROM {table} AS pending
WHERE status = 'pending' AND task_payload IS NOT NULL
AND (
-- Non-consolidation tasks: always claimable
operation_type != 'consolidation'
OR
-- Consolidation: only if no other consolidation processing for same bank
NOT EXISTS (
SELECT 1 FROM {table} AS processing
WHERE processing.bank_id = pending.bank_id
AND processing.operation_type = 'consolidation'
AND processing.status = 'processing'
)
)
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type != 'consolidation'
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
@@ -170,11 +222,39 @@ class WorkerPoller:
limit,
)
if not rows:
claimed_count = len(non_consolidation_rows)
remaining_limit = limit - claimed_count
# 2. Claim consolidation tasks (up to consolidation_limit and remaining_limit)
consolidation_rows = []
if consolidation_limit > 0 and remaining_limit > 0:
consolidation_rows = await conn.fetch(
f"""
SELECT operation_id, task_payload
FROM {table} AS pending
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = 'consolidation'
AND NOT EXISTS (
SELECT 1 FROM {table} AS processing
WHERE processing.bank_id = pending.bank_id
AND processing.operation_type = 'consolidation'
AND processing.status = 'processing'
)
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
min(consolidation_limit, remaining_limit),
)
all_rows = non_consolidation_rows + consolidation_rows
if not all_rows:
return []
# Claim the tasks by updating status and worker_id
operation_ids = [row["operation_id"] for row in rows]
operation_ids = [row["operation_id"] for row in all_rows]
await conn.execute(
f"""
UPDATE {table}
@@ -192,7 +272,7 @@ class WorkerPoller:
task_dict=json.loads(row["task_payload"]),
schema=schema,
)
for row in rows
for row in all_rows
]
async def _mark_completed(self, operation_id: str, schema: str | None):
@@ -258,18 +338,43 @@ class WorkerPoller:
logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})")
async def execute_task(self, task: ClaimedTask):
"""Execute a single task and update its status."""
"""Execute a single task as a background job (fire-and-forget)."""
task_type = task.task_dict.get("type", "unknown")
operation_type = task.task_dict.get("operation_type", "unknown")
bank_id = task.task_dict.get("bank_id", "unknown")
# Create background task
bg_task = asyncio.create_task(self._execute_task_inner(task))
# Track this task as active
async with self._in_flight_lock:
self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema)
self._active_tasks[task.operation_id] = (task_type, bank_id, task.schema, bg_task)
self._in_flight_count += 1
self._in_flight_by_type[operation_type] = self._in_flight_by_type.get(operation_type, 0) + 1
# Add cleanup callback
bg_task.add_done_callback(lambda _: asyncio.create_task(self._cleanup_task(task.operation_id, operation_type)))
async def _cleanup_task(self, operation_id: str, operation_type: str):
"""Remove task from tracking after completion."""
async with self._in_flight_lock:
if operation_id in self._active_tasks:
self._active_tasks.pop(operation_id, None)
self._in_flight_count -= 1
count = self._in_flight_by_type.get(operation_type, 0)
if count > 0:
self._in_flight_by_type[operation_type] = count - 1
if self._in_flight_by_type[operation_type] == 0:
del self._in_flight_by_type[operation_type]
async def _execute_task_inner(self, task: ClaimedTask):
"""Inner task execution with error handling."""
task_type = task.task_dict.get("type", "unknown")
bank_id = task.task_dict.get("bank_id", "unknown")
try:
schema_info = f", schema={task.schema}" if task.schema else ""
logger.debug(f"Executing task {task.operation_id} (type={task_type}, bank={bank_id}{schema_info})")
# Pass schema to executor so it can set the correct context
if task.schema:
task.task_dict["_schema"] = task.schema
await self._executor(task.task_dict)
@@ -279,10 +384,6 @@ class WorkerPoller:
error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
logger.error(f"Task {task.operation_id} failed: {e}")
await self._retry_or_fail(task.operation_id, error_msg, task.schema)
finally:
# Remove from active tasks
async with self._in_flight_lock:
self._active_tasks.pop(task.operation_id, None)
async def recover_own_tasks(self) -> int:
"""
@@ -325,59 +426,59 @@ class WorkerPoller:
async def run(self):
"""
Main polling loop.
Main polling loop with fire-and-forget task execution.
Continuously polls for pending tasks, claims them, and executes them
until shutdown is signaled.
If tenant_extension is configured, dynamically discovers schemas on each poll.
Continuously polls for pending tasks, spawns them as background tasks,
and immediately continues polling (up to slot limits).
"""
# Recover any tasks from a previous crash before starting
await self.recover_own_tasks()
logger.info(f"Worker {self._worker_id} starting polling loop")
logger.info(
f"Worker {self._worker_id} starting polling loop "
f"(max_slots={self._max_slots}, consolidation_max_slots={self._consolidation_max_slots})"
)
while not self._shutdown.is_set():
try:
# Claim a batch of tasks (across all tenant schemas if configured)
# Claim a batch of tasks (respecting slot limits)
tasks = await self.claim_batch()
if tasks:
# Log batch info
task_types: dict[str, int] = {}
schemas_seen: set[str | None] = set()
consolidation_count = 0
for task in tasks:
t = task.task_dict.get("type", "unknown")
op_type = task.task_dict.get("operation_type", "unknown")
task_types[t] = task_types.get(t, 0) + 1
schemas_seen.add(task.schema)
if op_type == "consolidation":
consolidation_count += 1
types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items())
schemas_str = ", ".join(s or "public" for s in schemas_seen)
logger.info(
f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str} (schemas: {schemas_str})"
f"Worker {self._worker_id} claimed {len(tasks)} tasks "
f"({consolidation_count} consolidation): {types_str} (schemas: {schemas_str})"
)
# Track in-flight tasks
async with self._in_flight_lock:
self._in_flight_count += len(tasks)
# Spawn tasks as background jobs (fire-and-forget)
for task in tasks:
await self.execute_task(task)
# Execute tasks concurrently
try:
await asyncio.gather(
*[self.execute_task(task) for task in tasks],
return_exceptions=True,
)
finally:
async with self._in_flight_lock:
self._in_flight_count -= len(tasks)
else:
# No tasks found, wait before polling again
try:
await asyncio.wait_for(
self._shutdown.wait(),
timeout=self._poll_interval_ms / 1000,
)
except asyncio.TimeoutError:
pass # Normal timeout, continue polling
# Continue immediately to claim more tasks (if slots available)
continue
# No tasks claimed (either no pending tasks or slots full)
# Wait before polling again
try:
await asyncio.wait_for(
self._shutdown.wait(),
timeout=self._poll_interval_ms / 1000,
)
except asyncio.TimeoutError:
pass # Normal timeout, continue polling
# Log progress stats periodically
await self._log_progress_if_due()
@@ -408,15 +509,27 @@ class WorkerPoller:
while asyncio.get_event_loop().time() - start_time < timeout:
async with self._in_flight_lock:
in_flight = self._in_flight_count
active_task_objects = [task_info[3] for task_info in self._active_tasks.values()]
if in_flight == 0:
logger.info(f"Worker {self._worker_id} graceful shutdown complete")
return
logger.info(f"Worker {self._worker_id} waiting for {in_flight} in-flight tasks")
await asyncio.sleep(0.5)
logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s")
# Wait for at least one task to complete
if active_task_objects:
done, _ = await asyncio.wait(active_task_objects, timeout=0.5, return_when=asyncio.FIRST_COMPLETED)
else:
await asyncio.sleep(0.5)
logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s, cancelling remaining tasks")
# Cancel remaining tasks
async with self._in_flight_lock:
for operation_id, (_, _, _, bg_task) in list(self._active_tasks.items()):
if not bg_task.done():
bg_task.cancel()
async def _log_progress_if_due(self):
"""Log progress stats every PROGRESS_LOG_INTERVAL seconds."""
@@ -427,14 +540,19 @@ class WorkerPoller:
self._last_progress_log = now
try:
# Get local active tasks (this worker only)
# Get local active tasks
async with self._in_flight_lock:
in_flight = self._in_flight_count
active_tasks = dict(self._active_tasks) # Copy to avoid holding lock
in_flight_by_type = dict(self._in_flight_by_type)
active_tasks = dict(self._active_tasks)
# Build local processing breakdown grouped by (op_type, bank_id)
consolidation_count = in_flight_by_type.get("consolidation", 0)
available_slots = self._max_slots - in_flight
available_consolidation_slots = self._consolidation_max_slots - consolidation_count
# Build local processing breakdown
task_groups: dict[tuple[str, str], int] = {}
for op_type, bank_id, _ in active_tasks.values():
for op_type, bank_id, _, _ in active_tasks.values():
key = (op_type, bank_id)
task_groups[key] = task_groups.get(key, 0) + 1
@@ -443,7 +561,7 @@ class WorkerPoller:
if len(processing_info) > 10:
processing_str += f" +{len(processing_info) - 10} more"
# Get global stats from DB across all schemas
# Get global stats from DB
schemas = await self._get_schemas()
global_pending = 0
all_worker_counts: dict[str, int] = {}
@@ -455,7 +573,6 @@ class WorkerPoller:
row = await conn.fetchrow(f"SELECT COUNT(*) as count FROM {table} WHERE status = 'pending'")
global_pending += row["count"] if row else 0
# Get processing breakdown by worker
worker_rows = await conn.fetch(
f"""
SELECT worker_id, COUNT(*) as count
@@ -468,7 +585,6 @@ class WorkerPoller:
wid = wr["worker_id"] or "unknown"
all_worker_counts[wid] = all_worker_counts.get(wid, 0) + wr["count"]
# Format other workers' processing counts
other_workers = []
for wid, cnt in all_worker_counts.items():
if wid != self._worker_id:
@@ -477,7 +593,9 @@ class WorkerPoller:
schemas_str = ", ".join(s or "public" for s in schemas)
logger.info(
f"[WORKER_STATS] worker={self._worker_id} in_flight={in_flight} | "
f"[WORKER_STATS] worker={self._worker_id} "
f"slots={in_flight}/{self._max_slots} (consolidation={consolidation_count}/{self._consolidation_max_slots}) | "
f"available={available_slots} (consolidation={available_consolidation_slots}) | "
f"global: pending={global_pending} (schemas: {schemas_str}) | "
f"others: {others_str} | "
f"my_active: {processing_str}"
+11 -11
View File
@@ -346,11 +346,11 @@ class TestConsolidationIntegration:
or when one directly updates another (e.g., location change).
Given:
- "Nicolò lives in Italy"
- "Nicolò moved to the US recently" (updates the living location)
- "Alex lives in Italy"
- "Alex moved to the US recently" (updates the living location)
The second fact should UPDATE the first, not create a separate observation.
But unrelated facts like "Nicolò works at Vectorize" should stay separate.
But unrelated facts like "Alex works at Vectorize" should stay separate.
"""
bank_id = f"test-consolidation-merge-{uuid.uuid4().hex[:8]}"
@@ -360,14 +360,14 @@ class TestConsolidationIntegration:
# Retain a memory about living location
await memory.retain_async(
bank_id=bank_id,
content="Nicolò lives in Italy.",
content="Alex lives in Italy.",
request_context=request_context,
)
# Retain an unrelated memory (different topic - should NOT merge)
await memory.retain_async(
bank_id=bank_id,
content="Nicolò works at Vectorize as an engineer.",
content="Alex works at Vectorize as an engineer.",
request_context=request_context,
)
@@ -384,7 +384,7 @@ class TestConsolidationIntegration:
# Add a memory that UPDATES the living location (should merge with first)
await memory.retain_async(
bank_id=bank_id,
content="Nicolò recently moved to the United States.",
content="Alex recently moved to the United States.",
request_context=request_context,
)
@@ -485,9 +485,9 @@ class TestConsolidationIntegration:
they should be merged into ONE observation that captures the change.
Example:
- "Nicolò loves pizza"
- "Nicolò hates pizza"
→ Should become: "Nicolò used to love pizza but now hates it" (or similar)
- "Alex loves pizza"
- "Alex hates pizza"
→ Should become: "Alex used to love pizza but now hates it" (or similar)
"""
bank_id = f"test-consolidation-contradict-{uuid.uuid4().hex[:8]}"
@@ -497,7 +497,7 @@ class TestConsolidationIntegration:
# Add initial fact
await memory.retain_async(
bank_id=bank_id,
content="Nicolò loves pizza.",
content="Alex loves pizza.",
request_context=request_context,
)
@@ -515,7 +515,7 @@ class TestConsolidationIntegration:
# Add contradicting fact (same person, same topic, opposite sentiment)
await memory.retain_async(
bank_id=bank_id,
content="Nicolò hates pizza.",
content="Alex hates pizza.",
request_context=request_context,
)
+212 -15
View File
@@ -156,7 +156,6 @@ class TestWorkerPoller:
pool=pool,
worker_id="test-worker-1",
executor=mock_executor,
batch_size=10,
)
claimed = await poller.claim_batch()
@@ -177,8 +176,8 @@ class TestWorkerPoller:
assert row["worker_id"] == "test-worker-1"
@pytest.mark.asyncio
async def test_claim_batch_respects_batch_size(self, pool, clean_operations):
"""Test that claim_batch respects the batch_size limit."""
async def test_claim_batch_respects_max_slots(self, pool, clean_operations):
"""Test that claim_batch respects the max_slots limit."""
from hindsight_api.worker import WorkerPoller
# Create 10 pending tasks
@@ -196,12 +195,11 @@ class TestWorkerPoller:
payload,
)
# Claim with batch_size=3
poller = WorkerPoller(
pool=pool,
worker_id="test-worker-1",
executor=lambda x: None,
batch_size=3,
max_slots=3, # Limit to 3 concurrent tasks
)
claimed = await poller.claim_batch()
@@ -238,11 +236,14 @@ class TestWorkerPoller:
executor=mock_executor,
)
# Execute the task
# Execute the task (fire-and-forget)
task_dict = json.loads(payload)
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
await poller.execute_task(claimed_task)
# Wait for background task to complete
completed = await poller.wait_for_active_tasks(timeout=5.0)
assert completed, "Task did not complete within timeout"
assert len(executed) == 1
# Verify task is marked as completed
@@ -283,11 +284,15 @@ class TestWorkerPoller:
max_retries=3,
)
# Execute (should fail and retry)
# Execute (should fail and retry) - fire-and-forget
task_dict = json.loads(payload)
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
await poller.execute_task(claimed_task)
# Wait for background task to complete
completed = await poller.wait_for_active_tasks(timeout=5.0)
assert completed, "Task did not complete within timeout"
# Verify task is back to pending with incremented retry_count
row = await pool.fetchrow(
"SELECT status, retry_count, worker_id FROM async_operations WHERE operation_id = $1",
@@ -327,11 +332,15 @@ class TestWorkerPoller:
max_retries=3,
)
# Execute (should fail permanently)
# Execute (should fail permanently) - fire-and-forget
task_dict = json.loads(payload)
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
await poller.execute_task(claimed_task)
# Wait for background task to complete
completed = await poller.wait_for_active_tasks(timeout=5.0)
assert completed, "Task did not complete within timeout"
# Verify task is marked as failed
row = await pool.fetchrow(
"SELECT status, error_message FROM async_operations WHERE operation_id = $1",
@@ -388,7 +397,6 @@ class TestWorkerPoller:
pool=pool,
worker_id="test-worker-1",
executor=lambda x: None,
batch_size=10,
)
claimed = await poller.claim_batch()
@@ -440,7 +448,6 @@ class TestWorkerPoller:
pool=pool,
worker_id="test-worker-1",
executor=lambda x: None,
batch_size=10,
)
claimed = await poller.claim_batch()
@@ -607,7 +614,6 @@ class TestConcurrentWorkers:
pool=pool,
worker_id=worker_id,
executor=lambda x: None,
batch_size=5, # Each worker tries to claim 5
)
claimed = await poller.claim_batch()
workers_claimed[worker_id] = [task.operation_id for task in claimed]
@@ -680,7 +686,6 @@ class TestConcurrentWorkers:
pool=pool,
worker_id="new-worker",
executor=lambda x: None,
batch_size=10,
)
claimed = await poller.claim_batch()
@@ -879,7 +884,6 @@ class TestDynamicTenantDiscovery:
pool=pool,
worker_id="test-worker-1",
executor=lambda x: None,
batch_size=10,
tenant_extension=mock_extension,
)
@@ -946,7 +950,6 @@ class TestDynamicTenantDiscovery:
pool=pool,
worker_id="test-worker-1",
executor=lambda x: None,
batch_size=10,
tenant_extension=dynamic_extension,
)
@@ -1008,7 +1011,6 @@ class TestDynamicTenantDiscovery:
pool=pool,
worker_id="test-worker-1",
executor=lambda x: None,
batch_size=10,
)
claimed = await poller.claim_batch()
@@ -1017,3 +1019,198 @@ class TestDynamicTenantDiscovery:
# All tasks should have schema=None (public)
for task in claimed:
assert task.schema is None
async def test_worker_fire_and_forget_nonblocking(pool, clean_operations):
"""
Test that worker continues polling while tasks run (fire-and-forget pattern).
This test verifies the FIX: With the old blocking behavior, the worker would
wait for all tasks in a batch to complete before claiming more. This test
would FAIL with the old code because tasks 3-4 wouldn't be claimed until
tasks 1-2 complete. With fire-and-forget, tasks 3-4 are claimed immediately.
"""
from hindsight_api.worker.poller import WorkerPoller
task_started = {} # operation_id -> Event (set when task starts)
task_canfinish = {} # operation_id -> Event (wait before finishing)
async def blocking_executor(task_dict: dict):
op_id = task_dict["operation_id"]
# Signal that this task has started
started = asyncio.Event()
task_started[op_id] = started
started.set()
# Block until we're told to finish
finish = asyncio.Event()
task_canfinish[op_id] = finish
await finish.wait()
poller = WorkerPoller(
pool=pool,
worker_id="test-worker",
executor=blocking_executor,
poll_interval_ms=50, # Fast polling
max_slots=10,
consolidation_max_slots=2,
)
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
# Submit initial 2 tasks
task_ids = []
for i in range(2):
op_id = uuid.uuid4()
task_ids.append(str(op_id))
payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
""",
op_id,
bank_id,
payload,
)
poll_task = asyncio.create_task(poller.run())
try:
# Wait for first 2 tasks to start executing (but not finish)
for i in range(100): # Try for up to 1 second
if len(task_started) >= 2:
break
await asyncio.sleep(0.01)
assert len(task_started) == 2, f"Expected 2 tasks started, got {len(task_started)}"
# Verify tasks are in_flight
async with poller._in_flight_lock:
assert poller._in_flight_count == 2
# NOW submit 2 more tasks WHILE the first 2 are still running
for i in range(2):
op_id = uuid.uuid4()
task_ids.append(str(op_id))
payload = json.dumps({"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
""",
op_id,
bank_id,
payload,
)
# KEY ASSERTION: Worker should claim tasks 3-4 WITHOUT waiting for 1-2 to finish
# This would FAIL with the old blocking behavior
for i in range(100): # Try for up to 1 second
if len(task_started) >= 4:
break
await asyncio.sleep(0.01)
assert len(task_started) == 4, (
f"Fire-and-forget FAILED: Expected 4 tasks started, got {len(task_started)}. "
"This means the worker blocked waiting for the first batch to complete."
)
# Verify all 4 tasks are in-flight
async with poller._in_flight_lock:
assert poller._in_flight_count == 4
# Clean up: allow all tasks to finish
for event in task_canfinish.values():
event.set()
finally:
# Ensure cleanup
for event in task_canfinish.values():
event.set()
await poller.shutdown_graceful(timeout=2.0)
try:
await asyncio.wait_for(poll_task, timeout=1.0)
except asyncio.CancelledError:
pass
async def test_worker_slot_limits_enforced(pool, clean_operations):
"""Test that worker respects max_slots and won't exceed the limit."""
from hindsight_api.worker.poller import WorkerPoller
tasks_started = set()
task_events = {}
async def controlled_executor(task_dict: dict):
op_id = task_dict["operation_id"]
tasks_started.add(op_id)
event = asyncio.Event()
task_events[op_id] = event
await event.wait()
poller = WorkerPoller(
pool=pool,
worker_id="test-worker",
executor=controlled_executor,
poll_interval_ms=50,
max_slots=3, # Only allow 3 concurrent tasks
consolidation_max_slots=1,
)
# Submit 10 tasks
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
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})
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
""",
op_id,
bank_id,
payload,
)
poll_task = asyncio.create_task(poller.run())
try:
# Wait for slots to fill
for i in range(100):
if len(tasks_started) >= 3:
break
await asyncio.sleep(0.01)
# Should have claimed exactly 3 tasks (slot limit)
assert len(tasks_started) == 3
# Wait to ensure no additional tasks are claimed
for i in range(30):
await asyncio.sleep(0.01)
assert len(tasks_started) == 3, "Worker exceeded slot limit!"
# Release tasks one by one and verify remaining are claimed
completed = 0
while completed < 10 and len(tasks_started) < 10:
# Release the next batch
events_to_release = list(task_events.values())[completed:completed+3]
for event in events_to_release:
event.set()
completed += len(events_to_release)
# Wait for new tasks to be claimed
for i in range(100):
if len(tasks_started) >= min(completed + 3, 10):
break
await asyncio.sleep(0.01)
assert len(tasks_started) == 10
finally:
for event in task_events.values():
event.set()
await poller.shutdown_graceful(timeout=2.0)
try:
await asyncio.wait_for(poll_task, timeout=1.0)
except asyncio.CancelledError:
pass
@@ -28,7 +28,7 @@ export async function GET(request: NextRequest) {
console.error("Graph API error:", response.error);
return NextResponse.json(
{ error: response.error || "Failed to fetch graph data" },
{ status: 500 },
{ status: 500 }
);
}
@@ -33,10 +33,8 @@ export const hindsightClient = new HindsightClient({
export const lowLevelClient = createClient(
createConfig({
baseUrl: DATAPLANE_URL,
headers: DATAPLANE_API_KEY
? { Authorization: `Bearer ${DATAPLANE_API_KEY}` }
: undefined,
}),
headers: DATAPLANE_API_KEY ? { Authorization: `Bearer ${DATAPLANE_API_KEY}` } : undefined,
})
);
/**
@@ -504,9 +504,10 @@ Configuration for background task processing. By default, the API processes task
| `HINDSIGHT_API_WORKER_ENABLED` | Enable internal worker in API process | `true` |
| `HINDSIGHT_API_WORKER_ID` | Unique worker identifier | hostname |
| `HINDSIGHT_API_WORKER_POLL_INTERVAL_MS` | Database polling interval in milliseconds | `500` |
| `HINDSIGHT_API_WORKER_BATCH_SIZE` | Tasks to claim per poll cycle | `10` |
| `HINDSIGHT_API_WORKER_MAX_RETRIES` | Max retries before marking task failed | `3` |
| `HINDSIGHT_API_WORKER_HTTP_PORT` | HTTP port for worker metrics/health (worker CLI only) | `8889` |
| `HINDSIGHT_API_WORKER_MAX_SLOTS` | Maximum concurrent tasks per worker | `10` |
| `HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS` | Maximum concurrent consolidation tasks per worker | `2` |
### Performance Optimization