tests/test_http_sessions.py

"""Verification fixes 2026-09-22 (F3), SEC-1: one client can no longer lock
every other client out of the HTTP transport.

With mcp 1.30.0 the session manager CanLex built had a global cap of 10,000
open sessions and no per-client limit, so one unauthenticated client could
fill it in about two minutes; every other client's `initialize` then got
503 'Too many open sessions' until the flood's sessions idled out (1800 s).
The critic's live probe, with the cap scaled from 10,000 to 5: one client's
six initializes [200, 200, 200, 200, 200, 503], another client's 503.

Now server.session_manager_with_timeout passes max_sessions=None (the 1.27.2
behaviour, bounded by the idle timeout) and server.http_app -- the app
__main__ serves -- adds canlex/ratelimit.py's SessionOpenLimit, on by
default. These tests drive the real stack (FastMCP's streamable-HTTP app,
the session manager, _Utf8Body, the cap) in-process with Starlette's
TestClient, the cap scaled down through CANLEX_SESSION_OPEN_LIMIT. Clients
are told apart by X-Forwarded-For, as Hugging Face's ingress supplies it.

    python -m unittest tests.test_http_sessions
"""
import io
import logging
import unittest
from contextlib import redirect_stderr

from mcp.server.fastmcp import FastMCP
from starlette.testclient import TestClient

from canlex import ratelimit, server

INIT = {"jsonrpc": "2.0", "id": 1, "method": "initialize",
        "params": {"protocolVersion": "2025-03-26", "capabilities": {},
                   "clientInfo": {"name": "t", "version": "1"}}}
ACCEPT = {"Accept": "application/json, text/event-stream"}
FLOODER, OFFICER = "203.0.113.7", "198.51.100.9"


def toy():
    srv = FastMCP("t")

    @srv.tool()
    def echo(word: str) -> str:
        return word
    return srv


class HttpSessionTests(unittest.TestCase):
    def setUp(self):
        for name in ("mcp", "httpx"):
            log = logging.getLogger(name)
            self.addCleanup(log.setLevel, log.level)
            log.setLevel(logging.WARNING)
        # _Utf8Body prints each new client tag once: keep the set as it was
        seen = set(server._seen_clients)
        self.addCleanup(lambda: (server._seen_clients.clear(),
                                 server._seen_clients.update(seen)))
        quiet = redirect_stderr(io.StringIO())
        quiet.__enter__()
        self.addCleanup(quiet.__exit__, None, None, None)

    def serve(self, srv, environ):
        err = io.StringIO()
        with redirect_stderr(err):
            app = server.http_app(srv, environ)
        client = TestClient(app, base_url="http://127.0.0.1:8000")
        client.__enter__()                       # runs the session manager's lifespan
        self.addCleanup(client.__exit__, None, None, None)
        return client, err.getvalue()

    @staticmethod
    def initialize(client, who):
        return client.post("/mcp", json=INIT, headers=dict(ACCEPT, **{"X-Forwarded-For": who}))

    def test_the_served_stack_has_no_global_cap_and_the_per_client_cap_on(self):
        srv = toy()
        _, err = self.serve(srv, {})
        self.assertIsNone(server.SESSION_MAX_OPEN)
        self.assertIsNone(srv._session_manager.max_sessions)
        self.assertEqual(srv._session_manager.session_idle_timeout, server.SESSION_IDLE_SECONDS)
        self.assertIn("no global session cap", err)
        self.assertIn(f"at most {ratelimit.SESSION_OPEN_LIMIT} new sessions per client", err)

    def test_a_flood_from_one_client_leaves_another_able_to_initialize(self):
        srv = toy()
        client, _ = self.serve(srv, {"CANLEX_SESSION_OPEN_LIMIT": "3"})
        flood = [self.initialize(client, FLOODER) for _ in range(8)]
        self.assertEqual([r.status_code for r in flood], [200] * 3 + [429] * 5)
        self.assertIn("Too many new sessions", flood[-1].text)
        officer = self.initialize(client, OFFICER)
        self.assertEqual(officer.status_code, 200)
        self.assertIn('"protocolVersion"', officer.text)
        self.assertEqual(len(srv._session_manager._server_instances), 4)
        # the flooder's own open sessions keep working: only new sessions are capped
        sid = flood[0].headers["mcp-session-id"]
        listed = client.post("/mcp", json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
                             headers=dict(ACCEPT, **{"X-Forwarded-For": FLOODER,
                                                     "mcp-session-id": sid,
                                                     "mcp-protocol-version": "2025-03-26"}))
        self.assertEqual(listed.status_code, 200)
        self.assertIn('"echo"', listed.text)

    def test_even_under_a_global_cap_one_client_cannot_fill_it(self):
        """The critic's scaled stand-in (a global cap of 5) with the per-client
        cap at 3: before, the flood filled the table and the officer got 503."""
        srv = toy()
        client, _ = self.serve(srv, {"CANLEX_SESSION_OPEN_LIMIT": "3"})
        srv._session_manager.max_sessions = 5
        self.assertEqual([self.initialize(client, FLOODER).status_code for _ in range(6)],
                         [200] * 3 + [429] * 3)
        self.assertEqual(self.initialize(client, OFFICER).status_code, 200)

    def test_turned_off_the_flood_reaches_the_manager(self):
        srv = toy()
        client, err = self.serve(srv, {"CANLEX_SESSION_OPEN_LIMIT": "0"})
        self.assertIn("no per-client cap on new sessions", err)
        self.assertEqual({self.initialize(client, FLOODER).status_code for _ in range(5)}, {200})


if __name__ == "__main__":
    unittest.main()