I think this works but the test is randomly hanging once I've removed sleeps and prints and I have no idea why?

Author: wreed4

tests/shared/test_streamable_http.py

@@ -172,7 +172,8 @@ async def handle_call_tool(name: str, args: dict) -> list[TextContent]:
                     related_request_id=ctx.request_id,  # need for stream association
                 )
 
-                await anyio.sleep(0.1)
+                # need to wait for long enough that the client can reliably stop the tool before this finishes
+                await anyio.sleep(1)
 
                 await ctx.session.send_log_message(
                     level="info",
@@ -181,6 +182,16 @@ async def handle_call_tool(name: str, args: dict) -> list[TextContent]:
                     related_request_id=ctx.request_id,
                 )
 
+                # Adding another message just to make it even less likely that this tool will exit before the client can stop it
+                await anyio.sleep(1)
+
+                await ctx.session.send_log_message(
+                    level="info",
+                    data="Tool is done",
+                    logger="tool",
+                    related_request_id=ctx.request_id,
+                )
+
                 return [TextContent(type="text", text="Completed!")]
 
             elif name == "test_sampling_tool":
@@ -1033,6 +1044,7 @@ async def test_streamablehttp_client_resumption(event_server):
     """Test client session to resume a long running tool."""
     _, server_url = event_server
 
+    print("what is happening")
     # Variables to track the state
     captured_session_id = None
     captured_resumption_token = None
@@ -1099,6 +1111,10 @@ async def run_tool():
                     await anyio.sleep(0.1)
                 tg.cancel_scope.cancel()
 
+    # Make sure we only have one notification.. otherwise the test is flaky
+    # More than one notification means the tool likely could have finished already and will not call the message handler again upon resumption
+    assert len(captured_notifications) == 1
+
     # Store pre notifications and clear the captured notifications
     # for the post-resumption check
     captured_notifications_pre = captured_notifications.copy()