Skip to content

Commit b7ab1cf

Browse files
Manage tools on a per session basis (#179)
* Add session specific methods for tools * Add session managment docs to README * Formatting * Remove TODO and update tests * Apply suggestions from code review Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> * Update tests * Update tests * Fix inaffassign * gitignore * Impl AddSessionTool * Fix potential race condition * Fixes * Fixes --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
1 parent 5a5781d commit b7ab1cf

File tree

6 files changed

+1208
-127
lines changed

6 files changed

+1208
-127
lines changed

.gitignore

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
11
.aider*
22
.env
3-
.idea
3+
.idea
4+
.opencode
5+
.claude

README.md

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,10 @@ MCP Go handles all the complex protocol details and server management, so you ca
9191
- [Tools](#tools)
9292
- [Prompts](#prompts)
9393
- [Examples](#examples)
94+
- [Extras](#extras)
95+
- [Session Management](#session-management)
96+
- [Request Hooks](#request-hooks)
97+
- [Tool Handler Middleware](#tool-handler-middleware)
9498
- [Contributing](#contributing)
9599
- [Prerequisites](#prerequisites)
96100
- [Installation](#installation-1)
@@ -516,6 +520,214 @@ For examples, see the `examples/` directory.
516520

517521
## Extras
518522

523+
### Session Management
524+
525+
MCP-Go provides a robust session management system that allows you to:
526+
- Maintain separate state for each connected client
527+
- Register and track client sessions
528+
- Send notifications to specific clients
529+
- Provide per-session tool customization
530+
531+
<details>
532+
<summary>Show Session Management Examples</summary>
533+
534+
#### Basic Session Handling
535+
536+
```go
537+
// Create a server with session capabilities
538+
s := server.NewMCPServer(
539+
"Session Demo",
540+
"1.0.0",
541+
server.WithToolCapabilities(true),
542+
)
543+
544+
// Implement your own ClientSession
545+
type MySession struct {
546+
id string
547+
notifChannel chan mcp.JSONRPCNotification
548+
isInitialized bool
549+
// Add custom fields for your application
550+
}
551+
552+
// Implement the ClientSession interface
553+
func (s *MySession) SessionID() string {
554+
return s.id
555+
}
556+
557+
func (s *MySession) NotificationChannel() chan<- mcp.JSONRPCNotification {
558+
return s.notifChannel
559+
}
560+
561+
func (s *MySession) Initialize() {
562+
s.isInitialized = true
563+
}
564+
565+
func (s *MySession) Initialized() bool {
566+
return s.isInitialized
567+
}
568+
569+
// Register a session
570+
session := &MySession{
571+
id: "user-123",
572+
notifChannel: make(chan mcp.JSONRPCNotification, 10),
573+
}
574+
if err := s.RegisterSession(context.Background(), session); err != nil {
575+
log.Printf("Failed to register session: %v", err)
576+
}
577+
578+
// Send notification to a specific client
579+
err := s.SendNotificationToSpecificClient(
580+
session.SessionID(),
581+
"notification/update",
582+
map[string]any{"message": "New data available!"},
583+
)
584+
if err != nil {
585+
log.Printf("Failed to send notification: %v", err)
586+
}
587+
588+
// Unregister session when done
589+
s.UnregisterSession(context.Background(), session.SessionID())
590+
```
591+
592+
#### Per-Session Tools
593+
594+
For more advanced use cases, you can implement the `SessionWithTools` interface to support per-session tool customization:
595+
596+
```go
597+
// Implement SessionWithTools interface for per-session tools
598+
type MyAdvancedSession struct {
599+
MySession // Embed the basic session
600+
sessionTools map[string]server.ServerTool
601+
}
602+
603+
// Implement additional methods for SessionWithTools
604+
func (s *MyAdvancedSession) GetSessionTools() map[string]server.ServerTool {
605+
return s.sessionTools
606+
}
607+
608+
func (s *MyAdvancedSession) SetSessionTools(tools map[string]server.ServerTool) {
609+
s.sessionTools = tools
610+
}
611+
612+
// Create and register a session with tools support
613+
advSession := &MyAdvancedSession{
614+
MySession: MySession{
615+
id: "user-456",
616+
notifChannel: make(chan mcp.JSONRPCNotification, 10),
617+
},
618+
sessionTools: make(map[string]server.ServerTool),
619+
}
620+
if err := s.RegisterSession(context.Background(), advSession); err != nil {
621+
log.Printf("Failed to register session: %v", err)
622+
}
623+
624+
// Add session-specific tools
625+
userSpecificTool := mcp.NewTool(
626+
"user_data",
627+
mcp.WithDescription("Access user-specific data"),
628+
)
629+
// You can use AddSessionTool (similar to AddTool)
630+
err := s.AddSessionTool(
631+
advSession.SessionID(),
632+
userSpecificTool,
633+
func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
634+
// This handler is only available to this specific session
635+
return mcp.NewToolResultText("User-specific data for " + advSession.SessionID()), nil
636+
},
637+
)
638+
if err != nil {
639+
log.Printf("Failed to add session tool: %v", err)
640+
}
641+
642+
// Or use AddSessionTools directly with ServerTool
643+
/*
644+
err := s.AddSessionTools(
645+
advSession.SessionID(),
646+
server.ServerTool{
647+
Tool: userSpecificTool,
648+
Handler: func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
649+
// This handler is only available to this specific session
650+
return mcp.NewToolResultText("User-specific data for " + advSession.SessionID()), nil
651+
},
652+
},
653+
)
654+
if err != nil {
655+
log.Printf("Failed to add session tool: %v", err)
656+
}
657+
*/
658+
659+
// Delete session-specific tools when no longer needed
660+
err = s.DeleteSessionTools(advSession.SessionID(), "user_data")
661+
if err != nil {
662+
log.Printf("Failed to delete session tool: %v", err)
663+
}
664+
```
665+
666+
#### Tool Filtering
667+
668+
You can also apply filters to control which tools are available to certain sessions:
669+
670+
```go
671+
// Add a tool filter that only shows tools with certain prefixes
672+
s := server.NewMCPServer(
673+
"Tool Filtering Demo",
674+
"1.0.0",
675+
server.WithToolCapabilities(true),
676+
server.WithToolFilter(func(ctx context.Context, tools []mcp.Tool) []mcp.Tool {
677+
// Get session from context
678+
session := server.ClientSessionFromContext(ctx)
679+
if session == nil {
680+
return tools // Return all tools if no session
681+
}
682+
683+
// Example: filter tools based on session ID prefix
684+
if strings.HasPrefix(session.SessionID(), "admin-") {
685+
// Admin users get all tools
686+
return tools
687+
} else {
688+
// Regular users only get tools with "public-" prefix
689+
var filteredTools []mcp.Tool
690+
for _, tool := range tools {
691+
if strings.HasPrefix(tool.Name, "public-") {
692+
filteredTools = append(filteredTools, tool)
693+
}
694+
}
695+
return filteredTools
696+
}
697+
}),
698+
)
699+
```
700+
701+
#### Working with Context
702+
703+
The session context is automatically passed to tool and resource handlers:
704+
705+
```go
706+
s.AddTool(mcp.NewTool("session_aware"), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
707+
// Get the current session from context
708+
session := server.ClientSessionFromContext(ctx)
709+
if session == nil {
710+
return mcp.NewToolResultError("No active session"), nil
711+
}
712+
713+
return mcp.NewToolResultText("Hello, session " + session.SessionID()), nil
714+
})
715+
716+
// When using handlers in HTTP/SSE servers, you need to pass the context with the session
717+
httpHandler := func(w http.ResponseWriter, r *http.Request) {
718+
// Get session from somewhere (like a cookie or header)
719+
session := getSessionFromRequest(r)
720+
721+
// Add session to context
722+
ctx := s.WithContext(r.Context(), session)
723+
724+
// Use this context when handling requests
725+
// ...
726+
}
727+
```
728+
729+
</details>
730+
519731
### Request Hooks
520732

521733
Hook into the request lifecycle by creating a `Hooks` object with your

server/errors.go

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package server
2+
3+
import (
4+
"errors"
5+
)
6+
7+
var (
8+
// Common server errors
9+
ErrUnsupported = errors.New("not supported")
10+
ErrResourceNotFound = errors.New("resource not found")
11+
ErrPromptNotFound = errors.New("prompt not found")
12+
ErrToolNotFound = errors.New("tool not found")
13+
14+
// Session-related errors
15+
ErrSessionNotFound = errors.New("session not found")
16+
ErrSessionExists = errors.New("session already exists")
17+
ErrSessionNotInitialized = errors.New("session not properly initialized")
18+
ErrSessionDoesNotSupportTools = errors.New("session does not support per-session tools")
19+
20+
// Notification-related errors
21+
ErrNotificationNotInitialized = errors.New("notification channel not initialized")
22+
ErrNotificationChannelBlocked = errors.New("notification channel full or blocked")
23+
)

0 commit comments

Comments
 (0)