Merge pull request #477 from sunerpy/main

Author: pottekkat

client/transport/streamable_http.go

@@ -114,6 +114,7 @@ type StreamableHTTP struct {
 
 	// OAuth support
 	oauthHandler *OAuthHandler
+	wg           sync.WaitGroup
 }
 
 // NewStreamableHTTP creates a new Streamable HTTP transport with the given server URL.
@@ -182,9 +183,10 @@ func (c *StreamableHTTP) Close() error {
 	sessionId := c.sessionID.Load().(string)
 	if sessionId != "" {
 		c.sessionID.Store("")
-
+		c.wg.Add(1)
 		// notify server session closed
 		go func() {
+			defer c.wg.Done()
 			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 			defer cancel()
 			req, err := http.NewRequestWithContext(ctx, http.MethodDelete, c.serverURL.String(), nil)
@@ -201,7 +203,7 @@ func (c *StreamableHTTP) Close() error {
 			res.Body.Close()
 		}()
 	}
-
+	c.wg.Wait()
 	return nil
 }
 
@@ -231,7 +233,6 @@ func (c *StreamableHTTP) SendRequest(
 	ctx context.Context,
 	request JSONRPCRequest,
 ) (*JSONRPCResponse, error) {
-
 	// Marshal request
 	requestBody, err := json.Marshal(request)
 	if err != nil {
@@ -316,7 +317,6 @@ func (c *StreamableHTTP) sendHTTP(
 	body io.Reader,
 	acceptType string,
 ) (resp *http.Response, err error) {
-
 	// Create HTTP request
 	req, err := http.NewRequestWithContext(ctx, method, c.serverURL.String(), body)
 	if err != nil {
@@ -374,7 +374,6 @@ func (c *StreamableHTTP) sendHTTP(
 // It returns the final result for the request once received, or an error.
 // If ignoreResponse is true, it won't return when a response messge is received. This is for continuous listening.
 func (c *StreamableHTTP) handleSSEResponse(ctx context.Context, reader io.ReadCloser, ignoreResponse bool) (*JSONRPCResponse, error) {
-
 	// Create a channel for this specific request
 	responseChan := make(chan *JSONRPCResponse, 1)
 
@@ -387,7 +386,6 @@ func (c *StreamableHTTP) handleSSEResponse(ctx context.Context, reader io.ReadCl
 		defer close(responseChan)
 
 		c.readSSE(ctx, reader, func(event, data string) {
-
 			// (unsupported: batching)
 
 			var message JSONRPCResponse
@@ -490,7 +488,6 @@ func (c *StreamableHTTP) readSSE(ctx context.Context, reader io.ReadCloser, hand
 }
 
 func (c *StreamableHTTP) SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error {
-
 	// Marshal request
 	requestBody, err := json.Marshal(notification)
 	if err != nil {
@@ -577,7 +574,6 @@ var (
 )
 
 func (c *StreamableHTTP) createGETConnectionToServer(ctx context.Context) error {
-
 	resp, err := c.sendHTTP(ctx, http.MethodGet, nil, "text/event-stream")
 	if err != nil {
 		return fmt.Errorf("failed to send request: %w", err)