Skip to content

Commit ec996fd

Browse files
GelembjukGelembjuk
authored
Add support of config for SSE servers. Add Authorization header support (#24)
* Add support of config for SSE servers. Add Authorization header support * Change the format of JSON for SSE servers to be similar to other tools --------- Co-authored-by: Roman Gelembjuk <Roman Gelembjuk>
1 parent c572625 commit ec996fd

File tree

4 files changed

+121
-28
lines changed

4 files changed

+121
-28
lines changed

‎cmd/mcp.go‎

Lines changed: 113 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@ import (
2121
"github.com/mark3labs/mcphost/pkg/llm"
2222
)
2323

24+
const (
25+
transportStdio = "stdio"
26+
transportSSE = "sse"
27+
)
28+
2429
var (
2530
// Tokyo Night theme colors
2631
tokyoPurple = lipgloss.Color("99") // #9d7cd8
@@ -60,15 +65,66 @@ var (
6065
)
6166

6267
type MCPConfig struct {
63-
MCPServers map[string]ServerConfig `json:"mcpServers"`
68+
MCPServers map[string]ServerConfigWrapper `json:"mcpServers"`
6469
}
6570

66-
type ServerConfig struct {
71+
type ServerConfig interface {
72+
GetType() string
73+
}
74+
75+
type STDIOServerConfig struct {
6776
Command string `json:"command"`
6877
Args []string `json:"args"`
6978
Env map[string]string `json:"env,omitempty"`
7079
}
7180

81+
func (s STDIOServerConfig) GetType() string {
82+
return transportStdio
83+
}
84+
85+
type SSEServerConfig struct {
86+
Url string `json:"url"`
87+
Headers []string `json:"headers,omitempty"`
88+
}
89+
90+
func (s SSEServerConfig) GetType() string {
91+
return transportSSE
92+
}
93+
94+
type ServerConfigWrapper struct {
95+
Config ServerConfig
96+
}
97+
98+
func (w *ServerConfigWrapper) UnmarshalJSON(data []byte) error {
99+
var typeField struct {
100+
Url string `json:"url"`
101+
}
102+
103+
if err := json.Unmarshal(data, &typeField); err != nil {
104+
return err
105+
}
106+
if typeField.Url != "" {
107+
// If the URL field is present, treat it as an SSE server
108+
var sse SSEServerConfig
109+
if err := json.Unmarshal(data, &sse); err != nil {
110+
return err
111+
}
112+
w.Config = sse
113+
} else {
114+
// Otherwise, treat it as a STDIOServerConfig
115+
var stdio STDIOServerConfig
116+
if err := json.Unmarshal(data, &stdio); err != nil {
117+
return err
118+
}
119+
w.Config = stdio
120+
}
121+
122+
return nil
123+
}
124+
func (w ServerConfigWrapper) MarshalJSON() ([]byte, error) {
125+
return json.Marshal(w.Config)
126+
}
127+
72128
func mcpToolsToAnthropicTools(
73129
serverName string,
74130
mcpTools []mcp.Tool,
@@ -108,7 +164,7 @@ func loadMCPConfig() (*MCPConfig, error) {
108164
if _, err := os.Stat(configPath); os.IsNotExist(err) {
109165
// Create default config
110166
defaultConfig := MCPConfig{
111-
MCPServers: make(map[string]ServerConfig),
167+
MCPServers: make(map[string]ServerConfigWrapper),
112168
}
113169

114170
// Create the file with default config
@@ -149,31 +205,45 @@ func createMCPClients(
149205
clients := make(map[string]mcpclient.MCPClient)
150206

151207
for name, server := range config.MCPServers {
152-
var env []string
153-
for k, v := range server.Env {
154-
env = append(env, fmt.Sprintf("%s=%s", k, v))
155-
}
156208
var client mcpclient.MCPClient
157209
var err error
158210

159-
if server.Command == "sse_server" {
160-
if len(server.Args) == 0 {
161-
return nil, fmt.Errorf(
162-
"no arguments provided for sse command",
163-
)
211+
if server.Config.GetType() == transportSSE {
212+
sseConfig := server.Config.(SSEServerConfig)
213+
214+
options := []mcpclient.ClientOption{}
215+
216+
if sseConfig.Headers != nil {
217+
// Parse headers from the config
218+
headers := make(map[string]string)
219+
for _, header := range sseConfig.Headers {
220+
parts := strings.SplitN(header, ":", 2)
221+
if len(parts) == 2 {
222+
key := strings.TrimSpace(parts[0])
223+
value := strings.TrimSpace(parts[1])
224+
headers[key] = value
225+
}
226+
}
227+
options = append(options, mcpclient.WithHeaders(headers))
164228
}
165229

166230
client, err = mcpclient.NewSSEMCPClient(
167-
server.Args[0],
231+
sseConfig.Url,
232+
options...,
168233
)
169234
if err == nil {
170235
err = client.(*mcpclient.SSEMCPClient).Start(context.Background())
171236
}
172237
} else {
238+
stdioConfig := server.Config.(STDIOServerConfig)
239+
var env []string
240+
for k, v := range stdioConfig.Env {
241+
env = append(env, fmt.Sprintf("%s=%s", k, v))
242+
}
173243
client, err = mcpclient.NewStdioMCPClient(
174-
server.Command,
244+
stdioConfig.Command,
175245
env,
176-
server.Args...)
246+
stdioConfig.Args...)
177247
}
178248
if err != nil {
179249
for _, c := range clients {
@@ -310,15 +380,37 @@ func handleServersCommand(config *MCPConfig) {
310380
} else {
311381
for name, server := range config.MCPServers {
312382
markdown.WriteString(fmt.Sprintf("# %s\n\n", name))
313-
markdown.WriteString("*Command*\n")
314-
markdown.WriteString(fmt.Sprintf("`%s`\n\n", server.Command))
315383

316-
markdown.WriteString("*Arguments*\n")
317-
if len(server.Args) > 0 {
318-
markdown.WriteString(fmt.Sprintf("`%s`\n", strings.Join(server.Args, " ")))
384+
if server.Config.GetType() == transportSSE {
385+
sseConfig := server.Config.(SSEServerConfig)
386+
markdown.WriteString("*Url*\n")
387+
markdown.WriteString(fmt.Sprintf("`%s`\n\n", sseConfig.Url))
388+
markdown.WriteString("*headers*\n")
389+
if sseConfig.Headers != nil {
390+
for _, header := range sseConfig.Headers {
391+
parts := strings.SplitN(header, ":", 2)
392+
if len(parts) == 2 {
393+
key := strings.TrimSpace(parts[0])
394+
markdown.WriteString("`" + key + ": [REDACTED]`\n")
395+
}
396+
}
397+
} else {
398+
markdown.WriteString("*None*\n")
399+
}
400+
319401
} else {
320-
markdown.WriteString("*None*\n")
402+
stdioConfig := server.Config.(STDIOServerConfig)
403+
markdown.WriteString("*Command*\n")
404+
markdown.WriteString(fmt.Sprintf("`%s`\n\n", stdioConfig.Command))
405+
406+
markdown.WriteString("*Arguments*\n")
407+
if len(stdioConfig.Args) > 0 {
408+
markdown.WriteString(fmt.Sprintf("`%s`\n", strings.Join(stdioConfig.Args, " ")))
409+
} else {
410+
markdown.WriteString("*None*\n")
411+
}
321412
}
413+
322414
markdown.WriteString("\n") // Add spacing between servers
323415
}
324416
}

‎cmd/root.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -431,10 +431,8 @@ func runPrompt(
431431
var resultText string
432432
// Handle array content directly since we know it's []interface{}
433433
for _, item := range toolResult.Content {
434-
if contentMap, ok := item.(map[string]interface{}); ok {
435-
if text, ok := contentMap["text"]; ok {
436-
resultText += fmt.Sprintf("%v ", text)
437-
}
434+
if contentMap, ok := item.(mcp.TextContent); ok {
435+
resultText += fmt.Sprintf("%v ", contentMap.Text)
438436
}
439437
}
440438

‎go.mod‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ require (
88
github.com/charmbracelet/lipgloss v1.0.0
99
github.com/charmbracelet/log v0.4.0
1010
github.com/google/generative-ai-go v0.19.0
11-
github.com/mark3labs/mcp-go v0.8.2
11+
github.com/mark3labs/mcp-go v0.20.0
1212
github.com/ollama/ollama v0.5.1
1313
github.com/spf13/cobra v1.8.1
1414
golang.org/x/term v0.30.0
@@ -37,6 +37,7 @@ require (
3737
github.com/gorilla/css v1.0.1 // indirect
3838
github.com/microcosm-cc/bluemonday v1.0.27 // indirect
3939
github.com/muesli/reflow v0.3.0 // indirect
40+
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
4041
github.com/yuin/goldmark v1.7.4 // indirect
4142
github.com/yuin/goldmark-emoji v1.0.3 // indirect
4243
go.opentelemetry.io/auto/sdk v1.1.0 // indirect

‎go.sum‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,8 @@ github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2
8686
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
8787
github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
8888
github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
89-
github.com/mark3labs/mcp-go v0.8.2 h1:OtqqXlRqjXs6zuMhf1uiuQ2iqBrhMGgLpDeVDUWMKFc=
90-
github.com/mark3labs/mcp-go v0.8.2/go.mod h1:cjMlBU0cv/cj9kjlgmRhoJ5JREdS7YX83xeIG9Ko/jE=
89+
github.com/mark3labs/mcp-go v0.20.0 h1:NYZDZ10GBKHVz4SdQ2tPFSDFQFKCTrTZJLn4wj6jAaw=
90+
github.com/mark3labs/mcp-go v0.20.0/go.mod h1:KmJndYv7GIgcPVwEKJjNcbhVQ+hJGJhrCCB/9xITzpE=
9191
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
9292
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
9393
github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4=
@@ -120,6 +120,8 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
120120
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
121121
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
122122
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
123+
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
124+
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
123125
github.com/yuin/goldmark v1.7.1/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E=
124126
github.com/yuin/goldmark v1.7.4 h1:BDXOHExt+A7gwPCJgPIIq7ENvceR7we7rOS9TNoLZeg=
125127
github.com/yuin/goldmark v1.7.4/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E=

0 commit comments

Comments
 (0)