File size: 3,362 Bytes
0f07ba7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | package mcp
import (
"context"
"net/http"
"os"
"os/exec"
"sync"
"time"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/pkg/signals"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/mudler/xlog"
)
type sessionCache struct {
mu sync.Mutex
cache map[string][]*mcp.ClientSession
}
var (
cache = sessionCache{
cache: make(map[string][]*mcp.ClientSession),
}
client = mcp.NewClient(&mcp.Implementation{Name: "LocalAI", Version: "v1.0.0"}, nil)
)
func SessionsFromMCPConfig(
name string,
remote config.MCPGenericConfig[config.MCPRemoteServers],
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
) ([]*mcp.ClientSession, error) {
cache.mu.Lock()
defer cache.mu.Unlock()
sessions, exists := cache.cache[name]
if exists {
return sessions, nil
}
allSessions := []*mcp.ClientSession{}
ctx, cancel := context.WithCancel(context.Background())
// Get the list of all the tools that the Agent will be esposed to
for _, server := range remote.Servers {
xlog.Debug("[MCP remote server] Configuration", "server", server)
// Create HTTP client with custom roundtripper for bearer token injection
httpClient := &http.Client{
Timeout: 360 * time.Second,
Transport: newBearerTokenRoundTripper(server.Token, http.DefaultTransport),
}
transport := &mcp.StreamableClientTransport{Endpoint: server.URL, HTTPClient: httpClient}
mcpSession, err := client.Connect(ctx, transport, nil)
if err != nil {
xlog.Error("Failed to connect to MCP server", "error", err, "url", server.URL)
continue
}
xlog.Debug("[MCP remote server] Connected to MCP server", "url", server.URL)
cache.cache[name] = append(cache.cache[name], mcpSession)
allSessions = append(allSessions, mcpSession)
}
for _, server := range stdio.Servers {
xlog.Debug("[MCP stdio server] Configuration", "server", server)
command := exec.Command(server.Command, server.Args...)
command.Env = os.Environ()
for key, value := range server.Env {
command.Env = append(command.Env, key+"="+value)
}
transport := &mcp.CommandTransport{Command: command}
mcpSession, err := client.Connect(ctx, transport, nil)
if err != nil {
xlog.Error("Failed to start MCP server", "error", err, "command", command)
continue
}
xlog.Debug("[MCP stdio server] Connected to MCP server", "command", command)
cache.cache[name] = append(cache.cache[name], mcpSession)
allSessions = append(allSessions, mcpSession)
}
signals.RegisterGracefulTerminationHandler(func() {
for _, session := range allSessions {
session.Close()
}
cancel()
})
return allSessions, nil
}
// bearerTokenRoundTripper is a custom roundtripper that injects a bearer token
// into HTTP requests
type bearerTokenRoundTripper struct {
token string
base http.RoundTripper
}
// RoundTrip implements the http.RoundTripper interface
func (rt *bearerTokenRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if rt.token != "" {
req.Header.Set("Authorization", "Bearer "+rt.token)
}
return rt.base.RoundTrip(req)
}
// newBearerTokenRoundTripper creates a new roundtripper that injects the given token
func newBearerTokenRoundTripper(token string, base http.RoundTripper) http.RoundTripper {
if base == nil {
base = http.DefaultTransport
}
return &bearerTokenRoundTripper{
token: token,
base: base,
}
}
|