examples: Fix SSE streaming issue by moving logging from HTTP to MCP middlewarre (#614)
In `examples/http`, previous HTTP-level logging middleware blocked SSE streaming because `handler.ServeHTTP()` does not return until the connection closes for GET requests (SSE streams). This prevented response headers from being sent to the client. Changes: - Replace HTTP-level logging handler with MCP middleware pattern - Implement logging at the MCP protocol layer using mcp.Middleware - Add middleware to server using AddReceivingMiddleware() - Simplify main.go by removing HTTP handler wrapping This allows SSE streams to work properly while still providing detailed logging for all MCP protocol-level method calls. Fixes #632
This commit is contained in:
parent
d297272705
commit
35dfcf21c9
2 changed files with 46 additions and 43 deletions
|
|
@ -5,47 +5,49 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
// responseWriter wraps http.ResponseWriter to capture the status code.
|
||||
type responseWriter struct {
|
||||
http.ResponseWriter
|
||||
statusCode int
|
||||
}
|
||||
// createLoggingMiddleware creates an MCP middleware that logs method calls.
|
||||
func createLoggingMiddleware() mcp.Middleware {
|
||||
return func(next mcp.MethodHandler) mcp.MethodHandler {
|
||||
return func(
|
||||
ctx context.Context,
|
||||
method string,
|
||||
req mcp.Request,
|
||||
) (mcp.Result, error) {
|
||||
start := time.Now()
|
||||
sessionID := req.GetSession().ID()
|
||||
|
||||
func (rw *responseWriter) WriteHeader(code int) {
|
||||
rw.statusCode = code
|
||||
rw.ResponseWriter.WriteHeader(code)
|
||||
}
|
||||
// Log request details.
|
||||
log.Printf("[REQUEST] Session: %s | Method: %s",
|
||||
sessionID,
|
||||
method)
|
||||
|
||||
func loggingHandler(handler http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
start := time.Now()
|
||||
// Call the actual handler.
|
||||
result, err := next(ctx, method, req)
|
||||
|
||||
// Create a response writer wrapper to capture status code.
|
||||
wrapped := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK}
|
||||
// Log response details.
|
||||
duration := time.Since(start)
|
||||
|
||||
// Log request details.
|
||||
log.Printf("[REQUEST] %s | %s | %s %s",
|
||||
start.Format(time.RFC3339),
|
||||
r.RemoteAddr,
|
||||
r.Method,
|
||||
r.URL.Path)
|
||||
if err != nil {
|
||||
log.Printf("[RESPONSE] Session: %s | Method: %s | Status: ERROR | Duration: %v | Error: %v",
|
||||
sessionID,
|
||||
method,
|
||||
duration,
|
||||
err)
|
||||
} else {
|
||||
log.Printf("[RESPONSE] Session: %s | Method: %s | Status: OK | Duration: %v",
|
||||
sessionID,
|
||||
method,
|
||||
duration)
|
||||
}
|
||||
|
||||
// Call the actual handler.
|
||||
handler.ServeHTTP(wrapped, r)
|
||||
|
||||
// Log response details.
|
||||
duration := time.Since(start)
|
||||
log.Printf("[RESPONSE] %s | %s | %s %s | Status: %d | Duration: %v",
|
||||
time.Now().Format(time.RFC3339),
|
||||
r.RemoteAddr,
|
||||
r.Method,
|
||||
r.URL.Path,
|
||||
wrapped.statusCode,
|
||||
duration)
|
||||
})
|
||||
return result, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -117,6 +117,9 @@ func runServer(url string) {
|
|||
Version: "1.0.0",
|
||||
}, nil)
|
||||
|
||||
// Add MCP-level logging middleware.
|
||||
server.AddReceivingMiddleware(createLoggingMiddleware())
|
||||
|
||||
// Add the cityTime tool.
|
||||
mcp.AddTool(server, &mcp.Tool{
|
||||
Name: "cityTime",
|
||||
|
|
@ -128,13 +131,11 @@ func runServer(url string) {
|
|||
return server
|
||||
}, nil)
|
||||
|
||||
handlerWithLogging := loggingHandler(handler)
|
||||
|
||||
log.Printf("MCP server listening on %s", url)
|
||||
log.Printf("Available tool: cityTime (cities: nyc, sf, boston)")
|
||||
|
||||
// Start the HTTP server with logging handler.
|
||||
if err := http.ListenAndServe(url, handlerWithLogging); err != nil {
|
||||
// Start the HTTP server.
|
||||
if err := http.ListenAndServe(url, handler); err != nil {
|
||||
log.Fatalf("Server failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue