11"""Client-side scaffold for story examples.
22
3- A story's ``client.py`` imports only from here. The ``Connect`` factory and
4- ``run_client`` ride the locked ``Client(transport, mode=...)`` surface; the one
5- volatile line is the stdio wrap (marked inline).
3+ A story's ``client.py`` imports ``Target`` (or ``TargetFactory``) for its ``main``
4+ signature and calls ``run_client(main)`` from ``__main__``. The story owns the
5+ ``Client(target, mode=...)`` construction; this module only decides WHICH target
6+ ``__main__`` hands it.
67"""
78
89from __future__ import annotations
910
1011import sys
1112import traceback
12- from collections .abc import AsyncIterator , Awaitable , Callable
13- from contextlib import AbstractAsyncContextManager , asynccontextmanager
13+ from collections .abc import Awaitable , Callable
1414from pathlib import Path
15- from typing import Any , Protocol
15+ from typing import Any , TypeAlias
16+ from urllib .parse import urlsplit
1617
1718import anyio
1819import httpx
1920
2021from mcp import StdioServerParameters , stdio_client
21- from mcp .client import Client
22+ from mcp .client import Transport
23+ from mcp .client .streamable_http import streamable_http_client
24+ from mcp .server import Server
25+ from mcp .server .mcpserver import MCPServer
2226from mcp .shared .version import LATEST_MODERN_VERSION
2327
24- Scenario = Callable [[ Client ], Awaitable [ None ]]
25- ScenarioWithConnect = Callable [[ Client , "Connect" ], Awaitable [ None ]]
26- AuthBuilder = Callable [[ httpx . AsyncClient ], httpx . Auth ]
27- """Builds an ``httpx.Auth`` bound to the in-process HTTP client (auth-story harness seam)."""
28+ if sys . version_info >= ( 3 , 11 ):
29+ import tomllib
30+ else :
31+ import tomli as tomllib
2832
33+ Target : TypeAlias = "Server[Any] | MCPServer | Transport | str"
34+ """Anything ``Client(...)`` accepts: an in-process server, a ``Transport``, or an HTTP URL."""
2935
30- class Connect (Protocol ):
31- """A factory yielding a connected ``Client``; accepts the same kwargs as ``Client``.
32-
33- ``auth`` is the HTTP-only escape hatch for auth stories: when given, the factory
34- builds a fresh ``httpx.AsyncClient`` against the same app, applies ``auth(http)``
35- to it, and wraps the result in ``streamable_http_client`` before entering ``Client``.
36- """
36+ TargetFactory = Callable [[], Target ]
37+ """Yields a FRESH target against the same server/app on every call (``multi_connection`` stories)."""
3738
38- def __call__ (self , * , auth : AuthBuilder | None = None , ** client_kw : Any ) -> AbstractAsyncContextManager [Client ]: ...
39+ AuthBuilder = Callable [[httpx .AsyncClient ], httpx .Auth ]
40+ """Builds an ``httpx.Auth`` bound to the in-process HTTP client (auth-story harness seam)."""
3941
4042
4143def argv_after (flag : str , * , default : str | None = None ) -> str :
@@ -48,67 +50,84 @@ def argv_after(flag: str, *, default: str | None = None) -> str:
4850 return default
4951
5052
51- def connect_from_args (file : str ) -> Connect :
52- """Build a ``Connect `` targeting the sibling server over the argv-selected transport.
53+ def target_from_args (file : str ) -> TargetFactory :
54+ """Build a ``TargetFactory `` for the sibling server over the argv-selected transport.
5355
54- ``--http <url>`` connects over streamable HTTP; ``--stdio`` (the default) spawns the
55- sibling ``server.py`` as a subprocess. ``--server <stem>`` selects ``<stem>.py``
56- (e.g. ``server_lowlevel``). ``--legacy`` pins the handshake era; otherwise the
57- modern era is used. ``file`` is the caller's ``__file__``.
56+ ``--http <url>`` targets that streamable-HTTP URL; ``--stdio`` (the default) spawns
57+ the sibling ``server.py`` as a fresh subprocess on each call. ``--server <stem>``
58+ selects ``<stem>.py`` (e.g. ``server_lowlevel``). ``file`` is the caller's ``__file__``.
5859 """
59- here = Path (file ).parent
60- server_stem = argv_after ("--server" , default = "server" )
61- # Never rely on the SDK's mode= default — be explicit. stdio is legacy-only until
62- # the SDK's stdio entry can negotiate the era; the modern arm is --http only for now.
6360 if "--http" in sys .argv :
64- mode = "legacy" if "--legacy" in sys .argv else LATEST_MODERN_VERSION
65- else :
66- mode = "legacy" # stdio gains a modern arm once serve_stdio() lands
67-
68- @asynccontextmanager
69- async def _connect (* , auth : AuthBuilder | None = None , ** client_kw : Any ) -> AsyncIterator [Client ]:
70- assert auth is None , "auth= via connect_from_args is not wired; auth stories own their __main__"
71- client_kw .setdefault ("mode" , mode )
72- target : Any
73- if "--http" in sys .argv :
74- target = argv_after ("--http" )
75- else :
76- params = StdioServerParameters (command = sys .executable , args = [str (here / f"{ server_stem } )])
77- target = stdio_client (params ) # becomes Client(params) once that overload lands
78- async with Client (target , ** client_kw ) as client :
79- yield client
80-
81- return _connect
82-
83-
84- def run_client (
85- scenario : Scenario | ScenarioWithConnect ,
86- * ,
87- connect : Connect ,
88- needs_connect : bool = False ,
89- ** client_kw : Any ,
90- ) -> None :
61+ url = argv_after ("--http" )
62+ return lambda : url
63+ # stdio is legacy-only until serve_stdio() lands; the modern arm is --http only for now.
64+ server = Path (file ).parent / f"{ argv_after ('--server' , default = 'server' )}
65+ params = StdioServerParameters (command = sys .executable , args = [str (server )])
66+ return lambda : stdio_client (params ) # becomes Client(params) once that overload lands
67+
68+
69+ def _story_cfg (name : str ) -> dict [str , Any ]:
70+ """The manifest entry for the story ``name`` with ``[defaults]`` applied."""
71+ manifest : dict [str , Any ] = tomllib .loads ((Path (__file__ ).parent / "manifest.toml" ).read_text ())
72+ return manifest ["defaults" ] | manifest ["story" ].get (name , {})
73+
74+
75+ def _authed_targets (url : str , http : httpx .AsyncClient ) -> TargetFactory :
76+ """Fresh streamable-HTTP transports over an already-authed ``httpx`` client."""
77+ return lambda : streamable_http_client (url , http_client = http )
78+
79+
80+ def run_client (main : Callable [..., Awaitable [None ]]) -> None :
9181 """Entry point for ``if __name__ == "__main__"`` in every ``client.py``.
9282
93- Runs ``scenario`` inside a connected client; prints ``OK:``/``FAIL:`` to stderr and
94- exits 0/1. ``needs_connect=True`` passes ``connect`` as the second argument so the
95- scenario can open additional clients.
83+ Builds the argv-selected target(s) for the story that defines ``main``, picks the
84+ era from argv, and calls ``main`` with an explicit ``mode=``. If the story module
85+ exports ``build_auth``, the ``--http`` target is routed through an ``httpx.AsyncClient``
86+ that carries the returned ``httpx.Auth``. Prints ``OK:``/``FAIL:`` to stderr, exits 0/1.
9687 """
97- file = getattr (scenario , "__globals__" , {}).get ("__file__" , "<unknown>" )
88+ globals_ = getattr (main , "__globals__" , {})
89+ file = str (globals_ .get ("__file__" , "<unknown>" ))
9890 name = Path (file ).parent .name
91+ cfg = _story_cfg (name )
92+ targets = target_from_args (file )
93+ build_auth : AuthBuilder | None = globals_ .get ("build_auth" )
9994 transport = "http" if "--http" in sys .argv else "stdio"
95+ # Never rely on the SDK's mode= default — be explicit. stdio is legacy-only until
96+ # the SDK's stdio entry can negotiate the era, so only --http gets a modern arm.
10097 era = "modern" if transport == "http" and "--legacy" not in sys .argv else "legacy"
101-
102- async def _main () -> None :
103- with anyio .fail_after (30 ):
104- async with connect (** client_kw ) as client :
105- if needs_connect :
106- await scenario (client , connect ) # type: ignore[call-arg]
98+ if cfg ["era" ] == "dual-in-body" :
99+ # The story pins its connection modes inside ``main`` itself, so hand it the
100+ # real-user "auto" default and let those in-body pins decide. A hard version pin
101+ # here would skip the discover probe and leave ``server_info`` blank.
102+ era = "in-body"
103+ mode = {"modern" : LATEST_MODERN_VERSION , "legacy" : "legacy" , "in-body" : "auto" }[era ]
104+
105+ async def _run () -> None :
106+ with anyio .fail_after (cfg ["timeout_s" ]):
107+ if not cfg ["needs_http" ] and (build_auth is None or transport != "http" ):
108+ await main (targets if cfg ["multi_connection" ] else targets (), mode = mode )
109+ return
110+ # Auth and needs_http stories want the raw httpx client underneath the transport:
111+ # build_auth threads an httpx.Auth onto it (Client(url, auth=...) doesn't exist
112+ # yet), and needs_http stories assert on raw responses, so root the client at the
113+ # server origin and relative paths like "/mcp" resolve.
114+ if transport != "http" :
115+ raise SystemExit (f"{ name } )
116+ url = argv_after ("--http" )
117+ parts = urlsplit (url )
118+ async with httpx .AsyncClient (base_url = f"{ parts .scheme } { parts .netloc } ) as http :
119+ make = targets
120+ if build_auth is not None :
121+ http .auth = build_auth (http )
122+ make = _authed_targets (url , http )
123+ target : Any = make if cfg ["multi_connection" ] else make ()
124+ if cfg ["needs_http" ]:
125+ await main (target , mode = mode , http = http )
107126 else :
108- await scenario ( client ) # type: ignore[call-arg]
127+ await main ( target , mode = mode )
109128
110129 try :
111- anyio .run (_main )
130+ anyio .run (_run )
112131 except Exception :
113132 print (f"FAIL: { name } { transport } { era } , file = sys .stderr )
114133 traceback .print_exc ()
0 commit comments