api/persistence.py

"""backend/api/persistence.py — Supabase persistence for agent task state (S359).

Design principles:
  - ALL writes are fire-and-forget: never block the SSE hot path.
  - Reads happen ONLY on lazy restore (when in-memory is empty, i.e. backend restarted).
  - Supabase unavailable → silent fallback to in-memory-only (no exception propagation).
  - Event buffer capped at MAX_EVENTS per task to avoid runaway storage.

GAP-P40D-FIX: aggiunto retry (max 2 tentativi, 0.3s sleep) per sb_upsert_task e
sb_append_event — stesse operazioni già protette in linter.py (P40-D pattern).
Il gap: su Railway cold-start le prime scritture Supabase cadono nel vuoto senza retry.
Il pattern è identico a quello in linter.py (lint_and_store) già confermato stabile.

Required Supabase tables (run backend/migrations/s359_task_persistence.sql once):
  agent_tasks        — task metadata (task_id, goal, status, max_steps, created_at, updated_at)
  agent_task_events  — SSE event buffer (task_id, event_index, event_data, created_at)
"""
import asyncio, time, json
from .state import safe_json_dumps as _sjd  # B11-FIX: surrogate-safe drop-in
from typing import Optional, Any

import logging
_logger = logging.getLogger("api.persistence")

MAX_EVENTS = 500  # max SSE frames persisted per task

_MAX_RETRY   = 2    # GAP-P40D-FIX: tentativi massimi per write Supabase
_RETRY_SLEEP = 0.3  # GAP-P40D-FIX: sleep tra tentativi (secondi)

# P0: per-run locks serialize compatible envelope updates in one worker.
_ENGINEERING_LOCKS: dict[str, asyncio.Lock] = {}
_ENGINEERING_LOCK_LAST_USED: dict[str, float] = {}
_ENGINEERING_LOCK_MAX = 256


def _engineering_lock(task_id: str) -> asyncio.Lock:
    lock = _ENGINEERING_LOCKS.get(task_id)
    if lock is None:
        lock = asyncio.Lock()
        _ENGINEERING_LOCKS[task_id] = lock
    _ENGINEERING_LOCK_LAST_USED[task_id] = time.monotonic()
    if len(_ENGINEERING_LOCKS) > _ENGINEERING_LOCK_MAX:
        for stale_id, _ in sorted(_ENGINEERING_LOCK_LAST_USED.items(), key=lambda item: item[1]):
            stale_lock = _ENGINEERING_LOCKS.get(stale_id)
            if stale_lock is not None and not stale_lock.locked() and stale_id != task_id:
                _ENGINEERING_LOCKS.pop(stale_id, None)
                _ENGINEERING_LOCK_LAST_USED.pop(stale_id, None)
                break
    return lock


# ── Write helpers (fire-and-forget, never raise) ───────────────────────────────

async def sb_upsert_task(
    task_id: str, goal: str, status: str, max_steps: int,
    context: list, created_at: int,
) -> None:
    """Persist or update task metadata. Silent if Supabase not configured.
    GAP-P40D-FIX: retry fino a _MAX_RETRY tentativi su errore transitorio."""
    from .state import _sb
    if not _sb:
        return
    now = int(time.time() * 1000)
    payload = {
        'task_id':    task_id,
        'goal':       goal[:1000],
        'status':     status,
        'max_steps':  max_steps,
        'context':    _sjd(context)[:8000],
        'created_at': created_at,
        'updated_at': now,
    }
    for _attempt in range(_MAX_RETRY):
        try:
            await asyncio.to_thread(
                lambda: _sb.table('agent_tasks').upsert(payload).execute()
            )
            return  # successo
        except Exception as e:
            if _attempt < _MAX_RETRY - 1:
                await asyncio.sleep(_RETRY_SLEEP)
            else:
                _logger.warning('[persist] upsert_task %s failed after %d attempts: %s',
                                task_id, _MAX_RETRY, e)


async def sb_update_status(task_id: str, status: str) -> None:
    """Update only the status column of an existing task."""
    from .state import _sb
    if not _sb:
        return
    now = int(time.time() * 1000)
    try:
        await asyncio.to_thread(
            lambda: _sb.table('agent_tasks')
                .update({'status': status, 'updated_at': now})
                .eq('task_id', task_id)
                .execute()
        )
    except Exception as e:
        _logger.warning('[persist] update_status %s→%s: %s', task_id, status, e)


