JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MCP: add Streamable HTTP transport endpoint (#4122)

Browse Source 
Co-authored-by: utkarsh240799 <utkarsh240799@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Lev Neiman
2026-03-25 04:06:37 -07:00
committed by GitHub
parent 924ac00c52
commit 13c7f84eec
3 changed files with 478 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
openmemory/api/app/mcp_server.py
+82 -1
View File
@@ -21,6 +21,8 @@ import json
import logging
import uuid
import anyio
from app.database import SessionLocal
from app.models import Memory, MemoryAccessLog, MemoryState, MemoryStatusHistory
from app.utils.db import get_user_and_app
@@ -31,6 +33,8 @@ from fastapi import FastAPI, Request
from fastapi.routing import APIRouter
from mcp.server.fastmcp import FastMCP
from mcp.server.sse import SseServerTransport
from mcp.server.streamable_http import StreamableHTTPServerTransport
from starlette.responses import Response
# Load environment variables
load_dotenv()
@@ -438,7 +442,10 @@ async def handle_sse(request: Request):
client_token = client_name_var.set(client_name or "")
try:
# Handle SSE connection
# NOTE: request._send is the raw ASGI `send` callable. Starlette does not
# expose it publicly, but the MCP SDK transports require the raw ASGI
# interface (scope, receive, send). This is the standard pattern from the
# MCP Python SDK examples.
async with sse.connect_sse(
request.scope,
request.receive,
@@ -485,6 +492,80 @@ async def handle_post_message(request: Request):
finally:
pass
@mcp_router.api_route("/{client_name}/http/{user_id}", methods=["POST", "GET", "DELETE"])
async def handle_streamable_http(request: Request):
"""Handle Streamable HTTP connections for a specific user and client.
Uses the Streamable HTTP transport (MCP spec 2025-03-26+) which replaces
the deprecated SSE transport. Runs in stateless mode — each request is
handled independently with no persistent session.
The transport writes its response directly to the ASGI ``send`` callable.
We intercept it via ``capture_send`` so we can return a proper ``Response``
to FastAPI — otherwise FastAPI would also try to send its own response,
causing a "double-response" bug.
"""
uid = request.path_params.get("user_id")
user_token = user_id_var.set(uid or "")
client_name = request.path_params.get("client_name")
client_token = client_name_var.set(client_name or "")
# Intercept the ASGI messages the transport sends so we can return them
# as a single Response to FastAPI. Without this, FastAPI would attempt to
# write its own response after the transport already wrote one.
response_started = False
response_status = 200
response_headers: list[tuple[bytes, bytes]] = []
response_body = bytearray()
async def capture_send(message):
nonlocal response_started, response_status
if message["type"] == "http.response.start":
response_started = True
response_status = message["status"]
response_headers.extend(message.get("headers", []))
elif message["type"] == "http.response.body":
response_body.extend(message.get("body", b""))
try:
transport = StreamableHTTPServerTransport(
mcp_session_id=None,
is_json_response_enabled=True,
)
async with anyio.create_task_group() as tg:
async def run_server(*, task_status=anyio.TASK_STATUS_IGNORED):
async with transport.connect() as (read_stream, write_stream):
task_status.started()
await mcp._mcp_server.run(
read_stream,
write_stream,
mcp._mcp_server.create_initialization_options(),
stateless=True,
)
await tg.start(run_server)
await transport.handle_request(request.scope, request.receive, capture_send)
await transport.terminate()
tg.cancel_scope.cancel()
finally:
user_id_var.reset(user_token)
client_name_var.reset(client_token)
if not response_started:
return Response(status_code=500, content=b"Transport did not produce a response")
# Header dict conversion is safe here: the MCP transport in stateless JSON
# mode only emits single-valued headers (Content-Type, Content-Length).
return Response(
content=bytes(response_body),
status_code=response_status,
headers={k.decode(): v.decode() for k, v in response_headers},
)
def setup_mcp_server(app: FastAPI):
"""Setup MCP server with the FastAPI application"""
mcp._mcp_server.name = "mem0-mcp-server"
openmemory/api/tests/__init__.py
View File
openmemory/api/tests/test_mcp_server.py
+396
View File
@@ -0,0 +1,396 @@
"""Tests for the MCP server endpoints (SSE and Streamable HTTP transports).
Covers the Streamable HTTP transport (MCP spec 2025-03-26+) and the legacy SSE
transport. Tests exercise the full JSON-RPC flow — initialize, tools/list,
tools/call — as well as error handling and context-variable isolation.
"""
import os
# Set dummy keys before any imports that trigger client initialization
os.environ.setdefault("OPENAI_API_KEY", "test-key")
import pytest
import pytest_asyncio
from httpx import ASGITransport, AsyncClient
from app.mcp_server import client_name_var, mcp, mcp_router, user_id_var
# MCP Streamable HTTP requires the Accept header to include application/json.
# Including text/event-stream as well satisfies GET (SSE) requests.
MCP_HEADERS = {"Accept": "application/json, text/event-stream"}
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture
def test_app():
"""Create a minimal FastAPI app with just the MCP router for testing."""
from fastapi import FastAPI
app = FastAPI()
app.include_router(mcp_router)
return app
@pytest_asyncio.fixture
async def client(test_app):
"""Async HTTP client wired to the test app via ASGI transport."""
transport = ASGITransport(app=test_app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
yield ac
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _jsonrpc(method: str, params: dict | None = None, req_id: int = 1) -> dict:
"""Build a JSON-RPC 2.0 request envelope."""
return {
"jsonrpc": "2.0",
"id": req_id,
"method": method,
"params": params or {},
}
def _initialize_payload(req_id: int = 1) -> dict:
return _jsonrpc(
"initialize",
{
"protocolVersion": "2025-03-26",
"capabilities": {},
"clientInfo": {"name": "test-client", "version": "0.1.0"},
},
req_id=req_id,
)
# ---------------------------------------------------------------------------
# Streamable HTTP — route existence & basic protocol
# ---------------------------------------------------------------------------
class TestStreamableHTTPBasic:
"""Verify the Streamable HTTP route is registered and responds."""
@pytest.mark.asyncio
async def test_post_initialize(self, client):
"""POST initialize should return a valid JSON-RPC result."""
resp = await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
assert resp.status_code == 200
data = resp.json()
assert data["jsonrpc"] == "2.0"
assert data["id"] == 1
assert "result" in data
result = data["result"]
assert "serverInfo" in result
assert "capabilities" in result
assert result["protocolVersion"] == "2025-03-26"
@pytest.mark.asyncio
async def test_delete_returns_method_not_allowed(self, client):
"""DELETE in stateless mode should return 405 (no session to terminate)."""
resp = await client.delete(
"/mcp/testclient/http/user1",
headers=MCP_HEADERS,
)
assert resp.status_code == 405
@pytest.mark.asyncio
async def test_missing_accept_header_returns_406(self, client):
"""POST without the required Accept header should return 406."""
resp = await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
)
assert resp.status_code == 406
@pytest.mark.asyncio
async def test_invalid_json_returns_400(self, client):
"""POST with unparseable body should return 400."""
resp = await client.post(
"/mcp/testclient/http/user1",
content=b"not json",
headers={**MCP_HEADERS, "Content-Type": "application/json"},
)
assert resp.status_code == 400
@pytest.mark.asyncio
async def test_route_not_found_for_wrong_path(self, client):
"""Requests to a non-existent path should 404."""
resp = await client.post(
"/mcp/testclient/nonexistent/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
assert resp.status_code == 404
# ---------------------------------------------------------------------------
# Streamable HTTP — full protocol flow
# ---------------------------------------------------------------------------
class TestStreamableHTTPProtocol:
"""End-to-end JSON-RPC flows over Streamable HTTP."""
@pytest.mark.asyncio
async def test_tools_list(self, client):
"""tools/list should return all registered MCP tools."""
init_resp = await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
assert init_resp.status_code == 200
resp = await client.post(
"/mcp/testclient/http/user1",
json=_jsonrpc("tools/list", req_id=2),
headers=MCP_HEADERS,
)
assert resp.status_code == 200
data = resp.json()
assert "result" in data
tool_names = {t["name"] for t in data["result"]["tools"]}
expected = {"add_memories", "search_memory", "list_memories",
"delete_memories", "delete_all_memories"}
assert expected.issubset(tool_names), f"Missing tools: {expected - tool_names}"
@pytest.mark.asyncio
async def test_tools_list_has_descriptions(self, client):
"""Every tool returned by tools/list should have a non-empty description."""
await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
resp = await client.post(
"/mcp/testclient/http/user1",
json=_jsonrpc("tools/list", req_id=2),
headers=MCP_HEADERS,
)
for tool in resp.json()["result"]["tools"]:
assert tool.get("description"), f"Tool {tool['name']} has no description"
@pytest.mark.asyncio
async def test_tools_list_has_input_schemas(self, client):
"""Every tool should declare an inputSchema."""
await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
resp = await client.post(
"/mcp/testclient/http/user1",
json=_jsonrpc("tools/list", req_id=2),
headers=MCP_HEADERS,
)
for tool in resp.json()["result"]["tools"]:
assert "inputSchema" in tool, f"Tool {tool['name']} missing inputSchema"
@pytest.mark.asyncio
async def test_call_unknown_tool_returns_error(self, client):
"""Calling a non-existent tool should return a JSON-RPC error."""
await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
resp = await client.post(
"/mcp/testclient/http/user1",
json=_jsonrpc("tools/call", {"name": "no_such_tool", "arguments": {}}, req_id=2),
headers=MCP_HEADERS,
)
assert resp.status_code == 200
data = resp.json()
assert "error" in data or (
"result" in data and data["result"].get("isError")
)
@pytest.mark.asyncio
async def test_unknown_jsonrpc_method(self, client):
"""An unknown JSON-RPC method should return an error."""
resp = await client.post(
"/mcp/testclient/http/user1",
json=_jsonrpc("nonexistent/method"),
headers=MCP_HEADERS,
)
assert resp.status_code in (200, 400)
@pytest.mark.asyncio
async def test_response_content_type_is_json(self, client):
"""Responses should have Content-Type: application/json."""
resp = await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
ct = resp.headers.get("content-type", "")
assert "application/json" in ct
# ---------------------------------------------------------------------------
# Streamable HTTP — context variable handling
# ---------------------------------------------------------------------------
class TestStreamableHTTPContext:
"""Verify that user_id and client_name context variables are set correctly."""
@pytest.mark.asyncio
async def test_context_vars_set_during_tool_call(self, client):
"""Context vars should reflect the path parameters during tool execution."""
captured = {}
@mcp.tool(name="__test_ctx", description="test only")
async def _capture(query: str = "") -> str:
captured["user_id"] = user_id_var.get(None)
captured["client_name"] = client_name_var.get(None)
return "ok"
try:
await client.post(
"/mcp/my-app/http/alice",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
resp = await client.post(
"/mcp/my-app/http/alice",
json=_jsonrpc("tools/call", {"name": "__test_ctx", "arguments": {}}, req_id=2),
headers=MCP_HEADERS,
)
assert resp.status_code == 200
assert captured.get("user_id") == "alice"
assert captured.get("client_name") == "my-app"
finally:
mcp._tool_manager._tools.pop("__test_ctx", None)
@pytest.mark.asyncio
async def test_different_users_are_isolated(self, client):
"""Sequential requests with different user_ids must not leak state."""
results = []
@mcp.tool(name="__test_uid_iso", description="test only")
async def _capture_uid(query: str = "") -> str:
results.append(user_id_var.get(None))
return "ok"
try:
for uid in ("userA", "userB", "userC"):
await client.post(
f"/mcp/app1/http/{uid}",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
await client.post(
f"/mcp/app1/http/{uid}",
json=_jsonrpc("tools/call", {"name": "__test_uid_iso", "arguments": {}}, req_id=2),
headers=MCP_HEADERS,
)
assert results == ["userA", "userB", "userC"]
finally:
mcp._tool_manager._tools.pop("__test_uid_iso", None)
@pytest.mark.asyncio
async def test_different_clients_are_isolated(self, client):
"""Sequential requests with different client_names must not leak state."""
results = []
@mcp.tool(name="__test_cn_iso", description="test only")
async def _capture_cn(query: str = "") -> str:
results.append(client_name_var.get(None))
return "ok"
try:
for cn in ("cursor", "windsurf", "claude"):
await client.post(
f"/mcp/{cn}/http/user1",
json=_initialize_payload(),
headers=MCP_HEADERS,
)
await client.post(
f"/mcp/{cn}/http/user1",
json=_jsonrpc("tools/call", {"name": "__test_cn_iso", "arguments": {}}, req_id=2),
headers=MCP_HEADERS,
)
assert results == ["cursor", "windsurf", "claude"]
finally:
mcp._tool_manager._tools.pop("__test_cn_iso", None)
# ---------------------------------------------------------------------------
# Streamable HTTP — response correctness
# ---------------------------------------------------------------------------
class TestStreamableHTTPResponses:
"""Verify that captured responses are returned correctly to the caller."""
@pytest.mark.asyncio
async def test_error_status_codes_are_preserved(self, client):
"""Transport error codes (e.g. 406) must be forwarded, not masked as 200."""
resp = await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(),
)
assert resp.status_code == 406
@pytest.mark.asyncio
async def test_delete_status_code_preserved(self, client):
"""DELETE 405 from stateless transport must not be masked."""
resp = await client.delete(
"/mcp/testclient/http/user1",
headers=MCP_HEADERS,
)
assert resp.status_code == 405
@pytest.mark.asyncio
async def test_multiple_sequential_requests(self, client):
"""Multiple requests in sequence should each get independent responses."""
for i in range(5):
resp = await client.post(
"/mcp/testclient/http/user1",
json=_initialize_payload(req_id=i + 1),
headers=MCP_HEADERS,
)
assert resp.status_code == 200
data = resp.json()
assert data["id"] == i + 1
assert "result" in data
@pytest.mark.asyncio
async def test_wrong_content_type_returns_error(self, client):
"""POST with wrong Content-Type should return an error status."""
resp = await client.post(
"/mcp/testclient/http/user1",
content=b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}',
headers={**MCP_HEADERS, "Content-Type": "text/plain"},
)
assert resp.status_code in (400, 415)
# ---------------------------------------------------------------------------
# Route registration
# ---------------------------------------------------------------------------
class TestRouteRegistration:
"""Verify all expected routes are registered in the router."""
def test_sse_route_is_registered(self, test_app):
routes = [r.path for r in test_app.routes if hasattr(r, "path")]
assert "/mcp/{client_name}/sse/{user_id}" in routes
def test_sse_post_messages_route_is_registered(self, test_app):
routes = [r.path for r in test_app.routes if hasattr(r, "path")]
assert "/mcp/messages/" in routes or "/mcp/{client_name}/sse/{user_id}/messages/" in routes
def test_streamable_http_route_is_registered(self, test_app):
routes = [r.path for r in test_app.routes if hasattr(r, "path")]
assert "/mcp/{client_name}/http/{user_id}" in routes