feat: allow to specify extra args: url, traffic policy, pooled for ngrok (#1345)

Co-authored-by: taskbot <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

cmd/thv/app/proxy_tunnel.go

@@ -36,7 +36,15 @@ Examples:
 
 Flags:
   --tunnel-provider string   The provider to use for the tunnel (e.g., "ngrok") - mandatory
-  --provider-args string     JSON object with provider-specific arguments (default "{}")
+  --provider-args string     JSON object with provider-specific arguments: auth-token (mandatory),
+  							 url, pooling, traffic-policy-file
+  --dry-run                  If set, only validate the configuration without starting the tunnel
+
+Examples:
+  thv proxy tunnel --tunnel-provider ngrok --provider-args '{"auth-token": "your-token",
+  "url": "https://example.com", "pooling": true}' http://localhost:8080 my-server
+  thv proxy tunnel --tunnel-provider ngrok --provider-args '{"auth-token": "your-token",
+  "traffic-policy-file": "/path/to/policy.yml"}' my-workload my-server
 `,
 	Args: cobra.ExactArgs(2),
 	RunE: proxyTunnelCmdFunc,

docs/cli/thv_proxy_tunnel.md

@@ -4,8 +4,12 @@ package ngrok
 import (
 	"context"
 	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
 
 	"golang.ngrok.com/ngrok/v2"
+	"gopkg.in/yaml.v2"
 
 	"github.com/stacklok/toolhive/pkg/logger"
 )
@@ -17,31 +21,67 @@ type TunnelProvider struct {
 
 // TunnelConfig holds configuration options for the ngrok tunnel provider.
 type TunnelConfig struct {
-	AuthToken string
-	Domain    string // Optional: specify custom domain
-	DryRun    bool
+	AuthToken      string
+	URL            string // Optional: specify custom URL
+	TrafficPolicy  string // Optional: specify traffic policy
+	PoolingEnabled bool   // Optional: enable pooling
+	DryRun         bool
+}
+
+// loadTrafficPolicyFile reads a YAML file, ensures it's .yml/.yaml,
+// validates its contents, and returns its text.
+func loadTrafficPolicyFile(path string) (string, error) {
+	ext := strings.ToLower(filepath.Ext(path))
+	if ext != ".yml" && ext != ".yaml" {
+		return "", fmt.Errorf("traffic policy file must be .yml or .yaml, got %q", ext)
+	}
+
+	cleanPath := filepath.Clean(path)
+	b, err := os.ReadFile(cleanPath)
+	if err != nil {
+		return "", fmt.Errorf("reading traffic policy file: %w", err)
+	}
+
+	var tmp any
+	if err := yaml.Unmarshal(b, &tmp); err != nil {
+		return "", fmt.Errorf("invalid YAML in traffic policy file: %w", err)
+	}
+
+	return string(b), nil
 }
 
 // ParseConfig parses the configuration for the ngrok tunnel provider from a map.
 func (p *TunnelProvider) ParseConfig(raw map[string]any) error {
-	token, ok := raw["ngrok-auth-token"].(string)
+	token, ok := raw["auth-token"].(string)
 	if !ok || token == "" {
-		return fmt.Errorf("ngrok-auth-token is required")
+		return fmt.Errorf("auth-token is required")
 	}
 
 	cfg := TunnelConfig{
 		AuthToken: token,
 	}
 
-	if domain, ok := raw["ngrok-domain"].(string); ok {
-		cfg.Domain = domain
+	// optional settings: url, traffic policy, pooling
+	if url, ok := raw["url"].(string); ok {
+		cfg.URL = url
+	}
+	if path, ok := raw["traffic-policy-file"].(string); ok && path != "" {
+		policyText, err := loadTrafficPolicyFile(path)
+		if err != nil {
+			return err
+		}
+		cfg.TrafficPolicy = policyText
+	}
+	if pooling, ok := raw["pooling"].(bool); ok {
+		cfg.PoolingEnabled = pooling
 	}
 
 	p.config = cfg
 
 	if dr, ok := raw["dry-run"].(bool); ok {
 		p.config.DryRun = dr
 	}
+
 	return nil
 }
 
@@ -69,8 +109,14 @@ func (p *TunnelProvider) StartTunnel(ctx context.Context, name, targetURI string
 	endpointOpts := []ngrok.EndpointOption{
 		ngrok.WithDescription("tunnel proxy for " + name),
 	}
-	if p.config.Domain != "" {
-		endpointOpts = append(endpointOpts, ngrok.WithURL(p.config.Domain))
+	if p.config.URL != "" {
+		endpointOpts = append(endpointOpts, ngrok.WithURL(p.config.URL))
+	}
+	if p.config.TrafficPolicy != "" {
+		endpointOpts = append(endpointOpts, ngrok.WithTrafficPolicy(p.config.TrafficPolicy))
+	}
+	if p.config.PoolingEnabled {
+		endpointOpts = append(endpointOpts, ngrok.WithPoolingEnabled(true))
 	}
 
 	forwarder, err := agent.Forward(ctx,

pkg/transport/tunnel/ngrok/tunnel_provider.go

@@ -111,14 +111,14 @@ var _ = Describe("Proxy Tunnel E2E", Serial, func() {
 				"definitely-not-a-workload",
 				proxyServerName,
 				"--tunnel-provider", "ngrok",
-				`--provider-args`, `{"ngrok-auth-token":"dummy","dry-run":true}`,
+				`--provider-args`, `{"auth-token":"dummy","dry-run":true}`,
 			).ExpectFailure()
 
 			// The exact text may vary a bit; cover both likely messages.
 			Expect(stderr).To(MatchRegexp(`failed to get workload|workload .* has empty URL`))
 		})
 
-		It("fails when ngrok args are incorrect (missing ngrok-auth-token)", func() {
+		It("fails when ngrok args are incorrect (missing auth-token)", func() {
 			osvServerURL, err := e2e.GetMCPServerURL(config, osvServerName)
 			Expect(err).ToNot(HaveOccurred())
 			base := mustBaseURL(osvServerURL)
@@ -133,7 +133,7 @@ var _ = Describe("Proxy Tunnel E2E", Serial, func() {
 			).ExpectFailure()
 
 			// ParseConfig should surface this
-			Expect(stderr).To(MatchRegexp(`invalid provider config:.*ngrok-auth-token is required`))
+			Expect(stderr).To(MatchRegexp(`invalid provider config:.*auth-token is required`))
 		})
 	})
 
@@ -144,7 +144,7 @@ var _ = Describe("Proxy Tunnel E2E", Serial, func() {
 			base := mustBaseURL(osvServerURL)
 
 			// Use dry-run to skip real network calls
-			argsJSON := `{"ngrok-auth-token":"dummy-token","dry-run":true}`
+			argsJSON := `{"auth-token":"dummy-token","dry-run":true}`
 
 			By("Starting the proxy tunnel (URL target, dry-run ngrok)")
 			proxyTunnelCmd = startProxyTunnel(config, proxyServerName, base, "ngrok", argsJSON)
@@ -165,7 +165,7 @@ var _ = Describe("Proxy Tunnel E2E", Serial, func() {
 		})
 
 		It("starts a tunnel when target is a workload name", func() {
-			argsJSON := `{"ngrok-auth-token":"dummy-token","dry-run":true}`
+			argsJSON := `{"auth-token":"dummy-token","dry-run":true}`
 
 			By("Starting the proxy tunnel (workload target, dry-run ngrok)")
 			proxyTunnelCmd = startProxyTunnel(config, proxyServerName, osvServerName, "ngrok", argsJSON)