Files
forge/apps/api/tests/test_mcp_streamable.py
nihalashetty ae67bff5a3 feat: deep-agent canvas, live observability, and multi-environment tooling
Self-hosted platform for building, testing, and shipping LangChain/LangGraph agents. Deep-agent sub-agents on the canvas, a live tracing/observability timeline, auto-provisioned built-in tools with import/export, per-environment tool variables, streamed evaluations, and per-user auth token forwarding.
2026-07-28 01:49:19 +05:30

125 lines
5.6 KiB
Python

"""Forge-as-an-MCP-server over the Streamable-HTTP transport.
The headline test drives the endpoint with the REAL `mcp` SDK client (the same protocol Claude
Desktop / Cursor / VS Code speak), routed at the in-process ASGI app — proving a native client can
initialize, list, and call tools with no `mcp-remote` proxy bridge. The rest pin the HTTP-level
contract: a POST that accepts SSE gets a `text/event-stream` reply, a plain-JSON POST still gets the
legacy JSON response, and the per-project auth applies to the streaming transport too.
"""
from __future__ import annotations
import json
import httpx
from forge.db.base import SessionLocal
from forge.main import create_app
from forge.models import Project, Tool
async def _seed_project_with_tool(tenant="t_stream", slug="mcp-stream", config=None) -> str:
from forge.services.tool_sets import ToolSetService
async with SessionLocal() as s:
proj = Project(tenant_id=tenant, name="Stream Proj", slug=slug, config=config or {})
s.add(proj)
await s.flush()
tool = Tool(tenant_id=tenant, project_id=proj.id, name="calculator", kind="builtin",
config={"builtin": "calculator", "description": "Evaluate arithmetic."})
s.add(tool)
await s.commit()
await s.refresh(proj)
await s.refresh(tool)
await ToolSetService.create(s, tenant, proj.id, name="General", tool_ids=[tool.id])
return proj.id
def _asgi_httpx_factory(app):
"""An httpx client factory (the shape `streamablehttp_client` expects) that routes the MCP
client's real HTTP traffic through the in-process ASGI app instead of the network."""
def make(*, headers=None, timeout=None, auth=None, **_):
return httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://test",
headers=headers, timeout=timeout, auth=auth,
)
return make
async def test_streamable_real_mcp_client_end_to_end():
"""A real MCP SDK client initializes, lists, and calls a tool over Streamable HTTP."""
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamablehttp_client
pid = await _seed_project_with_tool()
app = create_app()
url = f"http://test/v1/mcp/{pid}"
async with streamablehttp_client(url, httpx_client_factory=_asgi_httpx_factory(app)) as (read, write, _sid):
async with ClientSession(read, write) as session:
init = await session.initialize()
assert init.serverInfo.name.startswith("forge-")
tools = await session.list_tools()
assert "calculator" in [t.name for t in tools.tools]
res = await session.call_tool("calculator", {"expression": "6*7"})
assert res.isError is False
assert "42" in res.content[0].text
def _parse_sse_json(body: str) -> dict:
"""Pull the JSON-RPC payload out of a single-message `text/event-stream` response."""
for line in body.splitlines():
if line.startswith("data:"):
return json.loads(line[len("data:"):].strip())
raise AssertionError(f"no SSE data frame in response:\n{body}")
async def test_streamable_post_negotiates_sse():
"""A POST that accepts text/event-stream is answered with an SSE-framed JSON-RPC reply."""
pid = await _seed_project_with_tool(tenant="t_stream2", slug="mcp-stream2")
app = create_app()
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
r = await c.post(
f"/v1/mcp/{pid}",
headers={"Accept": "application/json, text/event-stream"},
json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {
"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "t", "version": "1"},
}},
)
assert r.status_code == 200
assert "text/event-stream" in r.headers["content-type"]
payload = _parse_sse_json(r.text)
assert payload["result"]["serverInfo"]["name"].startswith("forge-")
async def test_plain_json_post_still_uses_legacy_json():
"""A POST WITHOUT an SSE Accept stays on the legacy request/response path (application/json)."""
pid = await _seed_project_with_tool(tenant="t_stream3", slug="mcp-stream3")
app = create_app()
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
r = await c.post(f"/v1/mcp/{pid}", json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
assert r.status_code == 200
assert "application/json" in r.headers["content-type"]
assert "calculator" in [t["name"] for t in r.json()["result"]["tools"]]
async def test_streamable_transport_enforces_project_key():
"""The per-project mcp_api_key gates the streaming transport, not just the legacy path."""
pid = await _seed_project_with_tool(tenant="t_stream4", slug="mcp-stream4", config={"mcp_api_key": "sk-stream"})
app = create_app()
sse = {"Accept": "application/json, text/event-stream"}
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
init = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {
"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "t", "version": "1"},
}}
# no key -> 401 even on the streamable transport
r = await c.post(f"/v1/mcp/{pid}", headers=sse, json=init)
assert r.status_code == 401
# correct key -> streamed 200
r = await c.post(f"/v1/mcp/{pid}", headers={**sse, "Authorization": "Bearer sk-stream"}, json=init)
assert r.status_code == 200 and "text/event-stream" in r.headers["content-type"]