Initial release: iai-mcp v0.1.0

Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: XNLLLLH <XNLLLLH@users.noreply.github.com>

tests/test_socket_server_dispatch.py

@@ -0,0 +1,441 @@
+"""Plan 07-02 Wave 2 R1 acceptance: every dispatch method reachable over socket.
+
+Boots SocketServer (NEW per D7-07) against a tmp_path-isolated
+~/.iai-mcp/.daemon.sock equivalent and asserts that JSON-RPC 2.0 envelopes
+sent over the unix socket return the same response shape as the stdio path
+(R1, R6 backward-compat by construction).
+
+Reuses the short_socket_paths fixture pattern from test_daemon_dispatcher.py
+(AF_UNIX 104-byte cap mitigation via /tmp/iai-disp-<pid>-<n>/d.sock).
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def short_socket_paths(tmp_path, monkeypatch):
+    """Redirect LOCK_PATH + SOCKET_PATH + STATE_PATH to tmp_path.
+
+    AF_UNIX on macOS caps socket paths at ~104 bytes; pytest's tmp_path can
+    be too long under xdist. Use a short /tmp/iai-<pid>-<n>/ fallback for
+    the socket. The state file lives under tmp_path (regular filesystem).
+
+    Per D7-14: also point IAI_MCP_STORE at a tmp dir so MemoryStore()
+    constructed inside the test gets an isolated lancedb root.
+    """
+    from iai_mcp import concurrency, daemon_state
+
+    lock_path = tmp_path / ".lock"
+    sock_dir = Path(f"/tmp/iai-srvdisp-{os.getpid()}-{id(tmp_path)}")
+    sock_dir.mkdir(parents=True, exist_ok=True)
+    sock_path = sock_dir / "d.sock"
+    state_path = tmp_path / ".daemon-state.json"
+
+    monkeypatch.setattr(concurrency, "LOCK_PATH", lock_path)
+    monkeypatch.setattr(concurrency, "SOCKET_PATH", sock_path)
+    monkeypatch.setattr(daemon_state, "STATE_PATH", state_path)
+    # Per D7-14 isolate the lancedb store under tmp_path.
+    store_root = tmp_path / "lancedb_root"
+    store_root.mkdir(parents=True, exist_ok=True)
+    monkeypatch.setenv("IAI_MCP_STORE", str(store_root))
+
+    try:
+        yield lock_path, sock_path, state_path
+    finally:
+        try:
+            if sock_path.exists():
+                sock_path.unlink()
+        except OSError:
+            pass
+        try:
+            sock_dir.rmdir()
+        except OSError:
+            pass
+
+
+# ---------------------------------------------------------------------------
+# JSON-RPC helpers (reused by sibling Wave 2 test files)
+# ---------------------------------------------------------------------------
+
+
+async def _send_jsonrpc(
+    sock_path: Path,
+    method: str,
+    params: dict | None = None,
+    req_id: int | str = 1,
+    *,
+    timeout: float = 10.0,
+) -> dict:
+    """Per D7-01: send one JSON-RPC 2.0 envelope, read one response line.
+
+    Each call opens a fresh unix-stream connection (matches the per-connection
+    multiplexing pattern from D7-02; the daemon gives every client its own
+    coroutine).
+    """
+    reader, writer = await asyncio.wait_for(
+        asyncio.open_unix_connection(path=str(sock_path)),
+        timeout=timeout,
+    )
+    try:
+        envelope: dict = {"jsonrpc": "2.0", "id": req_id, "method": method}
+        if params is not None:
+            envelope["params"] = params
+        writer.write((json.dumps(envelope) + "\n").encode("utf-8"))
+        await writer.drain()
+        line = await asyncio.wait_for(reader.readline(), timeout=timeout)
+    finally:
+        writer.close()
+        try:
+            await writer.wait_closed()
+        except Exception:
+            pass
+    if not line:
+        raise AssertionError(f"daemon closed without reply (method={method})")
+    return json.loads(line.decode("utf-8"))
+
+
+async def _send_raw(sock_path: Path, raw_bytes: bytes, *, timeout: float = 5.0) -> dict:
+    """Send arbitrary bytes (used to test parse error path)."""
+    reader, writer = await asyncio.wait_for(
+        asyncio.open_unix_connection(path=str(sock_path)),
+        timeout=timeout,
+    )
+    try:
+        writer.write(raw_bytes)
+        await writer.drain()
+        line = await asyncio.wait_for(reader.readline(), timeout=timeout)
+    finally:
+        writer.close()
+        try:
+            await writer.wait_closed()
+        except Exception:
+            pass
+    if not line:
+        raise AssertionError("daemon closed without reply")
+    return json.loads(line.decode("utf-8"))
+
+
+async def _with_socket_server(sock_path: Path, store, coro_fn):
+    """Boot SocketServer + run coro_fn(sock_path, store), tear down cleanly.
+
+    Idle disabled (idle_secs=99999) so the test runner controls shutdown.
+    """
+    from iai_mcp.socket_server import SocketServer
+
+    srv = SocketServer(store, idle_secs=99999)
+    server_task = asyncio.create_task(srv.serve(socket_path=sock_path))
+
+    # Wait for bind (mirrors test_daemon_dispatcher.py:108-117).
+    for _ in range(250):
+        if sock_path.exists():
+            break
+        await asyncio.sleep(0.01)
+    if not sock_path.exists():
+        srv.shutdown_event.set()
+        try:
+            await asyncio.wait_for(server_task, timeout=5)
+        except Exception:
+            pass
+        raise AssertionError("socket never bound")
+
+    try:
+        result = await coro_fn(sock_path, store)
+    finally:
+        srv.shutdown_event.set()
+        try:
+            await asyncio.wait_for(server_task, timeout=5)
+        except Exception:
+            pass
+    return result
+
+
+# ---------------------------------------------------------------------------
+# R1 acceptance tests
+# ---------------------------------------------------------------------------
+
+
+def test_memory_recall_routed_over_socket(short_socket_paths):
+    """R1: memory_recall via socket returns the same shape as stdio path."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        return await _send_jsonrpc(
+            sock_path, "memory_recall",
+            {"cue": "phase 7 done", "budget_tokens": 400},
+            req_id=1,
+        )
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] == 1, resp
+    assert "result" in resp, resp
+    assert "hits" in resp["result"], resp["result"]
+
+
+def test_session_start_payload_routed(short_socket_paths):
+    """R1: session_start_payload via socket returns the assembler's dict."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        return await _send_jsonrpc(
+            sock_path, "session_start_payload", {}, req_id=2,
+        )
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] == 2, resp
+    assert "result" in resp, resp
+    # Empty store path returns the SessionStartPayload(empty) JSON shape;
+    # at minimum the wake_depth and l0/l1 keys must be present.
+    result = resp["result"]
+    assert "l0" in result and "l1" in result, result
+    assert "wake_depth" in result, result
+
+
+def test_profile_get_routed(short_socket_paths):
+    """R1: profile_get via socket returns the 11-knob registry dict.
+
+    Note: tools.ts wraps this as 'profile_get_set' with operation='get'/'set';
+    core.dispatch only knows 'profile_get' and 'profile_set' as primitives.
+    Tests against the core surface (D7-08: socket_server.py imports core.dispatch).
+    """
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        return await _send_jsonrpc(
+            sock_path, "profile_get", {"knob": None}, req_id=3,
+        )
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] == 3, resp
+    assert "result" in resp, resp
+    # profile.profile_get(None, _profile_state) returns the full knob dict;
+    # at least one canonical knob name must be present.
+    result = resp["result"]
+    assert isinstance(result, dict), result
+    # profile.profile_get(None, ...) returns
+    #   {'live': {<10 AUTIST + wake_depth>}, 'deferred': {}, 'total_knobs': 11}
+    # per src/iai_mcp/profile.py (Plan 07.12-02 removed AUTIST-02/08/11/12).
+    # keeps literal_preservation live.
+    assert "live" in result, result
+    assert "literal_preservation" in result["live"], result
+
+
+def test_topology_routed(short_socket_paths):
+    """R1: topology via socket returns the sigma snapshot or insufficient_data."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        return await _send_jsonrpc(sock_path, "topology", {}, req_id=4)
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] == 4, resp
+    assert "result" in resp, resp
+    result = resp["result"]
+    # Empty-store branch returns regime='insufficient_data' with N=0 and
+    # sigma=None; non-empty store returns numeric sigma. Both shapes carry
+    # a 'regime' key.
+    assert "regime" in result or "sigma" in result, result
+
+
+def test_invalid_jsonrpc_returns_minus_32600(short_socket_paths):
+    """R1: malformed envelope (jsonrpc='1.0') returns ERR_INVALID_REQUEST."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        # Bypass _send_jsonrpc helper (which always sends jsonrpc='2.0') by
+        # constructing the envelope manually.
+        bad = {"jsonrpc": "1.0", "id": 1, "method": "memory_recall"}
+        return await _send_raw(
+            sock_path, (json.dumps(bad) + "\n").encode("utf-8"),
+        )
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert "error" in resp, resp
+    assert resp["error"]["code"] == -32600, resp
+
+
+def test_unknown_method_returns_minus_32601(short_socket_paths):
+    """V3-03 fix: unknown method raises UnknownMethodError -> JSON-RPC -32601.
+
+    Pre-Phase-07.13: dispatch returned {"error": f"unknown method {method!r}"}
+    inside the result envelope. Post-fix: dispatch raises UnknownMethodError;
+    socket_server.handle catches it and emits {error: {code: -32601, ...}}.
+    """
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        return await _send_jsonrpc(
+            sock_path, "not_a_real_method", {}, req_id=5,
+        )
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] == 5, resp
+    assert "error" in resp, resp
+    assert "result" not in resp, resp                       # V3-03 tightening
+    assert resp["error"]["code"] == -32601, resp
+    assert "not_a_real_method" in resp["error"]["message"], resp
+
+
+def test_missing_required_param_returns_minus_32602(short_socket_paths):
+    """V3-04 fix: missing required param (e.g. memory_recall without 'cue')
+    raises KeyError inside dispatch -> JSON-RPC -32602 ERR_INVALID_PARAMS.
+
+    Pre-Phase-07.13: KeyError was mapped to -32601 ERR_METHOD_NOT_FOUND
+    (wrong code; "method not found" implies the route doesn't exist).
+    Post-fix: KeyError maps to -32602 with message 'missing required
+    param: <key>'.
+    """
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        # memory_recall consumes params["cue"] (required path) at core.py:213/249/273.
+        # Sending an empty params dict triggers KeyError on the first cue access.
+        return await _send_jsonrpc(
+            sock_path, "memory_recall", {}, req_id=6,
+        )
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] == 6, resp
+    assert "error" in resp, resp
+    assert "result" not in resp, resp
+    assert resp["error"]["code"] == -32602, resp
+    msg = resp["error"]["message"]
+    assert "missing required param" in msg, resp
+    assert "cue" in msg, resp
+
+
+def test_id_echoed_unchanged(short_socket_paths):
+    """D7-02: response.id matches request.id verbatim across types."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        r1 = await _send_jsonrpc(
+            sock_path, "session_start_payload", {}, req_id=1,
+        )
+        r2 = await _send_jsonrpc(
+            sock_path, "session_start_payload", {}, req_id=999,
+        )
+        r3 = await _send_jsonrpc(
+            sock_path, "session_start_payload", {}, req_id="some-string-id",
+        )
+        return r1, r2, r3
+
+    r1, r2, r3 = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert r1["id"] == 1, r1
+    assert r2["id"] == 999, r2
+    assert r3["id"] == "some-string-id", r3
+
+
+def test_unknown_method_does_not_crash_server(short_socket_paths):
+    """R1: an unknown method must not crash the server; the next call still works."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        # First call: unknown method — V3-03 fix: must surface as JSON-RPC error.
+        r_bad = await _send_jsonrpc(
+            sock_path, "definitely_not_a_method", {}, req_id=100,
+        )
+        # Second call must succeed against the same server.
+        r_good = await _send_jsonrpc(
+            sock_path, "session_start_payload", {}, req_id=101,
+        )
+        return r_bad, r_good
+
+    r_bad, r_good = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert r_bad["id"] == 100, r_bad
+    assert "error" in r_bad, r_bad                          # V3-03 tightening
+    assert "result" not in r_bad, r_bad                     # V3-03 tightening
+    assert r_good["id"] == 101, r_good
+    assert "result" in r_good, r_good
+
+
+def test_parse_error_returns_minus_32700(short_socket_paths):
+    """D7-01: malformed JSON → ERR_PARSE_ERROR with id=None per JSON-RPC 2.0 spec."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        return await _send_raw(sock_path, b"not valid json\n")
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] is None, resp
+    assert "error" in resp, resp
+    assert resp["error"]["code"] == -32700, resp
+
+
+def test_empty_params_defaults_to_empty_dict(short_socket_paths):
+    """D7-01: omitted params field → dispatch sees an empty dict, no crash."""
+    _, sock_path, _ = short_socket_paths
+    from iai_mcp.store import MemoryStore
+
+    store = MemoryStore()
+
+    async def _runner(sock_path, store):
+        # Pass params=None to _send_jsonrpc which omits the params key.
+        return await _send_jsonrpc(
+            sock_path, "session_start_payload", None, req_id=200,
+        )
+
+    resp = asyncio.run(_with_socket_server(sock_path, store, _runner))
+
+    assert resp["jsonrpc"] == "2.0", resp
+    assert resp["id"] == 200, resp
+    assert "result" in resp, resp