handler/tests/test_memory.py

"""Agent memory: the notes/links DAL, the /memory API, the bundled handler-memory MCP
server, and the SessionStart recall hook."""

from __future__ import annotations

import json

import pytest

from handler.db import repository as repo
from handler.db.engine import get_engine


@pytest.fixture
def seeded(env):
    """Two projects, one agent, and a small note graph."""
    with get_engine().begin() as conn:
        repo.create_project(conn, "proj", "/tmp/proj")
        repo.create_project(conn, "other", "/tmp/other")
        agent = repo.create_agent(conn, "proj", "api", "/tmp/proj", "working")
        n1 = repo.create_memory_note(
            conn, "Auth flow", "Tokens are minted in deps.py", "fact",
            project_id="proj", agent_id=agent["id"], tags=["auth"],
        )
        n2 = repo.create_memory_note(
            conn, "Use SQLite fallback", "Postgres is default, SQLite for tests",
            "decision", project_id="proj",
        )
        n3 = repo.create_memory_note(
            conn, "Global runbook", "How to rotate tokens", "runbook",
        )
        n4 = repo.create_memory_note(
            conn, "Other-project note", "Not visible from proj scope", "fact",
            project_id="other",
        )
        link = repo.create_memory_link(
            conn, n1["id"], n2["id"], relation="caused_by", agent_id=agent["id"]
        )
    return {"agent": agent, "n1": n1, "n2": n2, "n3": n3, "n4": n4, "link": link}


# --- DAL -------------------------------------------------------------------------------


def test_scoped_listing_and_search(seeded):
    with get_engine().begin() as conn:
        # Project scope = its notes + global, newest first; never another project's.
        notes = repo.list_memory_notes(conn, project_id="proj")
        ids = [n["id"] for n in notes]
        assert seeded["n4"]["id"] not in ids
        assert {seeded["n1"]["id"], seeded["n2"]["id"], seeded["n3"]["id"]} <= set(ids)

        # Every term must match; case-insensitive over title/body/kind.
        hits = repo.search_memory_notes(conn, "auth tokens", project_id="proj")
        assert [h["id"] for h in hits] == [seeded["n1"]["id"]]
        # Empty query = recent notes in scope.
        assert repo.search_memory_notes(conn, "", project_id="proj")
        # kind matches too.
        assert any(
            h["id"] == seeded["n3"]["id"]
            for h in repo.search_memory_notes(conn, "runbook", project_id="proj")
        )


def test_link_idempotent_and_graph_scope(seeded):
    with get_engine().begin() as conn:
        again = repo.create_memory_link(
            conn, seeded["n1"]["id"], seeded["n2"]["id"], relation="caused_by"
        )
        assert again["id"] == seeded["link"]["id"]  # re-asserting is not a new edge

        # A cross-scope link's edge drops out of a scoped graph when one endpoint is out.
        repo.create_memory_link(conn, seeded["n1"]["id"], seeded["n4"]["id"])
        graph = repo.memory_graph(conn, project_id="proj")
        graph_ids = {n["id"] for n in graph["notes"]}
        assert seeded["n4"]["id"] not in graph_ids
        assert [ln["id"] for ln in graph["links"]] == [seeded["link"]["id"]]


def test_delete_note_removes_edges(seeded):
    with get_engine().begin() as conn:
        assert repo.delete_memory_note(conn, seeded["n2"]["id"])
        assert repo.get_memory_note(conn, seeded["n2"]["id"]) is None
        assert repo.list_memory_links(conn, note_ids=[seeded["n1"]["id"]]) == []


def test_agent_delete_keeps_notes_project_delete_removes_them(seeded):
    with get_engine().begin() as conn:
        # Deleting the authoring agent orphans the attribution, never the knowledge.
        assert repo.delete_agent(conn, "proj", "api")
        note = repo.get_memory_note(conn, seeded["n1"]["id"])
        assert note is not None and note["agent_id"] is None
        link = repo.get_memory_link(conn, seeded["link"]["id"])
        assert link is not None and link["created_by_agent_id"] is None

        # Deleting the project takes its notes (and their edges); global notes stay.
        assert repo.delete_project(conn, "proj")
        assert repo.get_memory_note(conn, seeded["n1"]["id"]) is None
        assert repo.get_memory_note(conn, seeded["n2"]["id"]) is None
        assert repo.get_memory_note(conn, seeded["n3"]["id"]) is not None


# --- API -------------------------------------------------------------------------------


def test_api_note_crud_and_graph(client, auth, seeded):
    r = client.post(
        "/memory/notes",
        json={"title": "From the dashboard", "body": "operator wisdom", "kind": "gotcha"},
        headers=auth,
    )
    assert r.status_code == 201
    note = r.json()
    assert note["agent_id"] is None and note["project_id"] is None

    r = client.patch(f"/memory/notes/{note['id']}", json={"kind": "runbook"}, headers=auth)
    assert r.status_code == 200 and r.json()["kind"] == "runbook"

    r = client.get("/memory/notes?q=wisdom", headers=auth)
    assert [n["id"] for n in r.json()] == [note["id"]]

    r = client.get("/memory/graph?project_id=proj", headers=auth)
    graph = r.json()
    assert {n["id"] for n in graph["notes"]} >= {seeded["n1"]["id"], seeded["n3"]["id"]}
    assert graph["links"][0]["relation"] == "caused_by"

    assert client.delete(f"/memory/notes/{note['id']}", headers=auth).status_code == 200
    assert client.get(f"/memory/notes/{note['id']}", headers=auth).status_code == 404


