Skip to content

Commit 077f546

Browse files
authored
feat(MCPServer): support logging/setlevel request (#276)
* feat(MCPServer): support logging/setlevel request * update template file and adopt coderabbitai suggestion
1 parent 09c23b5 commit 077f546

File tree

10 files changed

+299
-16
lines changed

10 files changed

+299
-16
lines changed

mcp/types.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,10 @@ const (
5050
// https://modelcontextprotocol.io/specification/2024-11-05/server/tools/
5151
MethodToolsCall MCPMethod = "tools/call"
5252

53+
// MethodSetLogLevel configures the minimum log level for client
54+
// https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/logging
55+
MethodSetLogLevel MCPMethod = "logging/setLevel"
56+
5357
// MethodNotificationResourcesListChanged notifies when the list of available resources changes.
5458
// https://modelcontextprotocol.io/specification/2025-03-26/server/resources#list-changed-notification
5559
MethodNotificationResourcesListChanged = "notifications/resources/list_changed"

server/errors.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,11 @@ var (
1313
ErrToolNotFound = errors.New("tool not found")
1414

1515
// Session-related errors
16-
ErrSessionNotFound = errors.New("session not found")
17-
ErrSessionExists = errors.New("session already exists")
18-
ErrSessionNotInitialized = errors.New("session not properly initialized")
19-
ErrSessionDoesNotSupportTools = errors.New("session does not support per-session tools")
16+
ErrSessionNotFound = errors.New("session not found")
17+
ErrSessionExists = errors.New("session already exists")
18+
ErrSessionNotInitialized = errors.New("session not properly initialized")
19+
ErrSessionDoesNotSupportTools = errors.New("session does not support per-session tools")
20+
ErrSessionDoesNotSupportLogging = errors.New("session does not support setting logging level")
2021

2122
// Notification-related errors
2223
ErrNotificationNotInitialized = errors.New("notification channel not initialized")

server/hooks.go

Lines changed: 32 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

server/internal/gen/data.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,16 @@ var MCPRequestTypes = []MCPRequestType{
2727
HookName: "Ping",
2828
UnmarshalError: "invalid ping request",
2929
HandlerFunc: "handlePing",
30+
}, {
31+
MethodName: "MethodSetLogLevel",
32+
ParamType: "SetLevelRequest",
33+
ResultType: "EmptyResult",
34+
Group: "logging",
35+
GroupName: "Logging",
36+
GroupHookName: "Logging",
37+
HookName: "SetLevel",
38+
UnmarshalError: "invalid set level request",
39+
HandlerFunc: "handleSetLevel",
3040
}, {
3141
MethodName: "MethodResourcesList",
3242
ParamType: "ListResourcesRequest",

server/request_handler.go

Lines changed: 25 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

server/server.go

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ type serverCapabilities struct {
161161
tools *toolCapabilities
162162
resources *resourceCapabilities
163163
prompts *promptCapabilities
164-
logging bool
164+
logging *bool
165165
}
166166

167167
// resourceCapabilities defines the supported resource-related features
@@ -260,7 +260,7 @@ func WithToolCapabilities(listChanged bool) ServerOption {
260260
// WithLogging enables logging capabilities for the server
261261
func WithLogging() ServerOption {
262262
return func(s *MCPServer) {
263-
s.capabilities.logging = true
263+
s.capabilities.logging = mcp.ToBoolPtr(true)
264264
}
265265
}
266266

@@ -289,7 +289,7 @@ func NewMCPServer(
289289
tools: nil,
290290
resources: nil,
291291
prompts: nil,
292-
logging: false,
292+
logging: nil,
293293
},
294294
}
295295

@@ -521,7 +521,7 @@ func (s *MCPServer) handleInitialize(
521521
}
522522
}
523523

524-
if s.capabilities.logging {
524+
if s.capabilities.logging != nil && *s.capabilities.logging {
525525
capabilities.Logging = &struct{}{}
526526
}
527527

@@ -549,6 +549,49 @@ func (s *MCPServer) handlePing(
549549
return &mcp.EmptyResult{}, nil
550550
}
551551

552+
func (s *MCPServer) handleSetLevel(
553+
ctx context.Context,
554+
id any,
555+
request mcp.SetLevelRequest,
556+
) (*mcp.EmptyResult, *requestError) {
557+
clientSession := ClientSessionFromContext(ctx)
558+
if clientSession == nil || !clientSession.Initialized() {
559+
return nil, &requestError{
560+
id: id,
561+
code: mcp.INTERNAL_ERROR,
562+
err: ErrSessionNotInitialized,
563+
}
564+
}
565+
566+
sessionLogging, ok := clientSession.(SessionWithLogging)
567+
if !ok {
568+
return nil, &requestError{
569+
id: id,
570+
code: mcp.INTERNAL_ERROR,
571+
err: ErrSessionDoesNotSupportLogging,
572+
}
573+
}
574+
575+
level := request.Params.Level
576+
// Validate logging level
577+
switch level {
578+
case mcp.LoggingLevelDebug, mcp.LoggingLevelInfo, mcp.LoggingLevelNotice,
579+
mcp.LoggingLevelWarning, mcp.LoggingLevelError, mcp.LoggingLevelCritical,
580+
mcp.LoggingLevelAlert, mcp.LoggingLevelEmergency:
581+
// Valid level
582+
default:
583+
return nil, &requestError{
584+
id: id,
585+
code: mcp.INVALID_PARAMS,
586+
err: fmt.Errorf("invalid logging level '%s'", level),
587+
}
588+
}
589+
590+
sessionLogging.SetLogLevel(level)
591+
592+
return &mcp.EmptyResult{}, nil
593+
}
594+
552595
func listByPagination[T mcp.Named](
553596
ctx context.Context,
554597
s *MCPServer,

server/session.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,15 @@ type ClientSession interface {
1919
SessionID() string
2020
}
2121

22+
// SessionWithLogging is an extension of ClientSession that can receive log message notifications and set log level
23+
type SessionWithLogging interface {
24+
ClientSession
25+
// SetLogLevel sets the minimum log level
26+
SetLogLevel(level mcp.LoggingLevel)
27+
// GetLogLevel retrieves the minimum log level
28+
GetLogLevel() mcp.LoggingLevel
29+
}
30+
2231
// SessionWithTools is an extension of ClientSession that can store session-specific tool data
2332
type SessionWithTools interface {
2433
ClientSession

server/session_test.go

Lines changed: 128 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"encoding/json"
66
"errors"
77
"sync"
8+
"sync/atomic"
89
"testing"
910
"time"
1011

@@ -98,9 +99,47 @@ func (f *sessionTestClientWithTools) SetSessionTools(tools map[string]ServerTool
9899
f.sessionTools = toolsCopy
99100
}
100101

101-
// Verify that both implementations satisfy their respective interfaces
102-
var _ ClientSession = &sessionTestClient{}
103-
var _ SessionWithTools = &sessionTestClientWithTools{}
102+
// sessionTestClientWithTools implements the SessionWithLogging interface for testing
103+
type sessionTestClientWithLogging struct {
104+
sessionID string
105+
notificationChannel chan mcp.JSONRPCNotification
106+
initialized bool
107+
loggingLevel atomic.Value
108+
}
109+
110+
func (f *sessionTestClientWithLogging) SessionID() string {
111+
return f.sessionID
112+
}
113+
114+
func (f *sessionTestClientWithLogging) NotificationChannel() chan<- mcp.JSONRPCNotification {
115+
return f.notificationChannel
116+
}
117+
118+
func (f *sessionTestClientWithLogging) Initialize() {
119+
// set default logging level
120+
f.loggingLevel.Store(mcp.LoggingLevelError)
121+
f.initialized = true
122+
}
123+
124+
func (f *sessionTestClientWithLogging) Initialized() bool {
125+
return f.initialized
126+
}
127+
128+
func (f *sessionTestClientWithLogging) SetLogLevel(level mcp.LoggingLevel) {
129+
f.loggingLevel.Store(level)
130+
}
131+
132+
func (f *sessionTestClientWithLogging) GetLogLevel() mcp.LoggingLevel {
133+
level := f.loggingLevel.Load()
134+
return level.(mcp.LoggingLevel)
135+
}
136+
137+
// Verify that all implementations satisfy their respective interfaces
138+
var (
139+
_ ClientSession = (*sessionTestClient)(nil)
140+
_ SessionWithTools = (*sessionTestClientWithTools)(nil)
141+
_ SessionWithLogging = (*sessionTestClientWithLogging)(nil)
142+
)
104143

105144
func TestSessionWithTools_Integration(t *testing.T) {
106145
server := NewMCPServer("test-server", "1.0.0", WithToolCapabilities(true))
@@ -917,3 +956,89 @@ func TestMCPServer_ToolNotificationsDisabled(t *testing.T) {
917956
// Verify tool was deleted from session
918957
assert.Len(t, session.GetSessionTools(), 0)
919958
}
959+
960+
func TestMCPServer_SetLevelNotEnabled(t *testing.T) {
961+
// Create server without logging capability
962+
server := NewMCPServer("test-server", "1.0.0")
963+
964+
// Create and initialize a session
965+
sessionChan := make(chan mcp.JSONRPCNotification, 10)
966+
session := &sessionTestClientWithLogging{
967+
sessionID: "session-1",
968+
notificationChannel: sessionChan,
969+
}
970+
session.Initialize()
971+
972+
// Register the session
973+
err := server.RegisterSession(context.Background(), session)
974+
require.NoError(t, err)
975+
976+
// Try to set logging level when capability is disabled
977+
sessionCtx := server.WithContext(context.Background(), session)
978+
setRequest := map[string]any{
979+
"jsonrpc": "2.0",
980+
"id": 1,
981+
"method": "logging/setLevel",
982+
"params": map[string]any{
983+
"level": mcp.LoggingLevelCritical,
984+
},
985+
}
986+
requestBytes, err := json.Marshal(setRequest)
987+
require.NoError(t, err)
988+
989+
response := server.HandleMessage(sessionCtx, requestBytes)
990+
errorResponse, ok := response.(mcp.JSONRPCError)
991+
assert.True(t, ok)
992+
993+
// Verify we get a METHOD_NOT_FOUND error
994+
assert.NotNil(t, errorResponse.Error)
995+
assert.Equal(t, mcp.METHOD_NOT_FOUND, errorResponse.Error.Code)
996+
}
997+
998+
func TestMCPServer_SetLevel(t *testing.T) {
999+
server := NewMCPServer("test-server", "1.0.0", WithLogging())
1000+
1001+
// Create and initicalize a session
1002+
sessionChan := make(chan mcp.JSONRPCNotification, 10)
1003+
session := &sessionTestClientWithLogging{
1004+
sessionID: "session-1",
1005+
notificationChannel: sessionChan,
1006+
}
1007+
session.Initialize()
1008+
1009+
// Check default logging level
1010+
if session.GetLogLevel() != mcp.LoggingLevelError {
1011+
t.Errorf("Expected error level, got %v", session.GetLogLevel())
1012+
}
1013+
1014+
// Register the session
1015+
err := server.RegisterSession(context.Background(), session)
1016+
require.NoError(t, err)
1017+
1018+
// Set Logging level to critical
1019+
sessionCtx := server.WithContext(context.Background(), session)
1020+
setRequest := map[string]any{
1021+
"jsonrpc": "2.0",
1022+
"id": 1,
1023+
"method": "logging/setLevel",
1024+
"params": map[string]any{
1025+
"level": mcp.LoggingLevelCritical,
1026+
},
1027+
}
1028+
requestBytes, err := json.Marshal(setRequest)
1029+
if err != nil {
1030+
t.Fatalf("Failed to marshal tool request: %v", err)
1031+
}
1032+
1033+
response := server.HandleMessage(sessionCtx, requestBytes)
1034+
resp, ok := response.(mcp.JSONRPCResponse)
1035+
assert.True(t, ok)
1036+
1037+
_, ok = resp.Result.(mcp.EmptyResult)
1038+
assert.True(t, ok)
1039+
1040+
// Check logging level
1041+
if session.GetLogLevel() != mcp.LoggingLevelCritical {
1042+
t.Errorf("Expected critical level, got %v", session.GetLogLevel())
1043+
}
1044+
}

0 commit comments

Comments
 (0)