|
| 1 | +package main |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "fmt" |
| 6 | + "log" |
| 7 | + "net/url" |
| 8 | + "os" |
| 9 | + "os/signal" |
| 10 | + "path/filepath" |
| 11 | + "strings" |
| 12 | + "syscall" |
| 13 | + |
| 14 | + "github.com/mark3labs/mcp-go/client" |
| 15 | + "github.com/mark3labs/mcp-go/client/transport" |
| 16 | + "github.com/mark3labs/mcp-go/mcp" |
| 17 | +) |
| 18 | + |
| 19 | +// fileURI returns a file:// URI for both Unix and Windows absolute paths. |
| 20 | +func fileURI(p string) string { |
| 21 | + p = filepath.ToSlash(p) |
| 22 | + if !strings.HasPrefix(p, "/") { // e.g., "C:/Users/..." on Windows |
| 23 | + p = "/" + p |
| 24 | + } |
| 25 | + return (&url.URL{Scheme: "file", Path: p}).String() |
| 26 | +} |
| 27 | + |
| 28 | +// MockRootsHandler implements client.RootsHandler for demonstration. |
| 29 | +// In a real implementation, this would enumerate workspace/project roots. |
| 30 | +type MockRootsHandler struct{} |
| 31 | + |
| 32 | +// ListRoots implements client.RootsHandler by returning example workspace roots. |
| 33 | +func (h *MockRootsHandler) ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error) { |
| 34 | + home, err := os.UserHomeDir() |
| 35 | + if err != nil { |
| 36 | + log.Printf("Warning: failed to get home directory: %v", err) |
| 37 | + home = "/tmp" // fallback for demonstration |
| 38 | + } |
| 39 | + app := filepath.ToSlash(filepath.Join(home, "app")) |
| 40 | + proj := filepath.ToSlash(filepath.Join(home, "projects", "test-project")) |
| 41 | + result := &mcp.ListRootsResult{ |
| 42 | + Roots: []mcp.Root{ |
| 43 | + { |
| 44 | + Name: "app", |
| 45 | + URI: fileURI(app), |
| 46 | + }, |
| 47 | + { |
| 48 | + Name: "test-project", |
| 49 | + URI: fileURI(proj), |
| 50 | + }, |
| 51 | + }, |
| 52 | + } |
| 53 | + return result, nil |
| 54 | +} |
| 55 | + |
| 56 | +// main starts a mock MCP roots client that communicates with a subprocess over stdio. |
| 57 | +// It expects the server command as the first command-line argument, creates a stdio |
| 58 | +// transport and an MCP client with a MockRootsHandler, starts and initializes the |
| 59 | +// client, logs server info and available tools, notifies the server of root list |
| 60 | +// changes, invokes the "roots" tool and prints any text content returned, and |
| 61 | +// shuts down the client gracefully on SIGINT or SIGTERM. |
| 62 | +func main() { |
| 63 | + if len(os.Args) < 2 { |
| 64 | + log.Fatal("Usage: roots_client <server_command>") |
| 65 | + } |
| 66 | + |
| 67 | + serverCommand := os.Args[1] |
| 68 | + serverArgs := os.Args[2:] |
| 69 | + |
| 70 | + // Create stdio transport to communicate with the server |
| 71 | + stdio := transport.NewStdio(serverCommand, nil, serverArgs...) |
| 72 | + |
| 73 | + // Create roots handler |
| 74 | + rootsHandler := &MockRootsHandler{} |
| 75 | + |
| 76 | + // Create client with roots capability |
| 77 | + mcpClient := client.NewClient(stdio, client.WithRootsHandler(rootsHandler)) |
| 78 | + |
| 79 | + ctx := context.Background() |
| 80 | + |
| 81 | + // Start the client |
| 82 | + if err := mcpClient.Start(ctx); err != nil { |
| 83 | + log.Fatalf("Failed to start client: %v", err) |
| 84 | + } |
| 85 | + |
| 86 | + // Setup graceful shutdown |
| 87 | + sigChan := make(chan os.Signal, 1) |
| 88 | + signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) |
| 89 | + |
| 90 | + // Create a context that cancels on signal |
| 91 | + ctx, cancel := context.WithCancel(ctx) |
| 92 | + go func() { |
| 93 | + <-sigChan |
| 94 | + log.Println("Received shutdown signal, closing client...") |
| 95 | + cancel() |
| 96 | + }() |
| 97 | + |
| 98 | + // Move defer after error checking |
| 99 | + defer func() { |
| 100 | + if err := mcpClient.Close(); err != nil { |
| 101 | + log.Printf("Error closing client: %v", err) |
| 102 | + } |
| 103 | + }() |
| 104 | + |
| 105 | + // Initialize the connection |
| 106 | + initResult, err := mcpClient.Initialize(ctx, mcp.InitializeRequest{ |
| 107 | + Params: mcp.InitializeParams{ |
| 108 | + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, |
| 109 | + ClientInfo: mcp.Implementation{ |
| 110 | + Name: "roots-stdio-client", |
| 111 | + Version: "1.0.0", |
| 112 | + }, |
| 113 | + Capabilities: mcp.ClientCapabilities{ |
| 114 | + // Roots capability will be automatically added by WithRootsHandler |
| 115 | + }, |
| 116 | + }, |
| 117 | + }) |
| 118 | + if err != nil { |
| 119 | + log.Fatalf("Failed to initialize: %v", err) |
| 120 | + } |
| 121 | + |
| 122 | + log.Printf("Connected to server: %s v%s", initResult.ServerInfo.Name, initResult.ServerInfo.Version) |
| 123 | + log.Printf("Server capabilities: %+v", initResult.Capabilities) |
| 124 | + |
| 125 | + // list tools |
| 126 | + toolsResult, err := mcpClient.ListTools(ctx, mcp.ListToolsRequest{}) |
| 127 | + if err != nil { |
| 128 | + log.Fatalf("Failed to list tools: %v", err) |
| 129 | + } |
| 130 | + log.Printf("Available tools:") |
| 131 | + for _, tool := range toolsResult.Tools { |
| 132 | + log.Printf(" - %s: %s", tool.Name, tool.Description) |
| 133 | + } |
| 134 | + |
| 135 | + // call server tool |
| 136 | + request := mcp.CallToolRequest{} |
| 137 | + request.Params.Name = "roots" |
| 138 | + request.Params.Arguments = map[string]any{"testonly": "yes"} |
| 139 | + result, err := mcpClient.CallTool(ctx, request) |
| 140 | + if err != nil { |
| 141 | + log.Fatalf("failed to call tool roots: %v", err) |
| 142 | + } else if result.IsError { |
| 143 | + log.Printf("tool reported error") |
| 144 | + } else if len(result.Content) > 0 { |
| 145 | + resultStr := "" |
| 146 | + for _, content := range result.Content { |
| 147 | + switch tc := content.(type) { |
| 148 | + case mcp.TextContent: |
| 149 | + resultStr += fmt.Sprintf("%s\n", tc.Text) |
| 150 | + } |
| 151 | + } |
| 152 | + fmt.Printf("client call tool result: %s\n", resultStr) |
| 153 | + } |
| 154 | + |
| 155 | + // mock the root change |
| 156 | + if err := mcpClient.RootListChanges(ctx); err != nil { |
| 157 | + log.Printf("failed to notify root list change: %v", err) |
| 158 | + } |
| 159 | + |
| 160 | + // Keep running until cancelled by signal |
| 161 | + <-ctx.Done() |
| 162 | + log.Println("Client context cancelled") |
| 163 | +} |
0 commit comments