mirror of
https://github.com/BishopFox/sliver.git
synced 2026-11-08 08:04:14 -10:00
Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d4728af1ea | ||
|
|
d1b7d91b1a | ||
|
|
cac6dced34 | ||
|
|
a35bf0c1fc | ||
|
|
67616e50c1 | ||
|
|
54de8a2a97 | ||
|
|
20916c655c | ||
|
|
6c18201ed6 | ||
|
|
29165aff0d | ||
|
|
704285f28d | ||
|
|
a007872fbc |
No files matched your search
@@ -38,6 +38,8 @@ jobs:
|
||||
steps:
|
||||
- name: Check Out Code
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Go 1.25
|
||||
uses: actions/setup-go@v6
|
||||
|
||||
@@ -2,6 +2,10 @@
|
||||
/.devcontainer
|
||||
/.idea
|
||||
|
||||
/artifacts
|
||||
/artifacts/**
|
||||
|
||||
|
||||
/Dockerfile.*
|
||||
/*.sdk.tar.xz
|
||||
|
||||
|
||||
@@ -73,6 +73,18 @@ func GetConsoleLogsDir() string {
|
||||
return consoleLogsDir
|
||||
}
|
||||
|
||||
// GetMCPLogsDir - Get the Sliver client MCP logs dir ~/.sliver-client/logs/mcp/
|
||||
func GetMCPLogsDir() string {
|
||||
mcpLogsDir := filepath.Join(GetClientLogsDir(), "mcp")
|
||||
if _, err := os.Stat(mcpLogsDir); os.IsNotExist(err) {
|
||||
err = os.MkdirAll(mcpLogsDir, 0700)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}
|
||||
return mcpLogsDir
|
||||
}
|
||||
|
||||
func assetVersion() string {
|
||||
appDir := GetRootAppDir()
|
||||
data, err := os.ReadFile(filepath.Join(appDir, versionFileName))
|
||||
|
||||
@@ -69,6 +69,9 @@ func init() {
|
||||
// other command trees below, such as the implant one.
|
||||
rootCmd.AddCommand(consoleCmd(con))
|
||||
|
||||
// MCP stdio server.
|
||||
rootCmd.AddCommand(mcpCmd(con))
|
||||
|
||||
// Implant.
|
||||
// The implant command allows users to run commands on slivers from their
|
||||
// system shell. It makes use of pre-runners for connecting to the server
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
package cli
|
||||
|
||||
/*
|
||||
Sliver Implant Framework
|
||||
Copyright (C) 2023 Bishop Fox
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
*/
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/bishopfox/sliver/client/assets"
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
clientmcp "github.com/bishopfox/sliver/client/mcp"
|
||||
"github.com/bishopfox/sliver/client/transport"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
func mcpCmd(con *console.SliverClient) *cobra.Command {
|
||||
mcpCmd := &cobra.Command{
|
||||
Use: "mcp",
|
||||
Short: "Start the MCP stdio server",
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
return runMCP(cmd, con)
|
||||
},
|
||||
}
|
||||
mcpCmd.Flags().String("config", "", "path to client config file")
|
||||
return mcpCmd
|
||||
}
|
||||
|
||||
func runMCP(cmd *cobra.Command, con *console.SliverClient) error {
|
||||
configPath, _ := cmd.Flags().GetString("config")
|
||||
config, err := loadMCPClientConfig(configPath)
|
||||
if err != nil {
|
||||
fmt.Printf("%s\n", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
rpc, ln, err := transport.MTLSConnect(config)
|
||||
if err != nil {
|
||||
fmt.Printf("Connection to server failed %s\n", err)
|
||||
return nil
|
||||
}
|
||||
defer ln.Close()
|
||||
|
||||
con.Rpc = rpc
|
||||
go handleConnectionLost(ln)
|
||||
|
||||
cfg := clientmcp.DefaultConfig()
|
||||
if err := clientmcp.ServeStdio(cfg, con.Rpc); err != nil {
|
||||
fmt.Printf("MCP server error: %s\n", err)
|
||||
return nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func loadMCPClientConfig(path string) (*assets.ClientConfig, error) {
|
||||
if path != "" {
|
||||
configPath := path
|
||||
if !filepath.IsAbs(configPath) {
|
||||
if _, err := os.Stat(configPath); err != nil {
|
||||
configPath = filepath.Join(assets.GetConfigDir(), configPath)
|
||||
}
|
||||
}
|
||||
return assets.ReadConfig(configPath)
|
||||
}
|
||||
|
||||
configs := assets.GetConfigs()
|
||||
switch len(configs) {
|
||||
case 0:
|
||||
return nil, fmt.Errorf("no config files found at %s (use --config)", assets.GetConfigDir())
|
||||
case 1:
|
||||
for _, config := range configs {
|
||||
return config, nil
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("multiple configs found; use --config to select one")
|
||||
}
|
||||
return nil, fmt.Errorf("no config files found")
|
||||
}
|
||||
@@ -40,6 +40,7 @@ Shared registration glue, scaffolding, and helpers for Sliver client console com
|
||||
- `licenses/` – Implements the 'licenses' command group for the Sliver client console.
|
||||
- `loot/` – Implements the 'loot' command group for the Sliver client console.
|
||||
- `monitor/` – Implements the 'monitor' command group for the Sliver client console.
|
||||
- `mcp/` – Implements the 'mcp' command group for the Sliver client console.
|
||||
- `network/` – Implements the 'network' command group for the Sliver client console.
|
||||
- `operators/` – Implements the 'operators' command group for the Sliver client console.
|
||||
- `pivots/` – Implements the 'pivots' command group for the Sliver client console.
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"github.com/bishopfox/sliver/client/command/flags"
|
||||
"github.com/bishopfox/sliver/client/command/help"
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
consts "github.com/bishopfox/sliver/client/constants"
|
||||
clientmcp "github.com/bishopfox/sliver/client/mcp"
|
||||
"github.com/rsteube/carapace"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/pflag"
|
||||
)
|
||||
|
||||
// Commands returns the `mcp` command and its subcommands.
|
||||
func Commands(con *console.SliverClient) []*cobra.Command {
|
||||
defaultConfig := clientmcp.DefaultConfig()
|
||||
|
||||
mcpCmd := &cobra.Command{
|
||||
Use: consts.MCPStr,
|
||||
Short: "Manage the MCP server",
|
||||
Long: help.GetHelpFor([]string{consts.MCPStr}),
|
||||
GroupID: consts.NetworkHelpGroup,
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
McpCmd(cmd, con, args)
|
||||
},
|
||||
}
|
||||
|
||||
startCmd := &cobra.Command{
|
||||
Use: consts.StartStr,
|
||||
Short: "Start the MCP server",
|
||||
Long: help.GetHelpFor([]string{consts.MCPStr, consts.StartStr}),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
McpStartCmd(cmd, con, args)
|
||||
},
|
||||
}
|
||||
flags.Bind("", false, startCmd, func(f *pflag.FlagSet) {
|
||||
f.String("transport", string(defaultConfig.Transport), "mcp transport (http or sse)")
|
||||
f.String("listen", defaultConfig.ListenAddress, "listen address for the MCP server")
|
||||
f.String("name", defaultConfig.ServerName, "server name for MCP initialize")
|
||||
f.String("version", defaultConfig.ServerVersion, "server version for MCP initialize")
|
||||
})
|
||||
flags.BindFlagCompletions(startCmd, func(comp *carapace.ActionMap) {
|
||||
(*comp)["transport"] = carapace.ActionValues("http", "sse")
|
||||
})
|
||||
mcpCmd.AddCommand(startCmd)
|
||||
|
||||
stopCmd := &cobra.Command{
|
||||
Use: consts.StopStr,
|
||||
Short: "Stop the MCP server",
|
||||
Long: help.GetHelpFor([]string{consts.MCPStr, consts.StopStr}),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
McpStopCmd(cmd, con, args)
|
||||
},
|
||||
}
|
||||
mcpCmd.AddCommand(stopCmd)
|
||||
|
||||
consoleCmd := &cobra.Command{
|
||||
Use: "console",
|
||||
Aliases: []string{"repl"},
|
||||
Short: "Open an MCP console session",
|
||||
Long: help.GetHelpFor([]string{consts.MCPStr, "console"}),
|
||||
Run: func(cmd *cobra.Command, args []string) {
|
||||
McpConsoleCmd(cmd, con, args)
|
||||
},
|
||||
}
|
||||
flags.Bind("", false, consoleCmd, func(f *pflag.FlagSet) {
|
||||
f.String("transport", string(defaultConfig.Transport), "mcp transport (http or sse)")
|
||||
f.String("url", "", "mcp server URL (defaults to local server settings)")
|
||||
})
|
||||
flags.BindFlagCompletions(consoleCmd, func(comp *carapace.ActionMap) {
|
||||
(*comp)["transport"] = carapace.ActionValues("http", "sse")
|
||||
})
|
||||
mcpCmd.AddCommand(consoleCmd)
|
||||
|
||||
return []*cobra.Command{mcpCmd}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
clientmcp "github.com/bishopfox/sliver/client/mcp"
|
||||
"github.com/bishopfox/sliver/client/version"
|
||||
mcpclient "github.com/mark3labs/mcp-go/client"
|
||||
mcpapi "github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// McpConsoleCmd connects to an MCP server and starts a minimal console.
|
||||
func McpConsoleCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
|
||||
rawTransport, _ := cmd.Flags().GetString("transport")
|
||||
transport, err := clientmcp.ParseTransport(rawTransport)
|
||||
if err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
targetURL, _ := cmd.Flags().GetString("url")
|
||||
if targetURL == "" {
|
||||
cfg := clientmcp.GetStatus().Config.WithDefaults()
|
||||
cfg.Transport = transport
|
||||
targetURL, err = cfg.EndpointURL()
|
||||
if err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
client, err := newConsoleClient(transport, targetURL)
|
||||
if err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
if err := client.Start(ctx); err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
_, err = client.Initialize(ctx, mcpapi.InitializeRequest{
|
||||
Params: mcpapi.InitializeParams{
|
||||
ProtocolVersion: mcpapi.LATEST_PROTOCOL_VERSION,
|
||||
ClientInfo: mcpapi.Implementation{
|
||||
Name: "sliver-mcp-console",
|
||||
Version: version.Version,
|
||||
},
|
||||
Capabilities: mcpapi.ClientCapabilities{},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
con.PrintInfof("Connected to MCP server (%s)\n", targetURL)
|
||||
runConsoleREPL(ctx, client, cmd.InOrStdin(), cmd.OutOrStdout())
|
||||
}
|
||||
|
||||
func newConsoleClient(transport clientmcp.Transport, targetURL string) (*mcpclient.Client, error) {
|
||||
switch transport {
|
||||
case clientmcp.TransportHTTP:
|
||||
return mcpclient.NewStreamableHttpClient(targetURL)
|
||||
case clientmcp.TransportSSE:
|
||||
return mcpclient.NewSSEMCPClient(targetURL)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported transport %q", transport)
|
||||
}
|
||||
}
|
||||
|
||||
func runConsoleREPL(ctx context.Context, client *mcpclient.Client, in io.Reader, out io.Writer) {
|
||||
reader := bufio.NewReader(in)
|
||||
for {
|
||||
fmt.Fprint(out, "mcp> ")
|
||||
line, err := reader.ReadString('\n')
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) {
|
||||
fmt.Fprintln(out)
|
||||
return
|
||||
}
|
||||
fmt.Fprintf(out, "read error: %s\n", err)
|
||||
return
|
||||
}
|
||||
line = strings.TrimSpace(line)
|
||||
switch line {
|
||||
case "":
|
||||
continue
|
||||
case "exit", "quit":
|
||||
return
|
||||
case "help", "?":
|
||||
printConsoleHelp(out)
|
||||
case "tools":
|
||||
printConsoleTools(ctx, client, out)
|
||||
default:
|
||||
fmt.Fprintf(out, "unknown command: %s\n", line)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func printConsoleHelp(out io.Writer) {
|
||||
fmt.Fprintln(out, "commands: tools, help, exit")
|
||||
}
|
||||
|
||||
func printConsoleTools(ctx context.Context, client *mcpclient.Client, out io.Writer) {
|
||||
tools, err := client.ListTools(ctx, mcpapi.ListToolsRequest{})
|
||||
if err != nil {
|
||||
fmt.Fprintf(out, "tools error: %s\n", err)
|
||||
return
|
||||
}
|
||||
if len(tools.Tools) == 0 {
|
||||
fmt.Fprintln(out, "no tools available")
|
||||
return
|
||||
}
|
||||
for _, tool := range tools.Tools {
|
||||
if tool.Description != "" {
|
||||
fmt.Fprintf(out, "- %s: %s\n", tool.Name, tool.Description)
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(out, "- %s\n", tool.Name)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
clientmcp "github.com/bishopfox/sliver/client/mcp"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// McpCmd prints the current MCP server state.
|
||||
func McpCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
|
||||
status := clientmcp.GetStatus()
|
||||
state := "stopped"
|
||||
if status.Running {
|
||||
state = "running"
|
||||
}
|
||||
|
||||
con.Printf("Status: %s\n", state)
|
||||
con.Printf("Transport: %s\n", status.Config.Transport)
|
||||
con.Printf("Listen: %s\n", status.Config.ListenAddress)
|
||||
|
||||
endpoint, err := status.Config.EndpointURL()
|
||||
if err == nil {
|
||||
con.Printf("Endpoint: %s\n", endpoint)
|
||||
} else {
|
||||
con.PrintErrorf("Endpoint: %s\n", err)
|
||||
}
|
||||
|
||||
if status.Running && !status.StartedAt.IsZero() {
|
||||
uptime := time.Since(status.StartedAt).Truncate(time.Second)
|
||||
con.Printf("Uptime: %s\n", uptime)
|
||||
}
|
||||
if status.LastError != "" {
|
||||
con.PrintErrorf("Last error: %s\n", status.LastError)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
slivermcp "github.com/bishopfox/sliver/client/mcp"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
// McpStartCmd starts the local MCP server.
|
||||
func McpStartCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
|
||||
rawTransport, _ := cmd.Flags().GetString("transport")
|
||||
transport, err := slivermcp.ParseTransport(rawTransport)
|
||||
if err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
listen, _ := cmd.Flags().GetString("listen")
|
||||
name, _ := cmd.Flags().GetString("name")
|
||||
version, _ := cmd.Flags().GetString("version")
|
||||
|
||||
cfg := slivermcp.Config{
|
||||
Transport: transport,
|
||||
ListenAddress: listen,
|
||||
ServerName: name,
|
||||
ServerVersion: version,
|
||||
}.WithDefaults()
|
||||
|
||||
if err := slivermcp.Start(cfg, con.Rpc); err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
con.PrintInfof("Starting MCP server (%s) on %s\n", cfg.Transport, cfg.ListenAddress)
|
||||
endpoint, err := cfg.EndpointURL()
|
||||
if err == nil {
|
||||
con.PrintInfof("Endpoint: %s\n", endpoint)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/bishopfox/sliver/client/console"
|
||||
clientmcp "github.com/bishopfox/sliver/client/mcp"
|
||||
"github.com/spf13/cobra"
|
||||
)
|
||||
|
||||
const mcpStopTimeout = 5 * time.Second
|
||||
|
||||
// McpStopCmd stops the local MCP server.
|
||||
func McpStopCmd(cmd *cobra.Command, con *console.SliverClient, args []string) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), mcpStopTimeout)
|
||||
defer cancel()
|
||||
|
||||
if err := clientmcp.Stop(ctx); err != nil {
|
||||
con.PrintErrorf("%s\n", err)
|
||||
return
|
||||
}
|
||||
con.PrintInfof("MCP server stopped\n")
|
||||
}
|
||||
@@ -39,6 +39,7 @@ import (
|
||||
"github.com/bishopfox/sliver/client/command/jobs"
|
||||
"github.com/bishopfox/sliver/client/command/licenses"
|
||||
"github.com/bishopfox/sliver/client/command/loot"
|
||||
"github.com/bishopfox/sliver/client/command/mcp"
|
||||
"github.com/bishopfox/sliver/client/command/monitor"
|
||||
"github.com/bishopfox/sliver/client/command/operators"
|
||||
"github.com/bishopfox/sliver/client/command/reaction"
|
||||
@@ -106,6 +107,7 @@ func ServerCommands(con *client.SliverClient, serverCmds func() []*cobra.Command
|
||||
// C2 Network
|
||||
bind(consts.NetworkHelpGroup,
|
||||
jobs.Commands,
|
||||
mcp.Commands,
|
||||
websites.Commands,
|
||||
wireguard.Commands,
|
||||
c2profiles.Commands,
|
||||
|
||||
@@ -276,6 +276,7 @@ const (
|
||||
WgPortFwdStr = "wg-portfwd"
|
||||
MonitorStr = "monitor"
|
||||
MonitorConfigStr = "config"
|
||||
MCPStr = "mcp"
|
||||
SSHStr = "ssh"
|
||||
DLLHijackStr = "dllhijack"
|
||||
InteractiveStr = "interactive"
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/bishopfox/sliver/client/version"
|
||||
)
|
||||
|
||||
// Transport identifies the MCP server transport to expose.
|
||||
type Transport string
|
||||
|
||||
const (
|
||||
TransportHTTP Transport = "http"
|
||||
TransportSSE Transport = "sse"
|
||||
)
|
||||
|
||||
// Config controls the local MCP server settings.
|
||||
type Config struct {
|
||||
Transport Transport
|
||||
ListenAddress string
|
||||
ServerName string
|
||||
ServerVersion string
|
||||
}
|
||||
|
||||
// DefaultConfig returns the baseline MCP configuration.
|
||||
func DefaultConfig() Config {
|
||||
return Config{
|
||||
Transport: TransportSSE,
|
||||
ListenAddress: "127.0.0.1:8080",
|
||||
ServerName: "Sliver MCP",
|
||||
ServerVersion: version.Version,
|
||||
}
|
||||
}
|
||||
|
||||
// WithDefaults fills any empty fields with defaults.
|
||||
func (c Config) WithDefaults() Config {
|
||||
defaults := DefaultConfig()
|
||||
if c.Transport == "" {
|
||||
c.Transport = defaults.Transport
|
||||
}
|
||||
if c.ListenAddress == "" {
|
||||
c.ListenAddress = defaults.ListenAddress
|
||||
}
|
||||
if c.ServerName == "" {
|
||||
c.ServerName = defaults.ServerName
|
||||
}
|
||||
if c.ServerVersion == "" {
|
||||
c.ServerVersion = defaults.ServerVersion
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
// ParseTransport validates an input transport string.
|
||||
func ParseTransport(raw string) (Transport, error) {
|
||||
value := strings.ToLower(strings.TrimSpace(raw))
|
||||
if value == "" {
|
||||
return DefaultConfig().Transport, nil
|
||||
}
|
||||
switch Transport(value) {
|
||||
case TransportHTTP, TransportSSE:
|
||||
return Transport(value), nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported transport %q", raw)
|
||||
}
|
||||
}
|
||||
|
||||
// Validate ensures the configuration can be used for a server start.
|
||||
func (c Config) Validate() error {
|
||||
if c.Transport != TransportHTTP && c.Transport != TransportSSE {
|
||||
return fmt.Errorf("unsupported transport %q", c.Transport)
|
||||
}
|
||||
if c.ListenAddress == "" {
|
||||
return fmt.Errorf("listen address is required")
|
||||
}
|
||||
if _, _, err := net.SplitHostPort(c.ListenAddress); err != nil {
|
||||
return fmt.Errorf("invalid listen address %q: %w", c.ListenAddress, err)
|
||||
}
|
||||
if c.ServerName == "" {
|
||||
return fmt.Errorf("server name is required")
|
||||
}
|
||||
if c.ServerVersion == "" {
|
||||
return fmt.Errorf("server version is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// EndpointURL returns the base URL for clients to connect.
|
||||
func (c Config) EndpointURL() (string, error) {
|
||||
cfg := c.WithDefaults()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
host, port, err := net.SplitHostPort(cfg.ListenAddress)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if host == "" {
|
||||
host = "127.0.0.1"
|
||||
}
|
||||
|
||||
endpoint := url.URL{
|
||||
Scheme: "http",
|
||||
Host: net.JoinHostPort(host, port),
|
||||
}
|
||||
switch cfg.Transport {
|
||||
case TransportHTTP:
|
||||
endpoint.Path = "/mcp"
|
||||
case TransportSSE:
|
||||
endpoint.Path = "/sse"
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported transport %q", cfg.Transport)
|
||||
}
|
||||
return endpoint.String(), nil
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,67 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/bishopfox/sliver/client/assets"
|
||||
)
|
||||
|
||||
const mcpLogTimeFormat = "2006-01-02_15-04-05"
|
||||
|
||||
func newMCPLogger() (*log.Logger, *os.File, error) {
|
||||
logsDir := assets.GetMCPLogsDir()
|
||||
logPath := filepath.Join(logsDir, fmt.Sprintf("%s.log", time.Now().Format(mcpLogTimeFormat)))
|
||||
logFile, err := os.OpenFile(logPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0o600)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
logger := log.New(logFile, "", log.LstdFlags|log.Lshortfile)
|
||||
return logger, logFile, nil
|
||||
}
|
||||
|
||||
func (m *mcpManager) logf(format string, args ...any) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
logger := m.logger
|
||||
m.mu.Unlock()
|
||||
if logger != nil {
|
||||
logger.Printf(format, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SliverMCPServer) logf(format string, args ...any) {
|
||||
if s == nil || s.logger == nil {
|
||||
return
|
||||
}
|
||||
s.logger.Printf(format, args...)
|
||||
}
|
||||
|
||||
func (s *SliverMCPServer) logToolCall(tool, sessionID, beaconID string, extras ...string) {
|
||||
if s == nil || s.logger == nil {
|
||||
return
|
||||
}
|
||||
parts := make([]string, 0, 2+len(extras))
|
||||
if sessionID != "" {
|
||||
parts = append(parts, fmt.Sprintf("session_id=%s", sessionID))
|
||||
}
|
||||
if beaconID != "" {
|
||||
parts = append(parts, fmt.Sprintf("beacon_id=%s", beaconID))
|
||||
}
|
||||
for _, extra := range extras {
|
||||
if extra != "" {
|
||||
parts = append(parts, extra)
|
||||
}
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
s.logger.Printf("mcp tool=%s", tool)
|
||||
return
|
||||
}
|
||||
s.logger.Printf("mcp tool=%s %s", tool, strings.Join(parts, " "))
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/rpcpb"
|
||||
mcpserver "github.com/mark3labs/mcp-go/server"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAlreadyRunning = errors.New("mcp server already running")
|
||||
ErrNotRunning = errors.New("mcp server not running")
|
||||
)
|
||||
|
||||
// Status describes the running state of the MCP server.
|
||||
type Status struct {
|
||||
Running bool
|
||||
StartedAt time.Time
|
||||
Config Config
|
||||
LastError string
|
||||
}
|
||||
|
||||
type serverTransport interface {
|
||||
Start(addr string) error
|
||||
Shutdown(ctx context.Context) error
|
||||
}
|
||||
|
||||
type mcpManager struct {
|
||||
mu sync.Mutex
|
||||
cfg Config
|
||||
running bool
|
||||
startedAt time.Time
|
||||
lastErr string
|
||||
server *SliverMCPServer
|
||||
transport serverTransport
|
||||
done chan struct{}
|
||||
logger *log.Logger
|
||||
logFile *os.File
|
||||
}
|
||||
|
||||
func newManager() *mcpManager {
|
||||
return &mcpManager{
|
||||
cfg: DefaultConfig(),
|
||||
}
|
||||
}
|
||||
|
||||
var defaultManager = newManager()
|
||||
|
||||
// GetStatus returns the current MCP server state.
|
||||
func GetStatus() Status {
|
||||
return defaultManager.status()
|
||||
}
|
||||
|
||||
// Start launches the MCP server using the provided configuration.
|
||||
func Start(cfg Config, rpc rpcpb.SliverRPCClient) error {
|
||||
return defaultManager.start(cfg, rpc)
|
||||
}
|
||||
|
||||
// ServeStdio runs the MCP server over stdio using the provided configuration.
|
||||
func ServeStdio(cfg Config, rpc rpcpb.SliverRPCClient) error {
|
||||
cfg = cfg.WithDefaults()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
logger, logFile, err := newMCPLogger()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer logFile.Close()
|
||||
logger.Printf("starting mcp stdio server")
|
||||
srv := newServer(cfg, rpc, logger)
|
||||
err = mcpserver.ServeStdio(srv.server)
|
||||
if err != nil {
|
||||
logger.Printf("mcp stdio server error: %v", err)
|
||||
} else {
|
||||
logger.Printf("mcp stdio server stopped")
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// Stop shuts down the MCP server.
|
||||
func Stop(ctx context.Context) error {
|
||||
return defaultManager.stop(ctx)
|
||||
}
|
||||
|
||||
func (m *mcpManager) status() Status {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return Status{
|
||||
Running: m.running,
|
||||
StartedAt: m.startedAt,
|
||||
Config: m.cfg,
|
||||
LastError: m.lastErr,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mcpManager) start(cfg Config, rpc rpcpb.SliverRPCClient) error {
|
||||
cfg = cfg.WithDefaults()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
if m.running {
|
||||
m.mu.Unlock()
|
||||
return ErrAlreadyRunning
|
||||
}
|
||||
logger, logFile, err := newMCPLogger()
|
||||
if err != nil {
|
||||
m.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
mcpServer := newServer(cfg, rpc, logger)
|
||||
var transport serverTransport
|
||||
switch cfg.Transport {
|
||||
case TransportHTTP:
|
||||
transport = mcpserver.NewStreamableHTTPServer(mcpServer.server)
|
||||
case TransportSSE:
|
||||
transport = mcpserver.NewSSEServer(mcpServer.server)
|
||||
default:
|
||||
m.mu.Unlock()
|
||||
logFile.Close()
|
||||
return errors.New("unsupported transport")
|
||||
}
|
||||
|
||||
m.cfg = cfg
|
||||
m.running = true
|
||||
m.startedAt = time.Now()
|
||||
m.lastErr = ""
|
||||
m.server = mcpServer
|
||||
m.transport = transport
|
||||
m.done = make(chan struct{})
|
||||
m.logger = logger
|
||||
m.logFile = logFile
|
||||
addr := cfg.ListenAddress
|
||||
done := m.done
|
||||
m.mu.Unlock()
|
||||
|
||||
logger.Printf("starting mcp server transport=%s listen=%s", cfg.Transport, cfg.ListenAddress)
|
||||
go m.run(addr, transport, done)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mcpManager) run(addr string, transport serverTransport, done chan struct{}) {
|
||||
err := transport.Start(addr)
|
||||
if err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
m.logf("mcp server error: %v", err)
|
||||
m.mu.Lock()
|
||||
m.lastErr = err.Error()
|
||||
m.mu.Unlock()
|
||||
}
|
||||
m.logf("mcp server stopped")
|
||||
|
||||
m.mu.Lock()
|
||||
m.running = false
|
||||
m.transport = nil
|
||||
m.server = nil
|
||||
logFile := m.logFile
|
||||
m.logFile = nil
|
||||
m.logger = nil
|
||||
m.mu.Unlock()
|
||||
if logFile != nil {
|
||||
logFile.Close()
|
||||
}
|
||||
close(done)
|
||||
}
|
||||
|
||||
func (m *mcpManager) stop(ctx context.Context) error {
|
||||
m.mu.Lock()
|
||||
if !m.running {
|
||||
m.mu.Unlock()
|
||||
return ErrNotRunning
|
||||
}
|
||||
transport := m.transport
|
||||
done := m.done
|
||||
m.mu.Unlock()
|
||||
|
||||
if transport == nil {
|
||||
return ErrNotRunning
|
||||
}
|
||||
m.logf("mcp server shutdown requested")
|
||||
if err := transport.Shutdown(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
if done == nil {
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
|
||||
"github.com/bishopfox/sliver/protobuf/commonpb"
|
||||
"github.com/bishopfox/sliver/protobuf/rpcpb"
|
||||
mcpapi "github.com/mark3labs/mcp-go/mcp"
|
||||
mcpserver "github.com/mark3labs/mcp-go/server"
|
||||
)
|
||||
|
||||
const (
|
||||
listSessionsAndBeaconsToolName = "list_sessions_and_beacons"
|
||||
)
|
||||
|
||||
type sessionSummary struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Transport string `json:"transport"`
|
||||
RemoteAddress string `json:"remote_address"`
|
||||
Hostname string `json:"hostname"`
|
||||
Username string `json:"username"`
|
||||
PID int32 `json:"pid"`
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
Locale string `json:"locale"`
|
||||
LastCheckin int64 `json:"last_checkin"`
|
||||
IsDead bool `json:"is_dead"`
|
||||
Integrity string `json:"integrity"`
|
||||
}
|
||||
|
||||
type beaconSummary struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Transport string `json:"transport"`
|
||||
RemoteAddress string `json:"remote_address"`
|
||||
Hostname string `json:"hostname"`
|
||||
Username string `json:"username"`
|
||||
PID int32 `json:"pid"`
|
||||
OS string `json:"os"`
|
||||
Arch string `json:"arch"`
|
||||
Locale string `json:"locale"`
|
||||
LastCheckin int64 `json:"last_checkin"`
|
||||
NextCheckin int64 `json:"next_checkin"`
|
||||
Interval int64 `json:"interval"`
|
||||
Jitter int64 `json:"jitter"`
|
||||
TasksCount int64 `json:"tasks_count"`
|
||||
TasksCountCompleted int64 `json:"tasks_count_completed"`
|
||||
IsDead bool `json:"is_dead"`
|
||||
Integrity string `json:"integrity"`
|
||||
}
|
||||
|
||||
type listSessionsAndBeaconsResult struct {
|
||||
Sessions []sessionSummary `json:"sessions"`
|
||||
SessionsCount int `json:"sessions_count"`
|
||||
Beacons []beaconSummary `json:"beacons"`
|
||||
BeaconsCount int `json:"beacons_count"`
|
||||
}
|
||||
|
||||
// SliverMCPServer wraps the MCP server with Sliver RPC access for handlers.
|
||||
type SliverMCPServer struct {
|
||||
Rpc rpcpb.SliverRPCClient
|
||||
server *mcpserver.MCPServer
|
||||
logger *log.Logger
|
||||
}
|
||||
|
||||
func newServer(cfg Config, rpc rpcpb.SliverRPCClient, logger *log.Logger) *SliverMCPServer {
|
||||
base := mcpserver.NewMCPServer(
|
||||
cfg.ServerName,
|
||||
cfg.ServerVersion,
|
||||
mcpserver.WithToolCapabilities(false),
|
||||
)
|
||||
|
||||
listSessionsAndBeaconsTool := mcpapi.NewTool(
|
||||
listSessionsAndBeaconsToolName,
|
||||
mcpapi.WithDescription("List active sessions and beacons."),
|
||||
mcpapi.WithReadOnlyHintAnnotation(true),
|
||||
mcpapi.WithDestructiveHintAnnotation(false),
|
||||
mcpapi.WithIdempotentHintAnnotation(true),
|
||||
)
|
||||
lsTool := mcpapi.NewTool(
|
||||
lsToolName,
|
||||
mcpapi.WithDescription("List the contents of a remote directory."),
|
||||
mcpapi.WithInputSchema[lsArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(true),
|
||||
mcpapi.WithDestructiveHintAnnotation(false),
|
||||
mcpapi.WithIdempotentHintAnnotation(true),
|
||||
)
|
||||
cdTool := mcpapi.NewTool(
|
||||
cdToolName,
|
||||
mcpapi.WithDescription("Change the current directory on a remote session or beacon."),
|
||||
mcpapi.WithInputSchema[cdArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(false),
|
||||
mcpapi.WithDestructiveHintAnnotation(false),
|
||||
)
|
||||
catTool := mcpapi.NewTool(
|
||||
catToolName,
|
||||
mcpapi.WithDescription("Download and return the contents of a remote file."),
|
||||
mcpapi.WithInputSchema[catArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(true),
|
||||
mcpapi.WithDestructiveHintAnnotation(false),
|
||||
mcpapi.WithIdempotentHintAnnotation(true),
|
||||
)
|
||||
pwdTool := mcpapi.NewTool(
|
||||
pwdToolName,
|
||||
mcpapi.WithDescription("Return the current working directory of a remote session or beacon."),
|
||||
mcpapi.WithInputSchema[pwdArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(true),
|
||||
mcpapi.WithDestructiveHintAnnotation(false),
|
||||
mcpapi.WithIdempotentHintAnnotation(true),
|
||||
)
|
||||
rmTool := mcpapi.NewTool(
|
||||
rmToolName,
|
||||
mcpapi.WithDescription("Remove a file or directory on the remote target."),
|
||||
mcpapi.WithInputSchema[rmArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(false),
|
||||
mcpapi.WithDestructiveHintAnnotation(true),
|
||||
)
|
||||
mvTool := mcpapi.NewTool(
|
||||
mvToolName,
|
||||
mcpapi.WithDescription("Move or rename a file or directory on the remote target."),
|
||||
mcpapi.WithInputSchema[mvArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(false),
|
||||
mcpapi.WithDestructiveHintAnnotation(true),
|
||||
)
|
||||
cpTool := mcpapi.NewTool(
|
||||
cpToolName,
|
||||
mcpapi.WithDescription("Copy a file on the remote target."),
|
||||
mcpapi.WithInputSchema[cpArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(false),
|
||||
mcpapi.WithDestructiveHintAnnotation(true),
|
||||
)
|
||||
mkdirTool := mcpapi.NewTool(
|
||||
mkdirToolName,
|
||||
mcpapi.WithDescription("Create a directory on the remote target."),
|
||||
mcpapi.WithInputSchema[mkdirArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(false),
|
||||
mcpapi.WithDestructiveHintAnnotation(true),
|
||||
)
|
||||
chmodTool := mcpapi.NewTool(
|
||||
chmodToolName,
|
||||
mcpapi.WithDescription("Change file or directory permissions on the remote target."),
|
||||
mcpapi.WithInputSchema[chmodArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(false),
|
||||
mcpapi.WithDestructiveHintAnnotation(true),
|
||||
)
|
||||
chownTool := mcpapi.NewTool(
|
||||
chownToolName,
|
||||
mcpapi.WithDescription("Change file or directory ownership on the remote target."),
|
||||
mcpapi.WithInputSchema[chownArgs](),
|
||||
mcpapi.WithReadOnlyHintAnnotation(false),
|
||||
mcpapi.WithDestructiveHintAnnotation(true),
|
||||
)
|
||||
srv := &SliverMCPServer{
|
||||
Rpc: rpc,
|
||||
server: base,
|
||||
logger: logger,
|
||||
}
|
||||
srv.server.AddTool(listSessionsAndBeaconsTool, srv.listSessionsAndBeaconsHandler)
|
||||
srv.server.AddTool(lsTool, srv.lsHandler)
|
||||
srv.server.AddTool(cdTool, srv.cdHandler)
|
||||
srv.server.AddTool(catTool, srv.catHandler)
|
||||
srv.server.AddTool(pwdTool, srv.pwdHandler)
|
||||
srv.server.AddTool(rmTool, srv.rmHandler)
|
||||
srv.server.AddTool(mvTool, srv.mvHandler)
|
||||
srv.server.AddTool(cpTool, srv.cpHandler)
|
||||
srv.server.AddTool(mkdirTool, srv.mkdirHandler)
|
||||
srv.server.AddTool(chmodTool, srv.chmodHandler)
|
||||
srv.server.AddTool(chownTool, srv.chownHandler)
|
||||
return srv
|
||||
}
|
||||
|
||||
func (s *SliverMCPServer) listSessionsAndBeaconsHandler(ctx context.Context, _ mcpapi.CallToolRequest) (*mcpapi.CallToolResult, error) {
|
||||
if s.Rpc == nil {
|
||||
return mcpapi.NewToolResultError("rpc client not configured"), nil
|
||||
}
|
||||
|
||||
s.logToolCall(listSessionsAndBeaconsToolName, "", "")
|
||||
|
||||
sessionsResp, err := s.Rpc.GetSessions(ctx, &commonpb.Empty{})
|
||||
if err != nil {
|
||||
return mcpapi.NewToolResultErrorFromErr("failed to list sessions", err), nil
|
||||
}
|
||||
|
||||
beaconsResp, err := s.Rpc.GetBeacons(ctx, &commonpb.Empty{})
|
||||
if err != nil {
|
||||
return mcpapi.NewToolResultErrorFromErr("failed to list beacons", err), nil
|
||||
}
|
||||
|
||||
result := listSessionsAndBeaconsResult{
|
||||
Sessions: make([]sessionSummary, 0, len(sessionsResp.GetSessions())),
|
||||
Beacons: make([]beaconSummary, 0, len(beaconsResp.GetBeacons())),
|
||||
}
|
||||
|
||||
for _, session := range sessionsResp.GetSessions() {
|
||||
if session == nil {
|
||||
continue
|
||||
}
|
||||
result.Sessions = append(result.Sessions, sessionSummary{
|
||||
ID: session.ID,
|
||||
Name: session.Name,
|
||||
Transport: session.Transport,
|
||||
RemoteAddress: session.RemoteAddress,
|
||||
Hostname: session.Hostname,
|
||||
Username: session.Username,
|
||||
PID: session.PID,
|
||||
OS: session.OS,
|
||||
Arch: session.Arch,
|
||||
Locale: session.Locale,
|
||||
LastCheckin: session.LastCheckin,
|
||||
IsDead: session.IsDead,
|
||||
Integrity: session.Integrity,
|
||||
})
|
||||
}
|
||||
|
||||
for _, beacon := range beaconsResp.GetBeacons() {
|
||||
if beacon == nil {
|
||||
continue
|
||||
}
|
||||
result.Beacons = append(result.Beacons, beaconSummary{
|
||||
ID: beacon.ID,
|
||||
Name: beacon.Name,
|
||||
Transport: beacon.Transport,
|
||||
RemoteAddress: beacon.RemoteAddress,
|
||||
Hostname: beacon.Hostname,
|
||||
Username: beacon.Username,
|
||||
PID: beacon.PID,
|
||||
OS: beacon.OS,
|
||||
Arch: beacon.Arch,
|
||||
Locale: beacon.Locale,
|
||||
LastCheckin: beacon.LastCheckin,
|
||||
NextCheckin: beacon.NextCheckin,
|
||||
Interval: beacon.Interval,
|
||||
Jitter: beacon.Jitter,
|
||||
TasksCount: beacon.TasksCount,
|
||||
TasksCountCompleted: beacon.TasksCountCompleted,
|
||||
IsDead: beacon.IsDead,
|
||||
Integrity: beacon.Integrity,
|
||||
})
|
||||
}
|
||||
|
||||
result.SessionsCount = len(result.Sessions)
|
||||
result.BeaconsCount = len(result.Beacons)
|
||||
|
||||
return mcpapi.NewToolResultStructuredOnly(result), nil
|
||||
}
|
||||
@@ -68,6 +68,16 @@ require (
|
||||
tailscale.com v1.90.6
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/bahlo/generic-list-go v0.2.0 // indirect
|
||||
github.com/buger/jsonparser v1.1.1 // indirect
|
||||
github.com/invopop/jsonschema v0.13.0 // indirect
|
||||
github.com/mailru/easyjson v0.7.7 // indirect
|
||||
github.com/spf13/cast v1.7.1 // indirect
|
||||
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
|
||||
)
|
||||
|
||||
require (
|
||||
cloud.google.com/go/compute/metadata v0.7.0 // indirect
|
||||
filippo.io/edwards25519 v1.1.0 // indirect
|
||||
@@ -133,6 +143,7 @@ require (
|
||||
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
|
||||
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
|
||||
github.com/lxn/win v0.0.0-20210218163916-a377121e959e // indirect
|
||||
github.com/mark3labs/mcp-go v0.43.2
|
||||
github.com/mattn/go-colorable v0.1.14 // indirect
|
||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/mattn/go-localereader v0.0.1 // indirect
|
||||
|
||||
@@ -84,7 +84,11 @@ github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiE
|
||||
github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
|
||||
github.com/aymanbagabas/go-udiff v0.3.1 h1:LV+qyBQ2pqe0u42ZsUEtPiCaUoqgA9gYRDs3vj1nolY=
|
||||
github.com/aymanbagabas/go-udiff v0.3.1/go.mod h1:G0fsKmG+P6ylD0r6N/KgQD/nWzgfnl8ZBcNLgcbrw8E=
|
||||
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
|
||||
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
|
||||
github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
|
||||
github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
|
||||
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
|
||||
github.com/carapace-sh/carapace v1.9.0 h1:IuEP6YeeK4fhrH7MoOyL0I3GMhvROerkw6YGxkx8Fbw=
|
||||
github.com/carapace-sh/carapace v1.9.0/go.mod h1:Zs3DpsawpFm+8fEfRVYE9EzMBYH8u7S3Aw9yAAKX0JM=
|
||||
github.com/carapace-sh/carapace-shlex v1.1.0 h1:58QSJTUxghpa/rEn6+Z9UM/Hy0ipVHDUn9S1mp4i+ZE=
|
||||
@@ -253,6 +257,8 @@ github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2
|
||||
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
|
||||
github.com/insomniacslk/dhcp v0.0.0-20231206064809-8c70d406f6d2 h1:9K06NfxkBh25x56yVhWWlKFE8YpicaSfHwoV8SFbueA=
|
||||
github.com/insomniacslk/dhcp v0.0.0-20231206064809-8c70d406f6d2/go.mod h1:3A9PQ1cunSDF/1rbTq99Ts4pVnycWg+vlPkfeD2NLFI=
|
||||
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
|
||||
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
|
||||
@@ -274,6 +280,7 @@ github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD
|
||||
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
|
||||
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
|
||||
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
|
||||
github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
|
||||
github.com/jsimonetti/rtnetlink v1.4.2 h1:Df9w9TZ3npHTyDn0Ev9e1uzmN2odmXd0QX+J5GTEn90=
|
||||
github.com/jsimonetti/rtnetlink v1.4.2/go.mod h1:92s6LJdE+1iOrw+F2/RO7LYI2Qd8pPpFNNUYW06gcoM=
|
||||
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs=
|
||||
@@ -306,6 +313,10 @@ github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69
|
||||
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
|
||||
github.com/lxn/win v0.0.0-20210218163916-a377121e959e h1:H+t6A/QJMbhCSEH5rAuRxh+CtW96g0Or0Fxa9IKr4uc=
|
||||
github.com/lxn/win v0.0.0-20210218163916-a377121e959e/go.mod h1:KxxjdtRkfNoYDCUP5ryK7XJJNTnpC8atvtmTheChOtk=
|
||||
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
|
||||
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
|
||||
github.com/mark3labs/mcp-go v0.43.2 h1:21PUSlWWiSbUPQwXIJ5WKlETixpFpq+WBpbMGDSVy/I=
|
||||
github.com/mark3labs/mcp-go v0.43.2/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw=
|
||||
github.com/mattn/go-colorable v0.1.8/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
|
||||
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
|
||||
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
|
||||
@@ -399,6 +410,8 @@ github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6Mwd
|
||||
github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
|
||||
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
|
||||
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
|
||||
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
|
||||
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
|
||||
github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
|
||||
github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
|
||||
github.com/spf13/cobra v1.10.1 h1:lJeBwCfmrnXthfAupyUTzJ/J4Nc1RsHC/mSRU2dll/s=
|
||||
@@ -455,6 +468,8 @@ github.com/ulikunitz/xz v0.5.15 h1:9DNdB5s+SgV3bQ2ApL10xRc35ck0DuIX/isZvIk+ubY=
|
||||
github.com/ulikunitz/xz v0.5.15/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
|
||||
github.com/vishvananda/netns v0.0.5 h1:DfiHV+j8bA32MFM7bfEunvT8IAqQ/NzSJHtcmW5zdEY=
|
||||
github.com/vishvananda/netns v0.0.5/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM=
|
||||
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
|
||||
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
|
||||
github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
|
||||
github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
|
||||
github.com/xlab/treeprint v1.2.0 h1:HzHnuAF1plUN2zGlAFHbSQP2qJ0ZAD3XF5XD7OesXRQ=
|
||||
@@ -463,6 +478,8 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
|
||||
github.com/yiya1989/sshkrb5 v0.0.1 h1:oAKpXLBeZExjYDV9eTL2KypmXU+5oPrJtqDuR5JbWic=
|
||||
github.com/yiya1989/sshkrb5 v0.0.1/go.mod h1:r1bHcNa22x8rrg/Kj/cDF5bX6UtZrHiQgRFvj3Mnpl8=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
|
||||
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
|
||||
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
|
||||
|
||||
@@ -28,12 +28,15 @@ import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
|
||||
consts "github.com/bishopfox/sliver/client/constants"
|
||||
"github.com/bishopfox/sliver/protobuf/clientpb"
|
||||
"github.com/bishopfox/sliver/server/certs"
|
||||
"github.com/bishopfox/sliver/server/console/forms"
|
||||
"github.com/bishopfox/sliver/server/core"
|
||||
"github.com/bishopfox/sliver/server/db"
|
||||
"github.com/bishopfox/sliver/server/db/models"
|
||||
@@ -172,14 +175,54 @@ func NewOperatorConfig(operatorName string, lhost string, lport uint16, permissi
|
||||
return json.Marshal(config)
|
||||
}
|
||||
|
||||
func kickOperatorCmd(cmd *cobra.Command, _ []string) {
|
||||
func kickOperatorCmd(cmd *cobra.Command, args []string) {
|
||||
operator, _ := cmd.Flags().GetString("name")
|
||||
|
||||
if shouldPromptKickOperator(cmd, args) {
|
||||
names, err := operatorNames()
|
||||
if err != nil {
|
||||
fmt.Printf(Warn+"Failed to list operators: %v\n", err)
|
||||
return
|
||||
}
|
||||
if len(names) == 0 {
|
||||
fmt.Printf(Warn + "No operators found.\n")
|
||||
return
|
||||
}
|
||||
if err := forms.SelectOperator("Select operator to kick", names, &operator); err != nil {
|
||||
if errors.Is(err, forms.ErrUserAborted) {
|
||||
return
|
||||
}
|
||||
fmt.Printf(Warn+"Operator selection failed: %v\n", err)
|
||||
return
|
||||
}
|
||||
if err := cmd.Flags().Set("name", operator); err != nil {
|
||||
fmt.Printf(Warn+"Failed to set operator name: %v\n", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
operator = strings.TrimSpace(operator)
|
||||
if operator == "" {
|
||||
fmt.Printf(Warn + "Operator name required (use --name).\n")
|
||||
return
|
||||
}
|
||||
|
||||
exists, err := operatorExists(operator)
|
||||
if err != nil {
|
||||
fmt.Printf(Warn+"Failed to lookup operator %s: %v\n", operator, err)
|
||||
return
|
||||
}
|
||||
if !exists {
|
||||
fmt.Printf(Warn+"Operator %s does not exist.\n", operator)
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf(Info+"Removing auth token(s) for %s, please wait ... \n", operator)
|
||||
err := db.Session().Where(&models.Operator{
|
||||
err = db.Session().Where(&models.Operator{
|
||||
Name: operator,
|
||||
}).Delete(&models.Operator{}).Error
|
||||
if err != nil {
|
||||
fmt.Printf(Warn+"Failed to remove operator %s: %v\n", operator, err)
|
||||
return
|
||||
}
|
||||
transport.ClearTokenCache()
|
||||
@@ -192,6 +235,55 @@ func kickOperatorCmd(cmd *cobra.Command, _ []string) {
|
||||
fmt.Printf(Info+"Operator %s has been kicked out.\n", operator)
|
||||
}
|
||||
|
||||
func shouldPromptKickOperator(cmd *cobra.Command, args []string) bool {
|
||||
if len(args) != 0 {
|
||||
return false
|
||||
}
|
||||
return cmd.Flags().NFlag() == 0
|
||||
}
|
||||
|
||||
func operatorNames() ([]string, error) {
|
||||
operators, err := db.OperatorAll()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(operators))
|
||||
seen := make(map[string]struct{}, len(operators))
|
||||
for _, operator := range operators {
|
||||
if operator == nil {
|
||||
continue
|
||||
}
|
||||
name := strings.TrimSpace(operator.Name)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[name]; ok {
|
||||
continue
|
||||
}
|
||||
seen[name] = struct{}{}
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names, nil
|
||||
}
|
||||
|
||||
func operatorExists(name string) (bool, error) {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
return false, nil
|
||||
}
|
||||
err := db.Session().Where(&models.Operator{
|
||||
Name: name,
|
||||
}).First(&models.Operator{}).Error
|
||||
if err == nil {
|
||||
return true, nil
|
||||
}
|
||||
if errors.Is(err, db.ErrRecordNotFound) {
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
|
||||
func startMultiplayerModeCmd(cmd *cobra.Command, _ []string) {
|
||||
lhost, _ := cmd.Flags().GetString("lhost")
|
||||
lport, _ := cmd.Flags().GetUint16("lport")
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
package forms
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"github.com/charmbracelet/huh"
|
||||
)
|
||||
|
||||
const defaultSelectHeight = 10
|
||||
|
||||
// ErrUserAborted exposes the huh user-abort sentinel for callers.
|
||||
var ErrUserAborted = huh.ErrUserAborted
|
||||
|
||||
// SelectOperator prompts for a single operator selection.
|
||||
func SelectOperator(title string, operators []string, value *string) error {
|
||||
if value == nil {
|
||||
return errors.New("operator selection value is required")
|
||||
}
|
||||
if len(operators) == 0 {
|
||||
return errors.New("operator selection options are required")
|
||||
}
|
||||
|
||||
ensureSelectedValue(operators, value)
|
||||
|
||||
field := huh.NewSelect[string]().
|
||||
Title(title).
|
||||
Options(makeStringOptions(operators)...).
|
||||
Height(listHeight(len(operators))).
|
||||
Value(value).
|
||||
Validate(func(val string) error {
|
||||
if strings.TrimSpace(val) == "" {
|
||||
return errors.New("selection required")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
form := huh.NewForm(huh.NewGroup(field))
|
||||
return form.Run()
|
||||
}
|
||||
|
||||
func ensureSelectedValue(options []string, value *string) {
|
||||
if value == nil || len(options) == 0 {
|
||||
return
|
||||
}
|
||||
if strings.TrimSpace(*value) == "" {
|
||||
*value = options[0]
|
||||
return
|
||||
}
|
||||
for _, option := range options {
|
||||
if option == *value {
|
||||
return
|
||||
}
|
||||
}
|
||||
*value = options[0]
|
||||
}
|
||||
|
||||
func makeStringOptions(options []string) []huh.Option[string] {
|
||||
converted := make([]huh.Option[string], 0, len(options))
|
||||
for _, option := range options {
|
||||
converted = append(converted, huh.NewOption(option, option))
|
||||
}
|
||||
return converted
|
||||
}
|
||||
|
||||
func listHeight(count int) int {
|
||||
if count < 1 {
|
||||
return 1
|
||||
}
|
||||
if count > defaultSelectHeight {
|
||||
return defaultSelectHeight
|
||||
}
|
||||
return count
|
||||
}
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
Copyright (c) 2009 The Go Authors. All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are
|
||||
met:
|
||||
|
||||
* Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
* Redistributions in binary form must reproduce the above
|
||||
copyright notice, this list of conditions and the following disclaimer
|
||||
in the documentation and/or other materials provided with the
|
||||
distribution.
|
||||
* Neither the name of Google Inc. nor the names of its
|
||||
contributors may be used to endorse or promote products derived from
|
||||
this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
||||
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
||||
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
||||
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
||||
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
||||
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
||||
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
# generic-list-go [](https://github.com/bahlo/generic-list-go/actions/workflows/ci.yml)
|
||||
|
||||
Go [container/list](https://pkg.go.dev/container/list) but with generics.
|
||||
|
||||
The code is based on `container/list` in `go1.18beta2`.
|
||||
+235
@@ -0,0 +1,235 @@
|
||||
// Copyright 2009 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package list implements a doubly linked list.
|
||||
//
|
||||
// To iterate over a list (where l is a *List):
|
||||
// for e := l.Front(); e != nil; e = e.Next() {
|
||||
// // do something with e.Value
|
||||
// }
|
||||
//
|
||||
package list
|
||||
|
||||
// Element is an element of a linked list.
|
||||
type Element[T any] struct {
|
||||
// Next and previous pointers in the doubly-linked list of elements.
|
||||
// To simplify the implementation, internally a list l is implemented
|
||||
// as a ring, such that &l.root is both the next element of the last
|
||||
// list element (l.Back()) and the previous element of the first list
|
||||
// element (l.Front()).
|
||||
next, prev *Element[T]
|
||||
|
||||
// The list to which this element belongs.
|
||||
list *List[T]
|
||||
|
||||
// The value stored with this element.
|
||||
Value T
|
||||
}
|
||||
|
||||
// Next returns the next list element or nil.
|
||||
func (e *Element[T]) Next() *Element[T] {
|
||||
if p := e.next; e.list != nil && p != &e.list.root {
|
||||
return p
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Prev returns the previous list element or nil.
|
||||
func (e *Element[T]) Prev() *Element[T] {
|
||||
if p := e.prev; e.list != nil && p != &e.list.root {
|
||||
return p
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// List represents a doubly linked list.
|
||||
// The zero value for List is an empty list ready to use.
|
||||
type List[T any] struct {
|
||||
root Element[T] // sentinel list element, only &root, root.prev, and root.next are used
|
||||
len int // current list length excluding (this) sentinel element
|
||||
}
|
||||
|
||||
// Init initializes or clears list l.
|
||||
func (l *List[T]) Init() *List[T] {
|
||||
l.root.next = &l.root
|
||||
l.root.prev = &l.root
|
||||
l.len = 0
|
||||
return l
|
||||
}
|
||||
|
||||
// New returns an initialized list.
|
||||
func New[T any]() *List[T] { return new(List[T]).Init() }
|
||||
|
||||
// Len returns the number of elements of list l.
|
||||
// The complexity is O(1).
|
||||
func (l *List[T]) Len() int { return l.len }
|
||||
|
||||
// Front returns the first element of list l or nil if the list is empty.
|
||||
func (l *List[T]) Front() *Element[T] {
|
||||
if l.len == 0 {
|
||||
return nil
|
||||
}
|
||||
return l.root.next
|
||||
}
|
||||
|
||||
// Back returns the last element of list l or nil if the list is empty.
|
||||
func (l *List[T]) Back() *Element[T] {
|
||||
if l.len == 0 {
|
||||
return nil
|
||||
}
|
||||
return l.root.prev
|
||||
}
|
||||
|
||||
// lazyInit lazily initializes a zero List value.
|
||||
func (l *List[T]) lazyInit() {
|
||||
if l.root.next == nil {
|
||||
l.Init()
|
||||
}
|
||||
}
|
||||
|
||||
// insert inserts e after at, increments l.len, and returns e.
|
||||
func (l *List[T]) insert(e, at *Element[T]) *Element[T] {
|
||||
e.prev = at
|
||||
e.next = at.next
|
||||
e.prev.next = e
|
||||
e.next.prev = e
|
||||
e.list = l
|
||||
l.len++
|
||||
return e
|
||||
}
|
||||
|
||||
// insertValue is a convenience wrapper for insert(&Element{Value: v}, at).
|
||||
func (l *List[T]) insertValue(v T, at *Element[T]) *Element[T] {
|
||||
return l.insert(&Element[T]{Value: v}, at)
|
||||
}
|
||||
|
||||
// remove removes e from its list, decrements l.len
|
||||
func (l *List[T]) remove(e *Element[T]) {
|
||||
e.prev.next = e.next
|
||||
e.next.prev = e.prev
|
||||
e.next = nil // avoid memory leaks
|
||||
e.prev = nil // avoid memory leaks
|
||||
e.list = nil
|
||||
l.len--
|
||||
}
|
||||
|
||||
// move moves e to next to at.
|
||||
func (l *List[T]) move(e, at *Element[T]) {
|
||||
if e == at {
|
||||
return
|
||||
}
|
||||
e.prev.next = e.next
|
||||
e.next.prev = e.prev
|
||||
|
||||
e.prev = at
|
||||
e.next = at.next
|
||||
e.prev.next = e
|
||||
e.next.prev = e
|
||||
}
|
||||
|
||||
// Remove removes e from l if e is an element of list l.
|
||||
// It returns the element value e.Value.
|
||||
// The element must not be nil.
|
||||
func (l *List[T]) Remove(e *Element[T]) T {
|
||||
if e.list == l {
|
||||
// if e.list == l, l must have been initialized when e was inserted
|
||||
// in l or l == nil (e is a zero Element) and l.remove will crash
|
||||
l.remove(e)
|
||||
}
|
||||
return e.Value
|
||||
}
|
||||
|
||||
// PushFront inserts a new element e with value v at the front of list l and returns e.
|
||||
func (l *List[T]) PushFront(v T) *Element[T] {
|
||||
l.lazyInit()
|
||||
return l.insertValue(v, &l.root)
|
||||
}
|
||||
|
||||
// PushBack inserts a new element e with value v at the back of list l and returns e.
|
||||
func (l *List[T]) PushBack(v T) *Element[T] {
|
||||
l.lazyInit()
|
||||
return l.insertValue(v, l.root.prev)
|
||||
}
|
||||
|
||||
// InsertBefore inserts a new element e with value v immediately before mark and returns e.
|
||||
// If mark is not an element of l, the list is not modified.
|
||||
// The mark must not be nil.
|
||||
func (l *List[T]) InsertBefore(v T, mark *Element[T]) *Element[T] {
|
||||
if mark.list != l {
|
||||
return nil
|
||||
}
|
||||
// see comment in List.Remove about initialization of l
|
||||
return l.insertValue(v, mark.prev)
|
||||
}
|
||||
|
||||
// InsertAfter inserts a new element e with value v immediately after mark and returns e.
|
||||
// If mark is not an element of l, the list is not modified.
|
||||
// The mark must not be nil.
|
||||
func (l *List[T]) InsertAfter(v T, mark *Element[T]) *Element[T] {
|
||||
if mark.list != l {
|
||||
return nil
|
||||
}
|
||||
// see comment in List.Remove about initialization of l
|
||||
return l.insertValue(v, mark)
|
||||
}
|
||||
|
||||
// MoveToFront moves element e to the front of list l.
|
||||
// If e is not an element of l, the list is not modified.
|
||||
// The element must not be nil.
|
||||
func (l *List[T]) MoveToFront(e *Element[T]) {
|
||||
if e.list != l || l.root.next == e {
|
||||
return
|
||||
}
|
||||
// see comment in List.Remove about initialization of l
|
||||
l.move(e, &l.root)
|
||||
}
|
||||
|
||||
// MoveToBack moves element e to the back of list l.
|
||||
// If e is not an element of l, the list is not modified.
|
||||
// The element must not be nil.
|
||||
func (l *List[T]) MoveToBack(e *Element[T]) {
|
||||
if e.list != l || l.root.prev == e {
|
||||
return
|
||||
}
|
||||
// see comment in List.Remove about initialization of l
|
||||
l.move(e, l.root.prev)
|
||||
}
|
||||
|
||||
// MoveBefore moves element e to its new position before mark.
|
||||
// If e or mark is not an element of l, or e == mark, the list is not modified.
|
||||
// The element and mark must not be nil.
|
||||
func (l *List[T]) MoveBefore(e, mark *Element[T]) {
|
||||
if e.list != l || e == mark || mark.list != l {
|
||||
return
|
||||
}
|
||||
l.move(e, mark.prev)
|
||||
}
|
||||
|
||||
// MoveAfter moves element e to its new position after mark.
|
||||
// If e or mark is not an element of l, or e == mark, the list is not modified.
|
||||
// The element and mark must not be nil.
|
||||
func (l *List[T]) MoveAfter(e, mark *Element[T]) {
|
||||
if e.list != l || e == mark || mark.list != l {
|
||||
return
|
||||
}
|
||||
l.move(e, mark)
|
||||
}
|
||||
|
||||
// PushBackList inserts a copy of another list at the back of list l.
|
||||
// The lists l and other may be the same. They must not be nil.
|
||||
func (l *List[T]) PushBackList(other *List[T]) {
|
||||
l.lazyInit()
|
||||
for i, e := other.Len(), other.Front(); i > 0; i, e = i-1, e.Next() {
|
||||
l.insertValue(e.Value, l.root.prev)
|
||||
}
|
||||
}
|
||||
|
||||
// PushFrontList inserts a copy of another list at the front of list l.
|
||||
// The lists l and other may be the same. They must not be nil.
|
||||
func (l *List[T]) PushFrontList(other *List[T]) {
|
||||
l.lazyInit()
|
||||
for i, e := other.Len(), other.Back(); i > 0; i, e = i-1, e.Prev() {
|
||||
l.insertValue(e.Value, &l.root)
|
||||
}
|
||||
}
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
|
||||
*.test
|
||||
|
||||
*.out
|
||||
|
||||
*.mprof
|
||||
|
||||
.idea
|
||||
|
||||
vendor/github.com/buger/goterm/
|
||||
prof.cpu
|
||||
prof.mem
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
language: go
|
||||
arch:
|
||||
- amd64
|
||||
- ppc64le
|
||||
go:
|
||||
- 1.7.x
|
||||
- 1.8.x
|
||||
- 1.9.x
|
||||
- 1.10.x
|
||||
- 1.11.x
|
||||
script: go test -v ./.
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
FROM golang:1.6
|
||||
|
||||
RUN go get github.com/Jeffail/gabs
|
||||
RUN go get github.com/bitly/go-simplejson
|
||||
RUN go get github.com/pquerna/ffjson
|
||||
RUN go get github.com/antonholmquist/jason
|
||||
RUN go get github.com/mreiferson/go-ujson
|
||||
RUN go get -tags=unsafe -u github.com/ugorji/go/codec
|
||||
RUN go get github.com/mailru/easyjson
|
||||
|
||||
WORKDIR /go/src/github.com/buger/jsonparser
|
||||
ADD . /go/src/github.com/buger/jsonparser
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2016 Leonid Bugaev
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
SOURCE = parser.go
|
||||
CONTAINER = jsonparser
|
||||
SOURCE_PATH = /go/src/github.com/buger/jsonparser
|
||||
BENCHMARK = JsonParser
|
||||
BENCHTIME = 5s
|
||||
TEST = .
|
||||
DRUN = docker run -v `pwd`:$(SOURCE_PATH) -i -t $(CONTAINER)
|
||||
|
||||
build:
|
||||
docker build -t $(CONTAINER) .
|
||||
|
||||
race:
|
||||
$(DRUN) --env GORACE="halt_on_error=1" go test ./. $(ARGS) -v -race -timeout 15s
|
||||
|
||||
bench:
|
||||
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -benchtime $(BENCHTIME) -v
|
||||
|
||||
bench_local:
|
||||
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench . $(ARGS) -benchtime $(BENCHTIME) -v
|
||||
|
||||
profile:
|
||||
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -memprofile mem.mprof -v
|
||||
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -cpuprofile cpu.out -v
|
||||
$(DRUN) go test $(LDFLAGS) -test.benchmem -bench $(BENCHMARK) ./benchmark/ $(ARGS) -c
|
||||
|
||||
test:
|
||||
$(DRUN) go test $(LDFLAGS) ./ -run $(TEST) -timeout 10s $(ARGS) -v
|
||||
|
||||
fmt:
|
||||
$(DRUN) go fmt ./...
|
||||
|
||||
vet:
|
||||
$(DRUN) go vet ./.
|
||||
|
||||
bash:
|
||||
$(DRUN) /bin/bash
|
||||
+365
@@ -0,0 +1,365 @@
|
||||
[](https://goreportcard.com/report/github.com/buger/jsonparser) 
|
||||
# Alternative JSON parser for Go (10x times faster standard library)
|
||||
|
||||
It does not require you to know the structure of the payload (eg. create structs), and allows accessing fields by providing the path to them. It is up to **10 times faster** than standard `encoding/json` package (depending on payload size and usage), **allocates no memory**. See benchmarks below.
|
||||
|
||||
## Rationale
|
||||
Originally I made this for a project that relies on a lot of 3rd party APIs that can be unpredictable and complex.
|
||||
I love simplicity and prefer to avoid external dependecies. `encoding/json` requires you to know exactly your data structures, or if you prefer to use `map[string]interface{}` instead, it will be very slow and hard to manage.
|
||||
I investigated what's on the market and found that most libraries are just wrappers around `encoding/json`, there is few options with own parsers (`ffjson`, `easyjson`), but they still requires you to create data structures.
|
||||
|
||||
|
||||
Goal of this project is to push JSON parser to the performance limits and not sacrifice with compliance and developer user experience.
|
||||
|
||||
## Example
|
||||
For the given JSON our goal is to extract the user's full name, number of github followers and avatar.
|
||||
|
||||
```go
|
||||
import "github.com/buger/jsonparser"
|
||||
|
||||
...
|
||||
|
||||
data := []byte(`{
|
||||
"person": {
|
||||
"name": {
|
||||
"first": "Leonid",
|
||||
"last": "Bugaev",
|
||||
"fullName": "Leonid Bugaev"
|
||||
},
|
||||
"github": {
|
||||
"handle": "buger",
|
||||
"followers": 109
|
||||
},
|
||||
"avatars": [
|
||||
{ "url": "https://avatars1.githubusercontent.com/u/14009?v=3&s=460", "type": "thumbnail" }
|
||||
]
|
||||
},
|
||||
"company": {
|
||||
"name": "Acme"
|
||||
}
|
||||
}`)
|
||||
|
||||
// You can specify key path by providing arguments to Get function
|
||||
jsonparser.Get(data, "person", "name", "fullName")
|
||||
|
||||
// There is `GetInt` and `GetBoolean` helpers if you exactly know key data type
|
||||
jsonparser.GetInt(data, "person", "github", "followers")
|
||||
|
||||
// When you try to get object, it will return you []byte slice pointer to data containing it
|
||||
// In `company` it will be `{"name": "Acme"}`
|
||||
jsonparser.Get(data, "company")
|
||||
|
||||
// If the key doesn't exist it will throw an error
|
||||
var size int64
|
||||
if value, err := jsonparser.GetInt(data, "company", "size"); err == nil {
|
||||
size = value
|
||||
}
|
||||
|
||||
// You can use `ArrayEach` helper to iterate items [item1, item2 .... itemN]
|
||||
jsonparser.ArrayEach(data, func(value []byte, dataType jsonparser.ValueType, offset int, err error) {
|
||||
fmt.Println(jsonparser.Get(value, "url"))
|
||||
}, "person", "avatars")
|
||||
|
||||
// Or use can access fields by index!
|
||||
jsonparser.GetString(data, "person", "avatars", "[0]", "url")
|
||||
|
||||
// You can use `ObjectEach` helper to iterate objects { "key1":object1, "key2":object2, .... "keyN":objectN }
|
||||
jsonparser.ObjectEach(data, func(key []byte, value []byte, dataType jsonparser.ValueType, offset int) error {
|
||||
fmt.Printf("Key: '%s'\n Value: '%s'\n Type: %s\n", string(key), string(value), dataType)
|
||||
return nil
|
||||
}, "person", "name")
|
||||
|
||||
// The most efficient way to extract multiple keys is `EachKey`
|
||||
|
||||
paths := [][]string{
|
||||
[]string{"person", "name", "fullName"},
|
||||
[]string{"person", "avatars", "[0]", "url"},
|
||||
[]string{"company", "url"},
|
||||
}
|
||||
jsonparser.EachKey(data, func(idx int, value []byte, vt jsonparser.ValueType, err error){
|
||||
switch idx {
|
||||
case 0: // []string{"person", "name", "fullName"}
|
||||
...
|
||||
case 1: // []string{"person", "avatars", "[0]", "url"}
|
||||
...
|
||||
case 2: // []string{"company", "url"},
|
||||
...
|
||||
}
|
||||
}, paths...)
|
||||
|
||||
// For more information see docs below
|
||||
```
|
||||
|
||||
## Need to speedup your app?
|
||||
|
||||
I'm available for consulting and can help you push your app performance to the limits. Ping me at: leonsbox@gmail.com.
|
||||
|
||||
## Reference
|
||||
|
||||
Library API is really simple. You just need the `Get` method to perform any operation. The rest is just helpers around it.
|
||||
|
||||
You also can view API at [godoc.org](https://godoc.org/github.com/buger/jsonparser)
|
||||
|
||||
|
||||
### **`Get`**
|
||||
```go
|
||||
func Get(data []byte, keys ...string) (value []byte, dataType jsonparser.ValueType, offset int, err error)
|
||||
```
|
||||
Receives data structure, and key path to extract value from.
|
||||
|
||||
Returns:
|
||||
* `value` - Pointer to original data structure containing key value, or just empty slice if nothing found or error
|
||||
* `dataType` - Can be: `NotExist`, `String`, `Number`, `Object`, `Array`, `Boolean` or `Null`
|
||||
* `offset` - Offset from provided data structure where key value ends. Used mostly internally, for example for `ArrayEach` helper.
|
||||
* `err` - If the key is not found or any other parsing issue, it should return error. If key not found it also sets `dataType` to `NotExist`
|
||||
|
||||
Accepts multiple keys to specify path to JSON value (in case of quering nested structures).
|
||||
If no keys are provided it will try to extract the closest JSON value (simple ones or object/array), useful for reading streams or arrays, see `ArrayEach` implementation.
|
||||
|
||||
Note that keys can be an array indexes: `jsonparser.GetInt("person", "avatars", "[0]", "url")`, pretty cool, yeah?
|
||||
|
||||
### **`GetString`**
|
||||
```go
|
||||
func GetString(data []byte, keys ...string) (val string, err error)
|
||||
```
|
||||
Returns strings properly handing escaped and unicode characters. Note that this will cause additional memory allocations.
|
||||
|
||||
### **`GetUnsafeString`**
|
||||
If you need string in your app, and ready to sacrifice with support of escaped symbols in favor of speed. It returns string mapped to existing byte slice memory, without any allocations:
|
||||
```go
|
||||
s, _, := jsonparser.GetUnsafeString(data, "person", "name", "title")
|
||||
switch s {
|
||||
case 'CEO':
|
||||
...
|
||||
case 'Engineer'
|
||||
...
|
||||
...
|
||||
}
|
||||
```
|
||||
Note that `unsafe` here means that your string will exist until GC will free underlying byte slice, for most of cases it means that you can use this string only in current context, and should not pass it anywhere externally: through channels or any other way.
|
||||
|
||||
|
||||
### **`GetBoolean`**, **`GetInt`** and **`GetFloat`**
|
||||
```go
|
||||
func GetBoolean(data []byte, keys ...string) (val bool, err error)
|
||||
|
||||
func GetFloat(data []byte, keys ...string) (val float64, err error)
|
||||
|
||||
func GetInt(data []byte, keys ...string) (val int64, err error)
|
||||
```
|
||||
If you know the key type, you can use the helpers above.
|
||||
If key data type do not match, it will return error.
|
||||
|
||||
### **`ArrayEach`**
|
||||
```go
|
||||
func ArrayEach(data []byte, cb func(value []byte, dataType jsonparser.ValueType, offset int, err error), keys ...string)
|
||||
```
|
||||
Needed for iterating arrays, accepts a callback function with the same return arguments as `Get`.
|
||||
|
||||
### **`ObjectEach`**
|
||||
```go
|
||||
func ObjectEach(data []byte, callback func(key []byte, value []byte, dataType ValueType, offset int) error, keys ...string) (err error)
|
||||
```
|
||||
Needed for iterating object, accepts a callback function. Example:
|
||||
```go
|
||||
var handler func([]byte, []byte, jsonparser.ValueType, int) error
|
||||
handler = func(key []byte, value []byte, dataType jsonparser.ValueType, offset int) error {
|
||||
//do stuff here
|
||||
}
|
||||
jsonparser.ObjectEach(myJson, handler)
|
||||
```
|
||||
|
||||
|
||||
### **`EachKey`**
|
||||
```go
|
||||
func EachKey(data []byte, cb func(idx int, value []byte, dataType jsonparser.ValueType, err error), paths ...[]string)
|
||||
```
|
||||
When you need to read multiple keys, and you do not afraid of low-level API `EachKey` is your friend. It read payload only single time, and calls callback function once path is found. For example when you call multiple times `Get`, it has to process payload multiple times, each time you call it. Depending on payload `EachKey` can be multiple times faster than `Get`. Path can use nested keys as well!
|
||||
|
||||
```go
|
||||
paths := [][]string{
|
||||
[]string{"uuid"},
|
||||
[]string{"tz"},
|
||||
[]string{"ua"},
|
||||
[]string{"st"},
|
||||
}
|
||||
var data SmallPayload
|
||||
|
||||
jsonparser.EachKey(smallFixture, func(idx int, value []byte, vt jsonparser.ValueType, err error){
|
||||
switch idx {
|
||||
case 0:
|
||||
data.Uuid, _ = value
|
||||
case 1:
|
||||
v, _ := jsonparser.ParseInt(value)
|
||||
data.Tz = int(v)
|
||||
case 2:
|
||||
data.Ua, _ = value
|
||||
case 3:
|
||||
v, _ := jsonparser.ParseInt(value)
|
||||
data.St = int(v)
|
||||
}
|
||||
}, paths...)
|
||||
```
|
||||
|
||||
### **`Set`**
|
||||
```go
|
||||
func Set(data []byte, setValue []byte, keys ...string) (value []byte, err error)
|
||||
```
|
||||
Receives existing data structure, key path to set, and value to set at that key. *This functionality is experimental.*
|
||||
|
||||
Returns:
|
||||
* `value` - Pointer to original data structure with updated or added key value.
|
||||
* `err` - If any parsing issue, it should return error.
|
||||
|
||||
Accepts multiple keys to specify path to JSON value (in case of updating or creating nested structures).
|
||||
|
||||
Note that keys can be an array indexes: `jsonparser.Set(data, []byte("http://github.com"), "person", "avatars", "[0]", "url")`
|
||||
|
||||
### **`Delete`**
|
||||
```go
|
||||
func Delete(data []byte, keys ...string) value []byte
|
||||
```
|
||||
Receives existing data structure, and key path to delete. *This functionality is experimental.*
|
||||
|
||||
Returns:
|
||||
* `value` - Pointer to original data structure with key path deleted if it can be found. If there is no key path, then the whole data structure is deleted.
|
||||
|
||||
Accepts multiple keys to specify path to JSON value (in case of updating or creating nested structures).
|
||||
|
||||
Note that keys can be an array indexes: `jsonparser.Delete(data, "person", "avatars", "[0]", "url")`
|
||||
|
||||
|
||||
## What makes it so fast?
|
||||
* It does not rely on `encoding/json`, `reflection` or `interface{}`, the only real package dependency is `bytes`.
|
||||
* Operates with JSON payload on byte level, providing you pointers to the original data structure: no memory allocation.
|
||||
* No automatic type conversions, by default everything is a []byte, but it provides you value type, so you can convert by yourself (there is few helpers included).
|
||||
* Does not parse full record, only keys you specified
|
||||
|
||||
|
||||
## Benchmarks
|
||||
|
||||
There are 3 benchmark types, trying to simulate real-life usage for small, medium and large JSON payloads.
|
||||
For each metric, the lower value is better. Time/op is in nanoseconds. Values better than standard encoding/json marked as bold text.
|
||||
Benchmarks run on standard Linode 1024 box.
|
||||
|
||||
Compared libraries:
|
||||
* https://golang.org/pkg/encoding/json
|
||||
* https://github.com/Jeffail/gabs
|
||||
* https://github.com/a8m/djson
|
||||
* https://github.com/bitly/go-simplejson
|
||||
* https://github.com/antonholmquist/jason
|
||||
* https://github.com/mreiferson/go-ujson
|
||||
* https://github.com/ugorji/go/codec
|
||||
* https://github.com/pquerna/ffjson
|
||||
* https://github.com/mailru/easyjson
|
||||
* https://github.com/buger/jsonparser
|
||||
|
||||
#### TLDR
|
||||
If you want to skip next sections we have 2 winner: `jsonparser` and `easyjson`.
|
||||
`jsonparser` is up to 10 times faster than standard `encoding/json` package (depending on payload size and usage), and almost infinitely (literally) better in memory consumption because it operates with data on byte level, and provide direct slice pointers.
|
||||
`easyjson` wins in CPU in medium tests and frankly i'm impressed with this package: it is remarkable results considering that it is almost drop-in replacement for `encoding/json` (require some code generation).
|
||||
|
||||
It's hard to fully compare `jsonparser` and `easyjson` (or `ffson`), they a true parsers and fully process record, unlike `jsonparser` which parse only keys you specified.
|
||||
|
||||
If you searching for replacement of `encoding/json` while keeping structs, `easyjson` is an amazing choice. If you want to process dynamic JSON, have memory constrains, or more control over your data you should try `jsonparser`.
|
||||
|
||||
`jsonparser` performance heavily depends on usage, and it works best when you do not need to process full record, only some keys. The more calls you need to make, the slower it will be, in contrast `easyjson` (or `ffjson`, `encoding/json`) parser record only 1 time, and then you can make as many calls as you want.
|
||||
|
||||
With great power comes great responsibility! :)
|
||||
|
||||
|
||||
#### Small payload
|
||||
|
||||
Each test processes 190 bytes of http log as a JSON record.
|
||||
It should read multiple fields.
|
||||
https://github.com/buger/jsonparser/blob/master/benchmark/benchmark_small_payload_test.go
|
||||
|
||||
Library | time/op | bytes/op | allocs/op
|
||||
------ | ------- | -------- | -------
|
||||
encoding/json struct | 7879 | 880 | 18
|
||||
encoding/json interface{} | 8946 | 1521 | 38
|
||||
Jeffail/gabs | 10053 | 1649 | 46
|
||||
bitly/go-simplejson | 10128 | 2241 | 36
|
||||
antonholmquist/jason | 27152 | 7237 | 101
|
||||
github.com/ugorji/go/codec | 8806 | 2176 | 31
|
||||
mreiferson/go-ujson | **7008** | **1409** | 37
|
||||
a8m/djson | 3862 | 1249 | 30
|
||||
pquerna/ffjson | **3769** | **624** | **15**
|
||||
mailru/easyjson | **2002** | **192** | **9**
|
||||
buger/jsonparser | **1367** | **0** | **0**
|
||||
buger/jsonparser (EachKey API) | **809** | **0** | **0**
|
||||
|
||||
Winners are ffjson, easyjson and jsonparser, where jsonparser is up to 9.8x faster than encoding/json and 4.6x faster than ffjson, and slightly faster than easyjson.
|
||||
If you look at memory allocation, jsonparser has no rivals, as it makes no data copy and operates with raw []byte structures and pointers to it.
|
||||
|
||||
#### Medium payload
|
||||
|
||||
Each test processes a 2.4kb JSON record (based on Clearbit API).
|
||||
It should read multiple nested fields and 1 array.
|
||||
|
||||
https://github.com/buger/jsonparser/blob/master/benchmark/benchmark_medium_payload_test.go
|
||||
|
||||
| Library | time/op | bytes/op | allocs/op |
|
||||
| ------- | ------- | -------- | --------- |
|
||||
| encoding/json struct | 57749 | 1336 | 29 |
|
||||
| encoding/json interface{} | 79297 | 10627 | 215 |
|
||||
| Jeffail/gabs | 83807 | 11202 | 235 |
|
||||
| bitly/go-simplejson | 88187 | 17187 | 220 |
|
||||
| antonholmquist/jason | 94099 | 19013 | 247 |
|
||||
| github.com/ugorji/go/codec | 114719 | 6712 | 152 |
|
||||
| mreiferson/go-ujson | **56972** | 11547 | 270 |
|
||||
| a8m/djson | 28525 | 10196 | 198 |
|
||||
| pquerna/ffjson | **20298** | **856** | **20** |
|
||||
| mailru/easyjson | **10512** | **336** | **12** |
|
||||
| buger/jsonparser | **15955** | **0** | **0** |
|
||||
| buger/jsonparser (EachKey API) | **8916** | **0** | **0** |
|
||||
|
||||
The difference between ffjson and jsonparser in CPU usage is smaller, while the memory consumption difference is growing. On the other hand `easyjson` shows remarkable performance for medium payload.
|
||||
|
||||
`gabs`, `go-simplejson` and `jason` are based on encoding/json and map[string]interface{} and actually only helpers for unstructured JSON, their performance correlate with `encoding/json interface{}`, and they will skip next round.
|
||||
`go-ujson` while have its own parser, shows same performance as `encoding/json`, also skips next round. Same situation with `ugorji/go/codec`, but it showed unexpectedly bad performance for complex payloads.
|
||||
|
||||
|
||||
#### Large payload
|
||||
|
||||
Each test processes a 24kb JSON record (based on Discourse API)
|
||||
It should read 2 arrays, and for each item in array get a few fields.
|
||||
Basically it means processing a full JSON file.
|
||||
|
||||
https://github.com/buger/jsonparser/blob/master/benchmark/benchmark_large_payload_test.go
|
||||
|
||||
| Library | time/op | bytes/op | allocs/op |
|
||||
| --- | --- | --- | --- |
|
||||
| encoding/json struct | 748336 | 8272 | 307 |
|
||||
| encoding/json interface{} | 1224271 | 215425 | 3395 |
|
||||
| a8m/djson | 510082 | 213682 | 2845 |
|
||||
| pquerna/ffjson | **312271** | **7792** | **298** |
|
||||
| mailru/easyjson | **154186** | **6992** | **288** |
|
||||
| buger/jsonparser | **85308** | **0** | **0** |
|
||||
|
||||
`jsonparser` now is a winner, but do not forget that it is way more lightweight parser than `ffson` or `easyjson`, and they have to parser all the data, while `jsonparser` parse only what you need. All `ffjson`, `easysjon` and `jsonparser` have their own parsing code, and does not depend on `encoding/json` or `interface{}`, thats one of the reasons why they are so fast. `easyjson` also use a bit of `unsafe` package to reduce memory consuption (in theory it can lead to some unexpected GC issue, but i did not tested enough)
|
||||
|
||||
Also last benchmark did not included `EachKey` test, because in this particular case we need to read lot of Array values, and using `ArrayEach` is more efficient.
|
||||
|
||||
## Questions and support
|
||||
|
||||
All bug-reports and suggestions should go though Github Issues.
|
||||
|
||||
## Contributing
|
||||
|
||||
1. Fork it
|
||||
2. Create your feature branch (git checkout -b my-new-feature)
|
||||
3. Commit your changes (git commit -am 'Added some feature')
|
||||
4. Push to the branch (git push origin my-new-feature)
|
||||
5. Create new Pull Request
|
||||
|
||||
## Development
|
||||
|
||||
All my development happens using Docker, and repo include some Make tasks to simplify development.
|
||||
|
||||
* `make build` - builds docker image, usually can be called only once
|
||||
* `make test` - run tests
|
||||
* `make fmt` - run go fmt
|
||||
* `make bench` - run benchmarks (if you need to run only single benchmark modify `BENCHMARK` variable in make file)
|
||||
* `make profile` - runs benchmark and generate 3 files- `cpu.out`, `mem.mprof` and `benchmark.test` binary, which can be used for `go tool pprof`
|
||||
* `make bash` - enter container (i use it for running `go tool pprof` above)
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
package jsonparser
|
||||
|
||||
import (
|
||||
bio "bytes"
|
||||
)
|
||||
|
||||
// minInt64 '-9223372036854775808' is the smallest representable number in int64
|
||||
const minInt64 = `9223372036854775808`
|
||||
|
||||
// About 2x faster then strconv.ParseInt because it only supports base 10, which is enough for JSON
|
||||
func parseInt(bytes []byte) (v int64, ok bool, overflow bool) {
|
||||
if len(bytes) == 0 {
|
||||
return 0, false, false
|
||||
}
|
||||
|
||||
var neg bool = false
|
||||
if bytes[0] == '-' {
|
||||
neg = true
|
||||
bytes = bytes[1:]
|
||||
}
|
||||
|
||||
var b int64 = 0
|
||||
for _, c := range bytes {
|
||||
if c >= '0' && c <= '9' {
|
||||
b = (10 * v) + int64(c-'0')
|
||||
} else {
|
||||
return 0, false, false
|
||||
}
|
||||
if overflow = (b < v); overflow {
|
||||
break
|
||||
}
|
||||
v = b
|
||||
}
|
||||
|
||||
if overflow {
|
||||
if neg && bio.Equal(bytes, []byte(minInt64)) {
|
||||
return b, true, false
|
||||
}
|
||||
return 0, false, true
|
||||
}
|
||||
|
||||
if neg {
|
||||
return -v, true, false
|
||||
} else {
|
||||
return v, true, false
|
||||
}
|
||||
}
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
// +build appengine appenginevm
|
||||
|
||||
package jsonparser
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// See fastbytes_unsafe.go for explanation on why *[]byte is used (signatures must be consistent with those in that file)
|
||||
|
||||
func equalStr(b *[]byte, s string) bool {
|
||||
return string(*b) == s
|
||||
}
|
||||
|
||||
func parseFloat(b *[]byte) (float64, error) {
|
||||
return strconv.ParseFloat(string(*b), 64)
|
||||
}
|
||||
|
||||
func bytesToString(b *[]byte) string {
|
||||
return string(*b)
|
||||
}
|
||||
|
||||
func StringToBytes(s string) []byte {
|
||||
return []byte(s)
|
||||
}
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
// +build !appengine,!appenginevm
|
||||
|
||||
package jsonparser
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strconv"
|
||||
"unsafe"
|
||||
"runtime"
|
||||
)
|
||||
|
||||
//
|
||||
// The reason for using *[]byte rather than []byte in parameters is an optimization. As of Go 1.6,
|
||||
// the compiler cannot perfectly inline the function when using a non-pointer slice. That is,
|
||||
// the non-pointer []byte parameter version is slower than if its function body is manually
|
||||
// inlined, whereas the pointer []byte version is equally fast to the manually inlined
|
||||
// version. Instruction count in assembly taken from "go tool compile" confirms this difference.
|
||||
//
|
||||
// TODO: Remove hack after Go 1.7 release
|
||||
//
|
||||
func equalStr(b *[]byte, s string) bool {
|
||||
return *(*string)(unsafe.Pointer(b)) == s
|
||||
}
|
||||
|
||||
func parseFloat(b *[]byte) (float64, error) {
|
||||
return strconv.ParseFloat(*(*string)(unsafe.Pointer(b)), 64)
|
||||
}
|
||||
|
||||
// A hack until issue golang/go#2632 is fixed.
|
||||
// See: https://github.com/golang/go/issues/2632
|
||||
func bytesToString(b *[]byte) string {
|
||||
return *(*string)(unsafe.Pointer(b))
|
||||
}
|
||||
|
||||
func StringToBytes(s string) []byte {
|
||||
b := make([]byte, 0, 0)
|
||||
bh := (*reflect.SliceHeader)(unsafe.Pointer(&b))
|
||||
sh := (*reflect.StringHeader)(unsafe.Pointer(&s))
|
||||
bh.Data = sh.Data
|
||||
bh.Cap = sh.Len
|
||||
bh.Len = sh.Len
|
||||
runtime.KeepAlive(s)
|
||||
return b
|
||||
}
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
package jsonparser
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// JSON Unicode stuff: see https://tools.ietf.org/html/rfc7159#section-7
|
||||
|
||||
const supplementalPlanesOffset = 0x10000
|
||||
const highSurrogateOffset = 0xD800
|
||||
const lowSurrogateOffset = 0xDC00
|
||||
|
||||
const basicMultilingualPlaneReservedOffset = 0xDFFF
|
||||
const basicMultilingualPlaneOffset = 0xFFFF
|
||||
|
||||
func combineUTF16Surrogates(high, low rune) rune {
|
||||
return supplementalPlanesOffset + (high-highSurrogateOffset)<<10 + (low - lowSurrogateOffset)
|
||||
}
|
||||
|
||||
const badHex = -1
|
||||
|
||||
func h2I(c byte) int {
|
||||
switch {
|
||||
case c >= '0' && c <= '9':
|
||||
return int(c - '0')
|
||||
case c >= 'A' && c <= 'F':
|
||||
return int(c - 'A' + 10)
|
||||
case c >= 'a' && c <= 'f':
|
||||
return int(c - 'a' + 10)
|
||||
}
|
||||
return badHex
|
||||
}
|
||||
|
||||
// decodeSingleUnicodeEscape decodes a single \uXXXX escape sequence. The prefix \u is assumed to be present and
|
||||
// is not checked.
|
||||
// In JSON, these escapes can either come alone or as part of "UTF16 surrogate pairs" that must be handled together.
|
||||
// This function only handles one; decodeUnicodeEscape handles this more complex case.
|
||||
func decodeSingleUnicodeEscape(in []byte) (rune, bool) {
|
||||
// We need at least 6 characters total
|
||||
if len(in) < 6 {
|
||||
return utf8.RuneError, false
|
||||
}
|
||||
|
||||
// Convert hex to decimal
|
||||
h1, h2, h3, h4 := h2I(in[2]), h2I(in[3]), h2I(in[4]), h2I(in[5])
|
||||
if h1 == badHex || h2 == badHex || h3 == badHex || h4 == badHex {
|
||||
return utf8.RuneError, false
|
||||
}
|
||||
|
||||
// Compose the hex digits
|
||||
return rune(h1<<12 + h2<<8 + h3<<4 + h4), true
|
||||
}
|
||||
|
||||
// isUTF16EncodedRune checks if a rune is in the range for non-BMP characters,
|
||||
// which is used to describe UTF16 chars.
|
||||
// Source: https://en.wikipedia.org/wiki/Plane_(Unicode)#Basic_Multilingual_Plane
|
||||
func isUTF16EncodedRune(r rune) bool {
|
||||
return highSurrogateOffset <= r && r <= basicMultilingualPlaneReservedOffset
|
||||
}
|
||||
|
||||
func decodeUnicodeEscape(in []byte) (rune, int) {
|
||||
if r, ok := decodeSingleUnicodeEscape(in); !ok {
|
||||
// Invalid Unicode escape
|
||||
return utf8.RuneError, -1
|
||||
} else if r <= basicMultilingualPlaneOffset && !isUTF16EncodedRune(r) {
|
||||
// Valid Unicode escape in Basic Multilingual Plane
|
||||
return r, 6
|
||||
} else if r2, ok := decodeSingleUnicodeEscape(in[6:]); !ok { // Note: previous decodeSingleUnicodeEscape success guarantees at least 6 bytes remain
|
||||
// UTF16 "high surrogate" without manditory valid following Unicode escape for the "low surrogate"
|
||||
return utf8.RuneError, -1
|
||||
} else if r2 < lowSurrogateOffset {
|
||||
// Invalid UTF16 "low surrogate"
|
||||
return utf8.RuneError, -1
|
||||
} else {
|
||||
// Valid UTF16 surrogate pair
|
||||
return combineUTF16Surrogates(r, r2), 12
|
||||
}
|
||||
}
|
||||
|
||||
// backslashCharEscapeTable: when '\X' is found for some byte X, it is to be replaced with backslashCharEscapeTable[X]
|
||||
var backslashCharEscapeTable = [...]byte{
|
||||
'"': '"',
|
||||
'\\': '\\',
|
||||
'/': '/',
|
||||
'b': '\b',
|
||||
'f': '\f',
|
||||
'n': '\n',
|
||||
'r': '\r',
|
||||
't': '\t',
|
||||
}
|
||||
|
||||
// unescapeToUTF8 unescapes the single escape sequence starting at 'in' into 'out' and returns
|
||||
// how many characters were consumed from 'in' and emitted into 'out'.
|
||||
// If a valid escape sequence does not appear as a prefix of 'in', (-1, -1) to signal the error.
|
||||
func unescapeToUTF8(in, out []byte) (inLen int, outLen int) {
|
||||
if len(in) < 2 || in[0] != '\\' {
|
||||
// Invalid escape due to insufficient characters for any escape or no initial backslash
|
||||
return -1, -1
|
||||
}
|
||||
|
||||
// https://tools.ietf.org/html/rfc7159#section-7
|
||||
switch e := in[1]; e {
|
||||
case '"', '\\', '/', 'b', 'f', 'n', 'r', 't':
|
||||
// Valid basic 2-character escapes (use lookup table)
|
||||
out[0] = backslashCharEscapeTable[e]
|
||||
return 2, 1
|
||||
case 'u':
|
||||
// Unicode escape
|
||||
if r, inLen := decodeUnicodeEscape(in); inLen == -1 {
|
||||
// Invalid Unicode escape
|
||||
return -1, -1
|
||||
} else {
|
||||
// Valid Unicode escape; re-encode as UTF8
|
||||
outLen := utf8.EncodeRune(out, r)
|
||||
return inLen, outLen
|
||||
}
|
||||
}
|
||||
|
||||
return -1, -1
|
||||
}
|
||||
|
||||
// unescape unescapes the string contained in 'in' and returns it as a slice.
|
||||
// If 'in' contains no escaped characters:
|
||||
// Returns 'in'.
|
||||
// Else, if 'out' is of sufficient capacity (guaranteed if cap(out) >= len(in)):
|
||||
// 'out' is used to build the unescaped string and is returned with no extra allocation
|
||||
// Else:
|
||||
// A new slice is allocated and returned.
|
||||
func Unescape(in, out []byte) ([]byte, error) {
|
||||
firstBackslash := bytes.IndexByte(in, '\\')
|
||||
if firstBackslash == -1 {
|
||||
return in, nil
|
||||
}
|
||||
|
||||
// Get a buffer of sufficient size (allocate if needed)
|
||||
if cap(out) < len(in) {
|
||||
out = make([]byte, len(in))
|
||||
} else {
|
||||
out = out[0:len(in)]
|
||||
}
|
||||
|
||||
// Copy the first sequence of unescaped bytes to the output and obtain a buffer pointer (subslice)
|
||||
copy(out, in[:firstBackslash])
|
||||
in = in[firstBackslash:]
|
||||
buf := out[firstBackslash:]
|
||||
|
||||
for len(in) > 0 {
|
||||
// Unescape the next escaped character
|
||||
inLen, bufLen := unescapeToUTF8(in, buf)
|
||||
if inLen == -1 {
|
||||
return nil, MalformedStringEscapeError
|
||||
}
|
||||
|
||||
in = in[inLen:]
|
||||
buf = buf[bufLen:]
|
||||
|
||||
// Copy everything up until the next backslash
|
||||
nextBackslash := bytes.IndexByte(in, '\\')
|
||||
if nextBackslash == -1 {
|
||||
copy(buf, in)
|
||||
buf = buf[len(in):]
|
||||
break
|
||||
} else {
|
||||
copy(buf, in[:nextBackslash])
|
||||
buf = buf[nextBackslash:]
|
||||
in = in[nextBackslash:]
|
||||
}
|
||||
}
|
||||
|
||||
// Trim the out buffer to the amount that was actually emitted
|
||||
return out[:len(out)-len(buf)], nil
|
||||
}
|
||||
+117
@@ -0,0 +1,117 @@
|
||||
package jsonparser
|
||||
|
||||
func FuzzParseString(data []byte) int {
|
||||
r, err := ParseString(data)
|
||||
if err != nil || r == "" {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzEachKey(data []byte) int {
|
||||
paths := [][]string{
|
||||
{"name"},
|
||||
{"order"},
|
||||
{"nested", "a"},
|
||||
{"nested", "b"},
|
||||
{"nested2", "a"},
|
||||
{"nested", "nested3", "b"},
|
||||
{"arr", "[1]", "b"},
|
||||
{"arrInt", "[3]"},
|
||||
{"arrInt", "[5]"},
|
||||
{"nested"},
|
||||
{"arr", "["},
|
||||
{"a\n", "b\n"},
|
||||
}
|
||||
EachKey(data, func(idx int, value []byte, vt ValueType, err error) {}, paths...)
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzDelete(data []byte) int {
|
||||
Delete(data, "test")
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzSet(data []byte) int {
|
||||
_, err := Set(data, []byte(`"new value"`), "test")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzObjectEach(data []byte) int {
|
||||
_ = ObjectEach(data, func(key, value []byte, valueType ValueType, off int) error {
|
||||
return nil
|
||||
})
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzParseFloat(data []byte) int {
|
||||
_, err := ParseFloat(data)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzParseInt(data []byte) int {
|
||||
_, err := ParseInt(data)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzParseBool(data []byte) int {
|
||||
_, err := ParseBoolean(data)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzTokenStart(data []byte) int {
|
||||
_ = tokenStart(data)
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzGetString(data []byte) int {
|
||||
_, err := GetString(data, "test")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzGetFloat(data []byte) int {
|
||||
_, err := GetFloat(data, "test")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzGetInt(data []byte) int {
|
||||
_, err := GetInt(data, "test")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzGetBoolean(data []byte) int {
|
||||
_, err := GetBoolean(data, "test")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func FuzzGetUnsafeString(data []byte) int {
|
||||
_, err := GetUnsafeString(data, "test")
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return 1
|
||||
}
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
#!/bin/bash -eu
|
||||
|
||||
git clone https://github.com/dvyukov/go-fuzz-corpus
|
||||
zip corpus.zip go-fuzz-corpus/json/corpus/*
|
||||
|
||||
cp corpus.zip $OUT/fuzzparsestring_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzParseString fuzzparsestring
|
||||
|
||||
cp corpus.zip $OUT/fuzzeachkey_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzEachKey fuzzeachkey
|
||||
|
||||
cp corpus.zip $OUT/fuzzdelete_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzDelete fuzzdelete
|
||||
|
||||
cp corpus.zip $OUT/fuzzset_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzSet fuzzset
|
||||
|
||||
cp corpus.zip $OUT/fuzzobjecteach_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzObjectEach fuzzobjecteach
|
||||
|
||||
cp corpus.zip $OUT/fuzzparsefloat_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzParseFloat fuzzparsefloat
|
||||
|
||||
cp corpus.zip $OUT/fuzzparseint_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzParseInt fuzzparseint
|
||||
|
||||
cp corpus.zip $OUT/fuzzparsebool_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzParseBool fuzzparsebool
|
||||
|
||||
cp corpus.zip $OUT/fuzztokenstart_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzTokenStart fuzztokenstart
|
||||
|
||||
cp corpus.zip $OUT/fuzzgetstring_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzGetString fuzzgetstring
|
||||
|
||||
cp corpus.zip $OUT/fuzzgetfloat_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzGetFloat fuzzgetfloat
|
||||
|
||||
cp corpus.zip $OUT/fuzzgetint_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzGetInt fuzzgetint
|
||||
|
||||
cp corpus.zip $OUT/fuzzgetboolean_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzGetBoolean fuzzgetboolean
|
||||
|
||||
cp corpus.zip $OUT/fuzzgetunsafestring_seed_corpus.zip
|
||||
compile_go_fuzzer github.com/buger/jsonparser FuzzGetUnsafeString fuzzgetunsafestring
|
||||
|
||||
+1283
File diff suppressed because it is too large.
Load diff
+2
@@ -0,0 +1,2 @@
|
||||
vendor/
|
||||
.idea/
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
run:
|
||||
tests: true
|
||||
max-same-issues: 50
|
||||
|
||||
output:
|
||||
print-issued-lines: false
|
||||
|
||||
linters:
|
||||
enable:
|
||||
- gocyclo
|
||||
- gocritic
|
||||
- goconst
|
||||
- dupl
|
||||
- unconvert
|
||||
- goimports
|
||||
- unused
|
||||
- govet
|
||||
- nakedret
|
||||
- errcheck
|
||||
- revive
|
||||
- ineffassign
|
||||
- goconst
|
||||
- unparam
|
||||
- gofmt
|
||||
|
||||
linters-settings:
|
||||
vet:
|
||||
check-shadowing: true
|
||||
use-installed-packages: true
|
||||
dupl:
|
||||
threshold: 100
|
||||
goconst:
|
||||
min-len: 8
|
||||
min-occurrences: 3
|
||||
gocyclo:
|
||||
min-complexity: 20
|
||||
gocritic:
|
||||
disabled-checks:
|
||||
- ifElseChain
|
||||
gofmt:
|
||||
rewrite-rules:
|
||||
- pattern: "interface{}"
|
||||
replacement: "any"
|
||||
- pattern: "a[b:len(a)]"
|
||||
replacement: "a[b:]"
|
||||
|
||||
issues:
|
||||
max-per-linter: 0
|
||||
max-same: 0
|
||||
exclude-dirs:
|
||||
- resources
|
||||
- old
|
||||
exclude-files:
|
||||
- cmd/protopkg/main.go
|
||||
exclude-use-default: false
|
||||
exclude:
|
||||
# Captured by errcheck.
|
||||
- "^(G104|G204):"
|
||||
# Very commonly not checked.
|
||||
- 'Error return value of .(.*\.Help|.*\.MarkFlagRequired|(os\.)?std(out|err)\..*|.*Close|.*Flush|os\.Remove(All)?|.*Print(f|ln|)|os\.(Un)?Setenv). is not checked'
|
||||
# Weird error only seen on Kochiku...
|
||||
- "internal error: no range for"
|
||||
- 'exported method `.*\.(MarshalJSON|UnmarshalJSON|URN|Payload|GoString|Close|Provides|Requires|ExcludeFromHash|MarshalText|UnmarshalText|Description|Check|Poll|Severity)` should have comment or be unexported'
|
||||
- "composite literal uses unkeyed fields"
|
||||
- 'declaration of "err" shadows declaration'
|
||||
- "by other packages, and that stutters"
|
||||
- "Potential file inclusion via variable"
|
||||
- "at least one file in a package should have a package comment"
|
||||
- "bad syntax for struct tag pair"
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
Copyright (C) 2014 Alec Thomas
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of
|
||||
this software and associated documentation files (the "Software"), to deal in
|
||||
the Software without restriction, including without limitation the rights to
|
||||
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
|
||||
of the Software, and to permit persons to whom the Software is furnished to do
|
||||
so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+374
@@ -0,0 +1,374 @@
|
||||
# Go JSON Schema Reflection
|
||||
|
||||
[](https://github.com/invopop/jsonschema/actions/workflows/lint.yaml)
|
||||
[](https://github.com/invopop/jsonschema/actions/workflows/test.yaml)
|
||||
[](https://goreportcard.com/report/github.com/invopop/jsonschema)
|
||||
[](https://godoc.org/github.com/invopop/jsonschema)
|
||||
[](https://codecov.io/gh/invopop/jsonschema)
|
||||

|
||||
|
||||
This package can be used to generate [JSON Schemas](http://json-schema.org/latest/json-schema-validation.html) from Go types through reflection.
|
||||
|
||||
- Supports arbitrarily complex types, including `interface{}`, maps, slices, etc.
|
||||
- Supports json-schema features such as minLength, maxLength, pattern, format, etc.
|
||||
- Supports simple string and numeric enums.
|
||||
- Supports custom property fields via the `jsonschema_extras` struct tag.
|
||||
|
||||
This repository is a fork of the original [jsonschema](https://github.com/alecthomas/jsonschema) by [@alecthomas](https://github.com/alecthomas). At [Invopop](https://invopop.com) we use jsonschema as a cornerstone in our [GOBL library](https://github.com/invopop/gobl), and wanted to be able to continue building and adding features without taking up Alec's time. There have been a few significant changes that probably mean this version is a not compatible with with Alec's:
|
||||
|
||||
- The original was stuck on the draft-04 version of JSON Schema, we've now moved to the latest JSON Schema Draft 2020-12.
|
||||
- Schema IDs are added automatically from the current Go package's URL in order to be unique, and can be disabled with the `Anonymous` option.
|
||||
- Support for the `FullyQualifyTypeName` option has been removed. If you have conflicts, you should use multiple schema files with different IDs, set the `DoNotReference` option to true to hide definitions completely, or add your own naming strategy using the `Namer` property.
|
||||
- Support for `yaml` tags and related options has been dropped for the sake of simplification. There were a [few inconsistencies](https://github.com/invopop/jsonschema/pull/21) around this that have now been fixed.
|
||||
|
||||
## Versions
|
||||
|
||||
This project is still under v0 scheme, as per Go convention, breaking changes are likely. Please pin go modules to version tags or branches, and reach out if you think something can be improved.
|
||||
|
||||
Go version >= 1.18 is required as generics are now being used.
|
||||
|
||||
## Example
|
||||
|
||||
The following Go type:
|
||||
|
||||
```go
|
||||
type TestUser struct {
|
||||
ID int `json:"id"`
|
||||
Name string `json:"name" jsonschema:"title=the name,description=The name of a friend,example=joe,example=lucy,default=alex"`
|
||||
Friends []int `json:"friends,omitempty" jsonschema_description:"The list of IDs, omitted when empty"`
|
||||
Tags map[string]interface{} `json:"tags,omitempty" jsonschema_extras:"a=b,foo=bar,foo=bar1"`
|
||||
BirthDate time.Time `json:"birth_date,omitempty" jsonschema:"oneof_required=date"`
|
||||
YearOfBirth string `json:"year_of_birth,omitempty" jsonschema:"oneof_required=year"`
|
||||
Metadata interface{} `json:"metadata,omitempty" jsonschema:"oneof_type=string;array"`
|
||||
FavColor string `json:"fav_color,omitempty" jsonschema:"enum=red,enum=green,enum=blue"`
|
||||
}
|
||||
```
|
||||
|
||||
Results in following JSON Schema:
|
||||
|
||||
```go
|
||||
jsonschema.Reflect(&TestUser{})
|
||||
```
|
||||
|
||||
```json
|
||||
{
|
||||
"$schema": "https://json-schema.org/draft/2020-12/schema",
|
||||
"$id": "https://github.com/invopop/jsonschema_test/test-user",
|
||||
"$ref": "#/$defs/TestUser",
|
||||
"$defs": {
|
||||
"TestUser": {
|
||||
"oneOf": [
|
||||
{
|
||||
"required": ["birth_date"],
|
||||
"title": "date"
|
||||
},
|
||||
{
|
||||
"required": ["year_of_birth"],
|
||||
"title": "year"
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"id": {
|
||||
"type": "integer"
|
||||
},
|
||||
"name": {
|
||||
"type": "string",
|
||||
"title": "the name",
|
||||
"description": "The name of a friend",
|
||||
"default": "alex",
|
||||
"examples": ["joe", "lucy"]
|
||||
},
|
||||
"friends": {
|
||||
"items": {
|
||||
"type": "integer"
|
||||
},
|
||||
"type": "array",
|
||||
"description": "The list of IDs, omitted when empty"
|
||||
},
|
||||
"tags": {
|
||||
"type": "object",
|
||||
"a": "b",
|
||||
"foo": ["bar", "bar1"]
|
||||
},
|
||||
"birth_date": {
|
||||
"type": "string",
|
||||
"format": "date-time"
|
||||
},
|
||||
"year_of_birth": {
|
||||
"type": "string"
|
||||
},
|
||||
"metadata": {
|
||||
"oneOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "array"
|
||||
}
|
||||
]
|
||||
},
|
||||
"fav_color": {
|
||||
"type": "string",
|
||||
"enum": ["red", "green", "blue"]
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"type": "object",
|
||||
"required": ["id", "name"]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## YAML
|
||||
|
||||
Support for `yaml` tags has now been removed. If you feel very strongly about this, we've opened a discussion to hear your comments: https://github.com/invopop/jsonschema/discussions/28
|
||||
|
||||
The recommended approach if you need to deal with YAML data is to first convert to JSON. The [invopop/yaml](https://github.com/invopop/yaml) library will make this trivial.
|
||||
|
||||
## Configurable behaviour
|
||||
|
||||
The behaviour of the schema generator can be altered with parameters when a `jsonschema.Reflector`
|
||||
instance is created.
|
||||
|
||||
### ExpandedStruct
|
||||
|
||||
If set to `true`, makes the top level struct not to reference itself in the definitions. But type passed should be a struct type.
|
||||
|
||||
eg.
|
||||
|
||||
```go
|
||||
type GrandfatherType struct {
|
||||
FamilyName string `json:"family_name" jsonschema:"required"`
|
||||
}
|
||||
|
||||
type SomeBaseType struct {
|
||||
SomeBaseProperty int `json:"some_base_property"`
|
||||
// The jsonschema required tag is nonsensical for private and ignored properties.
|
||||
// Their presence here tests that the fields *will not* be required in the output
|
||||
// schema, even if they are tagged required.
|
||||
somePrivateBaseProperty string `json:"i_am_private" jsonschema:"required"`
|
||||
SomeIgnoredBaseProperty string `json:"-" jsonschema:"required"`
|
||||
SomeSchemaIgnoredProperty string `jsonschema:"-,required"`
|
||||
SomeUntaggedBaseProperty bool `jsonschema:"required"`
|
||||
someUnexportedUntaggedBaseProperty bool
|
||||
Grandfather GrandfatherType `json:"grand"`
|
||||
}
|
||||
```
|
||||
|
||||
will output:
|
||||
|
||||
```json
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft/2020-12/schema",
|
||||
"required": ["some_base_property", "grand", "SomeUntaggedBaseProperty"],
|
||||
"properties": {
|
||||
"SomeUntaggedBaseProperty": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"grand": {
|
||||
"$schema": "http://json-schema.org/draft/2020-12/schema",
|
||||
"$ref": "#/definitions/GrandfatherType"
|
||||
},
|
||||
"some_base_property": {
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
"$defs": {
|
||||
"GrandfatherType": {
|
||||
"required": ["family_name"],
|
||||
"properties": {
|
||||
"family_name": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"type": "object"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Using Go Comments
|
||||
|
||||
Writing a good schema with descriptions inside tags can become cumbersome and tedious, especially if you already have some Go comments around your types and field definitions. If you'd like to take advantage of these existing comments, you can use the `AddGoComments(base, path string)` method that forms part of the reflector to parse your go files and automatically generate a dictionary of Go import paths, types, and fields, to individual comments. These will then be used automatically as description fields, and can be overridden with a manual definition if needed.
|
||||
|
||||
Take a simplified example of a User struct which for the sake of simplicity we assume is defined inside this package:
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
// User is used as a base to provide tests for comments.
|
||||
type User struct {
|
||||
// Unique sequential identifier.
|
||||
ID int `json:"id" jsonschema:"required"`
|
||||
// Name of the user
|
||||
Name string `json:"name"`
|
||||
}
|
||||
```
|
||||
|
||||
To get the comments provided into your JSON schema, use a regular `Reflector` and add the go code using an import module URL and path. Fully qualified go module paths cannot be determined reliably by the `go/parser` library, so we need to introduce this manually:
|
||||
|
||||
```go
|
||||
r := new(Reflector)
|
||||
if err := r.AddGoComments("github.com/invopop/jsonschema", "./"); err != nil {
|
||||
// deal with error
|
||||
}
|
||||
s := r.Reflect(&User{})
|
||||
// output
|
||||
```
|
||||
|
||||
Expect the results to be similar to:
|
||||
|
||||
```json
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft/2020-12/schema",
|
||||
"$ref": "#/$defs/User",
|
||||
"$defs": {
|
||||
"User": {
|
||||
"required": ["id"],
|
||||
"properties": {
|
||||
"id": {
|
||||
"type": "integer",
|
||||
"description": "Unique sequential identifier."
|
||||
},
|
||||
"name": {
|
||||
"type": "string",
|
||||
"description": "Name of the user"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"type": "object",
|
||||
"description": "User is used as a base to provide tests for comments."
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Custom Key Naming
|
||||
|
||||
In some situations, the keys actually used to write files are different from Go structs'.
|
||||
|
||||
This is often the case when writing a configuration file to YAML or JSON from a Go struct, or when returning a JSON response for a Web API: APIs typically use snake_case, while Go uses PascalCase.
|
||||
|
||||
You can pass a `func(string) string` function to `Reflector`'s `KeyNamer` option to map Go field names to JSON key names and reflect the aforementioned transformations, without having to specify `json:"..."` on every struct field.
|
||||
|
||||
For example, consider the following struct
|
||||
|
||||
```go
|
||||
type User struct {
|
||||
GivenName string
|
||||
PasswordSalted []byte `json:"salted_password"`
|
||||
}
|
||||
```
|
||||
|
||||
We can transform field names to snake_case in the generated JSON schema:
|
||||
|
||||
```go
|
||||
r := new(jsonschema.Reflector)
|
||||
r.KeyNamer = strcase.SnakeCase // from package github.com/stoewer/go-strcase
|
||||
|
||||
r.Reflect(&User{})
|
||||
```
|
||||
|
||||
Will yield
|
||||
|
||||
```diff
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft/2020-12/schema",
|
||||
"$ref": "#/$defs/User",
|
||||
"$defs": {
|
||||
"User": {
|
||||
"properties": {
|
||||
- "GivenName": {
|
||||
+ "given_name": {
|
||||
"type": "string"
|
||||
},
|
||||
"salted_password": {
|
||||
"type": "string",
|
||||
"contentEncoding": "base64"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false,
|
||||
"type": "object",
|
||||
- "required": ["GivenName", "salted_password"]
|
||||
+ "required": ["given_name", "salted_password"]
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
As you can see, if a field name has a `json:""` tag set, the `key` argument to `KeyNamer` will have the value of that tag.
|
||||
|
||||
### Custom Type Definitions
|
||||
|
||||
Sometimes it can be useful to have custom JSON Marshal and Unmarshal methods in your structs that automatically convert for example a string into an object.
|
||||
|
||||
This library will recognize and attempt to call four different methods that help you adjust schemas to your specific needs:
|
||||
|
||||
- `JSONSchema() *Schema` - will prevent auto-generation of the schema so that you can provide your own definition.
|
||||
- `JSONSchemaExtend(schema *jsonschema.Schema)` - will be called _after_ the schema has been generated, allowing you to add or manipulate the fields easily.
|
||||
- `JSONSchemaAlias() any` - is called when reflecting the type of object and allows for an alternative to be used instead.
|
||||
- `JSONSchemaProperty(prop string) any` - will be called for every property inside a struct giving you the chance to provide an alternative object to convert into a schema.
|
||||
|
||||
Note that all of these methods **must** be defined on a non-pointer object for them to be called.
|
||||
|
||||
Take the following simplified example of a `CompactDate` that only includes the Year and Month:
|
||||
|
||||
```go
|
||||
type CompactDate struct {
|
||||
Year int
|
||||
Month int
|
||||
}
|
||||
|
||||
func (d *CompactDate) UnmarshalJSON(data []byte) error {
|
||||
if len(data) != 9 {
|
||||
return errors.New("invalid compact date length")
|
||||
}
|
||||
var err error
|
||||
d.Year, err = strconv.Atoi(string(data[1:5]))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
d.Month, err = strconv.Atoi(string(data[7:8]))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *CompactDate) MarshalJSON() ([]byte, error) {
|
||||
buf := new(bytes.Buffer)
|
||||
buf.WriteByte('"')
|
||||
buf.WriteString(fmt.Sprintf("%d-%02d", d.Year, d.Month))
|
||||
buf.WriteByte('"')
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
func (CompactDate) JSONSchema() *Schema {
|
||||
return &Schema{
|
||||
Type: "string",
|
||||
Title: "Compact Date",
|
||||
Description: "Short date that only includes year and month",
|
||||
Pattern: "^[0-9]{4}-[0-1][0-9]$",
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
The resulting schema generated for this struct would look like:
|
||||
|
||||
```json
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft/2020-12/schema",
|
||||
"$ref": "#/$defs/CompactDate",
|
||||
"$defs": {
|
||||
"CompactDate": {
|
||||
"pattern": "^[0-9]{4}-[0-1][0-9]$",
|
||||
"type": "string",
|
||||
"title": "Compact Date",
|
||||
"description": "Short date that only includes year and month"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
+76
@@ -0,0 +1,76 @@
|
||||
package jsonschema
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ID represents a Schema ID type which should always be a URI.
|
||||
// See draft-bhutton-json-schema-00 section 8.2.1
|
||||
type ID string
|
||||
|
||||
// EmptyID is used to explicitly define an ID with no value.
|
||||
const EmptyID ID = ""
|
||||
|
||||
// Validate is used to check if the ID looks like a proper schema.
|
||||
// This is done by parsing the ID as a URL and checking it has all the
|
||||
// relevant parts.
|
||||
func (id ID) Validate() error {
|
||||
u, err := url.Parse(id.String())
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid URL: %w", err)
|
||||
}
|
||||
if u.Hostname() == "" {
|
||||
return errors.New("missing hostname")
|
||||
}
|
||||
if !strings.Contains(u.Hostname(), ".") {
|
||||
return errors.New("hostname does not look valid")
|
||||
}
|
||||
if u.Path == "" {
|
||||
return errors.New("path is expected")
|
||||
}
|
||||
if u.Scheme != "https" && u.Scheme != "http" {
|
||||
return errors.New("unexpected schema")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Anchor sets the anchor part of the schema URI.
|
||||
func (id ID) Anchor(name string) ID {
|
||||
b := id.Base()
|
||||
return ID(b.String() + "#" + name)
|
||||
}
|
||||
|
||||
// Def adds or replaces a definition identifier.
|
||||
func (id ID) Def(name string) ID {
|
||||
b := id.Base()
|
||||
return ID(b.String() + "#/$defs/" + name)
|
||||
}
|
||||
|
||||
// Add appends the provided path to the id, and removes any
|
||||
// anchor data that might be there.
|
||||
func (id ID) Add(path string) ID {
|
||||
b := id.Base()
|
||||
if !strings.HasPrefix(path, "/") {
|
||||
path = "/" + path
|
||||
}
|
||||
return ID(b.String() + path)
|
||||
}
|
||||
|
||||
// Base removes any anchor information from the schema
|
||||
func (id ID) Base() ID {
|
||||
s := id.String()
|
||||
i := strings.LastIndex(s, "#")
|
||||
if i != -1 {
|
||||
s = s[0:i]
|
||||
}
|
||||
s = strings.TrimRight(s, "/")
|
||||
return ID(s)
|
||||
}
|
||||
|
||||
// String provides string version of ID
|
||||
func (id ID) String() string {
|
||||
return string(id)
|
||||
}
|
||||
+1148
File diff suppressed because it is too large.
Load diff
+146
@@ -0,0 +1,146 @@
|
||||
package jsonschema
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io/fs"
|
||||
gopath "path"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
|
||||
"go/ast"
|
||||
"go/doc"
|
||||
"go/parser"
|
||||
"go/token"
|
||||
)
|
||||
|
||||
type commentOptions struct {
|
||||
fullObjectText bool // use the first sentence only?
|
||||
}
|
||||
|
||||
// CommentOption allows for special configuration options when preparing Go
|
||||
// source files for comment extraction.
|
||||
type CommentOption func(*commentOptions)
|
||||
|
||||
// WithFullComment will configure the comment extraction to process to use an
|
||||
// object type's full comment text instead of just the synopsis.
|
||||
func WithFullComment() CommentOption {
|
||||
return func(o *commentOptions) {
|
||||
o.fullObjectText = true
|
||||
}
|
||||
}
|
||||
|
||||
// AddGoComments will update the reflectors comment map with all the comments
|
||||
// found in the provided source directories including sub-directories, in order to
|
||||
// generate a dictionary of comments associated with Types and Fields. The results
|
||||
// will be added to the `Reflect.CommentMap` ready to use with Schema "description"
|
||||
// fields.
|
||||
//
|
||||
// The `go/parser` library is used to extract all the comments and unfortunately doesn't
|
||||
// have a built-in way to determine the fully qualified name of a package. The `base`
|
||||
// parameter, the URL used to import that package, is thus required to be able to match
|
||||
// reflected types.
|
||||
//
|
||||
// When parsing type comments, by default we use the `go/doc`'s Synopsis method to extract
|
||||
// the first phrase only. Field comments, which tend to be much shorter, will include everything.
|
||||
// This behavior can be changed by using the `WithFullComment` option.
|
||||
func (r *Reflector) AddGoComments(base, path string, opts ...CommentOption) error {
|
||||
if r.CommentMap == nil {
|
||||
r.CommentMap = make(map[string]string)
|
||||
}
|
||||
co := new(commentOptions)
|
||||
for _, opt := range opts {
|
||||
opt(co)
|
||||
}
|
||||
|
||||
return r.extractGoComments(base, path, r.CommentMap, co)
|
||||
}
|
||||
|
||||
func (r *Reflector) extractGoComments(base, path string, commentMap map[string]string, opts *commentOptions) error {
|
||||
fset := token.NewFileSet()
|
||||
dict := make(map[string][]*ast.Package)
|
||||
err := filepath.Walk(path, func(path string, info fs.FileInfo, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if info.IsDir() {
|
||||
d, err := parser.ParseDir(fset, path, nil, parser.ParseComments)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, v := range d {
|
||||
// paths may have multiple packages, like for tests
|
||||
k := gopath.Join(base, path)
|
||||
dict[k] = append(dict[k], v)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for pkg, p := range dict {
|
||||
for _, f := range p {
|
||||
gtxt := ""
|
||||
typ := ""
|
||||
ast.Inspect(f, func(n ast.Node) bool {
|
||||
switch x := n.(type) {
|
||||
case *ast.TypeSpec:
|
||||
typ = x.Name.String()
|
||||
if !ast.IsExported(typ) {
|
||||
typ = ""
|
||||
} else {
|
||||
txt := x.Doc.Text()
|
||||
if txt == "" && gtxt != "" {
|
||||
txt = gtxt
|
||||
gtxt = ""
|
||||
}
|
||||
if !opts.fullObjectText {
|
||||
txt = doc.Synopsis(txt)
|
||||
}
|
||||
commentMap[fmt.Sprintf("%s.%s", pkg, typ)] = strings.TrimSpace(txt)
|
||||
}
|
||||
case *ast.Field:
|
||||
txt := x.Doc.Text()
|
||||
if txt == "" {
|
||||
txt = x.Comment.Text()
|
||||
}
|
||||
if typ != "" && txt != "" {
|
||||
for _, n := range x.Names {
|
||||
if ast.IsExported(n.String()) {
|
||||
k := fmt.Sprintf("%s.%s.%s", pkg, typ, n)
|
||||
commentMap[k] = strings.TrimSpace(txt)
|
||||
}
|
||||
}
|
||||
}
|
||||
case *ast.GenDecl:
|
||||
// remember for the next type
|
||||
gtxt = x.Doc.Text()
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Reflector) lookupComment(t reflect.Type, name string) string {
|
||||
if r.LookupComment != nil {
|
||||
if comment := r.LookupComment(t, name); comment != "" {
|
||||
return comment
|
||||
}
|
||||
}
|
||||
|
||||
if r.CommentMap == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
n := fullyQualifiedTypeName(t)
|
||||
if name != "" {
|
||||
n = n + "." + name
|
||||
}
|
||||
|
||||
return r.CommentMap[n]
|
||||
}
|
||||
+94
@@ -0,0 +1,94 @@
|
||||
package jsonschema
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
orderedmap "github.com/wk8/go-ordered-map/v2"
|
||||
)
|
||||
|
||||
// Version is the JSON Schema version.
|
||||
var Version = "https://json-schema.org/draft/2020-12/schema"
|
||||
|
||||
// Schema represents a JSON Schema object type.
|
||||
// RFC draft-bhutton-json-schema-00 section 4.3
|
||||
type Schema struct {
|
||||
// RFC draft-bhutton-json-schema-00
|
||||
Version string `json:"$schema,omitempty"` // section 8.1.1
|
||||
ID ID `json:"$id,omitempty"` // section 8.2.1
|
||||
Anchor string `json:"$anchor,omitempty"` // section 8.2.2
|
||||
Ref string `json:"$ref,omitempty"` // section 8.2.3.1
|
||||
DynamicRef string `json:"$dynamicRef,omitempty"` // section 8.2.3.2
|
||||
Definitions Definitions `json:"$defs,omitempty"` // section 8.2.4
|
||||
Comments string `json:"$comment,omitempty"` // section 8.3
|
||||
// RFC draft-bhutton-json-schema-00 section 10.2.1 (Sub-schemas with logic)
|
||||
AllOf []*Schema `json:"allOf,omitempty"` // section 10.2.1.1
|
||||
AnyOf []*Schema `json:"anyOf,omitempty"` // section 10.2.1.2
|
||||
OneOf []*Schema `json:"oneOf,omitempty"` // section 10.2.1.3
|
||||
Not *Schema `json:"not,omitempty"` // section 10.2.1.4
|
||||
// RFC draft-bhutton-json-schema-00 section 10.2.2 (Apply sub-schemas conditionally)
|
||||
If *Schema `json:"if,omitempty"` // section 10.2.2.1
|
||||
Then *Schema `json:"then,omitempty"` // section 10.2.2.2
|
||||
Else *Schema `json:"else,omitempty"` // section 10.2.2.3
|
||||
DependentSchemas map[string]*Schema `json:"dependentSchemas,omitempty"` // section 10.2.2.4
|
||||
// RFC draft-bhutton-json-schema-00 section 10.3.1 (arrays)
|
||||
PrefixItems []*Schema `json:"prefixItems,omitempty"` // section 10.3.1.1
|
||||
Items *Schema `json:"items,omitempty"` // section 10.3.1.2 (replaces additionalItems)
|
||||
Contains *Schema `json:"contains,omitempty"` // section 10.3.1.3
|
||||
// RFC draft-bhutton-json-schema-00 section 10.3.2 (sub-schemas)
|
||||
Properties *orderedmap.OrderedMap[string, *Schema] `json:"properties,omitempty"` // section 10.3.2.1
|
||||
PatternProperties map[string]*Schema `json:"patternProperties,omitempty"` // section 10.3.2.2
|
||||
AdditionalProperties *Schema `json:"additionalProperties,omitempty"` // section 10.3.2.3
|
||||
PropertyNames *Schema `json:"propertyNames,omitempty"` // section 10.3.2.4
|
||||
// RFC draft-bhutton-json-schema-validation-00, section 6
|
||||
Type string `json:"type,omitempty"` // section 6.1.1
|
||||
Enum []any `json:"enum,omitempty"` // section 6.1.2
|
||||
Const any `json:"const,omitempty"` // section 6.1.3
|
||||
MultipleOf json.Number `json:"multipleOf,omitempty"` // section 6.2.1
|
||||
Maximum json.Number `json:"maximum,omitempty"` // section 6.2.2
|
||||
ExclusiveMaximum json.Number `json:"exclusiveMaximum,omitempty"` // section 6.2.3
|
||||
Minimum json.Number `json:"minimum,omitempty"` // section 6.2.4
|
||||
ExclusiveMinimum json.Number `json:"exclusiveMinimum,omitempty"` // section 6.2.5
|
||||
MaxLength *uint64 `json:"maxLength,omitempty"` // section 6.3.1
|
||||
MinLength *uint64 `json:"minLength,omitempty"` // section 6.3.2
|
||||
Pattern string `json:"pattern,omitempty"` // section 6.3.3
|
||||
MaxItems *uint64 `json:"maxItems,omitempty"` // section 6.4.1
|
||||
MinItems *uint64 `json:"minItems,omitempty"` // section 6.4.2
|
||||
UniqueItems bool `json:"uniqueItems,omitempty"` // section 6.4.3
|
||||
MaxContains *uint64 `json:"maxContains,omitempty"` // section 6.4.4
|
||||
MinContains *uint64 `json:"minContains,omitempty"` // section 6.4.5
|
||||
MaxProperties *uint64 `json:"maxProperties,omitempty"` // section 6.5.1
|
||||
MinProperties *uint64 `json:"minProperties,omitempty"` // section 6.5.2
|
||||
Required []string `json:"required,omitempty"` // section 6.5.3
|
||||
DependentRequired map[string][]string `json:"dependentRequired,omitempty"` // section 6.5.4
|
||||
// RFC draft-bhutton-json-schema-validation-00, section 7
|
||||
Format string `json:"format,omitempty"`
|
||||
// RFC draft-bhutton-json-schema-validation-00, section 8
|
||||
ContentEncoding string `json:"contentEncoding,omitempty"` // section 8.3
|
||||
ContentMediaType string `json:"contentMediaType,omitempty"` // section 8.4
|
||||
ContentSchema *Schema `json:"contentSchema,omitempty"` // section 8.5
|
||||
// RFC draft-bhutton-json-schema-validation-00, section 9
|
||||
Title string `json:"title,omitempty"` // section 9.1
|
||||
Description string `json:"description,omitempty"` // section 9.1
|
||||
Default any `json:"default,omitempty"` // section 9.2
|
||||
Deprecated bool `json:"deprecated,omitempty"` // section 9.3
|
||||
ReadOnly bool `json:"readOnly,omitempty"` // section 9.4
|
||||
WriteOnly bool `json:"writeOnly,omitempty"` // section 9.4
|
||||
Examples []any `json:"examples,omitempty"` // section 9.5
|
||||
|
||||
Extras map[string]any `json:"-"`
|
||||
|
||||
// Special boolean representation of the Schema - section 4.3.2
|
||||
boolean *bool
|
||||
}
|
||||
|
||||
var (
|
||||
// TrueSchema defines a schema with a true value
|
||||
TrueSchema = &Schema{boolean: &[]bool{true}[0]}
|
||||
// FalseSchema defines a schema with a false value
|
||||
FalseSchema = &Schema{boolean: &[]bool{false}[0]}
|
||||
)
|
||||
|
||||
// Definitions hold schema definitions.
|
||||
// http://json-schema.org/latest/json-schema-validation.html#rfc.section.5.26
|
||||
// RFC draft-wright-json-schema-validation-00, section 5.26
|
||||
type Definitions map[string]*Schema
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
package jsonschema
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
orderedmap "github.com/wk8/go-ordered-map/v2"
|
||||
)
|
||||
|
||||
var matchFirstCap = regexp.MustCompile("(.)([A-Z][a-z]+)")
|
||||
var matchAllCap = regexp.MustCompile("([a-z0-9])([A-Z])")
|
||||
|
||||
// ToSnakeCase converts the provided string into snake case using dashes.
|
||||
// This is useful for Schema IDs and definitions to be coherent with
|
||||
// common JSON Schema examples.
|
||||
func ToSnakeCase(str string) string {
|
||||
snake := matchFirstCap.ReplaceAllString(str, "${1}-${2}")
|
||||
snake = matchAllCap.ReplaceAllString(snake, "${1}-${2}")
|
||||
return strings.ToLower(snake)
|
||||
}
|
||||
|
||||
// NewProperties is a helper method to instantiate a new properties ordered
|
||||
// map.
|
||||
func NewProperties() *orderedmap.OrderedMap[string, *Schema] {
|
||||
return orderedmap.New[string, *Schema]()
|
||||
}
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
Copyright (c) 2016 Mail.Ru Group
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
+278
@@ -0,0 +1,278 @@
|
||||
// Package buffer implements a buffer for serialization, consisting of a chain of []byte-s to
|
||||
// reduce copying and to allow reuse of individual chunks.
|
||||
package buffer
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// PoolConfig contains configuration for the allocation and reuse strategy.
|
||||
type PoolConfig struct {
|
||||
StartSize int // Minimum chunk size that is allocated.
|
||||
PooledSize int // Minimum chunk size that is reused, reusing chunks too small will result in overhead.
|
||||
MaxSize int // Maximum chunk size that will be allocated.
|
||||
}
|
||||
|
||||
var config = PoolConfig{
|
||||
StartSize: 128,
|
||||
PooledSize: 512,
|
||||
MaxSize: 32768,
|
||||
}
|
||||
|
||||
// Reuse pool: chunk size -> pool.
|
||||
var buffers = map[int]*sync.Pool{}
|
||||
|
||||
func initBuffers() {
|
||||
for l := config.PooledSize; l <= config.MaxSize; l *= 2 {
|
||||
buffers[l] = new(sync.Pool)
|
||||
}
|
||||
}
|
||||
|
||||
func init() {
|
||||
initBuffers()
|
||||
}
|
||||
|
||||
// Init sets up a non-default pooling and allocation strategy. Should be run before serialization is done.
|
||||
func Init(cfg PoolConfig) {
|
||||
config = cfg
|
||||
initBuffers()
|
||||
}
|
||||
|
||||
// putBuf puts a chunk to reuse pool if it can be reused.
|
||||
func putBuf(buf []byte) {
|
||||
size := cap(buf)
|
||||
if size < config.PooledSize {
|
||||
return
|
||||
}
|
||||
if c := buffers[size]; c != nil {
|
||||
c.Put(buf[:0])
|
||||
}
|
||||
}
|
||||
|
||||
// getBuf gets a chunk from reuse pool or creates a new one if reuse failed.
|
||||
func getBuf(size int) []byte {
|
||||
if size >= config.PooledSize {
|
||||
if c := buffers[size]; c != nil {
|
||||
v := c.Get()
|
||||
if v != nil {
|
||||
return v.([]byte)
|
||||
}
|
||||
}
|
||||
}
|
||||
return make([]byte, 0, size)
|
||||
}
|
||||
|
||||
// Buffer is a buffer optimized for serialization without extra copying.
|
||||
type Buffer struct {
|
||||
|
||||
// Buf is the current chunk that can be used for serialization.
|
||||
Buf []byte
|
||||
|
||||
toPool []byte
|
||||
bufs [][]byte
|
||||
}
|
||||
|
||||
// EnsureSpace makes sure that the current chunk contains at least s free bytes,
|
||||
// possibly creating a new chunk.
|
||||
func (b *Buffer) EnsureSpace(s int) {
|
||||
if cap(b.Buf)-len(b.Buf) < s {
|
||||
b.ensureSpaceSlow(s)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Buffer) ensureSpaceSlow(s int) {
|
||||
l := len(b.Buf)
|
||||
if l > 0 {
|
||||
if cap(b.toPool) != cap(b.Buf) {
|
||||
// Chunk was reallocated, toPool can be pooled.
|
||||
putBuf(b.toPool)
|
||||
}
|
||||
if cap(b.bufs) == 0 {
|
||||
b.bufs = make([][]byte, 0, 8)
|
||||
}
|
||||
b.bufs = append(b.bufs, b.Buf)
|
||||
l = cap(b.toPool) * 2
|
||||
} else {
|
||||
l = config.StartSize
|
||||
}
|
||||
|
||||
if l > config.MaxSize {
|
||||
l = config.MaxSize
|
||||
}
|
||||
b.Buf = getBuf(l)
|
||||
b.toPool = b.Buf
|
||||
}
|
||||
|
||||
// AppendByte appends a single byte to buffer.
|
||||
func (b *Buffer) AppendByte(data byte) {
|
||||
b.EnsureSpace(1)
|
||||
b.Buf = append(b.Buf, data)
|
||||
}
|
||||
|
||||
// AppendBytes appends a byte slice to buffer.
|
||||
func (b *Buffer) AppendBytes(data []byte) {
|
||||
if len(data) <= cap(b.Buf)-len(b.Buf) {
|
||||
b.Buf = append(b.Buf, data...) // fast path
|
||||
} else {
|
||||
b.appendBytesSlow(data)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Buffer) appendBytesSlow(data []byte) {
|
||||
for len(data) > 0 {
|
||||
b.EnsureSpace(1)
|
||||
|
||||
sz := cap(b.Buf) - len(b.Buf)
|
||||
if sz > len(data) {
|
||||
sz = len(data)
|
||||
}
|
||||
|
||||
b.Buf = append(b.Buf, data[:sz]...)
|
||||
data = data[sz:]
|
||||
}
|
||||
}
|
||||
|
||||
// AppendString appends a string to buffer.
|
||||
func (b *Buffer) AppendString(data string) {
|
||||
if len(data) <= cap(b.Buf)-len(b.Buf) {
|
||||
b.Buf = append(b.Buf, data...) // fast path
|
||||
} else {
|
||||
b.appendStringSlow(data)
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Buffer) appendStringSlow(data string) {
|
||||
for len(data) > 0 {
|
||||
b.EnsureSpace(1)
|
||||
|
||||
sz := cap(b.Buf) - len(b.Buf)
|
||||
if sz > len(data) {
|
||||
sz = len(data)
|
||||
}
|
||||
|
||||
b.Buf = append(b.Buf, data[:sz]...)
|
||||
data = data[sz:]
|
||||
}
|
||||
}
|
||||
|
||||
// Size computes the size of a buffer by adding sizes of every chunk.
|
||||
func (b *Buffer) Size() int {
|
||||
size := len(b.Buf)
|
||||
for _, buf := range b.bufs {
|
||||
size += len(buf)
|
||||
}
|
||||
return size
|
||||
}
|
||||
|
||||
// DumpTo outputs the contents of a buffer to a writer and resets the buffer.
|
||||
func (b *Buffer) DumpTo(w io.Writer) (written int, err error) {
|
||||
bufs := net.Buffers(b.bufs)
|
||||
if len(b.Buf) > 0 {
|
||||
bufs = append(bufs, b.Buf)
|
||||
}
|
||||
n, err := bufs.WriteTo(w)
|
||||
|
||||
for _, buf := range b.bufs {
|
||||
putBuf(buf)
|
||||
}
|
||||
putBuf(b.toPool)
|
||||
|
||||
b.bufs = nil
|
||||
b.Buf = nil
|
||||
b.toPool = nil
|
||||
|
||||
return int(n), err
|
||||
}
|
||||
|
||||
// BuildBytes creates a single byte slice with all the contents of the buffer. Data is
|
||||
// copied if it does not fit in a single chunk. You can optionally provide one byte
|
||||
// slice as argument that it will try to reuse.
|
||||
func (b *Buffer) BuildBytes(reuse ...[]byte) []byte {
|
||||
if len(b.bufs) == 0 {
|
||||
ret := b.Buf
|
||||
b.toPool = nil
|
||||
b.Buf = nil
|
||||
return ret
|
||||
}
|
||||
|
||||
var ret []byte
|
||||
size := b.Size()
|
||||
|
||||
// If we got a buffer as argument and it is big enough, reuse it.
|
||||
if len(reuse) == 1 && cap(reuse[0]) >= size {
|
||||
ret = reuse[0][:0]
|
||||
} else {
|
||||
ret = make([]byte, 0, size)
|
||||
}
|
||||
for _, buf := range b.bufs {
|
||||
ret = append(ret, buf...)
|
||||
putBuf(buf)
|
||||
}
|
||||
|
||||
ret = append(ret, b.Buf...)
|
||||
putBuf(b.toPool)
|
||||
|
||||
b.bufs = nil
|
||||
b.toPool = nil
|
||||
b.Buf = nil
|
||||
|
||||
return ret
|
||||
}
|
||||
|
||||
type readCloser struct {
|
||||
offset int
|
||||
bufs [][]byte
|
||||
}
|
||||
|
||||
func (r *readCloser) Read(p []byte) (n int, err error) {
|
||||
for _, buf := range r.bufs {
|
||||
// Copy as much as we can.
|
||||
x := copy(p[n:], buf[r.offset:])
|
||||
n += x // Increment how much we filled.
|
||||
|
||||
// Did we empty the whole buffer?
|
||||
if r.offset+x == len(buf) {
|
||||
// On to the next buffer.
|
||||
r.offset = 0
|
||||
r.bufs = r.bufs[1:]
|
||||
|
||||
// We can release this buffer.
|
||||
putBuf(buf)
|
||||
} else {
|
||||
r.offset += x
|
||||
}
|
||||
|
||||
if n == len(p) {
|
||||
break
|
||||
}
|
||||
}
|
||||
// No buffers left or nothing read?
|
||||
if len(r.bufs) == 0 {
|
||||
err = io.EOF
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (r *readCloser) Close() error {
|
||||
// Release all remaining buffers.
|
||||
for _, buf := range r.bufs {
|
||||
putBuf(buf)
|
||||
}
|
||||
// In case Close gets called multiple times.
|
||||
r.bufs = nil
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReadCloser creates an io.ReadCloser with all the contents of the buffer.
|
||||
func (b *Buffer) ReadCloser() io.ReadCloser {
|
||||
ret := &readCloser{0, append(b.bufs, b.Buf)}
|
||||
|
||||
b.bufs = nil
|
||||
b.toPool = nil
|
||||
b.Buf = nil
|
||||
|
||||
return ret
|
||||
}
|
||||
+405
@@ -0,0 +1,405 @@
|
||||
// Package jwriter contains a JSON writer.
|
||||
package jwriter
|
||||
|
||||
import (
|
||||
"io"
|
||||
"strconv"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/mailru/easyjson/buffer"
|
||||
)
|
||||
|
||||
// Flags describe various encoding options. The behavior may be actually implemented in the encoder, but
|
||||
// Flags field in Writer is used to set and pass them around.
|
||||
type Flags int
|
||||
|
||||
const (
|
||||
NilMapAsEmpty Flags = 1 << iota // Encode nil map as '{}' rather than 'null'.
|
||||
NilSliceAsEmpty // Encode nil slice as '[]' rather than 'null'.
|
||||
)
|
||||
|
||||
// Writer is a JSON writer.
|
||||
type Writer struct {
|
||||
Flags Flags
|
||||
|
||||
Error error
|
||||
Buffer buffer.Buffer
|
||||
NoEscapeHTML bool
|
||||
}
|
||||
|
||||
// Size returns the size of the data that was written out.
|
||||
func (w *Writer) Size() int {
|
||||
return w.Buffer.Size()
|
||||
}
|
||||
|
||||
// DumpTo outputs the data to given io.Writer, resetting the buffer.
|
||||
func (w *Writer) DumpTo(out io.Writer) (written int, err error) {
|
||||
return w.Buffer.DumpTo(out)
|
||||
}
|
||||
|
||||
// BuildBytes returns writer data as a single byte slice. You can optionally provide one byte slice
|
||||
// as argument that it will try to reuse.
|
||||
func (w *Writer) BuildBytes(reuse ...[]byte) ([]byte, error) {
|
||||
if w.Error != nil {
|
||||
return nil, w.Error
|
||||
}
|
||||
|
||||
return w.Buffer.BuildBytes(reuse...), nil
|
||||
}
|
||||
|
||||
// ReadCloser returns an io.ReadCloser that can be used to read the data.
|
||||
// ReadCloser also resets the buffer.
|
||||
func (w *Writer) ReadCloser() (io.ReadCloser, error) {
|
||||
if w.Error != nil {
|
||||
return nil, w.Error
|
||||
}
|
||||
|
||||
return w.Buffer.ReadCloser(), nil
|
||||
}
|
||||
|
||||
// RawByte appends raw binary data to the buffer.
|
||||
func (w *Writer) RawByte(c byte) {
|
||||
w.Buffer.AppendByte(c)
|
||||
}
|
||||
|
||||
// RawByte appends raw binary data to the buffer.
|
||||
func (w *Writer) RawString(s string) {
|
||||
w.Buffer.AppendString(s)
|
||||
}
|
||||
|
||||
// Raw appends raw binary data to the buffer or sets the error if it is given. Useful for
|
||||
// calling with results of MarshalJSON-like functions.
|
||||
func (w *Writer) Raw(data []byte, err error) {
|
||||
switch {
|
||||
case w.Error != nil:
|
||||
return
|
||||
case err != nil:
|
||||
w.Error = err
|
||||
case len(data) > 0:
|
||||
w.Buffer.AppendBytes(data)
|
||||
default:
|
||||
w.RawString("null")
|
||||
}
|
||||
}
|
||||
|
||||
// RawText encloses raw binary data in quotes and appends in to the buffer.
|
||||
// Useful for calling with results of MarshalText-like functions.
|
||||
func (w *Writer) RawText(data []byte, err error) {
|
||||
switch {
|
||||
case w.Error != nil:
|
||||
return
|
||||
case err != nil:
|
||||
w.Error = err
|
||||
case len(data) > 0:
|
||||
w.String(string(data))
|
||||
default:
|
||||
w.RawString("null")
|
||||
}
|
||||
}
|
||||
|
||||
// Base64Bytes appends data to the buffer after base64 encoding it
|
||||
func (w *Writer) Base64Bytes(data []byte) {
|
||||
if data == nil {
|
||||
w.Buffer.AppendString("null")
|
||||
return
|
||||
}
|
||||
w.Buffer.AppendByte('"')
|
||||
w.base64(data)
|
||||
w.Buffer.AppendByte('"')
|
||||
}
|
||||
|
||||
func (w *Writer) Uint8(n uint8) {
|
||||
w.Buffer.EnsureSpace(3)
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Uint16(n uint16) {
|
||||
w.Buffer.EnsureSpace(5)
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Uint32(n uint32) {
|
||||
w.Buffer.EnsureSpace(10)
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Uint(n uint) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Uint64(n uint64) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Int8(n int8) {
|
||||
w.Buffer.EnsureSpace(4)
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Int16(n int16) {
|
||||
w.Buffer.EnsureSpace(6)
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Int32(n int32) {
|
||||
w.Buffer.EnsureSpace(11)
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Int(n int) {
|
||||
w.Buffer.EnsureSpace(21)
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Int64(n int64) {
|
||||
w.Buffer.EnsureSpace(21)
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
|
||||
}
|
||||
|
||||
func (w *Writer) Uint8Str(n uint8) {
|
||||
w.Buffer.EnsureSpace(3)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Uint16Str(n uint16) {
|
||||
w.Buffer.EnsureSpace(5)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Uint32Str(n uint32) {
|
||||
w.Buffer.EnsureSpace(10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) UintStr(n uint) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Uint64Str(n uint64) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) UintptrStr(n uintptr) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Int8Str(n int8) {
|
||||
w.Buffer.EnsureSpace(4)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Int16Str(n int16) {
|
||||
w.Buffer.EnsureSpace(6)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Int32Str(n int32) {
|
||||
w.Buffer.EnsureSpace(11)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) IntStr(n int) {
|
||||
w.Buffer.EnsureSpace(21)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Int64Str(n int64) {
|
||||
w.Buffer.EnsureSpace(21)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Float32(n float32) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 32)
|
||||
}
|
||||
|
||||
func (w *Writer) Float32Str(n float32) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 32)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Float64(n float64) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, n, 'g', -1, 64)
|
||||
}
|
||||
|
||||
func (w *Writer) Float64Str(n float64) {
|
||||
w.Buffer.EnsureSpace(20)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 64)
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, '"')
|
||||
}
|
||||
|
||||
func (w *Writer) Bool(v bool) {
|
||||
w.Buffer.EnsureSpace(5)
|
||||
if v {
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, "true"...)
|
||||
} else {
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, "false"...)
|
||||
}
|
||||
}
|
||||
|
||||
const chars = "0123456789abcdef"
|
||||
|
||||
func getTable(falseValues ...int) [128]bool {
|
||||
table := [128]bool{}
|
||||
|
||||
for i := 0; i < 128; i++ {
|
||||
table[i] = true
|
||||
}
|
||||
|
||||
for _, v := range falseValues {
|
||||
table[v] = false
|
||||
}
|
||||
|
||||
return table
|
||||
}
|
||||
|
||||
var (
|
||||
htmlEscapeTable = getTable(0, 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, '"', '&', '<', '>', '\\')
|
||||
htmlNoEscapeTable = getTable(0, 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, '"', '\\')
|
||||
)
|
||||
|
||||
func (w *Writer) String(s string) {
|
||||
w.Buffer.AppendByte('"')
|
||||
|
||||
// Portions of the string that contain no escapes are appended as
|
||||
// byte slices.
|
||||
|
||||
p := 0 // last non-escape symbol
|
||||
|
||||
escapeTable := &htmlEscapeTable
|
||||
if w.NoEscapeHTML {
|
||||
escapeTable = &htmlNoEscapeTable
|
||||
}
|
||||
|
||||
for i := 0; i < len(s); {
|
||||
c := s[i]
|
||||
|
||||
if c < utf8.RuneSelf {
|
||||
if escapeTable[c] {
|
||||
// single-width character, no escaping is required
|
||||
i++
|
||||
continue
|
||||
}
|
||||
|
||||
w.Buffer.AppendString(s[p:i])
|
||||
switch c {
|
||||
case '\t':
|
||||
w.Buffer.AppendString(`\t`)
|
||||
case '\r':
|
||||
w.Buffer.AppendString(`\r`)
|
||||
case '\n':
|
||||
w.Buffer.AppendString(`\n`)
|
||||
case '\\':
|
||||
w.Buffer.AppendString(`\\`)
|
||||
case '"':
|
||||
w.Buffer.AppendString(`\"`)
|
||||
default:
|
||||
w.Buffer.AppendString(`\u00`)
|
||||
w.Buffer.AppendByte(chars[c>>4])
|
||||
w.Buffer.AppendByte(chars[c&0xf])
|
||||
}
|
||||
|
||||
i++
|
||||
p = i
|
||||
continue
|
||||
}
|
||||
|
||||
// broken utf
|
||||
runeValue, runeWidth := utf8.DecodeRuneInString(s[i:])
|
||||
if runeValue == utf8.RuneError && runeWidth == 1 {
|
||||
w.Buffer.AppendString(s[p:i])
|
||||
w.Buffer.AppendString(`\ufffd`)
|
||||
i++
|
||||
p = i
|
||||
continue
|
||||
}
|
||||
|
||||
// jsonp stuff - tab separator and line separator
|
||||
if runeValue == '\u2028' || runeValue == '\u2029' {
|
||||
w.Buffer.AppendString(s[p:i])
|
||||
w.Buffer.AppendString(`\u202`)
|
||||
w.Buffer.AppendByte(chars[runeValue&0xf])
|
||||
i += runeWidth
|
||||
p = i
|
||||
continue
|
||||
}
|
||||
i += runeWidth
|
||||
}
|
||||
w.Buffer.AppendString(s[p:])
|
||||
w.Buffer.AppendByte('"')
|
||||
}
|
||||
|
||||
const encode = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
|
||||
const padChar = '='
|
||||
|
||||
func (w *Writer) base64(in []byte) {
|
||||
|
||||
if len(in) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
w.Buffer.EnsureSpace(((len(in)-1)/3 + 1) * 4)
|
||||
|
||||
si := 0
|
||||
n := (len(in) / 3) * 3
|
||||
|
||||
for si < n {
|
||||
// Convert 3x 8bit source bytes into 4 bytes
|
||||
val := uint(in[si+0])<<16 | uint(in[si+1])<<8 | uint(in[si+2])
|
||||
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, encode[val>>18&0x3F], encode[val>>12&0x3F], encode[val>>6&0x3F], encode[val&0x3F])
|
||||
|
||||
si += 3
|
||||
}
|
||||
|
||||
remain := len(in) - si
|
||||
if remain == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Add the remaining small block
|
||||
val := uint(in[si+0]) << 16
|
||||
if remain == 2 {
|
||||
val |= uint(in[si+1]) << 8
|
||||
}
|
||||
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, encode[val>>18&0x3F], encode[val>>12&0x3F])
|
||||
|
||||
switch remain {
|
||||
case 2:
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, encode[val>>6&0x3F], byte(padChar))
|
||||
case 1:
|
||||
w.Buffer.Buf = append(w.Buffer.Buf, byte(padChar), byte(padChar))
|
||||
}
|
||||
}
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 Anthropic, PBC
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+710
@@ -0,0 +1,710 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/mark3labs/mcp-go/client/transport"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// Client implements the MCP client.
|
||||
type Client struct {
|
||||
transport transport.Interface
|
||||
|
||||
initialized bool
|
||||
notifications []func(mcp.JSONRPCNotification)
|
||||
notifyMu sync.RWMutex
|
||||
requestID atomic.Int64
|
||||
clientCapabilities mcp.ClientCapabilities
|
||||
serverCapabilities mcp.ServerCapabilities
|
||||
protocolVersion string
|
||||
samplingHandler SamplingHandler
|
||||
rootsHandler RootsHandler
|
||||
elicitationHandler ElicitationHandler
|
||||
}
|
||||
|
||||
type ClientOption func(*Client)
|
||||
|
||||
// WithClientCapabilities sets the client capabilities for the client.
|
||||
func WithClientCapabilities(capabilities mcp.ClientCapabilities) ClientOption {
|
||||
return func(c *Client) {
|
||||
c.clientCapabilities = capabilities
|
||||
}
|
||||
}
|
||||
|
||||
// WithSamplingHandler sets the sampling handler for the client.
|
||||
// When set, the client will declare sampling capability during initialization.
|
||||
func WithSamplingHandler(handler SamplingHandler) ClientOption {
|
||||
return func(c *Client) {
|
||||
c.samplingHandler = handler
|
||||
}
|
||||
}
|
||||
|
||||
// WithRootsHandler sets the roots handler for the client.
|
||||
// WithRootsHandler returns a ClientOption that sets the client's RootsHandler.
|
||||
// When provided, the client will declare the roots capability (ListChanged) during initialization.
|
||||
func WithRootsHandler(handler RootsHandler) ClientOption {
|
||||
return func(c *Client) {
|
||||
c.rootsHandler = handler
|
||||
}
|
||||
}
|
||||
|
||||
// WithElicitationHandler sets the elicitation handler for the client.
|
||||
// When set, the client will declare elicitation capability during initialization.
|
||||
func WithElicitationHandler(handler ElicitationHandler) ClientOption {
|
||||
return func(c *Client) {
|
||||
c.elicitationHandler = handler
|
||||
}
|
||||
}
|
||||
|
||||
// WithSession assumes a MCP Session has already been initialized
|
||||
func WithSession() ClientOption {
|
||||
return func(c *Client) {
|
||||
c.initialized = true
|
||||
}
|
||||
}
|
||||
|
||||
// NewClient creates a new MCP client with the given transport.
|
||||
// Usage:
|
||||
//
|
||||
// stdio := transport.NewStdio("./mcp_server", nil, "xxx")
|
||||
// client, err := NewClient(stdio)
|
||||
// if err != nil {
|
||||
// log.Fatalf("Failed to create client: %v", err)
|
||||
// }
|
||||
func NewClient(transport transport.Interface, options ...ClientOption) *Client {
|
||||
client := &Client{
|
||||
transport: transport,
|
||||
}
|
||||
|
||||
for _, opt := range options {
|
||||
opt(client)
|
||||
}
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
// Start initiates the connection to the server.
|
||||
// Must be called before using the client.
|
||||
func (c *Client) Start(ctx context.Context) error {
|
||||
if c.transport == nil {
|
||||
return fmt.Errorf("transport is nil")
|
||||
}
|
||||
|
||||
// Start is idempotent - transports handle being called multiple times
|
||||
err := c.transport.Start(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
c.transport.SetNotificationHandler(func(notification mcp.JSONRPCNotification) {
|
||||
c.notifyMu.RLock()
|
||||
defer c.notifyMu.RUnlock()
|
||||
for _, handler := range c.notifications {
|
||||
handler(notification)
|
||||
}
|
||||
})
|
||||
|
||||
// Set up request handler for bidirectional communication (e.g., sampling)
|
||||
if bidirectional, ok := c.transport.(transport.BidirectionalInterface); ok {
|
||||
bidirectional.SetRequestHandler(c.handleIncomingRequest)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close shuts down the client and closes the transport.
|
||||
func (c *Client) Close() error {
|
||||
return c.transport.Close()
|
||||
}
|
||||
|
||||
// OnNotification registers a handler function to be called when notifications are received.
|
||||
// Multiple handlers can be registered and will be called in the order they were added.
|
||||
func (c *Client) OnNotification(
|
||||
handler func(notification mcp.JSONRPCNotification),
|
||||
) {
|
||||
c.notifyMu.Lock()
|
||||
defer c.notifyMu.Unlock()
|
||||
c.notifications = append(c.notifications, handler)
|
||||
}
|
||||
|
||||
// OnConnectionLost registers a handler function to be called when the connection is lost.
|
||||
// This is useful for handling HTTP2 idle timeout disconnections that should not be treated as errors.
|
||||
func (c *Client) OnConnectionLost(handler func(error)) {
|
||||
type connectionLostSetter interface {
|
||||
SetConnectionLostHandler(func(error))
|
||||
}
|
||||
if setter, ok := c.transport.(connectionLostSetter); ok {
|
||||
setter.SetConnectionLostHandler(handler)
|
||||
}
|
||||
}
|
||||
|
||||
// sendRequest sends a JSON-RPC request to the server and waits for a response.
|
||||
// Returns the raw JSON response message or an error if the request fails.
|
||||
func (c *Client) sendRequest(
|
||||
ctx context.Context,
|
||||
method string,
|
||||
params any,
|
||||
header http.Header,
|
||||
) (*json.RawMessage, error) {
|
||||
if !c.initialized && method != "initialize" {
|
||||
return nil, fmt.Errorf("client not initialized")
|
||||
}
|
||||
|
||||
id := c.requestID.Add(1)
|
||||
|
||||
request := transport.JSONRPCRequest{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
ID: mcp.NewRequestId(id),
|
||||
Method: method,
|
||||
Params: params,
|
||||
Header: header,
|
||||
}
|
||||
|
||||
response, err := c.transport.SendRequest(ctx, request)
|
||||
if err != nil {
|
||||
return nil, transport.NewError(err)
|
||||
}
|
||||
|
||||
if response.Error != nil {
|
||||
return nil, response.Error.AsError()
|
||||
}
|
||||
|
||||
return &response.Result, nil
|
||||
}
|
||||
|
||||
// Initialize negotiates with the server.
|
||||
// Must be called after Start, and before any request methods.
|
||||
func (c *Client) Initialize(
|
||||
ctx context.Context,
|
||||
request mcp.InitializeRequest,
|
||||
) (*mcp.InitializeResult, error) {
|
||||
// Merge client capabilities with sampling capability if handler is configured
|
||||
capabilities := request.Params.Capabilities
|
||||
if c.samplingHandler != nil {
|
||||
capabilities.Sampling = &struct{}{}
|
||||
}
|
||||
if c.rootsHandler != nil {
|
||||
capabilities.Roots = &struct {
|
||||
ListChanged bool `json:"listChanged,omitempty"`
|
||||
}{
|
||||
ListChanged: true,
|
||||
}
|
||||
}
|
||||
// Add elicitation capability if handler is configured
|
||||
if c.elicitationHandler != nil {
|
||||
capabilities.Elicitation = &struct{}{}
|
||||
}
|
||||
|
||||
// Ensure we send a params object with all required fields
|
||||
params := struct {
|
||||
ProtocolVersion string `json:"protocolVersion"`
|
||||
ClientInfo mcp.Implementation `json:"clientInfo"`
|
||||
Capabilities mcp.ClientCapabilities `json:"capabilities"`
|
||||
}{
|
||||
ProtocolVersion: request.Params.ProtocolVersion,
|
||||
ClientInfo: request.Params.ClientInfo,
|
||||
Capabilities: capabilities,
|
||||
}
|
||||
|
||||
response, err := c.sendRequest(ctx, "initialize", params, request.Header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result mcp.InitializeResult
|
||||
if err := json.Unmarshal(*response, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
// Validate protocol version
|
||||
if !slices.Contains(mcp.ValidProtocolVersions, result.ProtocolVersion) {
|
||||
return nil, mcp.UnsupportedProtocolVersionError{Version: result.ProtocolVersion}
|
||||
}
|
||||
|
||||
// Store serverCapabilities and protocol version
|
||||
c.serverCapabilities = result.Capabilities
|
||||
c.protocolVersion = result.ProtocolVersion
|
||||
|
||||
// Set protocol version on HTTP transports
|
||||
if httpConn, ok := c.transport.(transport.HTTPConnection); ok {
|
||||
httpConn.SetProtocolVersion(result.ProtocolVersion)
|
||||
}
|
||||
|
||||
// Send initialized notification
|
||||
notification := mcp.JSONRPCNotification{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
Notification: mcp.Notification{
|
||||
Method: "notifications/initialized",
|
||||
},
|
||||
}
|
||||
|
||||
err = c.transport.SendNotification(ctx, notification)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf(
|
||||
"failed to send initialized notification: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
|
||||
c.initialized = true
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
func (c *Client) Ping(ctx context.Context) error {
|
||||
_, err := c.sendRequest(ctx, "ping", nil, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListResourcesByPage manually list resources by page.
|
||||
func (c *Client) ListResourcesByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourcesRequest,
|
||||
) (*mcp.ListResourcesResult, error) {
|
||||
result, err := listByPage[mcp.ListResourcesResult](ctx, c, request.PaginatedRequest, "resources/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) ListResources(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourcesRequest,
|
||||
) (*mcp.ListResourcesResult, error) {
|
||||
result, err := c.ListResourcesByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for result.NextCursor != "" {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
request.Params.Cursor = result.NextCursor
|
||||
newPageRes, err := c.ListResourcesByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Resources = append(result.Resources, newPageRes.Resources...)
|
||||
result.NextCursor = newPageRes.NextCursor
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) ListResourceTemplatesByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourceTemplatesRequest,
|
||||
) (*mcp.ListResourceTemplatesResult, error) {
|
||||
result, err := listByPage[mcp.ListResourceTemplatesResult](ctx, c, request.PaginatedRequest, "resources/templates/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) ListResourceTemplates(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourceTemplatesRequest,
|
||||
) (*mcp.ListResourceTemplatesResult, error) {
|
||||
result, err := c.ListResourceTemplatesByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for result.NextCursor != "" {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
request.Params.Cursor = result.NextCursor
|
||||
newPageRes, err := c.ListResourceTemplatesByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.ResourceTemplates = append(result.ResourceTemplates, newPageRes.ResourceTemplates...)
|
||||
result.NextCursor = newPageRes.NextCursor
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) ReadResource(
|
||||
ctx context.Context,
|
||||
request mcp.ReadResourceRequest,
|
||||
) (*mcp.ReadResourceResult, error) {
|
||||
response, err := c.sendRequest(ctx, "resources/read", request.Params, request.Header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return mcp.ParseReadResourceResult(response)
|
||||
}
|
||||
|
||||
func (c *Client) Subscribe(
|
||||
ctx context.Context,
|
||||
request mcp.SubscribeRequest,
|
||||
) error {
|
||||
_, err := c.sendRequest(ctx, "resources/subscribe", request.Params, request.Header)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Client) Unsubscribe(
|
||||
ctx context.Context,
|
||||
request mcp.UnsubscribeRequest,
|
||||
) error {
|
||||
_, err := c.sendRequest(ctx, "resources/unsubscribe", request.Params, request.Header)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Client) ListPromptsByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListPromptsRequest,
|
||||
) (*mcp.ListPromptsResult, error) {
|
||||
result, err := listByPage[mcp.ListPromptsResult](ctx, c, request.PaginatedRequest, "prompts/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) ListPrompts(
|
||||
ctx context.Context,
|
||||
request mcp.ListPromptsRequest,
|
||||
) (*mcp.ListPromptsResult, error) {
|
||||
result, err := c.ListPromptsByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for result.NextCursor != "" {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
request.Params.Cursor = result.NextCursor
|
||||
newPageRes, err := c.ListPromptsByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Prompts = append(result.Prompts, newPageRes.Prompts...)
|
||||
result.NextCursor = newPageRes.NextCursor
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) GetPrompt(
|
||||
ctx context.Context,
|
||||
request mcp.GetPromptRequest,
|
||||
) (*mcp.GetPromptResult, error) {
|
||||
response, err := c.sendRequest(ctx, "prompts/get", request.Params, request.Header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return mcp.ParseGetPromptResult(response)
|
||||
}
|
||||
|
||||
func (c *Client) ListToolsByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListToolsRequest,
|
||||
) (*mcp.ListToolsResult, error) {
|
||||
result, err := listByPage[mcp.ListToolsResult](ctx, c, request.PaginatedRequest, "tools/list")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) ListTools(
|
||||
ctx context.Context,
|
||||
request mcp.ListToolsRequest,
|
||||
) (*mcp.ListToolsResult, error) {
|
||||
result, err := c.ListToolsByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for result.NextCursor != "" {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
request.Params.Cursor = result.NextCursor
|
||||
newPageRes, err := c.ListToolsByPage(ctx, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result.Tools = append(result.Tools, newPageRes.Tools...)
|
||||
result.NextCursor = newPageRes.NextCursor
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *Client) CallTool(
|
||||
ctx context.Context,
|
||||
request mcp.CallToolRequest,
|
||||
) (*mcp.CallToolResult, error) {
|
||||
response, err := c.sendRequest(ctx, "tools/call", request.Params, request.Header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return mcp.ParseCallToolResult(response)
|
||||
}
|
||||
|
||||
func (c *Client) SetLevel(
|
||||
ctx context.Context,
|
||||
request mcp.SetLevelRequest,
|
||||
) error {
|
||||
_, err := c.sendRequest(ctx, "logging/setLevel", request.Params, request.Header)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Client) Complete(
|
||||
ctx context.Context,
|
||||
request mcp.CompleteRequest,
|
||||
) (*mcp.CompleteResult, error) {
|
||||
response, err := c.sendRequest(ctx, "completion/complete", request.Params, request.Header)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var result mcp.CompleteResult
|
||||
if err := json.Unmarshal(*response, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// RootListChanges sends a roots list-changed notification to the server.
|
||||
func (c *Client) RootListChanges(
|
||||
ctx context.Context,
|
||||
) error {
|
||||
// Send root list changes notification
|
||||
notification := mcp.JSONRPCNotification{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
Notification: mcp.Notification{
|
||||
Method: mcp.MethodNotificationRootsListChanged,
|
||||
},
|
||||
}
|
||||
|
||||
err := c.transport.SendNotification(ctx, notification)
|
||||
if err != nil {
|
||||
return fmt.Errorf(
|
||||
"failed to send root list change notification: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleIncomingRequest processes incoming requests from the server.
|
||||
// This is the main entry point for server-to-client requests like sampling and elicitation.
|
||||
func (c *Client) handleIncomingRequest(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
|
||||
switch request.Method {
|
||||
case string(mcp.MethodSamplingCreateMessage):
|
||||
return c.handleSamplingRequestTransport(ctx, request)
|
||||
case string(mcp.MethodElicitationCreate):
|
||||
return c.handleElicitationRequestTransport(ctx, request)
|
||||
case string(mcp.MethodPing):
|
||||
return c.handlePingRequestTransport(ctx, request)
|
||||
case string(mcp.MethodListRoots):
|
||||
return c.handleListRootsRequestTransport(ctx, request)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported request method: %s", request.Method)
|
||||
}
|
||||
}
|
||||
|
||||
// handleSamplingRequestTransport handles sampling requests at the transport level.
|
||||
func (c *Client) handleSamplingRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
|
||||
if c.samplingHandler == nil {
|
||||
return nil, fmt.Errorf("no sampling handler configured")
|
||||
}
|
||||
|
||||
// Parse the request parameters
|
||||
var params mcp.CreateMessageParams
|
||||
if request.Params != nil {
|
||||
paramsBytes, err := json.Marshal(request.Params)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal params: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal(paramsBytes, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal params: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Fix content parsing - HTTP transport unmarshals TextContent as map[string]any
|
||||
// Use the helper function to properly handle content from different transports
|
||||
for i := range params.Messages {
|
||||
if contentMap, ok := params.Messages[i].Content.(map[string]any); ok {
|
||||
// Parse the content map into a proper Content type
|
||||
content, err := mcp.ParseContent(contentMap)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse content for message %d: %w", i, err)
|
||||
}
|
||||
params.Messages[i].Content = content
|
||||
}
|
||||
}
|
||||
|
||||
// Create the MCP request
|
||||
mcpRequest := mcp.CreateMessageRequest{
|
||||
Request: mcp.Request{
|
||||
Method: string(mcp.MethodSamplingCreateMessage),
|
||||
},
|
||||
CreateMessageParams: params,
|
||||
}
|
||||
|
||||
// Call the sampling handler
|
||||
result, err := c.samplingHandler.CreateMessage(ctx, mcpRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Marshal the result
|
||||
resultBytes, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal result: %w", err)
|
||||
}
|
||||
|
||||
// Create the transport response
|
||||
response := transport.NewJSONRPCResultResponse(request.ID, json.RawMessage(resultBytes))
|
||||
|
||||
return response, nil
|
||||
}
|
||||
|
||||
// handleListRootsRequestTransport handles list roots requests at the transport level.
|
||||
func (c *Client) handleListRootsRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
|
||||
if c.rootsHandler == nil {
|
||||
return nil, fmt.Errorf("no roots handler configured")
|
||||
}
|
||||
|
||||
// Create the MCP request
|
||||
mcpRequest := mcp.ListRootsRequest{
|
||||
Request: mcp.Request{
|
||||
Method: string(mcp.MethodListRoots),
|
||||
},
|
||||
}
|
||||
|
||||
// Call the list roots handler
|
||||
result, err := c.rootsHandler.ListRoots(ctx, mcpRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Marshal the result
|
||||
resultBytes, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal result: %w", err)
|
||||
}
|
||||
|
||||
// Create the transport response
|
||||
response := transport.NewJSONRPCResultResponse(request.ID, json.RawMessage(resultBytes))
|
||||
|
||||
return response, nil
|
||||
}
|
||||
|
||||
// handleElicitationRequestTransport handles elicitation requests at the transport level.
|
||||
func (c *Client) handleElicitationRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
|
||||
if c.elicitationHandler == nil {
|
||||
return nil, fmt.Errorf("no elicitation handler configured")
|
||||
}
|
||||
|
||||
// Parse the request parameters
|
||||
var params mcp.ElicitationParams
|
||||
if request.Params != nil {
|
||||
paramsBytes, err := json.Marshal(request.Params)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal params: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal(paramsBytes, ¶ms); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal params: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Create the MCP request
|
||||
mcpRequest := mcp.ElicitationRequest{
|
||||
Request: mcp.Request{
|
||||
Method: string(mcp.MethodElicitationCreate),
|
||||
},
|
||||
Params: params,
|
||||
}
|
||||
|
||||
// Call the elicitation handler
|
||||
result, err := c.elicitationHandler.Elicit(ctx, mcpRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Marshal the result
|
||||
resultBytes, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal result: %w", err)
|
||||
}
|
||||
|
||||
// Create the transport response
|
||||
response := transport.NewJSONRPCResultResponse(request.ID, resultBytes)
|
||||
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func (c *Client) handlePingRequestTransport(ctx context.Context, request transport.JSONRPCRequest) (*transport.JSONRPCResponse, error) {
|
||||
b, _ := json.Marshal(&mcp.EmptyResult{})
|
||||
return transport.NewJSONRPCResultResponse(request.ID, b), nil
|
||||
}
|
||||
|
||||
func listByPage[T any](
|
||||
ctx context.Context,
|
||||
client *Client,
|
||||
request mcp.PaginatedRequest,
|
||||
method string,
|
||||
) (*T, error) {
|
||||
response, err := client.sendRequest(ctx, method, request.Params, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var result T
|
||||
if err := json.Unmarshal(*response, &result); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// Helper methods
|
||||
|
||||
// GetTransport gives access to the underlying transport layer.
|
||||
// Cast it to the specific transport type and obtain the other helper methods.
|
||||
func (c *Client) GetTransport() transport.Interface {
|
||||
return c.transport
|
||||
}
|
||||
|
||||
// GetServerCapabilities returns the server capabilities.
|
||||
func (c *Client) GetServerCapabilities() mcp.ServerCapabilities {
|
||||
return c.serverCapabilities
|
||||
}
|
||||
|
||||
// GetClientCapabilities returns the client capabilities.
|
||||
func (c *Client) GetClientCapabilities() mcp.ClientCapabilities {
|
||||
return c.clientCapabilities
|
||||
}
|
||||
|
||||
// GetSessionId returns the session ID of the transport.
|
||||
// If the transport does not support sessions, it returns an empty string.
|
||||
func (c *Client) GetSessionId() string {
|
||||
if c.transport == nil {
|
||||
return ""
|
||||
}
|
||||
return c.transport.GetSessionId()
|
||||
}
|
||||
|
||||
// IsInitialized returns true if the client has been initialized.
|
||||
func (c *Client) IsInitialized() bool {
|
||||
return c.initialized
|
||||
}
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// ElicitationHandler defines the interface for handling elicitation requests from servers.
|
||||
// Clients can implement this interface to request additional information from users.
|
||||
type ElicitationHandler interface {
|
||||
// Elicit handles an elicitation request from the server and returns the user's response.
|
||||
// The implementation should:
|
||||
// 1. Present the request message to the user
|
||||
// 2. Validate input against the requested schema
|
||||
// 3. Allow the user to accept, decline, or cancel
|
||||
// 4. Return the appropriate response
|
||||
Elicit(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error)
|
||||
}
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/mark3labs/mcp-go/client/transport"
|
||||
)
|
||||
|
||||
// NewStreamableHttpClient is a convenience method that creates a new streamable-http-based MCP client
|
||||
// with the given base URL. Returns an error if the URL is invalid.
|
||||
func NewStreamableHttpClient(baseURL string, options ...transport.StreamableHTTPCOption) (*Client, error) {
|
||||
trans, err := transport.NewStreamableHTTP(baseURL, options...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create SSE transport: %w", err)
|
||||
}
|
||||
clientOptions := make([]ClientOption, 0)
|
||||
sessionID := trans.GetSessionId()
|
||||
if sessionID != "" {
|
||||
clientOptions = append(clientOptions, WithSession())
|
||||
}
|
||||
return NewClient(trans, clientOptions...), nil
|
||||
}
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mark3labs/mcp-go/client/transport"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/mark3labs/mcp-go/server"
|
||||
)
|
||||
|
||||
// NewInProcessClient connect directly to a mcp server object in the same process
|
||||
func NewInProcessClient(server *server.MCPServer) (*Client, error) {
|
||||
inProcessTransport := transport.NewInProcessTransport(server)
|
||||
return NewClient(inProcessTransport), nil
|
||||
}
|
||||
|
||||
// NewInProcessClientWithSamplingHandler creates an in-process client with sampling support
|
||||
func NewInProcessClientWithSamplingHandler(server *server.MCPServer, handler SamplingHandler) (*Client, error) {
|
||||
// Create a wrapper that implements server.SamplingHandler
|
||||
serverHandler := &inProcessSamplingHandlerWrapper{handler: handler}
|
||||
|
||||
inProcessTransport := transport.NewInProcessTransportWithOptions(server,
|
||||
transport.WithSamplingHandler(serverHandler))
|
||||
|
||||
client := NewClient(inProcessTransport)
|
||||
client.samplingHandler = handler
|
||||
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// inProcessSamplingHandlerWrapper wraps client.SamplingHandler to implement server.SamplingHandler
|
||||
type inProcessSamplingHandlerWrapper struct {
|
||||
handler SamplingHandler
|
||||
}
|
||||
|
||||
func (w *inProcessSamplingHandlerWrapper) CreateMessage(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
|
||||
return w.handler.CreateMessage(ctx, request)
|
||||
}
|
||||
+109
@@ -0,0 +1,109 @@
|
||||
// Package client provides MCP (Model Context Protocol) client implementations.
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// MCPClient represents an MCP client interface
|
||||
type MCPClient interface {
|
||||
// Initialize sends the initial connection request to the server
|
||||
Initialize(
|
||||
ctx context.Context,
|
||||
request mcp.InitializeRequest,
|
||||
) (*mcp.InitializeResult, error)
|
||||
|
||||
// Ping checks if the server is alive
|
||||
Ping(ctx context.Context) error
|
||||
|
||||
// ListResourcesByPage manually list resources by page.
|
||||
ListResourcesByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourcesRequest,
|
||||
) (*mcp.ListResourcesResult, error)
|
||||
|
||||
// ListResources requests a list of available resources from the server
|
||||
ListResources(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourcesRequest,
|
||||
) (*mcp.ListResourcesResult, error)
|
||||
|
||||
// ListResourceTemplatesByPage manually list resource templates by page.
|
||||
ListResourceTemplatesByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourceTemplatesRequest,
|
||||
) (*mcp.ListResourceTemplatesResult,
|
||||
error)
|
||||
|
||||
// ListResourceTemplates requests a list of available resource templates from the server
|
||||
ListResourceTemplates(
|
||||
ctx context.Context,
|
||||
request mcp.ListResourceTemplatesRequest,
|
||||
) (*mcp.ListResourceTemplatesResult,
|
||||
error)
|
||||
|
||||
// ReadResource reads a specific resource from the server
|
||||
ReadResource(
|
||||
ctx context.Context,
|
||||
request mcp.ReadResourceRequest,
|
||||
) (*mcp.ReadResourceResult, error)
|
||||
|
||||
// Subscribe requests notifications for changes to a specific resource
|
||||
Subscribe(ctx context.Context, request mcp.SubscribeRequest) error
|
||||
|
||||
// Unsubscribe cancels notifications for a specific resource
|
||||
Unsubscribe(ctx context.Context, request mcp.UnsubscribeRequest) error
|
||||
|
||||
// ListPromptsByPage manually list prompts by page.
|
||||
ListPromptsByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListPromptsRequest,
|
||||
) (*mcp.ListPromptsResult, error)
|
||||
|
||||
// ListPrompts requests a list of available prompts from the server
|
||||
ListPrompts(
|
||||
ctx context.Context,
|
||||
request mcp.ListPromptsRequest,
|
||||
) (*mcp.ListPromptsResult, error)
|
||||
|
||||
// GetPrompt retrieves a specific prompt from the server
|
||||
GetPrompt(
|
||||
ctx context.Context,
|
||||
request mcp.GetPromptRequest,
|
||||
) (*mcp.GetPromptResult, error)
|
||||
|
||||
// ListToolsByPage manually list tools by page.
|
||||
ListToolsByPage(
|
||||
ctx context.Context,
|
||||
request mcp.ListToolsRequest,
|
||||
) (*mcp.ListToolsResult, error)
|
||||
|
||||
// ListTools requests a list of available tools from the server
|
||||
ListTools(
|
||||
ctx context.Context,
|
||||
request mcp.ListToolsRequest,
|
||||
) (*mcp.ListToolsResult, error)
|
||||
|
||||
// CallTool invokes a specific tool on the server
|
||||
CallTool(
|
||||
ctx context.Context,
|
||||
request mcp.CallToolRequest,
|
||||
) (*mcp.CallToolResult, error)
|
||||
|
||||
// SetLevel sets the logging level for the server
|
||||
SetLevel(ctx context.Context, request mcp.SetLevelRequest) error
|
||||
|
||||
// Complete requests completion options for a given argument
|
||||
Complete(
|
||||
ctx context.Context,
|
||||
request mcp.CompleteRequest,
|
||||
) (*mcp.CompleteResult, error)
|
||||
|
||||
// Close client connection and cleanup resources
|
||||
Close() error
|
||||
|
||||
// OnNotification registers a handler for notifications
|
||||
OnNotification(handler func(notification mcp.JSONRPCNotification))
|
||||
}
|
||||
+76
@@ -0,0 +1,76 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/mark3labs/mcp-go/client/transport"
|
||||
)
|
||||
|
||||
// OAuthConfig is a convenience type that wraps transport.OAuthConfig
|
||||
type OAuthConfig = transport.OAuthConfig
|
||||
|
||||
// Token is a convenience type that wraps transport.Token
|
||||
type Token = transport.Token
|
||||
|
||||
// TokenStore is a convenience type that wraps transport.TokenStore
|
||||
type TokenStore = transport.TokenStore
|
||||
|
||||
// MemoryTokenStore is a convenience type that wraps transport.MemoryTokenStore
|
||||
type MemoryTokenStore = transport.MemoryTokenStore
|
||||
|
||||
// NewMemoryTokenStore is a convenience function that wraps transport.NewMemoryTokenStore
|
||||
var NewMemoryTokenStore = transport.NewMemoryTokenStore
|
||||
|
||||
// NewOAuthStreamableHttpClient creates a new streamable-http-based MCP client with OAuth support.
|
||||
// Returns an error if the URL is invalid.
|
||||
func NewOAuthStreamableHttpClient(baseURL string, oauthConfig OAuthConfig, options ...transport.StreamableHTTPCOption) (*Client, error) {
|
||||
// Add OAuth option to the list of options
|
||||
options = append(options, transport.WithHTTPOAuth(oauthConfig))
|
||||
|
||||
trans, err := transport.NewStreamableHTTP(baseURL, options...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create HTTP transport: %w", err)
|
||||
}
|
||||
return NewClient(trans), nil
|
||||
}
|
||||
|
||||
// NewOAuthStreamableHttpClient creates a new streamable-http-based MCP client with OAuth support.
|
||||
// Returns an error if the URL is invalid.
|
||||
func NewOAuthSSEClient(baseURL string, oauthConfig OAuthConfig, options ...transport.ClientOption) (*Client, error) {
|
||||
// Add OAuth option to the list of options
|
||||
options = append(options, transport.WithOAuth(oauthConfig))
|
||||
|
||||
trans, err := transport.NewSSE(baseURL, options...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create SSE transport: %w", err)
|
||||
}
|
||||
return NewClient(trans), nil
|
||||
}
|
||||
|
||||
// GenerateCodeVerifier generates a code verifier for PKCE
|
||||
var GenerateCodeVerifier = transport.GenerateCodeVerifier
|
||||
|
||||
// GenerateCodeChallenge generates a code challenge from a code verifier
|
||||
var GenerateCodeChallenge = transport.GenerateCodeChallenge
|
||||
|
||||
// GenerateState generates a state parameter for OAuth
|
||||
var GenerateState = transport.GenerateState
|
||||
|
||||
// OAuthAuthorizationRequiredError is returned when OAuth authorization is required
|
||||
type OAuthAuthorizationRequiredError = transport.OAuthAuthorizationRequiredError
|
||||
|
||||
// IsOAuthAuthorizationRequiredError checks if an error is an OAuthAuthorizationRequiredError
|
||||
func IsOAuthAuthorizationRequiredError(err error) bool {
|
||||
var target *OAuthAuthorizationRequiredError
|
||||
return errors.As(err, &target)
|
||||
}
|
||||
|
||||
// GetOAuthHandler extracts the OAuthHandler from an OAuthAuthorizationRequiredError
|
||||
func GetOAuthHandler(err error) *transport.OAuthHandler {
|
||||
var oauthErr *OAuthAuthorizationRequiredError
|
||||
if errors.As(err, &oauthErr) {
|
||||
return oauthErr.Handler
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// RootsHandler defines the interface for handling roots requests from servers.
|
||||
// Clients can implement this interface to provide roots list to servers.
|
||||
type RootsHandler interface {
|
||||
// ListRoots handles a list root request from the server and returns the roots list.
|
||||
// The implementation should:
|
||||
// 1. Validate input against the requested schema
|
||||
// 2. Return the appropriate response
|
||||
ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error)
|
||||
}
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// SamplingHandler defines the interface for handling sampling requests from servers.
|
||||
// Clients can implement this interface to provide LLM sampling capabilities to servers.
|
||||
type SamplingHandler interface {
|
||||
// CreateMessage handles a sampling request from the server and returns the generated message.
|
||||
// The implementation should:
|
||||
// 1. Validate the request parameters
|
||||
// 2. Optionally prompt the user for approval (human-in-the-loop)
|
||||
// 3. Select an appropriate model based on preferences
|
||||
// 4. Generate the response using the selected model
|
||||
// 5. Return the result with model information and stop reason
|
||||
CreateMessage(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error)
|
||||
}
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/mark3labs/mcp-go/client/transport"
|
||||
)
|
||||
|
||||
func WithHeaders(headers map[string]string) transport.ClientOption {
|
||||
return transport.WithHeaders(headers)
|
||||
}
|
||||
|
||||
func WithHeaderFunc(headerFunc transport.HTTPHeaderFunc) transport.ClientOption {
|
||||
return transport.WithHeaderFunc(headerFunc)
|
||||
}
|
||||
|
||||
func WithHTTPClient(httpClient *http.Client) transport.ClientOption {
|
||||
return transport.WithHTTPClient(httpClient)
|
||||
}
|
||||
|
||||
// NewSSEMCPClient creates a new SSE-based MCP client with the given base URL.
|
||||
// Returns an error if the URL is invalid.
|
||||
func NewSSEMCPClient(baseURL string, options ...transport.ClientOption) (*Client, error) {
|
||||
sseTransport, err := transport.NewSSE(baseURL, options...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create SSE transport: %w", err)
|
||||
}
|
||||
return NewClient(sseTransport), nil
|
||||
}
|
||||
|
||||
// GetEndpoint returns the current endpoint URL for the SSE connection.
|
||||
//
|
||||
// Note: This method only works with SSE transport, or it will panic.
|
||||
func GetEndpoint(c *Client) *url.URL {
|
||||
t := c.GetTransport()
|
||||
sse := t.(*transport.SSE)
|
||||
return sse.GetEndpoint()
|
||||
}
|
||||
+58
@@ -0,0 +1,58 @@
|
||||
package client
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
|
||||
"github.com/mark3labs/mcp-go/client/transport"
|
||||
)
|
||||
|
||||
// NewStdioMCPClient creates a new stdio-based MCP client that communicates with a subprocess.
|
||||
// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
|
||||
// Returns an error if the subprocess cannot be started or the pipes cannot be created.
|
||||
//
|
||||
// NOTICE: NewStdioMCPClient will start the connection automatically.
|
||||
// This is for backward compatibility.
|
||||
func NewStdioMCPClient(
|
||||
command string,
|
||||
env []string,
|
||||
args ...string,
|
||||
) (*Client, error) {
|
||||
return NewStdioMCPClientWithOptions(command, env, args)
|
||||
}
|
||||
|
||||
// NewStdioMCPClientWithOptions creates a new stdio-based MCP client that communicates with a subprocess.
|
||||
// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
|
||||
// Optional configuration functions can be provided to customize the transport before it starts,
|
||||
// such as setting a custom command function.
|
||||
//
|
||||
// NOTICE: NewStdioMCPClientWithOptions automatically starts the underlying transport.
|
||||
// This is for backward compatibility.
|
||||
func NewStdioMCPClientWithOptions(
|
||||
command string,
|
||||
env []string,
|
||||
args []string,
|
||||
opts ...transport.StdioOption,
|
||||
) (*Client, error) {
|
||||
stdioTransport := transport.NewStdioWithOptions(command, env, args, opts...)
|
||||
|
||||
if err := stdioTransport.Start(context.Background()); err != nil {
|
||||
return nil, fmt.Errorf("failed to start stdio transport: %w", err)
|
||||
}
|
||||
|
||||
return NewClient(stdioTransport), nil
|
||||
}
|
||||
|
||||
// GetStderr returns a reader for the stderr output of the subprocess.
|
||||
// This can be used to capture error messages or logs from the subprocess.
|
||||
func GetStderr(c *Client) (io.Reader, bool) {
|
||||
t := c.GetTransport()
|
||||
|
||||
stdio, ok := t.(*transport.Stdio)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
return stdio.Stderr(), true
|
||||
}
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
package transport
|
||||
|
||||
// Common HTTP header constants used across transports
|
||||
const (
|
||||
HeaderKeySessionID = "Mcp-Session-Id"
|
||||
HeaderKeyProtocolVersion = "Mcp-Protocol-Version"
|
||||
)
|
||||
+22
@@ -0,0 +1,22 @@
|
||||
package transport
|
||||
|
||||
import "fmt"
|
||||
|
||||
// Error wraps a low-level transport error in a concrete type.
|
||||
type Error struct {
|
||||
Err error
|
||||
}
|
||||
|
||||
func (e *Error) Error() string {
|
||||
return fmt.Sprintf("transport error: %v", e.Err)
|
||||
}
|
||||
|
||||
func (e *Error) Unwrap() error {
|
||||
return e.Err
|
||||
}
|
||||
|
||||
func NewError(err error) *Error {
|
||||
return &Error{
|
||||
Err: err,
|
||||
}
|
||||
}
|
||||
+140
@@ -0,0 +1,140 @@
|
||||
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
|
||||
rootsHandler server.RootsHandler
|
||||
session *server.InProcessSession
|
||||
sessionID string
|
||||
|
||||
onNotification func(mcp.JSONRPCNotification)
|
||||
notifyMu sync.RWMutex
|
||||
started bool
|
||||
startedMu sync.Mutex
|
||||
}
|
||||
|
||||
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 WithRootsHandler(handler server.RootsHandler) InProcessOption {
|
||||
return func(t *InProcessTransport) {
|
||||
t.rootsHandler = 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 {
|
||||
c.startedMu.Lock()
|
||||
if c.started {
|
||||
c.startedMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
c.started = true
|
||||
c.startedMu.Unlock()
|
||||
|
||||
// Create and register session if we have handlers
|
||||
if c.samplingHandler != nil || c.elicitationHandler != nil || c.rootsHandler != nil {
|
||||
c.session = server.NewInProcessSessionWithHandlers(c.sessionID, c.samplingHandler, c.elicitationHandler, c.rootsHandler)
|
||||
if err := c.server.RegisterSession(ctx, c.session); err != nil {
|
||||
c.startedMu.Lock()
|
||||
c.started = false
|
||||
c.startedMu.Unlock()
|
||||
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)
|
||||
}
|
||||
var 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 ""
|
||||
}
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// HTTPHeaderFunc is a function that extracts header entries from the given context
|
||||
// and returns them as key-value pairs. This is typically used to add context values
|
||||
// as HTTP headers in outgoing requests.
|
||||
type HTTPHeaderFunc func(context.Context) map[string]string
|
||||
|
||||
// Interface for the transport layer.
|
||||
type Interface interface {
|
||||
// Start the connection. Start should only be called once.
|
||||
Start(ctx context.Context) error
|
||||
|
||||
// SendRequest sends a json RPC request and returns the response synchronously.
|
||||
SendRequest(ctx context.Context, request JSONRPCRequest) (*JSONRPCResponse, error)
|
||||
|
||||
// SendNotification sends a json RPC Notification to the server.
|
||||
SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error
|
||||
|
||||
// SetNotificationHandler sets the handler for notifications.
|
||||
// Any notification before the handler is set will be discarded.
|
||||
SetNotificationHandler(handler func(notification mcp.JSONRPCNotification))
|
||||
|
||||
// Close the connection.
|
||||
Close() error
|
||||
|
||||
// GetSessionId returns the session ID of the transport.
|
||||
GetSessionId() string
|
||||
}
|
||||
|
||||
// RequestHandler defines a function that handles incoming requests from the server.
|
||||
type RequestHandler func(ctx context.Context, request JSONRPCRequest) (*JSONRPCResponse, error)
|
||||
|
||||
// BidirectionalInterface extends Interface to support incoming requests from the server.
|
||||
// This is used for features like sampling where the server can send requests to the client.
|
||||
type BidirectionalInterface interface {
|
||||
Interface
|
||||
|
||||
// SetRequestHandler sets the handler for incoming requests from the server.
|
||||
// The handler should process the request and return a response.
|
||||
SetRequestHandler(handler RequestHandler)
|
||||
}
|
||||
|
||||
// HTTPConnection is a Transport that runs over HTTP and supports
|
||||
// protocol version headers.
|
||||
type HTTPConnection interface {
|
||||
Interface
|
||||
SetProtocolVersion(version string)
|
||||
}
|
||||
|
||||
type JSONRPCRequest struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID mcp.RequestId `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params any `json:"params,omitempty"`
|
||||
Header http.Header `json:"-"`
|
||||
}
|
||||
|
||||
// JSONRPCResponse represents a JSON-RPC 2.0 response message.
|
||||
// Use NewJSONRPCResultResponse to create a JSONRPCResponse with a result.
|
||||
// Use NewJSONRPCErrorResponse to create a JSONRPCResponse with an error.
|
||||
type JSONRPCResponse struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID mcp.RequestId `json:"id"`
|
||||
Result json.RawMessage `json:"result,omitempty"`
|
||||
Error *mcp.JSONRPCErrorDetails `json:"error,omitempty"`
|
||||
}
|
||||
+712
@@ -0,0 +1,712 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrNoToken is returned when no token is available in the token store
|
||||
var ErrNoToken = errors.New("no token available")
|
||||
|
||||
// OAuthConfig holds the OAuth configuration for the client
|
||||
type OAuthConfig struct {
|
||||
// ClientID is the OAuth client ID
|
||||
ClientID string
|
||||
// ClientSecret is the OAuth client secret (for confidential clients)
|
||||
ClientSecret string
|
||||
// RedirectURI is the redirect URI for the OAuth flow
|
||||
RedirectURI string
|
||||
// Scopes is the list of OAuth scopes to request
|
||||
Scopes []string
|
||||
// TokenStore is the storage for OAuth tokens
|
||||
TokenStore TokenStore
|
||||
// AuthServerMetadataURL is the URL to the OAuth server metadata
|
||||
// If empty, the client will attempt to discover it from the base URL
|
||||
AuthServerMetadataURL string
|
||||
// PKCEEnabled enables PKCE for the OAuth flow (recommended for public clients)
|
||||
PKCEEnabled bool
|
||||
// HTTPClient is an optional HTTP client to use for requests.
|
||||
// If nil, a default HTTP client with a 30 second timeout will be used.
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
// TokenStore is an interface for storing and retrieving OAuth tokens.
|
||||
//
|
||||
// Implementations must:
|
||||
// - Honor context cancellation and deadlines, returning context.Canceled
|
||||
// or context.DeadlineExceeded as appropriate
|
||||
// - Return ErrNoToken (or a sentinel error that wraps it) when no token
|
||||
// is available, rather than conflating this with other operational errors
|
||||
// - Properly propagate all other errors (database failures, I/O errors, etc.)
|
||||
// - Check ctx.Done() before performing operations and return ctx.Err() if cancelled
|
||||
type TokenStore interface {
|
||||
// GetToken returns the current token.
|
||||
// Returns ErrNoToken if no token is available.
|
||||
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
|
||||
// Returns other errors for operational failures (I/O, database, etc.).
|
||||
GetToken(ctx context.Context) (*Token, error)
|
||||
|
||||
// SaveToken saves a token.
|
||||
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
|
||||
// Returns other errors for operational failures (I/O, database, etc.).
|
||||
SaveToken(ctx context.Context, token *Token) error
|
||||
}
|
||||
|
||||
// Token represents an OAuth token
|
||||
type Token struct {
|
||||
// AccessToken is the OAuth access token
|
||||
AccessToken string `json:"access_token"`
|
||||
// TokenType is the type of token (usually "Bearer")
|
||||
TokenType string `json:"token_type"`
|
||||
// RefreshToken is the OAuth refresh token
|
||||
RefreshToken string `json:"refresh_token,omitempty"`
|
||||
// ExpiresIn is the number of seconds until the token expires
|
||||
ExpiresIn int64 `json:"expires_in,omitempty"`
|
||||
// Scope is the scope of the token
|
||||
Scope string `json:"scope,omitempty"`
|
||||
// ExpiresAt is the time when the token expires
|
||||
ExpiresAt time.Time `json:"expires_at,omitempty"`
|
||||
}
|
||||
|
||||
// IsExpired returns true if the token is expired
|
||||
func (t *Token) IsExpired() bool {
|
||||
if t.ExpiresAt.IsZero() {
|
||||
return false
|
||||
}
|
||||
return time.Now().After(t.ExpiresAt)
|
||||
}
|
||||
|
||||
// MemoryTokenStore is a simple in-memory token store
|
||||
type MemoryTokenStore struct {
|
||||
token *Token
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewMemoryTokenStore creates a new in-memory token store
|
||||
func NewMemoryTokenStore() *MemoryTokenStore {
|
||||
return &MemoryTokenStore{}
|
||||
}
|
||||
|
||||
// GetToken returns the current token.
|
||||
// Returns ErrNoToken if no token is available.
|
||||
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
|
||||
func (s *MemoryTokenStore) GetToken(ctx context.Context) (*Token, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
if s.token == nil {
|
||||
return nil, ErrNoToken
|
||||
}
|
||||
return s.token, nil
|
||||
}
|
||||
|
||||
// SaveToken saves a token.
|
||||
// Returns context.Canceled or context.DeadlineExceeded if ctx is cancelled.
|
||||
func (s *MemoryTokenStore) SaveToken(ctx context.Context, token *Token) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.token = token
|
||||
return nil
|
||||
}
|
||||
|
||||
// AuthServerMetadata represents the OAuth 2.0 Authorization Server Metadata
|
||||
type AuthServerMetadata struct {
|
||||
Issuer string `json:"issuer"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
RegistrationEndpoint string `json:"registration_endpoint,omitempty"`
|
||||
JwksURI string `json:"jwks_uri,omitempty"`
|
||||
ScopesSupported []string `json:"scopes_supported,omitempty"`
|
||||
ResponseTypesSupported []string `json:"response_types_supported"`
|
||||
GrantTypesSupported []string `json:"grant_types_supported,omitempty"`
|
||||
TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported,omitempty"`
|
||||
}
|
||||
|
||||
// OAuthHandler handles OAuth authentication for HTTP requests
|
||||
type OAuthHandler struct {
|
||||
config OAuthConfig
|
||||
httpClient *http.Client
|
||||
serverMetadata *AuthServerMetadata
|
||||
metadataFetchErr error
|
||||
metadataOnce sync.Once
|
||||
baseURL string
|
||||
|
||||
mu sync.RWMutex // Protects expectedState
|
||||
expectedState string // Expected state value for CSRF protection
|
||||
}
|
||||
|
||||
// NewOAuthHandler creates a new OAuth handler
|
||||
func NewOAuthHandler(config OAuthConfig) *OAuthHandler {
|
||||
if config.TokenStore == nil {
|
||||
config.TokenStore = NewMemoryTokenStore()
|
||||
}
|
||||
if config.HTTPClient == nil {
|
||||
config.HTTPClient = &http.Client{Timeout: 30 * time.Second}
|
||||
}
|
||||
|
||||
return &OAuthHandler{
|
||||
config: config,
|
||||
httpClient: config.HTTPClient,
|
||||
}
|
||||
}
|
||||
|
||||
// GetAuthorizationHeader returns the Authorization header value for a request
|
||||
func (h *OAuthHandler) GetAuthorizationHeader(ctx context.Context) (string, error) {
|
||||
token, err := h.getValidToken(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Some auth implementations are strict about token type
|
||||
tokenType := token.TokenType
|
||||
if tokenType == "bearer" {
|
||||
tokenType = "Bearer"
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s %s", tokenType, token.AccessToken), nil
|
||||
}
|
||||
|
||||
// getValidToken returns a valid token, refreshing if necessary
|
||||
func (h *OAuthHandler) getValidToken(ctx context.Context) (*Token, error) {
|
||||
token, err := h.config.TokenStore.GetToken(ctx)
|
||||
if err != nil && !errors.Is(err, ErrNoToken) {
|
||||
return nil, err
|
||||
}
|
||||
if err == nil && !token.IsExpired() && token.AccessToken != "" {
|
||||
return token, nil
|
||||
}
|
||||
|
||||
// If we have a refresh token, try to use it
|
||||
if err == nil && token.RefreshToken != "" {
|
||||
newToken, err := h.refreshToken(ctx, token.RefreshToken)
|
||||
if err == nil {
|
||||
return newToken, nil
|
||||
}
|
||||
// If refresh fails, continue to authorization flow
|
||||
}
|
||||
|
||||
// We need to get a new token through the authorization flow
|
||||
return nil, ErrOAuthAuthorizationRequired
|
||||
}
|
||||
|
||||
// refreshToken refreshes an OAuth token
|
||||
func (h *OAuthHandler) refreshToken(ctx context.Context, refreshToken string) (*Token, error) {
|
||||
metadata, err := h.getServerMetadata(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get server metadata: %w", err)
|
||||
}
|
||||
|
||||
data := url.Values{}
|
||||
data.Set("grant_type", "refresh_token")
|
||||
data.Set("refresh_token", refreshToken)
|
||||
data.Set("client_id", h.config.ClientID)
|
||||
if h.config.ClientSecret != "" {
|
||||
data.Set("client_secret", h.config.ClientSecret)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(
|
||||
ctx,
|
||||
http.MethodPost,
|
||||
metadata.TokenEndpoint,
|
||||
strings.NewReader(data.Encode()),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create refresh token request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := h.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to send refresh token request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, extractOAuthError(body, resp.StatusCode, "refresh token request failed")
|
||||
}
|
||||
|
||||
var tokenResp Token
|
||||
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to decode token response: %w", err)
|
||||
}
|
||||
|
||||
// Set expiration time
|
||||
if tokenResp.ExpiresIn > 0 {
|
||||
tokenResp.ExpiresAt = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
|
||||
}
|
||||
|
||||
// If no new refresh token is provided, keep the old one
|
||||
if tokenResp.RefreshToken == "" {
|
||||
tokenResp.RefreshToken = refreshToken
|
||||
}
|
||||
|
||||
// Save the token
|
||||
if err := h.config.TokenStore.SaveToken(ctx, &tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to save token: %w", err)
|
||||
}
|
||||
|
||||
return &tokenResp, nil
|
||||
}
|
||||
|
||||
// RefreshToken is a public wrapper for refreshToken
|
||||
func (h *OAuthHandler) RefreshToken(ctx context.Context, refreshToken string) (*Token, error) {
|
||||
return h.refreshToken(ctx, refreshToken)
|
||||
}
|
||||
|
||||
// GetClientID returns the client ID
|
||||
func (h *OAuthHandler) GetClientID() string {
|
||||
return h.config.ClientID
|
||||
}
|
||||
|
||||
// extractOAuthError attempts to parse an OAuth error response from the response body
|
||||
func extractOAuthError(body []byte, statusCode int, context string) error {
|
||||
// Try to parse the error as an OAuth error response
|
||||
var oauthErr OAuthError
|
||||
if err := json.Unmarshal(body, &oauthErr); err == nil && oauthErr.ErrorCode != "" {
|
||||
return fmt.Errorf("%s: %w", context, oauthErr)
|
||||
}
|
||||
|
||||
// If not a valid OAuth error, return the raw response
|
||||
return fmt.Errorf("%s with status %d: %s", context, statusCode, body)
|
||||
}
|
||||
|
||||
// GetClientSecret returns the client secret
|
||||
func (h *OAuthHandler) GetClientSecret() string {
|
||||
return h.config.ClientSecret
|
||||
}
|
||||
|
||||
// SetBaseURL sets the base URL for the API server
|
||||
func (h *OAuthHandler) SetBaseURL(baseURL string) {
|
||||
h.baseURL = baseURL
|
||||
}
|
||||
|
||||
// GetExpectedState returns the expected state value (for testing purposes)
|
||||
func (h *OAuthHandler) GetExpectedState() string {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return h.expectedState
|
||||
}
|
||||
|
||||
// SetExpectedState sets the expected state value.
|
||||
//
|
||||
// This can be useful if you cannot maintain an OAuthHandler
|
||||
// instance throughout the authentication flow; for example, if
|
||||
// the initialization and callback steps are handled in different
|
||||
// requests.
|
||||
//
|
||||
// In such cases, this should be called with the state value generated
|
||||
// during the initial authentication request (e.g. by GenerateState)
|
||||
// and included in the authorization URL.
|
||||
func (h *OAuthHandler) SetExpectedState(expectedState string) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.expectedState = expectedState
|
||||
}
|
||||
|
||||
// OAuthError represents a standard OAuth 2.0 error response
|
||||
type OAuthError struct {
|
||||
ErrorCode string `json:"error"`
|
||||
ErrorDescription string `json:"error_description,omitempty"`
|
||||
ErrorURI string `json:"error_uri,omitempty"`
|
||||
}
|
||||
|
||||
// Error implements the error interface
|
||||
func (e OAuthError) Error() string {
|
||||
if e.ErrorDescription != "" {
|
||||
return fmt.Sprintf("OAuth error: %s - %s", e.ErrorCode, e.ErrorDescription)
|
||||
}
|
||||
return fmt.Sprintf("OAuth error: %s", e.ErrorCode)
|
||||
}
|
||||
|
||||
// OAuthProtectedResource represents the response from /.well-known/oauth-protected-resource
|
||||
type OAuthProtectedResource struct {
|
||||
AuthorizationServers []string `json:"authorization_servers"`
|
||||
Resource string `json:"resource"`
|
||||
ResourceName string `json:"resource_name,omitempty"`
|
||||
}
|
||||
|
||||
// getServerMetadata fetches the OAuth server metadata
|
||||
func (h *OAuthHandler) getServerMetadata(ctx context.Context) (*AuthServerMetadata, error) {
|
||||
h.metadataOnce.Do(func() {
|
||||
// If AuthServerMetadataURL is explicitly provided, use it directly
|
||||
if h.config.AuthServerMetadataURL != "" {
|
||||
h.fetchMetadataFromURL(ctx, h.config.AuthServerMetadataURL)
|
||||
return
|
||||
}
|
||||
|
||||
// Try to discover the authorization server via OAuth Protected Resource
|
||||
// as per RFC 9728 (https://datatracker.ietf.org/doc/html/rfc9728)
|
||||
baseURL, err := h.extractBaseURL()
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to extract base URL: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Try to fetch the OAuth Protected Resource metadata
|
||||
protectedResourceURL := baseURL + "/.well-known/oauth-protected-resource"
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, protectedResourceURL, nil)
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to create protected resource request: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("MCP-Protocol-Version", "2025-03-26")
|
||||
|
||||
resp, err := h.httpClient.Do(req)
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to send protected resource request: %w", err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// If we can't get the protected resource metadata, try OAuth Authorization Server discovery
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
h.fetchMetadataFromURL(ctx, baseURL+"/.well-known/oauth-authorization-server")
|
||||
if h.serverMetadata != nil {
|
||||
return
|
||||
}
|
||||
// If that also fails, fall back to default endpoints
|
||||
metadata, err := h.getDefaultEndpoints(baseURL)
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to get default endpoints: %w", err)
|
||||
return
|
||||
}
|
||||
h.serverMetadata = metadata
|
||||
return
|
||||
}
|
||||
|
||||
// Parse the protected resource metadata
|
||||
var protectedResource OAuthProtectedResource
|
||||
if err := json.NewDecoder(resp.Body).Decode(&protectedResource); err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to decode protected resource response: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
// If no authorization servers are specified, fall back to default endpoints
|
||||
if len(protectedResource.AuthorizationServers) == 0 {
|
||||
metadata, err := h.getDefaultEndpoints(baseURL)
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to get default endpoints: %w", err)
|
||||
return
|
||||
}
|
||||
h.serverMetadata = metadata
|
||||
return
|
||||
}
|
||||
|
||||
// Use the first authorization server
|
||||
authServerURL := protectedResource.AuthorizationServers[0]
|
||||
|
||||
// Try OpenID Connect discovery first
|
||||
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/openid-configuration")
|
||||
if h.serverMetadata != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// If OpenID Connect discovery fails, try OAuth Authorization Server Metadata
|
||||
h.fetchMetadataFromURL(ctx, authServerURL+"/.well-known/oauth-authorization-server")
|
||||
if h.serverMetadata != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// If both discovery methods fail, use default endpoints based on the authorization server URL
|
||||
metadata, err := h.getDefaultEndpoints(authServerURL)
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to get default endpoints: %w", err)
|
||||
return
|
||||
}
|
||||
h.serverMetadata = metadata
|
||||
})
|
||||
|
||||
if h.metadataFetchErr != nil {
|
||||
return nil, h.metadataFetchErr
|
||||
}
|
||||
|
||||
return h.serverMetadata, nil
|
||||
}
|
||||
|
||||
// fetchMetadataFromURL fetches and parses OAuth server metadata from a URL
|
||||
func (h *OAuthHandler) fetchMetadataFromURL(ctx context.Context, metadataURL string) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, metadataURL, nil)
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to create metadata request: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("MCP-Protocol-Version", "2025-03-26")
|
||||
|
||||
resp, err := h.httpClient.Do(req)
|
||||
if err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to send metadata request: %w", err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
// If metadata discovery fails, don't set any metadata
|
||||
return
|
||||
}
|
||||
|
||||
var metadata AuthServerMetadata
|
||||
if err := json.NewDecoder(resp.Body).Decode(&metadata); err != nil {
|
||||
h.metadataFetchErr = fmt.Errorf("failed to decode metadata response: %w", err)
|
||||
return
|
||||
}
|
||||
|
||||
h.serverMetadata = &metadata
|
||||
}
|
||||
|
||||
// extractBaseURL extracts the base URL from the first request
|
||||
func (h *OAuthHandler) extractBaseURL() (string, error) {
|
||||
// If we have a base URL from a previous request, use it
|
||||
if h.baseURL != "" {
|
||||
return h.baseURL, nil
|
||||
}
|
||||
|
||||
// Otherwise, we need to infer it from the redirect URI
|
||||
if h.config.RedirectURI == "" {
|
||||
return "", fmt.Errorf("no base URL available and no redirect URI provided")
|
||||
}
|
||||
|
||||
// Parse the redirect URI to extract the authority
|
||||
parsedURL, err := url.Parse(h.config.RedirectURI)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to parse redirect URI: %w", err)
|
||||
}
|
||||
|
||||
// Use the scheme and host from the redirect URI
|
||||
baseURL := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host)
|
||||
return baseURL, nil
|
||||
}
|
||||
|
||||
// GetServerMetadata is a public wrapper for getServerMetadata
|
||||
func (h *OAuthHandler) GetServerMetadata(ctx context.Context) (*AuthServerMetadata, error) {
|
||||
return h.getServerMetadata(ctx)
|
||||
}
|
||||
|
||||
// getDefaultEndpoints returns default OAuth endpoints based on the base URL
|
||||
func (h *OAuthHandler) getDefaultEndpoints(baseURL string) (*AuthServerMetadata, error) {
|
||||
// Parse the base URL to extract the authority
|
||||
parsedURL, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse base URL: %w", err)
|
||||
}
|
||||
|
||||
// Discard any path component to get the authorization base URL
|
||||
parsedURL.Path = ""
|
||||
authBaseURL := parsedURL.String()
|
||||
|
||||
// Validate that the URL has a scheme and host
|
||||
if parsedURL.Scheme == "" || parsedURL.Host == "" {
|
||||
return nil, fmt.Errorf("invalid base URL: missing scheme or host in %q", baseURL)
|
||||
}
|
||||
|
||||
return &AuthServerMetadata{
|
||||
Issuer: authBaseURL,
|
||||
AuthorizationEndpoint: authBaseURL + "/authorize",
|
||||
TokenEndpoint: authBaseURL + "/token",
|
||||
RegistrationEndpoint: authBaseURL + "/register",
|
||||
}, nil
|
||||
}
|
||||
|
||||
// RegisterClient performs dynamic client registration
|
||||
func (h *OAuthHandler) RegisterClient(ctx context.Context, clientName string) error {
|
||||
metadata, err := h.getServerMetadata(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get server metadata: %w", err)
|
||||
}
|
||||
|
||||
if metadata.RegistrationEndpoint == "" {
|
||||
return errors.New("server does not support dynamic client registration")
|
||||
}
|
||||
|
||||
// Prepare registration request
|
||||
regRequest := map[string]any{
|
||||
"client_name": clientName,
|
||||
"redirect_uris": []string{h.config.RedirectURI},
|
||||
"token_endpoint_auth_method": "none", // For public clients
|
||||
"grant_types": []string{"authorization_code", "refresh_token"},
|
||||
"response_types": []string{"code"},
|
||||
"scope": strings.Join(h.config.Scopes, " "),
|
||||
}
|
||||
|
||||
// Add client_secret if this is a confidential client
|
||||
if h.config.ClientSecret != "" {
|
||||
regRequest["token_endpoint_auth_method"] = "client_secret_basic"
|
||||
}
|
||||
|
||||
reqBody, err := json.Marshal(regRequest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal registration request: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(
|
||||
ctx,
|
||||
http.MethodPost,
|
||||
metadata.RegistrationEndpoint,
|
||||
bytes.NewReader(reqBody),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create registration request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := h.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to send registration request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return extractOAuthError(body, resp.StatusCode, "registration request failed")
|
||||
}
|
||||
|
||||
var regResponse struct {
|
||||
ClientID string `json:"client_id"`
|
||||
ClientSecret string `json:"client_secret,omitempty"`
|
||||
}
|
||||
|
||||
if err := json.NewDecoder(resp.Body).Decode(®Response); err != nil {
|
||||
return fmt.Errorf("failed to decode registration response: %w", err)
|
||||
}
|
||||
|
||||
// Update the client configuration
|
||||
h.config.ClientID = regResponse.ClientID
|
||||
if regResponse.ClientSecret != "" {
|
||||
h.config.ClientSecret = regResponse.ClientSecret
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ErrInvalidState is returned when the state parameter doesn't match the expected value
|
||||
var ErrInvalidState = errors.New("invalid state parameter, possible CSRF attack")
|
||||
|
||||
// ProcessAuthorizationResponse processes the authorization response and exchanges the code for a token
|
||||
func (h *OAuthHandler) ProcessAuthorizationResponse(ctx context.Context, code, state, codeVerifier string) error {
|
||||
// Validate the state parameter to prevent CSRF attacks
|
||||
h.mu.Lock()
|
||||
expectedState := h.expectedState
|
||||
if expectedState == "" {
|
||||
h.mu.Unlock()
|
||||
return errors.New("no expected state found, authorization flow may not have been initiated properly")
|
||||
}
|
||||
|
||||
if state != expectedState {
|
||||
h.mu.Unlock()
|
||||
return ErrInvalidState
|
||||
}
|
||||
|
||||
// Clear the expected state after validation
|
||||
h.expectedState = ""
|
||||
h.mu.Unlock()
|
||||
|
||||
metadata, err := h.getServerMetadata(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get server metadata: %w", err)
|
||||
}
|
||||
|
||||
data := url.Values{}
|
||||
data.Set("grant_type", "authorization_code")
|
||||
data.Set("code", code)
|
||||
data.Set("client_id", h.config.ClientID)
|
||||
data.Set("redirect_uri", h.config.RedirectURI)
|
||||
|
||||
if h.config.ClientSecret != "" {
|
||||
data.Set("client_secret", h.config.ClientSecret)
|
||||
}
|
||||
|
||||
if h.config.PKCEEnabled && codeVerifier != "" {
|
||||
data.Set("code_verifier", codeVerifier)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(
|
||||
ctx,
|
||||
http.MethodPost,
|
||||
metadata.TokenEndpoint,
|
||||
strings.NewReader(data.Encode()),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create token request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := h.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to send token request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return extractOAuthError(body, resp.StatusCode, "token request failed")
|
||||
}
|
||||
|
||||
var tokenResp Token
|
||||
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
|
||||
return fmt.Errorf("failed to decode token response: %w", err)
|
||||
}
|
||||
|
||||
// Set expiration time
|
||||
if tokenResp.ExpiresIn > 0 {
|
||||
tokenResp.ExpiresAt = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
|
||||
}
|
||||
|
||||
// Save the token
|
||||
if err := h.config.TokenStore.SaveToken(ctx, &tokenResp); err != nil {
|
||||
return fmt.Errorf("failed to save token: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAuthorizationURL returns the URL for the authorization endpoint
|
||||
func (h *OAuthHandler) GetAuthorizationURL(ctx context.Context, state, codeChallenge string) (string, error) {
|
||||
metadata, err := h.getServerMetadata(ctx)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to get server metadata: %w", err)
|
||||
}
|
||||
|
||||
// Store the state for later validation
|
||||
h.SetExpectedState(state)
|
||||
|
||||
params := url.Values{}
|
||||
params.Set("response_type", "code")
|
||||
params.Set("client_id", h.config.ClientID)
|
||||
params.Set("redirect_uri", h.config.RedirectURI)
|
||||
params.Set("state", state)
|
||||
|
||||
if len(h.config.Scopes) > 0 {
|
||||
params.Set("scope", strings.Join(h.config.Scopes, " "))
|
||||
}
|
||||
|
||||
if h.config.PKCEEnabled && codeChallenge != "" {
|
||||
params.Set("code_challenge", codeChallenge)
|
||||
params.Set("code_challenge_method", "S256")
|
||||
}
|
||||
|
||||
return metadata.AuthorizationEndpoint + "?" + params.Encode(), nil
|
||||
}
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
// GenerateRandomString generates a random string of the specified length
|
||||
func GenerateRandomString(length int) (string, error) {
|
||||
bytes := make([]byte, length)
|
||||
if _, err := rand.Read(bytes); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(bytes)[:length], nil
|
||||
}
|
||||
|
||||
// GenerateCodeVerifier generates a code verifier for PKCE
|
||||
func GenerateCodeVerifier() (string, error) {
|
||||
// According to RFC 7636, the code verifier should be between 43 and 128 characters
|
||||
return GenerateRandomString(64)
|
||||
}
|
||||
|
||||
// GenerateCodeChallenge generates a code challenge from a code verifier
|
||||
func GenerateCodeChallenge(codeVerifier string) string {
|
||||
// SHA256 hash the code verifier
|
||||
hash := sha256.Sum256([]byte(codeVerifier))
|
||||
// Base64url encode the hash
|
||||
return base64.RawURLEncoding.EncodeToString(hash[:])
|
||||
}
|
||||
|
||||
// GenerateState generates a state parameter for OAuth
|
||||
func GenerateState() (string, error) {
|
||||
return GenerateRandomString(32)
|
||||
}
|
||||
|
||||
// ValidateRedirectURI validates that a redirect URI is secure
|
||||
func ValidateRedirectURI(redirectURI string) error {
|
||||
// According to the spec, redirect URIs must be either localhost URLs or HTTPS URLs
|
||||
if redirectURI == "" {
|
||||
return fmt.Errorf("redirect URI cannot be empty")
|
||||
}
|
||||
|
||||
// Parse the URL
|
||||
parsedURL, err := url.Parse(redirectURI)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid redirect URI: %w", err)
|
||||
}
|
||||
|
||||
// Check if it's a localhost URL
|
||||
if parsedURL.Scheme == "http" {
|
||||
hostname := parsedURL.Hostname()
|
||||
// Check for various forms of localhost
|
||||
if hostname == "localhost" || hostname == "127.0.0.1" || hostname == "::1" || hostname == "[::1]" {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("HTTP redirect URI must use localhost or 127.0.0.1")
|
||||
}
|
||||
|
||||
// Check if it's an HTTPS URL
|
||||
if parsedURL.Scheme == "https" {
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("redirect URI must use either HTTP with localhost or HTTPS")
|
||||
}
|
||||
+580
@@ -0,0 +1,580 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/mark3labs/mcp-go/util"
|
||||
)
|
||||
|
||||
// SSE implements the transport layer of the MCP protocol using Server-Sent Events (SSE).
|
||||
// It maintains a persistent HTTP connection to receive server-pushed events
|
||||
// while sending requests over regular HTTP POST calls. The client handles
|
||||
// automatic reconnection and message routing between requests and responses.
|
||||
type SSE struct {
|
||||
baseURL *url.URL
|
||||
endpoint *url.URL
|
||||
httpClient *http.Client
|
||||
responses map[string]chan *JSONRPCResponse
|
||||
mu sync.RWMutex
|
||||
onNotification func(mcp.JSONRPCNotification)
|
||||
notifyMu sync.RWMutex
|
||||
endpointChan chan struct{}
|
||||
headers map[string]string
|
||||
headerFunc HTTPHeaderFunc
|
||||
logger util.Logger
|
||||
|
||||
started atomic.Bool
|
||||
closed atomic.Bool
|
||||
cancelSSEStream context.CancelFunc
|
||||
protocolVersion atomic.Value // string
|
||||
onConnectionLost func(error)
|
||||
connectionLostMu sync.RWMutex
|
||||
|
||||
// OAuth support
|
||||
oauthHandler *OAuthHandler
|
||||
}
|
||||
|
||||
type ClientOption func(*SSE)
|
||||
|
||||
// WithSSELogger sets a custom logger for the SSE client.
|
||||
func WithSSELogger(logger util.Logger) ClientOption {
|
||||
return func(sc *SSE) {
|
||||
sc.logger = logger
|
||||
}
|
||||
}
|
||||
|
||||
func WithHeaders(headers map[string]string) ClientOption {
|
||||
return func(sc *SSE) {
|
||||
sc.headers = headers
|
||||
}
|
||||
}
|
||||
|
||||
func WithHeaderFunc(headerFunc HTTPHeaderFunc) ClientOption {
|
||||
return func(sc *SSE) {
|
||||
sc.headerFunc = headerFunc
|
||||
}
|
||||
}
|
||||
|
||||
func WithHTTPClient(httpClient *http.Client) ClientOption {
|
||||
return func(sc *SSE) {
|
||||
sc.httpClient = httpClient
|
||||
}
|
||||
}
|
||||
|
||||
func WithOAuth(config OAuthConfig) ClientOption {
|
||||
return func(sc *SSE) {
|
||||
sc.oauthHandler = NewOAuthHandler(config)
|
||||
}
|
||||
}
|
||||
|
||||
// NewSSE creates a new SSE-based MCP client with the given base URL.
|
||||
// Returns an error if the URL is invalid.
|
||||
func NewSSE(baseURL string, options ...ClientOption) (*SSE, error) {
|
||||
parsedURL, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid URL: %w", err)
|
||||
}
|
||||
|
||||
smc := &SSE{
|
||||
baseURL: parsedURL,
|
||||
httpClient: &http.Client{},
|
||||
responses: make(map[string]chan *JSONRPCResponse),
|
||||
endpointChan: make(chan struct{}),
|
||||
headers: make(map[string]string),
|
||||
logger: util.DefaultLogger(),
|
||||
}
|
||||
|
||||
for _, opt := range options {
|
||||
opt(smc)
|
||||
}
|
||||
|
||||
// If OAuth is configured, set the base URL for metadata discovery
|
||||
if smc.oauthHandler != nil {
|
||||
// Extract base URL from server URL for metadata discovery
|
||||
baseURL := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host)
|
||||
smc.oauthHandler.SetBaseURL(baseURL)
|
||||
}
|
||||
|
||||
return smc, nil
|
||||
}
|
||||
|
||||
// Start initiates the SSE connection to the server and waits for the endpoint information.
|
||||
// Returns an error if the connection fails or times out waiting for the endpoint.
|
||||
func (c *SSE) Start(ctx context.Context) error {
|
||||
if c.started.Load() {
|
||||
return nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
c.cancelSSEStream = cancel
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", c.baseURL.String(), nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
req.Header.Set("Cache-Control", "no-cache")
|
||||
req.Header.Set("Connection", "keep-alive")
|
||||
|
||||
// set custom http headers
|
||||
for k, v := range c.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
if c.headerFunc != nil {
|
||||
for k, v := range c.headerFunc(ctx) {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
}
|
||||
|
||||
// Add OAuth authorization if configured
|
||||
if c.oauthHandler != nil {
|
||||
authHeader, err := c.oauthHandler.GetAuthorizationHeader(ctx)
|
||||
if err != nil {
|
||||
// If we get an authorization error, return a specific error that can be handled by the client
|
||||
if err.Error() == "no valid token available, authorization required" {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("failed to get authorization header: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", authHeader)
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect to SSE stream: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
resp.Body.Close()
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
go c.readSSE(resp.Body)
|
||||
|
||||
// Wait for the endpoint to be received
|
||||
timeout := time.NewTimer(30 * time.Second)
|
||||
defer timeout.Stop()
|
||||
select {
|
||||
case <-c.endpointChan:
|
||||
// Endpoint received, proceed
|
||||
case <-ctx.Done():
|
||||
return fmt.Errorf("context cancelled while waiting for endpoint")
|
||||
case <-timeout.C: // Add a timeout
|
||||
cancel()
|
||||
return fmt.Errorf("timeout waiting for endpoint")
|
||||
}
|
||||
|
||||
c.started.Store(true)
|
||||
return nil
|
||||
}
|
||||
|
||||
// readSSE continuously reads the SSE stream and processes events.
|
||||
// It runs until the connection is closed or an error occurs.
|
||||
func (c *SSE) readSSE(reader io.ReadCloser) {
|
||||
defer reader.Close()
|
||||
|
||||
br := bufio.NewReader(reader)
|
||||
var event, data string
|
||||
|
||||
for {
|
||||
// when close or start's ctx cancel, the reader will be closed
|
||||
// and the for loop will break.
|
||||
line, err := br.ReadString('\n')
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
// Process any pending event before exit
|
||||
if data != "" {
|
||||
// If no event type is specified, use empty string (default event type)
|
||||
if event == "" {
|
||||
event = "message"
|
||||
}
|
||||
c.handleSSEEvent(event, data)
|
||||
}
|
||||
break
|
||||
}
|
||||
// Checking whether the connection was terminated due to NO_ERROR in HTTP2 based on RFC9113
|
||||
// Only handle NO_ERROR specially if onConnectionLost handler is set to maintain backward compatibility
|
||||
if strings.Contains(err.Error(), "NO_ERROR") {
|
||||
c.connectionLostMu.RLock()
|
||||
handler := c.onConnectionLost
|
||||
c.connectionLostMu.RUnlock()
|
||||
|
||||
if handler != nil {
|
||||
// This is not actually an error - HTTP2 idle timeout disconnection
|
||||
handler(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if !c.closed.Load() {
|
||||
c.logger.Errorf("SSE stream error: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Remove only newline markers
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
if line == "" {
|
||||
// Empty line means end of event
|
||||
if data != "" {
|
||||
// If no event type is specified, use empty string (default event type)
|
||||
if event == "" {
|
||||
event = "message"
|
||||
}
|
||||
c.handleSSEEvent(event, data)
|
||||
event = ""
|
||||
data = ""
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(line, "event:") {
|
||||
event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
|
||||
} else if strings.HasPrefix(line, "data:") {
|
||||
data = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleSSEEvent processes SSE events based on their type.
|
||||
// Handles 'endpoint' events for connection setup and 'message' events for JSON-RPC communication.
|
||||
func (c *SSE) handleSSEEvent(event, data string) {
|
||||
switch event {
|
||||
case "endpoint":
|
||||
endpoint, err := c.baseURL.Parse(data)
|
||||
if err != nil {
|
||||
c.logger.Errorf("Error parsing endpoint URL: %v", err)
|
||||
return
|
||||
}
|
||||
if endpoint.Host != c.baseURL.Host {
|
||||
c.logger.Errorf("Endpoint origin does not match connection origin")
|
||||
return
|
||||
}
|
||||
c.endpoint = endpoint
|
||||
close(c.endpointChan)
|
||||
|
||||
case "message":
|
||||
var baseMessage JSONRPCResponse
|
||||
if err := json.Unmarshal([]byte(data), &baseMessage); err != nil {
|
||||
c.logger.Errorf("Error unmarshaling message: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Handle notification
|
||||
if baseMessage.ID.IsNil() {
|
||||
var notification mcp.JSONRPCNotification
|
||||
if err := json.Unmarshal([]byte(data), ¬ification); err != nil {
|
||||
return
|
||||
}
|
||||
c.notifyMu.RLock()
|
||||
if c.onNotification != nil {
|
||||
c.onNotification(notification)
|
||||
}
|
||||
c.notifyMu.RUnlock()
|
||||
return
|
||||
}
|
||||
|
||||
// Create string key for map lookup
|
||||
idKey := baseMessage.ID.String()
|
||||
|
||||
c.mu.RLock()
|
||||
ch, exists := c.responses[idKey]
|
||||
c.mu.RUnlock()
|
||||
|
||||
if exists {
|
||||
ch <- &baseMessage
|
||||
c.mu.Lock()
|
||||
delete(c.responses, idKey)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *SSE) SetNotificationHandler(handler func(notification mcp.JSONRPCNotification)) {
|
||||
c.notifyMu.Lock()
|
||||
defer c.notifyMu.Unlock()
|
||||
c.onNotification = handler
|
||||
}
|
||||
|
||||
func (c *SSE) SetConnectionLostHandler(handler func(error)) {
|
||||
c.connectionLostMu.Lock()
|
||||
defer c.connectionLostMu.Unlock()
|
||||
c.onConnectionLost = handler
|
||||
}
|
||||
|
||||
// SendRequest sends a JSON-RPC request to the server and waits for a response.
|
||||
// Returns the raw JSON response message or an error if the request fails.
|
||||
func (c *SSE) SendRequest(
|
||||
ctx context.Context,
|
||||
request JSONRPCRequest,
|
||||
) (*JSONRPCResponse, error) {
|
||||
if !c.started.Load() {
|
||||
return nil, fmt.Errorf("transport not started yet")
|
||||
}
|
||||
if c.closed.Load() {
|
||||
return nil, fmt.Errorf("transport has been closed")
|
||||
}
|
||||
if c.endpoint == nil {
|
||||
return nil, fmt.Errorf("endpoint not received")
|
||||
}
|
||||
|
||||
// Marshal request
|
||||
requestBytes, err := json.Marshal(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
// Create HTTP request
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.endpoint.String(), bytes.NewReader(requestBytes))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
// Set headers
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
// Set protocol version header if negotiated
|
||||
if v := c.protocolVersion.Load(); v != nil {
|
||||
if version, ok := v.(string); ok && version != "" {
|
||||
req.Header.Set(HeaderKeyProtocolVersion, version)
|
||||
}
|
||||
}
|
||||
for k, v := range c.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
for k, v := range request.Header {
|
||||
if _, ok := req.Header[k]; !ok {
|
||||
req.Header[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
// Add OAuth authorization if configured
|
||||
if c.oauthHandler != nil {
|
||||
authHeader, err := c.oauthHandler.GetAuthorizationHeader(ctx)
|
||||
if err != nil {
|
||||
// If we get an authorization error, return a specific error that can be handled by the client
|
||||
if err.Error() == "no valid token available, authorization required" {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("failed to get authorization header: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", authHeader)
|
||||
}
|
||||
|
||||
if c.headerFunc != nil {
|
||||
for k, v := range c.headerFunc(ctx) {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
}
|
||||
|
||||
// Create string key for map lookup
|
||||
idKey := request.ID.String()
|
||||
|
||||
// Register response channel
|
||||
responseChan := make(chan *JSONRPCResponse, 1)
|
||||
c.mu.Lock()
|
||||
c.responses[idKey] = responseChan
|
||||
c.mu.Unlock()
|
||||
deleteResponseChan := func() {
|
||||
c.mu.Lock()
|
||||
delete(c.responses, idKey)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// Send request
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
deleteResponseChan()
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
|
||||
// Drain any outstanding io
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
// Check if we got an error response
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
deleteResponseChan()
|
||||
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
deleteResponseChan()
|
||||
return nil, ctx.Err()
|
||||
case response, ok := <-responseChan:
|
||||
if ok {
|
||||
return response, nil
|
||||
}
|
||||
return nil, fmt.Errorf("connection has been closed")
|
||||
}
|
||||
}
|
||||
|
||||
// Close shuts down the SSE client connection and cleans up any pending responses.
|
||||
// Returns an error if the shutdown process fails.
|
||||
func (c *SSE) Close() error {
|
||||
if !c.closed.CompareAndSwap(false, true) {
|
||||
return nil // Already closed
|
||||
}
|
||||
|
||||
if c.cancelSSEStream != nil {
|
||||
// It could stop the sse stream body, to quit the readSSE loop immediately
|
||||
// Also, it could quit start() immediately if not receiving the endpoint
|
||||
c.cancelSSEStream()
|
||||
}
|
||||
|
||||
// Clean up any pending responses
|
||||
c.mu.Lock()
|
||||
for _, ch := range c.responses {
|
||||
close(ch)
|
||||
}
|
||||
c.responses = make(map[string]chan *JSONRPCResponse)
|
||||
c.mu.Unlock()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetSessionId returns the session ID of the transport.
|
||||
// Since SSE does not maintain a session ID, it returns an empty string.
|
||||
func (c *SSE) GetSessionId() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// SetProtocolVersion sets the negotiated protocol version for this connection.
|
||||
func (c *SSE) SetProtocolVersion(version string) {
|
||||
c.protocolVersion.Store(version)
|
||||
}
|
||||
|
||||
// SendNotification sends a JSON-RPC notification to the server without expecting a response.
|
||||
func (c *SSE) SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error {
|
||||
if c.endpoint == nil {
|
||||
return fmt.Errorf("endpoint not received")
|
||||
}
|
||||
|
||||
notificationBytes, err := json.Marshal(notification)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal notification: %w", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(
|
||||
ctx,
|
||||
"POST",
|
||||
c.endpoint.String(),
|
||||
bytes.NewReader(notificationBytes),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create notification request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
// Set protocol version header if negotiated
|
||||
if v := c.protocolVersion.Load(); v != nil {
|
||||
if version, ok := v.(string); ok && version != "" {
|
||||
req.Header.Set(HeaderKeyProtocolVersion, version)
|
||||
}
|
||||
}
|
||||
// Set custom HTTP headers
|
||||
for k, v := range c.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
// Add OAuth authorization if configured
|
||||
if c.oauthHandler != nil {
|
||||
authHeader, err := c.oauthHandler.GetAuthorizationHeader(ctx)
|
||||
if err != nil {
|
||||
// If we get an authorization error, return a specific error that can be handled by the client
|
||||
if errors.Is(err, ErrOAuthAuthorizationRequired) {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("failed to get authorization header: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", authHeader)
|
||||
}
|
||||
|
||||
if c.headerFunc != nil {
|
||||
for k, v := range c.headerFunc(ctx) {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to send notification: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf(
|
||||
"notification failed with status %d: %s",
|
||||
resp.StatusCode,
|
||||
body,
|
||||
)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetEndpoint returns the current endpoint URL for the SSE connection.
|
||||
func (c *SSE) GetEndpoint() *url.URL {
|
||||
return c.endpoint
|
||||
}
|
||||
|
||||
// GetBaseURL returns the base URL set in the SSE constructor.
|
||||
func (c *SSE) GetBaseURL() *url.URL {
|
||||
return c.baseURL
|
||||
}
|
||||
|
||||
// GetOAuthHandler returns the OAuth handler if configured
|
||||
func (c *SSE) GetOAuthHandler() *OAuthHandler {
|
||||
return c.oauthHandler
|
||||
}
|
||||
|
||||
// IsOAuthEnabled returns true if OAuth is enabled
|
||||
func (c *SSE) IsOAuthEnabled() bool {
|
||||
return c.oauthHandler != nil
|
||||
}
|
||||
+475
@@ -0,0 +1,475 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/mark3labs/mcp-go/util"
|
||||
)
|
||||
|
||||
// Stdio implements the transport layer of the MCP protocol using stdio communication.
|
||||
// It launches a subprocess and communicates with it via standard input/output streams
|
||||
// using JSON-RPC messages. The client handles message routing between requests and
|
||||
// responses, and supports asynchronous notifications.
|
||||
type Stdio struct {
|
||||
command string
|
||||
args []string
|
||||
env []string
|
||||
|
||||
cmd *exec.Cmd
|
||||
cmdFunc CommandFunc
|
||||
stdin io.WriteCloser
|
||||
stdout *bufio.Reader
|
||||
stderr io.ReadCloser
|
||||
responses map[string]chan *JSONRPCResponse
|
||||
mu sync.RWMutex
|
||||
done chan struct{}
|
||||
onNotification func(mcp.JSONRPCNotification)
|
||||
notifyMu sync.RWMutex
|
||||
onRequest RequestHandler
|
||||
requestMu sync.RWMutex
|
||||
ctx context.Context
|
||||
ctxMu sync.RWMutex
|
||||
logger util.Logger
|
||||
started bool
|
||||
startedMu sync.Mutex
|
||||
}
|
||||
|
||||
// StdioOption defines a function that configures a Stdio transport instance.
|
||||
// Options can be used to customize the behavior of the transport before it starts,
|
||||
// such as setting a custom command function.
|
||||
type StdioOption func(*Stdio)
|
||||
|
||||
// CommandFunc is a factory function that returns a custom exec.Cmd used to launch the MCP subprocess.
|
||||
// It can be used to apply sandboxing, custom environment control, working directories, etc.
|
||||
type CommandFunc func(ctx context.Context, command string, env []string, args []string) (*exec.Cmd, error)
|
||||
|
||||
// WithCommandFunc sets a custom command factory function for the stdio transport.
|
||||
// The CommandFunc is responsible for constructing the exec.Cmd used to launch the subprocess,
|
||||
// allowing control over attributes like environment, working directory, and system-level sandboxing.
|
||||
func WithCommandFunc(f CommandFunc) StdioOption {
|
||||
return func(s *Stdio) {
|
||||
s.cmdFunc = f
|
||||
}
|
||||
}
|
||||
|
||||
// WithCommandLogger sets a custom logger for the stdio transport.
|
||||
func WithCommandLogger(logger util.Logger) StdioOption {
|
||||
return func(s *Stdio) {
|
||||
s.logger = logger
|
||||
}
|
||||
}
|
||||
|
||||
// NewIO returns a new stdio-based transport using existing input, output, and
|
||||
// logging streams instead of spawning a subprocess.
|
||||
// This is useful for testing and simulating client behavior.
|
||||
func NewIO(input io.Reader, output io.WriteCloser, logging io.ReadCloser) *Stdio {
|
||||
return &Stdio{
|
||||
stdin: output,
|
||||
stdout: bufio.NewReader(input),
|
||||
stderr: logging,
|
||||
|
||||
responses: make(map[string]chan *JSONRPCResponse),
|
||||
done: make(chan struct{}),
|
||||
ctx: context.Background(),
|
||||
logger: util.DefaultLogger(),
|
||||
}
|
||||
}
|
||||
|
||||
// NewStdio creates a new stdio transport to communicate with a subprocess.
|
||||
// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
|
||||
// Returns an error if the subprocess cannot be started or the pipes cannot be created.
|
||||
func NewStdio(
|
||||
command string,
|
||||
env []string,
|
||||
args ...string,
|
||||
) *Stdio {
|
||||
return NewStdioWithOptions(command, env, args)
|
||||
}
|
||||
|
||||
// NewStdioWithOptions creates a new stdio transport to communicate with a subprocess.
|
||||
// It launches the specified command with given arguments and sets up stdin/stdout pipes for communication.
|
||||
// Returns an error if the subprocess cannot be started or the pipes cannot be created.
|
||||
// Optional configuration functions can be provided to customize the transport before it starts,
|
||||
// such as setting a custom command factory.
|
||||
func NewStdioWithOptions(
|
||||
command string,
|
||||
env []string,
|
||||
args []string,
|
||||
opts ...StdioOption,
|
||||
) *Stdio {
|
||||
s := &Stdio{
|
||||
command: command,
|
||||
args: args,
|
||||
env: env,
|
||||
|
||||
responses: make(map[string]chan *JSONRPCResponse),
|
||||
done: make(chan struct{}),
|
||||
ctx: context.Background(),
|
||||
logger: util.DefaultLogger(),
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(s)
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
func (c *Stdio) Start(ctx context.Context) error {
|
||||
c.startedMu.Lock()
|
||||
if c.started {
|
||||
c.startedMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
c.started = true
|
||||
c.startedMu.Unlock()
|
||||
|
||||
// Store the context for use in request handling
|
||||
c.ctxMu.Lock()
|
||||
c.ctx = ctx
|
||||
c.ctxMu.Unlock()
|
||||
|
||||
if err := c.spawnCommand(ctx); err != nil {
|
||||
c.startedMu.Lock()
|
||||
c.started = false
|
||||
c.startedMu.Unlock()
|
||||
return err
|
||||
}
|
||||
|
||||
ready := make(chan struct{})
|
||||
go func() {
|
||||
close(ready)
|
||||
c.readResponses()
|
||||
}()
|
||||
<-ready
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// spawnCommand spawns a new process running the configured command, args, and env.
|
||||
// If an (optional) cmdFunc custom command factory function was configured, it will be used to construct the subprocess;
|
||||
// otherwise, the default behavior uses exec.CommandContext with the merged environment.
|
||||
// Initializes stdin, stdout, and stderr pipes for JSON-RPC communication.
|
||||
func (c *Stdio) spawnCommand(ctx context.Context) error {
|
||||
if c.command == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
var cmd *exec.Cmd
|
||||
var err error
|
||||
|
||||
// Standard behavior if no command func present.
|
||||
if c.cmdFunc == nil {
|
||||
cmd = exec.CommandContext(ctx, c.command, c.args...)
|
||||
cmd.Env = append(os.Environ(), c.env...)
|
||||
} else if cmd, err = c.cmdFunc(ctx, c.command, c.env, c.args); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
stdin, err := cmd.StdinPipe()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create stdin pipe: %w", err)
|
||||
}
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create stdout pipe: %w", err)
|
||||
}
|
||||
|
||||
stderr, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create stderr pipe: %w", err)
|
||||
}
|
||||
|
||||
c.cmd = cmd
|
||||
c.stdin = stdin
|
||||
c.stderr = stderr
|
||||
c.stdout = bufio.NewReader(stdout)
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return fmt.Errorf("failed to start command: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close shuts down the stdio client, closing the stdin pipe and waiting for the subprocess to exit.
|
||||
// Returns an error if there are issues closing stdin or waiting for the subprocess to terminate.
|
||||
func (c *Stdio) Close() error {
|
||||
select {
|
||||
case <-c.done:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
// cancel all in-flight request
|
||||
close(c.done)
|
||||
|
||||
if c.stdin != nil {
|
||||
if err := c.stdin.Close(); err != nil {
|
||||
return fmt.Errorf("failed to close stdin: %w", err)
|
||||
}
|
||||
}
|
||||
if c.stderr != nil {
|
||||
if err := c.stderr.Close(); err != nil {
|
||||
return fmt.Errorf("failed to close stderr: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if c.cmd != nil {
|
||||
return c.cmd.Wait()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetSessionId returns the session ID of the transport.
|
||||
// Since stdio does not maintain a session ID, it returns an empty string.
|
||||
func (c *Stdio) GetSessionId() string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// SetNotificationHandler sets the handler function to be called when a notification is received.
|
||||
// Only one handler can be set at a time; setting a new one replaces the previous handler.
|
||||
func (c *Stdio) SetNotificationHandler(
|
||||
handler func(notification mcp.JSONRPCNotification),
|
||||
) {
|
||||
c.notifyMu.Lock()
|
||||
defer c.notifyMu.Unlock()
|
||||
c.onNotification = handler
|
||||
}
|
||||
|
||||
// SetRequestHandler sets the handler function to be called when a request is received from the server.
|
||||
// This enables bidirectional communication for features like sampling.
|
||||
func (c *Stdio) SetRequestHandler(handler RequestHandler) {
|
||||
c.requestMu.Lock()
|
||||
defer c.requestMu.Unlock()
|
||||
c.onRequest = handler
|
||||
}
|
||||
|
||||
// readResponses continuously reads and processes responses from the server's stdout.
|
||||
// It handles both responses to requests and notifications, routing them appropriately.
|
||||
// Runs until the done channel is closed or an error occurs reading from stdout.
|
||||
func (c *Stdio) readResponses() {
|
||||
for {
|
||||
select {
|
||||
case <-c.done:
|
||||
return
|
||||
default:
|
||||
line, err := c.stdout.ReadString('\n')
|
||||
if err != nil {
|
||||
if err != io.EOF && !errors.Is(err, context.Canceled) {
|
||||
c.logger.Errorf("Error reading from stdout: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
// First try to parse as a generic message to check for ID field
|
||||
var baseMessage struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID *mcp.RequestId `json:"id,omitempty"`
|
||||
Method string `json:"method,omitempty"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(line), &baseMessage); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// If it has a method but no ID, it's a notification
|
||||
if baseMessage.Method != "" && baseMessage.ID == nil {
|
||||
var notification mcp.JSONRPCNotification
|
||||
if err := json.Unmarshal([]byte(line), ¬ification); err != nil {
|
||||
continue
|
||||
}
|
||||
c.notifyMu.RLock()
|
||||
if c.onNotification != nil {
|
||||
c.onNotification(notification)
|
||||
}
|
||||
c.notifyMu.RUnlock()
|
||||
continue
|
||||
}
|
||||
|
||||
// If it has a method and an ID, it's an incoming request
|
||||
if baseMessage.Method != "" && baseMessage.ID != nil {
|
||||
var request JSONRPCRequest
|
||||
if err := json.Unmarshal([]byte(line), &request); err == nil {
|
||||
c.handleIncomingRequest(request)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
// Otherwise, it's a response to our request
|
||||
var response JSONRPCResponse
|
||||
if err := json.Unmarshal([]byte(line), &response); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Create string key for map lookup
|
||||
idKey := response.ID.String()
|
||||
|
||||
c.mu.RLock()
|
||||
ch, exists := c.responses[idKey]
|
||||
c.mu.RUnlock()
|
||||
|
||||
if exists {
|
||||
ch <- &response
|
||||
c.mu.Lock()
|
||||
delete(c.responses, idKey)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SendRequest sends a JSON-RPC request to the server and waits for a response.
|
||||
// It creates a unique request ID, sends the request over stdin, and waits for
|
||||
// the corresponding response or context cancellation.
|
||||
// Returns the raw JSON response message or an error if the request fails.
|
||||
func (c *Stdio) SendRequest(
|
||||
ctx context.Context,
|
||||
request JSONRPCRequest,
|
||||
) (*JSONRPCResponse, error) {
|
||||
// Check if context is already canceled before doing any work
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
if c.stdin == nil {
|
||||
return nil, fmt.Errorf("stdio client not started")
|
||||
}
|
||||
|
||||
// Marshal request
|
||||
requestBytes, err := json.Marshal(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
requestBytes = append(requestBytes, '\n')
|
||||
|
||||
// Create string key for map lookup
|
||||
idKey := request.ID.String()
|
||||
|
||||
// Register response channel
|
||||
responseChan := make(chan *JSONRPCResponse, 1)
|
||||
c.mu.Lock()
|
||||
c.responses[idKey] = responseChan
|
||||
c.mu.Unlock()
|
||||
deleteResponseChan := func() {
|
||||
c.mu.Lock()
|
||||
delete(c.responses, idKey)
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// Send request
|
||||
if _, err := c.stdin.Write(requestBytes); err != nil {
|
||||
deleteResponseChan()
|
||||
return nil, fmt.Errorf("failed to write request: %w", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
deleteResponseChan()
|
||||
return nil, ctx.Err()
|
||||
case response := <-responseChan:
|
||||
return response, nil
|
||||
}
|
||||
}
|
||||
|
||||
// SendNotification sends a json RPC Notification to the server.
|
||||
func (c *Stdio) SendNotification(
|
||||
ctx context.Context,
|
||||
notification mcp.JSONRPCNotification,
|
||||
) error {
|
||||
if c.stdin == nil {
|
||||
return fmt.Errorf("stdio client not started")
|
||||
}
|
||||
|
||||
notificationBytes, err := json.Marshal(notification)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal notification: %w", err)
|
||||
}
|
||||
notificationBytes = append(notificationBytes, '\n')
|
||||
|
||||
if _, err := c.stdin.Write(notificationBytes); err != nil {
|
||||
return fmt.Errorf("failed to write notification: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleIncomingRequest processes incoming requests from the server.
|
||||
// It calls the registered request handler and sends the response back to the server.
|
||||
func (c *Stdio) handleIncomingRequest(request JSONRPCRequest) {
|
||||
c.requestMu.RLock()
|
||||
handler := c.onRequest
|
||||
c.requestMu.RUnlock()
|
||||
|
||||
if handler == nil {
|
||||
// Send error response if no handler is configured
|
||||
errorResponse := *NewJSONRPCErrorResponse(
|
||||
request.ID,
|
||||
mcp.METHOD_NOT_FOUND,
|
||||
"No request handler configured",
|
||||
nil,
|
||||
)
|
||||
c.sendResponse(errorResponse)
|
||||
return
|
||||
}
|
||||
|
||||
// Handle the request in a goroutine to avoid blocking
|
||||
go func() {
|
||||
c.ctxMu.RLock()
|
||||
ctx := c.ctx
|
||||
c.ctxMu.RUnlock()
|
||||
|
||||
// Check if context is already cancelled before processing
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
errorResponse := *NewJSONRPCErrorResponse(request.ID, mcp.INTERNAL_ERROR, ctx.Err().Error(), nil)
|
||||
c.sendResponse(errorResponse)
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
response, err := handler(ctx, request)
|
||||
if err != nil {
|
||||
errorResponse := *NewJSONRPCErrorResponse(request.ID, mcp.INTERNAL_ERROR, err.Error(), nil)
|
||||
c.sendResponse(errorResponse)
|
||||
return
|
||||
}
|
||||
|
||||
if response != nil {
|
||||
c.sendResponse(*response)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// sendResponse sends a response back to the server.
|
||||
func (c *Stdio) sendResponse(response JSONRPCResponse) {
|
||||
responseBytes, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
c.logger.Errorf("Error marshaling response: %v", err)
|
||||
return
|
||||
}
|
||||
responseBytes = append(responseBytes, '\n')
|
||||
|
||||
if _, err := c.stdin.Write(responseBytes); err != nil {
|
||||
c.logger.Errorf("Error writing response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Stderr returns a reader for the stderr output of the subprocess.
|
||||
// This can be used to capture error messages or logs from the subprocess.
|
||||
func (c *Stdio) Stderr() io.Reader {
|
||||
return c.stderr
|
||||
}
|
||||
+791
@@ -0,0 +1,791 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/mark3labs/mcp-go/util"
|
||||
)
|
||||
|
||||
type StreamableHTTPCOption func(*StreamableHTTP)
|
||||
|
||||
// WithContinuousListening enables receiving server-to-client notifications when no request is in flight.
|
||||
// In particular, if you want to receive global notifications from the server (like ToolListChangedNotification),
|
||||
// you should enable this option.
|
||||
//
|
||||
// It will establish a standalone long-live GET HTTP connection to the server.
|
||||
// https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#listening-for-messages-from-the-server
|
||||
// NOTICE: Even enabled, the server may not support this feature.
|
||||
func WithContinuousListening() StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.getListeningEnabled = true
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPClient sets a custom HTTP client on the StreamableHTTP transport.
|
||||
func WithHTTPBasicClient(client *http.Client) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.httpClient = client
|
||||
}
|
||||
}
|
||||
|
||||
func WithHTTPHeaders(headers map[string]string) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.headers = headers
|
||||
}
|
||||
}
|
||||
|
||||
func WithHTTPHeaderFunc(headerFunc HTTPHeaderFunc) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.headerFunc = headerFunc
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPTimeout sets the timeout for a HTTP request and stream.
|
||||
func WithHTTPTimeout(timeout time.Duration) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.httpClient.Timeout = timeout
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPOAuth enables OAuth authentication for the client.
|
||||
func WithHTTPOAuth(config OAuthConfig) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.oauthHandler = NewOAuthHandler(config)
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPLogger sets a custom logger for the StreamableHTTP transport.
|
||||
func WithHTTPLogger(logger util.Logger) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.logger = logger
|
||||
}
|
||||
}
|
||||
|
||||
// Deprecated: Use [WithHTTPLogger] instead.
|
||||
func WithLogger(logger util.Logger) StreamableHTTPCOption {
|
||||
return WithHTTPLogger(logger)
|
||||
}
|
||||
|
||||
// WithSession creates a client with a pre-configured session
|
||||
func WithSession(sessionID string) StreamableHTTPCOption {
|
||||
return func(sc *StreamableHTTP) {
|
||||
sc.sessionID.Store(sessionID)
|
||||
}
|
||||
}
|
||||
|
||||
// StreamableHTTP implements Streamable HTTP transport.
|
||||
//
|
||||
// It transmits JSON-RPC messages over individual HTTP requests. One message per request.
|
||||
// The HTTP response body can either be a single JSON-RPC response,
|
||||
// or an upgraded SSE stream that concludes with a JSON-RPC response for the same request.
|
||||
//
|
||||
// https://modelcontextprotocol.io/specification/2025-03-26/basic/transports
|
||||
//
|
||||
// The current implementation does not support the following features:
|
||||
// - resuming stream
|
||||
// (https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#resumability-and-redelivery)
|
||||
type StreamableHTTP struct {
|
||||
serverURL *url.URL
|
||||
httpClient *http.Client
|
||||
headers map[string]string
|
||||
headerFunc HTTPHeaderFunc
|
||||
logger util.Logger
|
||||
getListeningEnabled bool
|
||||
|
||||
sessionID atomic.Value // string
|
||||
protocolVersion atomic.Value // string
|
||||
|
||||
initialized chan struct{}
|
||||
initializedOnce sync.Once
|
||||
|
||||
notificationHandler func(mcp.JSONRPCNotification)
|
||||
notifyMu sync.RWMutex
|
||||
|
||||
// Request handler for incoming server-to-client requests (like sampling)
|
||||
requestHandler RequestHandler
|
||||
requestMu sync.RWMutex
|
||||
|
||||
closed chan struct{}
|
||||
|
||||
// OAuth support
|
||||
oauthHandler *OAuthHandler
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewStreamableHTTP creates a new Streamable HTTP transport with the given server URL.
|
||||
// Returns an error if the URL is invalid.
|
||||
func NewStreamableHTTP(serverURL string, options ...StreamableHTTPCOption) (*StreamableHTTP, error) {
|
||||
parsedURL, err := url.Parse(serverURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid URL: %w", err)
|
||||
}
|
||||
|
||||
smc := &StreamableHTTP{
|
||||
serverURL: parsedURL,
|
||||
httpClient: &http.Client{},
|
||||
headers: make(map[string]string),
|
||||
closed: make(chan struct{}),
|
||||
logger: util.DefaultLogger(),
|
||||
initialized: make(chan struct{}),
|
||||
}
|
||||
smc.sessionID.Store("") // set initial value to simplify later usage
|
||||
|
||||
for _, opt := range options {
|
||||
if opt != nil {
|
||||
opt(smc)
|
||||
}
|
||||
}
|
||||
|
||||
// If OAuth is configured, set the base URL for metadata discovery
|
||||
if smc.oauthHandler != nil {
|
||||
// Extract base URL from server URL for metadata discovery
|
||||
baseURL := fmt.Sprintf("%s://%s", parsedURL.Scheme, parsedURL.Host)
|
||||
smc.oauthHandler.SetBaseURL(baseURL)
|
||||
}
|
||||
|
||||
return smc, nil
|
||||
}
|
||||
|
||||
// Start initiates the HTTP connection to the server.
|
||||
func (c *StreamableHTTP) Start(ctx context.Context) error {
|
||||
// Start is idempotent - check if already initialized
|
||||
select {
|
||||
case <-c.initialized:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
|
||||
// For Streamable HTTP, we don't need to establish a persistent connection by default
|
||||
if c.getListeningEnabled {
|
||||
go func() {
|
||||
select {
|
||||
case <-c.initialized:
|
||||
ctx, cancel := c.contextAwareOfClientClose(ctx)
|
||||
defer cancel()
|
||||
c.listenForever(ctx)
|
||||
case <-c.closed:
|
||||
return
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close closes the all the HTTP connections to the server.
|
||||
func (c *StreamableHTTP) Close() error {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
// Cancel all in-flight requests
|
||||
close(c.closed)
|
||||
|
||||
sessionId := c.sessionID.Load().(string)
|
||||
if sessionId != "" {
|
||||
c.sessionID.Store("")
|
||||
c.wg.Add(1)
|
||||
// notify server session closed
|
||||
go func() {
|
||||
defer c.wg.Done()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, c.serverURL.String(), nil)
|
||||
if err != nil {
|
||||
c.logger.Errorf("failed to create close request: %v", err)
|
||||
return
|
||||
}
|
||||
req.Header.Set(HeaderKeySessionID, sessionId)
|
||||
// Set protocol version header if negotiated
|
||||
if v := c.protocolVersion.Load(); v != nil {
|
||||
if version, ok := v.(string); ok && version != "" {
|
||||
req.Header.Set(HeaderKeyProtocolVersion, version)
|
||||
}
|
||||
}
|
||||
res, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
c.logger.Errorf("failed to send close request: %v", err)
|
||||
return
|
||||
}
|
||||
res.Body.Close()
|
||||
}()
|
||||
}
|
||||
c.wg.Wait()
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetProtocolVersion sets the negotiated protocol version for this connection.
|
||||
func (c *StreamableHTTP) SetProtocolVersion(version string) {
|
||||
c.protocolVersion.Store(version)
|
||||
}
|
||||
|
||||
// ErrOAuthAuthorizationRequired is a sentinel error for OAuth authorization required
|
||||
var ErrOAuthAuthorizationRequired = errors.New("no valid token available, authorization required")
|
||||
|
||||
// OAuthAuthorizationRequiredError is returned when OAuth authorization is required
|
||||
type OAuthAuthorizationRequiredError struct {
|
||||
Handler *OAuthHandler
|
||||
}
|
||||
|
||||
func (e *OAuthAuthorizationRequiredError) Error() string {
|
||||
return ErrOAuthAuthorizationRequired.Error()
|
||||
}
|
||||
|
||||
func (e *OAuthAuthorizationRequiredError) Unwrap() error {
|
||||
return ErrOAuthAuthorizationRequired
|
||||
}
|
||||
|
||||
// SendRequest sends a JSON-RPC request to the server and waits for a response.
|
||||
// Returns the raw JSON response message or an error if the request fails.
|
||||
func (c *StreamableHTTP) SendRequest(
|
||||
ctx context.Context,
|
||||
request JSONRPCRequest,
|
||||
) (*JSONRPCResponse, error) {
|
||||
// Marshal request
|
||||
requestBody, err := json.Marshal(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
ctx, cancel := c.contextAwareOfClientClose(ctx)
|
||||
defer cancel()
|
||||
|
||||
resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(requestBody), "application/json, text/event-stream", request.Header)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrSessionTerminated) && request.Method == string(mcp.MethodInitialize) {
|
||||
// If the request is initialize, should not return a SessionTerminated error
|
||||
// It should be a genuine endpoint-routing issue.
|
||||
// ( Fall through to return StatusCode checking. )
|
||||
} else {
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Only proceed if we have a valid response.
|
||||
// When sendHTTP fails and resp is nil but method is mcp.MethodInitialize
|
||||
// defer resp.Body.Close() fails with nil pointer dereference.
|
||||
if resp == nil {
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Check if we got an error response
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
|
||||
// handle error response
|
||||
var errResponse JSONRPCResponse
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if err := json.Unmarshal(body, &errResponse); err == nil {
|
||||
return &errResponse, nil
|
||||
}
|
||||
return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
if request.Method == string(mcp.MethodInitialize) {
|
||||
// saved the received session ID in the response
|
||||
// empty session ID is allowed
|
||||
if sessionID := resp.Header.Get(HeaderKeySessionID); sessionID != "" {
|
||||
c.sessionID.Store(sessionID)
|
||||
}
|
||||
|
||||
c.initializedOnce.Do(func() {
|
||||
close(c.initialized)
|
||||
})
|
||||
}
|
||||
|
||||
// Handle different response types
|
||||
mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
|
||||
switch mediaType {
|
||||
case "application/json":
|
||||
// Single response
|
||||
var response JSONRPCResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
|
||||
return nil, fmt.Errorf("failed to decode response: %w", err)
|
||||
}
|
||||
|
||||
// should not be a notification
|
||||
if response.ID.IsNil() {
|
||||
return nil, fmt.Errorf("response should contain RPC id: %v", response)
|
||||
}
|
||||
|
||||
return &response, nil
|
||||
|
||||
case "text/event-stream":
|
||||
// Server is using SSE for streaming responses
|
||||
return c.handleSSEResponse(ctx, resp.Body, false)
|
||||
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected content type: %s", resp.Header.Get("Content-Type"))
|
||||
}
|
||||
}
|
||||
|
||||
func (c *StreamableHTTP) sendHTTP(
|
||||
ctx context.Context,
|
||||
method string,
|
||||
body io.Reader,
|
||||
acceptType string,
|
||||
header http.Header,
|
||||
) (resp *http.Response, err error) {
|
||||
// Create HTTP request
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.serverURL.String(), body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
// request headers
|
||||
if header != nil {
|
||||
req.Header = header
|
||||
}
|
||||
|
||||
// Set headers
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", acceptType)
|
||||
sessionID := c.sessionID.Load().(string)
|
||||
if sessionID != "" {
|
||||
req.Header.Set(HeaderKeySessionID, sessionID)
|
||||
}
|
||||
// Set protocol version header if negotiated
|
||||
if v := c.protocolVersion.Load(); v != nil {
|
||||
if version, ok := v.(string); ok && version != "" {
|
||||
req.Header.Set(HeaderKeyProtocolVersion, version)
|
||||
}
|
||||
}
|
||||
for k, v := range c.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
|
||||
// Add OAuth authorization if configured
|
||||
if c.oauthHandler != nil {
|
||||
authHeader, err := c.oauthHandler.GetAuthorizationHeader(ctx)
|
||||
if err != nil {
|
||||
// If we get an authorization error, return a specific error that can be handled by the client
|
||||
if errors.Is(err, ErrOAuthAuthorizationRequired) {
|
||||
return nil, &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("failed to get authorization header: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", authHeader)
|
||||
}
|
||||
|
||||
if c.headerFunc != nil {
|
||||
for k, v := range c.headerFunc(ctx) {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
}
|
||||
|
||||
// Send request
|
||||
resp, err = c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
|
||||
// universal handling for session terminated
|
||||
if resp.StatusCode == http.StatusNotFound {
|
||||
c.sessionID.CompareAndSwap(sessionID, "")
|
||||
return nil, ErrSessionTerminated
|
||||
}
|
||||
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// handleSSEResponse processes an SSE stream for a specific request.
|
||||
// It returns the final result for the request once received, or an error.
|
||||
// If ignoreResponse is true, it won't return when a response messge is received. This is for continuous listening.
|
||||
func (c *StreamableHTTP) handleSSEResponse(ctx context.Context, reader io.ReadCloser, ignoreResponse bool) (*JSONRPCResponse, error) {
|
||||
// Create a channel for this specific request
|
||||
responseChan := make(chan *JSONRPCResponse, 1)
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
// Start a goroutine to process the SSE stream
|
||||
go func() {
|
||||
// Ensure this goroutine respects the context
|
||||
defer close(responseChan)
|
||||
|
||||
c.readSSE(ctx, reader, func(event, data string) {
|
||||
// Try to unmarshal as a response first
|
||||
var message JSONRPCResponse
|
||||
if err := json.Unmarshal([]byte(data), &message); err != nil {
|
||||
c.logger.Infof("failed to unmarshal message (non-fatal): %v", err, "message", data)
|
||||
return
|
||||
}
|
||||
|
||||
// Handle notification
|
||||
if message.ID.IsNil() {
|
||||
var notification mcp.JSONRPCNotification
|
||||
if err := json.Unmarshal([]byte(data), ¬ification); err != nil {
|
||||
c.logger.Errorf("failed to unmarshal notification: %v", err)
|
||||
return
|
||||
}
|
||||
c.notifyMu.RLock()
|
||||
if c.notificationHandler != nil {
|
||||
c.notificationHandler(notification)
|
||||
}
|
||||
c.notifyMu.RUnlock()
|
||||
return
|
||||
}
|
||||
|
||||
// Check if this is actually a request from the server by looking for method field
|
||||
var rawMessage map[string]json.RawMessage
|
||||
if err := json.Unmarshal([]byte(data), &rawMessage); err == nil {
|
||||
if _, hasMethod := rawMessage["method"]; hasMethod && !message.ID.IsNil() {
|
||||
var request JSONRPCRequest
|
||||
if err := json.Unmarshal([]byte(data), &request); err == nil {
|
||||
// This is a request from the server
|
||||
c.handleIncomingRequest(ctx, request)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !ignoreResponse {
|
||||
responseChan <- &message
|
||||
}
|
||||
})
|
||||
}()
|
||||
|
||||
// Wait for the response or context cancellation
|
||||
select {
|
||||
case response := <-responseChan:
|
||||
if response == nil {
|
||||
return nil, fmt.Errorf("unexpected nil response")
|
||||
}
|
||||
return response, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// readSSE reads the SSE stream(reader) and calls the handler for each event and data pair.
|
||||
// It will end when the reader is closed (or the context is done).
|
||||
func (c *StreamableHTTP) readSSE(ctx context.Context, reader io.ReadCloser, handler func(event, data string)) {
|
||||
defer reader.Close()
|
||||
|
||||
br := bufio.NewReader(reader)
|
||||
var event, data string
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
line, err := br.ReadString('\n')
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
// Process any pending event before exit
|
||||
if data != "" {
|
||||
// If no event type is specified, use empty string (default event type)
|
||||
if event == "" {
|
||||
event = "message"
|
||||
}
|
||||
handler(event, data)
|
||||
}
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
c.logger.Errorf("SSE stream error: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Remove only newline markers
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
if line == "" {
|
||||
// Empty line means end of event
|
||||
if data != "" {
|
||||
// If no event type is specified, use empty string (default event type)
|
||||
if event == "" {
|
||||
event = "message"
|
||||
}
|
||||
handler(event, data)
|
||||
event = ""
|
||||
data = ""
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(line, "event:") {
|
||||
event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
|
||||
} else if strings.HasPrefix(line, "data:") {
|
||||
data = strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *StreamableHTTP) SendNotification(ctx context.Context, notification mcp.JSONRPCNotification) error {
|
||||
// Marshal request
|
||||
requestBody, err := json.Marshal(notification)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal notification: %w", err)
|
||||
}
|
||||
|
||||
// Create HTTP request
|
||||
ctx, cancel := c.contextAwareOfClientClose(ctx)
|
||||
defer cancel()
|
||||
|
||||
resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(requestBody), "application/json, text/event-stream", nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
// Handle OAuth unauthorized error
|
||||
if resp.StatusCode == http.StatusUnauthorized && c.oauthHandler != nil {
|
||||
return &OAuthAuthorizationRequiredError{
|
||||
Handler: c.oauthHandler,
|
||||
}
|
||||
}
|
||||
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf(
|
||||
"notification failed with status %d: %s",
|
||||
resp.StatusCode,
|
||||
body,
|
||||
)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *StreamableHTTP) SetNotificationHandler(handler func(mcp.JSONRPCNotification)) {
|
||||
c.notifyMu.Lock()
|
||||
defer c.notifyMu.Unlock()
|
||||
c.notificationHandler = handler
|
||||
}
|
||||
|
||||
// SetRequestHandler sets the handler for incoming requests from the server.
|
||||
func (c *StreamableHTTP) SetRequestHandler(handler RequestHandler) {
|
||||
c.requestMu.Lock()
|
||||
defer c.requestMu.Unlock()
|
||||
c.requestHandler = handler
|
||||
}
|
||||
|
||||
func (c *StreamableHTTP) GetSessionId() string {
|
||||
return c.sessionID.Load().(string)
|
||||
}
|
||||
|
||||
// GetOAuthHandler returns the OAuth handler if configured
|
||||
func (c *StreamableHTTP) GetOAuthHandler() *OAuthHandler {
|
||||
return c.oauthHandler
|
||||
}
|
||||
|
||||
// IsOAuthEnabled returns true if OAuth is enabled
|
||||
func (c *StreamableHTTP) IsOAuthEnabled() bool {
|
||||
return c.oauthHandler != nil
|
||||
}
|
||||
|
||||
func (c *StreamableHTTP) listenForever(ctx context.Context) {
|
||||
c.logger.Infof("listening to server forever")
|
||||
for {
|
||||
// Use the original context for continuous listening - no per-iteration timeout
|
||||
// The SSE connection itself will detect disconnections via the underlying HTTP transport,
|
||||
// and the context cancellation will propagate from the parent to stop listening gracefully.
|
||||
// We don't add an artificial timeout here because:
|
||||
// 1. Persistent SSE connections are meant to stay open indefinitely
|
||||
// 2. Network-level timeouts and keep-alives handle connection health
|
||||
// 3. Context cancellation (user-initiated or system shutdown) provides clean shutdown
|
||||
err := c.createGETConnectionToServer(ctx)
|
||||
if errors.Is(err, ErrGetMethodNotAllowed) {
|
||||
// server does not support listening
|
||||
c.logger.Errorf("server does not support listening")
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
c.logger.Errorf("failed to listen to server. retry in 1 second: %v", err)
|
||||
}
|
||||
|
||||
// Use context-aware sleep
|
||||
select {
|
||||
case <-time.After(retryInterval):
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
ErrSessionTerminated = fmt.Errorf("session terminated (404). need to re-initialize")
|
||||
ErrGetMethodNotAllowed = fmt.Errorf("GET method not allowed")
|
||||
|
||||
retryInterval = 1 * time.Second // a variable is convenient for testing
|
||||
)
|
||||
|
||||
func (c *StreamableHTTP) createGETConnectionToServer(ctx context.Context) error {
|
||||
resp, err := c.sendHTTP(ctx, http.MethodGet, nil, "text/event-stream", nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// Check if we got an error response
|
||||
if resp.StatusCode == http.StatusMethodNotAllowed {
|
||||
return ErrGetMethodNotAllowed
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return fmt.Errorf("request failed with status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
// handle SSE response
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType != "text/event-stream" {
|
||||
return fmt.Errorf("unexpected content type: %s", contentType)
|
||||
}
|
||||
|
||||
// When ignoreResponse is true, the function will never return expect context is done.
|
||||
// NOTICE: Due to the ambiguity of the specification, other SDKs may use the GET connection to transfer the response
|
||||
// messages. To be more compatible, we should handle this response, however, as the transport layer is message-based,
|
||||
// currently, there is no convenient way to handle this response.
|
||||
// So we ignore the response here. It's not a bug, but may be not compatible with other SDKs.
|
||||
_, err = c.handleSSEResponse(ctx, resp.Body, true)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to handle SSE response: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleIncomingRequest processes requests from the server (like sampling requests)
|
||||
func (c *StreamableHTTP) handleIncomingRequest(ctx context.Context, request JSONRPCRequest) {
|
||||
c.requestMu.RLock()
|
||||
handler := c.requestHandler
|
||||
c.requestMu.RUnlock()
|
||||
|
||||
if handler == nil {
|
||||
c.logger.Errorf("received request from server but no handler set: %s", request.Method)
|
||||
// Send method not found error
|
||||
errorResponse := NewJSONRPCErrorResponse(
|
||||
request.ID,
|
||||
mcp.METHOD_NOT_FOUND,
|
||||
fmt.Sprintf("no handler configured for method: %s", request.Method),
|
||||
nil,
|
||||
)
|
||||
c.sendResponseToServer(ctx, errorResponse)
|
||||
return
|
||||
}
|
||||
|
||||
// Handle the request in a goroutine to avoid blocking the SSE reader
|
||||
go func() {
|
||||
// Create a new context with timeout for request handling, respecting parent context
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
response, err := handler(requestCtx, request)
|
||||
if err != nil {
|
||||
c.logger.Errorf("error handling request %s: %v", request.Method, err)
|
||||
|
||||
// Determine appropriate JSON-RPC error code based on error type
|
||||
var errorCode int
|
||||
var errorMessage string
|
||||
|
||||
// Check for specific sampling-related errors
|
||||
if errors.Is(err, context.Canceled) {
|
||||
errorCode = mcp.REQUEST_INTERRUPTED
|
||||
errorMessage = "request was cancelled"
|
||||
} else if errors.Is(err, context.DeadlineExceeded) {
|
||||
errorCode = mcp.REQUEST_INTERRUPTED
|
||||
errorMessage = "request timed out"
|
||||
} else {
|
||||
// Generic error cases
|
||||
switch request.Method {
|
||||
case string(mcp.MethodSamplingCreateMessage):
|
||||
errorCode = mcp.INTERNAL_ERROR
|
||||
errorMessage = fmt.Sprintf("sampling request failed: %v", err)
|
||||
default:
|
||||
errorCode = mcp.INTERNAL_ERROR
|
||||
errorMessage = err.Error()
|
||||
}
|
||||
}
|
||||
|
||||
// Send error response
|
||||
errorResponse := NewJSONRPCErrorResponse(request.ID, errorCode, errorMessage, nil)
|
||||
c.sendResponseToServer(requestCtx, errorResponse)
|
||||
return
|
||||
}
|
||||
|
||||
if response != nil {
|
||||
c.sendResponseToServer(requestCtx, response)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// sendResponseToServer sends a response back to the server via HTTP POST
|
||||
func (c *StreamableHTTP) sendResponseToServer(ctx context.Context, response *JSONRPCResponse) {
|
||||
if response == nil {
|
||||
c.logger.Errorf("cannot send nil response to server")
|
||||
return
|
||||
}
|
||||
|
||||
responseBody, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
c.logger.Errorf("failed to marshal response: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
ctx, cancel := c.contextAwareOfClientClose(ctx)
|
||||
defer cancel()
|
||||
|
||||
resp, err := c.sendHTTP(ctx, http.MethodPost, bytes.NewReader(responseBody), "application/json, text/event-stream", nil)
|
||||
if err != nil {
|
||||
c.logger.Errorf("failed to send response to server: %v", err)
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
c.logger.Errorf("server rejected response with status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *StreamableHTTP) contextAwareOfClientClose(ctx context.Context) (context.Context, context.CancelFunc) {
|
||||
newCtx, cancel := context.WithCancel(ctx)
|
||||
go func() {
|
||||
select {
|
||||
case <-c.closed:
|
||||
cancel()
|
||||
case <-newCtx.Done():
|
||||
// The original context was canceled
|
||||
cancel()
|
||||
}
|
||||
}()
|
||||
return newCtx, cancel
|
||||
}
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
package transport
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// NewJSONRPCErrorResponse creates a new JSONRPCResponse with an error.
|
||||
func NewJSONRPCErrorResponse(id mcp.RequestId, code int, message string, data any) *JSONRPCResponse {
|
||||
details := mcp.NewJSONRPCErrorDetails(code, message, data)
|
||||
return &JSONRPCResponse{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Error: &details,
|
||||
}
|
||||
}
|
||||
|
||||
// NewJSONRPCResultResponse creates a new JSONRPCResponse with a result.
|
||||
func NewJSONRPCResultResponse(id mcp.RequestId, result json.RawMessage) *JSONRPCResponse {
|
||||
return &JSONRPCResponse{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Result: result,
|
||||
}
|
||||
}
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
package mcp
|
||||
|
||||
const (
|
||||
ContentTypeText = "text"
|
||||
ContentTypeImage = "image"
|
||||
ContentTypeAudio = "audio"
|
||||
ContentTypeLink = "resource_link"
|
||||
ContentTypeResource = "resource"
|
||||
)
|
||||
+85
@@ -0,0 +1,85 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Sentinel errors for common JSON-RPC error codes.
|
||||
var (
|
||||
// ErrParseError indicates a JSON parsing error (code: PARSE_ERROR).
|
||||
ErrParseError = errors.New("parse error")
|
||||
|
||||
// ErrInvalidRequest indicates an invalid JSON-RPC request (code: INVALID_REQUEST).
|
||||
ErrInvalidRequest = errors.New("invalid request")
|
||||
|
||||
// ErrMethodNotFound indicates the requested method does not exist (code: METHOD_NOT_FOUND).
|
||||
ErrMethodNotFound = errors.New("method not found")
|
||||
|
||||
// ErrInvalidParams indicates invalid method parameters (code: INVALID_PARAMS).
|
||||
ErrInvalidParams = errors.New("invalid params")
|
||||
|
||||
// ErrInternalError indicates an internal JSON-RPC error (code: INTERNAL_ERROR).
|
||||
ErrInternalError = errors.New("internal error")
|
||||
|
||||
// ErrRequestInterrupted indicates a request was cancelled or timed out (code: REQUEST_INTERRUPTED).
|
||||
ErrRequestInterrupted = errors.New("request interrupted")
|
||||
|
||||
// ErrResourceNotFound indicates a requested resource was not found (code: RESOURCE_NOT_FOUND).
|
||||
ErrResourceNotFound = errors.New("resource not found")
|
||||
)
|
||||
|
||||
// UnsupportedProtocolVersionError is returned when the server responds with
|
||||
// a protocol version that the client doesn't support.
|
||||
type UnsupportedProtocolVersionError struct {
|
||||
Version string
|
||||
}
|
||||
|
||||
func (e UnsupportedProtocolVersionError) Error() string {
|
||||
return fmt.Sprintf("unsupported protocol version: %q", e.Version)
|
||||
}
|
||||
|
||||
// Is implements the errors.Is interface for better error handling
|
||||
func (e UnsupportedProtocolVersionError) Is(target error) bool {
|
||||
_, ok := target.(UnsupportedProtocolVersionError)
|
||||
return ok
|
||||
}
|
||||
|
||||
// IsUnsupportedProtocolVersion checks if an error is an UnsupportedProtocolVersionError
|
||||
func IsUnsupportedProtocolVersion(err error) bool {
|
||||
_, ok := err.(UnsupportedProtocolVersionError)
|
||||
return ok
|
||||
}
|
||||
|
||||
// AsError maps JSONRPCErrorDetails to a Go error.
|
||||
// Returns sentinel errors wrapped with custom messages for known codes.
|
||||
// Defaults to a generic error with the original message when the code is not mapped.
|
||||
func (e *JSONRPCErrorDetails) AsError() error {
|
||||
var err error
|
||||
|
||||
switch e.Code {
|
||||
case PARSE_ERROR:
|
||||
err = ErrParseError
|
||||
case INVALID_REQUEST:
|
||||
err = ErrInvalidRequest
|
||||
case METHOD_NOT_FOUND:
|
||||
err = ErrMethodNotFound
|
||||
case INVALID_PARAMS:
|
||||
err = ErrInvalidParams
|
||||
case INTERNAL_ERROR:
|
||||
err = ErrInternalError
|
||||
case REQUEST_INTERRUPTED:
|
||||
err = ErrRequestInterrupted
|
||||
case RESOURCE_NOT_FOUND:
|
||||
err = ErrResourceNotFound
|
||||
default:
|
||||
return errors.New(e.Message)
|
||||
}
|
||||
|
||||
// Wrap the sentinel error with the custom message if it differs from the sentinel.
|
||||
if e.Message != "" && e.Message != err.Error() {
|
||||
return fmt.Errorf("%w: %s", err, e.Message)
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
+176
@@ -0,0 +1,176 @@
|
||||
package mcp
|
||||
|
||||
import "net/http"
|
||||
|
||||
/* Prompts */
|
||||
|
||||
// ListPromptsRequest is sent from the client to request a list of prompts and
|
||||
// prompt templates the server has.
|
||||
type ListPromptsRequest struct {
|
||||
PaginatedRequest
|
||||
Header http.Header `json:"-"`
|
||||
}
|
||||
|
||||
// ListPromptsResult is the server's response to a prompts/list request from
|
||||
// the client.
|
||||
type ListPromptsResult struct {
|
||||
PaginatedResult
|
||||
Prompts []Prompt `json:"prompts"`
|
||||
}
|
||||
|
||||
// GetPromptRequest is used by the client to get a prompt provided by the
|
||||
// server.
|
||||
type GetPromptRequest struct {
|
||||
Request
|
||||
Params GetPromptParams `json:"params"`
|
||||
Header http.Header `json:"-"`
|
||||
}
|
||||
|
||||
type GetPromptParams struct {
|
||||
// The name of the prompt or prompt template.
|
||||
Name string `json:"name"`
|
||||
// Arguments to use for templating the prompt.
|
||||
Arguments map[string]string `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
// GetPromptResult is the server's response to a prompts/get request from the
|
||||
// client.
|
||||
type GetPromptResult struct {
|
||||
Result
|
||||
// An optional description for the prompt.
|
||||
Description string `json:"description,omitempty"`
|
||||
Messages []PromptMessage `json:"messages"`
|
||||
}
|
||||
|
||||
// Prompt represents a prompt or prompt template that the server offers.
|
||||
// If Arguments is non-nil and non-empty, this indicates the prompt is a template
|
||||
// that requires argument values to be provided when calling prompts/get.
|
||||
// If Arguments is nil or empty, this is a static prompt that takes no arguments.
|
||||
type Prompt struct {
|
||||
// Meta is a metadata object that is reserved by MCP for storing additional information.
|
||||
Meta *Meta `json:"_meta,omitempty"`
|
||||
// The name of the prompt or prompt template.
|
||||
Name string `json:"name"`
|
||||
// An optional description of what this prompt provides
|
||||
Description string `json:"description,omitempty"`
|
||||
// A list of arguments to use for templating the prompt.
|
||||
// The presence of arguments indicates this is a template prompt.
|
||||
Arguments []PromptArgument `json:"arguments,omitempty"`
|
||||
}
|
||||
|
||||
// GetName returns the name of the prompt.
|
||||
func (p Prompt) GetName() string {
|
||||
return p.Name
|
||||
}
|
||||
|
||||
// PromptArgument describes an argument that a prompt template can accept.
|
||||
// When a prompt includes arguments, clients must provide values for all
|
||||
// required arguments when making a prompts/get request.
|
||||
type PromptArgument struct {
|
||||
// The name of the argument.
|
||||
Name string `json:"name"`
|
||||
// A human-readable description of the argument.
|
||||
Description string `json:"description,omitempty"`
|
||||
// Whether this argument must be provided.
|
||||
// If true, clients must include this argument when calling prompts/get.
|
||||
Required bool `json:"required,omitempty"`
|
||||
}
|
||||
|
||||
// Role represents the sender or recipient of messages and data in a
|
||||
// conversation.
|
||||
type Role string
|
||||
|
||||
const (
|
||||
RoleUser Role = "user"
|
||||
RoleAssistant Role = "assistant"
|
||||
)
|
||||
|
||||
// PromptMessage describes a message returned as part of a prompt.
|
||||
//
|
||||
// This is similar to `SamplingMessage`, but also supports the embedding of
|
||||
// resources from the MCP server.
|
||||
type PromptMessage struct {
|
||||
Role Role `json:"role"`
|
||||
Content Content `json:"content"` // Can be TextContent, ImageContent, AudioContent or EmbeddedResource
|
||||
}
|
||||
|
||||
// PromptListChangedNotification is an optional notification from the server
|
||||
// to the client, informing it that the list of prompts it offers has changed. This
|
||||
// may be issued by servers without any previous subscription from the client.
|
||||
type PromptListChangedNotification struct {
|
||||
Notification
|
||||
}
|
||||
|
||||
// PromptOption is a function that configures a Prompt.
|
||||
// It provides a flexible way to set various properties of a Prompt using the functional options pattern.
|
||||
type PromptOption func(*Prompt)
|
||||
|
||||
// ArgumentOption is a function that configures a PromptArgument.
|
||||
// It allows for flexible configuration of prompt arguments using the functional options pattern.
|
||||
type ArgumentOption func(*PromptArgument)
|
||||
|
||||
//
|
||||
// Core Prompt Functions
|
||||
//
|
||||
|
||||
// NewPrompt creates a new Prompt with the given name and options.
|
||||
// The prompt will be configured based on the provided options.
|
||||
// Options are applied in order, allowing for flexible prompt configuration.
|
||||
func NewPrompt(name string, opts ...PromptOption) Prompt {
|
||||
prompt := Prompt{
|
||||
Name: name,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&prompt)
|
||||
}
|
||||
|
||||
return prompt
|
||||
}
|
||||
|
||||
// WithPromptDescription adds a description to the Prompt.
|
||||
// The description should provide a clear, human-readable explanation of what the prompt does.
|
||||
func WithPromptDescription(description string) PromptOption {
|
||||
return func(p *Prompt) {
|
||||
p.Description = description
|
||||
}
|
||||
}
|
||||
|
||||
// WithArgument adds an argument to the prompt's argument list.
|
||||
// The argument will be configured based on the provided options.
|
||||
func WithArgument(name string, opts ...ArgumentOption) PromptOption {
|
||||
return func(p *Prompt) {
|
||||
arg := PromptArgument{
|
||||
Name: name,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&arg)
|
||||
}
|
||||
|
||||
if p.Arguments == nil {
|
||||
p.Arguments = make([]PromptArgument, 0)
|
||||
}
|
||||
p.Arguments = append(p.Arguments, arg)
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Argument Options
|
||||
//
|
||||
|
||||
// ArgumentDescription adds a description to a prompt argument.
|
||||
// The description should explain the purpose and expected values of the argument.
|
||||
func ArgumentDescription(desc string) ArgumentOption {
|
||||
return func(arg *PromptArgument) {
|
||||
arg.Description = desc
|
||||
}
|
||||
}
|
||||
|
||||
// RequiredArgument marks an argument as required in the prompt.
|
||||
// Required arguments must be provided when getting the prompt.
|
||||
func RequiredArgument() ArgumentOption {
|
||||
return func(arg *PromptArgument) {
|
||||
arg.Required = true
|
||||
}
|
||||
}
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
package mcp
|
||||
|
||||
import "github.com/yosida95/uritemplate/v3"
|
||||
|
||||
// ResourceOption is a function that configures a Resource.
|
||||
// It provides a flexible way to set various properties of a Resource using the functional options pattern.
|
||||
type ResourceOption func(*Resource)
|
||||
|
||||
// NewResource creates a new Resource with the given URI, name and options.
|
||||
// The resource will be configured based on the provided options.
|
||||
// Options are applied in order, allowing for flexible resource configuration.
|
||||
func NewResource(uri string, name string, opts ...ResourceOption) Resource {
|
||||
resource := Resource{
|
||||
URI: uri,
|
||||
Name: name,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&resource)
|
||||
}
|
||||
|
||||
return resource
|
||||
}
|
||||
|
||||
// WithResourceDescription adds a description to the Resource.
|
||||
// The description should provide a clear, human-readable explanation of what the resource represents.
|
||||
func WithResourceDescription(description string) ResourceOption {
|
||||
return func(r *Resource) {
|
||||
r.Description = description
|
||||
}
|
||||
}
|
||||
|
||||
// WithMIMEType sets the MIME type for the Resource.
|
||||
// This should indicate the format of the resource's contents.
|
||||
func WithMIMEType(mimeType string) ResourceOption {
|
||||
return func(r *Resource) {
|
||||
r.MIMEType = mimeType
|
||||
}
|
||||
}
|
||||
|
||||
// WithAnnotations adds annotations to the Resource.
|
||||
// Annotations can provide additional metadata about the resource's intended use.
|
||||
func WithAnnotations(audience []Role, priority float64) ResourceOption {
|
||||
return func(r *Resource) {
|
||||
if r.Annotations == nil {
|
||||
r.Annotations = &Annotations{}
|
||||
}
|
||||
r.Annotations.Audience = audience
|
||||
r.Annotations.Priority = priority
|
||||
}
|
||||
}
|
||||
|
||||
// ResourceTemplateOption is a function that configures a ResourceTemplate.
|
||||
// It provides a flexible way to set various properties of a ResourceTemplate using the functional options pattern.
|
||||
type ResourceTemplateOption func(*ResourceTemplate)
|
||||
|
||||
// NewResourceTemplate creates a new ResourceTemplate with the given URI template, name and options.
|
||||
// The template will be configured based on the provided options.
|
||||
// Options are applied in order, allowing for flexible template configuration.
|
||||
func NewResourceTemplate(uriTemplate string, name string, opts ...ResourceTemplateOption) ResourceTemplate {
|
||||
template := ResourceTemplate{
|
||||
URITemplate: &URITemplate{Template: uritemplate.MustNew(uriTemplate)},
|
||||
Name: name,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&template)
|
||||
}
|
||||
|
||||
return template
|
||||
}
|
||||
|
||||
// WithTemplateDescription adds a description to the ResourceTemplate.
|
||||
// The description should provide a clear, human-readable explanation of what resources this template represents.
|
||||
func WithTemplateDescription(description string) ResourceTemplateOption {
|
||||
return func(t *ResourceTemplate) {
|
||||
t.Description = description
|
||||
}
|
||||
}
|
||||
|
||||
// WithTemplateMIMEType sets the MIME type for the ResourceTemplate.
|
||||
// This should only be set if all resources matching this template will have the same type.
|
||||
func WithTemplateMIMEType(mimeType string) ResourceTemplateOption {
|
||||
return func(t *ResourceTemplate) {
|
||||
t.MIMEType = mimeType
|
||||
}
|
||||
}
|
||||
|
||||
// WithTemplateAnnotations adds annotations to the ResourceTemplate.
|
||||
// Annotations can provide additional metadata about the template's intended use.
|
||||
func WithTemplateAnnotations(audience []Role, priority float64) ResourceTemplateOption {
|
||||
return func(t *ResourceTemplate) {
|
||||
if t.Annotations == nil {
|
||||
t.Annotations = &Annotations{}
|
||||
}
|
||||
t.Annotations.Audience = audience
|
||||
t.Annotations.Priority = priority
|
||||
}
|
||||
}
|
||||
+1331
File diff suppressed because it is too large.
Load diff
+42
@@ -0,0 +1,42 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// TypedToolHandlerFunc is a function that handles a tool call with typed arguments
|
||||
type TypedToolHandlerFunc[T any] func(ctx context.Context, request CallToolRequest, args T) (*CallToolResult, error)
|
||||
|
||||
// StructuredToolHandlerFunc is a function that handles a tool call with typed arguments and returns structured output
|
||||
type StructuredToolHandlerFunc[TArgs any, TResult any] func(ctx context.Context, request CallToolRequest, args TArgs) (TResult, error)
|
||||
|
||||
// NewTypedToolHandler creates a ToolHandlerFunc that automatically binds arguments to a typed struct
|
||||
func NewTypedToolHandler[T any](handler TypedToolHandlerFunc[T]) func(ctx context.Context, request CallToolRequest) (*CallToolResult, error) {
|
||||
return func(ctx context.Context, request CallToolRequest) (*CallToolResult, error) {
|
||||
var args T
|
||||
if err := request.BindArguments(&args); err != nil {
|
||||
return NewToolResultError(fmt.Sprintf("failed to bind arguments: %v", err)), nil
|
||||
}
|
||||
return handler(ctx, request, args)
|
||||
}
|
||||
}
|
||||
|
||||
// NewStructuredToolHandler creates a ToolHandlerFunc that automatically binds arguments to a typed struct
|
||||
// and returns structured output. It automatically creates both structured and
|
||||
// text content (from the structured output) for backwards compatibility.
|
||||
func NewStructuredToolHandler[TArgs any, TResult any](handler StructuredToolHandlerFunc[TArgs, TResult]) func(ctx context.Context, request CallToolRequest) (*CallToolResult, error) {
|
||||
return func(ctx context.Context, request CallToolRequest) (*CallToolResult, error) {
|
||||
var args TArgs
|
||||
if err := request.BindArguments(&args); err != nil {
|
||||
return NewToolResultError(fmt.Sprintf("failed to bind arguments: %v", err)), nil
|
||||
}
|
||||
|
||||
result, err := handler(ctx, request, args)
|
||||
if err != nil {
|
||||
return NewToolResultError(fmt.Sprintf("tool execution failed: %v", err)), nil
|
||||
}
|
||||
|
||||
return NewToolResultStructuredOnly(result), nil
|
||||
}
|
||||
}
|
||||
+1252
File diff suppressed because it is too large.
Load diff
+979
@@ -0,0 +1,979 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/spf13/cast"
|
||||
)
|
||||
|
||||
// ClientRequest types
|
||||
var (
|
||||
_ ClientRequest = (*PingRequest)(nil)
|
||||
_ ClientRequest = (*InitializeRequest)(nil)
|
||||
_ ClientRequest = (*CompleteRequest)(nil)
|
||||
_ ClientRequest = (*SetLevelRequest)(nil)
|
||||
_ ClientRequest = (*GetPromptRequest)(nil)
|
||||
_ ClientRequest = (*ListPromptsRequest)(nil)
|
||||
_ ClientRequest = (*ListResourcesRequest)(nil)
|
||||
_ ClientRequest = (*ReadResourceRequest)(nil)
|
||||
_ ClientRequest = (*SubscribeRequest)(nil)
|
||||
_ ClientRequest = (*UnsubscribeRequest)(nil)
|
||||
_ ClientRequest = (*CallToolRequest)(nil)
|
||||
_ ClientRequest = (*ListToolsRequest)(nil)
|
||||
)
|
||||
|
||||
// ClientNotification types
|
||||
var (
|
||||
_ ClientNotification = (*CancelledNotification)(nil)
|
||||
_ ClientNotification = (*ProgressNotification)(nil)
|
||||
_ ClientNotification = (*InitializedNotification)(nil)
|
||||
_ ClientNotification = (*RootsListChangedNotification)(nil)
|
||||
)
|
||||
|
||||
// ClientResult types
|
||||
var (
|
||||
_ ClientResult = (*EmptyResult)(nil)
|
||||
_ ClientResult = (*CreateMessageResult)(nil)
|
||||
_ ClientResult = (*ListRootsResult)(nil)
|
||||
)
|
||||
|
||||
// ServerRequest types
|
||||
var (
|
||||
_ ServerRequest = (*PingRequest)(nil)
|
||||
_ ServerRequest = (*CreateMessageRequest)(nil)
|
||||
_ ServerRequest = (*ListRootsRequest)(nil)
|
||||
)
|
||||
|
||||
// ServerNotification types
|
||||
var (
|
||||
_ ServerNotification = (*CancelledNotification)(nil)
|
||||
_ ServerNotification = (*ProgressNotification)(nil)
|
||||
_ ServerNotification = (*LoggingMessageNotification)(nil)
|
||||
_ ServerNotification = (*ResourceUpdatedNotification)(nil)
|
||||
_ ServerNotification = (*ResourceListChangedNotification)(nil)
|
||||
_ ServerNotification = (*ToolListChangedNotification)(nil)
|
||||
_ ServerNotification = (*PromptListChangedNotification)(nil)
|
||||
)
|
||||
|
||||
// ServerResult types
|
||||
var (
|
||||
_ ServerResult = (*EmptyResult)(nil)
|
||||
_ ServerResult = (*InitializeResult)(nil)
|
||||
_ ServerResult = (*CompleteResult)(nil)
|
||||
_ ServerResult = (*GetPromptResult)(nil)
|
||||
_ ServerResult = (*ListPromptsResult)(nil)
|
||||
_ ServerResult = (*ListResourcesResult)(nil)
|
||||
_ ServerResult = (*ReadResourceResult)(nil)
|
||||
_ ServerResult = (*CallToolResult)(nil)
|
||||
_ ServerResult = (*ListToolsResult)(nil)
|
||||
)
|
||||
|
||||
// Helper functions for type assertions
|
||||
|
||||
// asType attempts to cast the given interface to the given type
|
||||
func asType[T any](content any) (*T, bool) {
|
||||
tc, ok := content.(T)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
return &tc, true
|
||||
}
|
||||
|
||||
// AsTextContent attempts to cast the given interface to TextContent
|
||||
func AsTextContent(content any) (*TextContent, bool) {
|
||||
return asType[TextContent](content)
|
||||
}
|
||||
|
||||
// AsImageContent attempts to cast the given interface to ImageContent
|
||||
func AsImageContent(content any) (*ImageContent, bool) {
|
||||
return asType[ImageContent](content)
|
||||
}
|
||||
|
||||
// AsAudioContent attempts to cast the given interface to AudioContent
|
||||
func AsAudioContent(content any) (*AudioContent, bool) {
|
||||
return asType[AudioContent](content)
|
||||
}
|
||||
|
||||
// AsEmbeddedResource attempts to cast the given interface to EmbeddedResource
|
||||
func AsEmbeddedResource(content any) (*EmbeddedResource, bool) {
|
||||
return asType[EmbeddedResource](content)
|
||||
}
|
||||
|
||||
// AsTextResourceContents attempts to cast the given interface to TextResourceContents
|
||||
func AsTextResourceContents(content any) (*TextResourceContents, bool) {
|
||||
return asType[TextResourceContents](content)
|
||||
}
|
||||
|
||||
// AsBlobResourceContents attempts to cast the given interface to BlobResourceContents
|
||||
func AsBlobResourceContents(content any) (*BlobResourceContents, bool) {
|
||||
return asType[BlobResourceContents](content)
|
||||
}
|
||||
|
||||
// Helper function for JSON-RPC
|
||||
|
||||
// NewJSONRPCResponse creates a new JSONRPCResponse with the given id and result.
|
||||
// NOTE: This function expects a Result struct, but JSONRPCResponse.Result is typed as `any`.
|
||||
// The Result struct wraps the actual result data with optional metadata.
|
||||
// For direct result assignment, use NewJSONRPCResultResponse instead.
|
||||
func NewJSONRPCResponse(id RequestId, result Result) JSONRPCResponse {
|
||||
return JSONRPCResponse{
|
||||
JSONRPC: JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Result: result,
|
||||
}
|
||||
}
|
||||
|
||||
// NewJSONRPCResultResponse creates a new JSONRPCResponse with the given id and result.
|
||||
// This function accepts any type for the result, matching the JSONRPCResponse.Result field type.
|
||||
func NewJSONRPCResultResponse(id RequestId, result any) JSONRPCResponse {
|
||||
return JSONRPCResponse{
|
||||
JSONRPC: JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Result: result,
|
||||
}
|
||||
}
|
||||
|
||||
// NewJSONRPCErrorDetails creates a new JSONRPCErrorDetails with the given code, message, and data.
|
||||
func NewJSONRPCErrorDetails(code int, message string, data any) JSONRPCErrorDetails {
|
||||
return JSONRPCErrorDetails{
|
||||
Code: code,
|
||||
Message: message,
|
||||
Data: data,
|
||||
}
|
||||
}
|
||||
|
||||
// NewJSONRPCError creates a new JSONRPCResponse with the given id, code, and message
|
||||
func NewJSONRPCError(
|
||||
id RequestId,
|
||||
code int,
|
||||
message string,
|
||||
data any,
|
||||
) JSONRPCError {
|
||||
return JSONRPCError{
|
||||
JSONRPC: JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Error: NewJSONRPCErrorDetails(code, message, data),
|
||||
}
|
||||
}
|
||||
|
||||
// NewProgressNotification
|
||||
// Helper function for creating a progress notification
|
||||
func NewProgressNotification(
|
||||
token ProgressToken,
|
||||
progress float64,
|
||||
total *float64,
|
||||
message *string,
|
||||
) ProgressNotification {
|
||||
notification := ProgressNotification{
|
||||
Notification: Notification{
|
||||
Method: "notifications/progress",
|
||||
},
|
||||
Params: struct {
|
||||
ProgressToken ProgressToken `json:"progressToken"`
|
||||
Progress float64 `json:"progress"`
|
||||
Total float64 `json:"total,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
}{
|
||||
ProgressToken: token,
|
||||
Progress: progress,
|
||||
},
|
||||
}
|
||||
if total != nil {
|
||||
notification.Params.Total = *total
|
||||
}
|
||||
if message != nil {
|
||||
notification.Params.Message = *message
|
||||
}
|
||||
return notification
|
||||
}
|
||||
|
||||
// NewLoggingMessageNotification
|
||||
// Helper function for creating a logging message notification
|
||||
func NewLoggingMessageNotification(
|
||||
level LoggingLevel,
|
||||
logger string,
|
||||
data any,
|
||||
) LoggingMessageNotification {
|
||||
return LoggingMessageNotification{
|
||||
Notification: Notification{
|
||||
Method: "notifications/message",
|
||||
},
|
||||
Params: struct {
|
||||
Level LoggingLevel `json:"level"`
|
||||
Logger string `json:"logger,omitempty"`
|
||||
Data any `json:"data"`
|
||||
}{
|
||||
Level: level,
|
||||
Logger: logger,
|
||||
Data: data,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewPromptMessage
|
||||
// Helper function to create a new PromptMessage
|
||||
func NewPromptMessage(role Role, content Content) PromptMessage {
|
||||
return PromptMessage{
|
||||
Role: role,
|
||||
Content: content,
|
||||
}
|
||||
}
|
||||
|
||||
// NewTextContent
|
||||
// Helper function to create a new TextContent
|
||||
func NewTextContent(text string) TextContent {
|
||||
return TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: text,
|
||||
}
|
||||
}
|
||||
|
||||
// NewImageContent
|
||||
// Helper function to create a new ImageContent
|
||||
func NewImageContent(data, mimeType string) ImageContent {
|
||||
return ImageContent{
|
||||
Type: ContentTypeImage,
|
||||
Data: data,
|
||||
MIMEType: mimeType,
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to create a new AudioContent
|
||||
func NewAudioContent(data, mimeType string) AudioContent {
|
||||
return AudioContent{
|
||||
Type: ContentTypeAudio,
|
||||
Data: data,
|
||||
MIMEType: mimeType,
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to create a new ResourceLink
|
||||
func NewResourceLink(uri, name, description, mimeType string) ResourceLink {
|
||||
return ResourceLink{
|
||||
Type: ContentTypeLink,
|
||||
URI: uri,
|
||||
Name: name,
|
||||
Description: description,
|
||||
MIMEType: mimeType,
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to create a new EmbeddedResource
|
||||
func NewEmbeddedResource(resource ResourceContents) EmbeddedResource {
|
||||
return EmbeddedResource{
|
||||
Type: ContentTypeResource,
|
||||
Resource: resource,
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultText creates a new CallToolResult with a text content
|
||||
func NewToolResultText(text string) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: text,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultJSON creates a new CallToolResult with a JSON content.
|
||||
func NewToolResultJSON[T any](data T) (*CallToolResult, error) {
|
||||
b, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unable to marshal JSON: %w", err)
|
||||
}
|
||||
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: string(b),
|
||||
},
|
||||
},
|
||||
StructuredContent: data,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// NewToolResultStructured creates a new CallToolResult with structured content.
|
||||
// It includes both the structured content and a text representation for backward compatibility.
|
||||
func NewToolResultStructured(structured any, fallbackText string) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: "text",
|
||||
Text: fallbackText,
|
||||
},
|
||||
},
|
||||
StructuredContent: structured,
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultStructuredOnly creates a new CallToolResult with structured
|
||||
// content and creates a JSON string fallback for backwards compatibility.
|
||||
// This is useful when you want to provide structured data without any specific text fallback.
|
||||
func NewToolResultStructuredOnly(structured any) *CallToolResult {
|
||||
var fallbackText string
|
||||
// Convert to JSON string for backward compatibility
|
||||
jsonBytes, err := json.Marshal(structured)
|
||||
if err != nil {
|
||||
fallbackText = fmt.Sprintf("Error serializing structured content: %v", err)
|
||||
} else {
|
||||
fallbackText = string(jsonBytes)
|
||||
}
|
||||
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: "text",
|
||||
Text: fallbackText,
|
||||
},
|
||||
},
|
||||
StructuredContent: structured,
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultImage creates a new CallToolResult with both text and image content
|
||||
func NewToolResultImage(text, imageData, mimeType string) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: text,
|
||||
},
|
||||
ImageContent{
|
||||
Type: ContentTypeImage,
|
||||
Data: imageData,
|
||||
MIMEType: mimeType,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultAudio creates a new CallToolResult with both text and audio content
|
||||
func NewToolResultAudio(text, imageData, mimeType string) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: text,
|
||||
},
|
||||
AudioContent{
|
||||
Type: ContentTypeAudio,
|
||||
Data: imageData,
|
||||
MIMEType: mimeType,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultResource creates a new CallToolResult with an embedded resource
|
||||
func NewToolResultResource(
|
||||
text string,
|
||||
resource ResourceContents,
|
||||
) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: text,
|
||||
},
|
||||
EmbeddedResource{
|
||||
Type: ContentTypeResource,
|
||||
Resource: resource,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultError creates a new CallToolResult with an error message.
|
||||
// Any errors that originate from the tool SHOULD be reported inside the result object.
|
||||
func NewToolResultError(text string) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: text,
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultErrorFromErr creates a new CallToolResult with an error message.
|
||||
// If an error is provided, its details will be appended to the text message.
|
||||
// Any errors that originate from the tool SHOULD be reported inside the result object.
|
||||
func NewToolResultErrorFromErr(text string, err error) *CallToolResult {
|
||||
if err != nil {
|
||||
text = fmt.Sprintf("%s: %v", text, err)
|
||||
}
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: text,
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
// NewToolResultErrorf creates a new CallToolResult with an error message.
|
||||
// The error message is formatted using the fmt package.
|
||||
// Any errors that originate from the tool SHOULD be reported inside the result object.
|
||||
func NewToolResultErrorf(format string, a ...any) *CallToolResult {
|
||||
return &CallToolResult{
|
||||
Content: []Content{
|
||||
TextContent{
|
||||
Type: ContentTypeText,
|
||||
Text: fmt.Sprintf(format, a...),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}
|
||||
}
|
||||
|
||||
// NewListResourcesResult creates a new ListResourcesResult
|
||||
func NewListResourcesResult(
|
||||
resources []Resource,
|
||||
nextCursor Cursor,
|
||||
) *ListResourcesResult {
|
||||
return &ListResourcesResult{
|
||||
PaginatedResult: PaginatedResult{
|
||||
NextCursor: nextCursor,
|
||||
},
|
||||
Resources: resources,
|
||||
}
|
||||
}
|
||||
|
||||
// NewListResourceTemplatesResult creates a new ListResourceTemplatesResult
|
||||
func NewListResourceTemplatesResult(
|
||||
templates []ResourceTemplate,
|
||||
nextCursor Cursor,
|
||||
) *ListResourceTemplatesResult {
|
||||
return &ListResourceTemplatesResult{
|
||||
PaginatedResult: PaginatedResult{
|
||||
NextCursor: nextCursor,
|
||||
},
|
||||
ResourceTemplates: templates,
|
||||
}
|
||||
}
|
||||
|
||||
// NewReadResourceResult creates a new ReadResourceResult with text content
|
||||
func NewReadResourceResult(text string) *ReadResourceResult {
|
||||
return &ReadResourceResult{
|
||||
Contents: []ResourceContents{
|
||||
TextResourceContents{
|
||||
Text: text,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// NewListPromptsResult creates a new ListPromptsResult
|
||||
func NewListPromptsResult(
|
||||
prompts []Prompt,
|
||||
nextCursor Cursor,
|
||||
) *ListPromptsResult {
|
||||
return &ListPromptsResult{
|
||||
PaginatedResult: PaginatedResult{
|
||||
NextCursor: nextCursor,
|
||||
},
|
||||
Prompts: prompts,
|
||||
}
|
||||
}
|
||||
|
||||
// NewGetPromptResult creates a new GetPromptResult
|
||||
func NewGetPromptResult(
|
||||
description string,
|
||||
messages []PromptMessage,
|
||||
) *GetPromptResult {
|
||||
return &GetPromptResult{
|
||||
Description: description,
|
||||
Messages: messages,
|
||||
}
|
||||
}
|
||||
|
||||
// NewListToolsResult creates a new ListToolsResult
|
||||
func NewListToolsResult(tools []Tool, nextCursor Cursor) *ListToolsResult {
|
||||
return &ListToolsResult{
|
||||
PaginatedResult: PaginatedResult{
|
||||
NextCursor: nextCursor,
|
||||
},
|
||||
Tools: tools,
|
||||
}
|
||||
}
|
||||
|
||||
// NewInitializeResult creates a new InitializeResult
|
||||
func NewInitializeResult(
|
||||
protocolVersion string,
|
||||
capabilities ServerCapabilities,
|
||||
serverInfo Implementation,
|
||||
instructions string,
|
||||
) *InitializeResult {
|
||||
return &InitializeResult{
|
||||
ProtocolVersion: protocolVersion,
|
||||
Capabilities: capabilities,
|
||||
ServerInfo: serverInfo,
|
||||
Instructions: instructions,
|
||||
}
|
||||
}
|
||||
|
||||
// FormatNumberResult
|
||||
// Helper for formatting numbers in tool results
|
||||
func FormatNumberResult(value float64) *CallToolResult {
|
||||
return NewToolResultText(fmt.Sprintf("%.2f", value))
|
||||
}
|
||||
|
||||
func ExtractString(data map[string]any, key string) string {
|
||||
if value, ok := data[key]; ok {
|
||||
if str, ok := value.(string); ok {
|
||||
return str
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func ParseAnnotations(data map[string]any) *Annotations {
|
||||
if data == nil {
|
||||
return nil
|
||||
}
|
||||
annotations := &Annotations{}
|
||||
if value, ok := data["priority"]; ok {
|
||||
annotations.Priority = cast.ToFloat64(value)
|
||||
}
|
||||
|
||||
if value, ok := data["audience"]; ok {
|
||||
for _, a := range cast.ToStringSlice(value) {
|
||||
a := Role(a)
|
||||
if a == RoleUser || a == RoleAssistant {
|
||||
annotations.Audience = append(annotations.Audience, a)
|
||||
}
|
||||
}
|
||||
}
|
||||
return annotations
|
||||
|
||||
}
|
||||
|
||||
func ExtractMap(data map[string]any, key string) map[string]any {
|
||||
if value, ok := data[key]; ok {
|
||||
if m, ok := value.(map[string]any); ok {
|
||||
return m
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func ParseContent(contentMap map[string]any) (Content, error) {
|
||||
contentType := ExtractString(contentMap, "type")
|
||||
|
||||
var annotations *Annotations
|
||||
if annotationsMap := ExtractMap(contentMap, "annotations"); annotationsMap != nil {
|
||||
annotations = ParseAnnotations(annotationsMap)
|
||||
}
|
||||
|
||||
switch contentType {
|
||||
case ContentTypeText:
|
||||
text := ExtractString(contentMap, "text")
|
||||
c := NewTextContent(text)
|
||||
c.Annotations = annotations
|
||||
return c, nil
|
||||
|
||||
case ContentTypeImage:
|
||||
data := ExtractString(contentMap, "data")
|
||||
mimeType := ExtractString(contentMap, "mimeType")
|
||||
if data == "" || mimeType == "" {
|
||||
return nil, fmt.Errorf("image data or mimeType is missing")
|
||||
}
|
||||
c := NewImageContent(data, mimeType)
|
||||
c.Annotations = annotations
|
||||
return c, nil
|
||||
|
||||
case ContentTypeAudio:
|
||||
data := ExtractString(contentMap, "data")
|
||||
mimeType := ExtractString(contentMap, "mimeType")
|
||||
if data == "" || mimeType == "" {
|
||||
return nil, fmt.Errorf("audio data or mimeType is missing")
|
||||
}
|
||||
c := NewAudioContent(data, mimeType)
|
||||
c.Annotations = annotations
|
||||
return c, nil
|
||||
|
||||
case ContentTypeLink:
|
||||
uri := ExtractString(contentMap, "uri")
|
||||
name := ExtractString(contentMap, "name")
|
||||
description := ExtractString(contentMap, "description")
|
||||
mimeType := ExtractString(contentMap, "mimeType")
|
||||
if uri == "" || name == "" {
|
||||
return nil, fmt.Errorf("resource_link uri or name is missing")
|
||||
}
|
||||
c := NewResourceLink(uri, name, description, mimeType)
|
||||
c.Annotations = annotations
|
||||
return c, nil
|
||||
|
||||
case ContentTypeResource:
|
||||
resourceMap := ExtractMap(contentMap, "resource")
|
||||
if resourceMap == nil {
|
||||
return nil, fmt.Errorf("resource is missing")
|
||||
}
|
||||
|
||||
resourceContents, err := ParseResourceContents(resourceMap)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
c := NewEmbeddedResource(resourceContents)
|
||||
c.Annotations = annotations
|
||||
return c, nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported content type: %s", contentType)
|
||||
}
|
||||
|
||||
func ParseGetPromptResult(rawMessage *json.RawMessage) (*GetPromptResult, error) {
|
||||
if rawMessage == nil {
|
||||
return nil, fmt.Errorf("response is nil")
|
||||
}
|
||||
|
||||
var jsonContent map[string]any
|
||||
if err := json.Unmarshal(*rawMessage, &jsonContent); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
result := GetPromptResult{}
|
||||
|
||||
meta, ok := jsonContent["_meta"]
|
||||
if ok {
|
||||
if metaMap, ok := meta.(map[string]any); ok {
|
||||
result.Meta = NewMetaFromMap(metaMap)
|
||||
}
|
||||
}
|
||||
|
||||
description, ok := jsonContent["description"]
|
||||
if ok {
|
||||
if descriptionStr, ok := description.(string); ok {
|
||||
result.Description = descriptionStr
|
||||
}
|
||||
}
|
||||
|
||||
messages, ok := jsonContent["messages"]
|
||||
if ok {
|
||||
messagesArr, ok := messages.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("messages is not an array")
|
||||
}
|
||||
|
||||
for _, message := range messagesArr {
|
||||
messageMap, ok := message.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("message is not an object")
|
||||
}
|
||||
|
||||
// Extract role
|
||||
roleStr := ExtractString(messageMap, "role")
|
||||
if roleStr == "" || (roleStr != string(RoleAssistant) && roleStr != string(RoleUser)) {
|
||||
return nil, fmt.Errorf("unsupported role: %s", roleStr)
|
||||
}
|
||||
|
||||
// Extract content
|
||||
contentMap, ok := messageMap["content"].(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("content is not an object")
|
||||
}
|
||||
|
||||
// Process content
|
||||
content, err := ParseContent(contentMap)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Append processed message
|
||||
result.Messages = append(result.Messages, NewPromptMessage(Role(roleStr), content))
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
func ParseCallToolResult(rawMessage *json.RawMessage) (*CallToolResult, error) {
|
||||
if rawMessage == nil {
|
||||
return nil, fmt.Errorf("response is nil")
|
||||
}
|
||||
|
||||
var jsonContent map[string]any
|
||||
if err := json.Unmarshal(*rawMessage, &jsonContent); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
var result CallToolResult
|
||||
|
||||
meta, ok := jsonContent["_meta"]
|
||||
if ok {
|
||||
if metaMap, ok := meta.(map[string]any); ok {
|
||||
result.Meta = NewMetaFromMap(metaMap)
|
||||
}
|
||||
}
|
||||
|
||||
isError, ok := jsonContent["isError"]
|
||||
if ok {
|
||||
if isErrorBool, ok := isError.(bool); ok {
|
||||
result.IsError = isErrorBool
|
||||
}
|
||||
}
|
||||
|
||||
contents, ok := jsonContent["content"]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("content is missing")
|
||||
}
|
||||
|
||||
contentArr, ok := contents.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("content is not an array")
|
||||
}
|
||||
|
||||
for _, content := range contentArr {
|
||||
// Extract content
|
||||
contentMap, ok := content.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("content is not an object")
|
||||
}
|
||||
|
||||
// Process content
|
||||
content, err := ParseContent(contentMap)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
result.Content = append(result.Content, content)
|
||||
}
|
||||
|
||||
// Handle structured content
|
||||
structuredContent, ok := jsonContent["structuredContent"]
|
||||
if ok {
|
||||
result.StructuredContent = structuredContent
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
func ParseResourceContents(contentMap map[string]any) (ResourceContents, error) {
|
||||
uri := ExtractString(contentMap, "uri")
|
||||
if uri == "" {
|
||||
return nil, fmt.Errorf("resource uri is missing")
|
||||
}
|
||||
|
||||
mimeType := ExtractString(contentMap, "mimeType")
|
||||
|
||||
meta := ExtractMap(contentMap, "_meta")
|
||||
|
||||
if _, present := contentMap["_meta"]; present && meta == nil {
|
||||
return nil, fmt.Errorf("_meta must be an object")
|
||||
}
|
||||
|
||||
if text := ExtractString(contentMap, "text"); text != "" {
|
||||
return TextResourceContents{
|
||||
Meta: meta,
|
||||
URI: uri,
|
||||
MIMEType: mimeType,
|
||||
Text: text,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if blob := ExtractString(contentMap, "blob"); blob != "" {
|
||||
return BlobResourceContents{
|
||||
Meta: meta,
|
||||
URI: uri,
|
||||
MIMEType: mimeType,
|
||||
Blob: blob,
|
||||
}, nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported resource type")
|
||||
}
|
||||
|
||||
func ParseReadResourceResult(rawMessage *json.RawMessage) (*ReadResourceResult, error) {
|
||||
if rawMessage == nil {
|
||||
return nil, fmt.Errorf("response is nil")
|
||||
}
|
||||
|
||||
var jsonContent map[string]any
|
||||
if err := json.Unmarshal(*rawMessage, &jsonContent); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
|
||||
}
|
||||
|
||||
var result ReadResourceResult
|
||||
|
||||
meta, ok := jsonContent["_meta"]
|
||||
if ok {
|
||||
if metaMap, ok := meta.(map[string]any); ok {
|
||||
result.Meta = NewMetaFromMap(metaMap)
|
||||
}
|
||||
}
|
||||
|
||||
contents, ok := jsonContent["contents"]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("contents is missing")
|
||||
}
|
||||
|
||||
contentArr, ok := contents.([]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("contents is not an array")
|
||||
}
|
||||
|
||||
for _, content := range contentArr {
|
||||
// Extract content
|
||||
contentMap, ok := content.(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("content is not an object")
|
||||
}
|
||||
|
||||
// Process content
|
||||
content, err := ParseResourceContents(contentMap)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
result.Contents = append(result.Contents, content)
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
func ParseArgument(request CallToolRequest, key string, defaultVal any) any {
|
||||
args := request.GetArguments()
|
||||
if _, ok := args[key]; !ok {
|
||||
return defaultVal
|
||||
} else {
|
||||
return args[key]
|
||||
}
|
||||
}
|
||||
|
||||
// ParseBoolean extracts and converts a boolean parameter from a CallToolRequest.
|
||||
// If the key is not found in the Arguments map, the defaultValue is returned.
|
||||
// The function uses cast.ToBool for conversion which handles various string representations
|
||||
// such as "true", "yes", "1", etc.
|
||||
func ParseBoolean(request CallToolRequest, key string, defaultValue bool) bool {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToBool(v)
|
||||
}
|
||||
|
||||
// ParseInt64 extracts and converts an int64 parameter from a CallToolRequest.
|
||||
// If the key is not found in the Arguments map, the defaultValue is returned.
|
||||
func ParseInt64(request CallToolRequest, key string, defaultValue int64) int64 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToInt64(v)
|
||||
}
|
||||
|
||||
// ParseInt32 extracts and converts an int32 parameter from a CallToolRequest.
|
||||
func ParseInt32(request CallToolRequest, key string, defaultValue int32) int32 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToInt32(v)
|
||||
}
|
||||
|
||||
// ParseInt16 extracts and converts an int16 parameter from a CallToolRequest.
|
||||
func ParseInt16(request CallToolRequest, key string, defaultValue int16) int16 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToInt16(v)
|
||||
}
|
||||
|
||||
// ParseInt8 extracts and converts an int8 parameter from a CallToolRequest.
|
||||
func ParseInt8(request CallToolRequest, key string, defaultValue int8) int8 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToInt8(v)
|
||||
}
|
||||
|
||||
// ParseInt extracts and converts an int parameter from a CallToolRequest.
|
||||
func ParseInt(request CallToolRequest, key string, defaultValue int) int {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToInt(v)
|
||||
}
|
||||
|
||||
// ParseUInt extracts and converts an uint parameter from a CallToolRequest.
|
||||
func ParseUInt(request CallToolRequest, key string, defaultValue uint) uint {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToUint(v)
|
||||
}
|
||||
|
||||
// ParseUInt64 extracts and converts an uint64 parameter from a CallToolRequest.
|
||||
func ParseUInt64(request CallToolRequest, key string, defaultValue uint64) uint64 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToUint64(v)
|
||||
}
|
||||
|
||||
// ParseUInt32 extracts and converts an uint32 parameter from a CallToolRequest.
|
||||
func ParseUInt32(request CallToolRequest, key string, defaultValue uint32) uint32 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToUint32(v)
|
||||
}
|
||||
|
||||
// ParseUInt16 extracts and converts an uint16 parameter from a CallToolRequest.
|
||||
func ParseUInt16(request CallToolRequest, key string, defaultValue uint16) uint16 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToUint16(v)
|
||||
}
|
||||
|
||||
// ParseUInt8 extracts and converts an uint8 parameter from a CallToolRequest.
|
||||
func ParseUInt8(request CallToolRequest, key string, defaultValue uint8) uint8 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToUint8(v)
|
||||
}
|
||||
|
||||
// ParseFloat32 extracts and converts a float32 parameter from a CallToolRequest.
|
||||
func ParseFloat32(request CallToolRequest, key string, defaultValue float32) float32 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToFloat32(v)
|
||||
}
|
||||
|
||||
// ParseFloat64 extracts and converts a float64 parameter from a CallToolRequest.
|
||||
func ParseFloat64(request CallToolRequest, key string, defaultValue float64) float64 {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToFloat64(v)
|
||||
}
|
||||
|
||||
// ParseString extracts and converts a string parameter from a CallToolRequest.
|
||||
func ParseString(request CallToolRequest, key string, defaultValue string) string {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToString(v)
|
||||
}
|
||||
|
||||
// ParseStringMap extracts and converts a string map parameter from a CallToolRequest.
|
||||
func ParseStringMap(request CallToolRequest, key string, defaultValue map[string]any) map[string]any {
|
||||
v := ParseArgument(request, key, defaultValue)
|
||||
return cast.ToStringMap(v)
|
||||
}
|
||||
|
||||
// ToBoolPtr returns a pointer to the given boolean value
|
||||
func ToBoolPtr(b bool) *bool {
|
||||
return &b
|
||||
}
|
||||
|
||||
// GetTextFromContent extracts text from a Content interface that might be a TextContent struct
|
||||
// or a map[string]any that was unmarshaled from JSON. This is useful when dealing with content
|
||||
// that comes from different transport layers that may handle JSON differently.
|
||||
//
|
||||
// This function uses fallback behavior for non-text content - it returns a string representation
|
||||
// via fmt.Sprintf for any content that cannot be extracted as text. This is a lossy operation
|
||||
// intended for convenience in logging and display scenarios.
|
||||
//
|
||||
// For strict type validation, use ParseContent() instead, which returns an error for invalid content.
|
||||
func GetTextFromContent(content any) string {
|
||||
switch c := content.(type) {
|
||||
case TextContent:
|
||||
return c.Text
|
||||
case map[string]any:
|
||||
// Handle JSON unmarshaled content
|
||||
if contentType, exists := c["type"]; exists && contentType == "text" {
|
||||
if text, exists := c["text"].(string); exists {
|
||||
return text
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("%v", content)
|
||||
case string:
|
||||
return c
|
||||
default:
|
||||
return fmt.Sprintf("%v", content)
|
||||
}
|
||||
}
|
||||
+7
@@ -0,0 +1,7 @@
|
||||
package server
|
||||
|
||||
// Common HTTP header constants used across server transports
|
||||
const (
|
||||
HeaderKeySessionID = "Mcp-Session-Id"
|
||||
HeaderKeyProtocolVersion = "Mcp-Protocol-Version"
|
||||
)
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
package server
|
||||
|
||||
type contextKey int
|
||||
|
||||
const (
|
||||
// This const is used as key for context value lookup
|
||||
requestHeader contextKey = iota
|
||||
)
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrNoActiveSession is returned when there is no active session in the context
|
||||
ErrNoActiveSession = errors.New("no active session")
|
||||
// ErrElicitationNotSupported is returned when the session does not support elicitation
|
||||
ErrElicitationNotSupported = errors.New("session does not support elicitation")
|
||||
)
|
||||
|
||||
// RequestElicitation sends an elicitation request to the client.
|
||||
// The client must have declared elicitation capability during initialization.
|
||||
// The session must implement SessionWithElicitation to support this operation.
|
||||
func (s *MCPServer) RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) {
|
||||
session := ClientSessionFromContext(ctx)
|
||||
if session == nil {
|
||||
return nil, ErrNoActiveSession
|
||||
}
|
||||
|
||||
// Check if the session supports elicitation requests
|
||||
if elicitationSession, ok := session.(SessionWithElicitation); ok {
|
||||
return elicitationSession.RequestElicitation(ctx, request)
|
||||
}
|
||||
|
||||
return nil, ErrElicitationNotSupported
|
||||
}
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
var (
|
||||
// Common server errors
|
||||
ErrUnsupported = errors.New("not supported")
|
||||
ErrResourceNotFound = errors.New("resource not found")
|
||||
ErrPromptNotFound = errors.New("prompt not found")
|
||||
ErrToolNotFound = errors.New("tool not found")
|
||||
|
||||
// Session-related errors
|
||||
ErrSessionNotFound = errors.New("session not found")
|
||||
ErrSessionExists = errors.New("session already exists")
|
||||
ErrSessionNotInitialized = errors.New("session not properly initialized")
|
||||
ErrSessionDoesNotSupportTools = errors.New("session does not support per-session tools")
|
||||
ErrSessionDoesNotSupportResources = errors.New("session does not support per-session resources")
|
||||
ErrSessionDoesNotSupportResourceTemplates = errors.New("session does not support resource templates")
|
||||
ErrSessionDoesNotSupportLogging = errors.New("session does not support setting logging level")
|
||||
|
||||
// Notification-related errors
|
||||
ErrNotificationNotInitialized = errors.New("notification channel not initialized")
|
||||
ErrNotificationChannelBlocked = errors.New("notification channel queue is full - client may not be processing notifications fast enough")
|
||||
)
|
||||
|
||||
// ErrDynamicPathConfig is returned when attempting to use static path methods with dynamic path configuration
|
||||
type ErrDynamicPathConfig struct {
|
||||
Method string
|
||||
}
|
||||
|
||||
func (e *ErrDynamicPathConfig) Error() string {
|
||||
return fmt.Sprintf("%s cannot be used with WithDynamicBasePath. Use dynamic path logic in your router.", e.Method)
|
||||
}
|
||||
+532
@@ -0,0 +1,532 @@
|
||||
// Code generated by `go generate`. DO NOT EDIT.
|
||||
// source: server/internal/gen/hooks.go.tmpl
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// OnRegisterSessionHookFunc is a hook that will be called when a new session is registered.
|
||||
type OnRegisterSessionHookFunc func(ctx context.Context, session ClientSession)
|
||||
|
||||
// OnUnregisterSessionHookFunc is a hook that will be called when a session is being unregistered.
|
||||
type OnUnregisterSessionHookFunc func(ctx context.Context, session ClientSession)
|
||||
|
||||
// BeforeAnyHookFunc is a function that is called after the request is
|
||||
// parsed but before the method is called.
|
||||
type BeforeAnyHookFunc func(ctx context.Context, id any, method mcp.MCPMethod, message any)
|
||||
|
||||
// OnSuccessHookFunc is a hook that will be called after the request
|
||||
// successfully generates a result, but before the result is sent to the client.
|
||||
type OnSuccessHookFunc func(ctx context.Context, id any, method mcp.MCPMethod, message any, result any)
|
||||
|
||||
// OnErrorHookFunc is a hook that will be called when an error occurs,
|
||||
// either during the request parsing or the method execution.
|
||||
//
|
||||
// Example usage:
|
||||
// ```
|
||||
//
|
||||
// hooks.AddOnError(func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) {
|
||||
// // Check for specific error types using errors.Is
|
||||
// if errors.Is(err, ErrUnsupported) {
|
||||
// // Handle capability not supported errors
|
||||
// log.Printf("Capability not supported: %v", err)
|
||||
// }
|
||||
//
|
||||
// // Use errors.As to get specific error types
|
||||
// var parseErr = &UnparsableMessageError{}
|
||||
// if errors.As(err, &parseErr) {
|
||||
// // Access specific methods/fields of the error type
|
||||
// log.Printf("Failed to parse message for method %s: %v",
|
||||
// parseErr.GetMethod(), parseErr.Unwrap())
|
||||
// // Access the raw message that failed to parse
|
||||
// rawMsg := parseErr.GetMessage()
|
||||
// }
|
||||
//
|
||||
// // Check for specific resource/prompt/tool errors
|
||||
// switch {
|
||||
// case errors.Is(err, ErrResourceNotFound):
|
||||
// log.Printf("Resource not found: %v", err)
|
||||
// case errors.Is(err, ErrPromptNotFound):
|
||||
// log.Printf("Prompt not found: %v", err)
|
||||
// case errors.Is(err, ErrToolNotFound):
|
||||
// log.Printf("Tool not found: %v", err)
|
||||
// }
|
||||
// })
|
||||
type OnErrorHookFunc func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error)
|
||||
|
||||
// OnRequestInitializationFunc is a function that called before handle diff request method
|
||||
// Should any errors arise during func execution, the service will promptly return the corresponding error message.
|
||||
type OnRequestInitializationFunc func(ctx context.Context, id any, message any) error
|
||||
|
||||
type OnBeforeInitializeFunc func(ctx context.Context, id any, message *mcp.InitializeRequest)
|
||||
type OnAfterInitializeFunc func(ctx context.Context, id any, message *mcp.InitializeRequest, result *mcp.InitializeResult)
|
||||
|
||||
type OnBeforePingFunc func(ctx context.Context, id any, message *mcp.PingRequest)
|
||||
type OnAfterPingFunc func(ctx context.Context, id any, message *mcp.PingRequest, result *mcp.EmptyResult)
|
||||
|
||||
type OnBeforeSetLevelFunc func(ctx context.Context, id any, message *mcp.SetLevelRequest)
|
||||
type OnAfterSetLevelFunc func(ctx context.Context, id any, message *mcp.SetLevelRequest, result *mcp.EmptyResult)
|
||||
|
||||
type OnBeforeListResourcesFunc func(ctx context.Context, id any, message *mcp.ListResourcesRequest)
|
||||
type OnAfterListResourcesFunc func(ctx context.Context, id any, message *mcp.ListResourcesRequest, result *mcp.ListResourcesResult)
|
||||
|
||||
type OnBeforeListResourceTemplatesFunc func(ctx context.Context, id any, message *mcp.ListResourceTemplatesRequest)
|
||||
type OnAfterListResourceTemplatesFunc func(ctx context.Context, id any, message *mcp.ListResourceTemplatesRequest, result *mcp.ListResourceTemplatesResult)
|
||||
|
||||
type OnBeforeReadResourceFunc func(ctx context.Context, id any, message *mcp.ReadResourceRequest)
|
||||
type OnAfterReadResourceFunc func(ctx context.Context, id any, message *mcp.ReadResourceRequest, result *mcp.ReadResourceResult)
|
||||
|
||||
type OnBeforeListPromptsFunc func(ctx context.Context, id any, message *mcp.ListPromptsRequest)
|
||||
type OnAfterListPromptsFunc func(ctx context.Context, id any, message *mcp.ListPromptsRequest, result *mcp.ListPromptsResult)
|
||||
|
||||
type OnBeforeGetPromptFunc func(ctx context.Context, id any, message *mcp.GetPromptRequest)
|
||||
type OnAfterGetPromptFunc func(ctx context.Context, id any, message *mcp.GetPromptRequest, result *mcp.GetPromptResult)
|
||||
|
||||
type OnBeforeListToolsFunc func(ctx context.Context, id any, message *mcp.ListToolsRequest)
|
||||
type OnAfterListToolsFunc func(ctx context.Context, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult)
|
||||
|
||||
type OnBeforeCallToolFunc func(ctx context.Context, id any, message *mcp.CallToolRequest)
|
||||
type OnAfterCallToolFunc func(ctx context.Context, id any, message *mcp.CallToolRequest, result *mcp.CallToolResult)
|
||||
|
||||
type Hooks struct {
|
||||
OnRegisterSession []OnRegisterSessionHookFunc
|
||||
OnUnregisterSession []OnUnregisterSessionHookFunc
|
||||
OnBeforeAny []BeforeAnyHookFunc
|
||||
OnSuccess []OnSuccessHookFunc
|
||||
OnError []OnErrorHookFunc
|
||||
OnRequestInitialization []OnRequestInitializationFunc
|
||||
OnBeforeInitialize []OnBeforeInitializeFunc
|
||||
OnAfterInitialize []OnAfterInitializeFunc
|
||||
OnBeforePing []OnBeforePingFunc
|
||||
OnAfterPing []OnAfterPingFunc
|
||||
OnBeforeSetLevel []OnBeforeSetLevelFunc
|
||||
OnAfterSetLevel []OnAfterSetLevelFunc
|
||||
OnBeforeListResources []OnBeforeListResourcesFunc
|
||||
OnAfterListResources []OnAfterListResourcesFunc
|
||||
OnBeforeListResourceTemplates []OnBeforeListResourceTemplatesFunc
|
||||
OnAfterListResourceTemplates []OnAfterListResourceTemplatesFunc
|
||||
OnBeforeReadResource []OnBeforeReadResourceFunc
|
||||
OnAfterReadResource []OnAfterReadResourceFunc
|
||||
OnBeforeListPrompts []OnBeforeListPromptsFunc
|
||||
OnAfterListPrompts []OnAfterListPromptsFunc
|
||||
OnBeforeGetPrompt []OnBeforeGetPromptFunc
|
||||
OnAfterGetPrompt []OnAfterGetPromptFunc
|
||||
OnBeforeListTools []OnBeforeListToolsFunc
|
||||
OnAfterListTools []OnAfterListToolsFunc
|
||||
OnBeforeCallTool []OnBeforeCallToolFunc
|
||||
OnAfterCallTool []OnAfterCallToolFunc
|
||||
}
|
||||
|
||||
func (c *Hooks) AddBeforeAny(hook BeforeAnyHookFunc) {
|
||||
c.OnBeforeAny = append(c.OnBeforeAny, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddOnSuccess(hook OnSuccessHookFunc) {
|
||||
c.OnSuccess = append(c.OnSuccess, hook)
|
||||
}
|
||||
|
||||
// AddOnError registers a hook function that will be called when an error occurs.
|
||||
// The error parameter contains the actual error object, which can be interrogated
|
||||
// using Go's error handling patterns like errors.Is and errors.As.
|
||||
//
|
||||
// Example:
|
||||
// ```
|
||||
// // Create a channel to receive errors for testing
|
||||
// errChan := make(chan error, 1)
|
||||
//
|
||||
// // Register hook to capture and inspect errors
|
||||
// hooks := &Hooks{}
|
||||
//
|
||||
// hooks.AddOnError(func(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) {
|
||||
// // For capability-related errors
|
||||
// if errors.Is(err, ErrUnsupported) {
|
||||
// // Handle capability not supported
|
||||
// errChan <- err
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// // For parsing errors
|
||||
// var parseErr = &UnparsableMessageError{}
|
||||
// if errors.As(err, &parseErr) {
|
||||
// // Handle unparsable message errors
|
||||
// fmt.Printf("Failed to parse %s request: %v\n",
|
||||
// parseErr.GetMethod(), parseErr.Unwrap())
|
||||
// errChan <- parseErr
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// // For resource/prompt/tool not found errors
|
||||
// if errors.Is(err, ErrResourceNotFound) ||
|
||||
// errors.Is(err, ErrPromptNotFound) ||
|
||||
// errors.Is(err, ErrToolNotFound) {
|
||||
// // Handle not found errors
|
||||
// errChan <- err
|
||||
// return
|
||||
// }
|
||||
//
|
||||
// // For other errors
|
||||
// errChan <- err
|
||||
// })
|
||||
//
|
||||
// server := NewMCPServer("test-server", "1.0.0", WithHooks(hooks))
|
||||
// ```
|
||||
func (c *Hooks) AddOnError(hook OnErrorHookFunc) {
|
||||
c.OnError = append(c.OnError, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeAny(ctx context.Context, id any, method mcp.MCPMethod, message any) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeAny {
|
||||
hook(ctx, id, method, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) onSuccess(ctx context.Context, id any, method mcp.MCPMethod, message any, result any) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnSuccess {
|
||||
hook(ctx, id, method, message, result)
|
||||
}
|
||||
}
|
||||
|
||||
// onError calls all registered error hooks with the error object.
|
||||
// The err parameter contains the actual error that occurred, which implements
|
||||
// the standard error interface and may be a wrapped error or custom error type.
|
||||
//
|
||||
// This allows consumer code to use Go's error handling patterns:
|
||||
// - errors.Is(err, ErrUnsupported) to check for specific sentinel errors
|
||||
// - errors.As(err, &customErr) to extract custom error types
|
||||
//
|
||||
// Common error types include:
|
||||
// - ErrUnsupported: When a capability is not enabled
|
||||
// - UnparsableMessageError: When request parsing fails
|
||||
// - ErrResourceNotFound: When a resource is not found
|
||||
// - ErrPromptNotFound: When a prompt is not found
|
||||
// - ErrToolNotFound: When a tool is not found
|
||||
func (c *Hooks) onError(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnError {
|
||||
hook(ctx, id, method, message, err)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) AddOnRegisterSession(hook OnRegisterSessionHookFunc) {
|
||||
c.OnRegisterSession = append(c.OnRegisterSession, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) RegisterSession(ctx context.Context, session ClientSession) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnRegisterSession {
|
||||
hook(ctx, session)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) AddOnUnregisterSession(hook OnUnregisterSessionHookFunc) {
|
||||
c.OnUnregisterSession = append(c.OnUnregisterSession, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) UnregisterSession(ctx context.Context, session ClientSession) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnUnregisterSession {
|
||||
hook(ctx, session)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) AddOnRequestInitialization(hook OnRequestInitializationFunc) {
|
||||
c.OnRequestInitialization = append(c.OnRequestInitialization, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) onRequestInitialization(ctx context.Context, id any, message any) error {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
for _, hook := range c.OnRequestInitialization {
|
||||
err := hook(ctx, id, message)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (c *Hooks) AddBeforeInitialize(hook OnBeforeInitializeFunc) {
|
||||
c.OnBeforeInitialize = append(c.OnBeforeInitialize, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterInitialize(hook OnAfterInitializeFunc) {
|
||||
c.OnAfterInitialize = append(c.OnAfterInitialize, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeInitialize(ctx context.Context, id any, message *mcp.InitializeRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodInitialize, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeInitialize {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterInitialize(ctx context.Context, id any, message *mcp.InitializeRequest, result *mcp.InitializeResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodInitialize, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterInitialize {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforePing(hook OnBeforePingFunc) {
|
||||
c.OnBeforePing = append(c.OnBeforePing, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterPing(hook OnAfterPingFunc) {
|
||||
c.OnAfterPing = append(c.OnAfterPing, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforePing(ctx context.Context, id any, message *mcp.PingRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodPing, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforePing {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterPing(ctx context.Context, id any, message *mcp.PingRequest, result *mcp.EmptyResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodPing, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterPing {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeSetLevel(hook OnBeforeSetLevelFunc) {
|
||||
c.OnBeforeSetLevel = append(c.OnBeforeSetLevel, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterSetLevel(hook OnAfterSetLevelFunc) {
|
||||
c.OnAfterSetLevel = append(c.OnAfterSetLevel, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeSetLevel(ctx context.Context, id any, message *mcp.SetLevelRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodSetLogLevel, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeSetLevel {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterSetLevel(ctx context.Context, id any, message *mcp.SetLevelRequest, result *mcp.EmptyResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodSetLogLevel, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterSetLevel {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeListResources(hook OnBeforeListResourcesFunc) {
|
||||
c.OnBeforeListResources = append(c.OnBeforeListResources, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterListResources(hook OnAfterListResourcesFunc) {
|
||||
c.OnAfterListResources = append(c.OnAfterListResources, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeListResources(ctx context.Context, id any, message *mcp.ListResourcesRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodResourcesList, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeListResources {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterListResources(ctx context.Context, id any, message *mcp.ListResourcesRequest, result *mcp.ListResourcesResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodResourcesList, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterListResources {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeListResourceTemplates(hook OnBeforeListResourceTemplatesFunc) {
|
||||
c.OnBeforeListResourceTemplates = append(c.OnBeforeListResourceTemplates, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterListResourceTemplates(hook OnAfterListResourceTemplatesFunc) {
|
||||
c.OnAfterListResourceTemplates = append(c.OnAfterListResourceTemplates, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeListResourceTemplates(ctx context.Context, id any, message *mcp.ListResourceTemplatesRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodResourcesTemplatesList, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeListResourceTemplates {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterListResourceTemplates(ctx context.Context, id any, message *mcp.ListResourceTemplatesRequest, result *mcp.ListResourceTemplatesResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodResourcesTemplatesList, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterListResourceTemplates {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeReadResource(hook OnBeforeReadResourceFunc) {
|
||||
c.OnBeforeReadResource = append(c.OnBeforeReadResource, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterReadResource(hook OnAfterReadResourceFunc) {
|
||||
c.OnAfterReadResource = append(c.OnAfterReadResource, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeReadResource(ctx context.Context, id any, message *mcp.ReadResourceRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodResourcesRead, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeReadResource {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterReadResource(ctx context.Context, id any, message *mcp.ReadResourceRequest, result *mcp.ReadResourceResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodResourcesRead, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterReadResource {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeListPrompts(hook OnBeforeListPromptsFunc) {
|
||||
c.OnBeforeListPrompts = append(c.OnBeforeListPrompts, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterListPrompts(hook OnAfterListPromptsFunc) {
|
||||
c.OnAfterListPrompts = append(c.OnAfterListPrompts, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeListPrompts(ctx context.Context, id any, message *mcp.ListPromptsRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodPromptsList, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeListPrompts {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterListPrompts(ctx context.Context, id any, message *mcp.ListPromptsRequest, result *mcp.ListPromptsResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodPromptsList, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterListPrompts {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeGetPrompt(hook OnBeforeGetPromptFunc) {
|
||||
c.OnBeforeGetPrompt = append(c.OnBeforeGetPrompt, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterGetPrompt(hook OnAfterGetPromptFunc) {
|
||||
c.OnAfterGetPrompt = append(c.OnAfterGetPrompt, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeGetPrompt(ctx context.Context, id any, message *mcp.GetPromptRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodPromptsGet, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeGetPrompt {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterGetPrompt(ctx context.Context, id any, message *mcp.GetPromptRequest, result *mcp.GetPromptResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodPromptsGet, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterGetPrompt {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeListTools(hook OnBeforeListToolsFunc) {
|
||||
c.OnBeforeListTools = append(c.OnBeforeListTools, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterListTools(hook OnAfterListToolsFunc) {
|
||||
c.OnAfterListTools = append(c.OnAfterListTools, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeListTools(ctx context.Context, id any, message *mcp.ListToolsRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodToolsList, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeListTools {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterListTools(ctx context.Context, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodToolsList, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterListTools {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
func (c *Hooks) AddBeforeCallTool(hook OnBeforeCallToolFunc) {
|
||||
c.OnBeforeCallTool = append(c.OnBeforeCallTool, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) AddAfterCallTool(hook OnAfterCallToolFunc) {
|
||||
c.OnAfterCallTool = append(c.OnAfterCallTool, hook)
|
||||
}
|
||||
|
||||
func (c *Hooks) beforeCallTool(ctx context.Context, id any, message *mcp.CallToolRequest) {
|
||||
c.beforeAny(ctx, id, mcp.MethodToolsCall, message)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnBeforeCallTool {
|
||||
hook(ctx, id, message)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Hooks) afterCallTool(ctx context.Context, id any, message *mcp.CallToolRequest, result *mcp.CallToolResult) {
|
||||
c.onSuccess(ctx, id, mcp.MethodToolsCall, message, result)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
for _, hook := range c.OnAfterCallTool {
|
||||
hook(ctx, id, message, result)
|
||||
}
|
||||
}
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
// HTTPContextFunc is a function that takes an existing context and the current
|
||||
// request and returns a potentially modified context based on the request
|
||||
// content. This can be used to inject context values from headers, for example.
|
||||
type HTTPContextFunc func(ctx context.Context, r *http.Request) context.Context
|
||||
+165
@@ -0,0 +1,165 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// SamplingHandler defines the interface for handling sampling requests from servers.
|
||||
type SamplingHandler interface {
|
||||
CreateMessage(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error)
|
||||
}
|
||||
|
||||
// ElicitationHandler defines the interface for handling elicitation requests from servers.
|
||||
type ElicitationHandler interface {
|
||||
Elicit(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error)
|
||||
}
|
||||
|
||||
// RootsHandler defines the interface for handling roots list requests from servers.
|
||||
type RootsHandler interface {
|
||||
ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error)
|
||||
}
|
||||
|
||||
type InProcessSession struct {
|
||||
sessionID string
|
||||
notifications chan mcp.JSONRPCNotification
|
||||
initialized atomic.Bool
|
||||
loggingLevel atomic.Value
|
||||
clientInfo atomic.Value
|
||||
clientCapabilities atomic.Value
|
||||
samplingHandler SamplingHandler
|
||||
elicitationHandler ElicitationHandler
|
||||
rootsHandler RootsHandler
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
func NewInProcessSession(sessionID string, samplingHandler SamplingHandler) *InProcessSession {
|
||||
return &InProcessSession{
|
||||
sessionID: sessionID,
|
||||
notifications: make(chan mcp.JSONRPCNotification, 100),
|
||||
samplingHandler: samplingHandler,
|
||||
}
|
||||
}
|
||||
|
||||
func NewInProcessSessionWithHandlers(sessionID string, samplingHandler SamplingHandler, elicitationHandler ElicitationHandler, rootsHandler RootsHandler) *InProcessSession {
|
||||
return &InProcessSession{
|
||||
sessionID: sessionID,
|
||||
notifications: make(chan mcp.JSONRPCNotification, 100),
|
||||
samplingHandler: samplingHandler,
|
||||
elicitationHandler: elicitationHandler,
|
||||
rootsHandler: rootsHandler,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *InProcessSession) SessionID() string {
|
||||
return s.sessionID
|
||||
}
|
||||
|
||||
func (s *InProcessSession) NotificationChannel() chan<- mcp.JSONRPCNotification {
|
||||
return s.notifications
|
||||
}
|
||||
|
||||
func (s *InProcessSession) Initialize() {
|
||||
s.loggingLevel.Store(mcp.LoggingLevelError)
|
||||
s.initialized.Store(true)
|
||||
}
|
||||
|
||||
func (s *InProcessSession) Initialized() bool {
|
||||
return s.initialized.Load()
|
||||
}
|
||||
|
||||
func (s *InProcessSession) GetClientInfo() mcp.Implementation {
|
||||
if value := s.clientInfo.Load(); value != nil {
|
||||
if clientInfo, ok := value.(mcp.Implementation); ok {
|
||||
return clientInfo
|
||||
}
|
||||
}
|
||||
return mcp.Implementation{}
|
||||
}
|
||||
|
||||
func (s *InProcessSession) SetClientInfo(clientInfo mcp.Implementation) {
|
||||
s.clientInfo.Store(clientInfo)
|
||||
}
|
||||
|
||||
func (s *InProcessSession) GetClientCapabilities() mcp.ClientCapabilities {
|
||||
if value := s.clientCapabilities.Load(); value != nil {
|
||||
if clientCapabilities, ok := value.(mcp.ClientCapabilities); ok {
|
||||
return clientCapabilities
|
||||
}
|
||||
}
|
||||
return mcp.ClientCapabilities{}
|
||||
}
|
||||
|
||||
func (s *InProcessSession) SetClientCapabilities(clientCapabilities mcp.ClientCapabilities) {
|
||||
s.clientCapabilities.Store(clientCapabilities)
|
||||
}
|
||||
|
||||
func (s *InProcessSession) SetLogLevel(level mcp.LoggingLevel) {
|
||||
s.loggingLevel.Store(level)
|
||||
}
|
||||
|
||||
func (s *InProcessSession) GetLogLevel() mcp.LoggingLevel {
|
||||
level := s.loggingLevel.Load()
|
||||
if level == nil {
|
||||
return mcp.LoggingLevelError
|
||||
}
|
||||
return level.(mcp.LoggingLevel)
|
||||
}
|
||||
|
||||
func (s *InProcessSession) RequestSampling(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
|
||||
s.mu.RLock()
|
||||
handler := s.samplingHandler
|
||||
s.mu.RUnlock()
|
||||
|
||||
if handler == nil {
|
||||
return nil, fmt.Errorf("no sampling handler available")
|
||||
}
|
||||
|
||||
return handler.CreateMessage(ctx, request)
|
||||
}
|
||||
|
||||
func (s *InProcessSession) RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) {
|
||||
s.mu.RLock()
|
||||
handler := s.elicitationHandler
|
||||
s.mu.RUnlock()
|
||||
|
||||
if handler == nil {
|
||||
return nil, fmt.Errorf("no elicitation handler available")
|
||||
}
|
||||
|
||||
return handler.Elicit(ctx, request)
|
||||
}
|
||||
|
||||
// ListRoots sends a list roots request to the client and waits for the response.
|
||||
// Returns an error if no roots handler is available.
|
||||
func (s *InProcessSession) ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error) {
|
||||
s.mu.RLock()
|
||||
handler := s.rootsHandler
|
||||
s.mu.RUnlock()
|
||||
|
||||
if handler == nil {
|
||||
return nil, fmt.Errorf("no roots handler available")
|
||||
}
|
||||
|
||||
return handler.ListRoots(ctx, request)
|
||||
}
|
||||
|
||||
// GenerateInProcessSessionID generates a unique session ID for inprocess clients
|
||||
func GenerateInProcessSessionID() string {
|
||||
return fmt.Sprintf("inprocess-%d", time.Now().UnixNano())
|
||||
}
|
||||
|
||||
// Ensure interface compliance
|
||||
var (
|
||||
_ ClientSession = (*InProcessSession)(nil)
|
||||
_ SessionWithLogging = (*InProcessSession)(nil)
|
||||
_ SessionWithClientInfo = (*InProcessSession)(nil)
|
||||
_ SessionWithSampling = (*InProcessSession)(nil)
|
||||
_ SessionWithElicitation = (*InProcessSession)(nil)
|
||||
_ SessionWithRoots = (*InProcessSession)(nil)
|
||||
)
|
||||
+339
@@ -0,0 +1,339 @@
|
||||
// Code generated by `go generate`. DO NOT EDIT.
|
||||
// source: server/internal/gen/request_handler.go.tmpl
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// HandleMessage processes an incoming JSON-RPC message and returns an appropriate response
|
||||
func (s *MCPServer) HandleMessage(
|
||||
ctx context.Context,
|
||||
message json.RawMessage,
|
||||
) mcp.JSONRPCMessage {
|
||||
// Add server to context
|
||||
ctx = context.WithValue(ctx, serverKey{}, s)
|
||||
var err *requestError
|
||||
|
||||
var baseMessage struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
Method mcp.MCPMethod `json:"method"`
|
||||
ID any `json:"id,omitempty"`
|
||||
Result any `json:"result,omitempty"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(message, &baseMessage); err != nil {
|
||||
return createErrorResponse(
|
||||
nil,
|
||||
mcp.PARSE_ERROR,
|
||||
"Failed to parse message",
|
||||
)
|
||||
}
|
||||
|
||||
// Check for valid JSONRPC version
|
||||
if baseMessage.JSONRPC != mcp.JSONRPC_VERSION {
|
||||
return createErrorResponse(
|
||||
baseMessage.ID,
|
||||
mcp.INVALID_REQUEST,
|
||||
"Invalid JSON-RPC version",
|
||||
)
|
||||
}
|
||||
|
||||
if baseMessage.ID == nil {
|
||||
var notification mcp.JSONRPCNotification
|
||||
if err := json.Unmarshal(message, ¬ification); err != nil {
|
||||
return createErrorResponse(
|
||||
nil,
|
||||
mcp.PARSE_ERROR,
|
||||
"Failed to parse notification",
|
||||
)
|
||||
}
|
||||
s.handleNotification(ctx, notification)
|
||||
return nil // Return nil for notifications
|
||||
}
|
||||
|
||||
if baseMessage.Result != nil {
|
||||
// this is a response to a request sent by the server (e.g. from a ping
|
||||
// sent due to WithKeepAlive option)
|
||||
return nil
|
||||
}
|
||||
|
||||
handleErr := s.hooks.onRequestInitialization(ctx, baseMessage.ID, message)
|
||||
if handleErr != nil {
|
||||
return createErrorResponse(
|
||||
baseMessage.ID,
|
||||
mcp.INVALID_REQUEST,
|
||||
handleErr.Error(),
|
||||
)
|
||||
}
|
||||
|
||||
// Get request header from ctx
|
||||
h := ctx.Value(requestHeader)
|
||||
headers, ok := h.(http.Header)
|
||||
|
||||
if headers == nil || !ok {
|
||||
headers = make(http.Header)
|
||||
}
|
||||
|
||||
switch baseMessage.Method {
|
||||
case mcp.MethodInitialize:
|
||||
var request mcp.InitializeRequest
|
||||
var result *mcp.InitializeResult
|
||||
if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeInitialize(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleInitialize(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterInitialize(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodPing:
|
||||
var request mcp.PingRequest
|
||||
var result *mcp.EmptyResult
|
||||
if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforePing(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handlePing(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterPing(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodSetLogLevel:
|
||||
var request mcp.SetLevelRequest
|
||||
var result *mcp.EmptyResult
|
||||
if s.capabilities.logging == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("logging %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeSetLevel(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleSetLevel(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterSetLevel(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodResourcesList:
|
||||
var request mcp.ListResourcesRequest
|
||||
var result *mcp.ListResourcesResult
|
||||
if s.capabilities.resources == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("resources %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeListResources(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleListResources(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterListResources(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodResourcesTemplatesList:
|
||||
var request mcp.ListResourceTemplatesRequest
|
||||
var result *mcp.ListResourceTemplatesResult
|
||||
if s.capabilities.resources == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("resources %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeListResourceTemplates(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleListResourceTemplates(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterListResourceTemplates(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodResourcesRead:
|
||||
var request mcp.ReadResourceRequest
|
||||
var result *mcp.ReadResourceResult
|
||||
if s.capabilities.resources == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("resources %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeReadResource(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleReadResource(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterReadResource(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodPromptsList:
|
||||
var request mcp.ListPromptsRequest
|
||||
var result *mcp.ListPromptsResult
|
||||
if s.capabilities.prompts == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("prompts %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeListPrompts(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleListPrompts(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterListPrompts(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodPromptsGet:
|
||||
var request mcp.GetPromptRequest
|
||||
var result *mcp.GetPromptResult
|
||||
if s.capabilities.prompts == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("prompts %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeGetPrompt(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleGetPrompt(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterGetPrompt(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodToolsList:
|
||||
var request mcp.ListToolsRequest
|
||||
var result *mcp.ListToolsResult
|
||||
if s.capabilities.tools == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("tools %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeListTools(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleListTools(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterListTools(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
case mcp.MethodToolsCall:
|
||||
var request mcp.CallToolRequest
|
||||
var result *mcp.CallToolResult
|
||||
if s.capabilities.tools == nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.METHOD_NOT_FOUND,
|
||||
err: fmt.Errorf("tools %w", ErrUnsupported),
|
||||
}
|
||||
} else if unmarshalErr := json.Unmarshal(message, &request); unmarshalErr != nil {
|
||||
err = &requestError{
|
||||
id: baseMessage.ID,
|
||||
code: mcp.INVALID_REQUEST,
|
||||
err: &UnparsableMessageError{message: message, err: unmarshalErr, method: baseMessage.Method},
|
||||
}
|
||||
} else {
|
||||
request.Header = headers
|
||||
s.hooks.beforeCallTool(ctx, baseMessage.ID, &request)
|
||||
result, err = s.handleToolCall(ctx, baseMessage.ID, request)
|
||||
}
|
||||
if err != nil {
|
||||
s.hooks.onError(ctx, baseMessage.ID, baseMessage.Method, &request, err)
|
||||
return err.ToJSONRPCError()
|
||||
}
|
||||
s.hooks.afterCallTool(ctx, baseMessage.ID, &request, result)
|
||||
return createResponse(baseMessage.ID, *result)
|
||||
default:
|
||||
return createErrorResponse(
|
||||
baseMessage.ID,
|
||||
mcp.METHOD_NOT_FOUND,
|
||||
fmt.Sprintf("Method %s not found", baseMessage.Method),
|
||||
)
|
||||
}
|
||||
}
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrNoClientSession is returned when there is no active client session in the context
|
||||
ErrNoClientSession = errors.New("no active client session")
|
||||
// ErrRootsNotSupported is returned when the session does not support roots
|
||||
ErrRootsNotSupported = errors.New("session does not support roots")
|
||||
)
|
||||
|
||||
// RequestRoots sends an list roots request to the client.
|
||||
// The client must have declared roots capability during initialization.
|
||||
// The session must implement SessionWithRoots to support this operation.
|
||||
func (s *MCPServer) RequestRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error) {
|
||||
session := ClientSessionFromContext(ctx)
|
||||
if session == nil {
|
||||
return nil, ErrNoClientSession
|
||||
}
|
||||
|
||||
// Check if the session supports roots requests
|
||||
if rootsSession, ok := session.(SessionWithRoots); ok {
|
||||
return rootsSession.ListRoots(ctx, request)
|
||||
}
|
||||
|
||||
return nil, ErrRootsNotSupported
|
||||
}
|
||||
+61
@@ -0,0 +1,61 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// EnableSampling enables sampling capabilities for the server.
|
||||
// This allows the server to send sampling requests to clients that support it.
|
||||
func (s *MCPServer) EnableSampling() {
|
||||
s.capabilitiesMu.Lock()
|
||||
defer s.capabilitiesMu.Unlock()
|
||||
|
||||
enabled := true
|
||||
s.capabilities.sampling = &enabled
|
||||
}
|
||||
|
||||
// RequestSampling sends a sampling request to the client.
|
||||
// The client must have declared sampling capability during initialization.
|
||||
func (s *MCPServer) RequestSampling(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
|
||||
session := ClientSessionFromContext(ctx)
|
||||
if session == nil {
|
||||
return nil, fmt.Errorf("no active session")
|
||||
}
|
||||
|
||||
// Check if the session supports sampling requests
|
||||
if samplingSession, ok := session.(SessionWithSampling); ok {
|
||||
return samplingSession.RequestSampling(ctx, request)
|
||||
}
|
||||
|
||||
// Check for inprocess sampling handler in context
|
||||
if handler := InProcessSamplingHandlerFromContext(ctx); handler != nil {
|
||||
return handler.CreateMessage(ctx, request)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("session does not support sampling")
|
||||
}
|
||||
|
||||
// SessionWithSampling extends ClientSession to support sampling requests.
|
||||
type SessionWithSampling interface {
|
||||
ClientSession
|
||||
RequestSampling(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error)
|
||||
}
|
||||
|
||||
// inProcessSamplingHandlerKey is the context key for storing inprocess sampling handler
|
||||
type inProcessSamplingHandlerKey struct{}
|
||||
|
||||
// WithInProcessSamplingHandler adds a sampling handler to the context for inprocess clients
|
||||
func WithInProcessSamplingHandler(ctx context.Context, handler SamplingHandler) context.Context {
|
||||
return context.WithValue(ctx, inProcessSamplingHandlerKey{}, handler)
|
||||
}
|
||||
|
||||
// InProcessSamplingHandlerFromContext retrieves the inprocess sampling handler from context
|
||||
func InProcessSamplingHandlerFromContext(ctx context.Context) SamplingHandler {
|
||||
if handler, ok := ctx.Value(inProcessSamplingHandlerKey{}).(SamplingHandler); ok {
|
||||
return handler
|
||||
}
|
||||
return nil
|
||||
}
|
||||
+1337
File diff suppressed because it is too large.
Load diff
+770
@@ -0,0 +1,770 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/url"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// ClientSession represents an active session that can be used by MCPServer to interact with client.
|
||||
type ClientSession interface {
|
||||
// Initialize marks session as fully initialized and ready for notifications
|
||||
Initialize()
|
||||
// Initialized returns if session is ready to accept notifications
|
||||
Initialized() bool
|
||||
// NotificationChannel provides a channel suitable for sending notifications to client.
|
||||
NotificationChannel() chan<- mcp.JSONRPCNotification
|
||||
// SessionID is a unique identifier used to track user session.
|
||||
SessionID() string
|
||||
}
|
||||
|
||||
// SessionWithLogging is an extension of ClientSession that can receive log message notifications and set log level
|
||||
type SessionWithLogging interface {
|
||||
ClientSession
|
||||
// SetLogLevel sets the minimum log level
|
||||
SetLogLevel(level mcp.LoggingLevel)
|
||||
// GetLogLevel retrieves the minimum log level
|
||||
GetLogLevel() mcp.LoggingLevel
|
||||
}
|
||||
|
||||
// SessionWithTools is an extension of ClientSession that can store session-specific tool data
|
||||
type SessionWithTools interface {
|
||||
ClientSession
|
||||
// GetSessionTools returns the tools specific to this session, if any
|
||||
// This method must be thread-safe for concurrent access
|
||||
GetSessionTools() map[string]ServerTool
|
||||
// SetSessionTools sets tools specific to this session
|
||||
// This method must be thread-safe for concurrent access
|
||||
SetSessionTools(tools map[string]ServerTool)
|
||||
}
|
||||
|
||||
// SessionWithResources is an extension of ClientSession that can store session-specific resource data
|
||||
type SessionWithResources interface {
|
||||
ClientSession
|
||||
// GetSessionResources returns the resources specific to this session, if any
|
||||
// This method must be thread-safe for concurrent access
|
||||
GetSessionResources() map[string]ServerResource
|
||||
// SetSessionResources sets resources specific to this session
|
||||
// This method must be thread-safe for concurrent access
|
||||
SetSessionResources(resources map[string]ServerResource)
|
||||
}
|
||||
|
||||
// SessionWithResourceTemplates is an extension of ClientSession that can store session-specific resource template data
|
||||
type SessionWithResourceTemplates interface {
|
||||
ClientSession
|
||||
// GetSessionResourceTemplates returns the resource templates specific to this session, if any
|
||||
// This method must be thread-safe for concurrent access
|
||||
GetSessionResourceTemplates() map[string]ServerResourceTemplate
|
||||
// SetSessionResourceTemplates sets resource templates specific to this session
|
||||
// This method must be thread-safe for concurrent access
|
||||
SetSessionResourceTemplates(templates map[string]ServerResourceTemplate)
|
||||
}
|
||||
|
||||
// SessionWithClientInfo is an extension of ClientSession that can store client info
|
||||
type SessionWithClientInfo interface {
|
||||
ClientSession
|
||||
// GetClientInfo returns the client information for this session
|
||||
GetClientInfo() mcp.Implementation
|
||||
// SetClientInfo sets the client information for this session
|
||||
SetClientInfo(clientInfo mcp.Implementation)
|
||||
// GetClientCapabilities returns the client capabilities for this session
|
||||
GetClientCapabilities() mcp.ClientCapabilities
|
||||
// SetClientCapabilities sets the client capabilities for this session
|
||||
SetClientCapabilities(clientCapabilities mcp.ClientCapabilities)
|
||||
}
|
||||
|
||||
// SessionWithElicitation is an extension of ClientSession that can send elicitation requests
|
||||
type SessionWithElicitation interface {
|
||||
ClientSession
|
||||
// RequestElicitation sends an elicitation request to the client and waits for response
|
||||
RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error)
|
||||
}
|
||||
|
||||
// SessionWithRoots is an extension of ClientSession that can send list roots requests
|
||||
type SessionWithRoots interface {
|
||||
ClientSession
|
||||
// ListRoots sends an list roots request to the client and waits for response
|
||||
ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error)
|
||||
}
|
||||
|
||||
// SessionWithStreamableHTTPConfig extends ClientSession to support streamable HTTP transport configurations
|
||||
type SessionWithStreamableHTTPConfig interface {
|
||||
ClientSession
|
||||
// UpgradeToSSEWhenReceiveNotification upgrades the client-server communication to SSE stream when the server
|
||||
// sends notifications to the client
|
||||
//
|
||||
// The protocol specification:
|
||||
// - If the server response contains any JSON-RPC notifications, it MUST either:
|
||||
// - Return Content-Type: text/event-stream to initiate an SSE stream, OR
|
||||
// - Return Content-Type: application/json for a single JSON object
|
||||
// - The client MUST support both response types.
|
||||
//
|
||||
// Reference: https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#sending-messages-to-the-server
|
||||
UpgradeToSSEWhenReceiveNotification()
|
||||
}
|
||||
|
||||
// clientSessionKey is the context key for storing current client notification channel.
|
||||
type clientSessionKey struct{}
|
||||
|
||||
// ClientSessionFromContext retrieves current client notification context from context.
|
||||
func ClientSessionFromContext(ctx context.Context) ClientSession {
|
||||
if session, ok := ctx.Value(clientSessionKey{}).(ClientSession); ok {
|
||||
return session
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// WithContext sets the current client session and returns the provided context
|
||||
func (s *MCPServer) WithContext(
|
||||
ctx context.Context,
|
||||
session ClientSession,
|
||||
) context.Context {
|
||||
return context.WithValue(ctx, clientSessionKey{}, session)
|
||||
}
|
||||
|
||||
// RegisterSession saves session that should be notified in case if some server attributes changed.
|
||||
func (s *MCPServer) RegisterSession(
|
||||
ctx context.Context,
|
||||
session ClientSession,
|
||||
) error {
|
||||
sessionID := session.SessionID()
|
||||
if _, exists := s.sessions.LoadOrStore(sessionID, session); exists {
|
||||
return ErrSessionExists
|
||||
}
|
||||
s.hooks.RegisterSession(ctx, session)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *MCPServer) buildLogNotification(notification mcp.LoggingMessageNotification) mcp.JSONRPCNotification {
|
||||
return mcp.JSONRPCNotification{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
Notification: mcp.Notification{
|
||||
Method: notification.Method,
|
||||
Params: mcp.NotificationParams{
|
||||
AdditionalFields: map[string]any{
|
||||
"level": notification.Params.Level,
|
||||
"logger": notification.Params.Logger,
|
||||
"data": notification.Params.Data,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MCPServer) SendLogMessageToClient(ctx context.Context, notification mcp.LoggingMessageNotification) error {
|
||||
session := ClientSessionFromContext(ctx)
|
||||
if session == nil || !session.Initialized() {
|
||||
return ErrNotificationNotInitialized
|
||||
}
|
||||
sessionLogging, ok := session.(SessionWithLogging)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportLogging
|
||||
}
|
||||
if !notification.Params.Level.ShouldSendTo(sessionLogging.GetLogLevel()) {
|
||||
return nil
|
||||
}
|
||||
return s.sendNotificationCore(ctx, session, s.buildLogNotification(notification))
|
||||
}
|
||||
|
||||
func (s *MCPServer) sendNotificationToAllClients(notification mcp.JSONRPCNotification) {
|
||||
s.sessions.Range(func(k, v any) bool {
|
||||
if session, ok := v.(ClientSession); ok && session.Initialized() {
|
||||
if sessionWithStreamableHTTPConfig, ok := session.(SessionWithStreamableHTTPConfig); ok {
|
||||
sessionWithStreamableHTTPConfig.UpgradeToSSEWhenReceiveNotification()
|
||||
}
|
||||
select {
|
||||
case session.NotificationChannel() <- notification:
|
||||
// Successfully sent notification
|
||||
default:
|
||||
// Channel is blocked, if there's an error hook, use it
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
err := ErrNotificationChannelBlocked
|
||||
// Copy hooks pointer to local variable to avoid race condition
|
||||
hooks := s.hooks
|
||||
go func(sessionID string, hooks *Hooks) {
|
||||
ctx := context.Background()
|
||||
// Use the error hook to report the blocked channel
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": notification.Method,
|
||||
"sessionID": sessionID,
|
||||
}, fmt.Errorf("notification channel blocked for session %s: %w", sessionID, err))
|
||||
}(session.SessionID(), hooks)
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
func (s *MCPServer) sendNotificationToSpecificClient(session ClientSession, notification mcp.JSONRPCNotification) error {
|
||||
// upgrades the client-server communication to SSE stream when the server sends notifications to the client
|
||||
if sessionWithStreamableHTTPConfig, ok := session.(SessionWithStreamableHTTPConfig); ok {
|
||||
sessionWithStreamableHTTPConfig.UpgradeToSSEWhenReceiveNotification()
|
||||
}
|
||||
select {
|
||||
case session.NotificationChannel() <- notification:
|
||||
return nil
|
||||
default:
|
||||
// Channel is blocked, if there's an error hook, use it
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
err := ErrNotificationChannelBlocked
|
||||
ctx := context.Background()
|
||||
// Copy hooks pointer to local variable to avoid race condition
|
||||
hooks := s.hooks
|
||||
go func(sID string, hooks *Hooks) {
|
||||
// Use the error hook to report the blocked channel
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": notification.Method,
|
||||
"sessionID": sID,
|
||||
}, fmt.Errorf("notification channel blocked for session %s: %w", sID, err))
|
||||
}(session.SessionID(), hooks)
|
||||
}
|
||||
return ErrNotificationChannelBlocked
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MCPServer) SendLogMessageToSpecificClient(sessionID string, notification mcp.LoggingMessageNotification) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
session, ok := sessionValue.(ClientSession)
|
||||
if !ok || !session.Initialized() {
|
||||
return ErrSessionNotInitialized
|
||||
}
|
||||
sessionLogging, ok := session.(SessionWithLogging)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportLogging
|
||||
}
|
||||
if !notification.Params.Level.ShouldSendTo(sessionLogging.GetLogLevel()) {
|
||||
return nil
|
||||
}
|
||||
return s.sendNotificationToSpecificClient(session, s.buildLogNotification(notification))
|
||||
}
|
||||
|
||||
// UnregisterSession removes from storage session that is shut down.
|
||||
func (s *MCPServer) UnregisterSession(
|
||||
ctx context.Context,
|
||||
sessionID string,
|
||||
) {
|
||||
sessionValue, ok := s.sessions.LoadAndDelete(sessionID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if session, ok := sessionValue.(ClientSession); ok {
|
||||
s.hooks.UnregisterSession(ctx, session)
|
||||
}
|
||||
}
|
||||
|
||||
// SendNotificationToAllClients sends a notification to all the currently active clients.
|
||||
func (s *MCPServer) SendNotificationToAllClients(
|
||||
method string,
|
||||
params map[string]any,
|
||||
) {
|
||||
notification := mcp.JSONRPCNotification{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
Notification: mcp.Notification{
|
||||
Method: method,
|
||||
Params: mcp.NotificationParams{
|
||||
AdditionalFields: params,
|
||||
},
|
||||
},
|
||||
}
|
||||
s.sendNotificationToAllClients(notification)
|
||||
}
|
||||
|
||||
// SendNotificationToClient sends a notification to the current client
|
||||
func (s *MCPServer) sendNotificationCore(
|
||||
ctx context.Context,
|
||||
session ClientSession,
|
||||
notification mcp.JSONRPCNotification,
|
||||
) error {
|
||||
// upgrades the client-server communication to SSE stream when the server sends notifications to the client
|
||||
if sessionWithStreamableHTTPConfig, ok := session.(SessionWithStreamableHTTPConfig); ok {
|
||||
sessionWithStreamableHTTPConfig.UpgradeToSSEWhenReceiveNotification()
|
||||
}
|
||||
select {
|
||||
case session.NotificationChannel() <- notification:
|
||||
return nil
|
||||
default:
|
||||
// Channel is blocked, if there's an error hook, use it
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
method := notification.Method
|
||||
err := ErrNotificationChannelBlocked
|
||||
// Copy hooks pointer to local variable to avoid race condition
|
||||
hooks := s.hooks
|
||||
go func(sessionID string, hooks *Hooks) {
|
||||
// Use the error hook to report the blocked channel
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": method,
|
||||
"sessionID": sessionID,
|
||||
}, fmt.Errorf("notification channel blocked for session %s: %w", sessionID, err))
|
||||
}(session.SessionID(), hooks)
|
||||
}
|
||||
return ErrNotificationChannelBlocked
|
||||
}
|
||||
}
|
||||
|
||||
// SendNotificationToClient sends a notification to the current client
|
||||
func (s *MCPServer) SendNotificationToClient(
|
||||
ctx context.Context,
|
||||
method string,
|
||||
params map[string]any,
|
||||
) error {
|
||||
session := ClientSessionFromContext(ctx)
|
||||
if session == nil || !session.Initialized() {
|
||||
return ErrNotificationNotInitialized
|
||||
}
|
||||
notification := mcp.JSONRPCNotification{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
Notification: mcp.Notification{
|
||||
Method: method,
|
||||
Params: mcp.NotificationParams{
|
||||
AdditionalFields: params,
|
||||
},
|
||||
},
|
||||
}
|
||||
return s.sendNotificationCore(ctx, session, notification)
|
||||
}
|
||||
|
||||
// SendNotificationToSpecificClient sends a notification to a specific client by session ID
|
||||
func (s *MCPServer) SendNotificationToSpecificClient(
|
||||
sessionID string,
|
||||
method string,
|
||||
params map[string]any,
|
||||
) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
session, ok := sessionValue.(ClientSession)
|
||||
if !ok || !session.Initialized() {
|
||||
return ErrSessionNotInitialized
|
||||
}
|
||||
notification := mcp.JSONRPCNotification{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
Notification: mcp.Notification{
|
||||
Method: method,
|
||||
Params: mcp.NotificationParams{
|
||||
AdditionalFields: params,
|
||||
},
|
||||
},
|
||||
}
|
||||
return s.sendNotificationToSpecificClient(session, notification)
|
||||
}
|
||||
|
||||
// AddSessionTool adds a tool for a specific session
|
||||
func (s *MCPServer) AddSessionTool(sessionID string, tool mcp.Tool, handler ToolHandlerFunc) error {
|
||||
return s.AddSessionTools(sessionID, ServerTool{Tool: tool, Handler: handler})
|
||||
}
|
||||
|
||||
// AddSessionTools adds tools for a specific session
|
||||
func (s *MCPServer) AddSessionTools(sessionID string, tools ...ServerTool) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
session, ok := sessionValue.(SessionWithTools)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportTools
|
||||
}
|
||||
|
||||
s.implicitlyRegisterToolCapabilities()
|
||||
|
||||
// Get existing tools (this should return a thread-safe copy)
|
||||
sessionTools := session.GetSessionTools()
|
||||
|
||||
// Create a new map to avoid concurrent modification issues
|
||||
newSessionTools := make(map[string]ServerTool, len(sessionTools)+len(tools))
|
||||
|
||||
// Copy existing tools
|
||||
for k, v := range sessionTools {
|
||||
newSessionTools[k] = v
|
||||
}
|
||||
|
||||
// Add new tools
|
||||
for _, tool := range tools {
|
||||
newSessionTools[tool.Tool.Name] = tool
|
||||
}
|
||||
|
||||
// Set the tools (this should be thread-safe)
|
||||
session.SetSessionTools(newSessionTools)
|
||||
|
||||
// It only makes sense to send tool notifications to initialized sessions --
|
||||
// if we're not initialized yet the client can't possibly have sent their
|
||||
// initial tools/list message.
|
||||
//
|
||||
// For initialized sessions, honor tools.listChanged, which is specifically
|
||||
// about whether notifications will be sent or not.
|
||||
// see <https://modelcontextprotocol.io/specification/2025-03-26/server/tools#capabilities>
|
||||
if session.Initialized() && s.capabilities.tools != nil && s.capabilities.tools.listChanged {
|
||||
// Send notification only to this session
|
||||
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
|
||||
// Log the error but don't fail the operation
|
||||
// The tools were successfully added, but notification failed
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
hooks := s.hooks
|
||||
go func(sID string, hooks *Hooks) {
|
||||
ctx := context.Background()
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": "notifications/tools/list_changed",
|
||||
"sessionID": sID,
|
||||
}, fmt.Errorf("failed to send notification after adding tools: %w", err))
|
||||
}(sessionID, hooks)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteSessionTools removes tools from a specific session
|
||||
func (s *MCPServer) DeleteSessionTools(sessionID string, names ...string) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
session, ok := sessionValue.(SessionWithTools)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportTools
|
||||
}
|
||||
|
||||
// Get existing tools (this should return a thread-safe copy)
|
||||
sessionTools := session.GetSessionTools()
|
||||
if sessionTools == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Create a new map to avoid concurrent modification issues
|
||||
newSessionTools := make(map[string]ServerTool, len(sessionTools))
|
||||
|
||||
// Copy existing tools except those being deleted
|
||||
for k, v := range sessionTools {
|
||||
newSessionTools[k] = v
|
||||
}
|
||||
|
||||
// Remove specified tools
|
||||
for _, name := range names {
|
||||
delete(newSessionTools, name)
|
||||
}
|
||||
|
||||
// Set the tools (this should be thread-safe)
|
||||
session.SetSessionTools(newSessionTools)
|
||||
|
||||
// It only makes sense to send tool notifications to initialized sessions --
|
||||
// if we're not initialized yet the client can't possibly have sent their
|
||||
// initial tools/list message.
|
||||
//
|
||||
// For initialized sessions, honor tools.listChanged, which is specifically
|
||||
// about whether notifications will be sent or not.
|
||||
// see <https://modelcontextprotocol.io/specification/2025-03-26/server/tools#capabilities>
|
||||
if session.Initialized() && s.capabilities.tools != nil && s.capabilities.tools.listChanged {
|
||||
// Send notification only to this session
|
||||
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/tools/list_changed", nil); err != nil {
|
||||
// Log the error but don't fail the operation
|
||||
// The tools were successfully deleted, but notification failed
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
hooks := s.hooks
|
||||
go func(sID string, hooks *Hooks) {
|
||||
ctx := context.Background()
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": "notifications/tools/list_changed",
|
||||
"sessionID": sID,
|
||||
}, fmt.Errorf("failed to send notification after deleting tools: %w", err))
|
||||
}(sessionID, hooks)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddSessionResource adds a resource for a specific session
|
||||
func (s *MCPServer) AddSessionResource(sessionID string, resource mcp.Resource, handler ResourceHandlerFunc) error {
|
||||
return s.AddSessionResources(sessionID, ServerResource{Resource: resource, Handler: handler})
|
||||
}
|
||||
|
||||
// AddSessionResources adds resources for a specific session
|
||||
func (s *MCPServer) AddSessionResources(sessionID string, resources ...ServerResource) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
session, ok := sessionValue.(SessionWithResources)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportResources
|
||||
}
|
||||
|
||||
// For session resources, we want listChanged enabled by default
|
||||
s.implicitlyRegisterCapabilities(
|
||||
func() bool { return s.capabilities.resources != nil },
|
||||
func() { s.capabilities.resources = &resourceCapabilities{listChanged: true} },
|
||||
)
|
||||
|
||||
// Get existing resources (this should return a thread-safe copy)
|
||||
sessionResources := session.GetSessionResources()
|
||||
|
||||
// Create a new map to avoid concurrent modification issues
|
||||
newSessionResources := make(map[string]ServerResource, len(sessionResources)+len(resources))
|
||||
|
||||
// Copy existing resources
|
||||
for k, v := range sessionResources {
|
||||
newSessionResources[k] = v
|
||||
}
|
||||
|
||||
// Add new resources with validation
|
||||
for _, resource := range resources {
|
||||
// Validate that URI is non-empty
|
||||
if resource.Resource.URI == "" {
|
||||
return fmt.Errorf("resource URI cannot be empty")
|
||||
}
|
||||
|
||||
// Validate that URI conforms to RFC 3986
|
||||
if _, err := url.ParseRequestURI(resource.Resource.URI); err != nil {
|
||||
return fmt.Errorf("invalid resource URI: %w", err)
|
||||
}
|
||||
|
||||
newSessionResources[resource.Resource.URI] = resource
|
||||
}
|
||||
|
||||
// Set the resources (this should be thread-safe)
|
||||
session.SetSessionResources(newSessionResources)
|
||||
|
||||
// It only makes sense to send resource notifications to initialized sessions --
|
||||
// if we're not initialized yet the client can't possibly have sent their
|
||||
// initial resources/list message.
|
||||
//
|
||||
// For initialized sessions, honor resources.listChanged, which is specifically
|
||||
// about whether notifications will be sent or not.
|
||||
// see <https://modelcontextprotocol.io/specification/2025-03-26/server/resources#capabilities>
|
||||
if session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
|
||||
// Send notification only to this session
|
||||
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
|
||||
// Log the error but don't fail the operation
|
||||
// The resources were successfully added, but notification failed
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
hooks := s.hooks
|
||||
go func(sID string, hooks *Hooks) {
|
||||
ctx := context.Background()
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": "notifications/resources/list_changed",
|
||||
"sessionID": sID,
|
||||
}, fmt.Errorf("failed to send notification after adding resources: %w", err))
|
||||
}(sessionID, hooks)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteSessionResources removes resources from a specific session
|
||||
func (s *MCPServer) DeleteSessionResources(sessionID string, uris ...string) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
session, ok := sessionValue.(SessionWithResources)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportResources
|
||||
}
|
||||
|
||||
// Get existing resources (this should return a thread-safe copy)
|
||||
sessionResources := session.GetSessionResources()
|
||||
if sessionResources == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Create a new map to avoid concurrent modification issues
|
||||
newSessionResources := make(map[string]ServerResource, len(sessionResources))
|
||||
|
||||
// Copy existing resources except those being deleted
|
||||
for k, v := range sessionResources {
|
||||
newSessionResources[k] = v
|
||||
}
|
||||
|
||||
// Remove specified resources and track if anything was actually deleted
|
||||
actuallyDeleted := false
|
||||
for _, uri := range uris {
|
||||
if _, exists := newSessionResources[uri]; exists {
|
||||
delete(newSessionResources, uri)
|
||||
actuallyDeleted = true
|
||||
}
|
||||
}
|
||||
|
||||
// Skip no-op write if nothing was actually deleted
|
||||
if !actuallyDeleted {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Set the resources (this should be thread-safe)
|
||||
session.SetSessionResources(newSessionResources)
|
||||
|
||||
// It only makes sense to send resource notifications to initialized sessions --
|
||||
// if we're not initialized yet the client can't possibly have sent their
|
||||
// initial resources/list message.
|
||||
//
|
||||
// For initialized sessions, honor resources.listChanged, which is specifically
|
||||
// about whether notifications will be sent or not.
|
||||
// see <https://modelcontextprotocol.io/specification/2025-03-26/server/resources#capabilities>
|
||||
// Only send notification if something was actually deleted
|
||||
if actuallyDeleted && session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
|
||||
// Send notification only to this session
|
||||
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
|
||||
// Log the error but don't fail the operation
|
||||
// The resources were successfully deleted, but notification failed
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
hooks := s.hooks
|
||||
go func(sID string, hooks *Hooks) {
|
||||
ctx := context.Background()
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": "notifications/resources/list_changed",
|
||||
"sessionID": sID,
|
||||
}, fmt.Errorf("failed to send notification after deleting resources: %w", err))
|
||||
}(sessionID, hooks)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddSessionResourceTemplate adds a resource template for a specific session
|
||||
func (s *MCPServer) AddSessionResourceTemplate(sessionID string, template mcp.ResourceTemplate, handler ResourceTemplateHandlerFunc) error {
|
||||
return s.AddSessionResourceTemplates(sessionID, ServerResourceTemplate{
|
||||
Template: template,
|
||||
Handler: handler,
|
||||
})
|
||||
}
|
||||
|
||||
// AddSessionResourceTemplates adds resource templates for a specific session
|
||||
func (s *MCPServer) AddSessionResourceTemplates(sessionID string, templates ...ServerResourceTemplate) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
session, ok := sessionValue.(SessionWithResourceTemplates)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportResourceTemplates
|
||||
}
|
||||
|
||||
// For session resource templates, enable listChanged by default
|
||||
// This is the same behavior as session resources
|
||||
s.implicitlyRegisterCapabilities(
|
||||
func() bool { return s.capabilities.resources != nil },
|
||||
func() { s.capabilities.resources = &resourceCapabilities{listChanged: true} },
|
||||
)
|
||||
|
||||
// Get existing templates (this returns a thread-safe copy)
|
||||
sessionTemplates := session.GetSessionResourceTemplates()
|
||||
|
||||
// Create a new map to avoid modifying the returned copy
|
||||
newTemplates := make(map[string]ServerResourceTemplate, len(sessionTemplates)+len(templates))
|
||||
|
||||
// Copy existing templates
|
||||
for k, v := range sessionTemplates {
|
||||
newTemplates[k] = v
|
||||
}
|
||||
|
||||
// Validate and add new templates
|
||||
for _, t := range templates {
|
||||
if t.Template.URITemplate == nil {
|
||||
return fmt.Errorf("resource template URITemplate cannot be nil")
|
||||
}
|
||||
raw := t.Template.URITemplate.Raw()
|
||||
if raw == "" {
|
||||
return fmt.Errorf("resource template URITemplate cannot be empty")
|
||||
}
|
||||
if t.Template.Name == "" {
|
||||
return fmt.Errorf("resource template name cannot be empty")
|
||||
}
|
||||
newTemplates[raw] = t
|
||||
}
|
||||
|
||||
// Set the new templates (this method must handle thread-safety)
|
||||
session.SetSessionResourceTemplates(newTemplates)
|
||||
|
||||
// Send notification if the session is initialized and listChanged is enabled
|
||||
if session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
|
||||
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
|
||||
// Log the error but don't fail the operation
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
hooks := s.hooks
|
||||
go func(sID string, hooks *Hooks) {
|
||||
ctx := context.Background()
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": "notifications/resources/list_changed",
|
||||
"sessionID": sID,
|
||||
}, fmt.Errorf("failed to send notification after adding resource templates: %w", err))
|
||||
}(sessionID, hooks)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteSessionResourceTemplates removes resource templates from a specific session
|
||||
func (s *MCPServer) DeleteSessionResourceTemplates(sessionID string, uriTemplates ...string) error {
|
||||
sessionValue, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
session, ok := sessionValue.(SessionWithResourceTemplates)
|
||||
if !ok {
|
||||
return ErrSessionDoesNotSupportResourceTemplates
|
||||
}
|
||||
|
||||
// Get existing templates (this returns a thread-safe copy)
|
||||
sessionTemplates := session.GetSessionResourceTemplates()
|
||||
|
||||
// Track if any were actually deleted
|
||||
deletedAny := false
|
||||
|
||||
// Create a new map without the deleted templates
|
||||
newTemplates := make(map[string]ServerResourceTemplate, len(sessionTemplates))
|
||||
for k, v := range sessionTemplates {
|
||||
newTemplates[k] = v
|
||||
}
|
||||
|
||||
// Delete specified templates
|
||||
for _, uriTemplate := range uriTemplates {
|
||||
if _, exists := newTemplates[uriTemplate]; exists {
|
||||
delete(newTemplates, uriTemplate)
|
||||
deletedAny = true
|
||||
}
|
||||
}
|
||||
|
||||
// Only update if something was actually deleted
|
||||
if deletedAny {
|
||||
// Set the new templates (this method must handle thread-safety)
|
||||
session.SetSessionResourceTemplates(newTemplates)
|
||||
|
||||
// Send notification if the session is initialized and listChanged is enabled
|
||||
if session.Initialized() && s.capabilities.resources != nil && s.capabilities.resources.listChanged {
|
||||
if err := s.SendNotificationToSpecificClient(sessionID, "notifications/resources/list_changed", nil); err != nil {
|
||||
// Log the error but don't fail the operation
|
||||
if s.hooks != nil && len(s.hooks.OnError) > 0 {
|
||||
hooks := s.hooks
|
||||
go func(sID string, hooks *Hooks) {
|
||||
ctx := context.Background()
|
||||
hooks.onError(ctx, nil, "notification", map[string]any{
|
||||
"method": "notifications/resources/list_changed",
|
||||
"sessionID": sID,
|
||||
}, fmt.Errorf("failed to send notification after deleting resource templates: %w", err))
|
||||
}(sessionID, hooks)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
+797
@@ -0,0 +1,797 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"path"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// sseSession represents an active SSE connection.
|
||||
type sseSession struct {
|
||||
done chan struct{}
|
||||
eventQueue chan string // Channel for queuing events
|
||||
sessionID string
|
||||
requestID atomic.Int64
|
||||
notificationChannel chan mcp.JSONRPCNotification
|
||||
initialized atomic.Bool
|
||||
loggingLevel atomic.Value
|
||||
tools sync.Map // stores session-specific tools
|
||||
resources sync.Map // stores session-specific resources
|
||||
resourceTemplates sync.Map // stores session-specific resource templates
|
||||
clientInfo atomic.Value // stores session-specific client info
|
||||
clientCapabilities atomic.Value // stores session-specific client capabilities
|
||||
}
|
||||
|
||||
// SSEContextFunc is a function that takes an existing context and the current
|
||||
// request and returns a potentially modified context based on the request
|
||||
// content. This can be used to inject context values from headers, for example.
|
||||
type SSEContextFunc func(ctx context.Context, r *http.Request) context.Context
|
||||
|
||||
// DynamicBasePathFunc allows the user to provide a function to generate the
|
||||
// base path for a given request and sessionID. This is useful for cases where
|
||||
// the base path is not known at the time of SSE server creation, such as when
|
||||
// using a reverse proxy or when the base path is dynamically generated. The
|
||||
// function should return the base path (e.g., "/mcp/tenant123").
|
||||
type DynamicBasePathFunc func(r *http.Request, sessionID string) string
|
||||
|
||||
func (s *sseSession) SessionID() string {
|
||||
return s.sessionID
|
||||
}
|
||||
|
||||
func (s *sseSession) NotificationChannel() chan<- mcp.JSONRPCNotification {
|
||||
return s.notificationChannel
|
||||
}
|
||||
|
||||
func (s *sseSession) Initialize() {
|
||||
// set default logging level
|
||||
s.loggingLevel.Store(mcp.LoggingLevelError)
|
||||
s.initialized.Store(true)
|
||||
}
|
||||
|
||||
func (s *sseSession) Initialized() bool {
|
||||
return s.initialized.Load()
|
||||
}
|
||||
|
||||
func (s *sseSession) SetLogLevel(level mcp.LoggingLevel) {
|
||||
s.loggingLevel.Store(level)
|
||||
}
|
||||
|
||||
func (s *sseSession) GetLogLevel() mcp.LoggingLevel {
|
||||
level := s.loggingLevel.Load()
|
||||
if level == nil {
|
||||
return mcp.LoggingLevelError
|
||||
}
|
||||
return level.(mcp.LoggingLevel)
|
||||
}
|
||||
|
||||
func (s *sseSession) GetSessionResources() map[string]ServerResource {
|
||||
resources := make(map[string]ServerResource)
|
||||
s.resources.Range(func(key, value any) bool {
|
||||
if resource, ok := value.(ServerResource); ok {
|
||||
resources[key.(string)] = resource
|
||||
}
|
||||
return true
|
||||
})
|
||||
return resources
|
||||
}
|
||||
|
||||
func (s *sseSession) SetSessionResources(resources map[string]ServerResource) {
|
||||
// Clear existing resources
|
||||
s.resources.Clear()
|
||||
|
||||
// Set new resources
|
||||
for name, resource := range resources {
|
||||
s.resources.Store(name, resource)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *sseSession) GetSessionResourceTemplates() map[string]ServerResourceTemplate {
|
||||
templates := make(map[string]ServerResourceTemplate)
|
||||
s.resourceTemplates.Range(func(key, value any) bool {
|
||||
if template, ok := value.(ServerResourceTemplate); ok {
|
||||
templates[key.(string)] = template
|
||||
}
|
||||
return true
|
||||
})
|
||||
return templates
|
||||
}
|
||||
|
||||
func (s *sseSession) SetSessionResourceTemplates(templates map[string]ServerResourceTemplate) {
|
||||
// Clear existing templates
|
||||
s.resourceTemplates.Clear()
|
||||
|
||||
// Set new templates
|
||||
for uriTemplate, template := range templates {
|
||||
s.resourceTemplates.Store(uriTemplate, template)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *sseSession) GetSessionTools() map[string]ServerTool {
|
||||
tools := make(map[string]ServerTool)
|
||||
s.tools.Range(func(key, value any) bool {
|
||||
if tool, ok := value.(ServerTool); ok {
|
||||
tools[key.(string)] = tool
|
||||
}
|
||||
return true
|
||||
})
|
||||
return tools
|
||||
}
|
||||
|
||||
func (s *sseSession) SetSessionTools(tools map[string]ServerTool) {
|
||||
// Clear existing tools
|
||||
s.tools.Clear()
|
||||
|
||||
// Set new tools
|
||||
for name, tool := range tools {
|
||||
s.tools.Store(name, tool)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *sseSession) GetClientInfo() mcp.Implementation {
|
||||
if value := s.clientInfo.Load(); value != nil {
|
||||
if clientInfo, ok := value.(mcp.Implementation); ok {
|
||||
return clientInfo
|
||||
}
|
||||
}
|
||||
return mcp.Implementation{}
|
||||
}
|
||||
|
||||
func (s *sseSession) SetClientInfo(clientInfo mcp.Implementation) {
|
||||
s.clientInfo.Store(clientInfo)
|
||||
}
|
||||
|
||||
func (s *sseSession) SetClientCapabilities(clientCapabilities mcp.ClientCapabilities) {
|
||||
s.clientCapabilities.Store(clientCapabilities)
|
||||
}
|
||||
|
||||
func (s *sseSession) GetClientCapabilities() mcp.ClientCapabilities {
|
||||
if value := s.clientCapabilities.Load(); value != nil {
|
||||
if clientCapabilities, ok := value.(mcp.ClientCapabilities); ok {
|
||||
return clientCapabilities
|
||||
}
|
||||
}
|
||||
return mcp.ClientCapabilities{}
|
||||
}
|
||||
|
||||
var (
|
||||
_ ClientSession = (*sseSession)(nil)
|
||||
_ SessionWithTools = (*sseSession)(nil)
|
||||
_ SessionWithResources = (*sseSession)(nil)
|
||||
_ SessionWithResourceTemplates = (*sseSession)(nil)
|
||||
_ SessionWithLogging = (*sseSession)(nil)
|
||||
_ SessionWithClientInfo = (*sseSession)(nil)
|
||||
)
|
||||
|
||||
// SSEServer implements a Server-Sent Events (SSE) based MCP server.
|
||||
// It provides real-time communication capabilities over HTTP using the SSE protocol.
|
||||
type SSEServer struct {
|
||||
server *MCPServer
|
||||
baseURL string
|
||||
basePath string
|
||||
appendQueryToMessageEndpoint bool
|
||||
useFullURLForMessageEndpoint bool
|
||||
messageEndpoint string
|
||||
sseEndpoint string
|
||||
sessions sync.Map
|
||||
srv *http.Server
|
||||
contextFunc SSEContextFunc
|
||||
dynamicBasePathFunc DynamicBasePathFunc
|
||||
|
||||
keepAlive bool
|
||||
keepAliveInterval time.Duration
|
||||
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// SSEOption defines a function type for configuring SSEServer
|
||||
type SSEOption func(*SSEServer)
|
||||
|
||||
// WithBaseURL sets the base URL for the SSE server
|
||||
func WithBaseURL(baseURL string) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
if baseURL != "" {
|
||||
u, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
return
|
||||
}
|
||||
// Check if the host is empty or only contains a port
|
||||
if u.Host == "" || strings.HasPrefix(u.Host, ":") {
|
||||
return
|
||||
}
|
||||
if len(u.Query()) > 0 {
|
||||
return
|
||||
}
|
||||
}
|
||||
s.baseURL = strings.TrimSuffix(baseURL, "/")
|
||||
}
|
||||
}
|
||||
|
||||
// WithStaticBasePath adds a new option for setting a static base path
|
||||
func WithStaticBasePath(basePath string) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.basePath = normalizeURLPath(basePath)
|
||||
}
|
||||
}
|
||||
|
||||
// WithBasePath adds a new option for setting a static base path.
|
||||
//
|
||||
// Deprecated: Use WithStaticBasePath instead. This will be removed in a future version.
|
||||
//
|
||||
//go:deprecated
|
||||
func WithBasePath(basePath string) SSEOption {
|
||||
return WithStaticBasePath(basePath)
|
||||
}
|
||||
|
||||
// WithDynamicBasePath accepts a function for generating the base path. This is
|
||||
// useful for cases where the base path is not known at the time of SSE server
|
||||
// creation, such as when using a reverse proxy or when the server is mounted
|
||||
// at a dynamic path.
|
||||
func WithDynamicBasePath(fn DynamicBasePathFunc) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
if fn != nil {
|
||||
s.dynamicBasePathFunc = func(r *http.Request, sid string) string {
|
||||
bp := fn(r, sid)
|
||||
return normalizeURLPath(bp)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithMessageEndpoint sets the message endpoint path
|
||||
func WithMessageEndpoint(endpoint string) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.messageEndpoint = endpoint
|
||||
}
|
||||
}
|
||||
|
||||
// WithAppendQueryToMessageEndpoint configures the SSE server to append the original request's
|
||||
// query parameters to the message endpoint URL that is sent to clients during the SSE connection
|
||||
// initialization. This is useful when you need to preserve query parameters from the initial
|
||||
// SSE connection request and carry them over to subsequent message requests, maintaining
|
||||
// context or authentication details across the communication channel.
|
||||
func WithAppendQueryToMessageEndpoint() SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.appendQueryToMessageEndpoint = true
|
||||
}
|
||||
}
|
||||
|
||||
// WithUseFullURLForMessageEndpoint controls whether the SSE server returns a complete URL (including baseURL)
|
||||
// or just the path portion for the message endpoint. Set to false when clients will concatenate
|
||||
// the baseURL themselves to avoid malformed URLs like "http://localhost/mcphttp://localhost/mcp/message".
|
||||
func WithUseFullURLForMessageEndpoint(useFullURLForMessageEndpoint bool) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.useFullURLForMessageEndpoint = useFullURLForMessageEndpoint
|
||||
}
|
||||
}
|
||||
|
||||
// WithSSEEndpoint sets the SSE endpoint path
|
||||
func WithSSEEndpoint(endpoint string) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.sseEndpoint = endpoint
|
||||
}
|
||||
}
|
||||
|
||||
// WithHTTPServer sets the HTTP server instance.
|
||||
// NOTE: When providing a custom HTTP server, you must handle routing yourself
|
||||
// If routing is not set up, the server will start but won't handle any MCP requests.
|
||||
func WithHTTPServer(srv *http.Server) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.srv = srv
|
||||
}
|
||||
}
|
||||
|
||||
func WithKeepAliveInterval(keepAliveInterval time.Duration) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.keepAlive = true
|
||||
s.keepAliveInterval = keepAliveInterval
|
||||
}
|
||||
}
|
||||
|
||||
func WithKeepAlive(keepAlive bool) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.keepAlive = keepAlive
|
||||
}
|
||||
}
|
||||
|
||||
// WithSSEContextFunc sets a function that will be called to customise the context
|
||||
// to the server using the incoming request.
|
||||
func WithSSEContextFunc(fn SSEContextFunc) SSEOption {
|
||||
return func(s *SSEServer) {
|
||||
s.contextFunc = fn
|
||||
}
|
||||
}
|
||||
|
||||
// NewSSEServer creates a new SSE server instance with the given MCP server and options.
|
||||
func NewSSEServer(server *MCPServer, opts ...SSEOption) *SSEServer {
|
||||
s := &SSEServer{
|
||||
server: server,
|
||||
sseEndpoint: "/sse",
|
||||
messageEndpoint: "/message",
|
||||
useFullURLForMessageEndpoint: true,
|
||||
keepAlive: false,
|
||||
keepAliveInterval: 10 * time.Second,
|
||||
}
|
||||
|
||||
// Apply all options
|
||||
for _, opt := range opts {
|
||||
opt(s)
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// NewTestServer creates a test server for testing purposes
|
||||
func NewTestServer(server *MCPServer, opts ...SSEOption) *httptest.Server {
|
||||
sseServer := NewSSEServer(server, opts...)
|
||||
|
||||
testServer := httptest.NewServer(sseServer)
|
||||
sseServer.baseURL = testServer.URL
|
||||
return testServer
|
||||
}
|
||||
|
||||
// Start begins serving SSE connections on the specified address.
|
||||
// It sets up HTTP handlers for SSE and message endpoints.
|
||||
func (s *SSEServer) Start(addr string) error {
|
||||
s.mu.Lock()
|
||||
if s.srv == nil {
|
||||
s.srv = &http.Server{
|
||||
Addr: addr,
|
||||
Handler: s,
|
||||
}
|
||||
} else {
|
||||
if s.srv.Addr == "" {
|
||||
s.srv.Addr = addr
|
||||
} else if s.srv.Addr != addr {
|
||||
return fmt.Errorf("conflicting listen address: WithHTTPServer(%q) vs Start(%q)", s.srv.Addr, addr)
|
||||
}
|
||||
}
|
||||
srv := s.srv
|
||||
s.mu.Unlock()
|
||||
|
||||
return srv.ListenAndServe()
|
||||
}
|
||||
|
||||
// Shutdown gracefully stops the SSE server, closing all active sessions
|
||||
// and shutting down the HTTP server.
|
||||
func (s *SSEServer) Shutdown(ctx context.Context) error {
|
||||
s.mu.RLock()
|
||||
srv := s.srv
|
||||
s.mu.RUnlock()
|
||||
|
||||
if srv != nil {
|
||||
s.sessions.Range(func(key, value any) bool {
|
||||
if session, ok := value.(*sseSession); ok {
|
||||
close(session.done)
|
||||
}
|
||||
s.sessions.Delete(key)
|
||||
return true
|
||||
})
|
||||
|
||||
return srv.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleSSE handles incoming SSE connection requests.
|
||||
// It sets up appropriate headers and creates a new session for the client.
|
||||
func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||||
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
sessionID := uuid.New().String()
|
||||
session := &sseSession{
|
||||
done: make(chan struct{}),
|
||||
eventQueue: make(chan string, 100), // Buffer for events
|
||||
sessionID: sessionID,
|
||||
notificationChannel: make(chan mcp.JSONRPCNotification, 100),
|
||||
}
|
||||
|
||||
s.sessions.Store(sessionID, session)
|
||||
defer s.sessions.Delete(sessionID)
|
||||
|
||||
if err := s.server.RegisterSession(r.Context(), session); err != nil {
|
||||
http.Error(
|
||||
w,
|
||||
fmt.Sprintf("Session registration failed: %v", err),
|
||||
http.StatusInternalServerError,
|
||||
)
|
||||
return
|
||||
}
|
||||
defer s.server.UnregisterSession(r.Context(), sessionID)
|
||||
|
||||
// Start notification handler for this session
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case notification := <-session.notificationChannel:
|
||||
eventData, err := json.Marshal(notification)
|
||||
if err == nil {
|
||||
select {
|
||||
case session.eventQueue <- fmt.Sprintf("event: message\ndata: %s\n\n", eventData):
|
||||
// Event queued successfully
|
||||
case <-session.done:
|
||||
return
|
||||
}
|
||||
}
|
||||
case <-session.done:
|
||||
return
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Start keep alive : ping
|
||||
if s.keepAlive {
|
||||
go func() {
|
||||
ticker := time.NewTicker(s.keepAliveInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
message := mcp.JSONRPCRequest{
|
||||
JSONRPC: "2.0",
|
||||
ID: mcp.NewRequestId(session.requestID.Add(1)),
|
||||
Request: mcp.Request{
|
||||
Method: "ping",
|
||||
},
|
||||
}
|
||||
messageBytes, _ := json.Marshal(message)
|
||||
pingMsg := fmt.Sprintf("event: message\ndata:%s\n\n", messageBytes)
|
||||
select {
|
||||
case session.eventQueue <- pingMsg:
|
||||
// Message sent successfully
|
||||
case <-session.done:
|
||||
return
|
||||
}
|
||||
case <-session.done:
|
||||
return
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Send the initial endpoint event
|
||||
endpoint := s.GetMessageEndpointForClient(r, sessionID)
|
||||
if s.appendQueryToMessageEndpoint && len(r.URL.RawQuery) > 0 {
|
||||
endpoint += "&" + r.URL.RawQuery
|
||||
}
|
||||
fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", endpoint)
|
||||
flusher.Flush()
|
||||
|
||||
// Main event loop - this runs in the HTTP handler goroutine
|
||||
for {
|
||||
select {
|
||||
case event := <-session.eventQueue:
|
||||
// Write the event to the response
|
||||
fmt.Fprint(w, event)
|
||||
flusher.Flush()
|
||||
case <-r.Context().Done():
|
||||
close(session.done)
|
||||
return
|
||||
case <-session.done:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// GetMessageEndpointForClient returns the appropriate message endpoint URL with session ID
|
||||
// for the given request. This is the canonical way to compute the message endpoint for a client.
|
||||
// It handles both dynamic and static path modes, and honors the WithUseFullURLForMessageEndpoint flag.
|
||||
func (s *SSEServer) GetMessageEndpointForClient(r *http.Request, sessionID string) string {
|
||||
basePath := s.basePath
|
||||
if s.dynamicBasePathFunc != nil {
|
||||
basePath = s.dynamicBasePathFunc(r, sessionID)
|
||||
}
|
||||
|
||||
endpointPath := normalizeURLPath(basePath, s.messageEndpoint)
|
||||
if s.useFullURLForMessageEndpoint && s.baseURL != "" {
|
||||
endpointPath = s.baseURL + endpointPath
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s?sessionId=%s", endpointPath, sessionID)
|
||||
}
|
||||
|
||||
// handleMessage processes incoming JSON-RPC messages from clients and sends responses
|
||||
// back through the SSE connection and 202 code to HTTP response.
|
||||
func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
s.writeJSONRPCError(w, nil, mcp.INVALID_REQUEST, "Method not allowed")
|
||||
return
|
||||
}
|
||||
|
||||
sessionID := r.URL.Query().Get("sessionId")
|
||||
if sessionID == "" {
|
||||
s.writeJSONRPCError(w, nil, mcp.INVALID_PARAMS, "Missing sessionId")
|
||||
return
|
||||
}
|
||||
sessionI, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
s.writeJSONRPCError(w, nil, mcp.INVALID_PARAMS, "Invalid session ID")
|
||||
return
|
||||
}
|
||||
session := sessionI.(*sseSession)
|
||||
|
||||
// Set the client context before handling the message
|
||||
ctx := s.server.WithContext(r.Context(), session)
|
||||
if s.contextFunc != nil {
|
||||
ctx = s.contextFunc(ctx, r)
|
||||
}
|
||||
|
||||
// Parse message as raw JSON
|
||||
var rawMessage json.RawMessage
|
||||
if err := json.NewDecoder(r.Body).Decode(&rawMessage); err != nil {
|
||||
s.writeJSONRPCError(w, nil, mcp.PARSE_ERROR, "Parse error")
|
||||
return
|
||||
}
|
||||
|
||||
// Create a context that preserves all values from parent ctx but won't be canceled when the parent is canceled.
|
||||
// this is required because the http ctx will be canceled when the client disconnects
|
||||
detachedCtx := context.WithoutCancel(ctx)
|
||||
|
||||
// quick return request, send 202 Accepted with no body, then deal the message and sent response via SSE
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
|
||||
// Create a new context for handling the message that will be canceled when the message handling is done
|
||||
messageCtx := context.WithValue(detachedCtx, requestHeader, r.Header)
|
||||
messageCtx, cancel := context.WithCancel(messageCtx)
|
||||
|
||||
go func(ctx context.Context) {
|
||||
defer cancel()
|
||||
// Use the context that will be canceled when session is done
|
||||
// Process message through MCPServer
|
||||
response := s.server.HandleMessage(ctx, rawMessage)
|
||||
// Only send response if there is one (not for notifications)
|
||||
if response != nil {
|
||||
var message string
|
||||
if eventData, err := json.Marshal(response); err != nil {
|
||||
// If there is an error marshalling the response, send a generic error response
|
||||
log.Printf("failed to marshal response: %v", err)
|
||||
message = "event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n"
|
||||
} else {
|
||||
message = fmt.Sprintf("event: message\ndata: %s\n\n", eventData)
|
||||
}
|
||||
|
||||
// Queue the event for sending via SSE
|
||||
select {
|
||||
case session.eventQueue <- message:
|
||||
// Event queued successfully
|
||||
case <-session.done:
|
||||
// Session is closed, don't try to queue
|
||||
default:
|
||||
// Queue is full, log this situation
|
||||
log.Printf("Event queue full for session %s", sessionID)
|
||||
}
|
||||
}
|
||||
}(messageCtx)
|
||||
}
|
||||
|
||||
// writeJSONRPCError writes a JSON-RPC error response with the given error details.
|
||||
func (s *SSEServer) writeJSONRPCError(
|
||||
w http.ResponseWriter,
|
||||
id any,
|
||||
code int,
|
||||
message string,
|
||||
) {
|
||||
response := createErrorResponse(id, code, message)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
if err := json.NewEncoder(w).Encode(response); err != nil {
|
||||
http.Error(
|
||||
w,
|
||||
fmt.Sprintf("Failed to encode response: %v", err),
|
||||
http.StatusInternalServerError,
|
||||
)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// SendEventToSession sends an event to a specific SSE session identified by sessionID.
|
||||
// Returns an error if the session is not found or closed.
|
||||
func (s *SSEServer) SendEventToSession(
|
||||
sessionID string,
|
||||
event any,
|
||||
) error {
|
||||
sessionI, ok := s.sessions.Load(sessionID)
|
||||
if !ok {
|
||||
return fmt.Errorf("session not found: %s", sessionID)
|
||||
}
|
||||
session := sessionI.(*sseSession)
|
||||
|
||||
eventData, err := json.Marshal(event)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Queue the event for sending via SSE
|
||||
select {
|
||||
case session.eventQueue <- fmt.Sprintf("event: message\ndata: %s\n\n", eventData):
|
||||
return nil
|
||||
case <-session.done:
|
||||
return fmt.Errorf("session closed")
|
||||
default:
|
||||
return fmt.Errorf("event queue full")
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SSEServer) GetUrlPath(input string) (string, error) {
|
||||
parse, err := url.Parse(input)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to parse URL %s: %w", input, err)
|
||||
}
|
||||
return parse.Path, nil
|
||||
}
|
||||
|
||||
func (s *SSEServer) CompleteSseEndpoint() (string, error) {
|
||||
if s.dynamicBasePathFunc != nil {
|
||||
return "", &ErrDynamicPathConfig{Method: "CompleteSseEndpoint"}
|
||||
}
|
||||
|
||||
path := normalizeURLPath(s.basePath, s.sseEndpoint)
|
||||
return s.baseURL + path, nil
|
||||
}
|
||||
|
||||
func (s *SSEServer) CompleteSsePath() string {
|
||||
path, err := s.CompleteSseEndpoint()
|
||||
if err != nil {
|
||||
return normalizeURLPath(s.basePath, s.sseEndpoint)
|
||||
}
|
||||
urlPath, err := s.GetUrlPath(path)
|
||||
if err != nil {
|
||||
return normalizeURLPath(s.basePath, s.sseEndpoint)
|
||||
}
|
||||
return urlPath
|
||||
}
|
||||
|
||||
func (s *SSEServer) CompleteMessageEndpoint() (string, error) {
|
||||
if s.dynamicBasePathFunc != nil {
|
||||
return "", &ErrDynamicPathConfig{Method: "CompleteMessageEndpoint"}
|
||||
}
|
||||
path := normalizeURLPath(s.basePath, s.messageEndpoint)
|
||||
return s.baseURL + path, nil
|
||||
}
|
||||
|
||||
func (s *SSEServer) CompleteMessagePath() string {
|
||||
path, err := s.CompleteMessageEndpoint()
|
||||
if err != nil {
|
||||
return normalizeURLPath(s.basePath, s.messageEndpoint)
|
||||
}
|
||||
urlPath, err := s.GetUrlPath(path)
|
||||
if err != nil {
|
||||
return normalizeURLPath(s.basePath, s.messageEndpoint)
|
||||
}
|
||||
return urlPath
|
||||
}
|
||||
|
||||
// SSEHandler returns an http.Handler for the SSE endpoint.
|
||||
//
|
||||
// This method allows you to mount the SSE handler at any arbitrary path
|
||||
// using your own router (e.g. net/http, gorilla/mux, chi, etc.). It is
|
||||
// intended for advanced scenarios where you want to control the routing or
|
||||
// support dynamic segments.
|
||||
//
|
||||
// IMPORTANT: When using this handler in advanced/dynamic mounting scenarios,
|
||||
// you must use the WithDynamicBasePath option to ensure the correct base path
|
||||
// is communicated to clients.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// // Advanced/dynamic:
|
||||
// sseServer := NewSSEServer(mcpServer,
|
||||
// WithDynamicBasePath(func(r *http.Request, sessionID string) string {
|
||||
// tenant := r.PathValue("tenant")
|
||||
// return "/mcp/" + tenant
|
||||
// }),
|
||||
// WithBaseURL("http://localhost:8080")
|
||||
// )
|
||||
// mux.Handle("/mcp/{tenant}/sse", sseServer.SSEHandler())
|
||||
// mux.Handle("/mcp/{tenant}/message", sseServer.MessageHandler())
|
||||
//
|
||||
// For non-dynamic cases, use ServeHTTP method instead.
|
||||
func (s *SSEServer) SSEHandler() http.Handler {
|
||||
return http.HandlerFunc(s.handleSSE)
|
||||
}
|
||||
|
||||
// MessageHandler returns an http.Handler for the message endpoint.
|
||||
//
|
||||
// This method allows you to mount the message handler at any arbitrary path
|
||||
// using your own router (e.g. net/http, gorilla/mux, chi, etc.). It is
|
||||
// intended for advanced scenarios where you want to control the routing or
|
||||
// support dynamic segments.
|
||||
//
|
||||
// IMPORTANT: When using this handler in advanced/dynamic mounting scenarios,
|
||||
// you must use the WithDynamicBasePath option to ensure the correct base path
|
||||
// is communicated to clients.
|
||||
//
|
||||
// Example usage:
|
||||
//
|
||||
// // Advanced/dynamic:
|
||||
// sseServer := NewSSEServer(mcpServer,
|
||||
// WithDynamicBasePath(func(r *http.Request, sessionID string) string {
|
||||
// tenant := r.PathValue("tenant")
|
||||
// return "/mcp/" + tenant
|
||||
// }),
|
||||
// WithBaseURL("http://localhost:8080")
|
||||
// )
|
||||
// mux.Handle("/mcp/{tenant}/sse", sseServer.SSEHandler())
|
||||
// mux.Handle("/mcp/{tenant}/message", sseServer.MessageHandler())
|
||||
//
|
||||
// For non-dynamic cases, use ServeHTTP method instead.
|
||||
func (s *SSEServer) MessageHandler() http.Handler {
|
||||
return http.HandlerFunc(s.handleMessage)
|
||||
}
|
||||
|
||||
// ServeHTTP implements the http.Handler interface.
|
||||
func (s *SSEServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if s.dynamicBasePathFunc != nil {
|
||||
http.Error(
|
||||
w,
|
||||
(&ErrDynamicPathConfig{Method: "ServeHTTP"}).Error(),
|
||||
http.StatusInternalServerError,
|
||||
)
|
||||
return
|
||||
}
|
||||
path := r.URL.Path
|
||||
// Use exact path matching rather than Contains
|
||||
ssePath := s.CompleteSsePath()
|
||||
if ssePath != "" && path == ssePath {
|
||||
s.handleSSE(w, r)
|
||||
return
|
||||
}
|
||||
messagePath := s.CompleteMessagePath()
|
||||
if messagePath != "" && path == messagePath {
|
||||
s.handleMessage(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
|
||||
// normalizeURLPath joins path elements like path.Join but ensures the
|
||||
// result always starts with a leading slash and never ends with a slash
|
||||
func normalizeURLPath(elem ...string) string {
|
||||
joined := path.Join(elem...)
|
||||
|
||||
// Ensure leading slash
|
||||
if !strings.HasPrefix(joined, "/") {
|
||||
joined = "/" + joined
|
||||
}
|
||||
|
||||
// Remove trailing slash if not just "/"
|
||||
if len(joined) > 1 && strings.HasSuffix(joined, "/") {
|
||||
joined = joined[:len(joined)-1]
|
||||
}
|
||||
|
||||
return joined
|
||||
}
|
||||
+877
@@ -0,0 +1,877 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
)
|
||||
|
||||
// StdioContextFunc is a function that takes an existing context and returns
|
||||
// a potentially modified context.
|
||||
// This can be used to inject context values from environment variables,
|
||||
// for example.
|
||||
type StdioContextFunc func(ctx context.Context) context.Context
|
||||
|
||||
// StdioServer wraps a MCPServer and handles stdio communication.
|
||||
// It provides a simple way to create command-line MCP servers that
|
||||
// communicate via standard input/output streams using JSON-RPC messages.
|
||||
type StdioServer struct {
|
||||
server *MCPServer
|
||||
errLogger *log.Logger
|
||||
contextFunc StdioContextFunc
|
||||
|
||||
// Thread-safe tool call processing
|
||||
toolCallQueue chan *toolCallWork
|
||||
workerWg sync.WaitGroup
|
||||
workerPoolSize int
|
||||
queueSize int
|
||||
writeMu sync.Mutex // Protects concurrent writes
|
||||
}
|
||||
|
||||
// toolCallWork represents a queued tool call request
|
||||
type toolCallWork struct {
|
||||
ctx context.Context
|
||||
message json.RawMessage
|
||||
writer io.Writer
|
||||
}
|
||||
|
||||
// StdioOption defines a function type for configuring StdioServer
|
||||
type StdioOption func(*StdioServer)
|
||||
|
||||
// WithErrorLogger sets the error logger for the server
|
||||
func WithErrorLogger(logger *log.Logger) StdioOption {
|
||||
return func(s *StdioServer) {
|
||||
s.errLogger = logger
|
||||
}
|
||||
}
|
||||
|
||||
// WithStdioContextFunc sets a function that will be called to customise the context
|
||||
// to the server. Note that the stdio server uses the same context for all requests,
|
||||
// so this function will only be called once per server instance.
|
||||
func WithStdioContextFunc(fn StdioContextFunc) StdioOption {
|
||||
return func(s *StdioServer) {
|
||||
s.contextFunc = fn
|
||||
}
|
||||
}
|
||||
|
||||
// WithWorkerPoolSize sets the number of workers for processing tool calls
|
||||
func WithWorkerPoolSize(size int) StdioOption {
|
||||
return func(s *StdioServer) {
|
||||
const maxWorkerPoolSize = 100
|
||||
if size > 0 && size <= maxWorkerPoolSize {
|
||||
s.workerPoolSize = size
|
||||
} else if size > maxWorkerPoolSize {
|
||||
s.errLogger.Printf("Worker pool size %d exceeds maximum (%d), using maximum", size, maxWorkerPoolSize)
|
||||
s.workerPoolSize = maxWorkerPoolSize
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WithQueueSize sets the size of the tool call queue
|
||||
func WithQueueSize(size int) StdioOption {
|
||||
return func(s *StdioServer) {
|
||||
const maxQueueSize = 10000
|
||||
if size > 0 && size <= maxQueueSize {
|
||||
s.queueSize = size
|
||||
} else if size > maxQueueSize {
|
||||
s.errLogger.Printf("Queue size %d exceeds maximum (%d), using maximum", size, maxQueueSize)
|
||||
s.queueSize = maxQueueSize
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// stdioSession is a static client session, since stdio has only one client.
|
||||
type stdioSession struct {
|
||||
notifications chan mcp.JSONRPCNotification
|
||||
initialized atomic.Bool
|
||||
loggingLevel atomic.Value
|
||||
clientInfo atomic.Value // stores session-specific client info
|
||||
clientCapabilities atomic.Value // stores session-specific client capabilities
|
||||
writer io.Writer // for sending requests to client
|
||||
requestID atomic.Int64 // for generating unique request IDs
|
||||
mu sync.RWMutex // protects writer
|
||||
pendingRequests map[int64]chan *samplingResponse // for tracking pending sampling requests
|
||||
pendingElicitations map[int64]chan *elicitationResponse // for tracking pending elicitation requests
|
||||
pendingRoots map[int64]chan *rootsResponse // for tracking pending list roots requests
|
||||
pendingMu sync.RWMutex // protects pendingRequests and pendingElicitations
|
||||
}
|
||||
|
||||
// samplingResponse represents a response to a sampling request
|
||||
type samplingResponse struct {
|
||||
result *mcp.CreateMessageResult
|
||||
err error
|
||||
}
|
||||
|
||||
// elicitationResponse represents a response to an elicitation request
|
||||
type elicitationResponse struct {
|
||||
result *mcp.ElicitationResult
|
||||
err error
|
||||
}
|
||||
|
||||
// rootsResponse represents a response to an list root request
|
||||
type rootsResponse struct {
|
||||
result *mcp.ListRootsResult
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *stdioSession) SessionID() string {
|
||||
return "stdio"
|
||||
}
|
||||
|
||||
func (s *stdioSession) NotificationChannel() chan<- mcp.JSONRPCNotification {
|
||||
return s.notifications
|
||||
}
|
||||
|
||||
func (s *stdioSession) Initialize() {
|
||||
// set default logging level
|
||||
s.loggingLevel.Store(mcp.LoggingLevelError)
|
||||
s.initialized.Store(true)
|
||||
}
|
||||
|
||||
func (s *stdioSession) Initialized() bool {
|
||||
return s.initialized.Load()
|
||||
}
|
||||
|
||||
func (s *stdioSession) GetClientInfo() mcp.Implementation {
|
||||
if value := s.clientInfo.Load(); value != nil {
|
||||
if clientInfo, ok := value.(mcp.Implementation); ok {
|
||||
return clientInfo
|
||||
}
|
||||
}
|
||||
return mcp.Implementation{}
|
||||
}
|
||||
|
||||
func (s *stdioSession) SetClientInfo(clientInfo mcp.Implementation) {
|
||||
s.clientInfo.Store(clientInfo)
|
||||
}
|
||||
|
||||
func (s *stdioSession) GetClientCapabilities() mcp.ClientCapabilities {
|
||||
if value := s.clientCapabilities.Load(); value != nil {
|
||||
if clientCapabilities, ok := value.(mcp.ClientCapabilities); ok {
|
||||
return clientCapabilities
|
||||
}
|
||||
}
|
||||
return mcp.ClientCapabilities{}
|
||||
}
|
||||
|
||||
func (s *stdioSession) SetClientCapabilities(clientCapabilities mcp.ClientCapabilities) {
|
||||
s.clientCapabilities.Store(clientCapabilities)
|
||||
}
|
||||
|
||||
func (s *stdioSession) SetLogLevel(level mcp.LoggingLevel) {
|
||||
s.loggingLevel.Store(level)
|
||||
}
|
||||
|
||||
func (s *stdioSession) GetLogLevel() mcp.LoggingLevel {
|
||||
level := s.loggingLevel.Load()
|
||||
if level == nil {
|
||||
return mcp.LoggingLevelError
|
||||
}
|
||||
return level.(mcp.LoggingLevel)
|
||||
}
|
||||
|
||||
// RequestSampling sends a sampling request to the client and waits for the response.
|
||||
func (s *stdioSession) RequestSampling(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
|
||||
s.mu.RLock()
|
||||
writer := s.writer
|
||||
s.mu.RUnlock()
|
||||
|
||||
if writer == nil {
|
||||
return nil, fmt.Errorf("no writer available for sending requests")
|
||||
}
|
||||
|
||||
// Generate a unique request ID
|
||||
id := s.requestID.Add(1)
|
||||
|
||||
// Create a response channel for this request
|
||||
responseChan := make(chan *samplingResponse, 1)
|
||||
s.pendingMu.Lock()
|
||||
s.pendingRequests[id] = responseChan
|
||||
s.pendingMu.Unlock()
|
||||
|
||||
// Cleanup function to remove the pending request
|
||||
cleanup := func() {
|
||||
s.pendingMu.Lock()
|
||||
delete(s.pendingRequests, id)
|
||||
s.pendingMu.Unlock()
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
// Create the JSON-RPC request
|
||||
jsonRPCRequest := struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID int64 `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params mcp.CreateMessageParams `json:"params"`
|
||||
}{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Method: string(mcp.MethodSamplingCreateMessage),
|
||||
Params: request.CreateMessageParams,
|
||||
}
|
||||
|
||||
// Marshal and send the request
|
||||
requestBytes, err := json.Marshal(jsonRPCRequest)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal sampling request: %w", err)
|
||||
}
|
||||
requestBytes = append(requestBytes, '\n')
|
||||
|
||||
if _, err := writer.Write(requestBytes); err != nil {
|
||||
return nil, fmt.Errorf("failed to write sampling request: %w", err)
|
||||
}
|
||||
|
||||
// Wait for the response or context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case response := <-responseChan:
|
||||
if response.err != nil {
|
||||
return nil, response.err
|
||||
}
|
||||
return response.result, nil
|
||||
}
|
||||
}
|
||||
|
||||
// ListRoots sends an list roots request to the client and waits for the response.
|
||||
func (s *stdioSession) ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error) {
|
||||
s.mu.RLock()
|
||||
writer := s.writer
|
||||
s.mu.RUnlock()
|
||||
|
||||
if writer == nil {
|
||||
return nil, fmt.Errorf("no writer available for sending requests")
|
||||
}
|
||||
|
||||
// Generate a unique request ID
|
||||
id := s.requestID.Add(1)
|
||||
|
||||
// Create a response channel for this request
|
||||
responseChan := make(chan *rootsResponse, 1)
|
||||
s.pendingMu.Lock()
|
||||
s.pendingRoots[id] = responseChan
|
||||
s.pendingMu.Unlock()
|
||||
|
||||
// Cleanup function to remove the pending request
|
||||
cleanup := func() {
|
||||
s.pendingMu.Lock()
|
||||
delete(s.pendingRoots, id)
|
||||
s.pendingMu.Unlock()
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
// Create the JSON-RPC request
|
||||
jsonRPCRequest := struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID int64 `json:"id"`
|
||||
Method string `json:"method"`
|
||||
}{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Method: string(mcp.MethodListRoots),
|
||||
}
|
||||
|
||||
// Marshal and send the request
|
||||
requestBytes, err := json.Marshal(jsonRPCRequest)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal list roots request: %w", err)
|
||||
}
|
||||
requestBytes = append(requestBytes, '\n')
|
||||
|
||||
if _, err := writer.Write(requestBytes); err != nil {
|
||||
return nil, fmt.Errorf("failed to write list roots request: %w", err)
|
||||
}
|
||||
|
||||
// Wait for the response or context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case response := <-responseChan:
|
||||
if response.err != nil {
|
||||
return nil, response.err
|
||||
}
|
||||
return response.result, nil
|
||||
}
|
||||
}
|
||||
|
||||
// RequestElicitation sends an elicitation request to the client and waits for the response.
|
||||
func (s *stdioSession) RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) {
|
||||
s.mu.RLock()
|
||||
writer := s.writer
|
||||
s.mu.RUnlock()
|
||||
|
||||
if writer == nil {
|
||||
return nil, fmt.Errorf("no writer available for sending requests")
|
||||
}
|
||||
|
||||
// Generate a unique request ID
|
||||
id := s.requestID.Add(1)
|
||||
|
||||
// Create a response channel for this request
|
||||
responseChan := make(chan *elicitationResponse, 1)
|
||||
s.pendingMu.Lock()
|
||||
s.pendingElicitations[id] = responseChan
|
||||
s.pendingMu.Unlock()
|
||||
|
||||
// Cleanup function to remove the pending request
|
||||
cleanup := func() {
|
||||
s.pendingMu.Lock()
|
||||
delete(s.pendingElicitations, id)
|
||||
s.pendingMu.Unlock()
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
// Create the JSON-RPC request
|
||||
jsonRPCRequest := struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID int64 `json:"id"`
|
||||
Method string `json:"method"`
|
||||
Params mcp.ElicitationParams `json:"params"`
|
||||
}{
|
||||
JSONRPC: mcp.JSONRPC_VERSION,
|
||||
ID: id,
|
||||
Method: string(mcp.MethodElicitationCreate),
|
||||
Params: request.Params,
|
||||
}
|
||||
|
||||
// Marshal and send the request
|
||||
requestBytes, err := json.Marshal(jsonRPCRequest)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal elicitation request: %w", err)
|
||||
}
|
||||
requestBytes = append(requestBytes, '\n')
|
||||
|
||||
if _, err := writer.Write(requestBytes); err != nil {
|
||||
return nil, fmt.Errorf("failed to write elicitation request: %w", err)
|
||||
}
|
||||
|
||||
// Wait for the response or context cancellation
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
case response := <-responseChan:
|
||||
if response.err != nil {
|
||||
return nil, response.err
|
||||
}
|
||||
return response.result, nil
|
||||
}
|
||||
}
|
||||
|
||||
// SetWriter sets the writer for sending requests to the client.
|
||||
func (s *stdioSession) SetWriter(writer io.Writer) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.writer = writer
|
||||
}
|
||||
|
||||
var (
|
||||
_ ClientSession = (*stdioSession)(nil)
|
||||
_ SessionWithLogging = (*stdioSession)(nil)
|
||||
_ SessionWithClientInfo = (*stdioSession)(nil)
|
||||
_ SessionWithSampling = (*stdioSession)(nil)
|
||||
_ SessionWithElicitation = (*stdioSession)(nil)
|
||||
_ SessionWithRoots = (*stdioSession)(nil)
|
||||
)
|
||||
|
||||
var stdioSessionInstance = stdioSession{
|
||||
notifications: make(chan mcp.JSONRPCNotification, 100),
|
||||
pendingRequests: make(map[int64]chan *samplingResponse),
|
||||
pendingElicitations: make(map[int64]chan *elicitationResponse),
|
||||
pendingRoots: make(map[int64]chan *rootsResponse),
|
||||
}
|
||||
|
||||
// NewStdioServer creates a new stdio server wrapper around an MCPServer.
|
||||
// It initializes the server with a default error logger that discards all output.
|
||||
func NewStdioServer(server *MCPServer) *StdioServer {
|
||||
return &StdioServer{
|
||||
server: server,
|
||||
errLogger: log.New(
|
||||
os.Stderr,
|
||||
"",
|
||||
log.LstdFlags,
|
||||
), // Default to discarding logs
|
||||
workerPoolSize: 5, // Default worker pool size
|
||||
queueSize: 100, // Default queue size
|
||||
}
|
||||
}
|
||||
|
||||
// SetErrorLogger configures where error messages from the StdioServer are logged.
|
||||
// The provided logger will receive all error messages generated during server operation.
|
||||
func (s *StdioServer) SetErrorLogger(logger *log.Logger) {
|
||||
s.errLogger = logger
|
||||
}
|
||||
|
||||
// SetContextFunc sets a function that will be called to customise the context
|
||||
// to the server. Note that the stdio server uses the same context for all requests,
|
||||
// so this function will only be called once per server instance.
|
||||
func (s *StdioServer) SetContextFunc(fn StdioContextFunc) {
|
||||
s.contextFunc = fn
|
||||
}
|
||||
|
||||
// handleNotifications continuously processes notifications from the session's notification channel
|
||||
// and writes them to the provided output. It runs until the context is cancelled.
|
||||
// Any errors encountered while writing notifications are logged but do not stop the handler.
|
||||
func (s *StdioServer) handleNotifications(ctx context.Context, stdout io.Writer) {
|
||||
for {
|
||||
select {
|
||||
case notification := <-stdioSessionInstance.notifications:
|
||||
if err := s.writeResponse(notification, stdout); err != nil {
|
||||
s.errLogger.Printf("Error writing notification: %v", err)
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// processInputStream continuously reads and processes messages from the input stream.
|
||||
// It handles EOF gracefully as a normal termination condition.
|
||||
// The function returns when either:
|
||||
// - The context is cancelled (returns context.Err())
|
||||
// - EOF is encountered (returns nil)
|
||||
// - An error occurs while reading or processing messages (returns the error)
|
||||
func (s *StdioServer) processInputStream(ctx context.Context, reader *bufio.Reader, stdout io.Writer) error {
|
||||
for {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
line, err := s.readNextLine(ctx, reader)
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
return nil
|
||||
}
|
||||
s.errLogger.Printf("Error reading input: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
if err := s.processMessage(ctx, line, stdout); err != nil {
|
||||
if err == io.EOF {
|
||||
return nil
|
||||
}
|
||||
s.errLogger.Printf("Error handling message: %v", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// toolCallWorker processes tool calls from the queue
|
||||
func (s *StdioServer) toolCallWorker(ctx context.Context) {
|
||||
defer s.workerWg.Done()
|
||||
|
||||
for {
|
||||
select {
|
||||
case work, ok := <-s.toolCallQueue:
|
||||
if !ok {
|
||||
// Channel closed, exit worker
|
||||
return
|
||||
}
|
||||
// Process the tool call
|
||||
response := s.server.HandleMessage(work.ctx, work.message)
|
||||
if response != nil {
|
||||
if err := s.writeResponse(response, work.writer); err != nil {
|
||||
s.errLogger.Printf("Error writing tool response: %v", err)
|
||||
}
|
||||
}
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// readNextLine reads a single line from the input reader in a context-aware manner.
|
||||
// It uses channels to make the read operation cancellable via context.
|
||||
// Returns the read line and any error encountered. If the context is cancelled,
|
||||
// returns an empty string and the context's error. EOF is returned when the input
|
||||
// stream is closed.
|
||||
func (s *StdioServer) readNextLine(ctx context.Context, reader *bufio.Reader) (string, error) {
|
||||
type result struct {
|
||||
line string
|
||||
err error
|
||||
}
|
||||
|
||||
resultCh := make(chan result, 1)
|
||||
|
||||
go func() {
|
||||
line, err := reader.ReadString('\n')
|
||||
resultCh <- result{line: line, err: err}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", nil
|
||||
case res := <-resultCh:
|
||||
return res.line, res.err
|
||||
}
|
||||
}
|
||||
|
||||
// Listen starts listening for JSON-RPC messages on the provided input and writes responses to the provided output.
|
||||
// It runs until the context is cancelled or an error occurs.
|
||||
// Returns an error if there are issues with reading input or writing output.
|
||||
func (s *StdioServer) Listen(
|
||||
ctx context.Context,
|
||||
stdin io.Reader,
|
||||
stdout io.Writer,
|
||||
) error {
|
||||
// Initialize the tool call queue
|
||||
s.toolCallQueue = make(chan *toolCallWork, s.queueSize)
|
||||
|
||||
// Set a static client context since stdio only has one client
|
||||
if err := s.server.RegisterSession(ctx, &stdioSessionInstance); err != nil {
|
||||
return fmt.Errorf("register session: %w", err)
|
||||
}
|
||||
defer s.server.UnregisterSession(ctx, stdioSessionInstance.SessionID())
|
||||
ctx = s.server.WithContext(ctx, &stdioSessionInstance)
|
||||
|
||||
// Set the writer for sending requests to the client
|
||||
stdioSessionInstance.SetWriter(stdout)
|
||||
|
||||
// Add in any custom context.
|
||||
if s.contextFunc != nil {
|
||||
ctx = s.contextFunc(ctx)
|
||||
}
|
||||
|
||||
reader := bufio.NewReader(stdin)
|
||||
|
||||
// Start worker pool for tool calls
|
||||
for i := 0; i < s.workerPoolSize; i++ {
|
||||
s.workerWg.Add(1)
|
||||
go s.toolCallWorker(ctx)
|
||||
}
|
||||
|
||||
// Start notification handler
|
||||
go s.handleNotifications(ctx, stdout)
|
||||
|
||||
// Process input stream
|
||||
err := s.processInputStream(ctx, reader, stdout)
|
||||
|
||||
// Shutdown workers gracefully
|
||||
close(s.toolCallQueue)
|
||||
s.workerWg.Wait()
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// processMessage handles a single JSON-RPC message and writes the response.
|
||||
// It parses the message, processes it through the wrapped MCPServer, and writes any response.
|
||||
// Returns an error if there are issues with message processing or response writing.
|
||||
func (s *StdioServer) processMessage(
|
||||
ctx context.Context,
|
||||
line string,
|
||||
writer io.Writer,
|
||||
) error {
|
||||
// If line is empty, likely due to ctx cancellation
|
||||
if len(line) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Parse the message as raw JSON
|
||||
var rawMessage json.RawMessage
|
||||
if err := json.Unmarshal([]byte(line), &rawMessage); err != nil {
|
||||
response := createErrorResponse(nil, mcp.PARSE_ERROR, "Parse error")
|
||||
return s.writeResponse(response, writer)
|
||||
}
|
||||
|
||||
// Check if this is a response to a sampling request
|
||||
if s.handleSamplingResponse(rawMessage) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if this is a response to an elicitation request
|
||||
if s.handleElicitationResponse(rawMessage) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if this is a response to an list roots request
|
||||
if s.handleListRootsResponse(rawMessage) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if this is a tool call that might need sampling (and thus should be processed concurrently)
|
||||
var baseMessage struct {
|
||||
Method string `json:"method"`
|
||||
}
|
||||
if json.Unmarshal(rawMessage, &baseMessage) == nil && baseMessage.Method == "tools/call" {
|
||||
// Queue tool calls for processing by workers
|
||||
select {
|
||||
case s.toolCallQueue <- &toolCallWork{
|
||||
ctx: ctx,
|
||||
message: rawMessage,
|
||||
writer: writer,
|
||||
}:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
default:
|
||||
// Queue is full, process synchronously as fallback
|
||||
s.errLogger.Printf("Tool call queue full, processing synchronously")
|
||||
response := s.server.HandleMessage(ctx, rawMessage)
|
||||
if response != nil {
|
||||
return s.writeResponse(response, writer)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// Handle other messages synchronously
|
||||
response := s.server.HandleMessage(ctx, rawMessage)
|
||||
|
||||
// Only write response if there is one (not for notifications)
|
||||
if response != nil {
|
||||
if err := s.writeResponse(response, writer); err != nil {
|
||||
return fmt.Errorf("failed to write response: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleSamplingResponse checks if the message is a response to a sampling request
|
||||
// and routes it to the appropriate pending request channel.
|
||||
func (s *StdioServer) handleSamplingResponse(rawMessage json.RawMessage) bool {
|
||||
return stdioSessionInstance.handleSamplingResponse(rawMessage)
|
||||
}
|
||||
|
||||
// handleSamplingResponse handles incoming sampling responses for this session
|
||||
func (s *stdioSession) handleSamplingResponse(rawMessage json.RawMessage) bool {
|
||||
// Try to parse as a JSON-RPC response
|
||||
var response struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.Number `json:"id"`
|
||||
Result json.RawMessage `json:"result,omitempty"`
|
||||
Error *struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(rawMessage, &response); err != nil {
|
||||
return false
|
||||
}
|
||||
// Parse the ID as int64
|
||||
idInt64, err := response.ID.Int64()
|
||||
if err != nil || (response.Result == nil && response.Error == nil) {
|
||||
return false
|
||||
}
|
||||
|
||||
// Look for a pending request with this ID
|
||||
s.pendingMu.RLock()
|
||||
responseChan, exists := s.pendingRequests[idInt64]
|
||||
s.pendingMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return false
|
||||
} // Parse and send the response
|
||||
samplingResp := &samplingResponse{}
|
||||
|
||||
if response.Error != nil {
|
||||
samplingResp.err = fmt.Errorf("sampling request failed: %s", response.Error.Message)
|
||||
} else {
|
||||
var result mcp.CreateMessageResult
|
||||
if err := json.Unmarshal(response.Result, &result); err != nil {
|
||||
samplingResp.err = fmt.Errorf("failed to unmarshal sampling response: %w", err)
|
||||
} else {
|
||||
// Parse content from map[string]any to proper Content type (TextContent, ImageContent, AudioContent)
|
||||
if contentMap, ok := result.Content.(map[string]any); ok {
|
||||
content, err := mcp.ParseContent(contentMap)
|
||||
if err != nil {
|
||||
samplingResp.err = fmt.Errorf("failed to parse sampling response content: %w", err)
|
||||
} else {
|
||||
result.Content = content
|
||||
samplingResp.result = &result
|
||||
}
|
||||
} else {
|
||||
samplingResp.result = &result
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Send the response (non-blocking)
|
||||
select {
|
||||
case responseChan <- samplingResp:
|
||||
default:
|
||||
// Channel is full or closed, ignore
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// handleElicitationResponse checks if the message is a response to an elicitation request
|
||||
// and routes it to the appropriate pending request channel.
|
||||
func (s *StdioServer) handleElicitationResponse(rawMessage json.RawMessage) bool {
|
||||
return stdioSessionInstance.handleElicitationResponse(rawMessage)
|
||||
}
|
||||
|
||||
// handleElicitationResponse handles incoming elicitation responses for this session
|
||||
func (s *stdioSession) handleElicitationResponse(rawMessage json.RawMessage) bool {
|
||||
// Try to parse as a JSON-RPC response
|
||||
var response struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.Number `json:"id"`
|
||||
Result json.RawMessage `json:"result,omitempty"`
|
||||
Error *struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(rawMessage, &response); err != nil {
|
||||
return false
|
||||
}
|
||||
// Parse the ID as int64
|
||||
id, err := response.ID.Int64()
|
||||
if err != nil || (response.Result == nil && response.Error == nil) {
|
||||
return false
|
||||
}
|
||||
|
||||
// Check if we have a pending elicitation request with this ID
|
||||
s.pendingMu.RLock()
|
||||
responseChan, exists := s.pendingElicitations[id]
|
||||
s.pendingMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return false
|
||||
}
|
||||
|
||||
// Parse and send the response
|
||||
elicitationResp := &elicitationResponse{}
|
||||
|
||||
if response.Error != nil {
|
||||
elicitationResp.err = fmt.Errorf("elicitation request failed: %s", response.Error.Message)
|
||||
} else {
|
||||
var result mcp.ElicitationResult
|
||||
if err := json.Unmarshal(response.Result, &result); err != nil {
|
||||
elicitationResp.err = fmt.Errorf("failed to unmarshal elicitation response: %w", err)
|
||||
} else {
|
||||
elicitationResp.result = &result
|
||||
}
|
||||
}
|
||||
|
||||
// Send the response (non-blocking)
|
||||
select {
|
||||
case responseChan <- elicitationResp:
|
||||
default:
|
||||
// Channel is full or closed, ignore
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// handleListRootsResponse checks if the message is a response to an list roots request
|
||||
// and routes it to the appropriate pending request channel.
|
||||
func (s *StdioServer) handleListRootsResponse(rawMessage json.RawMessage) bool {
|
||||
return stdioSessionInstance.handleListRootsResponse(rawMessage)
|
||||
}
|
||||
|
||||
// handleListRootsResponse handles incoming list root responses for this session
|
||||
func (s *stdioSession) handleListRootsResponse(rawMessage json.RawMessage) bool {
|
||||
// Try to parse as a JSON-RPC response
|
||||
var response struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.Number `json:"id"`
|
||||
Result json.RawMessage `json:"result,omitempty"`
|
||||
Error *struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(rawMessage, &response); err != nil {
|
||||
return false
|
||||
}
|
||||
// Parse the ID as int64
|
||||
id, err := response.ID.Int64()
|
||||
if err != nil || (response.Result == nil && response.Error == nil) {
|
||||
return false
|
||||
}
|
||||
|
||||
// Check if we have a pending list root request with this ID
|
||||
s.pendingMu.RLock()
|
||||
responseChan, exists := s.pendingRoots[id]
|
||||
s.pendingMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return false
|
||||
}
|
||||
|
||||
// Parse and send the response
|
||||
rootsResp := &rootsResponse{}
|
||||
|
||||
if response.Error != nil {
|
||||
rootsResp.err = fmt.Errorf("list root request failed: %s", response.Error.Message)
|
||||
} else {
|
||||
var result mcp.ListRootsResult
|
||||
if err := json.Unmarshal(response.Result, &result); err != nil {
|
||||
rootsResp.err = fmt.Errorf("failed to unmarshal list root response: %w", err)
|
||||
} else {
|
||||
rootsResp.result = &result
|
||||
}
|
||||
}
|
||||
|
||||
// Send the response (non-blocking)
|
||||
select {
|
||||
case responseChan <- rootsResp:
|
||||
default:
|
||||
// Channel is full or closed, ignore
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// writeResponse marshals and writes a JSON-RPC response message followed by a newline.
|
||||
// Returns an error if marshaling or writing fails.
|
||||
func (s *StdioServer) writeResponse(
|
||||
response mcp.JSONRPCMessage,
|
||||
writer io.Writer,
|
||||
) error {
|
||||
responseBytes, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Protect concurrent writes
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
|
||||
// Write response followed by newline
|
||||
if _, err := fmt.Fprintf(writer, "%s\n", responseBytes); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServeStdio is a convenience function that creates and starts a StdioServer with os.Stdin and os.Stdout.
|
||||
// It sets up signal handling for graceful shutdown on SIGTERM and SIGINT.
|
||||
// Returns an error if the server encounters any issues during operation.
|
||||
func ServeStdio(server *MCPServer, opts ...StdioOption) error {
|
||||
s := NewStdioServer(server)
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(s)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
// Set up signal handling
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, syscall.SIGTERM, syscall.SIGINT)
|
||||
|
||||
go func() {
|
||||
<-sigChan
|
||||
cancel()
|
||||
}()
|
||||
|
||||
return s.Listen(ctx, os.Stdin, os.Stdout)
|
||||
}
|
||||
+1434
File diff suppressed because it is too large.
Load diff
+33
@@ -0,0 +1,33 @@
|
||||
package util
|
||||
|
||||
import (
|
||||
"log"
|
||||
)
|
||||
|
||||
// Logger defines a minimal logging interface
|
||||
type Logger interface {
|
||||
Infof(format string, v ...any)
|
||||
Errorf(format string, v ...any)
|
||||
}
|
||||
|
||||
// --- Standard Library Logger Wrapper ---
|
||||
|
||||
// DefaultStdLogger implements Logger using the standard library's log.Logger.
|
||||
func DefaultLogger() Logger {
|
||||
return &stdLogger{
|
||||
logger: log.Default(),
|
||||
}
|
||||
}
|
||||
|
||||
// stdLogger wraps the standard library's log.Logger.
|
||||
type stdLogger struct {
|
||||
logger *log.Logger
|
||||
}
|
||||
|
||||
func (l *stdLogger) Infof(format string, v ...any) {
|
||||
l.logger.Printf("INFO: "+format, v...)
|
||||
}
|
||||
|
||||
func (l *stdLogger) Errorf(format string, v ...any) {
|
||||
l.logger.Printf("ERROR: "+format, v...)
|
||||
}
|
||||
Generated
+25
@@ -0,0 +1,25 @@
|
||||
# Compiled Object files, Static and Dynamic libs (Shared Objects)
|
||||
*.o
|
||||
*.a
|
||||
*.so
|
||||
|
||||
# Folders
|
||||
_obj
|
||||
_test
|
||||
|
||||
# Architecture specific extensions/prefixes
|
||||
*.[568vq]
|
||||
[568vq].out
|
||||
|
||||
*.cgo1.go
|
||||
*.cgo2.c
|
||||
_cgo_defun.c
|
||||
_cgo_gotypes.go
|
||||
_cgo_export.*
|
||||
|
||||
_testmain.go
|
||||
|
||||
*.exe
|
||||
*.test
|
||||
|
||||
*.bench
|
||||
+21
@@ -0,0 +1,21 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2014 Steve Francia
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
GOVERSION := $(shell go version | cut -d ' ' -f 3 | cut -d '.' -f 2)
|
||||
|
||||
.PHONY: check fmt lint test test-race vet test-cover-html help
|
||||
.DEFAULT_GOAL := help
|
||||
|
||||
check: test-race fmt vet lint ## Run tests and linters
|
||||
|
||||
test: ## Run tests
|
||||
go test ./...
|
||||
|
||||
test-race: ## Run tests with race detector
|
||||
go test -race ./...
|
||||
|
||||
fmt: ## Run gofmt linter
|
||||
ifeq "$(GOVERSION)" "12"
|
||||
@for d in `go list` ; do \
|
||||
if [ "`gofmt -l -s $$GOPATH/src/$$d | tee /dev/stderr`" ]; then \
|
||||
echo "^ improperly formatted go files" && echo && exit 1; \
|
||||
fi \
|
||||
done
|
||||
endif
|
||||
|
||||
lint: ## Run golint linter
|
||||
@for d in `go list` ; do \
|
||||
if [ "`golint $$d | tee /dev/stderr`" ]; then \
|
||||
echo "^ golint errors!" && echo && exit 1; \
|
||||
fi \
|
||||
done
|
||||
|
||||
vet: ## Run go vet linter
|
||||
@if [ "`go vet | tee /dev/stderr`" ]; then \
|
||||
echo "^ go vet errors!" && echo && exit 1; \
|
||||
fi
|
||||
|
||||
test-cover-html: ## Generate test coverage report
|
||||
go test -coverprofile=coverage.out -covermode=count
|
||||
go tool cover -func=coverage.out
|
||||
|
||||
help:
|
||||
@grep -E '^[a-zA-Z0-9_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}'
|
||||
+75
@@ -0,0 +1,75 @@
|
||||
# cast
|
||||
|
||||
[](https://github.com/spf13/cast/actions/workflows/test.yaml)
|
||||
[](https://pkg.go.dev/mod/github.com/spf13/cast)
|
||||

|
||||
[](https://goreportcard.com/report/github.com/spf13/cast)
|
||||
|
||||
Easy and safe casting from one type to another in Go
|
||||
|
||||
Don’t Panic! ... Cast
|
||||
|
||||
## What is Cast?
|
||||
|
||||
Cast is a library to convert between different go types in a consistent and easy way.
|
||||
|
||||
Cast provides simple functions to easily convert a number to a string, an
|
||||
interface into a bool, etc. Cast does this intelligently when an obvious
|
||||
conversion is possible. It doesn’t make any attempts to guess what you meant,
|
||||
for example you can only convert a string to an int when it is a string
|
||||
representation of an int such as “8”. Cast was developed for use in
|
||||
[Hugo](https://gohugo.io), a website engine which uses YAML, TOML or JSON
|
||||
for meta data.
|
||||
|
||||
## Why use Cast?
|
||||
|
||||
When working with dynamic data in Go you often need to cast or convert the data
|
||||
from one type into another. Cast goes beyond just using type assertion (though
|
||||
it uses that when possible) to provide a very straightforward and convenient
|
||||
library.
|
||||
|
||||
If you are working with interfaces to handle things like dynamic content
|
||||
you’ll need an easy way to convert an interface into a given type. This
|
||||
is the library for you.
|
||||
|
||||
If you are taking in data from YAML, TOML or JSON or other formats which lack
|
||||
full types, then Cast is the library for you.
|
||||
|
||||
## Usage
|
||||
|
||||
Cast provides a handful of To_____ methods. These methods will always return
|
||||
the desired type. **If input is provided that will not convert to that type, the
|
||||
0 or nil value for that type will be returned**.
|
||||
|
||||
Cast also provides identical methods To_____E. These return the same result as
|
||||
the To_____ methods, plus an additional error which tells you if it successfully
|
||||
converted. Using these methods you can tell the difference between when the
|
||||
input matched the zero value or when the conversion failed and the zero value
|
||||
was returned.
|
||||
|
||||
The following examples are merely a sample of what is available. Please review
|
||||
the code for a complete set.
|
||||
|
||||
### Example ‘ToString’:
|
||||
|
||||
cast.ToString("mayonegg") // "mayonegg"
|
||||
cast.ToString(8) // "8"
|
||||
cast.ToString(8.31) // "8.31"
|
||||
cast.ToString([]byte("one time")) // "one time"
|
||||
cast.ToString(nil) // ""
|
||||
|
||||
var foo interface{} = "one more time"
|
||||
cast.ToString(foo) // "one more time"
|
||||
|
||||
|
||||
### Example ‘ToInt’:
|
||||
|
||||
cast.ToInt(8) // 8
|
||||
cast.ToInt(8.31) // 8
|
||||
cast.ToInt("8") // 8
|
||||
cast.ToInt(true) // 1
|
||||
cast.ToInt(false) // 0
|
||||
|
||||
var eight interface{} = 8
|
||||
cast.ToInt(eight) // 8
|
||||
cast.ToInt(nil) // 0
|
||||
+176
@@ -0,0 +1,176 @@
|
||||
// Copyright © 2014 Steve Francia <[email protected]>.
|
||||
//
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package cast provides easy and safe casting in Go.
|
||||
package cast
|
||||
|
||||
import "time"
|
||||
|
||||
// ToBool casts an interface to a bool type.
|
||||
func ToBool(i interface{}) bool {
|
||||
v, _ := ToBoolE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToTime casts an interface to a time.Time type.
|
||||
func ToTime(i interface{}) time.Time {
|
||||
v, _ := ToTimeE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
func ToTimeInDefaultLocation(i interface{}, location *time.Location) time.Time {
|
||||
v, _ := ToTimeInDefaultLocationE(i, location)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToDuration casts an interface to a time.Duration type.
|
||||
func ToDuration(i interface{}) time.Duration {
|
||||
v, _ := ToDurationE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToFloat64 casts an interface to a float64 type.
|
||||
func ToFloat64(i interface{}) float64 {
|
||||
v, _ := ToFloat64E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToFloat32 casts an interface to a float32 type.
|
||||
func ToFloat32(i interface{}) float32 {
|
||||
v, _ := ToFloat32E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToInt64 casts an interface to an int64 type.
|
||||
func ToInt64(i interface{}) int64 {
|
||||
v, _ := ToInt64E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToInt32 casts an interface to an int32 type.
|
||||
func ToInt32(i interface{}) int32 {
|
||||
v, _ := ToInt32E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToInt16 casts an interface to an int16 type.
|
||||
func ToInt16(i interface{}) int16 {
|
||||
v, _ := ToInt16E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToInt8 casts an interface to an int8 type.
|
||||
func ToInt8(i interface{}) int8 {
|
||||
v, _ := ToInt8E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToInt casts an interface to an int type.
|
||||
func ToInt(i interface{}) int {
|
||||
v, _ := ToIntE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToUint casts an interface to a uint type.
|
||||
func ToUint(i interface{}) uint {
|
||||
v, _ := ToUintE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToUint64 casts an interface to a uint64 type.
|
||||
func ToUint64(i interface{}) uint64 {
|
||||
v, _ := ToUint64E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToUint32 casts an interface to a uint32 type.
|
||||
func ToUint32(i interface{}) uint32 {
|
||||
v, _ := ToUint32E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToUint16 casts an interface to a uint16 type.
|
||||
func ToUint16(i interface{}) uint16 {
|
||||
v, _ := ToUint16E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToUint8 casts an interface to a uint8 type.
|
||||
func ToUint8(i interface{}) uint8 {
|
||||
v, _ := ToUint8E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToString casts an interface to a string type.
|
||||
func ToString(i interface{}) string {
|
||||
v, _ := ToStringE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToStringMapString casts an interface to a map[string]string type.
|
||||
func ToStringMapString(i interface{}) map[string]string {
|
||||
v, _ := ToStringMapStringE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToStringMapStringSlice casts an interface to a map[string][]string type.
|
||||
func ToStringMapStringSlice(i interface{}) map[string][]string {
|
||||
v, _ := ToStringMapStringSliceE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToStringMapBool casts an interface to a map[string]bool type.
|
||||
func ToStringMapBool(i interface{}) map[string]bool {
|
||||
v, _ := ToStringMapBoolE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToStringMapInt casts an interface to a map[string]int type.
|
||||
func ToStringMapInt(i interface{}) map[string]int {
|
||||
v, _ := ToStringMapIntE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToStringMapInt64 casts an interface to a map[string]int64 type.
|
||||
func ToStringMapInt64(i interface{}) map[string]int64 {
|
||||
v, _ := ToStringMapInt64E(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToStringMap casts an interface to a map[string]interface{} type.
|
||||
func ToStringMap(i interface{}) map[string]interface{} {
|
||||
v, _ := ToStringMapE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToSlice casts an interface to a []interface{} type.
|
||||
func ToSlice(i interface{}) []interface{} {
|
||||
v, _ := ToSliceE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToBoolSlice casts an interface to a []bool type.
|
||||
func ToBoolSlice(i interface{}) []bool {
|
||||
v, _ := ToBoolSliceE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToStringSlice casts an interface to a []string type.
|
||||
func ToStringSlice(i interface{}) []string {
|
||||
v, _ := ToStringSliceE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToIntSlice casts an interface to a []int type.
|
||||
func ToIntSlice(i interface{}) []int {
|
||||
v, _ := ToIntSliceE(i)
|
||||
return v
|
||||
}
|
||||
|
||||
// ToDurationSlice casts an interface to a []time.Duration type.
|
||||
func ToDurationSlice(i interface{}) []time.Duration {
|
||||
v, _ := ToDurationSliceE(i)
|
||||
return v
|
||||
}
|
||||
Loaded 100 of 149 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user