Repository navigation
Fix: memory leak and exit problems - #586
PWZER wants to merge 5 commits into
Conversation
| "pydantic>=2.7.2,<3.0.0", | ||
| "starlette>=0.27", | ||
| "sse-starlette>=1.6.1", | ||
| "sse-starlette>=2.3.0", |
There was a problem hiding this comment.
Is this needed?
Sorry, something went wrong.
There was a problem hiding this comment.
Yes, sse-starlette supports client_close_handler_callable since version 2.3.0.
Sorry, something went wrong.
ihrpr
left a comment
There was a problem hiding this comment.
Unfortunately we can't upgrade sse-starlette in SDK now.
We can achieve similar functionality as client_close_handler_callable provides in this case by using disconnect_listener
Something like this should work:
async with anyio.create_task_group() as tg:
response = EventSourceResponse(
content=sse_stream_reader,
data_sender_callable=sse_writer,
)
async def disconnect_listener():
try:
await response.listen_for_disconnect(receive)
logger.error(f"Disconnect detected for session: {session_id}")
await client_close_handler() # defined earlier in the PR - name should be changed probably
except Exception as e:
logger.error(f"Error in disconnect listener: {e}")
tg.start_soon(disconnect_listener)
logger.debug("Starting SSE response task")
tg.start_soon(response, scope, receive, send)
...
yield (read_stream, write_stream)
Sorry, something went wrong.
| async def client_close_handler(message: Message) -> None: | ||
| await read_stream_writer.aclose() | ||
| await write_stream_reader.aclose() | ||
| del self._read_stream_writers[session_id] |
There was a problem hiding this comment.
| del self._read_stream_writers[session_id] | |
| self._read_stream_writers.pop(session_id, None) |
Sorry, something went wrong.
| scope: Scope, | ||
| receive: Receive, | ||
| send: Send, | ||
| callback: Callable[[], None] | None = None, |
There was a problem hiding this comment.
The callback function only is called when client connection close so the naming is not full-meaning . And I don't think it should callback here because connect_sse yield
Sorry, something went wrong.
|
Closing this PR as stale. We can't update sse-starlette for now. Thank you for your contribution and working on improving this project. |
Sorry, something went wrong.
Motivation and Context
Here handle_sse calls app.run() every time it processes a sse request, but app.run() keeps running and does not stop automatically, which causes handle_sse to never end, and finally causes a memory leak.
python-sdk/examples/servers/simple-tool/mcp_simple_tool/server.py
Lines 67 to 73 in b4c7db6
And this is also the key reason for the following exit problems issues#541, Fixed.
INFO: Waiting for background tasks to complete. (CTRL+C to force quit)Types of changes
Checklist