def test_api_validation(client, auth, seeded):
    n1 = seeded["n1"]["id"]
    # Unknown project scope on create; self-links; missing endpoints.
    r = client.post(
        "/memory/notes",
        json={"title": "x", "body": "y", "project_id": "nope"},
        headers=auth,
    )
    assert r.status_code == 400
    r = client.post(
        "/memory/links", json={"src_note_id": n1, "dst_note_id": n1}, headers=auth
    )
    assert r.status_code == 400
    r = client.post(
        "/memory/links", json={"src_note_id": n1, "dst_note_id": 999999}, headers=auth
    )
    assert r.status_code == 404
    # Bad kind is a 422 straight from the schema.
    r = client.post("/memory/notes", json={"title": "x", "body": "y", "kind": "poem"}, headers=auth)
    assert r.status_code == 422
    # No token, no memory.
    assert client.get("/memory/notes").status_code == 401


# --- MCP server ------------------------------------------------------------------------


def _rpc(server, method, params=None, msg_id=1):
    from handler.mcpserver import handle_message

    return handle_message(server, {"jsonrpc": "2.0", "id": msg_id, "method": method,
                                   "params": params or {}})


def _call_tool(server, name, args):
    resp = _rpc(server, "tools/call", {"name": name, "arguments": args})
    result = resp["result"]
    return result["isError"], json.loads(result["content"][0]["text"]) if not result[
        "isError"
    ] else result["content"][0]["text"]


def test_mcp_protocol_basics(seeded):
    from handler.mcpserver import MemoryServer, handle_message

    server = MemoryServer(agent_id=seeded["agent"]["id"], project_id="proj")
    init = _rpc(server, "initialize", {"protocolVersion": "2025-06-18"})
    assert init["result"]["serverInfo"]["name"] == "handler-memory"
    assert init["result"]["protocolVersion"] == "2025-06-18"
    # Notifications get no response; unknown methods error.
    assert handle_message(server, {"jsonrpc": "2.0", "method": "notifications/initialized"}) is None
    assert _rpc(server, "resources/list")["error"]["code"] == -32601
    tools = _rpc(server, "tools/list")["result"]["tools"]
    assert {t["name"] for t in tools} == {
        "memory_search", "memory_get", "memory_save", "memory_link",
    }


def test_mcp_tools_roundtrip(seeded):
    from handler.mcpserver import MemoryServer

    server = MemoryServer(agent_id=seeded["agent"]["id"], project_id="proj")

    err, saved = _call_tool(server, "memory_save", {
        "title": "Worker heartbeats", "body": "Reaper marks stale workers crashed",
        "kind": "fact", "tags": ["workers"],
    })
    assert not err and saved["created"]
    assert saved["note"]["project_id"] == "proj"
    assert saved["note"]["agent_id"] == seeded["agent"]["id"]

    err, found = _call_tool(server, "memory_search", {"query": "heartbeats"})
    assert not err and found["count"] == 1

    # Search sees global notes; never the other project's.
    err, found = _call_tool(server, "memory_search", {"query": "rotate tokens"})
    assert not err and found["count"] == 1
    err, found = _call_tool(server, "memory_search", {"query": "Other-project"})
    assert not err and found["count"] == 0

    err, linked = _call_tool(server, "memory_link", {
        "src_note_id": saved["note"]["id"], "dst_note_id": seeded["n1"]["id"],
        "relation": "relates_to",
    })
    assert not err and linked["link"]["created_by_agent_id"] == seeded["agent"]["id"]

    err, got = _call_tool(server, "memory_get", {"note_id": saved["note"]["id"]})
    assert not err
    assert got["links"][0]["other_title"] == "Auth flow"

    # Update in place via note_id.
    err, updated = _call_tool(server, "memory_save", {
        "note_id": saved["note"]["id"], "title": "Worker heartbeats",
        "body": "expanded", "kind": "gotcha",
    })
    assert not err and updated["updated"]

    # Tool errors come back in-band, not as protocol errors.
    err, msg = _call_tool(server, "memory_get", {"note_id": 424242})
    assert err and "not found" in msg


def test_mcp_global_save(seeded):
    from handler.mcpserver import MemoryServer

    server = MemoryServer(agent_id=None, project_id="proj")
    err, saved = _call_tool(server, "memory_save", {
        "title": "For everyone", "body": "x", "global": True,
    })
    assert not err and saved["note"]["project_id"] is None


# --- SessionStart recall hook + launch wiring ------------------------------------------


def test_session_start_hook_injects_notes(seeded, capsys):
    from handler.hooks import memory_ctx
    from handler.hooks.context import HookInput, Identity

    ident = Identity(seeded["agent"]["id"], "proj", "api", "/tmp/proj")
    with get_engine().begin() as conn:
        result = memory_ctx.handle(conn, ident, HookInput(raw={}, event="session_start"))
    ctx = result["hookSpecificOutput"]["additionalContext"]
    assert result["hookSpecificOutput"]["hookEventName"] == "SessionStart"
    assert "Auth flow" in ctx and "Global runbook" in ctx
    assert "Other-project note" not in ctx
    assert "memory_search" in ctx  # the pointer at the tools
    # The hook wrote its JSON response to stdout for claude to consume.
    assert json.loads(capsys.readouterr().out)["hookSpecificOutput"]


def test_settings_wire_session_start_and_memory_allow(env):
    from handler.control import settings_gen

    settings = settings_gen.build_settings()
    hook = settings["hooks"]["SessionStart"][0]["hooks"][0]["command"]
    assert hook.endswith("-m handler.hooks session_start")
    assert "mcp__handler-memory" in settings["permissions"]["allow"]