internal/mcp: implement logging

Provide a ServerSession method that directly implements the spec
method.

Also provide a custom slog.Handler that writes a notification on each
log call. It uses a slog.JSONHandler to format the message as JSON, then
adds the additional fields of LoggingMessageParams.

I decided that the direct method should handle levels, since that is
in the spec, but that rate-limiting is an extra feature, so is only
provided in the slog handler.

Change-Id: I1190ca124e992e61fb633ac52ee3492c8e6d99c9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/674498
Auto-Submit: Jonathan Amsterdam <jba@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>

mcp/server.go

@@ -299,11 +299,11 @@ func (s *Server) Run(ctx context.Context, t Transport) error {
 // bind implements the binder[*ServerSession] interface, so that Servers can
 // be connected using [connect].
 func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession {
-	cc := &ServerSession{conn: conn, server: s}
+	ss := &ServerSession{conn: conn, server: s}
 	s.mu.Lock()
-	s.sessions = append(s.sessions, cc)
+	s.sessions = append(s.sessions, ss)
 	s.mu.Unlock()
-	return cc
+	return ss
 }
 
 // disconnect implements the binder[*ServerSession] interface, so that
@@ -341,17 +341,17 @@ func (s *Server) callRootsListChangedHandler(ctx context.Context, ss *ServerSess
 // Call [ServerSession.Close] to close the connection, or await client
 // termination with [ServerSession.Wait].
 type ServerSession struct {
-	server *Server
-	conn   *jsonrpc2.Connection
-
+	server           *Server
+	conn             *jsonrpc2.Connection
 	mu               sync.Mutex
+	logLevel         LoggingLevel
 	initializeParams *InitializeParams
 	initialized      bool
 }
 
 // Ping pings the client.
 func (ss *ServerSession) Ping(ctx context.Context, _ *PingParams) error {
-	return call(ctx, ss.conn, "ping", (*PingParams)(nil), nil)
+	return call(ctx, ss.conn, methodPing, (*PingParams)(nil), nil)
 }
 
 // ListRoots lists the client roots.
@@ -364,6 +364,25 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag
 	return standardCall[CreateMessageResult](ctx, ss.conn, methodCreateMessage, params)
 }
 
+// LoggingMessage sends a logging message to the client.
+// The message is not sent if the client has not called SetLevel, or if its level
+// is below that of the last SetLevel.
+func (ss *ServerSession) LoggingMessage(ctx context.Context, params *LoggingMessageParams) error {
+	ss.mu.Lock()
+	logLevel := ss.logLevel
+	ss.mu.Unlock()
+	if logLevel == "" {
+		// The spec is unclear, but seems to imply that no log messages are sent until the client
+		// sets the level.
+		// TODO(jba): read other SDKs, possibly file an issue.
+		return nil
+	}
+	if compareLevels(params.Level, logLevel) < 0 {
+		return nil
+	}
+	return ss.conn.Notify(ctx, notificationLoggingMessage, params)
+}
+
 // AddMiddleware wraps the server's current method handler using the provided
 // middleware. Middleware is applied from right to left, so that the first one
 // is executed first.
@@ -386,6 +405,7 @@ var serverMethodInfos = map[string]methodInfo[ServerSession]{
 	methodCallTool:               newMethodInfo(serverMethod((*Server).callTool)),
 	methodListResources:          newMethodInfo(serverMethod((*Server).listResources)),
 	methodReadResource:           newMethodInfo(serverMethod((*Server).readResource)),
+	methodSetLevel:               newMethodInfo(sessionMethod((*ServerSession).setLevel)),
 	notificationInitialized:      newMethodInfo(serverMethod((*Server).callInitializedHandler)),
 	notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler)),
 }
@@ -457,6 +477,7 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
 			Resources: &resourceCapabilities{
 				ListChanged: true,
 			},
+			Logging: &loggingCapabilities{},
 		},
 		Instructions: ss.server.opts.Instructions,
 		ServerInfo: &implementation{
@@ -470,6 +491,13 @@ func (ss *ServerSession) ping(context.Context, *PingParams) (struct{}, error) {
 	return struct{}{}, nil
 }
 
+func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (struct{}, error) {
+	ss.mu.Lock()
+	defer ss.mu.Unlock()
+	ss.logLevel = params.Level
+	return struct{}{}, nil
+}
+
 // Close performs a graceful shutdown of the connection, preventing new
 // requests from being handled, and waiting for ongoing requests to return.
 // Close then terminates the connection.