@@ -6,8 +6,11 @@ import (
6 "encoding/json"
6 "encoding/json"
9 "log/slog"
10 "log/slog"
@@ -19,9 +22,7 @@ import (
19 "github.com/charmbracelet/crush/internal/permission"
22 "github.com/charmbracelet/crush/internal/permission"
20 "github.com/charmbracelet/crush/internal/pubsub"
23 "github.com/charmbracelet/crush/internal/pubsub"
21 "github.com/charmbracelet/crush/internal/version"
24 "github.com/charmbracelet/crush/internal/version"
22− "github.com/mark3labs/mcp-go/client"
25+ "github.com/modelcontextprotocol/go-sdk/mcp"
23− "github.com/mark3labs/mcp-go/client/transport"
24− "github.com/mark3labs/mcp-go/mcp"
27// MCPState represents the current state of an MCP client
28// MCPState represents the current state of an MCP client
@@ -71,7 +72,7 @@ type MCPClientInfo struct {
71 Name string
72 Name string
72 State MCPState
73 State MCPState
73 Error error
74 Error error
74− Client *client.Client
75+ Client *mcp.ClientSession
75 ToolCount int
76 ToolCount int
76 ConnectedAt time.Time
77 ConnectedAt time.Time
@@ -80,14 +81,14 @@ var (
80 mcpToolsOnce sync.Once
81 mcpToolsOnce sync.Once
81 mcpTools = csync.NewMap[string, tools.BaseTool]()
82 mcpTools = csync.NewMap[string, tools.BaseTool]()
82 mcpClient2Tools = csync.NewMap[string, []tools.BaseTool]()
83 mcpClient2Tools = csync.NewMap[string, []tools.BaseTool]()
83− mcpClients = csync.NewMap[string, *client.Client]()
84+ mcpClients = csync.NewMap[string, *mcp.ClientSession]()
84 mcpStates = csync.NewMap[string, MCPClientInfo]()
85 mcpStates = csync.NewMap[string, MCPClientInfo]()
85 mcpBroker = pubsub.NewBroker[MCPEvent]()
86 mcpBroker = pubsub.NewBroker[MCPEvent]()
88type McpTool struct {
89type McpTool struct {
89 mcpName string
90 mcpName string
90− tool mcp.Tool
91+ tool *mcp.Tool
91 permissions permission.Service
92 permissions permission.Service
92 workingDir string
93 workingDir string
@@ -97,14 +98,9 @@ func (b *McpTool) Name() string {
99func (b *McpTool) Info() tools.ToolInfo {
100func (b *McpTool) Info() tools.ToolInfo {
100− required := b.tool.InputSchema.Required
101+ input := b.tool.InputSchema.(map[string]any)
101− if required == nil {
102− required = make([]string, 0)
102+ required, _ := input["required"].([]string)
103− }
103+ parameters, _ := input["properties"].(map[string]any)
104− parameters := b.tool.InputSchema.Properties
105− if parameters == nil {
106− parameters = make(map[string]any)
108 return tools.ToolInfo{
104 return tools.ToolInfo{
109 Name: fmt.Sprintf("mcp_%s_%s", b.mcpName, b.tool.Name),
105 Name: fmt.Sprintf("mcp_%s_%s", b.mcpName, b.tool.Name),
110 Description: b.tool.Description,
106 Description: b.tool.Description,
@@ -123,11 +119,9 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
123 if err != nil {
119 if err != nil {
124 return tools.NewTextErrorResponse(err.Error()), nil
120 return tools.NewTextErrorResponse(err.Error()), nil
126− result, err := c.CallTool(ctx, mcp.CallToolRequest{
127− Params: mcp.CallToolParams{
122+ result, err := c.CallTool(ctx, &mcp.CallToolParams{
128− Name: toolName,
123+ Name: toolName,
129− Arguments: args,
124+ Arguments: args,
132 if err != nil {
126 if err != nil {
133 return tools.NewTextErrorResponse(err.Error()), nil
127 return tools.NewTextErrorResponse(err.Error()), nil
@@ -135,8 +129,8 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
136 output := make([]string, 0, len(result.Content))
130 output := make([]string, 0, len(result.Content))
137 for _, v := range result.Content {
131 for _, v := range result.Content {
138− if v, ok := v.(mcp.TextContent); ok {
132+ if vv, ok := v.(*mcp.TextContent); ok {
139− output = append(output, v.Text)
133+ output = append(output, vv.Text)
141 output = append(output, fmt.Sprintf("%v", v))
135 output = append(output, fmt.Sprintf("%v", v))
@@ -144,8 +138,8 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
144 return tools.NewTextResponse(strings.Join(output, "\n")), nil
138 return tools.NewTextResponse(strings.Join(output, "\n")), nil
147−func getOrRenewClient(ctx context.Context, name string) (*client.Client, error) {
141+func getOrRenewClient(ctx context.Context, name string) (*mcp.ClientSession, error) {
148− c, ok := mcpClients.Get(name)
142+ sess, ok := mcpClients.Get(name)
150 return nil, fmt.Errorf("mcp '%s' not available", name)
144 return nil, fmt.Errorf("mcp '%s' not available", name)
@@ -157,20 +151,20 @@ func getOrRenewClient(ctx context.Context, name string) (*client.Client, error)
157 timeout := mcpTimeout(m)
151 timeout := mcpTimeout(m)
158 pingCtx, cancel := context.WithTimeout(ctx, timeout)
152 pingCtx, cancel := context.WithTimeout(ctx, timeout)
159 defer cancel()
153 defer cancel()
160− err := c.Ping(pingCtx)
154+ err := sess.Ping(pingCtx, nil)
161 if err == nil {
155 if err == nil {
162− return c, nil
156+ return sess, nil
164 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, state.ToolCount)
158 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, state.ToolCount)
166− c, err = createAndInitializeClient(ctx, name, m, cfg.Resolver())
160+ sess, err = createMCPSession(ctx, name, m, cfg.Resolver())
167 if err != nil {
161 if err != nil {
168 return nil, err
162 return nil, err
171− updateMCPState(name, MCPStateConnected, nil, c, state.ToolCount)
165+ updateMCPState(name, MCPStateConnected, nil, sess, state.ToolCount)
172− mcpClients.Set(name, c)
166+ mcpClients.Set(name, sess)
173− return c, nil
167+ return sess, nil
176func (b *McpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolResponse, error) {
170func (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
197 return runTool(ctx, b.mcpName, b.tool.Name, params.Input)
191 return runTool(ctx, b.mcpName, b.tool.Name, params.Input)
200−func getTools(ctx context.Context, name string, permissions permission.Service, c *client.Client, workingDir string) ([]tools.BaseTool, error) {
194+func getTools(ctx context.Context, name string, permissions permission.Service, c *mcp.ClientSession, workingDir string) ([]tools.BaseTool, error) {
201− result, err := c.ListTools(ctx, mcp.ListToolsRequest{})
195+ result, err := c.ListTools(ctx, &mcp.ListToolsParams{})
202 if err != nil {
196 if err != nil {
203 return nil, err
197 return nil, err
@@ -230,7 +224,7 @@ func GetMCPState(name string) (MCPClientInfo, bool) {
232// updateMCPState updates the state of an MCP client and publishes an event
226// 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) {
234 info := MCPClientInfo{
228 info := MCPClientInfo{
235 Name: name,
229 Name: name,
236 State: state,
230 State: state,
@@ -257,19 +251,14 @@ func updateMCPState(name string, state MCPState, err error, client *client.Clien
260−// publishMCPEventToolsListChanged publishes a tool list changed event
261−func publishMCPEventToolsListChanged(name string) {
262− mcpBroker.Publish(pubsub.UpdatedEvent, MCPEvent{
263− Type: MCPEventToolsListChanged,
268// CloseMCPClients closes all MCP clients. This should be called during application shutdown.
254// CloseMCPClients closes all MCP clients. This should be called during application shutdown.
269func CloseMCPClients() error {
255func CloseMCPClients() error {
270 var errs []error
256 var errs []error
271 for name, c := range mcpClients.Seq2() {
257 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" {
273 errs = append(errs, fmt.Errorf("close mcp: %s: %w", name, err))
262 errs = append(errs, fmt.Errorf("close mcp: %s: %w", name, err))
@@ -277,16 +266,6 @@ func CloseMCPClients() error {
277 return errors.Join(errs...)
266 return errors.Join(errs...)
280−var mcpInitRequest = mcp.InitializeRequest{
281− Params: mcp.InitializeParams{
282− ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
283− ClientInfo: mcp.Implementation{
285− Version: version.Version,
290func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) {
269func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) {
291 var wg sync.WaitGroup
270 var wg sync.WaitGroup
292 // Initialize states for all configured MCPs
271 // Initialize states for all configured MCPs
@@ -322,7 +301,7 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con
322 ctx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
301 ctx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
323 defer cancel()
302 defer cancel()
325− c, err := createAndInitializeClient(ctx, name, m, cfg.Resolver())
304+ c, err := createMCPSession(ctx, name, m, cfg.Resolver())
326 if err != nil {
305 if err != nil {
@@ -359,49 +338,46 @@ func updateMcpTools(mcpName string, tools []tools.BaseTool) {
362−func createAndInitializeClient(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*client.Client, error) {
341+func createMCPSession(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*mcp.ClientSession, error) {
363− c, err := createMcpClient(name, m, resolver)
342+ timeout := mcpTimeout(m)
343+ mcpCtx, cancel := context.WithCancel(ctx)
344+ cancelTimer := time.AfterFunc(timeout, cancel)
346+ transport, err := createMCPTransport(mcpCtx, m, resolver)
364 if err != nil {
347 if err != nil {
365 updateMCPState(name, MCPStateError, err, nil, 0)
348 updateMCPState(name, MCPStateError, err, nil, 0)
366 slog.Error("error creating mcp client", "error", err, "name", name)
349 slog.Error("error creating mcp client", "error", err, "name", name)
367 return nil, err
350 return nil, err
370− c.OnNotification(func(n mcp.JSONRPCNotification) {
353+ client := mcp.NewClient(
371− slog.Debug("Received MCP notification", "name", name, "notification", n)
354+ &mcp.Implementation{
372− switch n.Method {
355+ Name: "crush",
373− case "notifications/tools/list_changed":
356+ Version: version.Version,
374− publishMCPEventToolsListChanged(name)
357+ Title: "Crush",
376− slog.Debug("Unhandled MCP notification", "name", name, "method", n.Method)
359+ &mcp.ClientOptions{
377− }
360+ ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) {
361+ mcpBroker.Publish(pubsub.UpdatedEvent, MCPEvent{
362+ Type: MCPEventToolsListChanged,
380− // XXX: ideally we should be able to use context.WithTimeout here, but,
366+ KeepAlive: time.Minute * 10,
381− // the SSE MCP client will start failing once that context is canceled.
367+ },
382− timeout := mcpTimeout(m)
368+ )
383− mcpCtx, cancel := context.WithCancel(ctx)
384− cancelTimer := time.AfterFunc(timeout, cancel)
386− if err := c.Start(mcpCtx); err != nil {
370+ session, err := client.Connect(mcpCtx, transport, nil)
387 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
372 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
388 slog.Error("error starting mcp client", "error", err, "name", name)
373 slog.Error("error starting mcp client", "error", err, "name", name)
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)
399 return nil, err
375 return nil, err
402 cancelTimer.Stop()
378 cancelTimer.Stop()
403 slog.Info("Initialized mcp client", "name", name)
379 slog.Info("Initialized mcp client", "name", name)
404− return c, nil
380+ return session, nil
407func maybeTimeoutErr(err error, timeout time.Duration) error {
383func maybeTimeoutErr(err error, timeout time.Duration) error {
@@ -411,7 +387,7 @@ func maybeTimeoutErr(err error, timeout time.Duration) error {
411 return err
387 return err
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) {
415 switch m.Type {
391 switch m.Type {
416 case config.MCPStdio:
392 case config.MCPStdio:
417 command, err := resolver.ResolveValue(m.Command)
393 command, err := resolver.ResolveValue(m.Command)
@@ -421,44 +397,51 @@ func createMcpClient(name string, m config.MCPConfig, resolver config.VariableRe
421 if strings.TrimSpace(command) == "" {
397 if strings.TrimSpace(command) == "" {
422 return nil, fmt.Errorf("mcp stdio config requires a non-empty 'command' field")
398 return nil, fmt.Errorf("mcp stdio config requires a non-empty 'command' field")
424− return client.NewStdioMCPClientWithOptions(
425− home.Long(command),
400+ cmd := exec.CommandContext(ctx, home.Long(command), m.Args...)
426− m.ResolvedEnv(),
401+ cmd.Env = m.ResolvedEnv()
427− m.Args,
402+ return &mcp.CommandTransport{
428− transport.WithCommandLogger(mcpLogger{name: name}),
403+ Command: cmd,
430 case config.MCPHttp:
405 case config.MCPHttp:
431 if strings.TrimSpace(m.URL) == "" {
406 if strings.TrimSpace(m.URL) == "" {
432 return nil, fmt.Errorf("mcp http config requires a non-empty 'url' field")
407 return nil, fmt.Errorf("mcp http config requires a non-empty 'url' field")
434− return client.NewStreamableHttpClient(
409+ client := &http.Client{
410+ Transport: &headerRoundTripper{
411+ headers: m.ResolvedHeaders(),
414+ return &mcp.StreamableClientTransport{
435− m.URL,
415+ Endpoint: m.URL,
436− transport.WithHTTPHeaders(m.ResolvedHeaders()),
416+ HTTPClient: client,
437− transport.WithHTTPLogger(mcpLogger{name: name}),
417+ }, nil
439− case config.MCPSse:
418+ case config.MCPSSE:
440 if strings.TrimSpace(m.URL) == "" {
419 if strings.TrimSpace(m.URL) == "" {
441 return nil, fmt.Errorf("mcp sse config requires a non-empty 'url' field")
420 return nil, fmt.Errorf("mcp sse config requires a non-empty 'url' field")
443− return client.NewSSEMCPClient(
422+ client := &http.Client{
423+ Transport: &headerRoundTripper{
424+ headers: m.ResolvedHeaders(),
427+ return &mcp.SSEClientTransport{
444− m.URL,
428+ Endpoint: m.URL,
445− client.WithHeaders(m.ResolvedHeaders()),
429+ HTTPClient: client,
446− transport.WithSSELogger(mcpLogger{name: name}),
430+ }, nil
449 return nil, fmt.Errorf("unsupported mcp type: %s", m.Type)
432 return nil, fmt.Errorf("unsupported mcp type: %s", m.Type)
453−// for MCP's clients.
436+type headerRoundTripper struct {
454−type mcpLogger struct{ name string }
437+ headers map[string]string
456−func (l mcpLogger) Errorf(format string, v ...any) {
457− slog.Error(fmt.Sprintf(format, v...), "name", l.name)
460−func (l mcpLogger) Infof(format string, v ...any) {
440+func (rt headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
461− slog.Info(fmt.Sprintf(format, v...), "name", l.name)
441+ for k, v := range rt.headers {
442+ req.Header.Set(k, v)
444+ return http.DefaultTransport.RoundTrip(req)
464func mcpTimeout(m config.MCPConfig) time.Duration {
447func mcpTimeout(m config.MCPConfig) time.Duration {