address todo comments

Author: IAmSurajBobade

mcp/streamable.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"io"
 	"iter"
+	"log/slog"
 	"math"
 	"math/rand/v2"
 	"net/http"
@@ -39,6 +40,7 @@ const (
 type StreamableHTTPHandler struct {
 	getServer func(*http.Request) *Server
 	opts      StreamableHTTPOptions
+	logger    *slog.Logger
 
 	onTransportDeletion func(sessionID string) // for testing only
 
@@ -67,6 +69,10 @@ type StreamableHTTPOptions struct {
 	//
 	// [§2.1.5]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#sending-messages-to-the-server
 	JSONResponse bool
+
+	// Logger specifies the logger to use.
+	// If nil, do not log.
+	Logger *slog.Logger
 }
 
 // NewStreamableHTTPHandler returns a new [StreamableHTTPHandler].
@@ -82,6 +88,12 @@ func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *Strea
 	if opts != nil {
 		h.opts = *opts
 	}
+
+	if h.opts.Logger == nil { // ensure we have a logger
+		h.opts.Logger = ensureLogger(nil)
+	}
+	h.logger = h.opts.Logger
+
 	return h
 }
 
@@ -367,6 +379,8 @@ type StreamableServerTransport struct {
 	// StreamableHTTPOptions.JSONResponse is exported.
 	jsonResponse bool
 
+	logger *slog.Logger
+
 	// connection is non-nil if and only if the transport has been connected.
 	connection *streamableServerConn
 }
@@ -381,6 +395,7 @@ func (t *StreamableServerTransport) Connect(ctx context.Context) (Connection, er
 		stateless:      t.Stateless,
 		eventStore:     t.EventStore,
 		jsonResponse:   t.jsonResponse,
+		logger:         t.logger,
 		incoming:       make(chan jsonrpc.Message, 10),
 		done:           make(chan struct{}),
 		streams:        make(map[string]*stream),
@@ -407,6 +422,8 @@ type streamableServerConn struct {
 	jsonResponse bool
 	eventStore   EventStore
 
+	logger *slog.Logger
+
 	incoming chan jsonrpc.Message // messages from the client to the server
 
 	mu sync.Mutex // guards all fields below
@@ -754,7 +771,7 @@ func (c *streamableServerConn) respondSSE(stream *stream, w http.ResponseWriter,
 		}
 		if _, err := writeEvent(w, e); err != nil {
 			// Connection closed or broken.
-			// TODO(#170): log when we add server-side logging.
+			c.logger.Warn("error writing event", "error", err)
 			return false
 		}
 		writes++
@@ -773,7 +790,13 @@ func (c *streamableServerConn) respondSSE(stream *stream, w http.ResponseWriter,
 				// simplify.
 				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 			} else {
-				// TODO(#170): log when we add server-side logging
+				if ctx.Err() != nil {
+					// Client disconnected or cancelled the request.
+					c.logger.Info("stream context done", "error", ctx.Err())
+				} else {
+					// Some other error.
+					c.logger.Warn("error receiving message", "error", err)
+				}
 			}
 			return
 		}