Repository: mark3labs/mcp-go

Issue:
ToolCall request returns 504 Gateway Timeout via Nginx Ingress for Long-Running Requests
## Problem Description

We observe intermittent 504 errors when handling long-running requests through Nginx Ingress, despite the SSE backend functioning properly. Below are the key details:

## Environment

### 1. Backend Service:

* mcp-server deployed in Kubernetes
* Used `server.WithKeepAliveInterval(10 * time.Second)` for SSE connections
* Processing logic  requires 2+ minutes to complete

### 2. Gateway Layer:

* Nginx Ingress Controller default timeout settings (60s)

## Observed Behavior

**client by `python-sdk`:**
* Successful results via /sse endpoint.
* the ToolCall request receive 504 Gateway Timeout after 60s, though backend logs confirm continued processing.
<img width="1616" alt="Image" src="https://github.com/user-attachments/assets/a2daf91a-aeee-4376-8da4-3fc455c418f6" />

**client by `mcp-go`:**

* timeout, no result return

![Image](https://github.com/user-attachments/assets/d7125eaf-0e4a-4f7c-b732-ba4304ee3e0c)

Hints from issue comments:
@Hwting: > ## Problem Description
> We observe intermittent 504 errors when handling long-running requests through Nginx Ingress, despite the SSE backend functioning properly. Below are the key details:
> 
> ## Environment
> ### 1. Backend Service:
> * mcp-server deployed in Kubernetes
> * Used `server.WithKeepAliveInterval(10 * time.Second)` for SSE connections
> * Processing logic  requires 2+ minutes to complete
> 
> ### 2. Gateway Layer:
> * Nginx Ingress Controller default timeout settings (60s)
> 
> ## Observed Behavior
> **client by `python-sdk`:**
> 
> * Successful results via /sse endpoint.
> * the ToolCall request receive 504 Gateway Timeout after 60s, though backend logs confirm continued processing.
> 
> <img alt="Image" width="1616" src="https://private-user-images.githubusercontent.com/79098669/434622390-a2daf91a-aeee-4376-8da4-3fc455c418f6.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3NDUxMjY3NzQsIm5iZiI6MTc0NTEyNjQ3NCwicGF0aCI6Ii83OTA5ODY2OS80MzQ2MjIzOTAtYTJkYWY5MWEtYWVlZS00Mzc2LThkYTQtM2ZjNDU1YzQxOGY2LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNTA0MjAlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjUwNDIwVDA1MjExNFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWYzOTIyNDIyNjg1NTRlZjgwNTIxN2RmMmM1MmJlMWYyZGJiZjRmYWQzOGM1MDQzY2U4NmU4MjQ5ZDQ0NGIwNzYmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0In0.ZCmUb8EB_zFTaWXKX6TaXVyPmRjQ0dCrR0CbwL6MOOo">
> **client by `mcp-go`:**
> 
> * timeout, no result return
> 
> ![Image](https://github.com/user-attachments/assets/d7125eaf-0e4a-4f7c-b732-ba4304ee3e0c)

@CeerDecy  How is SSE URL customization implemented? Why can't I modify it and only write the IP address? How to deploy a remote MCP server and access it through FQDNS?@CeerDecy: solved

Solution patch:
diff --git a/.gitignore b/.gitignore
index 5430d3b..b575ab6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,5 @@
 .aider*
 .env
-.idea
\ No newline at end of file
+.idea
+.opencode
+.claude
diff --git a/README.md b/README.md
index 332ab3d..594d49c 100644
--- a/README.md
+++ b/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
diff --git a/client/sse.go b/client/sse.go
index c26744a..d6eaf10 100644
--- a/client/sse.go
+++ b/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) {
diff --git a/client/transport/sse.go b/client/transport/sse.go
index 501036d..eda9446 100644
--- a/client/transport/sse.go
+++ b/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)
 	}
 
diff --git a/client/transport/stdio.go b/client/transport/stdio.go
index ffa0dff..3d9d832 100644
--- a/client/transport/stdio.go
+++ b/client/transport/stdio.go
@@ -133,7 +133,7 @@ func (c *Stdio) Close() error {
 	}
 	// cancel all in-flight request
 	close(c.done)
-	
+
 	if err := c.stdin.Close(); err != nil {
 		return fmt.Errorf("failed to close stdin: %w", err)
 	}
@@ -148,8 +148,8 @@ func (c *Stdio) Close() error {
 	return nil
 }
 
-// OnNotification registers a handler function to be called when notifications are received.
-// Multiple handlers can be registered and will be called in the order they were added.
+// SetNotificationHandler sets the handler function to be called when a notification is received.
+// Only one handler can be set at a time; setting a new one replaces the previous handler.
 func (c *Stdio) SetNotificationHandler(
 	handler func(notification mcp.JSONRPCNotification),
 ) {
@@ -208,7 +208,7 @@ func (c *Stdio) readResponses() {
 	}
 }
 
-// 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.
 // It creates a unique request ID, sends the request over stdin, and waits for
 // the corresponding response or context cancellation.
 // Returns the raw JSON response message or an error if the request fails.
