honcho/tests/crud/test_session_scope_clauses.py

"""Observer session scope is enforced in SQL, not as a fetched name list.

A peer's session-membership count is unbounded. Materializing it turns every
session name into its own bind parameter, and the PostgreSQL wire protocol caps
parameters at 65535 per statement, so a peer in enough sessions yields a
statement the driver refuses to serialize. `get_observation_context` applies the
scope twice in one statement, which halves that ceiling to ~32.7k sessions.

These tests assert on compiled SQL and never execute a statement.
"""

from typing import Any

import pytest
from sqlalchemy import Select, select
from sqlalchemy.dialects import postgresql

from src import models
from src.crud.message import resolve_session_scope_clauses
from src.utils.agent_tools import get_observation_context


def _compile(stmt: Select[Any]) -> tuple[str, dict[str, Any]]:
    compiled = stmt.compile(
        dialect=postgresql.dialect(),
        compile_kwargs={"render_postcompile": True},
    )
    return str(compiled), dict(compiled.params)


class _FakeResult:
    def scalars(self) -> "_FakeResult":
        return self

    def all(self) -> list[Any]:
        return []


class _CapturingDB:
    """Captures statements instead of executing them."""

    def __init__(self) -> None:
        self.statements: list[Any] = []

    async def execute(self, stmt: Any) -> _FakeResult:
        self.statements.append(stmt)
        return _FakeResult()


@pytest.mark.parametrize(
    ("session_allowlist", "observer", "expect_exists", "expected_params"),
    [
        # Observer only: membership becomes a correlated EXISTS, so only the
        # workspace and peer are bound — never the session names, which is what
        # keeps the parameter count from growing with membership.
        (None, "observer-peer", True, ["observer-peer", "workspace"]),
        # Allowlist only: caller-supplied and therefore bounded, so IN is fine.
        (["s1", "s2"], None, False, ["s1", "s2"]),
        # Both: the EXISTS is intersected with the bounded IN.
        (["s1"], "observer-peer", True, ["observer-peer", "s1", "workspace"]),
        # Neither: unrestricted, nothing filtered and nothing bound.
        (None, None, False, []),
    ],
    ids=["observer-only", "allowlist-only", "observer-and-allowlist", "unrestricted"],
)
def test_scope_clause_shape(
    session_allowlist: list[str] | None,
    observer: str | None,
    expect_exists: bool,
    expected_params: list[str],
) -> None:
    clauses, deny = resolve_session_scope_clauses(
        "workspace",
        None,
        session_allowlist,
        observer,
        models.Message.session_name,
    )

    assert not deny

    sql, params = _compile(select(models.Message.public_id).where(*clauses))

    assert ("EXISTS" in sql.upper()) is expect_exists
    assert ("session_peers" in sql) is expect_exists
    assert sorted(params.values()) == expected_params


@pytest.mark.parametrize(
    ("session_name", "session_allowlist", "observer"),
    [
        # An empty allowlist fails closed rather than matching everything.
        (None, [], "observer-peer"),
        (None, [], None),
        # A pinned session the allowlist forbids fails closed.
        ("s9", ["s1", "s2"], "observer-peer"),
    ],
    ids=[
        "empty-allowlist-with-observer",
        "empty-allowlist-without-observer",
        "pinned-session-not-in-allowlist",
    ],
)
def test_scope_fails_closed(
    session_name: str | None,
    session_allowlist: list[str] | None,
    observer: str | None,
) -> None:
    clauses, deny = resolve_session_scope_clauses(
        "workspace",
        session_name,
        session_allowlist,
        observer,
        models.Message.session_name,
    )

    assert deny
    assert clauses == []


@pytest.mark.asyncio
async def test_get_observation_context_scope_costs_no_per_session_parameters() -> None:
    """The statement's parameter count depends on message_ids, not membership."""
    db = _CapturingDB()
    message_ids = [f"msg-{i}" for i in range(5)]

    await get_observation_context(
        db,  # pyright: ignore[reportArgumentType]
        "workspace",
        None,
        message_ids,
        observer="observer-peer",
    )

    assert len(db.statements) == 1
    sql, params = _compile(db.statements[0])

    # The scope must be applied to *both* the CTE and the outer select — a
    # materialized list would have cost two parameters per session there.
    # Count the subquery's FROM rather than EXISTS: the adjacency check is also
    # an EXISTS, so counting those would pass with the CTE's scope missing.
    assert sql.count("FROM public.session_peers") == 2
    # Both must correlate to the enclosing `messages` row. An uncorrelated
    # subquery compiles just as happily and would silently drop scoping
    # instead of enforcing it.
    assert sql.count("session_peers.session_name = public.messages.session_name") == 2

    # Everything bound is either a message id, the workspace, the peer, or the
    # ±1 adjacency window — nothing that scales with the peer's session count.
    expected = set(message_ids) | {"workspace", "observer-peer", -1, 1}
    assert set(params.values()) <= expected


@pytest.mark.asyncio
async def test_get_observation_context_denies_without_querying() -> None:
    """Fail-closed scopes must not reach the database at all."""
    db = _CapturingDB()

    result = await get_observation_context(
        db,  # pyright: ignore[reportArgumentType]
        "workspace",
        None,
        ["msg-1"],
        observer="observer-peer",
        session_allowlist=[],
    )

    assert result == []
    assert db.statements == []