chopratejas--headroom/tests/test_memory_tool_session_sticky.py

"""Session-sticky memory tool injection tests for PR-A7 (closes P0-6).

The cache-killer pattern this guards against (guide §6.3 #2):

  * Mid-session toggle: memory enabled in turn N injects `memory_save` /
    `memory_search` tool definitions into `body["tools"]`. Turn N+1
    disables memory; tool list shrinks; prefix bytes hash differently;
    prefix-cache misses; full prompt re-runs at provider cost.

  * Tool definition drift across deploys: the same logical tool is
    injected but the bytes differ (key insertion order, schema bump,
    description tweak). Even with the tool list intact, prefix bytes
    change.

The fix:

  * `SessionToolTracker`: bounded LRU keyed by (provider, session_id)
    storing GOLDEN tool-definition bytes from the first injection.
    Subsequent turns of that session always replay those bytes — even
    when memory is disabled mid-session (sticky-on per §6.3 #2).

  * `apply_session_sticky_memory_tools`: single coordination point
    used at every memory injection site (Anthropic custom + native,
    OpenAI Chat-Completions + Responses + WS).

Operator opt-in `HEADROOM_TOOL_INJECTION_STICKY=disabled` short-
circuits the tracker (per-turn decision flows through verbatim — the
broken behavior). That mode is loud and explicit per realignment build
constraint #4 — NOT a silent fallback. It exists for diagnostic shadow
tracing and emergency rollback only.
"""

from __future__ import annotations

import hashlib
import json
import threading
from pathlib import Path
from typing import Any

import pytest

from headroom.proxy.helpers import (
    SessionToolTracker,
    _reset_session_tool_tracker_for_test,
    apply_session_sticky_memory_tools,
    get_session_tool_tracker,
    get_tool_injection_sticky_mode,
    get_tool_tracker_max_sessions,
    serialize_tool_definition_canonical,
)
from headroom.proxy.memory_handler import MemoryConfig, MemoryHandler

FIXTURES_DIR = Path(__file__).parent / "fixtures" / "memory_tool_definitions"


# ---------------------------------------------------------------------------
# Test isolation
# ---------------------------------------------------------------------------


@pytest.fixture(autouse=True)
def _isolate_tracker(monkeypatch: pytest.MonkeyPatch) -> None:
    """Reset env + tracker singleton between tests."""
    monkeypatch.delenv("HEADROOM_TOOL_INJECTION_STICKY", raising=False)
    monkeypatch.delenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", raising=False)
    _reset_session_tool_tracker_for_test()
    yield
    _reset_session_tool_tracker_for_test()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _anthropic_memory_defs() -> list[dict[str, Any]]:
    h = MemoryHandler(MemoryConfig(enabled=True, inject_tools=True))
    return h.compute_memory_tool_definitions("anthropic")


def _openai_memory_defs() -> list[dict[str, Any]]:
    h = MemoryHandler(MemoryConfig(enabled=True, inject_tools=True))
    return h.compute_memory_tool_definitions("openai")


def _names_in(tools: list[dict[str, Any]]) -> set[str]:
    out: set[str] = set()
    for t in tools:
        n = t.get("name") or (t.get("function") or {}).get("name") or t.get("type")
        if n:
            out.add(n)
    return out


# ---------------------------------------------------------------------------
# `SessionToolTracker` direct unit tests
# ---------------------------------------------------------------------------


def test_should_inject_false_for_unknown_session() -> None:
    tracker = SessionToolTracker(max_sessions=10)
    assert tracker.should_inject("anthropic", "s-1") is False


def test_record_then_should_inject_true() -> None:
    tracker = SessionToolTracker(max_sessions=10)
    tracker.record_injection(
        provider="anthropic",
        session_id="s-1",
        tool_name="memory_save",
        tool_definition_bytes=b'{"name":"memory_save"}',
    )
    assert tracker.should_inject("anthropic", "s-1") is True


