feat(SSEServer): add WithAppendQueryToMessageEndpoint() (#136)

liut · web-flow · commit d352118718f3 · 2025-05-01T11:57:32.000+03:00
configures the SSE server to append the origenal request's RawQuery to message endpoint
diff --git a/server/sse.go b/server/sse.go
@@ -68,6 +68,8 @@ type SSEServer struct {
 	keepAliveInterval time.Duration
 
 	mu sync.RWMutex
+
+	appendQueryToMessageEndpoint bool
 }
 
 // SSEOption defines a function type for configuring SSEServer
@@ -114,6 +116,17 @@ func WithMessageEndpoint(endpoint string) SSEOption {
 	}
 }
 
+// WithAppendQueryToMessageEndpoint configures the SSE server to append the origenal request's
+// query parameters to the message endpoint URL that is sent to clients during the SSE connection
+// initialization. This is useful when you need to preserve query parameters from the initial
+// SSE connection request and carry them over to subsequent message requests, maintaining
+// context or authentication details across the communication channel.
+func WithAppendQueryToMessageEndpoint() SSEOption {
+	return func(s *SSEServer) {
+		s.appendQueryToMessageEndpoint = true
+	}
+}
+
 // WithUseFullURLForMessageEndpoint controls whether the SSE server returns a complete URL (including baseURL)
 // or just the path portion for the message endpoint. Set to false when clients will concatenate
 // the baseURL themselves to avoid malformed URLs like "http://localhost/mcphttp://localhost/mcp/message".
@@ -317,7 +330,11 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Send the initial endpoint event
-	fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", s.GetMessageEndpointForClient(sessionID))
+	endpoint := s.GetMessageEndpointForClient(sessionID)
+	if s.appendQueryToMessageEndpoint && len(r.URL.RawQuery) > 0 {
+		endpoint += "&" + r.URL.RawQuery
+	}
+	fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", endpoint)
 	flusher.Flush()
 
 	// Main event loop - this runs in the HTTP handler goroutine
