Skip to content

Commit 60a5065

Browse files
committed
revert sse and test entirely
1 parent 03601e6 commit 60a5065

File tree

2 files changed

+31
-51
lines changed

2 files changed

+31
-51
lines changed

src/mcp/server/sse.py

Lines changed: 10 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -120,29 +120,15 @@ async def sse_writer():
120120
}
121121
)
122122

123-
# Ensure all streams are properly closed
124-
async with read_stream, write_stream, read_stream_writer, sse_stream_reader:
125-
async with anyio.create_task_group() as tg:
126-
response = EventSourceResponse(
127-
content=sse_stream_reader, data_sender_callable=sse_writer
128-
)
129-
logger.debug("Starting SSE response task")
130-
tg.start_soon(response, scope, receive, send)
131-
132-
try:
133-
logger.debug("Yielding read and write streams")
134-
yield (read_stream, write_stream)
135-
finally:
136-
# Cleanup when connection closes
137-
logger.debug(f"Cleaning up SSE session {session_id}")
138-
try:
139-
# Remove session from tracking dictionary
140-
if session_id in self._read_stream_writers:
141-
del self._read_stream_writers[session_id]
142-
# Cancel any remaining tasks in the task group
143-
tg.cancel_scope.cancel()
144-
except Exception as e:
145-
logger.error(f"Error during SSE cleanup: {e}")
123+
async with anyio.create_task_group() as tg:
124+
response = EventSourceResponse(
125+
content=sse_stream_reader, data_sender_callable=sse_writer
126+
)
127+
logger.debug("Starting SSE response task")
128+
tg.start_soon(response, scope, receive, send)
129+
130+
logger.debug("Yielding read and write streams")
131+
yield (read_stream, write_stream)
146132

147133
async def handle_post_message(
148134
self, scope: Scope, receive: Receive, send: Send
@@ -186,4 +172,4 @@ async def handle_post_message(
186172
logger.debug(f"Sending message to writer: {message}")
187173
response = Response("Accepted", status_code=202)
188174
await response(scope, receive, send)
189-
await writer.send(message)
175+
await writer.send(message)

tests/shared/test_sse.py

Lines changed: 21 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -142,30 +142,28 @@ def server(server_port: int) -> Generator[None, None, None]:
142142

143143
yield
144144

145-
print("shutting down server gracefully")
146-
# Try graceful shutdown first
147-
proc.terminate()
148-
try:
149-
proc.join(timeout=5)
150-
except Exception:
151-
print("Graceful shutdown failed, forcing kill")
152-
proc.kill()
153-
proc.join(timeout=2)
154-
145+
print("killing server")
146+
# Signal the server to stop
147+
proc.kill()
148+
proc.join(timeout=2)
155149
if proc.is_alive():
156150
print("server process failed to terminate")
157-
proc.kill() # Force kill as last resort
158151

159152

153+
@pytest.fixture()
154+
async def http_client(server, server_url) -> AsyncGenerator[httpx.AsyncClient, None]:
155+
"""Create test client"""
156+
async with httpx.AsyncClient(base_url=server_url) as client:
157+
yield client
158+
160159

160+
# Tests
161161
@pytest.mark.anyio
162-
@pytest.mark.skip(
163-
"fails in CI, but works locally. Need to investigate why."
164-
)
165-
async def test_raw_sse_connection(server, server_url) -> None:
162+
async def test_raw_sse_connection(http_client: httpx.AsyncClient) -> None:
166163
"""Test the SSE connection establishment simply with an HTTP client."""
167-
try:
168-
async with httpx.AsyncClient(base_url=server_url) as http_client:
164+
async with anyio.create_task_group():
165+
166+
async def connection_test() -> None:
169167
async with http_client.stream("GET", "/sse") as response:
170168
assert response.status_code == 200
171169
assert (
@@ -183,11 +181,12 @@ async def test_raw_sse_connection(server, server_url) -> None:
183181
return
184182
line_number += 1
185183

186-
except Exception as e:
187-
pytest.fail(f"{e}")
184+
# Add timeout to prevent test from hanging if it fails
185+
with anyio.fail_after(3):
186+
await connection_test()
188187

189-
@pytest.mark.anyio
190188

189+
@pytest.mark.anyio
191190
async def test_sse_client_basic_connection(server: None, server_url: str) -> None:
192191
async with sse_client(server_url + "/sse") as streams:
193192
async with ClientSession(*streams) as session:
@@ -200,6 +199,7 @@ async def test_sse_client_basic_connection(server: None, server_url: str) -> Non
200199
ping_result = await session.send_ping()
201200
assert isinstance(ping_result, EmptyResult)
202201

202+
203203
@pytest.fixture
204204
async def initialized_sse_client_session(
205205
server, server_url: str
@@ -211,9 +211,6 @@ async def initialized_sse_client_session(
211211

212212

213213
@pytest.mark.anyio
214-
@pytest.mark.skip(
215-
"fails in CI, but works locally. Need to investigate why."
216-
)
217214
async def test_sse_client_happy_request_and_response(
218215
initialized_sse_client_session: ClientSession,
219216
) -> None:
@@ -225,9 +222,6 @@ async def test_sse_client_happy_request_and_response(
225222

226223

227224
@pytest.mark.anyio
228-
@pytest.mark.skip(
229-
"fails in CI, but works locally. Need to investigate why."
230-
)
231225
async def test_sse_client_exception_handling(
232226
initialized_sse_client_session: ClientSession,
233227
) -> None:
@@ -255,4 +249,4 @@ async def test_sse_client_timeout(
255249
# we should receive an error here
256250
return
257251

258-
pytest.fail("the client should have timed out and returned an error already")
252+
pytest.fail("the client should have timed out and returned an error already")

0 commit comments

Comments
 (0)