def test_get_golden_definitions_returns_recorded_bytes() -> None:
    tracker = SessionToolTracker(max_sessions=10)
    tracker.record_injection(
        provider="anthropic",
        session_id="s-1",
        tool_name="memory_save",
        tool_definition_bytes=b'{"name":"memory_save","x":1}',
    )
    tracker.record_injection(
        provider="anthropic",
        session_id="s-1",
        tool_name="memory_search",
        tool_definition_bytes=b'{"name":"memory_search","x":2}',
    )
    golden = tracker.get_golden_definitions("anthropic", "s-1")
    assert golden is not None
    assert [name for name, _ in golden] == ["memory_save", "memory_search"]
    assert golden[0][1] == b'{"name":"memory_save","x":1}'
    assert golden[1][1] == b'{"name":"memory_search","x":2}'


def test_record_first_write_wins_on_duplicate_name() -> None:
    """Re-recording the same tool name is a no-op (prevents drift mid-session)."""
    tracker = SessionToolTracker(max_sessions=10)
    tracker.record_injection(
        provider="anthropic",
        session_id="s-1",
        tool_name="memory_save",
        tool_definition_bytes=b"original",
    )
    tracker.record_injection(
        provider="anthropic",
        session_id="s-1",
        tool_name="memory_save",
        tool_definition_bytes=b"new-bytes",
    )
    golden = tracker.get_golden_definitions("anthropic", "s-1") or []
    assert golden == [("memory_save", b"original")]


def test_provider_isolation_anthropic_vs_openai_same_session_id() -> None:
    """Same session_id under two providers keeps independent state."""
    tracker = SessionToolTracker(max_sessions=10)
    tracker.record_injection(
        provider="anthropic",
        session_id="shared",
        tool_name="memory_save",
        tool_definition_bytes=b"anthropic-bytes",
    )
    tracker.record_injection(
        provider="openai",
        session_id="shared",
        tool_name="memory_save",
        tool_definition_bytes=b"openai-bytes",
    )
    a_golden = tracker.get_golden_definitions("anthropic", "shared") or []
    o_golden = tracker.get_golden_definitions("openai", "shared") or []
    assert a_golden == [("memory_save", b"anthropic-bytes")]
    assert o_golden == [("memory_save", b"openai-bytes")]


def test_lru_eviction_at_max_sessions() -> None:
    """Bounded LRU pops oldest session when overflowing."""
    tracker = SessionToolTracker(max_sessions=2)
    tracker.record_injection(
        provider="anthropic",
        session_id="s-1",
        tool_name="memory_save",
        tool_definition_bytes=b"a",
    )
    tracker.record_injection(
        provider="anthropic",
        session_id="s-2",
        tool_name="memory_save",
        tool_definition_bytes=b"b",
    )
    assert tracker.active_sessions == 2

    # Touch s-1 so s-2 becomes the LRU.
    assert tracker.should_inject("anthropic", "s-1") is True

    # Add s-3: pops s-2.
    tracker.record_injection(
        provider="anthropic",
        session_id="s-3",
        tool_name="memory_save",
        tool_definition_bytes=b"c",
    )
    assert tracker.active_sessions == 2
    assert tracker.should_inject("anthropic", "s-2") is False
    assert tracker.should_inject("anthropic", "s-1") is True
    assert tracker.should_inject("anthropic", "s-3") is True


def test_max_sessions_invalid_raises() -> None:
    with pytest.raises(ValueError):
        SessionToolTracker(max_sessions=0)
    with pytest.raises(ValueError):
        SessionToolTracker(max_sessions=-1)