@@ -261,7 +261,7 @@ func (c *Stdio) SendNotification(
 	if c.stdin == nil {
 		return fmt.Errorf("stdio client not started")
 	}
-	
+
 	notificationBytes, err := json.Marshal(notification)
 	if err != nil {
 		return fmt.Errorf("failed to marshal notification: %w", err)
diff --git a/client/transport/streamable_http.go b/client/transport/streamable_http.go
index 0e771c5..7f3a695 100644
--- a/client/transport/streamable_http.go
+++ b/client/transport/streamable_http.go
@@ -131,7 +131,7 @@ const (
 	headerKeySessionID = "Mcp-Session-Id"
 )
 
-// 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 *StreamableHTTP) SendRequest(
 	ctx context.Context,
@@ -356,6 +356,7 @@ func (c *StreamableHTTP) SendNotification(ctx context.Context, notification mcp.
 
 	// Set headers
 	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json, text/event-stream")	
 	if sessionID := c.sessionID.Load(); sessionID != "" {
 		req.Header.Set(headerKeySessionID, sessionID.(string))
 	}
diff --git a/mcp/tools.go b/mcp/tools.go
index b62cdd5..d4fde48 100644
--- a/mcp/tools.go
+++ b/mcp/tools.go
@@ -113,6 +113,23 @@ type ToolInputSchema struct {
 	Required   []string               `json:"required,omitempty"`
 }
 
+// MarshalJSON implements the json.Marshaler interface for ToolInputSchema.
+func (tis ToolInputSchema) MarshalJSON() ([]byte, error) {
+	m := make(map[string]interface{})
+	m["type"] = tis.Type
+
+	// Marshal Properties to '{}' rather than `nil` when its length equals zero
+	if tis.Properties != nil {
+		m["properties"] = tis.Properties
+	}
+
+	if len(tis.Required) > 0 {
+		m["required"] = tis.Required
+	}
+
+	return json.Marshal(m)
+}
+
 type ToolAnnotation struct {
 	// Human-readable title for the tool
 	Title string `json:"title,omitempty"`
diff --git a/mcp/types.go b/mcp/types.go
index 714eb48..c79baae 100644
--- a/mcp/types.go
+++ b/mcp/types.go
@@ -11,53 +11,53 @@ import (
 type MCPMethod string
 
 const (
-	// Initiates connection and negotiates protocol capabilities.
+	// MethodInitialize initiates connection and negotiates protocol capabilities.
 	// https://modelcontextprotocol.io/specification/2024-11-05/basic/lifecycle/#initialization
 	MethodInitialize MCPMethod = "initialize"
 
-	// Verifies connection liveness between client and server.
+	// MethodPing verifies connection liveness between client and server.
 	// https://modelcontextprotocol.io/specification/2024-11-05/basic/utilities/ping/
 	MethodPing MCPMethod = "ping"
 
-	// Lists all available server resources.
+	// MethodResourcesList lists all available server resources.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/resources/
 	MethodResourcesList MCPMethod = "resources/list"
 
-	// Provides URI templates for constructing resource URIs.
+	// MethodResourcesTemplatesList provides URI templates for constructing resource URIs.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/resources/
 	MethodResourcesTemplatesList MCPMethod = "resources/templates/list"
 
-	// Retrieves content of a specific resource by URI.
+	// MethodResourcesRead retrieves content of a specific resource by URI.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/resources/
 	MethodResourcesRead MCPMethod = "resources/read"
 
-	// Lists all available prompt templates.
+	// MethodPromptsList lists all available prompt templates.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/prompts/
 	MethodPromptsList MCPMethod = "prompts/list"
 
-	// Retrieves a specific prompt template with filled parameters.
+	// MethodPromptsGet retrieves a specific prompt template with filled parameters.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/prompts/
 	MethodPromptsGet MCPMethod = "prompts/get"
 
-	// Lists all available executable tools.
+	// MethodToolsList lists all available executable tools.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/tools/
 	MethodToolsList MCPMethod = "tools/list"
 
-	// Invokes a specific tool with provided parameters.
+	// MethodToolsCall invokes a specific tool with provided parameters.
 	// https://modelcontextprotocol.io/specification/2024-11-05/server/tools/
 	MethodToolsCall MCPMethod = "tools/call"
 
-	// Notifies when the list of available resources changes.
+	// MethodNotificationResourcesListChanged notifies when the list of available resources changes.
 	// https://modelcontextprotocol.io/specification/2025-03-26/server/resources#list-changed-notification
 	MethodNotificationResourcesListChanged = "notifications/resources/list_changed"
 
 	MethodNotificationResourceUpdated = "notifications/resources/updated"
 
-	// Notifies when the list of available prompt templates changes.
+	// MethodNotificationPromptsListChanged notifies when the list of available prompt templates changes.
 	// https://modelcontextprotocol.io/specification/2025-03-26/server/prompts#list-changed-notification
 	MethodNotificationPromptsListChanged = "notifications/prompts/list_changed"
 
-	// Notifies when the list of available tools changes.
+	// MethodNotificationToolsListChanged notifies when the list of available tools changes.
 	// https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/list_changed/
 	MethodNotificationToolsListChanged = "notifications/tools/list_changed"
 )
@@ -132,7 +132,7 @@ type NotificationParams struct {
 }
 
 // MarshalJSON implements custom JSON marshaling
-func (p NotificationParams) MarshalJSON() ([]byte, error) {
+func (p *NotificationParams) MarshalJSON() ([]byte, error) {
 	// Create a map to hold all fields
 	m := make(map[string]interface{})
 
@@ -863,7 +863,6 @@ type RootsListChangedNotification struct {
 	Notification
 }
 
-/* Client messages */
 // ClientRequest represents any request that can be sent from client to server.
 type ClientRequest interface{}
 
@@ -873,7 +872,6 @@ type ClientNotification interface{}
 // ClientResult represents any result that can be sent from client to server.
 type ClientResult interface{}
 
-/* Server messages */
 // ServerRequest represents any request that can be sent from server to client.
 type ServerRequest interface{}
 
diff --git a/mcp/utils.go b/mcp/utils.go
index 333d657..250357f 100644
--- a/mcp/utils.go
+++ b/mcp/utils.go
@@ -3,6 +3,7 @@ package mcp
 import (
 	"encoding/json"
 	"fmt"
+
 	"github.com/spf13/cast"
 )
 
@@ -125,6 +126,7 @@ func NewJSONRPCError(
 	}
 }
 
+// NewProgressNotification
 // Helper function for creating a progress notification
 func NewProgressNotification(
 	token ProgressToken,
@@ -155,6 +157,7 @@ func NewProgressNotification(
 	return notification
 }
 
+// NewLoggingMessageNotification
 // Helper function for creating a logging message notification
 func NewLoggingMessageNotification(
 	level LoggingLevel,
@@ -177,6 +180,7 @@ func NewLoggingMessageNotification(
 	}
 }
 
+// NewPromptMessage
 // Helper function to create a new PromptMessage
 func NewPromptMessage(role Role, content Content) PromptMessage {
 	return PromptMessage{
@@ -185,6 +189,7 @@ func NewPromptMessage(role Role, content Content) PromptMessage {
 	}
 }
 
+// NewTextContent
 // Helper function to create a new TextContent
 func NewTextContent(text string) TextContent {
 	return TextContent{
@@ -193,6 +198,7 @@ func NewTextContent(text string) TextContent {
 	}
 }
 
+// NewImageContent
 // Helper function to create a new ImageContent
 func NewImageContent(data, mimeType string) ImageContent {
 	return ImageContent{
@@ -202,6 +208,7 @@ func NewImageContent(data, mimeType string) ImageContent {
 	}
 }
 
+// NewEmbeddedResource
 // Helper function to create a new EmbeddedResource
 func NewEmbeddedResource(resource ResourceContents) EmbeddedResource {
 	return EmbeddedResource{
@@ -376,6 +383,7 @@ func NewInitializeResult(
 	}
 }
 
+// FormatNumberResult
 // Helper for formatting numbers in tool results
 func FormatNumberResult(value float64) *CallToolResult {
 	return NewToolResultText(fmt.Sprintf("%.2f", value))
@@ -433,6 +441,10 @@ func ParseContent(contentMap map[string]any) (Content, error) {
 }
 
 func ParseGetPromptResult(rawMessage *json.RawMessage) (*GetPromptResult, error) {
+	if rawMessage == nil {
+		return nil, fmt.Errorf("response is nil")
+	}
+
 	var jsonContent map[string]any
 	if err := json.Unmarshal(*rawMessage, &jsonContent); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
@@ -495,6 +507,10 @@ func ParseGetPromptResult(rawMessage *json.RawMessage) (*GetPromptResult, error)
 }
 
 func ParseCallToolResult(rawMessage *json.RawMessage) (*CallToolResult, error) {
+	if rawMessage == nil {
+		return nil, fmt.Errorf("response is nil")
+	}
+
 	var jsonContent map[string]any
 	if err := json.Unmarshal(*rawMessage, &jsonContent); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
@@ -573,6 +589,10 @@ func ParseResourceContents(contentMap map[string]any) (ResourceContents, error)
 }
 
 func ParseReadResourceResult(rawMessage *json.RawMessage) (*ReadResourceResult, error) {
+	if rawMessage == nil {
+		return nil, fmt.Errorf("response is nil")
+	}
+
 	var jsonContent map[string]any
 	if err := json.Unmarshal(*rawMessage, &jsonContent); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
diff --git a/server/errors.go b/server/errors.go
new file mode 100644
index 0000000..7ced5cf
--- /dev/null
+++ b/server/errors.go
@@ -0,0 +1,23 @@
+package server
+
+import (
+	"errors"
+)
+
+var (
+	// Common server errors
+	ErrUnsupported      = errors.New("not supported")
+	ErrResourceNotFound = errors.New("resource not found")
+	ErrPromptNotFound   = errors.New("prompt not found")
+	ErrToolNotFound     = errors.New("tool not found")
+
+	// Session-related errors
+	ErrSessionNotFound            = errors.New("session not found")
+	ErrSessionExists              = errors.New("session already exists")
+	ErrSessionNotInitialized      = errors.New("session not properly initialized")
+	ErrSessionDoesNotSupportTools = errors.New("session does not support per-session tools")
+
+	// Notification-related errors
+	ErrNotificationNotInitialized = errors.New("notification channel not initialized")
+	ErrNotificationChannelBlocked = errors.New("notification channel full or blocked")
+)
diff --git a/server/server.go b/server/server.go
index 47b7185..95831eb 100644
--- a/server/server.go
+++ b/server/server.go
@@ -5,7 +5,6 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"reflect"
 	"sort"
@@ -44,31 +43,22 @@ type ToolHandlerFunc func(ctx context.Context, request mcp.CallToolRequest) (*mc
 // ToolHandlerMiddleware is a middleware function that wraps a ToolHandlerFunc.
 type ToolHandlerMiddleware func(ToolHandlerFunc) ToolHandlerFunc
 
+// ToolFilterFunc is a function that filters tools based on context, typically using session information.
+type ToolFilterFunc func(ctx context.Context, tools []mcp.Tool) []mcp.Tool
+
 // ServerTool combines a Tool with its ToolHandlerFunc.
 type ServerTool struct {
 	Tool    mcp.Tool
 	Handler ToolHandlerFunc
 }
 
-// ClientSession represents an active session that can be used by MCPServer to interact with client.
-type ClientSession interface {
-	// Initialize marks session as fully initialized and ready for notifications
-	Initialize()
-	// Initialized returns if session is ready to accept notifications
-	Initialized() bool
-	// NotificationChannel provides a channel suitable for sending notifications to client.
-	NotificationChannel() chan<- mcp.JSONRPCNotification
-	// SessionID is a unique identifier used to track user session.
-	SessionID() string
-}
-
-// clientSessionKey is the context key for storing current client notification channel.
-type clientSessionKey struct{}
+// serverKey is the context key for storing the server instance
+type serverKey struct{}
 
-// ClientSessionFromContext retrieves current client notification context from context.
-func ClientSessionFromContext(ctx context.Context) ClientSession {
-	if session, ok := ctx.Value(clientSessionKey{}).(ClientSession); ok {
-		return session
+// ServerFromContext retrieves the MCPServer instance from a context
+func ServerFromContext(ctx context.Context) *MCPServer {
+	if srv, ok := ctx.Value(serverKey{}).(*MCPServer); ok {
+		return srv
 	}
 	return nil
 }
@@ -128,13 +118,6 @@ func (e *requestError) Unwrap() error {
 	return e.err
 }
 
-var (
-	ErrUnsupported      = errors.New("not supported")
-	ErrResourceNotFound = errors.New("resource not found")
-	ErrPromptNotFound   = errors.New("prompt not found")
-	ErrToolNotFound     = errors.New("tool not found")
-)
-
 // NotificationHandlerFunc handles incoming notifications.
 type NotificationHandlerFunc func(ctx context.Context, notification mcp.JSONRPCNotification)
 
@@ -148,6 +131,7 @@ type MCPServer struct {
 	middlewareMu           sync.RWMutex
 	notificationHandlersMu sync.RWMutex
 	capabilitiesMu         sync.RWMutex
+	toolFiltersMu          sync.RWMutex
 
 	name                   string
 	version                string
@@ -158,6 +142,7 @@ type MCPServer struct {
 	promptHandlers         map[string]PromptHandlerFunc
 	tools                  map[string]ServerTool
 	toolHandlerMiddlewares []ToolHandlerMiddleware
+	toolFilters            []ToolFilterFunc
 	notificationHandlers   map[string]NotificationHandlerFunc
 	capabilities           serverCapabilities
 	paginationLimit        *int
@@ -165,17 +150,6 @@ type MCPServer struct {
 	hooks                  *Hooks
 }
 
-// serverKey is the context key for storing the server instance
-type serverKey struct{}
-
-// ServerFromContext retrieves the MCPServer instance from a context
-func ServerFromContext(ctx context.Context) *MCPServer {
-	if srv, ok := ctx.Value(serverKey{}).(*MCPServer); ok {
-		return srv
-	}
-	return nil
-}
-
 // WithPaginationLimit sets the pagination limit for the server.
 func WithPaginationLimit(limit int) ServerOption {
 	return func(s *MCPServer) {
@@ -183,92 +157,6 @@ func WithPaginationLimit(limit int) ServerOption {
 	}
 }
 
-// WithContext sets the current client session and returns the provided context
-func (s *MCPServer) WithContext(
-	ctx context.Context,
-	session ClientSession,
-) context.Context {
-	return context.WithValue(ctx, clientSessionKey{}, session)
-}
-
-// RegisterSession saves session that should be notified in case if some server attributes changed.
-func (s *MCPServer) RegisterSession(
-	ctx context.Context,
-	session ClientSession,
-) error {
-	sessionID := session.SessionID()
-	if _, exists := s.sessions.LoadOrStore(sessionID, session); exists {
-		return fmt.Errorf("session %s is already registered", sessionID)
-	}
-	s.hooks.RegisterSession(ctx, session)
-	return nil
-}
-
-// UnregisterSession removes from storage session that is shut down.
-func (s *MCPServer) UnregisterSession(
-	ctx context.Context,
-	sessionID string,
-) {
-	session, _ := s.sessions.LoadAndDelete(sessionID)
-	s.hooks.UnregisterSession(ctx, session.(ClientSession))
-}
-
-// SendNotificationToAllClients sends a notification to all the currently active clients.
-func (s *MCPServer) SendNotificationToAllClients(
-	method string,
-	params map[string]any,
-) {
-	notification := mcp.JSONRPCNotification{
-		JSONRPC: mcp.JSONRPC_VERSION,
-		Notification: mcp.Notification{
-			Method: method,
-			Params: mcp.NotificationParams{
-				AdditionalFields: params,
-			},
-		},
-	}
-
-	s.sessions.Range(func(k, v any) bool {
-		if session, ok := v.(ClientSession); ok && session.Initialized() {
-			select {
-			case session.NotificationChannel() <- notification:
-			default:
-				// TODO: log blocked channel in the future versions
-			}
-		}
-		return true
-	})
-}
-
-// SendNotificationToClient sends a notification to the current client
-func (s *MCPServer) SendNotificationToClient(
-	ctx context.Context,
-	method string,
-	params map[string]any,
-) error {
-	session := ClientSessionFromContext(ctx)
-	if session == nil || !session.Initialized() {
-		return fmt.Errorf("notification channel not initialized")
-	}
-
-	notification := mcp.JSONRPCNotification{
-		JSONRPC: mcp.JSONRPC_VERSION,
-		Notification: mcp.Notification{
-			Method: method,
-			Params: mcp.NotificationParams{
-				AdditionalFields: params,
-			},
-		},
-	}
-
-	select {
-	case session.NotificationChannel() <- notification:
-		return nil
-	default:
-		return fmt.Errorf("notification channel full or blocked")
-	}
-}
-
 // serverCapabilities defines the supported features of the MCP server
 type serverCapabilities struct {
 	tools     *toolCapabilities
@@ -316,6 +204,17 @@ func WithToolHandlerMiddleware(
 	}
 }
 
+// WithToolFilter adds a filter function that will be applied to tools before they are returned in list_tools
+func WithToolFilter(
+	toolFilter ToolFilterFunc,
+) ServerOption {
+	return func(s *MCPServer) {
+		s.toolFiltersMu.Lock()
+		s.toolFilters = append(s.toolFilters, toolFilter)
+		s.toolFiltersMu.Unlock()
+	}
+}
+
 // WithRecovery adds a middleware that recovers from panics in tool handlers.
 func WithRecovery() ServerOption {
 	return WithToolHandlerMiddleware(func(next ToolHandlerFunc) ToolHandlerFunc {
@@ -421,11 +320,11 @@ func (s *MCPServer) AddResource(
 	}
 
 	s.resourcesMu.Lock()
-	defer s.resourcesMu.Unlock()
 	s.resources[resource.URI] = resourceEntry{
 		resource: resource,
 		handler:  handler,
 	}
+	s.resourcesMu.Unlock()
 
 	// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
 	if s.capabilities.resources.listChanged {
@@ -464,13 +363,12 @@ func (s *MCPServer) AddResourceTemplate(
 		s.capabilitiesMu.RUnlock()
 	}
 
-
 	s.resourcesMu.Lock()
-	defer s.resourcesMu.Unlock()
 	s.resourceTemplates[template.URITemplate.Raw()] = resourceTemplateEntry{
 		template: template,
 		handler:  handler,
 	}
+	s.resourcesMu.Unlock()
 
 	// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification
 	if s.capabilities.resources.listChanged {
@@ -495,9 +393,9 @@ func (s *MCPServer) AddPrompt(prompt mcp.Prompt, handler PromptHandlerFunc) {
 	}
 
 	s.promptsMu.Lock()
-	defer s.promptsMu.Unlock()
 	s.prompts[prompt.Name] = prompt
 	s.promptHandlers[prompt.Name] = handler
+	s.promptsMu.Unlock()
 
 	// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification.
 	if s.capabilities.prompts.listChanged {
@@ -867,6 +765,7 @@ func (s *MCPServer) handleListTools(
 	id interface{},
 	request mcp.ListToolsRequest,
 ) (*mcp.ListToolsResult, *requestError) {
+	// Get the base tools from the server
 	s.toolsMu.RLock()
 	tools := make([]mcp.Tool, 0, len(s.tools))
 
@@ -885,6 +784,49 @@ func (s *MCPServer) handleListTools(
 	}
 	s.toolsMu.RUnlock()
 
+	// Check if there are session-specific tools
+	session := ClientSessionFromContext(ctx)
+	if session != nil {
+		if sessionWithTools, ok := session.(SessionWithTools); ok {
+			if sessionTools := sessionWithTools.GetSessionTools(); sessionTools != nil {
+				// Override or add session-specific tools
+				// We need to create a map first to merge the tools properly
+				toolMap := make(map[string]mcp.Tool)
+
+				// Add global tools first
+				for _, tool := range tools {
+					toolMap[tool.Name] = tool
+				}
+
+				// Then override with session-specific tools
+				for name, serverTool := range sessionTools {
+					toolMap[name] = serverTool.Tool
+				}
+
+				// Convert back to slice
+				tools = make([]mcp.Tool, 0, len(toolMap))
+				for _, tool := range toolMap {
+					tools = append(tools, tool)
+				}
+
+				// Sort again to maintain consistent ordering
+				sort.Slice(tools, func(i, j int) bool {
+					return tools[i].Name < tools[j].Name
+				})
+			}
+		}
+	}
+
+	// Apply tool filters if any are defined
+	s.toolFiltersMu.RLock()
+	if len(s.toolFilters) > 0 {
+		for _, filter := range s.toolFilters {
+			tools = filter(ctx, tools)
+		}
+	}
+	s.toolFiltersMu.RUnlock()
+
+	// Apply pagination
 	toolsToReturn, nextCursor, err := listByPagination[mcp.Tool](ctx, s, request.Params.Cursor, tools)
 	if err != nil {
 		return nil, &requestError{
@@ -893,6 +835,7 @@ func (s *MCPServer) handleListTools(
 			err:  err,
 		}
 	}
+
 	result := mcp.ListToolsResult{
 		Tools: toolsToReturn,
 		PaginatedResult: mcp.PaginatedResult{
@@ -907,9 +850,29 @@ func (s *MCPServer) handleToolCall(
 	id interface{},
 	request mcp.CallToolRequest,
 ) (*mcp.CallToolResult, *requestError) {
-	s.toolsMu.RLock()
-	tool, ok := s.tools[request.Params.Name]
-	s.toolsMu.RUnlock()
+	// First check session-specific tools
+	var tool ServerTool
+	var ok bool
+
+	session := ClientSessionFromContext(ctx)
+	if session != nil {
+		if sessionWithTools, ok := session.(SessionWithTools); ok {
+			if sessionTools := sessionWithTools.GetSessionTools(); sessionTools != nil {
+				var sessionOk bool
+				tool, sessionOk = sessionTools[request.Params.Name]
+				if sessionOk {
+					ok = true
+				}
+			}
+		}
+	}
+
+	// If not found in session tools, check global tools
+	if !ok {
+		s.toolsMu.RLock()
+		tool, ok = s.tools[request.Params.Name]
+		s.toolsMu.RUnlock()
+	}
 
 	if !ok {
 		return nil, &requestError{
diff --git a/server/session.go b/server/session.go
new file mode 100644
index 0000000..68ae0bb
--- /dev/null
+++ b/server/session.go
@@ -0,0 +1,318 @@
+package server
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/mark3labs/mcp-go/mcp"
+)
+
+// ClientSession represents an active session that can be used by MCPServer to interact with client.
+type ClientSession interface {
+	// Initialize marks session as fully initialized and ready for notifications
+	Initialize()
+	// Initialized returns if session is ready to accept notifications
+	Initialized() bool
+	// NotificationChannel provides a channel suitable for sending notifications to client.
+	NotificationChannel() chan<- mcp.JSONRPCNotification
+	// SessionID is a unique identifier used to track user session.
+	SessionID() string
+}
+
+// SessionWithTools is an extension of ClientSession that can store session-specific tool data
+type SessionWithTools interface {
+	ClientSession
+	// GetSessionTools returns the tools specific to this session, if any
+	// This method must be thread-safe for concurrent access
+	GetSessionTools() map[string]ServerTool
+	// SetSessionTools sets tools specific to this session
+	// This method must be thread-safe for concurrent access
+	SetSessionTools(tools map[string]ServerTool)
+}
+
+// clientSessionKey is the context key for storing current client notification channel.
+type clientSessionKey struct{}
+
+// ClientSessionFromContext retrieves current client notification context from context.
+func ClientSessionFromContext(ctx context.Context) ClientSession {
+	if session, ok := ctx.Value(clientSessionKey{}).(ClientSession); ok {
+		return session
+	}
+	return nil
+}
+
+// WithContext sets the current client session and returns the provided context
+func (s *MCPServer) WithContext(
+	ctx context.Context,
+	session ClientSession,
+) context.Context {
+	return context.WithValue(ctx, clientSessionKey{}, session)
+}
+
+// RegisterSession saves session that should be notified in case if some server attributes changed.
+func (s *MCPServer) RegisterSession(
+	ctx context.Context,
+	session ClientSession,
+) error {
+	sessionID := session.SessionID()
+	if _, exists := s.sessions.LoadOrStore(sessionID, session); exists {
+		return ErrSessionExists
+	}
+	s.hooks.RegisterSession(ctx, session)
+	return nil
+}
+
+// UnregisterSession removes from storage session that is shut down.
+func (s *MCPServer) UnregisterSession(
+	ctx context.Context,
+	sessionID string,
+) {
+	sessionValue, ok := s.sessions.LoadAndDelete(sessionID)
+	if !ok {
+		return
+	}
+	if session, ok := sessionValue.(ClientSession); ok {
+		s.hooks.UnregisterSession(ctx, session)
+	}
+}
+
+// SendNotificationToAllClients sends a notification to all the currently active clients.
+func (s *MCPServer) SendNotificationToAllClients(
+	method string,
+	params map[string]any,
+) {
+	notification := mcp.JSONRPCNotification{
+		JSONRPC: mcp.JSONRPC_VERSION,
+		Notification: mcp.Notification{
+			Method: method,
+			Params: mcp.NotificationParams{
+				AdditionalFields: params,
+			},
+		},
+	}
+
+	s.sessions.Range(func(k, v any) bool {
+		if session, ok := v.(ClientSession); ok && session.Initialized() {
+			select {
+			case session.NotificationChannel() <- notification:
+				// Successfully sent notification
+			default:
+				// Channel is blocked, if there's an error hook, use it
+				if s.hooks != nil && len(s.hooks.OnError) > 0 {
+					err := ErrNotificationChannelBlocked
+					// Copy hooks pointer to local variable to avoid race condition
+					hooks := s.hooks
+					go func(sessionID string, hooks *Hooks) {
+						ctx := context.Background()
+						// Use the error hook to report the blocked channel
+						hooks.onError(ctx, nil, "notification", map[string]interface{}{
+							"method":    method,
+							"sessionID": sessionID,
+						}, fmt.Errorf("notification channel blocked for session %s: %w", sessionID, err))
+					}(session.SessionID(), hooks)
+				}
+			}
+		}
+		return true
+	})
+}
+
+// SendNotificationToClient sends a notification to the current client
+func (s *MCPServer) SendNotificationToClient(
+	ctx context.Context,
+	method string,
+	params map[string]any,
+) error {
+	session := ClientSessionFromContext(ctx)
+	if session == nil || !session.Initialized() {
+		return ErrNotificationNotInitialized
+	}
+
+	notification := mcp.JSONRPCNotification{
+		JSONRPC: mcp.JSONRPC_VERSION,
+		Notification: mcp.Notification{
+			Method: method,
+			Params: mcp.NotificationParams{
+				AdditionalFields: params,
+			},
+		},
+	}
+
+	select {
+	case session.NotificationChannel() <- notification:
+		return nil
+	default:
+		// Channel is blocked, if there's an error hook, use it
+		if s.hooks != nil && len(s.hooks.OnError) > 0 {
+			err := ErrNotificationChannelBlocked
+			// Copy hooks pointer to local variable to avoid race condition
+			hooks := s.hooks
+			go func(sessionID string, hooks *Hooks) {
+				// Use the error hook to report the blocked channel
+				hooks.onError(ctx, nil, "notification", map[string]interface{}{
+					"method":    method,
+					"sessionID": sessionID,
+				}, fmt.Errorf("notification channel blocked for session %s: %w", sessionID, err))
+			}(session.SessionID(), hooks)
+		}
+		return ErrNotificationChannelBlocked
+	}
+}
+
+// SendNotificationToSpecificClient sends a notification to a specific client by session ID
+func (s *MCPServer) SendNotificationToSpecificClient(
+	sessionID string,
+	method string,
+	params map[string]any,
+) error {
+	sessionValue, ok := s.sessions.Load(sessionID)
+	if !ok {
+		return ErrSessionNotFound
+	}
+
+	session, ok := sessionValue.(ClientSession)
+	if !ok || !session.Initialized() {
+		return ErrSessionNotInitialized
+	}
+
+	notification := mcp.JSONRPCNotification{
+		JSONRPC: mcp.JSONRPC_VERSION,
+		Notification: mcp.Notification{
+			Method: method,
+			Params: mcp.NotificationParams{
+				AdditionalFields: params,
+			},
+		},
+	}
+
+	select {
+	case session.NotificationChannel() <- notification:
+		return nil
+	default:
+		// Channel is blocked, if there's an error hook, use it
+		if s.hooks != nil && len(s.hooks.OnError) > 0 {
+			err := ErrNotificationChannelBlocked
+			ctx := context.Background()
+			// Copy hooks pointer to local variable to avoid race condition
+			hooks := s.hooks
+			go func(sID string, hooks *Hooks) {
+				// Use the error hook to report the blocked channel
+				hooks.onError(ctx, nil, "notification", map[string]interface{}{
+					"method":    method,
+					"sessionID": sID,
+				}, fmt.Errorf("notification channel blocked for session %s: %w", sID, err))
+			}(sessionID, hooks)
+		}
+		return ErrNotificationChannelBlocked
+	}
+}
+
+// AddSessionTool adds a tool for a specific session
+func (s *MCPServer) AddSessionTool(sessionID string, tool mcp.Tool, handler ToolHandlerFunc) error {
+	return s.AddSessionTools(sessionID, ServerTool{Tool: tool, Handler: handler})
+}
+
+// AddSessionTools adds tools for a specific session
+func (s *MCPServer) AddSessionTools(sessionID string, tools ...ServerTool) error {
+	sessionValue, ok := s.sessions.Load(sessionID)
+	if !ok {
+		return ErrSessionNotFound
+	}
+
+	session, ok := sessionValue.(SessionWithTools)
+	if !ok {
+		return ErrSessionDoesNotSupportTools
+	}
+
+	// Get existing tools (this should return a thread-safe copy)
+	sessionTools := session.GetSessionTools()
+
+	// Create a new map to avoid concurrent modification issues
+	newSessionTools := make(map[string]ServerTool, len(sessionTools)+len(tools))
+
+	// Copy existing tools
+	if sessionTools != nil {
+		for k, v := range sessionTools {
+			newSessionTools[k] = v
+		}
+	}
+
+	// Add new tools
+	for _, tool := range tools {
+		newSessionTools[tool.Tool.Name] = tool
+	}
+
+	// Set the tools (this should be thread-safe)
+	session.SetSessionTools(newSessionTools)
+
+	// Send notification only to this session
+	if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
+		// Log the error but don't fail the operation
+		// The tools were successfully added, but notification failed
+		if s.hooks != nil && len(s.hooks.OnError) > 0 {
+			hooks := s.hooks
+			go func(sID string, hooks *Hooks) {
+				ctx := context.Background()
+				hooks.onError(ctx, nil, "notification", map[string]interface{}{
+					"method":    "notifications/tools/list_changed",
+					"sessionID": sID,
+				}, fmt.Errorf("failed to send notification after adding tools: %w", err))
+			}(sessionID, hooks)
+		}
+	}
+
+	return nil
+}
+
+// DeleteSessionTools removes tools from a specific session
+func (s *MCPServer) DeleteSessionTools(sessionID string, names ...string) error {
+	sessionValue, ok := s.sessions.Load(sessionID)
+	if !ok {
+		return ErrSessionNotFound
+	}
+
+	session, ok := sessionValue.(SessionWithTools)
+	if !ok {
+		return ErrSessionDoesNotSupportTools
+	}
+
+	// Get existing tools (this should return a thread-safe copy)
+	sessionTools := session.GetSessionTools()
+	if sessionTools == nil {
+		return nil
+	}
+
+	// Create a new map to avoid concurrent modification issues
+	newSessionTools := make(map[string]ServerTool, len(sessionTools))
+
+	// Copy existing tools except those being deleted
+	for k, v := range sessionTools {
+		newSessionTools[k] = v
+	}
+
+	// Remove specified tools
+	for _, name := range names {
+		delete(newSessionTools, name)
+	}
+
+	// Set the tools (this should be thread-safe)
+	session.SetSessionTools(newSessionTools)
+
+	// Send notification only to this session
+	if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
+		// Log the error but don't fail the operation
+		// The tools were successfully deleted, but notification failed
+		if s.hooks != nil && len(s.hooks.OnError) > 0 {
+			hooks := s.hooks
+			go func(sID string, hooks *Hooks) {
+				ctx := context.Background()
+				hooks.onError(ctx, nil, "notification", map[string]interface{}{
+					"method":    "notifications/tools/list_changed",
+					"sessionID": sID,
+				}, fmt.Errorf("failed to send notification after deleting tools: %w", err))
+			}(sessionID, hooks)
+		}
+	}
+
+	return nil
+}
diff --git a/server/sse.go b/server/sse.go
index 9a41915..382664d 100644
--- a/server/sse.go
+++ b/server/sse.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -68,6 +69,8 @@ type SSEServer struct {
 	keepAliveInterval time.Duration
 
 	mu sync.RWMutex
+
+	appendQueryToMessageEndpoint bool
 }
 
 // SSEOption defines a function type for configuring SSEServer
@@ -96,7 +99,7 @@ func WithBaseURL(baseURL string) SSEOption {
 	}
 }
 
-// Add a new option for setting base path
+// WithBasePath adds a new option for setting base path
 func WithBasePath(basePath string) SSEOption {
 	return func(s *SSEServer) {
 		// Ensure the path starts with / and doesn't end with /
@@ -114,6 +117,17 @@ func WithMessageEndpoint(endpoint string) SSEOption {
 	}
 }
 
+// WithAppendQueryToMessageEndpoint configures the SSE server to append the original request's
+// query parameters to the message endpoint URL that is sent to clients during the SSE connection
+// initialization. This is useful when you need to preserve query parameters from the initial
+// SSE connection request and carry them over to subsequent message requests, maintaining
+// context or authentication details across the communication channel.
+func WithAppendQueryToMessageEndpoint() SSEOption {
+	return func(s *SSEServer) {
+		s.appendQueryToMessageEndpoint = true
+	}
+}
+
 // WithUseFullURLForMessageEndpoint controls whether the SSE server returns a complete URL (including baseURL)
 // or just the path portion for the message endpoint. Set to false when clients will concatenate
 // the baseURL themselves to avoid malformed URLs like "http://localhost/mcphttp://localhost/mcp/message".
@@ -150,7 +164,7 @@ func WithKeepAlive(keepAlive bool) SSEOption {
 	}
 }
 
-// WithContextFunc sets a function that will be called to customise the context
+// WithSSEContextFunc sets a function that will be called to customise the context
 // to the server using the incoming request.
 func WithSSEContextFunc(fn SSEContextFunc) SSEOption {
 	return func(s *SSEServer) {
@@ -190,12 +204,21 @@ func NewTestServer(server *MCPServer, opts ...SSEOption) *httptest.Server {
 // It sets up HTTP handlers for SSE and message endpoints.
 func (s *SSEServer) Start(addr string) error {
 	s.mu.Lock()
-	s.srv = &http.Server{
-		Addr:    addr,
-		Handler: s,
-	}
-	s.mu.Unlock()
+	defer s.mu.Unlock()
 
+	if s.srv == nil {
+		s.srv = &http.Server{
+			Addr:		addr,
+			Handler:	s,
+		}
+	} else {
+		if s.srv.Addr == "" {
+			s.srv.Addr = addr
+		} else if s.srv.Addr != addr {
+			return fmt.Errorf("conflicting listen address: WithHTTPServer(%q) vs Start(%q)", s.srv.Addr, addr)
+		}
+	}
+	
 	return s.srv.ListenAndServe()
 }
 
@@ -308,7 +331,11 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// Send the initial endpoint event
-	fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", s.GetMessageEndpointForClient(sessionID))
+	endpoint := s.GetMessageEndpointForClient(sessionID)
+	if s.appendQueryToMessageEndpoint && len(r.URL.RawQuery) > 0 {
+		endpoint += "&" + r.URL.RawQuery
+	}
+	fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", endpoint)
 	flusher.Flush()
 
 	// Main event loop - this runs in the HTTP handler goroutine
@@ -338,7 +365,7 @@ func (s *SSEServer) GetMessageEndpointForClient(sessionID string) string {
 }
 
 // handleMessage processes incoming JSON-RPC messages from clients and sends responses
-// back through both the SSE connection and HTTP response.
+// back through the SSE connection and 202 code to HTTP response.
 func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodPost {
 		s.writeJSONRPCError(w, nil, mcp.INVALID_REQUEST, "Method not allowed")
@@ -370,31 +397,37 @@ func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Process message through MCPServer
-	response := s.server.HandleMessage(ctx, rawMessage)
+	// quick return request, send 202 Accepted with no body, then deal the message and sent response via SSE
+	w.WriteHeader(http.StatusAccepted)
 
-	// Only send response if there is one (not for notifications)
-	if response != nil {
-		eventData, _ := json.Marshal(response)
+	go func() {
+		// Process message through MCPServer
+		response := s.server.HandleMessage(ctx, rawMessage)
+
+		// Only send response if there is one (not for notifications)
+		if response != nil {
+			var message string
+			if eventData, err := json.Marshal(response); err != nil {
+				// If there is an error marshalling the response, send a generic error response
+				log.Printf("failed to marshal response: %v", err)
+				message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n")
+				return
+			} else {
+				message = fmt.Sprintf("event: message\ndata: %s\n\n", eventData)
+			}
 
-		// Queue the event for sending via SSE
-		select {
-		case session.eventQueue <- fmt.Sprintf("event: message\ndata: %s\n\n", eventData):
-			// Event queued successfully
-		case <-session.done:
-			// Session is closed, don't try to queue
-		default:
-			// Queue is full, could log this
+			// Queue the event for sending via SSE
+			select {
+			case session.eventQueue <- message:
+				// Event queued successfully
+			case <-session.done:
+				// Session is closed, don't try to queue
+			default:
+				// Queue is full, log this situation
+				log.Printf("Event queue full for session %s", sessionID)
+			}
 		}
-
-		// Send HTTP response
-		w.Header().Set("Content-Type", "application/json")
-		w.WriteHeader(http.StatusAccepted)
-		json.NewEncoder(w).Encode(response)
-	} else {
-		// For notifications, just send 202 Accepted with no body
-		w.WriteHeader(http.StatusAccepted)
-	}
+	}()
 }
 
 // writeJSONRPCError writes a JSON-RPC error response with the given error details.
diff --git a/server/stdio.go b/server/stdio.go
index 0de9f34..79407e5 100644
--- a/server/stdio.go
+++ b/server/stdio.go
@@ -40,7 +40,7 @@ func WithErrorLogger(logger *log.Logger) StdioOption {
 	}
 }
 
-// WithContextFunc sets a function that will be called to customise the context
+// WithStdioContextFunc sets a function that will be called to customise the context
 // to the server. Note that the stdio server uses the same context for all requests,
 // so this function will only be called once per server instance.
 func WithStdioContextFunc(fn StdioContextFunc) StdioOption {



Test patch:
diff --git a/client/stdio_test.go b/client/stdio_test.go
index 700398d..7bffa3b 100644
--- a/client/stdio_test.go
+++ b/client/stdio_test.go
@@ -26,6 +26,9 @@ func compileTestServer(outputPath string) error {
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("compilation failed: %v\nOutput: %s", err, output)
 	}
+	if _, err := os.Stat(outputPath); os.IsNotExist(err) {
+		return fmt.Errorf("mock server binary not found at %s after compilation", outputPath)
+	}
 	return nil
 }
 
diff --git a/client/transport/sse_test.go b/client/transport/sse_test.go
index 0c4dff6..b8b59d0 100644
--- a/client/transport/sse_test.go
+++ b/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()
diff --git a/server/session_test.go b/server/session_test.go
new file mode 100644
index 0000000..d1d0bc7
--- /dev/null
+++ b/server/session_test.go
@@ -0,0 +1,562 @@
+package server
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// sessionTestClient implements the basic ClientSession interface for testing
+type sessionTestClient struct {
+	sessionID           string
+	notificationChannel chan mcp.JSONRPCNotification
+	initialized         bool
+}
+
+func (f sessionTestClient) SessionID() string {
+	return f.sessionID
+}
+
+func (f sessionTestClient) NotificationChannel() chan<- mcp.JSONRPCNotification {
+	return f.notificationChannel
+}
+
+// Initialize marks the session as initialized
+// This implementation properly sets the initialized flag to true
+// as required by the interface contract
+func (f *sessionTestClient) Initialize() {
+	f.initialized = true
+}
+
+// Initialized returns whether the session has been initialized
+func (f sessionTestClient) Initialized() bool {
+	return f.initialized
+}
+
+// sessionTestClientWithTools implements the SessionWithTools interface for testing
+type sessionTestClientWithTools struct {
+	sessionID           string
+	notificationChannel chan mcp.JSONRPCNotification
+	initialized         bool
+	sessionTools        map[string]ServerTool
+	mu                  sync.RWMutex // Mutex to protect concurrent access to sessionTools
+}
+
+func (f *sessionTestClientWithTools) SessionID() string {
+	return f.sessionID
+}
+
+func (f *sessionTestClientWithTools) NotificationChannel() chan<- mcp.JSONRPCNotification {
+	return f.notificationChannel
+}
+
+func (f *sessionTestClientWithTools) Initialize() {
+	f.initialized = true
+}
+
+func (f *sessionTestClientWithTools) Initialized() bool {
+	return f.initialized
+}
+
+func (f *sessionTestClientWithTools) GetSessionTools() map[string]ServerTool {
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+
+	// Return a copy of the map to prevent concurrent modification
+	if f.sessionTools == nil {
+		return nil
+	}
+
+	toolsCopy := make(map[string]ServerTool, len(f.sessionTools))
+	for k, v := range f.sessionTools {
+		toolsCopy[k] = v
+	}
+	return toolsCopy
+}
+
+func (f *sessionTestClientWithTools) SetSessionTools(tools map[string]ServerTool) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	// Create a copy of the map to prevent concurrent modification
+	if tools == nil {
+		f.sessionTools = nil
+		return
+	}
+
+	toolsCopy := make(map[string]ServerTool, len(tools))
+	for k, v := range tools {
+		toolsCopy[k] = v
+	}
+	f.sessionTools = toolsCopy
+}
+
+// Verify that both implementations satisfy their respective interfaces
+var _ ClientSession = &sessionTestClient{}
+var _ SessionWithTools = &sessionTestClientWithTools{}
+
+func TestSessionWithTools_Integration(t *testing.T) {
+	server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
+
+	// Create session-specific tools
+	sessionTool := ServerTool{
+		Tool: mcp.NewTool("session-tool"),
+		Handler: func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			return mcp.NewToolResultText("session-tool result"), nil
+		},
+	}
+
+	// Create a session with tools
+	session := &sessionTestClientWithTools{
+		sessionID:           "session-1",
+		notificationChannel: make(chan mcp.JSONRPCNotification, 10),
+		initialized:         true,
+		sessionTools: map[string]ServerTool{
+			"session-tool": sessionTool,
+		},
+	}
+
+	// Register the session
+	err := server.RegisterSession(context.Background(), session)
+	require.NoError(t, err)
+
+	// Test that we can access the session-specific tool
+	testReq := mcp.CallToolRequest{}
+	testReq.Params.Name = "session-tool"
+	testReq.Params.Arguments = map[string]interface{}{}
+
+	// Call using session context
+	sessionCtx := server.WithContext(context.Background(), session)
+
+	// Check if the session was stored in the context correctly
+	s := ClientSessionFromContext(sessionCtx)
+	require.NotNil(t, s, "Session should be available from context")
+	assert.Equal(t, session.SessionID(), s.SessionID(), "Session ID should match")
+
+	// Check if the session can be cast to SessionWithTools
+	swt, ok := s.(SessionWithTools)
+	require.True(t, ok, "Session should implement SessionWithTools")
+
+	// Check if the tools are accessible
+	tools := swt.GetSessionTools()
+	require.NotNil(t, tools, "Session tools should be available")
+	require.Contains(t, tools, "session-tool", "Session should have session-tool")
+
+	// Test session tool access with session context
+	t.Run("test session tool access", func(t *testing.T) {
+		// First test directly getting the tool from session tools
+		tool, exists := tools["session-tool"]
+		require.True(t, exists, "Session tool should exist in the map")
+		require.NotNil(t, tool, "Session tool should not be nil")
+
+		// Now test calling directly with the handler
+		result, err := tool.Handler(sessionCtx, testReq)
+		require.NoError(t, err, "No error calling session tool handler directly")
+		require.NotNil(t, result, "Result should not be nil")
+		require.Len(t, result.Content, 1, "Result should have one content item")
+
+		textContent, ok := result.Content[0].(mcp.TextContent)
+		require.True(t, ok, "Content should be TextContent")
+		assert.Equal(t, "session-tool result", textContent.Text, "Result text should match")
+	})
+}
+
+func TestMCPServer_ToolsWithSessionTools(t *testing.T) {
+	// Basic test to verify that session-specific tools are returned correctly in a tools list
+	server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
+
+	// Add global tools
+	server.AddTools(
+		ServerTool{Tool: mcp.NewTool("global-tool-1")},
+		ServerTool{Tool: mcp.NewTool("global-tool-2")},
+	)
+
+	// Create a session with tools
+	session := &sessionTestClientWithTools{
+		sessionID:           "session-1",
+		notificationChannel: make(chan mcp.JSONRPCNotification, 10),
+		initialized:         true,
+		sessionTools: map[string]ServerTool{
+			"session-tool-1": {Tool: mcp.NewTool("session-tool-1")},
+			"global-tool-1":  {Tool: mcp.NewTool("global-tool-1", mcp.WithDescription("Overridden"))},
+		},
+	}
+
+	// Register the session
+	err := server.RegisterSession(context.Background(), session)
+	require.NoError(t, err)
+
+	// List tools with session context
+	sessionCtx := server.WithContext(context.Background(), session)
+	resp := server.HandleMessage(sessionCtx, []byte(`{
+		"jsonrpc": "2.0",
+		"id": 1,
+		"method": "tools/list"
+	}`))
+
+	jsonResp, ok := resp.(mcp.JSONRPCResponse)
+	require.True(t, ok, "Response should be a JSONRPCResponse")
+
+	result, ok := jsonResp.Result.(mcp.ListToolsResult)
+	require.True(t, ok, "Result should be a ListToolsResult")
+
+	// Should have 3 tools - 2 global tools (one overridden) and 1 session-specific tool
+	assert.Len(t, result.Tools, 3, "Should have 3 tools")
+
+	// Find the overridden tool and verify its description
+	var found bool
+	for _, tool := range result.Tools {
+		if tool.Name == "global-tool-1" {
+			assert.Equal(t, "Overridden", tool.Description, "Global tool should be overridden")
+			found = true
+			break
+		}
+	}
+	assert.True(t, found, "Should find the overridden global tool")
+}
+
+func TestMCPServer_AddSessionTools(t *testing.T) {
+	server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
+	ctx := context.Background()
+
+	// Create a session
+	sessionChan := make(chan mcp.JSONRPCNotification, 10)
+	session := &sessionTestClientWithTools{
+		sessionID:           "session-1",
+		notificationChannel: sessionChan,
+		initialized:         true,
+	}
+
+	// Register the session
+	err := server.RegisterSession(ctx, session)
+	require.NoError(t, err)
+
+	// Add session-specific tools
+	err = server.AddSessionTools(session.SessionID(),
+		ServerTool{Tool: mcp.NewTool("session-tool")},
+	)
+	require.NoError(t, err)
+
+	// Check that notification was sent
+	select {
+	case notification := <-sessionChan:
+		assert.Equal(t, "notifications/tools/list_changed", notification.Method)
+	case <-time.After(100 * time.Millisecond):
+		t.Error("Expected notification not received")
+	}
+
+	// Verify tool was added to session
+	assert.Len(t, session.GetSessionTools(), 1)
+	assert.Contains(t, session.GetSessionTools(), "session-tool")
+}
+
+func TestMCPServer_AddSessionTool(t *testing.T) {
+	server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
+	ctx := context.Background()
+
+	// Create a session
+	sessionChan := make(chan mcp.JSONRPCNotification, 10)
+	session := &sessionTestClientWithTools{
+		sessionID:           "session-1",
+		notificationChannel: sessionChan,
+		initialized:         true,
+	}
+
+	// Register the session
+	err := server.RegisterSession(ctx, session)
+	require.NoError(t, err)
+
+	// Add session-specific tool using the new helper method
+	err = server.AddSessionTool(
+		session.SessionID(),
+		mcp.NewTool("session-tool-helper"),
+		func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			return mcp.NewToolResultText("helper result"), nil
+		},
+	)
+	require.NoError(t, err)
+
+	// Check that notification was sent
+	select {
+	case notification := <-sessionChan:
+		assert.Equal(t, "notifications/tools/list_changed", notification.Method)
+	case <-time.After(100 * time.Millisecond):
+		t.Error("Expected notification not received")
+	}
+
+	// Verify tool was added to session
+	assert.Len(t, session.GetSessionTools(), 1)
+	assert.Contains(t, session.GetSessionTools(), "session-tool-helper")
+}
+
+func TestMCPServer_DeleteSessionTools(t *testing.T) {
+	server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
+	ctx := context.Background()
+
+	// Create a session with tools
+	sessionChan := make(chan mcp.JSONRPCNotification, 10)
+	session := &sessionTestClientWithTools{
+		sessionID:           "session-1",
+		notificationChannel: sessionChan,
+		initialized:         true,
+		sessionTools: map[string]ServerTool{
+			"session-tool-1": {
+				Tool: mcp.NewTool("session-tool-1"),
+			},
+			"session-tool-2": {
+				Tool: mcp.NewTool("session-tool-2"),
+			},
+		},
+	}
+
+	// Register the session
+	err := server.RegisterSession(ctx, session)
+	require.NoError(t, err)
+
+	// Delete one of the session tools
+	err = server.DeleteSessionTools(session.SessionID(), "session-tool-1")
+	require.NoError(t, err)
+
+	// Check that notification was sent
+	select {
+	case notification := <-sessionChan:
+		assert.Equal(t, "notifications/tools/list_changed", notification.Method)
+	case <-time.After(100 * time.Millisecond):
+		t.Error("Expected notification not received")
+	}
+
+	// Verify tool was removed from session
+	assert.Len(t, session.GetSessionTools(), 1)
+	assert.NotContains(t, session.GetSessionTools(), "session-tool-1")
+	assert.Contains(t, session.GetSessionTools(), "session-tool-2")
+}
+
+func TestMCPServer_ToolFiltering(t *testing.T) {
+	// Create a filter that filters tools by prefix
+	filterByPrefix := func(prefix string) ToolFilterFunc {
+		return func(ctx context.Context, tools []mcp.Tool) []mcp.Tool {
+			var filtered []mcp.Tool
+			for _, tool := range tools {
+				if len(tool.Name) >= len(prefix) && tool.Name[:len(prefix)] == prefix {
+					filtered = append(filtered, tool)
+				}
+			}
+			return filtered
+		}
+	}
+
+	// Create a server with a tool filter
+	server := NewMCPServer("test-server", "1.0.0",
+		WithToolCapabilities(true),
+		WithToolFilter(filterByPrefix("allow-")),
+	)
+
+	// Add tools with different prefixes
+	server.AddTools(
+		ServerTool{Tool: mcp.NewTool("allow-tool-1")},
+		ServerTool{Tool: mcp.NewTool("allow-tool-2")},
+		ServerTool{Tool: mcp.NewTool("deny-tool-1")},
+		ServerTool{Tool: mcp.NewTool("deny-tool-2")},
+	)
+
+	// Create a session with tools
+	session := &sessionTestClientWithTools{
+		sessionID:           "session-1",
+		notificationChannel: make(chan mcp.JSONRPCNotification, 10),
+		initialized:         true,
+		sessionTools: map[string]ServerTool{
+			"allow-session-tool": {
+				Tool: mcp.NewTool("allow-session-tool"),
+			},
+			"deny-session-tool": {
+				Tool: mcp.NewTool("deny-session-tool"),
+			},
+		},
+	}
+
+	// Register the session
+	err := server.RegisterSession(context.Background(), session)
+	require.NoError(t, err)
+
+	// List tools with session context
+	sessionCtx := server.WithContext(context.Background(), session)
+	response := server.HandleMessage(sessionCtx, []byte(`{
+		"jsonrpc": "2.0",
+		"id": 1,
+		"method": "tools/list"
+	}`))
+	resp, ok := response.(mcp.JSONRPCResponse)
+	require.True(t, ok)
+
+	result, ok := resp.Result.(mcp.ListToolsResult)
+	require.True(t, ok)
+
+	// Should only include tools with the "allow-" prefix
+	assert.Len(t, result.Tools, 3)
+
+	// Verify all tools start with "allow-"
+	for _, tool := range result.Tools {
+		assert.True(t, len(tool.Name) >= 6 && tool.Name[:6] == "allow-",
+			"Tool should start with 'allow-', got: %s", tool.Name)
+	}
+}
+
+func TestMCPServer_SendNotificationToSpecificClient(t *testing.T) {
+	server := NewMCPServer("test-server", "1.0.0")
+
+	session1Chan := make(chan mcp.JSONRPCNotification, 10)
+	session1 := &sessionTestClient{
+		sessionID:           "session-1",
+		notificationChannel: session1Chan,
+	}
+	session1.Initialize()
+
+	session2Chan := make(chan mcp.JSONRPCNotification, 10)
+	session2 := &sessionTestClient{
+		sessionID:           "session-2",
+		notificationChannel: session2Chan,
+	}
+	session2.Initialize()
+
+	session3 := &sessionTestClient{
+		sessionID:           "session-3",
+		notificationChannel: make(chan mcp.JSONRPCNotification, 10),
+		initialized:         false, // Not initialized - deliberately not calling Initialize()
+	}
+
+	// Register sessions
+	err := server.RegisterSession(context.Background(), session1)
+	require.NoError(t, err)
+	err = server.RegisterSession(context.Background(), session2)
+	require.NoError(t, err)
+	err = server.RegisterSession(context.Background(), session3)
+	require.NoError(t, err)
+
+	// Send notification to session 1
+	err = server.SendNotificationToSpecificClient(session1.SessionID(), "test-method", map[string]any{
+		"data": "test-data",
+	})
+	require.NoError(t, err)
+
+	// Check that only session 1 received the notification
+	select {
+	case notification := <-session1Chan:
+		assert.Equal(t, "test-method", notification.Method)
+		assert.Equal(t, "test-data", notification.Params.AdditionalFields["data"])
+	case <-time.After(100 * time.Millisecond):
+		t.Error("Expected notification not received by session 1")
+	}
+
+	// Verify session 2 did not receive notification
+	select {
+	case notification := <-session2Chan:
+		t.Errorf("Unexpected notification received by session 2: %v", notification)
+	case <-time.After(100 * time.Millisecond):
+		// Expected, no notification for session 2
+	}
+
+	// Test sending to non-existent session
+	err = server.SendNotificationToSpecificClient("non-existent", "test-method", nil)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "not found")
+
+	// Test sending to uninitialized session
+	err = server.SendNotificationToSpecificClient(session3.SessionID(), "test-method", nil)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "not properly initialized")
+}
+
+func TestMCPServer_NotificationChannelBlocked(t *testing.T) {
+	// Set up a hooks object to capture error notifications
+	var mu sync.Mutex
+	errorCaptured := false
+	errorSessionID := ""
+	errorMethod := ""
+
+	hooks := &Hooks{}
+	hooks.AddOnError(func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) {
+		mu.Lock()
+		defer mu.Unlock()
+
+		errorCaptured = true
+		// Extract session ID and method from the error message metadata
+		if msgMap, ok := message.(map[string]interface{}); ok {
+			if sid, ok := msgMap["sessionID"].(string); ok {
+				errorSessionID = sid
+			}
+			if m, ok := msgMap["method"].(string); ok {
+				errorMethod = m
+			}
+		}
+		// Verify the error is a notification channel blocked error
+		assert.True(t, errors.Is(err, ErrNotificationChannelBlocked))
+	})
+
+	// Create a server with hooks
+	server := NewMCPServer("test-server", "1.0.0", WithHooks(hooks))
+
+	// Create a session with a very small buffer that will get blocked
+	smallBufferChan := make(chan mcp.JSONRPCNotification, 1)
+	session := &sessionTestClient{
+		sessionID:           "blocked-session",
+		notificationChannel: smallBufferChan,
+	}
+	session.Initialize()
+
+	// Register the session
+	err := server.RegisterSession(context.Background(), session)
+	require.NoError(t, err)
+
+	// Fill the buffer first to ensure it gets blocked
+	server.SendNotificationToSpecificClient(session.SessionID(), "first-message", nil)
+
+	// This will cause the buffer to block
+	err = server.SendNotificationToSpecificClient(session.SessionID(), "blocked-message", nil)
+	assert.Error(t, err)
+	assert.Equal(t, ErrNotificationChannelBlocked, err)
+
+	// Wait a bit for the goroutine to execute
+	time.Sleep(10 * time.Millisecond)
+
+	// Verify the error was logged via hooks
+	mu.Lock()
+	localErrorCaptured := errorCaptured
+	localErrorSessionID := errorSessionID
+	localErrorMethod := errorMethod
+	mu.Unlock()
+
+	assert.True(t, localErrorCaptured, "Error hook should have been called")
+	assert.Equal(t, "blocked-session", localErrorSessionID, "Session ID should be captured in the error hook")
+	assert.Equal(t, "blocked-message", localErrorMethod, "Method should be captured in the error hook")
+
+	// Also test SendNotificationToAllClients with a blocked channel
+	// Reset the captured data
+	mu.Lock()
+	errorCaptured = false
+	errorSessionID = ""
+	errorMethod = ""
+	mu.Unlock()
+
+	// Send to all clients (which includes our blocked one)
+	server.SendNotificationToAllClients("broadcast-message", nil)
+
+	// Wait a bit for the goroutine to execute
+	time.Sleep(10 * time.Millisecond)
+
+	// Verify the error was logged via hooks
+	mu.Lock()
+	localErrorCaptured = errorCaptured
+	localErrorSessionID = errorSessionID
+	localErrorMethod = errorMethod
+	mu.Unlock()
+
+	assert.True(t, localErrorCaptured, "Error hook should have been called for broadcast")
+	assert.Equal(t, "blocked-session", localErrorSessionID, "Session ID should be captured in the error hook")
+	assert.Equal(t, "broadcast-message", localErrorMethod, "Method should be captured in the error hook")
+}
diff --git a/server/sse_test.go b/server/sse_test.go
index 93474c6..a1ce01f 100644
--- a/server/sse_test.go
+++ b/server/sse_test.go
@@ -61,13 +61,10 @@ func TestSSEServer(t *testing.T) {
 		defer sseResp.Body.Close()
 
 		// Read the endpoint event
-		buf := make([]byte, 1024)
-		n, err := sseResp.Body.Read(buf)
+		endpointEvent, err := readSeeEvent(sseResp)
 		if err != nil {
 			t.Fatalf("Failed to read SSE response: %v", err)
 		}
-
-		endpointEvent := string(buf[:n])
 		if !strings.Contains(endpointEvent, "event: endpoint") {
 			t.Fatalf("Expected endpoint event, got: %s", endpointEvent)
 		}
@@ -109,19 +106,6 @@ func TestSSEServer(t *testing.T) {
 		if resp.StatusCode != http.StatusAccepted {
 			t.Errorf("Expected status 202, got %d", resp.StatusCode)
 		}
-
-		// Verify response
-		var response map[string]interface{}
-		if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
-			t.Fatalf("Failed to decode response: %v", err)
-		}
-
-		if response["jsonrpc"] != "2.0" {
-			t.Errorf("Expected jsonrpc 2.0, got %v", response["jsonrpc"])
-		}
-		if response["id"].(float64) != 1 {
-			t.Errorf("Expected id 1, got %v", response["id"])
-		}
 	})
 
 	t.Run("Can handle multiple sessions", func(t *testing.T) {
@@ -210,8 +194,17 @@ func TestSSEServer(t *testing.T) {
 				}
 				defer resp.Body.Close()
 
+				endpointEvent, err = readSeeEvent(sseResp)
+				if err != nil {
+					t.Fatalf("Failed to read SSE response: %v", err)
+				}
+				respFromSee := strings.TrimSpace(
+					strings.Split(strings.Split(endpointEvent, "data: ")[1], "\n")[0],
+				)
+
+				fmt.Printf("========> %v", respFromSee)
 				var response map[string]interface{}
-				if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+				if err := json.NewDecoder(strings.NewReader(respFromSee)).Decode(&response); err != nil {
 					t.Errorf(
 						"Session %d: Failed to decode response: %v",
 						sessionNum,
@@ -588,13 +581,10 @@ func TestSSEServer(t *testing.T) {
 		defer sseResp.Body.Close()
 
 		// Read the endpoint event
-		buf := make([]byte, 1024)
-		n, err := sseResp.Body.Read(buf)
+		endpointEvent, err := readSeeEvent(sseResp)
 		if err != nil {
 			t.Fatalf("Failed to read SSE response: %v", err)
 		}
-
-		endpointEvent := string(buf[:n])
 		messageURL := strings.TrimSpace(
 			strings.Split(strings.Split(endpointEvent, "data: ")[1], "\n")[0],
 		)
@@ -634,8 +624,16 @@ func TestSSEServer(t *testing.T) {
 		}
 
 		// Verify response
+		endpointEvent, err = readSeeEvent(sseResp)
+		if err != nil {
+			t.Fatalf("Failed to read SSE response: %v", err)
+		}
+		respFromSee := strings.TrimSpace(
+			strings.Split(strings.Split(endpointEvent, "data: ")[1], "\n")[0],
+		)
+
 		var response map[string]interface{}
-		if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+		if err := json.NewDecoder(strings.NewReader(respFromSee)).Decode(&response); err != nil {
 			t.Fatalf("Failed to decode response: %v", err)
 		}
 
@@ -673,8 +671,17 @@ func TestSSEServer(t *testing.T) {
 		}
 		defer resp.Body.Close()
 
+		endpointEvent, err = readSeeEvent(sseResp)
+		if err != nil {
+			t.Fatalf("Failed to read SSE response: %v", err)
+		}
+
+		respFromSee = strings.TrimSpace(
+			strings.Split(strings.Split(endpointEvent, "data: ")[1], "\n")[0],
+		)
+
 		response = make(map[string]interface{})
-		if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+		if err := json.NewDecoder(strings.NewReader(respFromSee)).Decode(&response); err != nil {
 			t.Fatalf("Failed to decode response: %v", err)
 		}
 
@@ -855,3 +862,12 @@ func TestSSEServer(t *testing.T) {
 		}
 	})
 }
+
+func readSeeEvent(sseResp *http.Response) (string, error) {
+	buf := make([]byte, 1024)
+	n, err := sseResp.Body.Read(buf)
+	if err != nil {
+		return "", err
+	}
+	return string(buf[:n]), nil
+}