cmd/mcpcurl/main.go (347 lines of code) (raw):

package main import ( "bytes" "crypto/rand" "encoding/json" "fmt" "io" "math/big" "os" "os/exec" "slices" "strings" "github.com/spf13/cobra" "github.com/spf13/viper" ) type ( // SchemaResponse represents the top-level response containing tools SchemaResponse struct { Result Result `json:"result"` JSONRPC string `json:"jsonrpc"` ID int `json:"id"` } // Result contains the list of available tools Result struct { Tools []Tool `json:"tools"` } // Tool represents a single command with its schema Tool struct { Name string `json:"name"` Description string `json:"description"` InputSchema InputSchema `json:"inputSchema"` } // InputSchema defines the structure of a tool's input parameters InputSchema struct { Type string `json:"type"` Properties map[string]Property `json:"properties"` Required []string `json:"required"` AdditionalProperties bool `json:"additionalProperties"` Schema string `json:"$schema"` } // Property defines a single parameter's type and constraints Property struct { Type string `json:"type"` Description string `json:"description"` Enum []string `json:"enum,omitempty"` Minimum *float64 `json:"minimum,omitempty"` Maximum *float64 `json:"maximum,omitempty"` Items *PropertyItem `json:"items,omitempty"` } // PropertyItem defines the type of items in an array property PropertyItem struct { Type string `json:"type"` Properties map[string]Property `json:"properties,omitempty"` Required []string `json:"required,omitempty"` AdditionalProperties bool `json:"additionalProperties,omitempty"` } // JSONRPCRequest represents a JSON-RPC 2.0 request JSONRPCRequest struct { JSONRPC string `json:"jsonrpc"` ID int `json:"id"` Method string `json:"method"` Params RequestParams `json:"params"` } // RequestParams contains the tool name and arguments RequestParams struct { Name string `json:"name"` Arguments map[string]interface{} `json:"arguments"` } // Define structure to match the response format Content struct { Type string `json:"type"` Text string `json:"text"` } ResponseResult struct { Content []Content `json:"content"` } Response struct { Result ResponseResult `json:"result"` JSONRPC string `json:"jsonrpc"` ID int `json:"id"` } ) var ( // Create root command rootCmd = &cobra.Command{ Use: "mcpcurl", Short: "CLI tool with dynamically generated commands", Long: "A CLI tool for interacting with MCP API based on dynamically loaded schemas", PersistentPreRunE: func(cmd *cobra.Command, _ []string) error { // Skip validation for help and completion commands if cmd.Name() == "help" || cmd.Name() == "completion" { return nil } // Check if the required global flag is provided serverCmd, _ := cmd.Flags().GetString("stdio-server-cmd") if serverCmd == "" { return fmt.Errorf("--stdio-server-cmd is required") } return nil }, } // Add schema command schemaCmd = &cobra.Command{ Use: "schema", Short: "Fetch schema from MCP server", Long: "Fetches the tools schema from the MCP server specified by --stdio-server-cmd", RunE: func(cmd *cobra.Command, _ []string) error { serverCmd, _ := cmd.Flags().GetString("stdio-server-cmd") if serverCmd == "" { return fmt.Errorf("--stdio-server-cmd is required") } // Build the JSON-RPC request for tools/list jsonRequest, err := buildJSONRPCRequest("tools/list", "", nil) if err != nil { return fmt.Errorf("failed to build JSON-RPC request: %w", err) } // Execute the server command and pass the JSON-RPC request response, err := executeServerCommand(serverCmd, jsonRequest) if err != nil { return fmt.Errorf("error executing server command: %w", err) } // Output the response fmt.Println(response) return nil }, } // Create the tools command toolsCmd = &cobra.Command{ Use: "tools", Short: "Access available tools", Long: "Contains all dynamically generated tool commands from the schema", } ) func main() { rootCmd.AddCommand(schemaCmd) // Add global flag for stdio server command rootCmd.PersistentFlags().String("stdio-server-cmd", "", "Shell command to invoke MCP server via stdio (required)") _ = rootCmd.MarkPersistentFlagRequired("stdio-server-cmd") // Add global flag for pretty printing rootCmd.PersistentFlags().Bool("pretty", true, "Pretty print MCP response (only for JSON or JSONL responses)") // Add the tools command to the root command rootCmd.AddCommand(toolsCmd) // Execute the root command once to parse flags _ = rootCmd.ParseFlags(os.Args[1:]) // Get pretty flag prettyPrint, err := rootCmd.Flags().GetBool("pretty") if err != nil { _, _ = fmt.Fprintf(os.Stderr, "Error getting pretty flag: %v\n", err) os.Exit(1) } // Get server command serverCmd, err := rootCmd.Flags().GetString("stdio-server-cmd") if err == nil && serverCmd != "" { // Fetch schema from server jsonRequest, err := buildJSONRPCRequest("tools/list", "", nil) if err == nil { response, err := executeServerCommand(serverCmd, jsonRequest) if err == nil { // Parse the schema response var schemaResp SchemaResponse if err := json.Unmarshal([]byte(response), &schemaResp); err == nil { // Add all the generated commands as subcommands of tools for _, tool := range schemaResp.Result.Tools { addCommandFromTool(toolsCmd, &tool, prettyPrint) } } } } } // Execute if err := rootCmd.Execute(); err != nil { _, _ = fmt.Fprintf(os.Stderr, "Error executing command: %v\n", err) os.Exit(1) } } // addCommandFromTool creates a cobra command from a tool schema func addCommandFromTool(toolsCmd *cobra.Command, tool *Tool, prettyPrint bool) { // Create command from tool cmd := &cobra.Command{ Use: tool.Name, Short: tool.Description, Run: func(cmd *cobra.Command, _ []string) { // Build a map of arguments from flags arguments, err := buildArgumentsMap(cmd, tool) if err != nil { _, _ = fmt.Fprintf(os.Stderr, "failed to build arguments map: %v\n", err) return } jsonData, err := buildJSONRPCRequest("tools/call", tool.Name, arguments) if err != nil { _, _ = fmt.Fprintf(os.Stderr, "failed to build JSONRPC request: %v\n", err) return } // Execute the server command serverCmd, err := cmd.Flags().GetString("stdio-server-cmd") if err != nil { _, _ = fmt.Fprintf(os.Stderr, "failed to get stdio-server-cmd: %v\n", err) return } response, err := executeServerCommand(serverCmd, jsonData) if err != nil { _, _ = fmt.Fprintf(os.Stderr, "error executing server command: %v\n", err) return } if err := printResponse(response, prettyPrint); err != nil { _, _ = fmt.Fprintf(os.Stderr, "error printing response: %v\n", err) return } }, } // Initialize viper for this command viperInit := func() { viper.Reset() viper.AutomaticEnv() viper.SetEnvPrefix(strings.ToUpper(tool.Name)) viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_")) } // We'll call the init function directly instead of with cobra.OnInitialize // to avoid conflicts between commands viperInit() // Add flags based on schema properties for name, prop := range tool.InputSchema.Properties { isRequired := slices.Contains(tool.InputSchema.Required, name) // Enhance description to indicate if parameter is optional description := prop.Description if !isRequired { description += " (optional)" } switch prop.Type { case "string": cmd.Flags().String(name, "", description) if len(prop.Enum) > 0 { // Add validation in PreRun for enum values cmd.PreRunE = func(cmd *cobra.Command, _ []string) error { for flagName, property := range tool.InputSchema.Properties { if len(property.Enum) > 0 { value, _ := cmd.Flags().GetString(flagName) if value != "" && !slices.Contains(property.Enum, value) { return fmt.Errorf("%s must be one of: %s", flagName, strings.Join(property.Enum, ", ")) } } } return nil } } case "number": cmd.Flags().Float64(name, 0, description) case "boolean": cmd.Flags().Bool(name, false, description) case "array": if prop.Items != nil { if prop.Items.Type == "string" { cmd.Flags().StringSlice(name, []string{}, description) } else if prop.Items.Type == "object" { // For complex objects in arrays, we'll use a JSON string that users can provide cmd.Flags().String(name+"-json", "", description+" (provide as JSON array)") } } } if isRequired { _ = cmd.MarkFlagRequired(name) } // Bind flag to viper _ = viper.BindPFlag(name, cmd.Flags().Lookup(name)) } // Add command to root toolsCmd.AddCommand(cmd) } // buildArgumentsMap extracts flag values into a map of arguments func buildArgumentsMap(cmd *cobra.Command, tool *Tool) (map[string]interface{}, error) { arguments := make(map[string]interface{}) for name, prop := range tool.InputSchema.Properties { switch prop.Type { case "string": if value, _ := cmd.Flags().GetString(name); value != "" { arguments[name] = value } case "number": if value, _ := cmd.Flags().GetFloat64(name); value != 0 { arguments[name] = value } case "boolean": // For boolean, we need to check if it was explicitly set if cmd.Flags().Changed(name) { value, _ := cmd.Flags().GetBool(name) arguments[name] = value } case "array": if prop.Items != nil { if prop.Items.Type == "string" { if values, _ := cmd.Flags().GetStringSlice(name); len(values) > 0 { arguments[name] = values } } else if prop.Items.Type == "object" { if jsonStr, _ := cmd.Flags().GetString(name + "-json"); jsonStr != "" { var jsonArray []interface{} if err := json.Unmarshal([]byte(jsonStr), &jsonArray); err != nil { return nil, fmt.Errorf("error parsing JSON for %s: %w", name, err) } arguments[name] = jsonArray } } } } } return arguments, nil } // buildJSONRPCRequest creates a JSON-RPC request with the given tool name and arguments func buildJSONRPCRequest(method, toolName string, arguments map[string]interface{}) (string, error) { id, err := rand.Int(rand.Reader, big.NewInt(10000)) if err != nil { return "", fmt.Errorf("failed to generate random ID: %w", err) } request := JSONRPCRequest{ JSONRPC: "2.0", ID: int(id.Int64()), // Random ID between 0 and 9999 Method: method, Params: RequestParams{ Name: toolName, Arguments: arguments, }, } jsonData, err := json.Marshal(request) if err != nil { return "", fmt.Errorf("failed to marshal JSON request: %w", err) } return string(jsonData), nil } // executeServerCommand runs the specified command, sends the JSON request to stdin, // and returns the response from stdout func executeServerCommand(cmdStr, jsonRequest string) (string, error) { // Split the command string into command and arguments cmdParts := strings.Fields(cmdStr) if len(cmdParts) == 0 { return "", fmt.Errorf("empty command") } cmd := exec.Command(cmdParts[0], cmdParts[1:]...) //nolint:gosec //mcpcurl is a test command that needs to execute arbitrary shell commands // Setup stdin pipe stdin, err := cmd.StdinPipe() if err != nil { return "", fmt.Errorf("failed to create stdin pipe: %w", err) } // Setup stdout and stderr pipes var stdout, stderr bytes.Buffer cmd.Stdout = &stdout cmd.Stderr = &stderr // Start the command if err := cmd.Start(); err != nil { return "", fmt.Errorf("failed to start command: %w", err) } // Write the JSON request to stdin if _, err := io.WriteString(stdin, jsonRequest+"\n"); err != nil { return "", fmt.Errorf("failed to write to stdin: %w", err) } _ = stdin.Close() // Wait for the command to complete if err := cmd.Wait(); err != nil { return "", fmt.Errorf("command failed: %w, stderr: %s", err, stderr.String()) } return stdout.String(), nil } func printResponse(response string, prettyPrint bool) error { if !prettyPrint { fmt.Println(response) return nil } // Parse the JSON response var resp Response if err := json.Unmarshal([]byte(response), &resp); err != nil { return fmt.Errorf("failed to parse JSON: %w", err) } // Extract text from content items of type "text" for _, content := range resp.Result.Content { if content.Type == "text" { var textContentObj map[string]interface{} err := json.Unmarshal([]byte(content.Text), &textContentObj) if err == nil { prettyText, err := json.MarshalIndent(textContentObj, "", " ") if err != nil { return fmt.Errorf("failed to pretty print text content: %w", err) } fmt.Println(string(prettyText)) continue } // Fallback parsing as JSONL var textContentList []map[string]interface{} if err := json.Unmarshal([]byte(content.Text), &textContentList); err != nil { return fmt.Errorf("failed to parse text content as a list: %w", err) } prettyText, err := json.MarshalIndent(textContentList, "", " ") if err != nil { return fmt.Errorf("failed to pretty print array content: %w", err) } fmt.Println(string(prettyText)) } } // If no text content found, print the original response if len(resp.Result.Content) == 0 { fmt.Println(response) } return nil }