def test_blank_provider_or_session_raises() -> None:
    tracker = SessionToolTracker(max_sessions=10)
    with pytest.raises(ValueError):
        tracker.should_inject("", "s")
    with pytest.raises(ValueError):
        tracker.should_inject("anthropic", "")
    with pytest.raises(ValueError):
        tracker.record_injection(
            provider="",
            session_id="s",
            tool_name="x",
            tool_definition_bytes=b"y",
        )
    with pytest.raises(ValueError):
        tracker.record_injection(
            provider="anthropic",
            session_id="s",
            tool_name="",
            tool_definition_bytes=b"y",
        )
    with pytest.raises(ValueError):
        tracker.record_injection(
            provider="anthropic",
            session_id="s",
            tool_name="x",
            tool_definition_bytes=b"",
        )


def test_thread_safe_concurrent_access() -> None:
    """N threads on same session: no exceptions, all pinned bytes survive."""
    tracker = SessionToolTracker(max_sessions=10)
    n_threads = 16
    iterations = 50
    errors: list[BaseException] = []

    def worker(thread_idx: int) -> None:
        try:
            for i in range(iterations):
                tracker.record_injection(
                    provider="anthropic",
                    session_id="shared",
                    tool_name=f"t{thread_idx}-i{i}",
                    tool_definition_bytes=f"bytes-{thread_idx}-{i}".encode(),
                )
                # Concurrent reads.
                tracker.should_inject("anthropic", "shared")
                tracker.get_golden_definitions("anthropic", "shared")
        except BaseException as e:  # noqa: BLE001
            errors.append(e)

    threads = [threading.Thread(target=worker, args=(idx,)) for idx in range(n_threads)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    assert errors == []
    golden = tracker.get_golden_definitions("anthropic", "shared") or []
    names = {name for name, _ in golden}
    expected = {f"t{idx}-i{i}" for idx in range(n_threads) for i in range(iterations)}
    assert expected.issubset(names)


def test_singleton_returns_same_instance() -> None:
    a = get_session_tool_tracker()
    b = get_session_tool_tracker()
    assert a is b


def test_singleton_reset_replaces_instance() -> None:
    a = get_session_tool_tracker()
    _reset_session_tool_tracker_for_test()
    b = get_session_tool_tracker()
    assert a is not b


# ---------------------------------------------------------------------------
# Env vars
# ---------------------------------------------------------------------------


def test_max_sessions_env_var_default() -> None:
    assert get_tool_tracker_max_sessions() == 1000


def test_max_sessions_env_var_custom(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "42")
    assert get_tool_tracker_max_sessions() == 42


def test_max_sessions_env_var_invalid_raises(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "0")
    with pytest.raises(ValueError):
        get_tool_tracker_max_sessions()
    monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "-3")
    with pytest.raises(ValueError):
        get_tool_tracker_max_sessions()
    monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "not-int")
    with pytest.raises(ValueError):
        get_tool_tracker_max_sessions()


def test_sticky_mode_default_enabled() -> None:
    assert get_tool_injection_sticky_mode() == "enabled"


def test_sticky_mode_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "disabled")
    assert get_tool_injection_sticky_mode() == "disabled"


def test_sticky_mode_invalid_raises(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "yolo")
    with pytest.raises(ValueError, match="HEADROOM_TOOL_INJECTION_STICKY"):
        get_tool_injection_sticky_mode()


# ---------------------------------------------------------------------------
# `apply_session_sticky_memory_tools` integration
# ---------------------------------------------------------------------------


def test_injection_in_turn_1_repeats_in_turn_2_same_session_anthropic() -> None:
    """Core sticky-on guarantee for Anthropic provider."""
    defs = _anthropic_memory_defs()
    assert len(defs) >= 2

    # Turn 1: memory enabled — first-time injection.
    tools1, was1 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="s-1",
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was1 is True
    names1 = _names_in(tools1)
    assert "memory_save" in names1
    assert "memory_search" in names1

    # Turn 2: memory STILL enabled — bytes match turn 1.
    tools2, was2 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="s-1",
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was2 is True
    # Same set of memory tools.
    assert _names_in(tools2) == names1


