@@ -6,8 +6,11 @@ import (
66 "encoding/json"
77 "errors"
88 "fmt"
9+ "io"
910 "log/slog"
1011 "maps"
12+ "net/http"
13+ "os/exec"
1114 "strings"
1215 "sync"
1316 "time"
@@ -19,9 +22,7 @@ import (
1922 "github.com/charmbracelet/crush/internal/permission"
2023 "github.com/charmbracelet/crush/internal/pubsub"
2124 "github.com/charmbracelet/crush/internal/version"
22− "github.com/mark3labs/mcp-go/client"
23− "github.com/mark3labs/mcp-go/client/transport"
24− "github.com/mark3labs/mcp-go/mcp"
25+ "github.com/modelcontextprotocol/go-sdk/mcp"
2526)
2627
2728// MCPState represents the current state of an MCP client
@@ -71,7 +72,7 @@ type MCPClientInfo struct {
7172 Name string
7273 State MCPState
7374 Error error
74− Client *client.Client
75+ Client *mcp.ClientSession
7576 ToolCount int
7677 ConnectedAt time.Time
7778}
@@ -80,14 +81,14 @@ var (
8081 mcpToolsOnce sync.Once
8182 mcpTools = csync.NewMap[string, tools.BaseTool]()
8283 mcpClient2Tools = csync.NewMap[string, []tools.BaseTool]()
83− mcpClients = csync.NewMap[string, *client.Client]()
84+ mcpClients = csync.NewMap[string, *mcp.ClientSession]()
8485 mcpStates = csync.NewMap[string, MCPClientInfo]()
8586 mcpBroker = pubsub.NewBroker[MCPEvent]()
8687)
8788
8889type McpTool struct {
8990 mcpName string
90− tool mcp.Tool
91+ tool *mcp.Tool
9192 permissions permission.Service
9293 workingDir string
9394}
@@ -97,14 +98,9 @@ func (b *McpTool) Name() string {
9798}
9899
99100func (b *McpTool) Info() tools.ToolInfo {
100− required := b.tool.InputSchema.Required
101− if required == nil {
102− required = make([]string, 0)
103− }
104− parameters := b.tool.InputSchema.Properties
105− if parameters == nil {
106− parameters = make(map[string]any)
107− }
101+ input := b.tool.InputSchema.(map[string]any)
102+ required, _ := input["required"].([]string)
103+ parameters, _ := input["properties"].(map[string]any)
108104 return tools.ToolInfo{
109105 Name: fmt.Sprintf("mcp_%s_%s", b.mcpName, b.tool.Name),
110106 Description: b.tool.Description,
@@ -123,11 +119,9 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
123119 if err != nil {
124120 return tools.NewTextErrorResponse(err.Error()), nil
125121 }
126− result, err := c.CallTool(ctx, mcp.CallToolRequest{
127− Params: mcp.CallToolParams{
128− Name: toolName,
129− Arguments: args,
130− },
122+ result, err := c.CallTool(ctx, &mcp.CallToolParams{
123+ Name: toolName,
124+ Arguments: args,
131125 })
132126 if err != nil {
133127 return tools.NewTextErrorResponse(err.Error()), nil
@@ -135,8 +129,8 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
135129
136130 output := make([]string, 0, len(result.Content))
137131 for _, v := range result.Content {
138− if v, ok := v.(mcp.TextContent); ok {
139− output = append(output, v.Text)
132+ if vv, ok := v.(*mcp.TextContent); ok {
133+ output = append(output, vv.Text)
140134 } else {
141135 output = append(output, fmt.Sprintf("%v", v))
142136 }
@@ -144,8 +138,8 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
144138 return tools.NewTextResponse(strings.Join(output, "\n")), nil
145139}
146140
147−func getOrRenewClient(ctx context.Context, name string) (*client.Client, error) {
148− c, ok := mcpClients.Get(name)
141+func getOrRenewClient(ctx context.Context, name string) (*mcp.ClientSession, error) {
142+ sess, ok := mcpClients.Get(name)
149143 if !ok {
150144 return nil, fmt.Errorf("mcp '%s' not available", name)
151145 }
@@ -157,20 +151,20 @@ func getOrRenewClient(ctx context.Context, name string) (*client.Client, error)
157151 timeout := mcpTimeout(m)
158152 pingCtx, cancel := context.WithTimeout(ctx, timeout)
159153 defer cancel()
160− err := c.Ping(pingCtx)
154+ err := sess.Ping(pingCtx, nil)
161155 if err == nil {
162− return c, nil
156+ return sess, nil
163157 }
164158 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, state.ToolCount)
165159
166− c, err = createAndInitializeClient(ctx, name, m, cfg.Resolver())
160+ sess, err = createMCPSession(ctx, name, m, cfg.Resolver())
167161 if err != nil {
168162 return nil, err
169163 }
170164
171− updateMCPState(name, MCPStateConnected, nil, c, state.ToolCount)
172− mcpClients.Set(name, c)
173− return c, nil
165+ updateMCPState(name, MCPStateConnected, nil, sess, state.ToolCount)
166+ mcpClients.Set(name, sess)
167+ return sess, nil
174168}
175169
176170func (b *McpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolResponse, error) {
@@ -197,8 +191,8 @@ func (b *McpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolRes
197191 return runTool(ctx, b.mcpName, b.tool.Name, params.Input)
198192}
199193
200−func getTools(ctx context.Context, name string, permissions permission.Service, c *client.Client, workingDir string) ([]tools.BaseTool, error) {
201− result, err := c.ListTools(ctx, mcp.ListToolsRequest{})
194+func getTools(ctx context.Context, name string, permissions permission.Service, c *mcp.ClientSession, workingDir string) ([]tools.BaseTool, error) {
195+ result, err := c.ListTools(ctx, &mcp.ListToolsParams{})
202196 if err != nil {
203197 return nil, err
204198 }
@@ -230,7 +224,7 @@ func GetMCPState(name string) (MCPClientInfo, bool) {
230224}
231225
232226// updateMCPState updates the state of an MCP client and publishes an event
233−func updateMCPState(name string, state MCPState, err error, client *client.Client, toolCount int) {
227+func updateMCPState(name string, state MCPState, err error, client *mcp.ClientSession, toolCount int) {
234228 info := MCPClientInfo{
235229 Name: name,
236230 State: state,
@@ -257,19 +251,14 @@ func updateMCPState(name string, state MCPState, err error, client *client.Clien
257251 })
258252}
259253
260−// publishMCPEventToolsListChanged publishes a tool list changed event
261−func publishMCPEventToolsListChanged(name string) {
262− mcpBroker.Publish(pubsub.UpdatedEvent, MCPEvent{
263− Type: MCPEventToolsListChanged,
264− Name: name,
265− })
266−}
267−
268254// CloseMCPClients closes all MCP clients. This should be called during application shutdown.
269255func CloseMCPClients() error {
270256 var errs []error
271257 for name, c := range mcpClients.Seq2() {
272− if err := c.Close(); err != nil {
258+ if err := c.Close(); err != nil &&
259+ !errors.Is(err, io.EOF) &&
260+ !errors.Is(err, context.Canceled) &&
261+ err.Error() != "signal: killed" {
273262 errs = append(errs, fmt.Errorf("close mcp: %s: %w", name, err))
274263 }
275264 }
@@ -277,16 +266,6 @@ func CloseMCPClients() error {
277266 return errors.Join(errs...)
278267}
279268
280−var mcpInitRequest = mcp.InitializeRequest{
281− Params: mcp.InitializeParams{
282− ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
283− ClientInfo: mcp.Implementation{
284− Name: "Crush",
285− Version: version.Version,
286− },
287− },
288−}
289−
290269func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) {
291270 var wg sync.WaitGroup
292271 // Initialize states for all configured MCPs
@@ -322,7 +301,7 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con
322301 ctx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
323302 defer cancel()
324303
325− c, err := createAndInitializeClient(ctx, name, m, cfg.Resolver())
304+ c, err := createMCPSession(ctx, name, m, cfg.Resolver())
326305 if err != nil {
327306 return
328307 }
@@ -359,49 +338,46 @@ func updateMcpTools(mcpName string, tools []tools.BaseTool) {
359338 }
360339}
361340
362−func createAndInitializeClient(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*client.Client, error) {
363− c, err := createMcpClient(name, m, resolver)
341+func createMCPSession(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*mcp.ClientSession, error) {
342+ timeout := mcpTimeout(m)
343+ mcpCtx, cancel := context.WithCancel(ctx)
344+ cancelTimer := time.AfterFunc(timeout, cancel)
345+
346+ transport, err := createMCPTransport(mcpCtx, m, resolver)
364347 if err != nil {
365348 updateMCPState(name, MCPStateError, err, nil, 0)
366349 slog.Error("error creating mcp client", "error", err, "name", name)
367350 return nil, err
368351 }
369352
370− c.OnNotification(func(n mcp.JSONRPCNotification) {
371− slog.Debug("Received MCP notification", "name", name, "notification", n)
372− switch n.Method {
373− case "notifications/tools/list_changed":
374− publishMCPEventToolsListChanged(name)
375− default:
376− slog.Debug("Unhandled MCP notification", "name", name, "method", n.Method)
377− }
378− })
379−
380− // XXX: ideally we should be able to use context.WithTimeout here, but,
381− // the SSE MCP client will start failing once that context is canceled.
382− timeout := mcpTimeout(m)
383− mcpCtx, cancel := context.WithCancel(ctx)
384− cancelTimer := time.AfterFunc(timeout, cancel)
353+ client := mcp.NewClient(
354+ &mcp.Implementation{
355+ Name: "crush",
356+ Version: version.Version,
357+ Title: "Crush",
358+ },
359+ &mcp.ClientOptions{
360+ ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) {
361+ mcpBroker.Publish(pubsub.UpdatedEvent, MCPEvent{
362+ Type: MCPEventToolsListChanged,
363+ Name: name,
364+ })
365+ },
366+ KeepAlive: time.Minute * 10,
367+ },
368+ )
385369
386− if err := c.Start(mcpCtx); err != nil {
370+ session, err := client.Connect(mcpCtx, transport, nil)
371+ if err != nil {
387372 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
388373 slog.Error("error starting mcp client", "error", err, "name", name)
389− _ = c.Close()
390− cancel()
391− return nil, err
392− }
393−
394− if _, err := c.Initialize(mcpCtx, mcpInitRequest); err != nil {
395− updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
396− slog.Error("error initializing mcp client", "error", err, "name", name)
397− _ = c.Close()
398374 cancel()
399375 return nil, err
400376 }
401377
402378 cancelTimer.Stop()
403379 slog.Info("Initialized mcp client", "name", name)
404− return c, nil
380+ return session, nil
405381}
406382
407383func maybeTimeoutErr(err error, timeout time.Duration) error {
@@ -411,7 +387,7 @@ func maybeTimeoutErr(err error, timeout time.Duration) error {
411387 return err
412388}
413389
414−func createMcpClient(name string, m config.MCPConfig, resolver config.VariableResolver) (*client.Client, error) {
390+func createMCPTransport(ctx context.Context, m config.MCPConfig, resolver config.VariableResolver) (mcp.Transport, error) {
415391 switch m.Type {
416392 case config.MCPStdio:
417393 command, err := resolver.ResolveValue(m.Command)
@@ -421,44 +397,51 @@ func createMcpClient(name string, m config.MCPConfig, resolver config.VariableRe
421397 if strings.TrimSpace(command) == "" {
422398 return nil, fmt.Errorf("mcp stdio config requires a non-empty 'command' field")
423399 }
424− return client.NewStdioMCPClientWithOptions(
425− home.Long(command),
426− m.ResolvedEnv(),
427− m.Args,
428− transport.WithCommandLogger(mcpLogger{name: name}),
429− )
400+ cmd := exec.CommandContext(ctx, home.Long(command), m.Args...)
401+ cmd.Env = m.ResolvedEnv()
402+ return &mcp.CommandTransport{
403+ Command: cmd,
404+ }, nil
430405 case config.MCPHttp:
431406 if strings.TrimSpace(m.URL) == "" {
432407 return nil, fmt.Errorf("mcp http config requires a non-empty 'url' field")
433408 }
434− return client.NewStreamableHttpClient(
435− m.URL,
436− transport.WithHTTPHeaders(m.ResolvedHeaders()),
437− transport.WithHTTPLogger(mcpLogger{name: name}),
438− )
439− case config.MCPSse:
409+ client := &http.Client{
410+ Transport: &headerRoundTripper{
411+ headers: m.ResolvedHeaders(),
412+ },
413+ }
414+ return &mcp.StreamableClientTransport{
415+ Endpoint: m.URL,
416+ HTTPClient: client,
417+ }, nil
418+ case config.MCPSSE:
440419 if strings.TrimSpace(m.URL) == "" {
441420 return nil, fmt.Errorf("mcp sse config requires a non-empty 'url' field")
442421 }
443− return client.NewSSEMCPClient(
444− m.URL,
445− client.WithHeaders(m.ResolvedHeaders()),
446− transport.WithSSELogger(mcpLogger{name: name}),
447− )
422+ client := &http.Client{
423+ Transport: &headerRoundTripper{
424+ headers: m.ResolvedHeaders(),
425+ },
426+ }
427+ return &mcp.SSEClientTransport{
428+ Endpoint: m.URL,
429+ HTTPClient: client,
430+ }, nil
448431 default:
449432 return nil, fmt.Errorf("unsupported mcp type: %s", m.Type)
450433 }
451434}
452435
453−// for MCP's clients.
454−type mcpLogger struct{ name string }
455−
456−func (l mcpLogger) Errorf(format string, v ...any) {
457− slog.Error(fmt.Sprintf(format, v...), "name", l.name)
436+type headerRoundTripper struct {
437+ headers map[string]string
458438}
459439
460−func (l mcpLogger) Infof(format string, v ...any) {
461− slog.Info(fmt.Sprintf(format, v...), "name", l.name)
440+func (rt headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
441+ for k, v := range rt.headers {
442+ req.Header.Set(k, v)
443+ }
444+ return http.DefaultTransport.RoundTrip(req)
462445}
463446
464447func mcpTimeout(m config.MCPConfig) time.Duration {