93 lines
2.7 KiB
Go
93 lines
2.7 KiB
Go
package mcpclient
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"os"
|
|
"os/exec"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
)
|
|
|
|
func discover(ctx context.Context, c ServerConfig, sec Secrets) ([]ToolInfo, error) {
|
|
tools := []ToolInfo{}
|
|
err := withSession(ctx, c, sec, func(session *mcp.ClientSession) error {
|
|
for tool, err := range session.Tools(ctx, nil) {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(tools) >= 1000 {
|
|
return fmt.Errorf("too many tools")
|
|
}
|
|
schema, err := json.MarshalIndent(tool.InputSchema, "", " ")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tools = append(tools, ToolInfo{Name: tool.Name, Description: tool.Description, InputSchema: string(schema)})
|
|
}
|
|
return nil
|
|
})
|
|
return tools, err
|
|
}
|
|
|
|
func withSession(ctx context.Context, c ServerConfig, sec Secrets, use func(*mcp.ClientSession) error) error {
|
|
bundled := c.Transport == "filesystem"
|
|
if bundled {
|
|
var err error
|
|
c, err = filesystemCommand(c)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
client := mcp.NewClient(&mcp.Implementation{Name: "iTuoke", Version: "1.0.0"}, &mcp.ClientOptions{Logger: slog.New(slog.NewTextHandler(io.Discard, nil))})
|
|
var transport mcp.Transport
|
|
if c.Transport == "stdio" {
|
|
cmd := exec.CommandContext(ctx, c.Command, c.Args...)
|
|
cmd.Dir = c.WorkingDirectory
|
|
cmd.Env = os.Environ()
|
|
if bundled {
|
|
filtered := cmd.Env[:0]
|
|
for _, entry := range cmd.Env {
|
|
key, _, _ := strings.Cut(entry, "=")
|
|
if !strings.EqualFold(key, "NODE_OPTIONS") && !strings.EqualFold(key, "NODE_PATH") {
|
|
filtered = append(filtered, entry)
|
|
}
|
|
}
|
|
cmd.Env = filtered
|
|
}
|
|
// Remove inherited duplicates (case-insensitively for Windows), then override.
|
|
for key, value := range sec.Env {
|
|
filtered := cmd.Env[:0]
|
|
for _, entry := range cmd.Env {
|
|
k, _, _ := strings.Cut(entry, "=")
|
|
if !strings.EqualFold(k, key) {
|
|
filtered = append(filtered, entry)
|
|
}
|
|
}
|
|
cmd.Env = append(filtered, key+"="+value)
|
|
}
|
|
configureCommand(cmd)
|
|
cmd.WaitDelay = time.Second
|
|
transport = &mcp.CommandTransport{Command: cmd, TerminateDuration: 500 * time.Millisecond}
|
|
} else {
|
|
base := http.DefaultTransport.(*http.Transport).Clone()
|
|
defer base.CloseIdleConnections()
|
|
transport = &mcp.StreamableClientTransport{
|
|
Endpoint: c.URL, MaxRetries: -1, DisableStandaloneSSE: true, MaxEventSize: 4 * 1024 * 1024,
|
|
HTTPClient: &http.Client{Timeout: time.Duration(c.TimeoutSeconds) * time.Second, Transport: headerTransport{base: base, headers: sec.Headers, ctx: ctx}, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }},
|
|
}
|
|
}
|
|
session, err := client.Connect(ctx, transport, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer session.Close()
|
|
return use(session)
|
|
}
|