mirror of
https://github.com/0xWheatyz/handler.git
synced 2026-08-30 13:26:25 +00:00
dc80e7bc28
The distilled knowledge layer over the raw log/transcript history, inspired by TencentDB-Agent-Memory's memory-hub model, adapted to handler's invariants (state lives only in the database; workers stay stateless and disposable): - memory_notes + memory_links tables (migration 0014) with a portable DAL: scoped listing/search (project + global), idempotent linking, and a one-read graph. Deleting an agent nulls attribution but keeps its notes; deleting a project removes its notes and edges, leaving global knowledge intact. - /memory API routes: reads on the normal token, note/link authoring on the admin token, plus GET /memory/graph for the dashboard. - Bundled handler-memory MCP server (python -m handler.mcpserver), a dependency-free stdio JSON-RPC implementation injected into every launch's --mcp-config ahead of the DB connectors and allowlisted in generated settings, exposing memory_search / memory_get / memory_save / memory_link. Identity and DATABASE_URL arrive via the spawn env, same as hooks. - SessionStart recall hook: injects the most recent notes in scope as additional context at session start, best-effort, never blocking. - Memory page in the web UI: a hand-rolled force-directed SVG graph of the note web (colored by kind, hover highlights, click-through details), plus note/link authoring — the app's first visualization, no chart dependency. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01WYqkoYPX8NAo1V2KyXr1pk
269 lines
11 KiB
Python
269 lines
11 KiB
Python
"""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"]
|