"""Repository for reconciliation sweeps over stuck durability rows.""" from __future__ import annotations from sqlalchemy import Connection, text from application.storage.db.base_repository import row_to_dict class ReconciliationRepository: """Sweeps and terminal writes for the reconciler beat task.""" def __init__(self, conn: Connection) -> None: self._conn = conn def find_and_lock_stuck_messages( self, *, age_minutes: int = 5, limit: int = 100, ) -> list[dict]: """Lock stuck pending/streaming messages skipping live resumes. Staleness rides on the **later of** ``cm.timestamp`` (creation) and ``message_metadata.last_heartbeat_at`` (route heartbeat). An in-flight stream that re-stamps the heartbeat each minute stays out of the sweep; reconciler-side writes deliberately don't touch either column so the per-row attempts counter advances across ticks. Liveness exemption covers both ``pending`` (paused waiting for resume) and ``resuming`` (actively executing) ``pending_tool_state`` rows so a paused message survives until the PT row's own TTL retires it. """ result = self._conn.execute( text( """ SELECT cm.id, cm.conversation_id, cm.user_id, cm.timestamp, cm.message_metadata FROM conversation_messages cm WHERE cm.status IN ('pending', 'streaming') AND cm.timestamp < now() - make_interval(mins => :age) AND COALESCE( (cm.message_metadata->>'last_heartbeat_at')::timestamptz, cm.timestamp ) < now() - make_interval(mins => :age) AND NOT EXISTS ( SELECT 1 FROM pending_tool_state pts WHERE pts.conversation_id = cm.conversation_id AND pts.user_id = cm.user_id AND ( (pts.status = 'pending' AND pts.expires_at > now()) OR (pts.status = 'resuming' AND pts.resumed_at > now() - interval '10 minutes') ) ) ORDER BY cm.timestamp ASC LIMIT :limit FOR UPDATE OF cm SKIP LOCKED """ ), {"age": age_minutes, "limit": limit}, ) return [row_to_dict(r) for r in result.fetchall()] def find_and_lock_proposed_tool_calls( self, *, age_minutes: int = 5, limit: int = 100, ) -> list[dict]: """Lock tool_call_attempts that never advanced past ``proposed``.""" result = self._conn.execute( text( """ SELECT call_id, message_id, tool_id, tool_name, action_name, arguments, attempted_at, updated_at FROM tool_call_attempts WHERE status = 'proposed' AND attempted_at < now() - make_interval(mins => :age) ORDER BY attempted_at ASC LIMIT :limit FOR UPDATE SKIP LOCKED """ ), {"age": age_minutes, "limit": limit}, ) return [row_to_dict(r) for r in result.fetchall()] def find_and_lock_executed_tool_calls( self, *, age_minutes: int = 15, limit: int = 100, ) -> list[dict]: """Lock tool_call_attempts stuck in ``executed`` past confirm window.""" result = self._conn.execute( text( """ SELECT call_id, message_id, tool_id, tool_name, action_name, arguments, result, attempted_at, updated_at FROM tool_call_attempts WHERE status = 'executed' AND updated_at < now() - make_interval(mins => :age) ORDER BY updated_at ASC LIMIT :limit FOR UPDATE SKIP LOCKED """ ), {"age": age_minutes, "limit": limit}, ) return [row_to_dict(r) for r in result.fetchall()] def find_and_lock_stalled_ingests( self, *, age_minutes: int = 30, limit: int = 100, ) -> list[dict]: """Lock still-active ingest checkpoints with a silent heartbeat. The ``status = 'active'`` filter skips rows already escalated to ``'stalled'``, so a dead ingest is alerted once, not every tick. """ result = self._conn.execute( text( """ SELECT icp.source_id, icp.total_chunks, icp.embedded_chunks, icp.last_index, icp.last_updated, s.user_id, s.name AS source_name FROM ingest_chunk_progress icp LEFT JOIN sources s ON s.id = icp.source_id WHERE icp.last_updated < now() - make_interval(mins => :age) AND icp.embedded_chunks < icp.total_chunks AND icp.status = 'active' ORDER BY icp.last_updated ASC LIMIT :limit FOR UPDATE OF icp SKIP LOCKED """ ), {"age": age_minutes, "limit": limit}, ) return [row_to_dict(r) for r in result.fetchall()] def mark_ingest_stalled(self, source_id: str) -> bool: """Escalate a stalled checkpoint to terminal ``status='stalled'``. Drops the row out of the sweep so the reconciler alerts once; ``init_progress`` flips it back to ``'active'`` on reingest. """ result = self._conn.execute( text( "UPDATE ingest_chunk_progress SET status = 'stalled' " "WHERE source_id = CAST(:sid AS uuid)" ), {"sid": str(source_id)}, ) return result.rowcount > 0 def increment_message_reconcile_attempts(self, message_id: str) -> int: """Bump ``message_metadata.reconcile_attempts`` and return the new count.""" result = self._conn.execute( text( """ UPDATE conversation_messages SET message_metadata = jsonb_set( COALESCE(message_metadata, '{}'::jsonb), '{reconcile_attempts}', to_jsonb( COALESCE( (message_metadata->>'reconcile_attempts')::int, 0 ) + 1 ) ) WHERE id = CAST(:message_id AS uuid) RETURNING (message_metadata->>'reconcile_attempts')::int AS new_count """ ), {"message_id": message_id}, ) row = result.fetchone() return int(row[0]) if row is not None else 0 def mark_message_failed(self, message_id: str, *, error: str) -> bool: """Flip a message to ``status='failed'`` and stash ``error`` in metadata.""" result = self._conn.execute( text( """ UPDATE conversation_messages SET status = 'failed', message_metadata = jsonb_set( COALESCE(message_metadata, '{}'::jsonb), '{error}', to_jsonb(CAST(:error AS text)) ) WHERE id = CAST(:message_id AS uuid) """ ), {"message_id": message_id, "error": error}, ) return result.rowcount > 0 def mark_tool_call_failed(self, call_id: str, *, error: str) -> bool: """Flip a tool_call_attempts row to ``failed`` with ``error``.""" result = self._conn.execute( text( "UPDATE tool_call_attempts SET status = 'failed', " "error = :error WHERE call_id = :call_id" ), {"call_id": call_id, "error": error}, ) return result.rowcount > 0 def find_stuck_idempotency_pending( self, *, max_attempts: int, lease_grace_seconds: int = 60, limit: int = 100, ) -> list[dict]: """Lock ``task_dedup`` rows abandoned past the lease + retry budget. A row is "stuck" when: - ``status='pending'`` (lease was claimed but never finalised) - ``lease_expires_at`` is past by at least ``lease_grace_seconds`` (the heartbeat thread is gone — the lease isn't going to come back) - ``attempt_count >= max_attempts`` (the poison-loop guard should already have escalated this; if it hasn't, the wrapper died before getting there) These rows would otherwise sit in ``pending`` until the 24 h TTL aged them out, blocking same-key retries via ``_lookup_completed`` returning None for the whole window. """ result = self._conn.execute( text( """ SELECT idempotency_key, task_name, task_id, attempt_count, lease_owner_id, lease_expires_at, created_at FROM task_dedup WHERE status = 'pending' AND lease_expires_at IS NOT NULL AND lease_expires_at < now() - make_interval(secs => :grace) AND attempt_count >= :max_attempts ORDER BY created_at ASC LIMIT :limit FOR UPDATE SKIP LOCKED """ ), { "max_attempts": int(max_attempts), "grace": int(lease_grace_seconds), "limit": int(limit), }, ) return [row_to_dict(r) for r in result.fetchall()] def mark_idempotency_pending_failed( self, key: str, *, error: str, ) -> bool: """Promote a stuck pending ``task_dedup`` row to ``failed``.""" from application.storage.db.serialization import PGNativeJSONEncoder import json result = self._conn.execute( text( """ UPDATE task_dedup SET status = 'failed', result_json = CAST(:result AS jsonb), lease_owner_id = NULL, lease_expires_at = NULL WHERE idempotency_key = :key AND status = 'pending' """ ), { "key": key, "result": json.dumps( { "success": False, "error": error, "reconciled": True, }, cls=PGNativeJSONEncoder, ), }, ) return result.rowcount > 0