Files
openshift-mcp-server/vendor/github.com/mark3labs/mcp-go/client/transport/inprocess.go
2025-09-25 16:40:49 +03:00

121 lines
3.1 KiB
Go

package transport
import (
"context"
"encoding/json"
"fmt"
"sync"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
type InProcessTransport struct {
server *server.MCPServer
samplingHandler server.SamplingHandler
elicitationHandler server.ElicitationHandler
session *server.InProcessSession
sessionID string
onNotification func(mcp.JSONRPCNotification)
notifyMu sync.RWMutex
}
type InProcessOption func(*InProcessTransport)
func WithSamplingHandler(handler server.SamplingHandler) InProcessOption {
return func(t *InProcessTransport) {
t.samplingHandler = handler
}
}
func WithElicitationHandler(handler server.ElicitationHandler) InProcessOption {
return func(t *InProcessTransport) {
t.elicitationHandler = handler
}
}
func NewInProcessTransport(server *server.MCPServer) *InProcessTransport {
return &InProcessTransport{
server: server,
}
}
func NewInProcessTransportWithOptions(server *server.MCPServer, opts ...InProcessOption) *InProcessTransport {
t := &InProcessTransport{
server: server,
sessionID: server.GenerateInProcessSessionID(),
}
for _, opt := range opts {
opt(t)
}
return t
}
func (c *InProcessTransport) Start(ctx context.Context) error {
// Create and register session if we have handlers
if c.samplingHandler != nil || c.elicitationHandler != nil {
c.session = server.NewInProcessSessionWithHandlers(c.sessionID, c.samplingHandler, c.elicitationHandler)
if err := c.server.RegisterSession(ctx, c.session); err != nil {
return fmt.Errorf("failed to register session: %w", err)
}
}
return nil
}
func (c *InProcessTransport) SendRequest(ctx context.Context, request JSONRPCRequest) (*JSONRPCResponse, error) {
requestBytes, err := json.Marshal(request)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
}
requestBytes = append(requestBytes, '\n')
// Add session to context if available
if c.session != nil {
ctx = c.server.WithContext(ctx, c.session)
}
respMessage := c.server.HandleMessage(ctx, requestBytes)
respByte, err := json.Marshal(respMessage)
if err != nil {
return nil, fmt.Errorf("failed to marshal response message: %w", err)
}
rpcResp := JSONRPCResponse{}
err = json.Unmarshal(respByte, &rpcResp)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal response message: %w", err)
}
return &rpcResp, nil
}
func (c *InProcessTransport) SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error {
notificationBytes, err := json.Marshal(notification)
if err != nil {
return fmt.Errorf("failed to marshal notification: %w", err)
}
notificationBytes = append(notificationBytes, '\n')
c.server.HandleMessage(ctx, notificationBytes)
return nil
}
func (c *InProcessTransport) SetNotificationHandler(handler func(notification mcp.JSONRPCNotification)) {
c.notifyMu.Lock()
defer c.notifyMu.Unlock()
c.onNotification = handler
}
func (c *InProcessTransport) Close() error {
if c.session != nil {
c.server.UnregisterSession(context.Background(), c.sessionID)
}
return nil
}
func (c *InProcessTransport) GetSessionId() string {
return ""
}