atlas-iac/services/hermes/scripts/migrate_api_session_lineage.py

#!/usr/bin/env python3
"""Nest known legacy agent API workers under their originating objective."""

from __future__ import annotations

import os
import re
import sqlite3
import time
from pathlib import Path


STATE_DB = Path("/opt/data/state.db")
LEGACY_CASSANDRA_PARENT = "20260812_042739_d6ee37"
LEGACY_CASSANDRA_WORKERS = {
    "api-d947ebb51d09eb97": "Cassandra handoff existence check",
    "api-bd98c6145593e9e9": "Cassandra handoff heading check",
    "api-cb00bdf2f7cd3889": "Cassandra worktree file check",
    "api-91b5c2de1759380b": "Agent hosted route check",
    "api-0d1d15e79c619f7b": "Cassandra file-tool check",
    "api-5550b3556424fbdc": "Cassandra handoff summary",
    "api-c288b9f3024a91b8": "Cassandra stale-worktree check",
    "api-058e51802b58ee3f": "Cassandra verification worker",
}
LEGACY_ORPHANED_SMOKE_SESSIONS = {
    "e17f2d888689": "Archived agent workspace smoke test",
    "api-cff2ed44d3b0ab61": "Archived operator Codex packing probe",
    "run_efce779230534721b6e0a3115782fae6": "Archived model availability probe",
    "run_aab885acce7e4daca65e935cc684ffc6": "Archived Vault path probe",
    "run_e689e61419a842219c5fdcc789e27c2f": "Archived secret isolation probe",
}
TRIAGE_PARENT = "automated-triage"
TRIAGE_PARENT_TITLE = "Automated triage"
TRIAGE_MESSAGE_PREFIXES = (
    "A static-analysis finding, not a build failure.",
    "Use $triage-titan-test-failures.",
)


def close_stale_api_workers(connection: sqlite3.Connection) -> int:
    """Close parent-linked API runs that survived a prior gateway lifetime."""
    columns = {
        str(row[1])
        for row in connection.execute("PRAGMA table_info(sessions)").fetchall()
    }
    if not {"ended_at", "end_reason"}.issubset(columns):
        return 0
    cursor = connection.execute(
        """
        UPDATE sessions
           SET ended_at = ?,
               end_reason = 'api_run_recovered_stale'
         WHERE source = 'api_server'
           AND parent_session_id IS NOT NULL
           AND ended_at IS NULL
        """,
        (time.time(),),
    )
    return cursor.rowcount


def _triage_title(message: str, session_id: str) -> str:
    """Create a concise label from Ariadne's stable incident contract."""
    match = re.search(r"(?:for|Analyze) incident ([^\s.]+)(?:\.|\s)", message)
    if not match:
        return f"Automated triage run · {session_id[-8:]}"
    parts = match.group(1).split("/")
    label = " · ".join(parts[1:3]) if len(parts) >= 3 else match.group(1)
    return f"Sonar · {label} · {session_id[-8:]}"


def _is_automated_triage(message: str) -> bool:
    """Recognize only the stable Ariadne automation contracts."""
    return message.startswith(TRIAGE_MESSAGE_PREFIXES)


def migrate_triage_group(connection: sqlite3.Connection) -> int:
    """Nest only recognized automated Sonar runs under one durable parent."""
    rows = connection.execute(
        """
        SELECT s.id, s.started_at, m.content
          FROM sessions AS s
          JOIN messages AS m ON m.session_id = s.id
         WHERE s.source = 'api_server'
           AND s.parent_session_id IS NULL
           AND s.archived = 0
           AND m.role = 'user'
         ORDER BY s.started_at
        """
    ).fetchall()
    rows = [row for row in rows if _is_automated_triage(str(row[2] or ""))]
    if not rows:
        return 0
    started_at = min(float(row[1] or time.time()) for row in rows)
    connection.execute(
        """
        INSERT OR IGNORE INTO sessions (id, source, started_at, title, archived)
        VALUES (?, 'api_server', ?, ?, 0)
        """,
        (TRIAGE_PARENT, started_at, TRIAGE_PARENT_TITLE),
    )
    changed = 0
    for session_id, _, message in rows:
        cursor = connection.execute(
            """
            UPDATE sessions
               SET parent_session_id = ?, title = ?
             WHERE id = ?
               AND parent_session_id IS NULL
            """,
            (TRIAGE_PARENT, _triage_title(str(message or ""), session_id), session_id),
        )
        changed += cursor.rowcount
    return changed


def migrate(path: Path = STATE_DB, *, group_triage: bool = False) -> int:
    """Apply idempotent, transcript-preserving lineage corrections."""
    if not path.is_file():
        return 0
    changed = 0
    with sqlite3.connect(path) as connection:
        changed += close_stale_api_workers(connection)
        parent = connection.execute(
            "SELECT id FROM sessions WHERE id = ?", (LEGACY_CASSANDRA_PARENT,)
        ).fetchone()
        if parent:
            for session_id, title in LEGACY_CASSANDRA_WORKERS.items():
                cursor = connection.execute(
                    """
                    UPDATE sessions
                       SET parent_session_id = ?,
                           title = ?
                     WHERE id = ?
                       AND source = 'api_server'
                       AND parent_session_id IS NULL
                    """,
                    (LEGACY_CASSANDRA_PARENT, title, session_id),
                )
                changed += cursor.rowcount
        for session_id, title in LEGACY_ORPHANED_SMOKE_SESSIONS.items():
            cursor = connection.execute(
                """
                UPDATE sessions
                   SET archived = 1,
                       title = ?
                 WHERE id = ?
                   AND source = 'api_server'
                   AND parent_session_id IS NULL
                   AND archived = 0
                """,
                (title, session_id),
            )
            changed += cursor.rowcount
        if group_triage:
            changed += migrate_triage_group(connection)
    return changed


if __name__ == "__main__":
    group_triage = os.environ.get("HERMES_API_DEFAULT_PARENT_SESSION_ID") == TRIAGE_PARENT
    print(f"migrated {migrate(group_triage=group_triage)} legacy API sessions")