From 00aa3c6995cb190a6cea0fb91f64bb3c1198ad2c Mon Sep 17 00:00:00 2001 From: Ed Zynda Date: Fri, 27 Jun 2025 16:30:18 +0300 Subject: [PATCH] Env substitution (#93) * compact mode * tweaks * tweaks * tweaks * fix streaming in tool calls * impl env sub * fix --- AGENTS.md | 85 ++++- README.md | 151 ++++++-- cmd/root.go | 126 ++++++- cmd/script.go | 59 ++- cmd/script_integration_test.go | 284 ++++++++++++++ examples/scripts/env-substitution-script.sh | 42 +++ internal/agent/agent.go | 5 + internal/config/config.go | 12 +- internal/config/integration_test.go | 227 +++++++++++ internal/config/substitution.go | 108 ++++++ internal/config/substitution_test.go | 395 ++++++++++++++++++++ internal/tools/mcp.go | 9 + 12 files changed, 1407 insertions(+), 96 deletions(-) create mode 100644 cmd/script_integration_test.go create mode 100644 examples/scripts/env-substitution-script.sh create mode 100644 internal/config/integration_test.go create mode 100644 internal/config/substitution.go create mode 100644 internal/config/substitution_test.go diff --git a/AGENTS.md b/AGENTS.md index cd1695da..4d8e145e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,6 +8,31 @@ - **Format**: `go fmt ./...` - **Dependencies**: `go mod tidy` +## UI and Output Modes +MCPHost supports multiple output modes for different use cases: + +### Standard Mode (Default) +- Full-featured terminal UI with rich styling and formatting +- Interactive message display with proper spacing and visual hierarchy +- Suitable for regular interactive usage + +### Compact Mode (`--compact`) +- **Location**: `internal/ui/compact_renderer.go`, `internal/ui/cli.go` +- **Purpose**: Simplified output format without fancy styling for better readability in automation contexts +- **Features**: + - Single-line message format where possible + - Minimal visual styling and spacing + - Consistent symbol-based prefixes (🧑, 🤖, 🔧, etc.) + - Optimized for scripting and log parsing +- **Usage**: Add `--compact` flag to any mcphost command +- **Note**: Has no effect when combined with `--quiet` (mutually exclusive) + +### Quiet Mode (`--quiet`) +- Suppresses all UI elements except the AI response +- Only works with `--prompt` (non-interactive mode) +- Ideal for shell scripting and piping output to other commands +- **Note**: When `--quiet` is used, `--compact` has no effect since no UI elements are shown + ## Code Style Guidelines - **Package structure**: `pkg/` for reusable packages, `cmd/` for CLI commands - **Imports**: Standard library first, then third-party, then local packages with blank lines between groups @@ -83,33 +108,57 @@ MCPHost includes several builtin MCP servers for common functionality: - **Features**: HTTP/HTTPS support, HTML parsing, markdown conversion, size limits - **Security**: 5MB response limit, configurable timeouts, localhost-aware HTTPS upgrade +## Environment Variable Substitution System +MCPHost includes a comprehensive environment variable substitution system for secure configuration management. + +### Implementation Details +- **Location**: `internal/config/substitution.go` (core logic), `internal/config/substitution_test.go` (unit tests) +- **Purpose**: Replace `${env://VAR}` and `${env://VAR:-default}` patterns with environment variable values +- **Integration**: Works in both config files and script frontmatter/prompts + +### Substitution Components +- **EnvSubstituter**: Handles environment variable substitution +- **ArgsSubstituter**: Handles script argument substitution (refactored from existing code) +- **Shared Parsing Logic**: `parseVariableWithDefault()` function used by both substituters + +### Processing Order +1. **Config Loading**: `cmd/root.go` → `loadConfigWithEnvSubstitution()` → env substitution → YAML/JSON parsing +2. **Script Mode**: `cmd/script.go` → `parseScriptContent()` → env substitution → YAML parsing → args substitution + +### Security Features +- **No Shell Execution**: Direct environment variable lookup using `os.Getenv()` +- **Error Handling**: Clear error messages for missing required variables +- **Validation**: Regex-based pattern matching with comprehensive validation + ## Script System MCPHost includes a powerful script system for automation and reusable workflows. ### Script Features - **Location**: `cmd/script.go` (main implementation), `cmd/script_test.go` (comprehensive tests) -- **Purpose**: Execute YAML-based automation scripts with variable substitution +- **Purpose**: Execute YAML-based automation scripts with dual variable substitution - **Format**: YAML frontmatter + prompt content in single executable files ### Variable Substitution System -- **Required Variables**: `${variable}` - Must be provided via `--args:variable value` -- **Optional Variables**: `${variable:-default}` - Uses default if not provided +- **Environment Variables**: `${env://VAR}` and `${env://VAR:-default}` - Processed first +- **Script Arguments**: `${variable}` and `${variable:-default}` - Processed after environment variables - **Features**: - Bash-style default syntax for familiarity - - Empty defaults supported: `${var:-}` + - Empty defaults supported: `${var:-}` and `${env://VAR:-}` - Complex defaults: paths, URLs, commands with spaces - Full backward compatibility with existing scripts -- **Implementation**: Regex-based parsing with comprehensive validation +- **Implementation**: Dual-phase substitution with shared parsing logic ### Script Examples - **Location**: `examples/scripts/` directory -- **Demo Script**: `default-values-demo.sh` - Showcases new default values feature -- **Usage Examples**: Multiple scenarios from simple defaults to complex overrides +- **Demo Script**: `default-values-demo.sh` - Showcases script argument default values +- **Env Substitution Script**: `env-substitution-script.sh` - Demonstrates environment variable usage +- **Usage Examples**: Multiple scenarios from simple defaults to complex environment/args combinations ### Testing -- **Test Coverage**: 25+ test cases covering all variable scenarios -- **Edge Cases**: Empty defaults, complex values, mixed required/optional variables -- **Backward Compatibility**: Ensures existing scripts continue working unchanged +- **Unit Tests**: 25+ test cases in `internal/config/substitution_test.go` covering all substitution scenarios +- **Integration Tests**: `internal/config/integration_test.go` and `cmd/script_integration_test.go` +- **Edge Cases**: Empty defaults, complex values, mixed required/optional variables, processing order +- **Backward Compatibility**: Ensures existing scripts and configs continue working unchanged ## Authentication System MCPHost includes optional OAuth authentication for Anthropic Claude as an alternative to API keys. @@ -124,6 +173,22 @@ MCPHost includes optional OAuth authentication for Anthropic Claude as an altern - **Features**: PKCE security, automatic token refresh, encrypted storage, browser-based flow - **Priority**: OAuth credentials > API keys (environment variables/flags) +## Recent Features + +### Environment Variable Substitution (New) +- **Feature**: Added support for `${env://VAR}` and `${env://VAR:-default}` syntax in config files and scripts +- **Implementation**: + - `internal/config/substitution.go`: Core substitution logic with shared parsing + - `cmd/root.go`: Config loading integration with `loadConfigWithEnvSubstitution()` + - `cmd/script.go`: Script parsing integration with dual-phase substitution +- **Security**: Environment variables processed safely without shell execution +- **Compatibility**: Full backward compatibility with existing configurations and scripts +- **Testing**: Comprehensive test suite with 25+ unit tests and integration tests + +### Processing Flow +1. **Config Files**: Raw content → Env substitution → YAML/JSON parsing → Viper config +2. **Scripts**: Raw content → Env substitution → YAML parsing → Args substitution → Final config + ## Recent Bug Fixes ### Legacy MCP Server Configuration Fix diff --git a/README.md b/README.md index 7f5807bc..fa0b09bb 100644 --- a/README.md +++ b/README.md @@ -94,6 +94,42 @@ MCPHost will automatically create a configuration file in your home directory if You can also specify a custom location using the `--config` flag. +### Environment Variable Substitution + +MCPHost supports environment variable substitution in both config files and script frontmatter using the syntax: +- **`${env://VAR}`** - Required environment variable (fails if not set) +- **`${env://VAR:-default}`** - Optional environment variable with default value + +This allows you to keep sensitive information like API keys in environment variables while maintaining flexible configuration. + +**Example:** +```yaml +mcpServers: + github: + type: local + command: ["docker", "run", "-i", "--rm", "-e", "GITHUB_PERSONAL_ACCESS_TOKEN=${env://GITHUB_TOKEN}", "ghcr.io/github/github-mcp-server"] + environment: + DEBUG: "${env://DEBUG:-false}" + LOG_LEVEL: "${env://LOG_LEVEL:-info}" + +model: "${env://MODEL:-anthropic:claude-sonnet-4-20250514}" +provider-api-key: "${env://OPENAI_API_KEY}" # Required - will fail if not set +``` + +**Usage:** +```bash +# Set required environment variables +export GITHUB_TOKEN="ghp_your_token_here" +export OPENAI_API_KEY="your_openai_key" + +# Optionally override defaults +export DEBUG="true" +export MODEL="openai:gpt-4" + +# Run mcphost +mcphost +``` + ### Simplified Configuration Schema MCPHost now supports a simplified configuration schema with three server types: @@ -105,19 +141,28 @@ For local MCP servers that run commands on your machine: "mcpServers": { "filesystem": { "type": "local", - "command": ["npx", "@modelcontextprotocol/server-filesystem", "/tmp"], + "command": ["npx", "@modelcontextprotocol/server-filesystem", "${env://WORK_DIR:-/tmp}"], "environment": { - "DEBUG": "true", - "LOG_LEVEL": "info" + "DEBUG": "${env://DEBUG:-false}", + "LOG_LEVEL": "${env://LOG_LEVEL:-info}", + "API_TOKEN": "${env://FS_API_TOKEN}" }, "allowedTools": ["read_file", "write_file"], "excludedTools": ["delete_file"] }, + "github": { + "type": "local", + "command": ["docker", "run", "-i", "--rm", "-e", "GITHUB_PERSONAL_ACCESS_TOKEN=${env://GITHUB_TOKEN}", "ghcr.io/github/github-mcp-server"], + "environment": { + "DEBUG": "${env://DEBUG:-false}" + } + }, "sqlite": { "type": "local", - "command": ["uvx", "mcp-server-sqlite", "--db-path", "/tmp/foo.db"], + "command": ["uvx", "mcp-server-sqlite", "--db-path", "${env://DB_PATH:-/tmp/foo.db}"], "environment": { - "SQLITE_DEBUG": "1" + "SQLITE_DEBUG": "${env://DEBUG:-0}", + "DATABASE_URL": "${env://DATABASE_URL:-sqlite:///tmp/foo.db}" } } } @@ -138,12 +183,12 @@ For remote MCP servers accessible via HTTP: "mcpServers": { "websearch": { "type": "remote", - "url": "https://api.example.com/mcp", - "headers": ["Authorization: Bearer your-api-token"] + "url": "${env://WEBSEARCH_URL:-https://api.example.com/mcp}", + "headers": ["Authorization: Bearer ${env://WEBSEARCH_TOKEN}"] }, "weather": { "type": "remote", - "url": "https://weather-mcp.example.com" + "url": "${env://WEATHER_URL:-https://weather-mcp.example.com}" } } } @@ -165,7 +210,7 @@ For builtin MCP servers that run in-process for optimal performance: "type": "builtin", "name": "fs", "options": { - "allowed_directories": ["/tmp", "/home/user/documents"] + "allowed_directories": ["${env://WORK_DIR:-/tmp}", "${env://HOME}/documents"] }, "allowedTools": ["read_file", "write_file", "list_directory"] }, @@ -414,7 +459,12 @@ Each in their own environment. Give me the URL of each app #### Variable Substitution -Scripts support variable substitution using `${variable}` syntax with optional default values. Variables can be provided via command line arguments: +Scripts support both environment variable substitution and script argument substitution: + +1. **Environment Variables**: `${env://VAR}` and `${env://VAR:-default}` - Processed first +2. **Script Arguments**: `${variable}` and `${variable:-default}` - Processed after environment variables + +Variables can be provided via command line arguments: ```bash # Script with variables @@ -423,38 +473,55 @@ mcphost script myscript.sh --args:directory /tmp --args:name "John" ##### Variable Syntax -MCPHost supports two variable syntaxes: +MCPHost supports these variable syntaxes: -1. **Required Variables**: `${variable}` - Must be provided via `--args:variable value` -2. **Optional Variables with Defaults**: `${variable:-default}` - Uses default if not provided +1. **Required Environment Variables**: `${env://VAR}` - Must be set in environment +2. **Optional Environment Variables**: `${env://VAR:-default}` - Uses default if not set +3. **Required Script Arguments**: `${variable}` - Must be provided via `--args:variable value` +4. **Optional Script Arguments**: `${variable:-default}` - Uses default if not provided -Example script with mixed variables: +Example script with mixed environment variables and script arguments: ```yaml #!/usr/bin/env -S mcphost script --- mcpServers: + github: + type: "local" + command: ["gh", "api"] + environment: + GITHUB_TOKEN: "${env://GITHUB_TOKEN}" + DEBUG: "${env://DEBUG:-false}" + filesystem: type: "local" - command: ["npx", "-y", "@modelcontextprotocol/server-filesystem", "${directory:-/tmp}"] + command: ["npx", "-y", "@modelcontextprotocol/server-filesystem", "${env://WORK_DIR:-/tmp}"] + +model: "${env://MODEL:-anthropic:claude-sonnet-4-20250514}" --- -Hello ${name:-World}! Please list the files in ${directory:-/tmp} and tell me about them. -Use the ${command:-ls} command to show file details. +Hello ${name:-World}! Please list ${repo_type:-public} repositories for user ${username}. +Working directory is ${env://WORK_DIR:-/tmp}. +Use the ${command:-gh} command to fetch ${count:-10} repositories. ``` ##### Usage Examples ```bash -# Uses all defaults: name="World", directory="/tmp", command="ls" +# Set environment variables first +export GITHUB_TOKEN="ghp_your_token_here" +export DEBUG="true" +export WORK_DIR="/home/user/projects" + +# Uses env vars and defaults: name="World", repo_type="public", command="gh", count="10" mcphost script myscript.sh -# Override specific variables -mcphost script myscript.sh --args:name "John" +# Override specific script arguments +mcphost script myscript.sh --args:name "John" --args:username "alice" -# Override multiple variables -mcphost script myscript.sh --args:name "John" --args:directory "/home/john" +# Override multiple script arguments +mcphost script myscript.sh --args:name "John" --args:username "alice" --args:repo_type "private" -# Mix of provided and default values -mcphost script myscript.sh --args:name "Alice" --args:command "ls -la" +# Mix of env vars, provided args, and default values +mcphost script myscript.sh --args:name "Alice" --args:command "gh api" --args:count "5" ``` ##### Default Value Features @@ -464,7 +531,11 @@ mcphost script myscript.sh --args:name "Alice" --args:command "ls -la" - **Spaces in defaults**: `${msg:-Hello World}` - Supports spaces in default values - **Backward compatibility**: Existing `${variable}` syntax continues to work unchanged -**Important**: Only variables without defaults (e.g., `${directory}`, `${name}`) are required. Variables with defaults are optional and will use their default value if not provided via `--args:variable value` syntax. +**Important**: +- Environment variables without defaults (e.g., `${env://GITHUB_TOKEN}`) are required and must be set in the environment +- Script arguments without defaults (e.g., `${username}`) are required and must be provided via `--args:variable value` syntax +- Variables with defaults are optional and will use their default value if not provided +- Environment variables are processed first, then script arguments #### Script Features @@ -552,7 +623,10 @@ mcphost --model openai: \ # Single prompt with full UI mcphost -p "List files in the current directory" -# Quiet mode for scripting (only AI response output) +# Compact mode for cleaner output without fancy styling +mcphost -p "List files in the current directory" --compact + +# Quiet mode for scripting (only AI response output, no UI elements) mcphost -p "What is the capital of France?" --quiet # Use in shell scripts @@ -573,6 +647,7 @@ mcphost -p "Generate a random UUID" --quiet | tr '[:lower:]' '[:upper:]' - `-m, --model string`: Model to use (format: provider:model) (default "anthropic:claude-sonnet-4-20250514") - `-p, --prompt string`: **Run in non-interactive mode with the given prompt** - `--quiet`: **Suppress all output except the AI response (only works with --prompt)** +- `--compact`: **Enable compact output mode without fancy styling (ideal for scripting and automation)** - `--stream`: Enable streaming responses (default: true, use `--stream=false` to disable) ### Authentication Subcommands @@ -704,12 +779,34 @@ curl -X POST http://localhost:8080/process \ ``` ### Tips for Scripting -- Use `--quiet` flag to get clean output suitable for parsing +- Use `--quiet` flag to get clean output suitable for parsing (only AI response, no UI) +- Use `--compact` flag for simplified output without fancy styling (when you want to see UI elements) +- Note: `--compact` and `--quiet` are mutually exclusive - `--compact` has no effect with `--quiet` +- **Use environment variables for sensitive data** like API keys instead of hardcoding them +- **Use `${env://VAR}` syntax** in config files and scripts for environment variable substitution - Combine with standard Unix tools (`grep`, `awk`, `sed`, etc.) - Set appropriate timeouts for long-running operations - Handle errors appropriately in your scripts - Use environment variables for API keys in production +#### Environment Variable Best Practices +```bash +# Set sensitive variables in environment +export GITHUB_TOKEN="ghp_your_token_here" +export OPENAI_API_KEY="your_openai_key" +export DATABASE_URL="postgresql://user:pass@localhost/db" + +# Use in config files +mcpServers: + github: + environment: + GITHUB_TOKEN: "${env://GITHUB_TOKEN}" + DEBUG: "${env://DEBUG:-false}" + +# Use in scripts +mcphost script my-script.sh --args:username alice +``` + ## MCP Server Compatibility 🔌 MCPHost can work with any MCP-compliant server. For examples and reference implementations, see the [MCP Servers Repository](https://github.com/modelcontextprotocol/servers). diff --git a/cmd/root.go b/cmd/root.go index 4d64e4ad..7693d707 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -7,6 +7,7 @@ import ( "io" "log" "os" + "path/filepath" "strings" "github.com/cloudwego/eino/schema" @@ -98,10 +99,7 @@ func Execute(v string) { func initConfig() { if configFile != "" { // Use config file from the flag - viper.SetConfigFile(configFile) - - // Try to read the specified config file - if err := viper.ReadInConfig(); err != nil { + if err := loadConfigWithEnvSubstitution(configFile); err != nil { fmt.Fprintf(os.Stderr, "Error reading config file '%s': %v\n", configFile, err) os.Exit(1) } @@ -119,24 +117,40 @@ func initConfig() { os.Exit(1) } - // Search config in home directory with name ".mcphost" (without extension) - viper.AddConfigPath(home) - viper.SetConfigName(".mcphost") - viper.SetConfigType("yaml") + // Try to load config files with environment variable substitution + configLoaded := false + configPaths := []struct { + name string + types []string + }{ + {".mcphost", []string{"yml", "yaml", "json"}}, + {".mcp", []string{"yml", "yaml", "json"}}, + } - // Also try JSON format - if err := viper.ReadInConfig(); err != nil { - viper.SetConfigType("json") - if err := viper.ReadInConfig(); err != nil { - // Try legacy .mcp files - viper.SetConfigName(".mcp") - viper.SetConfigType("yaml") - if err := viper.ReadInConfig(); err != nil { - viper.SetConfigType("json") - viper.ReadInConfig() // Ignore error if no config found + for _, configPath := range configPaths { + for _, configType := range configPath.types { + fullPath := filepath.Join(home, configPath.name+"."+configType) + if _, err := os.Stat(fullPath); err == nil { + if err := loadConfigWithEnvSubstitution(fullPath); err != nil { + // Only exit on environment variable substitution errors + // Other errors should be handled gracefully + if strings.Contains(err.Error(), "environment variable substitution failed") { + fmt.Fprintf(os.Stderr, "Error reading config file '%s': %v\n", fullPath, err) + os.Exit(1) + } + // For other errors, continue trying other config files + continue + } + configLoaded = true + break } } + if configLoaded { + break + } } + + // If no config file was loaded, continue without error (optional config) } // Set environment variable prefix @@ -144,6 +158,32 @@ func initConfig() { viper.AutomaticEnv() } +// loadConfigWithEnvSubstitution loads a config file with environment variable substitution +func loadConfigWithEnvSubstitution(configPath string) error { + // Read raw config file content + rawContent, err := os.ReadFile(configPath) + if err != nil { + return fmt.Errorf("failed to read config file: %v", err) + } + + // Apply environment variable substitution + substituter := &config.EnvSubstituter{} + processedContent, err := substituter.SubstituteEnvVars(string(rawContent)) + if err != nil { + return fmt.Errorf("config env substitution failed: %v", err) + } + + // Determine config type from file extension + configType := "yaml" + if strings.HasSuffix(configPath, ".json") { + configType = "json" + } + + // Use viper to parse the processed content + viper.SetConfigType(configType) + return viper.ReadConfig(strings.NewReader(processedContent)) +} + func init() { cobra.OnInitialize(initConfig) @@ -348,6 +388,9 @@ func runNormalMode(ctx context.Context) error { // Set the model name for consistent display cli.SetModelName(modelName) + // Set the model name for consistent display + cli.SetModelName(modelName) + // Set up usage tracking for supported providers if len(parts) == 2 { provider := parts[0] @@ -413,6 +456,53 @@ func runNormalMode(ctx context.Context) error { debugConfig["provider-api-key"] = "[SET]" } + // Add MCP server configuration for debugging + if len(mcpConfig.MCPServers) > 0 { + mcpServers := make(map[string]any) + loadedServers := mcpAgent.GetLoadedServerNames() + loadedServerSet := make(map[string]bool) + for _, name := range loadedServers { + loadedServerSet[name] = true + } + + for name, server := range mcpConfig.MCPServers { + serverInfo := map[string]any{ + "type": server.Type, + "status": "failed", // Default to failed + } + + // Mark as loaded if it's in the loaded servers list + if loadedServerSet[name] { + serverInfo["status"] = "loaded" + } + + if len(server.Command) > 0 { + serverInfo["command"] = server.Command + } + if len(server.Environment) > 0 { + // Mask sensitive environment variables + maskedEnv := make(map[string]string) + for k, v := range server.Environment { + if strings.Contains(strings.ToLower(k), "token") || + strings.Contains(strings.ToLower(k), "key") || + strings.Contains(strings.ToLower(k), "secret") { + maskedEnv[k] = "[MASKED]" + } else { + maskedEnv[k] = v + } + } + serverInfo["environment"] = maskedEnv + } + if server.URL != "" { + serverInfo["url"] = server.URL + } + if server.Name != "" { + serverInfo["name"] = server.Name + } + mcpServers[name] = serverInfo + } + debugConfig["mcpServers"] = mcpServers + } cli.DisplayDebugConfig(debugConfig) } } diff --git a/cmd/script.go b/cmd/script.go index 34eb2bfb..82f8fc9e 100644 --- a/cmd/script.go +++ b/cmd/script.go @@ -277,13 +277,24 @@ func readRemainingLines(scanner *bufio.Scanner) string { // parseScriptContent parses the content to extract YAML frontmatter and prompt func parseScriptContent(content string, variables map[string]string) (*config.Config, error) { - // First, validate that all declared variables are provided - if err := validateVariables(content, variables); err != nil { + // STEP 1: Apply environment variable substitution FIRST + envSubstituter := &config.EnvSubstituter{} + processedContent, err := envSubstituter.SubstituteEnvVars(content) + if err != nil { + return nil, fmt.Errorf("script env substitution failed: %v", err) + } + + // STEP 2: Validate that all declared script variables are provided + if err := validateVariables(processedContent, variables); err != nil { return nil, err } - // Substitute variables in the content - content = substituteVariables(content, variables) + // STEP 3: Apply script args substitution + argsSubstituter := config.NewArgsSubstituter(variables) + content, err = argsSubstituter.SubstituteArgs(processedContent) + if err != nil { + return nil, fmt.Errorf("script args substitution failed: %v", err) + } lines := strings.Split(content, "\n") @@ -430,38 +441,16 @@ func validateVariables(content string, variables map[string]string) error { } // substituteVariables replaces ${variable} and ${variable:-default} patterns with their values +// This function is kept for backward compatibility but now uses the shared ArgsSubstituter func substituteVariables(content string, variables map[string]string) string { - // Pattern matches both ${varname} and ${varname:-default} - re := regexp.MustCompile(`\$\{([^}:]+)(?::-([^}]*))?\}`) - - return re.ReplaceAllStringFunc(content, func(match string) string { - // Parse the match to extract variable name and default value - submatches := re.FindStringSubmatch(match) - if len(submatches) < 2 { - return match // Should not happen, but safety check - } - - varName := submatches[1] - // Check if the original match contains the :- pattern - hasDefault := strings.Contains(match, ":-") - defaultValue := "" - if hasDefault && len(submatches) >= 3 { - defaultValue = submatches[2] // Can be empty string - } - - // Look up the variable value - if value, exists := variables[varName]; exists { - return value - } - - // Use default value if available - if hasDefault { - return defaultValue - } - - // If no value and no default, leave it as is - return match - }) + substituter := config.NewArgsSubstituter(variables) + result, err := substituter.SubstituteArgs(content) + if err != nil { + // For backward compatibility, if there's an error, return the original content + // This maintains the existing behavior where missing variables were left as-is + return content + } + return result } // runScriptMode executes the script using the unified agentic loop diff --git a/cmd/script_integration_test.go b/cmd/script_integration_test.go new file mode 100644 index 00000000..4b36852f --- /dev/null +++ b/cmd/script_integration_test.go @@ -0,0 +1,284 @@ +package cmd + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestScriptWithEnvAndArgsSubstitution(t *testing.T) { + // Create a temporary script file with both env vars and script args + tempDir := t.TempDir() + scriptPath := filepath.Join(tempDir, "test-script.sh") + + scriptContent := `#!/usr/bin/env -S mcphost script +--- +mcpServers: + github: + type: local + command: ["gh", "api"] + environment: + GITHUB_TOKEN: "${env://GITHUB_TOKEN}" + DEBUG: "${env://DEBUG:-false}" + + filesystem: + type: builtin + name: fs + options: + allowed_directories: ["${env://WORK_DIR:-/tmp}"] + +model: "${env://MODEL:-anthropic:claude-sonnet-4-20250514}" +debug: ${env://DEBUG:-false} +--- +List ${repo_type:-public} repositories for user ${username}. +Use the GitHub API to fetch ${count:-10} repositories. +Working directory is ${env://WORK_DIR:-/tmp}. +` + + err := os.WriteFile(scriptPath, []byte(scriptContent), 0644) + if err != nil { + t.Fatalf("Failed to write test script: %v", err) + } + + // Set up environment variables + os.Setenv("GITHUB_TOKEN", "ghp_test_token") + os.Setenv("DEBUG", "true") + os.Setenv("WORK_DIR", "/home/user/projects") + defer func() { + os.Unsetenv("GITHUB_TOKEN") + os.Unsetenv("DEBUG") + os.Unsetenv("WORK_DIR") + }() + + // Set up script arguments + variables := map[string]string{ + "username": "alice", + "repo_type": "private", + } + + // Parse the script + scriptConfig, err := parseScriptFile(scriptPath, variables) + if err != nil { + t.Fatalf("Failed to parse script: %v", err) + } + + // Verify environment variable substitution in MCP servers + githubServer, exists := scriptConfig.MCPServers["github"] + if !exists { + t.Fatal("GitHub server not found in script config") + } + + if githubServer.Environment["GITHUB_TOKEN"] != "ghp_test_token" { + t.Errorf("Expected GITHUB_TOKEN=ghp_test_token, got %s", githubServer.Environment["GITHUB_TOKEN"]) + } + if githubServer.Environment["DEBUG"] != "true" { + t.Errorf("Expected DEBUG=true, got %s", githubServer.Environment["DEBUG"]) + } + + // Verify environment variable substitution in builtin server options + fsServer, exists := scriptConfig.MCPServers["filesystem"] + if !exists { + t.Fatal("Filesystem server not found in script config") + } + + allowedDirs, ok := fsServer.Options["allowed_directories"].([]interface{}) + if !ok { + t.Fatal("allowed_directories should be an array") + } + if len(allowedDirs) != 1 || allowedDirs[0] != "/home/user/projects" { + t.Errorf("Expected allowed_directories=[/home/user/projects], got %v", allowedDirs) + } + + // Verify global config values + if scriptConfig.Model != "anthropic:claude-sonnet-4-20250514" { + t.Errorf("Expected model=anthropic:claude-sonnet-4-20250514, got %s", scriptConfig.Model) + } + if !scriptConfig.Debug { + t.Error("Expected debug=true") + } + + // Verify script args substitution in prompt + expectedPrompt := `List private repositories for user alice. +Use the GitHub API to fetch 10 repositories. +Working directory is /home/user/projects.` + + if strings.TrimSpace(scriptConfig.Prompt) != strings.TrimSpace(expectedPrompt) { + t.Errorf("Expected prompt:\n%s\nGot:\n%s", expectedPrompt, scriptConfig.Prompt) + } +} + +func TestScriptWithMissingRequiredEnvVar(t *testing.T) { + // Create a script with a required environment variable + tempDir := t.TempDir() + scriptPath := filepath.Join(tempDir, "test-script.sh") + + scriptContent := `#!/usr/bin/env -S mcphost script +--- +mcpServers: + github: + type: local + command: ["gh", "api"] + environment: + GITHUB_TOKEN: "${env://REQUIRED_TOKEN}" +--- +Test script with required env var. +` + + err := os.WriteFile(scriptPath, []byte(scriptContent), 0644) + if err != nil { + t.Fatalf("Failed to write test script: %v", err) + } + + // Make sure the environment variable is not set + os.Unsetenv("REQUIRED_TOKEN") + + // Parse the script - should fail + variables := map[string]string{} + _, err = parseScriptFile(scriptPath, variables) + if err == nil { + t.Fatal("Expected error for missing required environment variable") + } + + if !strings.Contains(err.Error(), "required environment variable REQUIRED_TOKEN not set") { + t.Errorf("Expected error about missing REQUIRED_TOKEN, got: %v", err) + } +} + +func TestScriptWithMissingRequiredArg(t *testing.T) { + // Create a script with a required script argument + tempDir := t.TempDir() + scriptPath := filepath.Join(tempDir, "test-script.sh") + + scriptContent := `#!/usr/bin/env -S mcphost script +--- +mcpServers: + github: + type: local + command: ["gh", "api"] + environment: + GITHUB_TOKEN: "${env://GITHUB_TOKEN:-default_token}" +--- +List repositories for user ${required_username}. +` + + err := os.WriteFile(scriptPath, []byte(scriptContent), 0644) + if err != nil { + t.Fatalf("Failed to write test script: %v", err) + } + + // Parse the script without providing the required argument + variables := map[string]string{} + _, err = parseScriptFile(scriptPath, variables) + if err == nil { + t.Fatal("Expected error for missing required script argument") + } + + if !strings.Contains(err.Error(), "required_username") { + t.Errorf("Expected error about missing required_username, got: %v", err) + } +} + +func TestScriptProcessingOrder(t *testing.T) { + // Test that env substitution happens before args substitution + tempDir := t.TempDir() + scriptPath := filepath.Join(tempDir, "test-script.sh") + + // This script tests that env vars are processed first, then script args + scriptContent := `#!/usr/bin/env -S mcphost script +--- +mcpServers: + test: + type: local + command: ["echo", "${env://BASE_PATH:-/tmp}/${path_suffix}"] +--- +Base path is ${env://BASE_PATH:-/tmp} and suffix is ${path_suffix:-default}. +` + + err := os.WriteFile(scriptPath, []byte(scriptContent), 0644) + if err != nil { + t.Fatalf("Failed to write test script: %v", err) + } + + // Set environment variable + os.Setenv("BASE_PATH", "/home/user") + defer os.Unsetenv("BASE_PATH") + + // Set script argument + variables := map[string]string{ + "path_suffix": "documents", + } + + // Parse the script + scriptConfig, err := parseScriptFile(scriptPath, variables) + if err != nil { + t.Fatalf("Failed to parse script: %v", err) + } + + // Verify that both substitutions worked correctly + testServer := scriptConfig.MCPServers["test"] + expectedCommand := []string{"echo", "/home/user/documents"} + + if len(testServer.Command) != len(expectedCommand) { + t.Errorf("Expected command length %d, got %d", len(expectedCommand), len(testServer.Command)) + } + + for i, expected := range expectedCommand { + if i < len(testServer.Command) && testServer.Command[i] != expected { + t.Errorf("Expected command[%d] = %s, got %s", i, expected, testServer.Command[i]) + } + } + + // Verify prompt substitution + expectedPrompt := "Base path is /home/user and suffix is documents." + if strings.TrimSpace(scriptConfig.Prompt) != expectedPrompt { + t.Errorf("Expected prompt: %s\nGot: %s", expectedPrompt, strings.TrimSpace(scriptConfig.Prompt)) + } +} + +func TestScriptBackwardCompatibility(t *testing.T) { + // Test that existing scripts without env vars still work + tempDir := t.TempDir() + scriptPath := filepath.Join(tempDir, "test-script.sh") + + scriptContent := `#!/usr/bin/env -S mcphost script +--- +mcpServers: + filesystem: + type: builtin + name: fs + options: + allowed_directories: ["/tmp"] + +model: "anthropic:claude-sonnet-4-20250514" +--- +List files in ${directory:-/tmp} for user ${username}. +` + + err := os.WriteFile(scriptPath, []byte(scriptContent), 0644) + if err != nil { + t.Fatalf("Failed to write test script: %v", err) + } + + // Set script arguments + variables := map[string]string{ + "username": "bob", + } + + // Parse the script + scriptConfig, err := parseScriptFile(scriptPath, variables) + if err != nil { + t.Fatalf("Failed to parse script: %v", err) + } + + // Verify that script args substitution still works + expectedPrompt := "List files in /tmp for user bob." + if strings.TrimSpace(scriptConfig.Prompt) != expectedPrompt { + t.Errorf("Expected prompt: %s\nGot: %s", expectedPrompt, strings.TrimSpace(scriptConfig.Prompt)) + } + + // Verify that config is unchanged + if scriptConfig.Model != "anthropic:claude-sonnet-4-20250514" { + t.Errorf("Expected model unchanged, got %s", scriptConfig.Model) + } +} diff --git a/examples/scripts/env-substitution-script.sh b/examples/scripts/env-substitution-script.sh new file mode 100644 index 00000000..15785ee8 --- /dev/null +++ b/examples/scripts/env-substitution-script.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env -S mcphost script +--- +# Example script demonstrating both environment variable and script argument substitution +# Environment variables are processed first, then script arguments + +mcpServers: + github: + type: local + command: ["gh", "api"] + environment: + GITHUB_TOKEN: "${env://GITHUB_TOKEN}" + DEBUG: "${env://DEBUG:-false}" + + filesystem: + type: builtin + name: fs + options: + allowed_directories: ["${env://WORK_DIR:-/tmp}"] + +model: "${env://MODEL:-anthropic:claude-sonnet-4-20250514}" +debug: ${env://DEBUG:-false} +--- +List ${repo_type:-public} repositories for user ${username}. +Use the GitHub API to fetch ${count:-10} repositories. +Working directory is ${env://WORK_DIR:-/tmp}. + +# Usage: +# 1. Set environment variables: +# export GITHUB_TOKEN="ghp_your_token_here" +# export DEBUG="true" +# export WORK_DIR="/home/user/projects" +# +# 2. Run with script arguments: +# mcphost script env-substitution-script.sh --args:username alice --args:repo_type private --args:count 5 +# +# This will: +# - Use GITHUB_TOKEN from environment +# - Set DEBUG=true from environment +# - Use WORK_DIR=/home/user/projects from environment +# - Use username=alice from script args +# - Use repo_type=private from script args +# - Use count=5 from script args \ No newline at end of file diff --git a/internal/agent/agent.go b/internal/agent/agent.go index ce0cba43..c69c8045 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -251,6 +251,11 @@ func (a *Agent) GetLoadingMessage() string { return a.loadingMessage } +// GetLoadedServerNames returns the names of successfully loaded MCP servers +func (a *Agent) GetLoadedServerNames() []string { + return a.toolManager.GetLoadedServerNames() +} + // generateWithCancellationAndStreaming calls the LLM with ESC key cancellation support and streaming callbacks func (a *Agent) generateWithCancellationAndStreaming(ctx context.Context, messages []*schema.Message, toolInfos []*schema.ToolInfo, streamingCallback StreamingResponseHandler) (*schema.Message, error) { // Check if streaming is enabled diff --git a/internal/config/config.go b/internal/config/config.go index 68ad030a..8d86dcc8 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -16,8 +16,8 @@ type MCPServerConfig struct { URL string `json:"url,omitempty"` Name string `json:"name,omitempty"` // For builtin servers Options map[string]any `json:"options,omitempty"` // For builtin servers - AllowedTools []string `json:"allowedTools,omitempty"` - ExcludedTools []string `json:"excludedTools,omitempty"` + AllowedTools []string `json:"allowedTools,omitempty" yaml:"allowedTools,omitempty"` + ExcludedTools []string `json:"excludedTools,omitempty" yaml:"excludedTools,omitempty"` // Legacy fields for backward compatibility Transport string `json:"transport,omitempty"` @@ -37,8 +37,8 @@ func (s *MCPServerConfig) UnmarshalJSON(data []byte) error { Headers []string `json:"headers,omitempty"` Name string `json:"name,omitempty"` Options map[string]any `json:"options,omitempty"` - AllowedTools []string `json:"allowedTools,omitempty"` - ExcludedTools []string `json:"excludedTools,omitempty"` + AllowedTools []string `json:"allowedTools,omitempty" yaml:"allowedTools,omitempty"` + ExcludedTools []string `json:"excludedTools,omitempty" yaml:"excludedTools,omitempty"` } // Also try legacy format @@ -49,8 +49,8 @@ func (s *MCPServerConfig) UnmarshalJSON(data []byte) error { Env map[string]any `json:"env,omitempty"` URL string `json:"url,omitempty"` Headers []string `json:"headers,omitempty"` - AllowedTools []string `json:"allowedTools,omitempty"` - ExcludedTools []string `json:"excludedTools,omitempty"` + AllowedTools []string `json:"allowedTools,omitempty" yaml:"allowedTools,omitempty"` + ExcludedTools []string `json:"excludedTools,omitempty" yaml:"excludedTools,omitempty"` } // Try new format first diff --git a/internal/config/integration_test.go b/internal/config/integration_test.go new file mode 100644 index 00000000..67d00b9f --- /dev/null +++ b/internal/config/integration_test.go @@ -0,0 +1,227 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/spf13/viper" +) + +func TestConfigLoadingWithEnvSubstitution(t *testing.T) { + // Create a temporary config file with environment variables + tempDir := t.TempDir() + configPath := filepath.Join(tempDir, "test-config.yaml") + + configContent := ` +mcpServers: + github: + type: local + command: ["docker", "run", "-i", "--rm", "-e", "GITHUB_PERSONAL_ACCESS_TOKEN=${env://GITHUB_TOKEN}", "ghcr.io/github/github-mcp-server"] + environment: + DEBUG: "${env://DEBUG:-false}" + LOG_LEVEL: "${env://LOG_LEVEL:-info}" + + database: + type: local + command: ["python", "db-server.py"] + environment: + DATABASE_URL: "${env://DATABASE_URL:-sqlite:///tmp/default.db}" + API_KEY: "${env://DB_API_KEY}" + +model: "${env://MODEL:-anthropic:claude-sonnet-4-20250514}" +provider-api-key: "${env://OPENAI_API_KEY:-}" +debug: ${env://DEBUG:-false} +` + + err := os.WriteFile(configPath, []byte(configContent), 0644) + if err != nil { + t.Fatalf("Failed to write test config: %v", err) + } + + // Set up environment variables + os.Setenv("GITHUB_TOKEN", "ghp_test_token") + os.Setenv("DEBUG", "true") + os.Setenv("DB_API_KEY", "secret_key") + defer func() { + os.Unsetenv("GITHUB_TOKEN") + os.Unsetenv("DEBUG") + os.Unsetenv("DB_API_KEY") + }() + + // Test the loadConfigWithEnvSubstitution function + viper.Reset() // Clear any existing config + err = loadConfigWithEnvSubstitution(configPath) + if err != nil { + t.Fatalf("Failed to load config with env substitution: %v", err) + } + + // Verify that environment variables were substituted correctly + var config Config + err = viper.Unmarshal(&config) + if err != nil { + t.Fatalf("Failed to unmarshal config: %v", err) + } + + // Check GitHub server config + githubServer, exists := config.MCPServers["github"] + if !exists { + t.Fatal("GitHub server not found in config") + } + + // Check that GITHUB_TOKEN was substituted in command + expectedCommand := []string{"docker", "run", "-i", "--rm", "-e", "GITHUB_PERSONAL_ACCESS_TOKEN=ghp_test_token", "ghcr.io/github/github-mcp-server"} + if len(githubServer.Command) != len(expectedCommand) { + t.Errorf("Expected command length %d, got %d", len(expectedCommand), len(githubServer.Command)) + } + for i, expected := range expectedCommand { + if i < len(githubServer.Command) && githubServer.Command[i] != expected { + t.Errorf("Expected command[%d] = %s, got %s", i, expected, githubServer.Command[i]) + } + } + + // Check environment variables (viper converts keys to lowercase) + if githubServer.Environment["debug"] != "true" { + t.Errorf("Expected debug=true, got %s", githubServer.Environment["debug"]) + } + if githubServer.Environment["log_level"] != "info" { + t.Errorf("Expected log_level=info, got %s", githubServer.Environment["log_level"]) + } + + // Check database server config + dbServer, exists := config.MCPServers["database"] + if !exists { + t.Fatal("Database server not found in config") + } + + if dbServer.Environment["database_url"] != "sqlite:///tmp/default.db" { + t.Errorf("Expected database_url=sqlite:///tmp/default.db, got %s", dbServer.Environment["database_url"]) + } + if dbServer.Environment["api_key"] != "secret_key" { + t.Errorf("Expected api_key=secret_key, got %s", dbServer.Environment["api_key"]) + } + + // Check global config values + if config.Model != "anthropic:claude-sonnet-4-20250514" { + t.Errorf("Expected model=anthropic:claude-sonnet-4-20250514, got %s", config.Model) + } + if !config.Debug { + t.Error("Expected debug=true") + } +} + +func TestConfigLoadingWithMissingRequiredEnvVar(t *testing.T) { + // Create a temporary config file with a required environment variable + tempDir := t.TempDir() + configPath := filepath.Join(tempDir, "test-config.yaml") + + configContent := ` +mcpServers: + github: + type: local + command: ["gh", "api"] + environment: + GITHUB_TOKEN: "${env://REQUIRED_GITHUB_TOKEN}" +` + + err := os.WriteFile(configPath, []byte(configContent), 0644) + if err != nil { + t.Fatalf("Failed to write test config: %v", err) + } + + // Make sure the environment variable is not set + os.Unsetenv("REQUIRED_GITHUB_TOKEN") + + // Test that loading fails with a clear error + viper.Reset() + err = loadConfigWithEnvSubstitution(configPath) + if err == nil { + t.Fatal("Expected error for missing required environment variable") + } + + if !strings.Contains(err.Error(), "required environment variable REQUIRED_GITHUB_TOKEN not set") { + t.Errorf("Expected error about missing REQUIRED_GITHUB_TOKEN, got: %v", err) + } +} + +func TestJSONConfigWithEnvSubstitution(t *testing.T) { + // Test JSON format as well + tempDir := t.TempDir() + configPath := filepath.Join(tempDir, "test-config.json") + + configContent := `{ + "mcpServers": { + "test": { + "type": "local", + "command": ["echo", "${env://TEST_VALUE:-hello}"], + "environment": { + "VAR1": "${env://VAR1:-default1}", + "VAR2": "${env://VAR2}" + } + } + }, + "model": "${env://MODEL:-anthropic:claude-sonnet-4-20250514}" +}` + + err := os.WriteFile(configPath, []byte(configContent), 0644) + if err != nil { + t.Fatalf("Failed to write test config: %v", err) + } + + // Set one environment variable, leave others to use defaults + os.Setenv("VAR2", "custom_value") + defer os.Unsetenv("VAR2") + + viper.Reset() + err = loadConfigWithEnvSubstitution(configPath) + if err != nil { + t.Fatalf("Failed to load JSON config with env substitution: %v", err) + } + + var config Config + err = viper.Unmarshal(&config) + if err != nil { + t.Fatalf("Failed to unmarshal config: %v", err) + } + + testServer := config.MCPServers["test"] + expectedCommand := []string{"echo", "hello"} + if len(testServer.Command) != len(expectedCommand) || testServer.Command[1] != "hello" { + t.Errorf("Expected command %v, got %v", expectedCommand, testServer.Command) + } + + if testServer.Environment["var1"] != "default1" { + t.Errorf("Expected var1=default1, got %s", testServer.Environment["var1"]) + } + if testServer.Environment["var2"] != "custom_value" { + t.Errorf("Expected var2=custom_value, got %s", testServer.Environment["var2"]) + } +} + +// Helper function to simulate the loadConfigWithEnvSubstitution function +// This is needed because the function is in cmd/root.go and we can't easily test it directly +func loadConfigWithEnvSubstitution(configPath string) error { + // Read raw config file content + rawContent, err := os.ReadFile(configPath) + if err != nil { + return err + } + + // Apply environment variable substitution + substituter := &EnvSubstituter{} + processedContent, err := substituter.SubstituteEnvVars(string(rawContent)) + if err != nil { + return err + } + + // Determine config type from file extension + configType := "yaml" + if strings.HasSuffix(configPath, ".json") { + configType = "json" + } + + // Use viper to parse the processed content + viper.SetConfigType(configType) + return viper.ReadConfig(strings.NewReader(processedContent)) +} diff --git a/internal/config/substitution.go b/internal/config/substitution.go new file mode 100644 index 00000000..567a826f --- /dev/null +++ b/internal/config/substitution.go @@ -0,0 +1,108 @@ +package config + +import ( + "fmt" + "os" + "regexp" + "strings" +) + +// Variable substitution patterns +var ( + envVarPattern = regexp.MustCompile(`\$\{env://([A-Za-z_][A-Za-z0-9_]*)(:-([^}]*))?\}`) + scriptArgsPattern = regexp.MustCompile(`\$\{([A-Za-z_][A-Za-z0-9_]*)(:-([^}]*))?\}`) +) + +// parseVariableWithDefault extracts variable name and default value +// Works for both ${var:-default} and ${env://var:-default} patterns +func parseVariableWithDefault(varPart string) (varName, defaultValue string, hasDefault bool) { + // Handle the case where varPart is like "VAR:-default" or just "VAR" + if strings.Contains(varPart, ":-") { + parts := strings.SplitN(varPart, ":-", 2) + return parts[0], parts[1], true + } + return varPart, "", false +} + +// EnvSubstituter handles environment variable substitution +type EnvSubstituter struct{} + +// SubstituteEnvVars replaces ${env://VAR} and ${env://VAR:-default} patterns with environment variables +func (e *EnvSubstituter) SubstituteEnvVars(content string) (string, error) { + var errors []string + + result := envVarPattern.ReplaceAllStringFunc(content, func(match string) string { + // Extract the variable part from ${env://VAR:-default} + // Remove ${env:// prefix and } suffix + varPart := strings.TrimPrefix(strings.TrimSuffix(match, "}"), "${env://") + + varName, defaultValue, hasDefault := parseVariableWithDefault(varPart) + + if envValue := os.Getenv(varName); envValue != "" { + return envValue + } + + if hasDefault { + return defaultValue + } + + errors = append(errors, fmt.Sprintf("required environment variable %s not set in %s", varName, match)) + return match // Keep original if error + }) + + if len(errors) > 0 { + return "", fmt.Errorf("environment variable substitution failed: %s", strings.Join(errors, ", ")) + } + + return result, nil +} + +// ArgsSubstituter handles script argument substitution +type ArgsSubstituter struct { + args map[string]string +} + +// NewArgsSubstituter creates a new args substituter with the given arguments +func NewArgsSubstituter(args map[string]string) *ArgsSubstituter { + return &ArgsSubstituter{args: args} +} + +// SubstituteArgs replaces ${VAR} and ${VAR:-default} patterns with script arguments +func (a *ArgsSubstituter) SubstituteArgs(content string) (string, error) { + var errors []string + + result := scriptArgsPattern.ReplaceAllStringFunc(content, func(match string) string { + // Extract the variable part from ${VAR:-default} + // Remove ${ prefix and } suffix + varPart := strings.TrimPrefix(strings.TrimSuffix(match, "}"), "${") + + varName, defaultValue, hasDefault := parseVariableWithDefault(varPart) + + if argValue, exists := a.args[varName]; exists { + return argValue + } + + if hasDefault { + return defaultValue + } + + errors = append(errors, fmt.Sprintf("required script argument '%s' not set in %s", varName, match)) + return match // Keep original if error + }) + + if len(errors) > 0 { + return "", fmt.Errorf("script argument substitution failed: %s", strings.Join(errors, ", ")) + } + + return result, nil +} + +// HasEnvVars checks if content contains environment variable patterns +func HasEnvVars(content string) bool { + return envVarPattern.MatchString(content) +} + +// HasScriptArgs checks if content contains script argument patterns +func HasScriptArgs(content string) bool { + return scriptArgsPattern.MatchString(content) +} diff --git a/internal/config/substitution_test.go b/internal/config/substitution_test.go new file mode 100644 index 00000000..be635c70 --- /dev/null +++ b/internal/config/substitution_test.go @@ -0,0 +1,395 @@ +package config + +import ( + "os" + "testing" +) + +func TestParseVariableWithDefault(t *testing.T) { + tests := []struct { + name string + input string + expectedVar string + expectedDefault string + expectedHasDefault bool + }{ + { + name: "variable without default", + input: "GITHUB_TOKEN", + expectedVar: "GITHUB_TOKEN", + expectedDefault: "", + expectedHasDefault: false, + }, + { + name: "variable with default", + input: "DEBUG:-false", + expectedVar: "DEBUG", + expectedDefault: "false", + expectedHasDefault: true, + }, + { + name: "variable with empty default", + input: "OPTIONAL:-", + expectedVar: "OPTIONAL", + expectedDefault: "", + expectedHasDefault: true, + }, + { + name: "variable with complex default", + input: "DATABASE_URL:-sqlite:///tmp/default.db", + expectedVar: "DATABASE_URL", + expectedDefault: "sqlite:///tmp/default.db", + expectedHasDefault: true, + }, + { + name: "variable with default containing colon", + input: "URL:-https://api.example.com:8080/path", + expectedVar: "URL", + expectedDefault: "https://api.example.com:8080/path", + expectedHasDefault: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + varName, defaultValue, hasDefault := parseVariableWithDefault(tt.input) + + if varName != tt.expectedVar { + t.Errorf("Expected var name %s, got %s", tt.expectedVar, varName) + } + if defaultValue != tt.expectedDefault { + t.Errorf("Expected default value %s, got %s", tt.expectedDefault, defaultValue) + } + if hasDefault != tt.expectedHasDefault { + t.Errorf("Expected hasDefault %v, got %v", tt.expectedHasDefault, hasDefault) + } + }) + } +} + +func TestEnvSubstituter_SubstituteEnvVars(t *testing.T) { + tests := []struct { + name string + input string + envVars map[string]string + expected string + expectError bool + }{ + { + name: "basic env substitution", + input: `{"token": "${env://GITHUB_TOKEN}"}`, + envVars: map[string]string{"GITHUB_TOKEN": "ghp_123"}, + expected: `{"token": "ghp_123"}`, + }, + { + name: "env with default value used", + input: `{"debug": "${env://DEBUG:-false}"}`, + envVars: map[string]string{}, + expected: `{"debug": "false"}`, + }, + { + name: "env with default value overridden", + input: `{"debug": "${env://DEBUG:-false}"}`, + envVars: map[string]string{"DEBUG": "true"}, + expected: `{"debug": "true"}`, + }, + { + name: "env with empty default", + input: `{"optional": "${env://OPTIONAL:-}"}`, + envVars: map[string]string{}, + expected: `{"optional": ""}`, + }, + { + name: "multiple env vars in same string", + input: `{"url": "${env://HOST:-localhost}:${env://PORT:-8080}"}`, + envVars: map[string]string{"HOST": "example.com"}, + expected: `{"url": "example.com:8080"}`, + }, + { + name: "mixed env and script args (env processed first)", + input: `{"token": "${env://TOKEN:-default}", "name": "${username}"}`, + envVars: map[string]string{}, + expected: `{"token": "default", "name": "${username}"}`, + }, + { + name: "complex default with special characters", + input: `{"db": "${env://DATABASE_URL:-sqlite:///tmp/default.db?cache=shared&mode=rwc}"}`, + envVars: map[string]string{}, + expected: `{"db": "sqlite:///tmp/default.db?cache=shared&mode=rwc"}`, + }, + { + name: "no env vars in content", + input: `{"normal": "value", "script": "${arg}"}`, + envVars: map[string]string{}, + expected: `{"normal": "value", "script": "${arg}"}`, + }, + { + name: "missing required env var", + input: `{"token": "${env://REQUIRED_TOKEN}"}`, + envVars: map[string]string{}, + expectError: true, + }, + { + name: "multiple missing required env vars", + input: `{"token": "${env://TOKEN1}", "key": "${env://TOKEN2}"}`, + envVars: map[string]string{}, + expectError: true, + }, + { + name: "yaml format", + input: "token: ${env://GITHUB_TOKEN:-default}\ndebug: ${env://DEBUG:-false}", + envVars: map[string]string{"GITHUB_TOKEN": "ghp_456"}, + expected: "token: ghp_456\ndebug: false", + }, + { + name: "env var with underscores and numbers", + input: `{"var": "${env://MY_VAR_123}"}`, + envVars: map[string]string{"MY_VAR_123": "test_value"}, + expected: `{"var": "test_value"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Set up environment variables + originalEnv := make(map[string]string) + for k, v := range tt.envVars { + originalEnv[k] = os.Getenv(k) + os.Setenv(k, v) + } + + // Clean up environment variables after test + defer func() { + for k := range tt.envVars { + if originalValue, existed := originalEnv[k]; existed { + os.Setenv(k, originalValue) + } else { + os.Unsetenv(k) + } + } + }() + + substituter := &EnvSubstituter{} + result, err := substituter.SubstituteEnvVars(tt.input) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error but got none") + } + } else { + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if result != tt.expected { + t.Errorf("Expected:\n%s\nGot:\n%s", tt.expected, result) + } + } + }) + } +} + +func TestArgsSubstituter_SubstituteArgs(t *testing.T) { + tests := []struct { + name string + input string + args map[string]string + expected string + expectError bool + }{ + { + name: "basic args substitution", + input: `{"name": "${username}"}`, + args: map[string]string{"username": "john"}, + expected: `{"name": "john"}`, + }, + { + name: "args with default value used", + input: `{"type": "${repo_type:-public}"}`, + args: map[string]string{}, + expected: `{"type": "public"}`, + }, + { + name: "args with default value overridden", + input: `{"type": "${repo_type:-public}"}`, + args: map[string]string{"repo_type": "private"}, + expected: `{"type": "private"}`, + }, + { + name: "args with empty default", + input: `{"optional": "${optional_arg:-}"}`, + args: map[string]string{}, + expected: `{"optional": ""}`, + }, + { + name: "multiple args in same string", + input: `{"message": "Hello ${name:-World}, you have ${count:-0} messages"}`, + args: map[string]string{"name": "Alice"}, + expected: `{"message": "Hello Alice, you have 0 messages"}`, + }, + { + name: "no args in content", + input: `{"normal": "value", "env": "${env://TOKEN}"}`, + args: map[string]string{}, + expected: `{"normal": "value", "env": "${env://TOKEN}"}`, + }, + { + name: "missing required arg", + input: `{"name": "${required_name}"}`, + args: map[string]string{}, + expectError: true, + }, + { + name: "multiple missing required args", + input: `{"name": "${name}", "id": "${id}"}`, + args: map[string]string{}, + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + substituter := NewArgsSubstituter(tt.args) + result, err := substituter.SubstituteArgs(tt.input) + + if tt.expectError { + if err == nil { + t.Errorf("Expected error but got none") + } + } else { + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if result != tt.expected { + t.Errorf("Expected:\n%s\nGot:\n%s", tt.expected, result) + } + } + }) + } +} + +func TestHasEnvVars(t *testing.T) { + tests := []struct { + name string + content string + expected bool + }{ + { + name: "has env vars", + content: `{"token": "${env://GITHUB_TOKEN}"}`, + expected: true, + }, + { + name: "has env vars with default", + content: `{"debug": "${env://DEBUG:-false}"}`, + expected: true, + }, + { + name: "no env vars", + content: `{"name": "${username}", "normal": "value"}`, + expected: false, + }, + { + name: "empty content", + content: "", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := HasEnvVars(tt.content) + if result != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, result) + } + }) + } +} + +func TestHasScriptArgs(t *testing.T) { + tests := []struct { + name string + content string + expected bool + }{ + { + name: "has script args", + content: `{"name": "${username}"}`, + expected: true, + }, + { + name: "has script args with default", + content: `{"type": "${repo_type:-public}"}`, + expected: true, + }, + { + name: "no script args", + content: `{"token": "${env://GITHUB_TOKEN}", "normal": "value"}`, + expected: false, + }, + { + name: "empty content", + content: "", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := HasScriptArgs(tt.content) + if result != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, result) + } + }) + } +} + +func TestIntegrationEnvAndArgsSubstitution(t *testing.T) { + // Test that env substitution and args substitution work together correctly + input := `{ + "token": "${env://GITHUB_TOKEN:-default_token}", + "user": "${username}", + "repo": "${repo_name:-my-repo}", + "debug": "${env://DEBUG:-false}" + }` + + // Set up environment + os.Setenv("GITHUB_TOKEN", "ghp_real_token") + defer os.Unsetenv("GITHUB_TOKEN") + + // Step 1: Apply env substitution + envSubstituter := &EnvSubstituter{} + afterEnv, err := envSubstituter.SubstituteEnvVars(input) + if err != nil { + t.Fatalf("Env substitution failed: %v", err) + } + + expectedAfterEnv := `{ + "token": "ghp_real_token", + "user": "${username}", + "repo": "${repo_name:-my-repo}", + "debug": "false" + }` + + if afterEnv != expectedAfterEnv { + t.Errorf("After env substitution, expected:\n%s\nGot:\n%s", expectedAfterEnv, afterEnv) + } + + // Step 2: Apply args substitution + args := map[string]string{"username": "alice"} + argsSubstituter := NewArgsSubstituter(args) + final, err := argsSubstituter.SubstituteArgs(afterEnv) + if err != nil { + t.Fatalf("Args substitution failed: %v", err) + } + + expectedFinal := `{ + "token": "ghp_real_token", + "user": "alice", + "repo": "my-repo", + "debug": "false" + }` + + if final != expectedFinal { + t.Errorf("After args substitution, expected:\n%s\nGot:\n%s", expectedFinal, final) + } +} diff --git a/internal/tools/mcp.go b/internal/tools/mcp.go index ea7813f1..c8492e23 100644 --- a/internal/tools/mcp.go +++ b/internal/tools/mcp.go @@ -196,6 +196,15 @@ func (m *MCPToolManager) GetTools() []tool.BaseTool { return m.tools } +// GetLoadedServerNames returns the names of successfully loaded MCP servers +func (m *MCPToolManager) GetLoadedServerNames() []string { + var names []string + for serverName := range m.clients { + names = append(names, serverName) + } + return names +} + // Close closes all MCP clients func (m *MCPToolManager) Close() error { for name, client := range m.clients {