def test_injection_in_turn_1_repeats_in_turn_2_same_session_openai() -> None:
    defs = _openai_memory_defs()
    assert len(defs) >= 2

    tools1, was1 = apply_session_sticky_memory_tools(
        provider="openai",
        session_id="o-1",
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was1 is True
    names1 = _names_in(tools1)
    assert "memory_save" in names1
    assert "memory_search" in names1

    tools2, was2 = apply_session_sticky_memory_tools(
        provider="openai",
        session_id="o-1",
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was2 is True
    assert _names_in(tools2) == names1


def test_byte_equal_tool_definition_across_turns() -> None:
    """The injected tool list serialization is BYTE-equal turn 1 vs turn 2.

    Pin the bytes via the golden snapshot fixture.
    """
    defs = _anthropic_memory_defs()

    tools1, _ = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="bytestable-1",
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    tools2, _ = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="bytestable-1",
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )

    # Byte-equality (the cache-stable invariant).
    bytes1 = b"".join(serialize_tool_definition_canonical(t) for t in tools1)
    bytes2 = b"".join(serialize_tool_definition_canonical(t) for t in tools2)
    assert bytes1 == bytes2

    # Match the golden fixture (computed via the pinned helper).
    fixture = json.loads((FIXTURES_DIR / "anthropic.json").read_text())
    fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
    assert bytes1 == fixture_bytes


def test_byte_equal_tool_definition_across_turns_openai() -> None:
    defs = _openai_memory_defs()
    tools1, _ = apply_session_sticky_memory_tools(
        provider="openai",
        session_id="bytestable-2",
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    tools2, _ = apply_session_sticky_memory_tools(
        provider="openai",
        session_id="bytestable-2",
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    bytes1 = b"".join(serialize_tool_definition_canonical(t) for t in tools1)
    bytes2 = b"".join(serialize_tool_definition_canonical(t) for t in tools2)
    assert bytes1 == bytes2

    fixture = json.loads((FIXTURES_DIR / "openai.json").read_text())
    fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
    assert bytes1 == fixture_bytes


def test_memory_disabled_after_inject_still_injects() -> None:
    """Turn 1 injects; turn 2 has memory disabled; turn 2 still injects golden bytes."""
    defs = _anthropic_memory_defs()

    # Turn 1: memory enabled.
    tools1, was1 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="s-cancel-1",
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was1 is True
    names1 = _names_in(tools1)

    # Turn 2: memory DISABLED for this turn (e.g. inject_tools flag flipped).
    # `inject_this_turn=False` AND `memory_tools_to_inject=[]` mimic the
    # caller's behavior under disabled-memory: nothing fresh to inject.
    tools2, was2 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="s-cancel-1",
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=[],
        inject_this_turn=False,
    )
    # Sticky-on: the golden bytes are still injected even though caller
    # passed nothing this turn.
    assert was2 is True
    assert _names_in(tools2) == names1

    # Bytes match.
    bytes1 = b"".join(serialize_tool_definition_canonical(t) for t in tools1)
    bytes2 = b"".join(serialize_tool_definition_canonical(t) for t in tools2)
    assert bytes1 == bytes2


def test_different_sessions_independent() -> None:
    """Session A injects; session B doesn't; verify isolation."""
    defs = _anthropic_memory_defs()

    # Session A: inject.
    _, was_a = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="A",
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was_a is True

    # Session B: NO inject this turn, no prior history.
    tools_b, was_b = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="B",
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=[],
        inject_this_turn=False,
    )
    assert was_b is False
    assert _names_in(tools_b) == set()

    # Session A still has its golden state.
    tools_a2, was_a2 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="A",
        request_id="r-3",
        existing_tools=[],
        memory_tools_to_inject=[],
        inject_this_turn=False,
    )
    assert was_a2 is True
    assert "memory_save" in _names_in(tools_a2)


