hermes-webui/tests/test_issue2628_cli_sessions_perf.py

"""Regression coverage for capped CLI/agent session sidebar scans (#2628)."""

import pathlib
import sqlite3
import time

import pytest

import api.agent_sessions as agent_sessions

_REAL_SQLITE_CONNECT = sqlite3.connect
REPO_ROOT = pathlib.Path(__file__).resolve().parents[1]


def _make_state_db(path, *, sessions=80, messages_per_session=3, create_messages_index=True, source="cli", session_source="cli"):
    conn = sqlite3.connect(str(path))
    conn.executescript(
        """
        CREATE TABLE sessions (
            id TEXT PRIMARY KEY,
            source TEXT,
            session_source TEXT,
            title TEXT,
            model TEXT,
            started_at REAL NOT NULL,
            message_count INTEGER DEFAULT 0,
            parent_session_id TEXT,
            ended_at REAL,
            end_reason TEXT
        );
        CREATE INDEX idx_sessions_started ON sessions(started_at);
        CREATE TABLE messages (
            id TEXT PRIMARY KEY,
            session_id TEXT,
            role TEXT,
            content TEXT,
            timestamp REAL
        );
        """
    )
    base = time.time() - sessions
    for i in range(sessions):
        sid = f"cli_perf_{i:04d}"
        started = base + i
        conn.execute(
            """
            INSERT INTO sessions
            (id, source, session_source, title, model, started_at, message_count, parent_session_id, ended_at, end_reason)
            VALUES (?, ?, ?, ?, 'openai/gpt-5', ?, ?, NULL, NULL, NULL)
            """,
            (sid, source, session_source, sid, started, messages_per_session),
        )
        for j in range(messages_per_session):
            conn.execute(
                "INSERT INTO messages (id, session_id, role, content, timestamp) VALUES (?, ?, ?, 'hello', ?)",
                (f"msg_{i:04d}_{j:02d}", sid, "user" if j == 0 else "assistant", started + j / 10),
            )
    if create_messages_index:
        conn.execute("CREATE INDEX idx_messages_session ON messages(session_id, timestamp)")
    conn.commit()
    conn.close()


def _newest_first_reference_ids(db_path, *, include_sources=None, exclude_sources=("webui",)):
    where_clauses = ["s.source IS NOT NULL"]
    params = []
    if include_sources:
        placeholders = ", ".join("?" for _ in include_sources)
        where_clauses.append(f"s.source IN ({placeholders})")
        params.extend(include_sources)
    if exclude_sources:
        placeholders = ", ".join("?" for _ in exclude_sources)
        where_clauses.append(f"s.source NOT IN ({placeholders})")
        params.extend(exclude_sources)
    conn = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True)
    conn.row_factory = sqlite3.Row
    try:
        rows = conn.execute(
            f"""
            SELECT s.id
            FROM sessions s
            LEFT JOIN messages m ON m.session_id = s.id
            WHERE {' AND '.join(where_clauses)}
            GROUP BY s.id
            ORDER BY COALESCE(MAX(m.timestamp), s.started_at) DESC
            """,
            params,
        ).fetchall()
        return [row["id"] for row in rows]
    finally:
        conn.close()


def _execute_candidate_ordering_baseline_sql(
    db_path,
    candidate_limit,
    *,
    budget_ops,
    interval=1,
    include_sources=None,
    exclude_sources=("webui",),
):
    where_clauses = ["s.source IS NOT NULL"]
    params = []
    if include_sources:
        placeholders = ", ".join("?" for _ in include_sources)
        where_clauses.append(f"s.source IN ({placeholders})")
        params.extend(include_sources)
    if exclude_sources:
        placeholders = ", ".join("?" for _ in exclude_sources)
        where_clauses.append(f"s.source NOT IN ({placeholders})")
        params.extend(exclude_sources)

    def _on_progress():
        nonlocal steps
        steps += 1
        return 1 if steps > budget_ops else 0

    steps = 0
    conn = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True)
    conn.set_progress_handler(_on_progress, interval)
    try:
        return conn.execute(
            f"""
            WITH candidates AS (
                SELECT s.id
                FROM sessions s
                WHERE {' AND '.join(where_clauses)}
                ORDER BY COALESCE(
                    (SELECT MAX(mx.timestamp) FROM messages mx WHERE mx.session_id = s.id),
                    s.started_at
                ) DESC,
                s.started_at DESC
                LIMIT ?
            )
            SELECT s.id
            FROM sessions s
            JOIN candidates c ON c.id = s.id
            LEFT JOIN messages m ON m.session_id = s.id
            GROUP BY s.id
            ORDER BY COALESCE(MAX(m.timestamp), s.started_at) DESC
            """,
            [*params, candidate_limit],
        ).fetchall()
    finally:
        conn.close()


def _make_connect_with_progress_budget(*, budget_ops, interval=1):
    def _connect(database, *_, **__):
        steps = {"count": 0}
        database_uri = str(database)
        if database_uri.startswith("file:"):
            target_uri = database_uri
        else:
            target_uri = f"file:{database_uri}?mode=ro&immutable=1"

        def _on_progress():
            steps["count"] += 1
            return 1 if steps["count"] > budget_ops else 0

        conn = _REAL_SQLITE_CONNECT(
            target_uri,
            uri=True,
        )
        conn.set_progress_handler(_on_progress, interval)
        return conn

    return _connect


