|
| 1 | +package forge |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "context" |
| 6 | + "crypto/sha256" |
| 7 | + "fmt" |
| 8 | + "log" |
| 9 | + "os" |
| 10 | + "os/exec" |
| 11 | + "path/filepath" |
| 12 | + "runtime" |
| 13 | + "strings" |
| 14 | + |
| 15 | + "github.com/mark3labs/mcp-go/mcp" |
| 16 | + "github.com/mark3labs/mcp-go/server" |
| 17 | +) |
| 18 | + |
| 19 | +// CtxAuthKey is used as a key for storing auth tokens in context |
| 20 | +type CtxAuthKey struct{} |
| 21 | + |
| 22 | +// CreateMCPServer creates and configures an MCP server with all tools registered |
| 23 | +func CreateMCPServer(appConfig *AppConfig, version string) (*server.MCPServer, error) { |
| 24 | + // Init MCP server |
| 25 | + srv := server.NewMCPServer(appConfig.Config.Name, version) |
| 26 | + |
| 27 | + // Discover & register tools |
| 28 | + if err := RegisterTools(srv, appConfig.Config, appConfig.ConfigDir, appConfig.IsDebug); err != nil { |
| 29 | + return nil, fmt.Errorf("registering tools: %w", err) |
| 30 | + } |
| 31 | + |
| 32 | + return srv, nil |
| 33 | +} |
| 34 | + |
| 35 | +// RegisterTools discovers and registers all tools from the config directory |
| 36 | +func RegisterTools(srv *server.MCPServer, cfg *ForgeConfig, configDir string, isDebug bool) error { |
| 37 | + // Discover & register tools |
| 38 | + files, err := filepath.Glob(filepath.Join(configDir, "*.yaml")) |
| 39 | + if err != nil { |
| 40 | + return fmt.Errorf("error discovering tools: %w", err) |
| 41 | + } |
| 42 | + |
| 43 | + for _, f := range files { |
| 44 | + if filepath.Base(f) == "forge.yaml" { |
| 45 | + continue |
| 46 | + } |
| 47 | + |
| 48 | + tcfg, err := LoadToolConfig(f) |
| 49 | + if err != nil { |
| 50 | + fmt.Fprintf(os.Stderr, "Warning: skipping %s: %v\n", f, err) |
| 51 | + continue |
| 52 | + } |
| 53 | + |
| 54 | + opts := []mcp.ToolOption{ |
| 55 | + mcp.WithDescription(tcfg.Description), |
| 56 | + } |
| 57 | + |
| 58 | + // Add annotations if specified |
| 59 | + if tcfg.Annotations.Title != "" { |
| 60 | + opts = append(opts, mcp.WithTitleAnnotation(tcfg.Annotations.Title)) |
| 61 | + } |
| 62 | + if tcfg.Annotations.ReadOnlyHint != nil { |
| 63 | + opts = append(opts, mcp.WithReadOnlyHintAnnotation(*tcfg.Annotations.ReadOnlyHint)) |
| 64 | + } |
| 65 | + if tcfg.Annotations.DestructiveHint != nil { |
| 66 | + opts = append(opts, mcp.WithDestructiveHintAnnotation(*tcfg.Annotations.DestructiveHint)) |
| 67 | + } |
| 68 | + if tcfg.Annotations.IdempotentHint != nil { |
| 69 | + opts = append(opts, mcp.WithIdempotentHintAnnotation(*tcfg.Annotations.IdempotentHint)) |
| 70 | + } |
| 71 | + if tcfg.Annotations.OpenWorldHint != nil { |
| 72 | + opts = append(opts, mcp.WithOpenWorldHintAnnotation(*tcfg.Annotations.OpenWorldHint)) |
| 73 | + } |
| 74 | + |
| 75 | + valid := true |
| 76 | + for _, inp := range tcfg.Inputs { |
| 77 | + pOpts := []mcp.PropertyOption{mcp.Description(inp.Description)} |
| 78 | + if inp.Required { |
| 79 | + pOpts = append(pOpts, mcp.Required()) |
| 80 | + } |
| 81 | + switch inp.Type { |
| 82 | + case "string": |
| 83 | + opts = append(opts, mcp.WithString(inp.Name, pOpts...)) |
| 84 | + case "number": |
| 85 | + opts = append(opts, mcp.WithNumber(inp.Name, pOpts...)) |
| 86 | + default: |
| 87 | + fmt.Fprintf(os.Stderr, "Warning: unsupported type %q in %s\n", inp.Type, tcfg.Name) |
| 88 | + valid = false |
| 89 | + } |
| 90 | + } |
| 91 | + if !valid { |
| 92 | + continue |
| 93 | + } |
| 94 | + |
| 95 | + tool := mcp.NewTool(tcfg.Name, opts...) |
| 96 | + srv.AddTool(tool, makeHandler(*cfg, *tcfg, isDebug)) |
| 97 | + } |
| 98 | + |
| 99 | + return nil |
| 100 | +} |
| 101 | + |
| 102 | +// makeHandler produces a ToolHandler for the given configs |
| 103 | +func makeHandler(cfg ForgeConfig, tcfg ToolConfig, isDebug bool) server.ToolHandlerFunc { |
| 104 | + return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { |
| 105 | + // 1. Gather variables |
| 106 | + vars := map[string]interface{}{} |
| 107 | + args := req.GetArguments() |
| 108 | + for _, inp := range tcfg.Inputs { |
| 109 | + val, ok := args[inp.Name] |
| 110 | + if !ok && inp.Required { |
| 111 | + return mcp.NewToolResultError(fmt.Sprintf("missing required argument: %s", inp.Name)), nil |
| 112 | + } |
| 113 | + vars[inp.Name] = val |
| 114 | + } |
| 115 | + |
| 116 | + // 2. Get the token |
| 117 | + token := "" |
| 118 | + if cfg.TokenCommand != "" { |
| 119 | + var cmd *exec.Cmd |
| 120 | + // Use the appropriate shell based on the OS |
| 121 | + if runtime.GOOS == "windows" { |
| 122 | + cmd = exec.Command("cmd", "/C", cfg.TokenCommand) |
| 123 | + } else { |
| 124 | + // Assume Unix-like shell for macOS, Linux, etc. |
| 125 | + cmd = exec.Command("sh", "-c", cfg.TokenCommand) |
| 126 | + } |
| 127 | + |
| 128 | + // Build merged environment: start with os.Environ() if passthrough, else start empty, |
| 129 | + // then overlay values from cfg.Env to ensure overrides. |
| 130 | + var envList []string |
| 131 | + if cfg.EnvPassthrough { |
| 132 | + envList = os.Environ() |
| 133 | + } else { |
| 134 | + envList = []string{} |
| 135 | + } |
| 136 | + |
| 137 | + for key, value := range cfg.Env { |
| 138 | + // Remove any existing entries for this key |
| 139 | + prefix := key + "=" |
| 140 | + filtered := envList[:0] |
| 141 | + for _, e := range envList { |
| 142 | + if !strings.HasPrefix(e, prefix) { |
| 143 | + filtered = append(filtered, e) |
| 144 | + } |
| 145 | + } |
| 146 | + envList = append(filtered, fmt.Sprintf("%s=%s", key, value)) |
| 147 | + } |
| 148 | + |
| 149 | + cmd.Env = envList |
| 150 | + |
| 151 | + if isDebug { |
| 152 | + log.Printf("Executing token command: %s", cfg.TokenCommand) |
| 153 | + if len(cmd.Env) > 0 { |
| 154 | + log.Printf("Environment variables: %v", cmd.Env) |
| 155 | + } |
| 156 | + } |
| 157 | + |
| 158 | + // Only get a token if the command is specified |
| 159 | + out, err := cmd.Output() |
| 160 | + if err != nil { |
| 161 | + // Include stderr in the error message if available |
| 162 | + errMsg := "token_command failed" |
| 163 | + if exitErr, ok := err.(*exec.ExitError); ok { |
| 164 | + // Combine exit error message and stderr for better context |
| 165 | + stderr := string(bytes.TrimSpace(exitErr.Stderr)) |
| 166 | + if stderr != "" { |
| 167 | + errMsg = fmt.Sprintf("%s: %v Stderr: %s", errMsg, exitErr, stderr) |
| 168 | + } else { |
| 169 | + errMsg = fmt.Sprintf("%s: %v", errMsg, exitErr) |
| 170 | + } |
| 171 | + } |
| 172 | + // Return nil error for MCP result error |
| 173 | + return mcp.NewToolResultErrorFromErr(errMsg, err), nil |
| 174 | + } |
| 175 | + token = "Bearer " + string(bytes.TrimSpace(out)) |
| 176 | + |
| 177 | + if isDebug { |
| 178 | + log.Printf("Obtained token (sha256): %x\n", sha256.Sum256([]byte(token))) |
| 179 | + } |
| 180 | + } else { |
| 181 | + // No token command specified, proceed with pass through token |
| 182 | + token, _ = ctx.Value(CtxAuthKey{}).(string) |
| 183 | + |
| 184 | + if isDebug { |
| 185 | + log.Printf("Pass through token (sha256): %x\n", sha256.Sum256([]byte(token))) |
| 186 | + } |
| 187 | + } |
| 188 | + |
| 189 | + // 3. Call GraphQL |
| 190 | + res, err := ExecuteGraphQL(cfg.URL, tcfg.Query, vars, token, isDebug) |
| 191 | + if err != nil { |
| 192 | + // Return error result to MCP instead of terminating |
| 193 | + return mcp.NewToolResultErrorFromErr("GraphQL execution failed", err), nil |
| 194 | + } |
| 195 | + |
| 196 | + // 4. Return raw JSON |
| 197 | + return mcp.NewToolResultText(string(res)), nil |
| 198 | + } |
| 199 | +} |
0 commit comments