def test_disabled_mode_passes_through_per_turn_decision(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """`HEADROOM_TOOL_INJECTION_STICKY=disabled` reverts to per-turn behavior.

    This is the broken behavior — explicit operator opt-in only. Turn 1
    injects; turn 2 with `inject_this_turn=False` does NOT replay (the
    sticky guarantee is bypassed).
    """
    monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "disabled")
    defs = _anthropic_memory_defs()

    tools1, was1 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="s-disabled",
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was1 is True
    assert "memory_save" in _names_in(tools1)

    # Turn 2: caller says don't inject. Disabled mode → tracker bypassed.
    tools2, was2 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="s-disabled",
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=[],
        inject_this_turn=False,
    )
    assert was2 is False
    assert _names_in(tools2) == set()


def test_existing_tool_with_memory_name_not_double_injected() -> None:
    """If client already has a tool by the same name, do not re-append it."""
    defs = _anthropic_memory_defs()
    client_tools: list[dict[str, Any]] = [
        {"name": "memory_save", "description": "client's own", "input_schema": {}}
    ]

    tools, _ = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id="s-dedup",
        request_id="r-1",
        existing_tools=client_tools,
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    # Exactly one tool named "memory_save".
    save_count = sum(1 for t in tools if t.get("name") == "memory_save")
    assert save_count == 1


def test_no_session_id_falls_back_to_per_turn(caplog: pytest.LogCaptureFixture) -> None:
    """`session_id=None` (e.g. WS pre-session) bypasses the tracker."""
    defs = _anthropic_memory_defs()
    tools1, was1 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id=None,
        request_id="r-1",
        existing_tools=[],
        memory_tools_to_inject=defs,
        inject_this_turn=True,
    )
    assert was1 is True
    assert "memory_save" in _names_in(tools1)

    # Without session_id we can't replay across turns.
    tools2, was2 = apply_session_sticky_memory_tools(
        provider="anthropic",
        session_id=None,
        request_id="r-2",
        existing_tools=[],
        memory_tools_to_inject=[],
        inject_this_turn=False,
    )
    assert was2 is False


def test_unknown_provider_raises() -> None:
    with pytest.raises(ValueError, match="unsupported provider"):
        apply_session_sticky_memory_tools(
            provider="gemini",  # type: ignore[arg-type]
            session_id="s",
            request_id="r",
            existing_tools=[],
            memory_tools_to_inject=[],
            inject_this_turn=True,
        )


# ---------------------------------------------------------------------------
# Golden fixture pinning
# ---------------------------------------------------------------------------


def test_anthropic_fixture_matches_helper_output() -> None:
    """Fixture file pins the canonical bytes — regenerate if this fails."""
    fixture = json.loads((FIXTURES_DIR / "anthropic.json").read_text())
    assert fixture["provider"] == "anthropic"

    helper_defs = _anthropic_memory_defs()
    helper_bytes = b"".join(serialize_tool_definition_canonical(t) for t in helper_defs)
    fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
    assert helper_bytes == fixture_bytes, (
        "Anthropic memory tool definitions drifted from golden fixture. "
        "If intentional, regenerate "
        "tests/fixtures/memory_tool_definitions/anthropic.json. "
        f"Helper SHA-256: {hashlib.sha256(helper_bytes).hexdigest()} "
        f"Fixture SHA-256: {hashlib.sha256(fixture_bytes).hexdigest()}"
    )


def test_openai_fixture_matches_helper_output() -> None:
    fixture = json.loads((FIXTURES_DIR / "openai.json").read_text())
    assert fixture["provider"] == "openai"

    helper_defs = _openai_memory_defs()
    helper_bytes = b"".join(serialize_tool_definition_canonical(t) for t in helper_defs)
    fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
    assert helper_bytes == fixture_bytes, (
        "OpenAI memory tool definitions drifted from golden fixture. "
        "If intentional, regenerate "
        "tests/fixtures/memory_tool_definitions/openai.json. "
        f"Helper SHA-256: {hashlib.sha256(helper_bytes).hexdigest()} "
        f"Fixture SHA-256: {hashlib.sha256(fixture_bytes).hexdigest()}"
    )