merge main and fix conflict

Author: dugenkui03

.gitignore

@@ -1,3 +1,5 @@
 .aider*
 .env
-.idea
+.idea
+.opencode
+.claude

README.md

@@ -91,6 +91,10 @@ MCP Go handles all the complex protocol details and server management, so you ca
   - [Tools](#tools)
   - [Prompts](#prompts)
 - [Examples](#examples)
+- [Extras](#extras)
+  - [Session Management](#session-management)
+  - [Request Hooks](#request-hooks)
+  - [Tool Handler Middleware](#tool-handler-middleware)
 - [Contributing](#contributing)
   - [Prerequisites](#prerequisites)
   - [Installation](#installation-1)
@@ -516,6 +520,214 @@ For examples, see the `examples/` directory.
 
 ## Extras
 
+### Session Management
+
+MCP-Go provides a robust session management system that allows you to:
+- Maintain separate state for each connected client
+- Register and track client sessions
+- Send notifications to specific clients
+- Provide per-session tool customization
+
+<details>
+<summary>Show Session Management Examples</summary>
+
+#### Basic Session Handling
+
+```go
+// Create a server with session capabilities
+s := server.NewMCPServer(
+    "Session Demo",
+    "1.0.0",
+    server.WithToolCapabilities(true),
+)
+
+// Implement your own ClientSession
+type MySession struct {
+    id           string
+    notifChannel chan mcp.JSONRPCNotification
+    isInitialized bool
+    // Add custom fields for your application
+}
+
+// Implement the ClientSession interface
+func (s *MySession) SessionID() string {
+    return s.id
+}
+
+func (s *MySession) NotificationChannel() chan<- mcp.JSONRPCNotification {
+    return s.notifChannel
+}
+
+func (s *MySession) Initialize() {
+    s.isInitialized = true
+}
+
+func (s *MySession) Initialized() bool {
+    return s.isInitialized
+}
+
+// Register a session
+session := &MySession{
+    id:           "user-123",
+    notifChannel: make(chan mcp.JSONRPCNotification, 10),
+}
+if err := s.RegisterSession(context.Background(), session); err != nil {
+    log.Printf("Failed to register session: %v", err)
+}
+
+// Send notification to a specific client
+err := s.SendNotificationToSpecificClient(
+    session.SessionID(),
+    "notification/update",
+    map[string]any{"message": "New data available!"},
+)
+if err != nil {
+    log.Printf("Failed to send notification: %v", err)
+}
+
+// Unregister session when done
+s.UnregisterSession(context.Background(), session.SessionID())
+```
+
+#### Per-Session Tools
+
+For more advanced use cases, you can implement the `SessionWithTools` interface to support per-session tool customization:
+
+```go
+// Implement SessionWithTools interface for per-session tools
+type MyAdvancedSession struct {
+    MySession  // Embed the basic session
+    sessionTools map[string]server.ServerTool
+}
+
+// Implement additional methods for SessionWithTools
+func (s *MyAdvancedSession) GetSessionTools() map[string]server.ServerTool {
+    return s.sessionTools
+}
+
+func (s *MyAdvancedSession) SetSessionTools(tools map[string]server.ServerTool) {
+    s.sessionTools = tools
+}
+
+// Create and register a session with tools support
+advSession := &MyAdvancedSession{
+    MySession: MySession{
+        id:           "user-456",
+        notifChannel: make(chan mcp.JSONRPCNotification, 10),
+    },
+    sessionTools: make(map[string]server.ServerTool),
+}
+if err := s.RegisterSession(context.Background(), advSession); err != nil {
+    log.Printf("Failed to register session: %v", err)
+}
+
+// Add session-specific tools
+userSpecificTool := mcp.NewTool(
+    "user_data",
+    mcp.WithDescription("Access user-specific data"),
+)
+// You can use AddSessionTool (similar to AddTool)
+err := s.AddSessionTool(
+    advSession.SessionID(),
+    userSpecificTool,
+    func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+        // This handler is only available to this specific session
+        return mcp.NewToolResultText("User-specific data for " + advSession.SessionID()), nil
+    },
+)
+if err != nil {
+    log.Printf("Failed to add session tool: %v", err)
+}
+
+// Or use AddSessionTools directly with ServerTool
+/*
+err := s.AddSessionTools(
+    advSession.SessionID(),
+    server.ServerTool{
+        Tool: userSpecificTool,
+        Handler: func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+            // This handler is only available to this specific session
+            return mcp.NewToolResultText("User-specific data for " + advSession.SessionID()), nil
+        },
+    },
+)
+if err != nil {
+    log.Printf("Failed to add session tool: %v", err)
+}
+*/
+
+// Delete session-specific tools when no longer needed
+err = s.DeleteSessionTools(advSession.SessionID(), "user_data")
+if err != nil {
+    log.Printf("Failed to delete session tool: %v", err)
+}
+```
+
+#### Tool Filtering
+
+You can also apply filters to control which tools are available to certain sessions:
+
+```go
+// Add a tool filter that only shows tools with certain prefixes
+s := server.NewMCPServer(
+    "Tool Filtering Demo",
+    "1.0.0",
+    server.WithToolCapabilities(true),
+    server.WithToolFilter(func(ctx context.Context, tools []mcp.Tool) []mcp.Tool {
+        // Get session from context
+        session := server.ClientSessionFromContext(ctx)
+        if session == nil {
+            return tools // Return all tools if no session
+        }
+        
+        // Example: filter tools based on session ID prefix
+        if strings.HasPrefix(session.SessionID(), "admin-") {
+            // Admin users get all tools
+            return tools
+        } else {
+            // Regular users only get tools with "public-" prefix
+            var filteredTools []mcp.Tool
+            for _, tool := range tools {
+                if strings.HasPrefix(tool.Name, "public-") {
+                    filteredTools = append(filteredTools, tool)
+                }
+            }
+            return filteredTools
+        }
+    }),
+)
+```
+
+#### Working with Context
+
+The session context is automatically passed to tool and resource handlers:
+
+```go
+s.AddTool(mcp.NewTool("session_aware"), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+    // Get the current session from context
+    session := server.ClientSessionFromContext(ctx)
+    if session == nil {
+        return mcp.NewToolResultError("No active session"), nil
+    }
+    
+    return mcp.NewToolResultText("Hello, session " + session.SessionID()), nil
+})
+
+// When using handlers in HTTP/SSE servers, you need to pass the context with the session
+httpHandler := func(w http.ResponseWriter, r *http.Request) {
+    // Get session from somewhere (like a cookie or header)
+    session := getSessionFromRequest(r)
+    
+    // Add session to context
+    ctx := s.WithContext(r.Context(), session)
+    
+    // Use this context when handling requests
+    // ...
+}
+```
+
+</details>
+
 ### Request Hooks
 
 Hook into the request lifecycle by creating a `Hooks` object with your

client/sse.go

@@ -2,14 +2,20 @@ package client
 
 import (
 	"fmt"
-	"github.com/mark3labs/mcp-go/client/transport"
+	"net/http"
 	"net/url"
+
+	"github.com/mark3labs/mcp-go/client/transport"
 )
 
 func WithHeaders(headers map[string]string) transport.ClientOption {
 	return transport.WithHeaders(headers)
 }
 
+func WithHTTPClient(httpClient *http.Client) transport.ClientOption {
+	return transport.WithHTTPClient(httpClient)
+}
+
 // 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) {

client/stdio_test.go

@@ -7,7 +7,7 @@ import (
 	"log/slog"
 	"os"
 	"os/exec"
-	"path/filepath"
+	"runtime"
 	"sync"
 	"testing"
 	"time"
@@ -19,21 +19,41 @@ func compileTestServer(outputPath string) error {
 	cmd := exec.Command(
 		"go",
 		"build",
+		"-buildmode=pie",
 		"-o",
 		outputPath,
 		"../testdata/mockstdio_server.go",
 	)
+	tmpCache, _ := os.MkdirTemp("", "gocache")
+	cmd.Env = append(os.Environ(), "GOCACHE="+tmpCache)
+
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("compilation failed: %v\nOutput: %s", err, output)
 	}
+	// Verify the binary was actually created
+	if _, err := os.Stat(outputPath); os.IsNotExist(err) {
+		return fmt.Errorf("mock server binary not found at %s after compilation", outputPath)
+	}
 	return nil
 }
 
 func TestStdioMCPClient(t *testing.T) {
-	// Compile mock server
-	mockServerPath := filepath.Join(os.TempDir(), "mockstdio_server")
-	if err := compileTestServer(mockServerPath); err != nil {
-		t.Fatalf("Failed to compile mock server: %v", err)
+	// Create a temporary file for the mock server
+	tempFile, err := os.CreateTemp("", "mockstdio_server")
+	if err != nil {
+		t.Fatalf("Failed to create temp file: %v", err)
+	}
+	tempFile.Close()
+	mockServerPath := tempFile.Name()
+	
+	// Add .exe suffix on Windows
+	if runtime.GOOS == "windows" {
+		os.Remove(mockServerPath) // Remove the empty file first
+		mockServerPath += ".exe"
+	}
+	
+	if compileErr := compileTestServer(mockServerPath); compileErr != nil {
+		t.Fatalf("Failed to compile mock server: %v", compileErr)
 	}
 	defer os.Remove(mockServerPath)
 

client/transport/sse.go

@@ -45,6 +45,12 @@ func WithHeaders(headers map[string]string) ClientOption {
 	}
 }
 
+func WithHTTPClient(httpClient *http.Client) ClientOption {
+	return func(sc *SSE) {
+		sc.httpClient = httpClient
+	}
+}
+
 // NewSSE creates a new SSE-based MCP client with the given base URL.
 // Returns an error if the URL is invalid.
 func NewSSE(baseURL string, options ...ClientOption) (*SSE, error) {
@@ -226,7 +232,7 @@ func (c *SSE) SetNotificationHandler(handler func(notification mcp.JSONRPCNotifi
 	c.onNotification = handler
 }
 
-// sendRequest sends a JSON-RPC request to the server and waits for a response.
+// SendRequest sends a JSON-RPC request to the server and waits for a response.
 // Returns the raw JSON response message or an error if the request fails.
 func (c *SSE) SendRequest(
 	ctx context.Context,
@@ -278,13 +284,19 @@ func (c *SSE) SendRequest(
 		deleteResponseChan()
 		return nil, fmt.Errorf("failed to send request: %w", err)
 	}
-	defer resp.Body.Close()
+
+	// Drain any outstanding io
+	body, err := io.ReadAll(resp.Body)
+	resp.Body.Close()
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
+	}
 
 	// Check if we got an error response
 	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
 		deleteResponseChan()
 
-		body, _ := io.ReadAll(resp.Body)
 		return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, body)
 	}
 

client/transport/sse_test.go

@@ -415,6 +415,31 @@ func TestSSEErrors(t *testing.T) {
 		}
 	})
 
+	t.Run("WithHTTPClient", func(t *testing.T) {
+		// Create a custom client with a very short timeout
+		customClient := &http.Client{Timeout: 1 * time.Nanosecond}
+
+		url, closeF := startMockSSEEchoServer()
+		defer closeF()
+		// Initialize SSE transport with the custom HTTP client
+		trans, err := NewSSE(url, WithHTTPClient(customClient))
+		if err != nil {
+			t.Fatalf("Failed to create SSE with custom client: %v", err)
+		}
+
+		// Starting should immediately error due to timeout
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		err = trans.Start(ctx)
+		if err == nil {
+			t.Error("Expected Start to fail with custom timeout, got nil")
+		}
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Errorf("Expected error 'context deadline exceeded', got '%s'", err.Error())
+		}
+		trans.Close()
+	})
+
 	t.Run("RequestBeforeStart", func(t *testing.T) {
 		url, closeF := startMockSSEEchoServer()
 		defer closeF()