def _make_connect_with_progress_counter(*, interval=1):
    steps = {"count": 0}

    def _connect(database, *_, **__):
        database_uri = str(database)
        if database_uri.startswith("file:"):
            target_uri = database_uri
        else:
            target_uri = f"file:{database_uri}?mode=ro&immutable=1"

        def _on_progress():
            steps["count"] += 1
            return 0

        conn = _REAL_SQLITE_CONNECT(
            target_uri,
            uri=True,
        )
        conn.set_progress_handler(_on_progress, interval)
        return conn

    return _connect, steps


def test_importable_agent_rows_push_sidebar_limit_into_sql(tmp_path):
    """A capped sidebar scan should not aggregate the entire state.db first."""
    db = tmp_path / "state.db"
    _make_state_db(db, sessions=120, messages_per_session=5)

    rows = agent_sessions.read_importable_agent_session_rows(db, limit=20, exclude_sources=("webui",))

    assert len(rows) == 20
    assert [row["id"] for row in rows][:3] == ["cli_perf_0119", "cli_perf_0118", "cli_perf_0117"]
    assert {row["actual_message_count"] for row in rows} == {5}

    src = (REPO_ROOT / "api" / "agent_sessions.py").read_text()
    assert "WITH candidates AS" in src
    assert "JOIN candidates c ON c.id = s.id" in src
    assert "latest_messages AS" in src
    assert "LEFT JOIN latest_messages lm ON lm.session_id = s.id" in src
    assert 'included == ("cron",)' in src
    assert "not messages_index_present" in src
    assert "PRAGMA index_list(messages)" in src
    assert "CREATE INDEX IF NOT EXISTS idx_messages_session" in src
    assert "_CRON_PREAGGREGATE_CANDIDATE_ORDER_MIN_MESSAGES" not in src
    assert "MAX(mx.timestamp) FROM messages mx WHERE mx.session_id = s.id" in src
    assert "candidate_limit = max(result_limit * 8, result_limit)" in src


def test_importable_agent_rows_candidate_ordering_stays_under_progress_budget(tmp_path, monkeypatch):
    """Cron-only missing-index scans should fail under the old shape budget, then pass after pre-aggregation."""
    db = tmp_path / "state.db"
    _make_state_db(
        db,
        sessions=120,
        messages_per_session=900,
        create_messages_index=False,
        source="cron",
        session_source="cron",
    )
    reference_ids = _newest_first_reference_ids(db, include_sources=("cron",), exclude_sources=None)
    candidate_limit = max(20 * 8, 20)
    progress_interval = 100

    original_connect = agent_sessions.sqlite3.connect
    connect_with_progress_counter, progress_counter = _make_connect_with_progress_counter(
        interval=progress_interval
    )
    monkeypatch.setattr(agent_sessions.sqlite3, "connect", connect_with_progress_counter)
    try:
        measured_rows = agent_sessions.read_importable_agent_session_rows(
            db,
            limit=20,
            exclude_sources=None,
            include_sources=("cron",),
        )
        assert [row["id"] for row in measured_rows] == reference_ids[:20]
    finally:
        # Keep this helper isolated; the baseline must still run without the
        # counting handler to validate raw cost differences.
        monkeypatch.setattr(agent_sessions.sqlite3, "connect", original_connect)

    # Give the head path a small deterministic margin, then require the old
    # correlated query to exceed the same budget on the missing-index branch.
    progress_budget_ops = max(progress_counter["count"] + 200, 1)

    with pytest.raises(sqlite3.OperationalError, match="interrupted"):
        _execute_candidate_ordering_baseline_sql(
            db,
            candidate_limit,
            budget_ops=progress_budget_ops,
            interval=progress_interval,
            include_sources=("cron",),
            exclude_sources=None,
        )

    monkeypatch.setattr(
        agent_sessions.sqlite3,
        "connect",
        _make_connect_with_progress_budget(
            budget_ops=progress_budget_ops,
            interval=progress_interval,
        ),
    )

    rows = agent_sessions.read_importable_agent_session_rows(
        db,
        limit=20,
        exclude_sources=None,
        include_sources=("cron",),
    )
    assert [row["id"] for row in rows] == reference_ids[:20]


def test_importable_agent_rows_limit_includes_resumed_old_session(tmp_path):
    """The capped candidate window must not hide old sessions resumed recently."""
    db = tmp_path / "state.db"
    _make_state_db(db, sessions=200, messages_per_session=1)

    old_started = time.time() - 60 * 60 * 24 * 30
    recent_activity = time.time() + 60
    conn = sqlite3.connect(str(db))
    conn.execute(
        """
        INSERT INTO sessions
        (id, source, session_source, title, model, started_at, message_count, parent_session_id, ended_at, end_reason)
        VALUES ('cli_resumed_old', 'cli', 'cli', 'Old resumed session', 'openai/gpt-5', ?, 2, NULL, NULL, NULL)
        """,
        (old_started,),
    )
    conn.execute(
        "INSERT INTO messages (id, session_id, role, content, timestamp) VALUES ('old_msg_1', 'cli_resumed_old', 'user', 'old hello', ?)",
        (old_started,),
    )
    conn.execute(
        "INSERT INTO messages (id, session_id, role, content, timestamp) VALUES ('old_msg_2', 'cli_resumed_old', 'assistant', 'recent reply', ?)",
        (recent_activity,),
    )
    conn.commit()
    conn.close()

    rows = agent_sessions.read_importable_agent_session_rows(db, limit=20, exclude_sources=("webui",))

    assert rows[0]["id"] == "cli_resumed_old"
    assert rows[0]["actual_message_count"] == 2


def test_importable_agent_rows_zero_limit_skips_query_work(tmp_path):
    db = tmp_path / "state.db"
    _make_state_db(db, sessions=5, messages_per_session=1)

    assert agent_sessions.read_importable_agent_session_rows(db, limit=0, exclude_sources=("webui",)) == []