async def sb_append_event(task_id: str, event_index: int, event_data: str) -> None:
    """Append one SSE event string to the persistent buffer. Capped at MAX_EVENTS.
    GAP-P40D-FIX: retry fino a _MAX_RETRY tentativi su errore transitorio."""
    from .state import _sb
    if not _sb or event_index > MAX_EVENTS:
        return
    now = int(time.time() * 1000)
    payload = {
        'task_id':     task_id,
        'event_index': event_index,
        'event_data':  event_data,
        'created_at':  now,
    }
    for _attempt in range(_MAX_RETRY):
        try:
            await asyncio.to_thread(
                lambda: _sb.table('agent_task_events').upsert(payload).execute()
            )
            return  # successo
        except Exception as e:
            if _attempt < _MAX_RETRY - 1:
                await asyncio.sleep(_RETRY_SLEEP)
            else:
                _logger.warning('[persist] append_event %s#%d failed after %d attempts: %s',
                                task_id, event_index, _MAX_RETRY, e)


async def sb_restore_task(task_id: str) -> Optional[dict]:
    """Restore task metadata from Supabase on backend restart."""
    from .state import _sb
    if not _sb:
        return None
    try:
        res = await asyncio.to_thread(
            lambda: _sb.table('agent_tasks').select('*').eq('task_id', task_id).limit(1).execute()
        )
        return res.data[0] if res.data else None
    except Exception as e:
        _logger.warning('[persist] restore_task %s: %s', task_id, e)
        return None


async def sb_get_events(task_id: str) -> list[str]:
    """Retrieve all persisted SSE events for a task, ordered by index."""
    from .state import _sb
    if not _sb:
        return []
    try:
        res = await asyncio.to_thread(
            lambda: _sb.table('agent_task_events')
                .select('event_index, event_data')
                .eq('task_id', task_id)
                .order('event_index')
                .limit(MAX_EVENTS)
                .execute()
        )
        return [r['event_data'] for r in (res.data or [])]
    except Exception as e:
        _logger.warning('[persist] get_events %s: %s', task_id, e)
        return []


async def sb_delete_task_events(task_id: str) -> None:
    """Delete all events for a task (post-replay cleanup)."""
    from .state import _sb
    if not _sb:
        return
    try:
        await asyncio.to_thread(
            lambda: _sb.table('agent_task_events').delete().eq('task_id', task_id).execute()
        )
    except Exception as e:
        _logger.debug('[persist] delete_events %s: %s', task_id, e)


async def sb_list_tasks(limit: int = 50) -> list[dict]:
    """List recent tasks from Supabase."""
    from .state import _sb
    if not _sb:
        return []
    try:
        res = await asyncio.to_thread(
            lambda: _sb.table('agent_tasks')
                .select('task_id, goal, status, created_at, updated_at')
                .order('created_at', desc=True)
                .limit(limit)
                .execute()
        )
        return res.data or []
    except Exception as e:
        _logger.warning('[persist] list_tasks: %s', e)
        return []


# ── Checkpoint helpers (S359: task state snapshots) ───────────────────────────

async def sb_save_checkpoint(task_id: str, step: int, checkpoint_data: dict) -> None:
    """Save a legacy checkpoint while preserving a valid EngineeringState envelope."""
    from .state import _sb
    if not _sb:
        return
    now = int(time.time() * 1000)
    try:
        payload = dict(checkpoint_data) if isinstance(checkpoint_data, dict) else {}
        # A legacy save must not erase the shadow/canary envelope written by the
        # adapter. Read/merge under the same per-task lock used by its writer.
        lock = _engineering_lock(task_id)
        async with lock:
            current = await sb_get_checkpoint(task_id)
            current_engineering = current.get('engineering_state') if isinstance(current, dict) else None
            if isinstance(current_engineering, dict) and 'engineering_state' not in payload:
                payload['engineering_state'] = current_engineering
            serialized = _sjd(payload)
            if len(serialized) > 16000:
                _logger.debug('[persist] save_checkpoint %s#%d skipped: payload exceeds size limit', task_id, step)
                return
            await asyncio.to_thread(
                lambda: _sb.table('agent_tasks')
                    .update({'checkpoint': serialized, 'updated_at': now})
                    .eq('task_id', task_id)
                    .execute()
            )
    except Exception as e:
        _logger.debug('[persist] save_checkpoint %s#%d: %s', task_id, step, e)


_ENGINEERING_DEBOUNCE_CACHE: dict[str, dict[str, Any]] = {}
_ENGINEERING_LAST_FLUSH_TS: dict[str, float] = {}
DEBOUNCE_INTERVAL_SEC = 2.0

