feat: allow to set a custom logger in the SSE and STDIO clients (#525)

* feat: allow to set a custom logger in the SSE client

So it logs to a logger instead of stdout.

Signed-off-by: Carlos Alexandro Becker <[email protected]>

* fix: tests, docs, stdio logger

* chore: lint

---------

Signed-off-by: Carlos Alexandro Becker <[email protected]>

Author: caarlos0

client/sse.go

@@ -23,12 +23,10 @@ func WithHTTPClient(httpClient *http.Client) transport.ClientOption {
 // NewSSEMCPClient creates a new SSE-based MCP client with the given base URL.
 // Returns an error if the URL is invalid.
 func NewSSEMCPClient(baseURL string, options ...transport.ClientOption) (*Client, error) {
-
 	sseTransport, err := transport.NewSSE(baseURL, options...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create SSE transport: %w", err)
 	}
-
 	return NewClient(sseTransport), nil
 }
 

client/transport/sse.go

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/util"
 )
 
 // SSE implements the transport layer of the MCP protocol using Server-Sent Events (SSE).
@@ -33,6 +34,7 @@ type SSE struct {
 	endpointChan   chan struct{}
 	headers        map[string]string
 	headerFunc     HTTPHeaderFunc
+	logger         util.Logger
 
 	started           atomic.Bool
 	closed            atomic.Bool
@@ -47,6 +49,13 @@ type SSE struct {
 
 type ClientOption func(*SSE)
 
+// WithSSELogger sets a custom logger for the SSE client.
+func WithSSELogger(logger util.Logger) ClientOption {
+	return func(sc *SSE) {
+		sc.logger = logger
+	}
+}
+
 func WithHeaders(headers map[string]string) ClientOption {
 	return func(sc *SSE) {
 		sc.headers = headers
@@ -85,6 +94,7 @@ func NewSSE(baseURL string, options ...ClientOption) (*SSE, error) {
 		responses:    make(map[string]chan *JSONRPCResponse),
 		endpointChan: make(chan struct{}),
 		headers:      make(map[string]string),
+		logger:       util.DefaultLogger(),
 	}
 
 	for _, opt := range options {
@@ -104,7 +114,6 @@ func NewSSE(baseURL string, options ...ClientOption) (*SSE, error) {
 // Start initiates the SSE connection to the server and waits for the endpoint information.
 // Returns an error if the connection fails or times out waiting for the endpoint.
 func (c *SSE) Start(ctx context.Context) error {
-
 	if c.started.Load() {
 		return fmt.Errorf("has already started")
 	}
@@ -113,7 +122,6 @@ func (c *SSE) Start(ctx context.Context) error {
 	c.cancelSSEStream = cancel
 
 	req, err := http.NewRequestWithContext(ctx, "GET", c.baseURL.String(), nil)
-
 	if err != nil {
 		return fmt.Errorf("failed to create request: %w", err)
 	}
@@ -220,7 +228,7 @@ func (c *SSE) readSSE(reader io.ReadCloser) {
 				}
 			}
 			if !c.closed.Load() {
-				fmt.Printf("SSE stream error: %v\n", err)
+				c.logger.Errorf("SSE stream error: %v", err)
 			}
 			return
 		}
@@ -256,11 +264,11 @@ func (c *SSE) handleSSEEvent(event, data string) {
 	case "endpoint":
 		endpoint, err := c.baseURL.Parse(data)
 		if err != nil {
-			fmt.Printf("Error parsing endpoint URL: %v\n", err)
+			c.logger.Errorf("Error parsing endpoint URL: %v", err)
 			return
 		}
 		if endpoint.Host != c.baseURL.Host {
-			fmt.Printf("Endpoint origin does not match connection origin\n")
+			c.logger.Errorf("Endpoint origin does not match connection origin")
 			return
 		}
 		c.endpoint = endpoint
@@ -269,7 +277,7 @@ func (c *SSE) handleSSEEvent(event, data string) {
 	case "message":
 		var baseMessage JSONRPCResponse
 		if err := json.Unmarshal([]byte(data), &baseMessage); err != nil {
-			fmt.Printf("Error unmarshaling message: %v\n", err)
+			c.logger.Errorf("Error unmarshaling message: %v", err)
 			return
 		}
 
@@ -321,7 +329,6 @@ func (c *SSE) SendRequest(
 	ctx context.Context,
 	request JSONRPCRequest,
 ) (*JSONRPCResponse, error) {
-
 	if !c.started.Load() {
 		return nil, fmt.Errorf("transport not started yet")
 	}