mcp: improve error messages from Wait for streamable clients

Previously, the error message received from ClientSession.Wait would
only report the closeErr, which would often be nil even if the client
transport was broken. Wait should return the reason the session
terminated, if abnormal.

I'm not sure of the exact semantics of this, but surely returning nil is
less useful than returning a meaningful non-nil error. We can refine our
handling of errors once we have more feedback.

Also add a test for client termination on HTTP server shutdown,
described in #265. This should work as long as (1) the session is
stateful (with a hanging GET), or (2) the session is stateless but the
client has a keepalive ping.

Also: don't send DELETE if the session was terminated with 404; +test.

Fixes #265

mcp/streamable_test.go

@@ -190,6 +190,80 @@ func TestStreamableTransports(t *testing.T) {
 	}
 }
 
+func TestStreamableServerShutdown(t *testing.T) {
+	ctx := context.Background()
+
+	// This test checks that closing the streamable HTTP server actually results
+	// in client session termination, provided one of following holds:
+	//  1. The server is stateful, and therefore the hanging GET fails the connection.
+	//  2. The server is stateless, and the client uses a KeepAlive.
+	tests := []struct {
+		name                 string
+		stateless, keepalive bool
+	}{
+		{"stateful", false, false},
+		{"stateless with keepalive", true, true},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			server := NewServer(testImpl, nil)
+			// Add a tool, just so we can check things are working.
+			AddTool(server, &Tool{Name: "greet"}, sayHi)
+
+			handler := NewStreamableHTTPHandler(
+				func(req *http.Request) *Server { return server },
+				&StreamableHTTPOptions{Stateless: test.stateless})
+
+			// When we shut down the server, we need to explicitly close ongoing
+			// connections. Otherwise, the hanging GET may never terminate.
+			httpServer := httptest.NewUnstartedServer(handler)
+			httpServer.Config.RegisterOnShutdown(func() {
+				for session := range server.Sessions() {
+					session.Close()
+				}
+			})
+			httpServer.Start()
+			defer httpServer.Close()
+
+			// Connect and run a tool.
+			var opts ClientOptions
+			if test.keepalive {
+				opts.KeepAlive = 50 * time.Millisecond
+			}
+			client := NewClient(testImpl, &opts)
+			clientSession, err := client.Connect(ctx, &StreamableClientTransport{
+				Endpoint:   httpServer.URL,
+				MaxRetries: -1, // avoid slow tests during exponential retries
+			}, nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			params := &CallToolParams{
+				Name:      "greet",
+				Arguments: map[string]any{"name": "foo"},
+			}
+			// Verify that we can call a tool.
+			if _, err := clientSession.CallTool(ctx, params); err != nil {
+				t.Fatalf("CallTool() failed: %v", err)
+			}
+
+			// Shut down the server. Sessions should terminate.
+			go func() {
+				if err := httpServer.Config.Shutdown(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) {
+					t.Errorf("closing http server: %v", err)
+				}
+			}()
+
+			// Wait may return an error (after all, the connection failed), but it
+			// should not hang.
+			t.Log("Client waiting")
+			_ = clientSession.Wait()
+		})
+	}
+}
+
 // TestClientReplay verifies that the client can recover from a mid-stream
 // network failure and receive replayed messages (if replay is configured). It
 // uses a proxy that is killed and restarted to simulate a recoverable network