refactor: connect-first stream lifecycle for sse and streamable_http (#2292)
Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
This commit is contained in:
Max Isbey
GitHub
parent
abfb482246
commit
75a80b6f07
2 changed files with 161 additions and 153 deletions
|
|
@ -57,108 +57,101 @@ async def sse_client(
|
|||
write_stream: MemoryObjectSendStream[SessionMessage]
|
||||
write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
|
||||
|
||||
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
|
||||
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
|
||||
logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}")
|
||||
async with httpx_client_factory(
|
||||
headers=headers, auth=auth, timeout=httpx.Timeout(timeout, read=sse_read_timeout)
|
||||
) as client:
|
||||
async with aconnect_sse(client, "GET", url) as event_source:
|
||||
event_source.response.raise_for_status()
|
||||
logger.debug("SSE connection established")
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
try:
|
||||
logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}")
|
||||
async with httpx_client_factory(
|
||||
headers=headers, auth=auth, timeout=httpx.Timeout(timeout, read=sse_read_timeout)
|
||||
) as client:
|
||||
async with aconnect_sse(
|
||||
client,
|
||||
"GET",
|
||||
url,
|
||||
) as event_source:
|
||||
event_source.response.raise_for_status()
|
||||
logger.debug("SSE connection established")
|
||||
read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
|
||||
write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
|
||||
|
||||
async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
|
||||
try:
|
||||
async for sse in event_source.aiter_sse(): # pragma: no branch
|
||||
logger.debug(f"Received SSE event: {sse.event}")
|
||||
match sse.event:
|
||||
case "endpoint":
|
||||
endpoint_url = urljoin(url, sse.data)
|
||||
logger.debug(f"Received endpoint URL: {endpoint_url}")
|
||||
async def sse_reader(task_status: TaskStatus[str] = anyio.TASK_STATUS_IGNORED):
|
||||
try:
|
||||
async for sse in event_source.aiter_sse(): # pragma: no branch
|
||||
logger.debug(f"Received SSE event: {sse.event}")
|
||||
match sse.event:
|
||||
case "endpoint":
|
||||
endpoint_url = urljoin(url, sse.data)
|
||||
logger.debug(f"Received endpoint URL: {endpoint_url}")
|
||||
|
||||
url_parsed = urlparse(url)
|
||||
endpoint_parsed = urlparse(endpoint_url)
|
||||
if ( # pragma: no cover
|
||||
url_parsed.netloc != endpoint_parsed.netloc
|
||||
or url_parsed.scheme != endpoint_parsed.scheme
|
||||
):
|
||||
error_msg = ( # pragma: no cover
|
||||
f"Endpoint origin does not match connection origin: {endpoint_url}"
|
||||
)
|
||||
logger.error(error_msg) # pragma: no cover
|
||||
raise ValueError(error_msg) # pragma: no cover
|
||||
|
||||
if on_session_created:
|
||||
session_id = _extract_session_id_from_endpoint(endpoint_url)
|
||||
if session_id:
|
||||
on_session_created(session_id)
|
||||
|
||||
task_status.started(endpoint_url)
|
||||
|
||||
case "message":
|
||||
# Skip empty data (keep-alive pings)
|
||||
if not sse.data:
|
||||
continue
|
||||
try:
|
||||
message = types.jsonrpc_message_adapter.validate_json(
|
||||
sse.data, by_name=False
|
||||
)
|
||||
logger.debug(f"Received server message: {message}")
|
||||
except Exception as exc: # pragma: no cover
|
||||
logger.exception("Error parsing server message") # pragma: no cover
|
||||
await read_stream_writer.send(exc) # pragma: no cover
|
||||
continue # pragma: no cover
|
||||
|
||||
session_message = SessionMessage(message)
|
||||
await read_stream_writer.send(session_message)
|
||||
case _: # pragma: no cover
|
||||
logger.warning(f"Unknown SSE event: {sse.event}") # pragma: no cover
|
||||
except SSEError as sse_exc: # pragma: lax no cover
|
||||
logger.exception("Encountered SSE exception")
|
||||
raise sse_exc
|
||||
except Exception as exc: # pragma: lax no cover
|
||||
logger.exception("Error in sse_reader")
|
||||
await read_stream_writer.send(exc)
|
||||
finally:
|
||||
await read_stream_writer.aclose()
|
||||
|
||||
async def post_writer(endpoint_url: str):
|
||||
try:
|
||||
async with write_stream_reader:
|
||||
async for session_message in write_stream_reader:
|
||||
logger.debug(f"Sending client message: {session_message}")
|
||||
response = await client.post(
|
||||
endpoint_url,
|
||||
json=session_message.message.model_dump(
|
||||
by_alias=True,
|
||||
mode="json",
|
||||
exclude_unset=True,
|
||||
),
|
||||
url_parsed = urlparse(url)
|
||||
endpoint_parsed = urlparse(endpoint_url)
|
||||
if ( # pragma: no cover
|
||||
url_parsed.netloc != endpoint_parsed.netloc
|
||||
or url_parsed.scheme != endpoint_parsed.scheme
|
||||
):
|
||||
error_msg = ( # pragma: no cover
|
||||
f"Endpoint origin does not match connection origin: {endpoint_url}"
|
||||
)
|
||||
response.raise_for_status()
|
||||
logger.debug(f"Client message sent successfully: {response.status_code}")
|
||||
except Exception: # pragma: lax no cover
|
||||
logger.exception("Error in post_writer")
|
||||
finally:
|
||||
await write_stream.aclose()
|
||||
logger.error(error_msg) # pragma: no cover
|
||||
raise ValueError(error_msg) # pragma: no cover
|
||||
|
||||
endpoint_url = await tg.start(sse_reader)
|
||||
logger.debug(f"Starting post writer with endpoint URL: {endpoint_url}")
|
||||
tg.start_soon(post_writer, endpoint_url)
|
||||
if on_session_created:
|
||||
session_id = _extract_session_id_from_endpoint(endpoint_url)
|
||||
if session_id:
|
||||
on_session_created(session_id)
|
||||
|
||||
try:
|
||||
yield read_stream, write_stream
|
||||
finally:
|
||||
tg.cancel_scope.cancel()
|
||||
finally:
|
||||
await read_stream_writer.aclose()
|
||||
await write_stream.aclose()
|
||||
await read_stream.aclose()
|
||||
await write_stream_reader.aclose()
|
||||
task_status.started(endpoint_url)
|
||||
|
||||
case "message":
|
||||
# Skip empty data (keep-alive pings)
|
||||
if not sse.data:
|
||||
continue
|
||||
try:
|
||||
message = types.jsonrpc_message_adapter.validate_json(sse.data, by_name=False)
|
||||
logger.debug(f"Received server message: {message}")
|
||||
except Exception as exc: # pragma: no cover
|
||||
logger.exception("Error parsing server message") # pragma: no cover
|
||||
await read_stream_writer.send(exc) # pragma: no cover
|
||||
continue # pragma: no cover
|
||||
|
||||
session_message = SessionMessage(message)
|
||||
await read_stream_writer.send(session_message)
|
||||
case _: # pragma: no cover
|
||||
logger.warning(f"Unknown SSE event: {sse.event}") # pragma: no cover
|
||||
except SSEError as sse_exc: # pragma: lax no cover
|
||||
logger.exception("Encountered SSE exception")
|
||||
raise sse_exc
|
||||
except Exception as exc: # pragma: lax no cover
|
||||
logger.exception("Error in sse_reader")
|
||||
await read_stream_writer.send(exc)
|
||||
finally:
|
||||
await read_stream_writer.aclose()
|
||||
|
||||
async def post_writer(endpoint_url: str):
|
||||
try:
|
||||
async with write_stream_reader, write_stream:
|
||||
async for session_message in write_stream_reader:
|
||||
logger.debug(f"Sending client message: {session_message}")
|
||||
response = await client.post(
|
||||
endpoint_url,
|
||||
json=session_message.message.model_dump(
|
||||
by_alias=True,
|
||||
mode="json",
|
||||
exclude_unset=True,
|
||||
),
|
||||
)
|
||||
response.raise_for_status()
|
||||
logger.debug(f"Client message sent successfully: {response.status_code}")
|
||||
except Exception: # pragma: lax no cover
|
||||
logger.exception("Error in post_writer")
|
||||
|
||||
# On Python 3.14, coverage.py reports a phantom branch arc on this
|
||||
# line (->yield) when nested two async-with levels deep. The branch
|
||||
# is the unreachable "did __aexit__ suppress?" arm for memory streams.
|
||||
async with ( # pragma: no branch
|
||||
read_stream_writer,
|
||||
read_stream,
|
||||
write_stream,
|
||||
write_stream_reader,
|
||||
anyio.create_task_group() as tg,
|
||||
):
|
||||
endpoint_url = await tg.start(sse_reader)
|
||||
logger.debug(f"Starting post writer with endpoint URL: {endpoint_url}")
|
||||
tg.start_soon(post_writer, endpoint_url)
|
||||
|
||||
yield read_stream, write_stream
|
||||
tg.cancel_scope.cancel()
|
||||
|
|
|
|||
|
|
@ -440,7 +440,7 @@ class StreamableHTTPTransport:
|
|||
) -> None:
|
||||
"""Handle writing requests to the server."""
|
||||
try:
|
||||
async with write_stream_reader:
|
||||
async with write_stream_reader, read_stream_writer, write_stream:
|
||||
async for session_message in write_stream_reader:
|
||||
message = session_message.message
|
||||
metadata = (
|
||||
|
|
@ -480,9 +480,6 @@ class StreamableHTTPTransport:
|
|||
|
||||
except Exception: # pragma: lax no cover
|
||||
logger.exception("Error in post_writer")
|
||||
finally:
|
||||
await read_stream_writer.aclose()
|
||||
await write_stream.aclose()
|
||||
|
||||
async def terminate_session(self, client: httpx.AsyncClient) -> None:
|
||||
"""Terminate the session by sending a DELETE request."""
|
||||
|
|
@ -533,9 +530,6 @@ async def streamable_http_client(
|
|||
Example:
|
||||
See examples/snippets/clients/ for usage patterns.
|
||||
"""
|
||||
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
|
||||
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
|
||||
|
||||
# Determine if we need to create and manage the client
|
||||
client_provided = http_client is not None
|
||||
client = http_client
|
||||
|
|
@ -546,36 +540,40 @@ async def streamable_http_client(
|
|||
|
||||
transport = StreamableHTTPTransport(url)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
try:
|
||||
logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
|
||||
logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
|
||||
|
||||
async with contextlib.AsyncExitStack() as stack:
|
||||
# Only manage client lifecycle if we created it
|
||||
if not client_provided:
|
||||
await stack.enter_async_context(client)
|
||||
async with contextlib.AsyncExitStack() as stack:
|
||||
# Only manage client lifecycle if we created it
|
||||
if not client_provided:
|
||||
await stack.enter_async_context(client)
|
||||
|
||||
def start_get_stream() -> None:
|
||||
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
|
||||
read_stream_writer, read_stream = anyio.create_memory_object_stream[SessionMessage | Exception](0)
|
||||
write_stream, write_stream_reader = anyio.create_memory_object_stream[SessionMessage](0)
|
||||
|
||||
tg.start_soon(
|
||||
transport.post_writer,
|
||||
client,
|
||||
write_stream_reader,
|
||||
read_stream_writer,
|
||||
write_stream,
|
||||
start_get_stream,
|
||||
tg,
|
||||
)
|
||||
async with (
|
||||
read_stream_writer,
|
||||
read_stream,
|
||||
write_stream,
|
||||
write_stream_reader,
|
||||
anyio.create_task_group() as tg,
|
||||
):
|
||||
|
||||
try:
|
||||
yield read_stream, write_stream
|
||||
finally:
|
||||
if transport.session_id and terminate_on_close:
|
||||
await transport.terminate_session(client)
|
||||
tg.cancel_scope.cancel()
|
||||
finally:
|
||||
await read_stream_writer.aclose()
|
||||
await write_stream.aclose()
|
||||
await read_stream.aclose()
|
||||
await write_stream_reader.aclose()
|
||||
def start_get_stream() -> None:
|
||||
tg.start_soon(transport.handle_get_stream, client, read_stream_writer)
|
||||
|
||||
tg.start_soon(
|
||||
transport.post_writer,
|
||||
client,
|
||||
write_stream_reader,
|
||||
read_stream_writer,
|
||||
write_stream,
|
||||
start_get_stream,
|
||||
tg,
|
||||
)
|
||||
|
||||
try:
|
||||
yield read_stream, write_stream
|
||||
finally:
|
||||
if transport.session_id and terminate_on_close:
|
||||
await transport.terminate_session(client)
|
||||
tg.cancel_scope.cancel()
|
||||
|
|
|
|||
Loading…
Reference in a new issue