88
99from __future__ import annotations
1010
11+ import socket
1112import sys
1213import traceback
13- from collections .abc import Awaitable , Callable
14+ from collections .abc import AsyncIterator , Awaitable , Callable
15+ from contextlib import AsyncExitStack , asynccontextmanager
1416from pathlib import Path
1517from typing import Any , TypeAlias
1618from urllib .parse import urlsplit
@@ -50,22 +52,75 @@ def argv_after(flag: str, *, default: str | None = None) -> str:
5052 return default
5153
5254
53- def target_from_args (file : str ) -> TargetFactory :
54- """Build a ``TargetFactory`` for the sibling server over the argv-selected transport .
55+ def target_from_args (file : str , url : str | None ) -> TargetFactory :
56+ """Build a ``TargetFactory`` for the sibling server of the ``client.py`` at ``file`` .
5557
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__``.
58+ ``url`` (already resolved by ``run_client``) targets that streamable-HTTP endpoint; ``None``
59+ spawns ``<stem>.py`` over stdio per call, ``<stem>`` from ``--server`` (default ``server``).
5960 """
60- if "--http" in sys .argv :
61- url = argv_after ("--http" )
61+ if url is not None :
6262 return lambda : url
6363 # stdio is legacy-only until serve_stdio() lands; the modern arm is --http only for now.
6464 server = Path (file ).parent / f"{ argv_after ('--server' , default = 'server' )}
6565 params = StdioServerParameters (command = sys .executable , args = [str (server )])
6666 return lambda : stdio_client (params ) # becomes Client(params) once that overload lands
6767
6868
69+ def _explicit_http_url () -> str | None :
70+ """The URL token after ``--http``, or ``None`` when the flag stands alone (self-host)."""
71+ rest = sys .argv [sys .argv .index ("--http" ) + 1 :]
72+ return rest [0 ] if rest and not rest [0 ].startswith ("-" ) else None
73+
74+
75+ def _free_port () -> int :
76+ """An OS-assigned free TCP port, released for the server subprocess to re-bind."""
77+ with socket .socket () as sock :
78+ sock .bind (("127.0.0.1" , 0 ))
79+ return sock .getsockname ()[1 ]
80+
81+
82+ async def _accepting (port : int ) -> bool :
83+ """Whether something accepts a TCP connect on ``127.0.0.1:port`` right now."""
84+ try :
85+ stream = await anyio .connect_tcp ("127.0.0.1" , port )
86+ except OSError :
87+ return False
88+ await stream .aclose ()
89+ return True
90+
91+
92+ @asynccontextmanager
93+ async def _self_hosted (name : str , cfg : dict [str , Any ]) -> AsyncIterator [str ]:
94+ """Serve the story's sibling server from a subprocess on a port this process owns; yield its URL.
95+
96+ Readiness is the first accepted TCP connect (bounded by ``run_client``'s
97+ ``anyio.fail_after``); exiting terminates the subprocess. Nothing to background or kill.
98+ A subprocess that dies before serving, or a ``fixed_port`` someone else already holds,
99+ is a loud ``SystemExit`` rather than a hang or a run against the wrong server.
100+ """
101+ port : int = cfg ["fixed_port" ] or _free_port ()
102+ if cfg ["fixed_port" ] and await _accepting (port ):
103+ # The readiness probe below can't tell our child from a server already on the
104+ # story's pinned port, so a foreign listener would be tested in its place.
105+ raise SystemExit (
106+ f"{ name } { port }
107+ f"stop it, or connect to it with --http <url>"
108+ )
109+ module = f"stories.{ name } { argv_after ('--server' , default = 'server' )}
110+ serve = ["--http" ] if cfg ["server_export" ] == "factory" else []
111+ argv = [sys .executable , "-m" , module , * serve , "--port" , str (port )]
112+ async with await anyio .open_process (argv , stdout = None , stderr = None ) as server :
113+ try :
114+ while server .returncode is None and not await _accepting (port ):
115+ await anyio .sleep (0.05 )
116+ if server .returncode is not None :
117+ raise SystemExit (f"{ module } { server .returncode } { port } )
118+ yield f"http://127.0.0.1:{ port } { cfg ['mcp_path' ]}
119+ finally :
120+ if server .returncode is None :
121+ server .terminate ()
122+
123+
69124def _story_cfg (name : str ) -> dict [str , Any ]:
70125 """The manifest entry for the story ``name`` with ``[defaults]`` applied."""
71126 manifest : dict [str , Any ] = tomllib .loads ((Path (__file__ ).parent / "manifest.toml" ).read_text ())
@@ -80,23 +135,24 @@ def _authed_targets(url: str, http: httpx.AsyncClient) -> TargetFactory:
80135def run_client (main : Callable [..., Awaitable [None ]]) -> None :
81136 """Entry point for ``if __name__ == "__main__"`` in every ``client.py``.
82137
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.
138+ Resolves the argv target — stdio (the default), ``--http <url>`` for a server you run, or
139+ bare ``--http`` to self-host the sibling server in a subprocess it owns — and calls ``main``
140+ with an explicit ``mode=``. A ``build_auth`` export auths the HTTP target. ``OK``/``FAIL``, exit 0/1.
87141 """
88142 globals_ = getattr (main , "__globals__" , {})
89143 file = str (globals_ .get ("__file__" , "<unknown>" ))
90144 name = Path (file ).parent .name
91145 cfg = _story_cfg (name )
92- targets = target_from_args (file )
93146 build_auth : AuthBuilder | None = globals_ .get ("build_auth" )
94147 transport = "http" if "--http" in sys .argv else "stdio"
95148 if cfg ["server_export" ] == "app" and transport != "http" :
96149 raise SystemExit (
97- f"{ name } start its server, then run :\n "
98- f" python -m stories.{ name } http://127.0.0.1:8000 { cfg [ 'mcp_path' ] }
150+ f"{ name } self-host it over HTTP :\n "
151+ f" python -m stories.{ name }
99152 )
153+ if cfg ["needs_http" ] and transport != "http" :
154+ raise SystemExit (f"{ name } )
155+ explicit_url = _explicit_http_url () if transport == "http" else None
100156 # The era is an axis of the story matrix, so ``mode=`` is always passed explicitly
101157 # even though it often matches the ``Client`` default of "auto". stdio is legacy-only
102158 # until the SDK's stdio entry can negotiate the era, so only --http gets a modern arm.
@@ -110,18 +166,21 @@ def run_client(main: Callable[..., Awaitable[None]]) -> None:
110166
111167 async def _run () -> None :
112168 with anyio .fail_after (cfg ["timeout_s" ]):
113- if not cfg ["needs_http" ] and (build_auth is None or transport != "http" ):
114- await main (targets if cfg ["multi_connection" ] else targets (), mode = mode )
115- return
116- # Auth and needs_http stories want the raw httpx client underneath the transport:
117- # build_auth threads an httpx.Auth onto it (Client(url, auth=...) doesn't exist
118- # yet), and needs_http stories assert on raw responses, so root the client at the
119- # server origin and relative paths like "/mcp" resolve.
120- if transport != "http" :
121- raise SystemExit (f"{ name } )
122- url = argv_after ("--http" )
123- parts = urlsplit (url )
124- async with httpx .AsyncClient (base_url = f"{ parts .scheme } { parts .netloc } ) as http :
169+ async with AsyncExitStack () as stack :
170+ url = explicit_url
171+ if transport == "http" and url is None :
172+ url = await stack .enter_async_context (_self_hosted (name , cfg ))
173+ targets = target_from_args (file , url )
174+ if url is None or (build_auth is None and not cfg ["needs_http" ]):
175+ await main (targets if cfg ["multi_connection" ] else targets (), mode = mode )
176+ return
177+ # Auth and needs_http stories want the raw httpx client underneath the transport:
178+ # build_auth threads an httpx.Auth onto it (Client(url, auth=...) doesn't exist
179+ # yet), and needs_http stories assert on raw responses, so root the client at the
180+ # server origin and relative paths like "/mcp" resolve.
181+ parts = urlsplit (url )
182+ base = f"{ parts .scheme } { parts .netloc }
183+ http = await stack .enter_async_context (httpx .AsyncClient (base_url = base ))
125184 make = targets
126185 if build_auth is not None :
127186 http .auth = build_auth (http )
0 commit comments