Skip to content

Commit 459e921

Browse files
ezynda3ezynda3
committed
Handle notifications
1 parent 8a19ae4 commit 459e921

File tree

5 files changed

+109
-48
lines changed

5 files changed

+109
-48
lines changed

‎examples/everything/main.go

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -311,23 +311,21 @@ func (s *MCPServer) handleSendNotification(
311311
ctx context.Context,
312312
arguments map[string]interface{},
313313
) (*mcp.CallToolResult, error) {
314-
log.Printf("Sending notification to client...")
315314

316315
server := server.ServerFromContext(ctx)
317316

318317
err := server.SendNotificationToClient(
319-
"notifications/initialized",
318+
"notifications/progress",
320319
map[string]interface{}{
321-
"message": "This is a test notification",
322-
"timestamp": time.Now().String(),
320+
"progress": 10,
321+
"total": 10,
322+
"progressToken": 0,
323323
},
324324
)
325325
if err != nil {
326326
return nil, fmt.Errorf("failed to send notification: %w", err)
327327
}
328328

329-
log.Printf("Notification sent")
330-
331329
return &mcp.CallToolResult{
332330
Content: []interface{}{
333331
mcp.TextContent{

‎mcp/types.go

Lines changed: 68 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// MCP is a protocol for communication between LLM-powered applications and their supporting services.
33
package mcp
44

5+
import "encoding/json"
6+
57
/* JSON-RPC types */
68

79
// JSONRPCMessage represents either a JSONRPCRequest, JSONRPCNotification, JSONRPCResponse, or JSONRPCError
@@ -34,13 +36,73 @@ type Request struct {
3436
} `json:"params,omitempty"`
3537
}
3638

39+
type Params map[string]interface{}
40+
3741
type Notification struct {
38-
Method string `json:"method"`
39-
Params struct {
40-
// This parameter name is reserved by MCP to allow clients and
41-
// servers to attach additional metadata to their notifications.
42-
Meta map[string]interface{} `json:"_meta,omitempty"`
43-
} `json:"params,omitempty"`
42+
Method string `json:"method"`
43+
Params NotificationParams `json:"params,omitempty"`
44+
}
45+
46+
type NotificationParams struct {
47+
// This parameter name is reserved by MCP to allow clients and
48+
// servers to attach additional metadata to their notifications.
49+
Meta map[string]interface{} `json:"_meta,omitempty"`
50+
51+
// Additional fields can be added to this map
52+
AdditionalFields map[string]interface{} `json:"-"`
53+
}
54+
55+
// MarshalJSON implements custom JSON marshaling
56+
func (p NotificationParams) MarshalJSON() ([]byte, error) {
57+
// Create a map to hold all fields
58+
m := make(map[string]interface{})
59+
60+
// Add Meta if it exists
61+
if p.Meta != nil {
62+
m["_meta"] = p.Meta
63+
}
64+
65+
// Add all additional fields
66+
for k, v := range p.AdditionalFields {
67+
// Ensure we don't override the _meta field
68+
if k != "_meta" {
69+
m[k] = v
70+
}
71+
}
72+
73+
return json.Marshal(m)
74+
}
75+
76+
// UnmarshalJSON implements custom JSON unmarshaling
77+
func (p *NotificationParams) UnmarshalJSON(data []byte) error {
78+
// Create a map to hold all fields
79+
var m map[string]interface{}
80+
if err := json.Unmarshal(data, &m); err != nil {
81+
return err
82+
}
83+
84+
// Initialize maps if they're nil
85+
if p.Meta == nil {
86+
p.Meta = make(map[string]interface{})
87+
}
88+
if p.AdditionalFields == nil {
89+
p.AdditionalFields = make(map[string]interface{})
90+
}
91+
92+
// Process all fields
93+
for k, v := range m {
94+
if k == "_meta" {
95+
// Handle Meta field
96+
if meta, ok := v.(map[string]interface{}); ok {
97+
p.Meta = meta
98+
}
99+
} else {
100+
// Handle additional fields
101+
p.AdditionalFields[k] = v
102+
}
103+
}
104+
105+
return nil
44106
}
45107

46108
type Result struct {

‎server/server.go

Lines changed: 33 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -61,12 +61,12 @@ type MCPServer struct {
6161
resourceTemplates map[string]resourceTemplateEntry
6262
prompts map[string]mcp.Prompt
6363
promptHandlers map[string]PromptHandlerFunc
64-
tools map[string]mcp.Tool
65-
toolHandlers map[string]ToolHandlerFunc
64+
tools map[string]mcp.Tool
65+
toolHandlers map[string]ToolHandlerFunc
6666
notificationHandlers map[string]NotificationHandlerFunc
67-
capabilities serverCapabilities
68-
notifications chan ServerNotification
69-
currentClient NotificationContext
67+
capabilities serverCapabilities
68+
notifications chan ServerNotification
69+
currentClient NotificationContext
7070
}
7171

7272
// serverKey is the context key for storing the server instance
@@ -81,13 +81,19 @@ func ServerFromContext(ctx context.Context) *MCPServer {
8181
}
8282

8383
// WithContext sets the current client context and returns the provided context
84-
func (s *MCPServer) WithContext(ctx context.Context, notifCtx NotificationContext) context.Context {
84+
func (s *MCPServer) WithContext(
85+
ctx context.Context,
86+
notifCtx NotificationContext,
87+
) context.Context {
8588
s.currentClient = notifCtx
8689
return ctx
8790
}
8891

8992
// SendNotificationToClient sends a notification to the current client
90-
func (s *MCPServer) SendNotificationToClient(method string, params interface{}) error {
93+
func (s *MCPServer) SendNotificationToClient(
94+
method string,
95+
params map[string]interface{},
96+
) error {
9197
if s.notifications == nil {
9298
return fmt.Errorf("notification channel not initialized")
9399
}
@@ -96,16 +102,11 @@ func (s *MCPServer) SendNotificationToClient(method string, params interface{})
96102
JSONRPC: mcp.JSONRPC_VERSION,
97103
Notification: mcp.Notification{
98104
Method: method,
105+
Params: mcp.NotificationParams{
106+
AdditionalFields: params,
107+
},
99108
},
100109
}
101-
102-
if params != nil {
103-
var notificationParams struct {
104-
Meta map[string]interface{} `json:"_meta,omitempty"`
105-
}
106-
notificationParams.Meta = map[string]interface{}{"params": params}
107-
notification.Notification.Params = notificationParams
108-
}
109110

110111
select {
111112
case s.notifications <- ServerNotification{
@@ -168,16 +169,16 @@ func NewMCPServer(
168169
opts ...ServerOption,
169170
) *MCPServer {
170171
s := &MCPServer{
171-
resources: make(map[string]resourceEntry),
172-
resourceTemplates: make(map[string]resourceTemplateEntry),
173-
prompts: make(map[string]mcp.Prompt),
174-
promptHandlers: make(map[string]PromptHandlerFunc),
175-
tools: make(map[string]mcp.Tool),
176-
toolHandlers: make(map[string]ToolHandlerFunc),
177-
name: name,
178-
version: version,
172+
resources: make(map[string]resourceEntry),
173+
resourceTemplates: make(map[string]resourceTemplateEntry),
174+
prompts: make(map[string]mcp.Prompt),
175+
promptHandlers: make(map[string]PromptHandlerFunc),
176+
tools: make(map[string]mcp.Tool),
177+
toolHandlers: make(map[string]ToolHandlerFunc),
178+
name: name,
179+
version: version,
179180
notificationHandlers: make(map[string]NotificationHandlerFunc),
180-
notifications: make(chan ServerNotification, 100),
181+
notifications: make(chan ServerNotification, 100),
181182
}
182183

183184
for _, opt := range opts {
@@ -424,7 +425,10 @@ func (s *MCPServer) AddTool(tool mcp.Tool, handler ToolHandlerFunc) {
424425
}
425426

426427
// AddNotificationHandler registers a new handler for incoming notifications
427-
func (s *MCPServer) AddNotificationHandler(method string, handler NotificationHandlerFunc) {
428+
func (s *MCPServer) AddNotificationHandler(
429+
method string,
430+
handler NotificationHandlerFunc,
431+
) {
428432
s.notificationHandlers[method] = handler
429433
}
430434

@@ -658,7 +662,10 @@ func (s *MCPServer) handleToolCall(
658662
return createResponse(id, result)
659663
}
660664

661-
func (s *MCPServer) handleNotification(ctx context.Context, notification mcp.JSONRPCNotification) mcp.JSONRPCMessage {
665+
func (s *MCPServer) handleNotification(
666+
ctx context.Context,
667+
notification mcp.JSONRPCNotification,
668+
) mcp.JSONRPCMessage {
662669
if handler, ok := s.notificationHandlers[notification.Method]; ok {
663670
handler(ctx, notification)
664671
}

‎server/sse.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
145145
s.baseURL,
146146
sessionID,
147147
)
148-
fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", messageEndpoint)
148+
fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", messageEndpoint)
149149
flusher.Flush()
150150

151151
<-r.Context().Done()

‎server/stdio.go

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ func NewStdioServer(server *MCPServer) *StdioServer {
2828
return &StdioServer{
2929
server: server,
3030
errLogger: log.New(
31-
io.Discard,
31+
os.Stderr,
3232
"",
3333
log.LstdFlags,
3434
), // Default to discarding logs
@@ -64,17 +64,11 @@ func (s *StdioServer) Listen(
6464
case serverNotification := <-s.server.notifications:
6565
// Only handle notifications for stdio client
6666
if serverNotification.Context.ClientID == "stdio" {
67-
notificationBytes, err := json.Marshal(
67+
err := s.writeResponse(
6868
serverNotification.Notification,
69+
stdout,
6970
)
7071
if err != nil {
71-
s.errLogger.Printf(
72-
"Error marshaling notification: %v",
73-
err,
74-
)
75-
continue
76-
}
77-
if _, err := fmt.Fprintf(stdout, "%s\n", notificationBytes); err != nil {
7872
s.errLogger.Printf(
7973
"Error writing notification: %v",
8074
err,

0 commit comments

Comments
 (0)