Skip to content

Manage tools on a per session basis #179

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Apr 29, 2025
4 changes: 3 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
.aider*
.env
.idea
.idea
.opencode
.claude
212 changes: 212 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️

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
Expand Down
23 changes: 23 additions & 0 deletions server/errors.go
Original file line number Diff line number Diff line change
@@ -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")
)
Loading