diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a6269e..9ce98dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ IMPROVEMENTS * Set custom User-Agent header for TFE API requests to enable tracking MCP server usage separately from other go-tfe clients [268](https://github.com/hashicorp/terraform-mcp-server/pull/268) * Adding a new cli flags `--log-level` to set the desired log level for the server logs and `--log-format` for the logs formatting [286](https://github.com/hashicorp/terraform-mcp-server/pull/286) +* Add `--pagination-limit` flag and `MCP_PAGINATION_LIMIT` env var to configure page sizes for list operations FIXES diff --git a/cmd/terraform-mcp-server/config_test.go b/cmd/terraform-mcp-server/config_test.go index 8206320..3c81c40 100644 --- a/cmd/terraform-mcp-server/config_test.go +++ b/cmd/terraform-mcp-server/config_test.go @@ -325,6 +325,122 @@ func TestInitLoggerWithFormat(t *testing.T) { } } +func TestGetPaginationLimit(t *testing.T) { + tests := []struct { + name string + envValue string + flagValue int + expected int + description string + }{ + { + name: "env var takes precedence", + envValue: "50", + flagValue: 100, + expected: 50, + description: "MCP_PAGINATION_LIMIT env var should override --pagination-limit flag", + }, + { + name: "flag used when env not set", + envValue: "", + flagValue: 25, + expected: 25, + description: "--pagination-limit flag should be used when MCP_PAGINATION_LIMIT is not set", + }, + { + name: "default when neither set", + envValue: "", + flagValue: 0, + expected: 0, + description: "should default to 0 when neither env nor flag is set", + }, + { + name: "invalid env falls back to default", + envValue: "invalid", + flagValue: 0, + expected: 0, + description: "invalid MCP_PAGINATION_LIMIT should fall back to default", + }, + { + name: "negative env falls back to default", + envValue: "-5", + flagValue: 0, + expected: 0, + description: "negative MCP_PAGINATION_LIMIT should fall back to default", + }, + { + name: "negative flag falls back to default", + envValue: "", + flagValue: -10, + expected: 0, + description: "negative --pagination-limit should fall back to default", + }, + { + name: "valid env overrides negative flag", + envValue: "100", + flagValue: -5, + expected: 100, + description: "valid MCP_PAGINATION_LIMIT should be used even if flag is negative", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Save and restore original env var + originalEnv := os.Getenv("MCP_PAGINATION_LIMIT") + defer func() { + if originalEnv != "" { + os.Setenv("MCP_PAGINATION_LIMIT", originalEnv) + } else { + os.Unsetenv("MCP_PAGINATION_LIMIT") + } + }() + + // Set up test environment + if tt.envValue != "" { + os.Setenv("MCP_PAGINATION_LIMIT", tt.envValue) + } else { + os.Unsetenv("MCP_PAGINATION_LIMIT") + } + + // Create a test command with the flag + cmd := &cobra.Command{} + cmd.Flags().Int("pagination-limit", tt.flagValue, "test flag") + + // Test the function + limit := getPaginationLimit(cmd) + assert.Equal(t, tt.expected, limit, tt.description) + }) + } +} + +func TestGetPaginationLimitWithNilCommand(t *testing.T) { + // Save and restore original env var + originalEnv := os.Getenv("MCP_PAGINATION_LIMIT") + defer func() { + if originalEnv != "" { + os.Setenv("MCP_PAGINATION_LIMIT", originalEnv) + } else { + os.Unsetenv("MCP_PAGINATION_LIMIT") + } + }() + + // Test with nil command and no env var + os.Unsetenv("MCP_PAGINATION_LIMIT") + limit := getPaginationLimit(nil) + assert.Equal(t, 0, limit, "expected default 0 with nil command") + + // Test with nil command but env var set + os.Setenv("MCP_PAGINATION_LIMIT", "75") + limit = getPaginationLimit(nil) + assert.Equal(t, 75, limit, "expected 75 from env var with nil command") + + // Test with nil command and invalid env var + os.Setenv("MCP_PAGINATION_LIMIT", "not-a-number") + limit = getPaginationLimit(nil) + assert.Equal(t, 0, limit, "expected default 0 with invalid env var and nil command") +} + func TestGetLogFormat(t *testing.T) { tests := []struct { name string diff --git a/cmd/terraform-mcp-server/init.go b/cmd/terraform-mcp-server/init.go index 700121d..dafd1c8 100644 --- a/cmd/terraform-mcp-server/init.go +++ b/cmd/terraform-mcp-server/init.go @@ -11,6 +11,7 @@ import ( "net/http" "os" "path" + "strconv" "strings" "time" @@ -51,8 +52,9 @@ var ( } enabledToolsets := getToolsetsFromCmd(cmd.Root(), logger) + paginationLimit := getPaginationLimit(cmd.Root()) - if err := runStdioServer(logger, enabledToolsets); err != nil { + if err := runStdioServer(logger, enabledToolsets, paginationLimit); err != nil { stdlog.Fatal("failed to run stdio server:", err) } }, @@ -89,8 +91,9 @@ var ( } enabledToolsets := getToolsetsFromCmd(cmd.Root(), logger) + paginationLimit := getPaginationLimit(cmd.Root()) - if err := runHTTPServer(logger, host, port, endpointPath, enabledToolsets); err != nil { + if err := runHTTPServer(logger, host, port, endpointPath, enabledToolsets, paginationLimit); err != nil { stdlog.Fatal("failed to run streamableHTTP server:", err) } }, @@ -117,6 +120,7 @@ func init() { rootCmd.PersistentFlags().String("log-format", "text", "Log format (text or json)") rootCmd.PersistentFlags().String("toolsets", "all", toolsets.GenerateToolsetsHelp()) rootCmd.PersistentFlags().String("tools", "", toolsets.GenerateToolsHelp()) + rootCmd.PersistentFlags().Int("pagination-limit", 0, "Pagination limit for list operations (0 = use default)") // Add StreamableHTTP command flags (avoid 'h' shorthand conflict with help) streamableHTTPCmd.Flags().String("transport-host", "127.0.0.1", "Host to bind to") @@ -192,6 +196,36 @@ func getLogFormat(cmd *cobra.Command) string { return "text" } +// getPaginationLimit determines the pagination limit from environment variable or CLI flag +func getPaginationLimit(cmd *cobra.Command) int { + // Check environment variable first + if envLimit := os.Getenv("MCP_PAGINATION_LIMIT"); envLimit != "" { + limit, err := strconv.Atoi(envLimit) + if err != nil { + stdlog.Printf("Warning: invalid MCP_PAGINATION_LIMIT '%s', using default\n", envLimit) + return 0 + } + if limit < 0 { + stdlog.Printf("Warning: MCP_PAGINATION_LIMIT cannot be negative, using default\n") + return 0 + } + return limit + } + + // Check CLI flag + if cmd != nil { + if limit, err := cmd.Flags().GetInt("pagination-limit"); err == nil { + if limit < 0 { + stdlog.Printf("Warning: --pagination-limit cannot be negative, using default\n") + return 0 + } + return limit + } + } + + return 0 +} + func initLogger(outPath string, level log.Level, format string) (*log.Logger, error) { logger := log.New() logger.SetLevel(level) diff --git a/cmd/terraform-mcp-server/main.go b/cmd/terraform-mcp-server/main.go index 6095a0d..b5943aa 100644 --- a/cmd/terraform-mcp-server/main.go +++ b/cmd/terraform-mcp-server/main.go @@ -25,27 +25,27 @@ import ( //go:embed instructions.md var instructions string -func runHTTPServer(logger *log.Logger, host string, port string, endpointPath string, enabledToolsets []string) error { +func runHTTPServer(logger *log.Logger, host string, port string, endpointPath string, enabledToolsets []string, paginationLimit int) error { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - hcServer := NewServer(version.Version, logger, enabledToolsets) + hcServer := NewServer(version.Version, logger, enabledToolsets, paginationLimit) registerToolsAndResources(hcServer, logger, enabledToolsets) return streamableHTTPServerInit(ctx, hcServer, logger, host, port, endpointPath) } -func runStdioServer(logger *log.Logger, enabledToolsets []string) error { +func runStdioServer(logger *log.Logger, enabledToolsets []string, paginationLimit int) error { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - hcServer := NewServer(version.Version, logger, enabledToolsets) + hcServer := NewServer(version.Version, logger, enabledToolsets, paginationLimit) registerToolsAndResources(hcServer, logger, enabledToolsets) return serverInit(ctx, hcServer, logger) } -func NewServer(version string, logger *log.Logger, enabledToolsets []string, opts ...server.ServerOption) *server.MCPServer { +func NewServer(version string, logger *log.Logger, enabledToolsets []string, paginationLimit int, opts ...server.ServerOption) *server.MCPServer { // Create rate limiting middleware with environment-based configuration rateLimitConfig := client.LoadRateLimitConfigFromEnv() rateLimitMiddleware := client.NewRateLimitMiddleware(rateLimitConfig, logger) @@ -58,6 +58,12 @@ func NewServer(version string, logger *log.Logger, enabledToolsets []string, opt server.WithToolHandlerMiddleware(rateLimitMiddleware.Middleware()), server.WithElicitation(), } + + if paginationLimit > 0 { + defaultOpts = append(defaultOpts, server.WithPaginationLimit(paginationLimit)) + logger.Infof("Pagination limit set to %d", paginationLimit) + } + opts = append(defaultOpts, opts...) // Create hooks for session management @@ -164,8 +170,9 @@ func runDefaultCommand(cmd *cobra.Command, _ []string) { // Get toolsets from the command that was passed in enabledToolsets := getToolsetsFromCmd(cmd, logger) + paginationLimit := getPaginationLimit(cmd) - if err := runStdioServer(logger, enabledToolsets); err != nil { + if err := runStdioServer(logger, enabledToolsets, paginationLimit); err != nil { stdlog.Fatal("failed to run stdio server:", err) } } @@ -186,8 +193,9 @@ func main() { } enabledToolsets := getToolsetsFromCmd(rootCmd, logger) + paginationLimit := getPaginationLimit(rootCmd) - if err := runHTTPServer(logger, host, port, endpointPath, enabledToolsets); err != nil { + if err := runHTTPServer(logger, host, port, endpointPath, enabledToolsets, paginationLimit); err != nil { stdlog.Fatal("failed to run StreamableHTTP server:", err) } return