@@ -36,13 +36,6 @@ type sseSession struct {
3636// content. This can be used to inject context values from headers, for example. 
3737type  SSEContextFunc  func (ctx  context.Context , r  * http.Request ) context.Context 
3838
39- // DynamicBasePathFunc allows the user to provide a function to generate the 
40- // base path for a given request and sessionID. This is useful for cases where 
41- // the base path is not known at the time of SSE server creation, such as when 
42- // using a reverse proxy or when the base path is dynamically generated. The 
43- // function should return the base path (e.g., "/mcp/tenant123"). 
44- type  DynamicBasePathFunc  func (r  * http.Request , sessionID  string ) string 
45- 
4639func  (s  * sseSession ) SessionID () string  {
4740	return  s .sessionID 
4841}
@@ -100,7 +93,7 @@ type SSEServer struct {
10093	sseEndpoint                   string 
10194	sessions                      sync.Map 
10295	srv                           * http.Server 
103- 	contextFunc                   SSEContextFunc 
96+ 	contextFunc                   HTTPContextFunc 
10497	dynamicBasePathFunc           DynamicBasePathFunc 
10598
10699	keepAlive          bool 
@@ -109,37 +102,41 @@ type SSEServer struct {
109102	mu  sync.RWMutex 
110103}
111104
112- // SSEOption defines a function type for configuring SSEServer  
113- type   SSEOption   func (* SSEServer )
105+ // Ensure SSEServer implements httpTransportConfigurable  
106+ var   _   httpTransportConfigurable   =   (* SSEServer )( nil )
114107
115- // WithBaseURL sets the base URL for the SSE server 
116- func  WithBaseURL (baseURL  string ) SSEOption  {
117- 	return  func (s  * SSEServer ) {
118- 		if  baseURL  !=  ""  {
119- 			u , err  :=  url .Parse (baseURL )
120- 			if  err  !=  nil  {
121- 				return 
122- 			}
123- 			if  u .Scheme  !=  "http"  &&  u .Scheme  !=  "https"  {
124- 				return 
125- 			}
126- 			// Check if the host is empty or only contains a port 
127- 			if  u .Host  ==  ""  ||  strings .HasPrefix (u .Host , ":" ) {
128- 				return 
129- 			}
130- 			if  len (u .Query ()) >  0  {
131- 				return 
132- 			}
108+ func  (s  * SSEServer ) setBasePath (basePath  string ) {
109+ 	s .basePath  =  normalizeURLPath (basePath )
110+ }
111+ 
112+ func  (s  * SSEServer ) setDynamicBasePath (fn  DynamicBasePathFunc ) {
113+ 	if  fn  !=  nil  {
114+ 		s .dynamicBasePathFunc  =  func (r  * http.Request , sid  string ) string  {
115+ 			bp  :=  fn (r , sid )
116+ 			return  normalizeURLPath (bp )
133117		}
134- 		s .baseURL  =  strings .TrimSuffix (baseURL , "/" )
135118	}
136119}
137120
138- // WithStaticBasePath adds a new option for setting a static base path 
139- func  WithStaticBasePath (basePath  string ) SSEOption  {
140- 	return  func (s  * SSEServer ) {
141- 		s .basePath  =  normalizeURLPath (basePath )
142- 	}
121+ func  (s  * SSEServer ) setKeepAliveInterval (interval  time.Duration ) {
122+ 	s .keepAlive  =  true 
123+ 	s .keepAliveInterval  =  interval 
124+ }
125+ 
126+ func  (s  * SSEServer ) setKeepAlive (keepAlive  bool ) {
127+ 	s .keepAlive  =  keepAlive 
128+ }
129+ 
130+ func  (s  * SSEServer ) setContextFunc (fn  HTTPContextFunc ) {
131+ 	s .contextFunc  =  fn 
132+ }
133+ 
134+ func  (s  * SSEServer ) setHTTPServer (srv  * http.Server ) {
135+ 	s .srv  =  srv 
136+ }
137+ 
138+ func  (s  * SSEServer ) setBaseURL (baseURL  string ) {
139+ 	s .baseURL  =  baseURL 
143140}
144141
145142// WithBasePath adds a new option for setting a static base path. 
@@ -151,26 +148,11 @@ func WithBasePath(basePath string) SSEOption {
151148	return  WithStaticBasePath (basePath )
152149}
153150
154- // WithDynamicBasePath accepts a function for generating the base path. This is 
155- // useful for cases where the base path is not known at the time of SSE server 
156- // creation, such as when using a reverse proxy or when the server is mounted 
157- // at a dynamic path. 
158- func  WithDynamicBasePath (fn  DynamicBasePathFunc ) SSEOption  {
159- 	return  func (s  * SSEServer ) {
160- 		if  fn  !=  nil  {
161- 			s .dynamicBasePathFunc  =  func (r  * http.Request , sid  string ) string  {
162- 				bp  :=  fn (r , sid )
163- 				return  normalizeURLPath (bp )
164- 			}
165- 		}
166- 	}
167- }
168- 
169151// WithMessageEndpoint sets the message endpoint path 
170152func  WithMessageEndpoint (endpoint  string ) SSEOption  {
171- 	return  func (s  * SSEServer ) {
153+ 	return  sseOption ( func (s  * SSEServer ) {
172154		s .messageEndpoint  =  endpoint 
173- 	}
155+ 	}) 
174156}
175157
176158// WithAppendQueryToMessageEndpoint configures the SSE server to append the original request's 
@@ -179,53 +161,37 @@ func WithMessageEndpoint(endpoint string) SSEOption {
179161// SSE connection request and carry them over to subsequent message requests, maintaining 
180162// context or authentication details across the communication channel. 
181163func  WithAppendQueryToMessageEndpoint () SSEOption  {
182- 	return  func (s  * SSEServer ) {
164+ 	return  sseOption ( func (s  * SSEServer ) {
183165		s .appendQueryToMessageEndpoint  =  true 
184- 	}
166+ 	}) 
185167}
186168
187169// WithUseFullURLForMessageEndpoint controls whether the SSE server returns a complete URL (including baseURL) 
188170// or just the path portion for the message endpoint. Set to false when clients will concatenate 
189171// the baseURL themselves to avoid malformed URLs like "http://localhost/mcphttp://localhost/mcp/message". 
190172func  WithUseFullURLForMessageEndpoint (useFullURLForMessageEndpoint  bool ) SSEOption  {
191- 	return  func (s  * SSEServer ) {
173+ 	return  sseOption ( func (s  * SSEServer ) {
192174		s .useFullURLForMessageEndpoint  =  useFullURLForMessageEndpoint 
193- 	}
175+ 	}) 
194176}
195177
196178// WithSSEEndpoint sets the SSE endpoint path 
197179func  WithSSEEndpoint (endpoint  string ) SSEOption  {
198- 	return  func (s  * SSEServer ) {
180+ 	return  sseOption ( func (s  * SSEServer ) {
199181		s .sseEndpoint  =  endpoint 
200- 	}
201- }
202- 
203- // WithHTTPServer sets the HTTP server instance 
204- func  WithHTTPServer (srv  * http.Server ) SSEOption  {
205- 	return  func (s  * SSEServer ) {
206- 		s .srv  =  srv 
207- 	}
208- }
209- 
210- func  WithKeepAliveInterval (keepAliveInterval  time.Duration ) SSEOption  {
211- 	return  func (s  * SSEServer ) {
212- 		s .keepAlive  =  true 
213- 		s .keepAliveInterval  =  keepAliveInterval 
214- 	}
215- }
216- 
217- func  WithKeepAlive (keepAlive  bool ) SSEOption  {
218- 	return  func (s  * SSEServer ) {
219- 		s .keepAlive  =  keepAlive 
220- 	}
182+ 	})
221183}
222184
223185// WithSSEContextFunc sets a function that will be called to customise the context 
224186// to the server using the incoming request. 
187+ // 
188+ // Deprecated: Use WithContextFunc instead. This will be removed in a future version. 
189+ // 
190+ //go:deprecated 
225191func  WithSSEContextFunc (fn  SSEContextFunc ) SSEOption  {
226- 	return  func (s  * SSEServer ) {
227- 		s . contextFunc   =   fn 
228- 	}
192+ 	return  sseOption ( func (s  * SSEServer ) {
193+ 		WithHTTPContextFunc ( HTTPContextFunc ( fn )). applyToSSE ( s ) 
194+ 	}) 
229195}
230196
231197// NewSSEServer creates a new SSE server instance with the given MCP server and options. 
@@ -241,16 +207,15 @@ func NewSSEServer(server *MCPServer, opts ...SSEOption) *SSEServer {
241207
242208	// Apply all options 
243209	for  _ , opt  :=  range  opts  {
244- 		opt (s )
210+ 		opt . applyToSSE (s )
245211	}
246212
247213	return  s 
248214}
249215
250- // NewTestServer creates a test server for testing purposes 
216+ // NewTestServer creates a test server for testing purposes.  
251217func  NewTestServer (server  * MCPServer , opts  ... SSEOption ) * httptest.Server  {
252218	sseServer  :=  NewSSEServer (server , opts ... )
253- 
254219	testServer  :=  httptest .NewServer (sseServer )
255220	sseServer .baseURL  =  testServer .URL 
256221	return  testServer 
0 commit comments