async def sb_save_engineering_state(task_id: str, envelope: dict, force: bool = False) -> None:
    """Merge a validated EngineeringState envelope with debouncing and monotone revision check."""
    from .state import _sb
    if not _sb or not task_id:
        return
    try:
        from agents.engineering_state import EngineeringState
        validated = EngineeringState.from_snapshot(envelope).snapshot()
    except Exception as exc:
        _logger.debug('[persist] engineering state rejected: %s', type(exc).__name__)
        return

    lock = _engineering_lock(task_id)
    async with lock:
        current = await sb_get_checkpoint(task_id)
        current = current if isinstance(current, dict) else {}
        current_engineering = current.get('engineering_state')
        try:
            current_revision = int(current_engineering.get('revision', -1)) if isinstance(current_engineering, dict) else -1
        except (TypeError, ValueError):
            current_revision = -1
        incoming_revision = int(validated.get('revision', -1))
        if current_revision > incoming_revision:
            _logger.debug('[persist] engineering state conflict %s: remote revision %d > %d', task_id, current_revision, incoming_revision)
            return
    now_t = time.time()
    _ENGINEERING_DEBOUNCE_CACHE[task_id] = validated
    if not force and task_id in _ENGINEERING_LAST_FLUSH_TS:
        if now_t - _ENGINEERING_LAST_FLUSH_TS[task_id] < DEBOUNCE_INTERVAL_SEC:
            return

    _ENGINEERING_LAST_FLUSH_TS[task_id] = now_t
    to_flush = _ENGINEERING_DEBOUNCE_CACHE.get(task_id, validated)

    async with lock:
        current = await sb_get_checkpoint(task_id)
        current = current if isinstance(current, dict) else {}
        current_engineering = current.get('engineering_state')
        try:
            current_revision = int(current_engineering.get('revision', -1)) if isinstance(current_engineering, dict) else -1
        except (TypeError, ValueError):
            current_revision = -1
        incoming_revision = int(to_flush.get('revision', -1))
        if current_revision > incoming_revision and not force:
            return
        merged = dict(current)
        merged['engineering_state'] = to_flush
        serialized = _sjd(merged)
        if len(serialized) > 16000:
            return
        now = int(time.time() * 1000)
        try:
            await asyncio.to_thread(
                lambda: _sb.table('agent_tasks')
                    .update({'checkpoint': serialized, 'updated_at': now})
                    .eq('task_id', task_id)
                    .execute()
            )
        except Exception as exc:
            _logger.debug('[persist] save_engineering_state %s: %s', task_id, exc)


async def sb_get_checkpoint(task_id: str) -> Optional[dict]:
    """Retrieve latest checkpoint for a task."""
    from .state import _sb
    if not _sb:
        return None
    try:
        res = await asyncio.to_thread(
            lambda: _sb.table('agent_tasks')
                .select('checkpoint')
                .eq('task_id', task_id)
                .limit(1)
                .execute()
        )
        if res.data and res.data[0].get('checkpoint'):
            raw = res.data[0]['checkpoint']
            return json.loads(raw) if isinstance(raw, str) else raw
        return None
    except Exception as e:
        _logger.debug('[persist] get_checkpoint %s: %s', task_id, e)
        return None


# ── Handoff helpers (BG-4: cross-session context handoff) ─────────────────────

async def sb_restore_handoff_context(session_id: str) -> Optional[dict]:
    """Restore cross-session handoff context."""
    from .state import _sb
    if not _sb:
        return None
    try:
        res = await asyncio.to_thread(
            lambda: _sb.table('handoff_contexts')
                .select('*').limit(500)  # BUGFIX: LIMIT 500 — senza limit OOM su stato persistente grande
                .eq('session_id', session_id)
                .limit(1)
                .execute()
        )
        return res.data[0] if res.data else None
    except Exception as e:
        _logger.debug('[persist] restore_handoff %s: %s', session_id, e)
        return None


async def sb_upsert_handoff(session_id: str, context: dict) -> None:
    """Save handoff context for cross-session resume."""
    from .state import _sb
    if not _sb:
        return
    now = int(time.time() * 1000)
    try:
        await asyncio.to_thread(
            lambda: _sb.table('handoff_contexts').upsert({
                'session_id': session_id,
                'context':    _sjd(context)[:16000],
                'updated_at': now,
            }, on_conflict='session_id').execute()
        )
    except Exception as e:
        _logger.debug('[persist] upsert_handoff %s: %s', session_id, e)


async def sb_delete_handoff(session_id: str) -> None:
    """Remove handoff context after successful resume."""
    from .state import _sb
    if not _sb:
        return
    try:
        await asyncio.to_thread(
            lambda: _sb.table('handoff_contexts').delete().eq('session_id', session_id).execute()
        )
    except Exception as e:
        _logger.debug('[persist] delete_handoff %s: %s', session_id, e)