diff --git a/pyproject.toml b/pyproject.toml index 5640f35..956d9c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,12 +30,12 @@ dependencies = [ "sse-starlette>=1.6.1", "pydantic-settings>=2.5.2", "uvicorn>=0.23.1", - "websockets>=15.0.1", ] [project.optional-dependencies] rich = ["rich>=13.9.4"] cli = ["typer>=0.12.4", "python-dotenv>=1.0.0"] +ws = ["websockets>=15.0.1"] [project.scripts] mcp = "mcp.cli:app [cli]" diff --git a/src/mcp/client/websocket.py b/src/mcp/client/websocket.py index 62c777f..1e2026b 100644 --- a/src/mcp/client/websocket.py +++ b/src/mcp/client/websocket.py @@ -56,8 +56,9 @@ async def websocket_client(url: str) -> AsyncGenerator[ except ValidationError as exc: # If JSON parse or model validation fails, send the exception await read_stream_writer.send(exc) - except (anyio.ClosedResourceError, Exception): + except (anyio.ClosedResourceError, Exception) as e: await ws.close() + raise e async def ws_writer(): """ @@ -71,8 +72,9 @@ async def websocket_client(url: str) -> AsyncGenerator[ by_alias=True, mode="json", exclude_none=True ) await ws.send(json.dumps(msg_dict)) - except (anyio.ClosedResourceError, Exception): + except (anyio.ClosedResourceError, Exception) as e: await ws.close() + raise e async with anyio.create_task_group() as tg: # Start reader and writer tasks