diff --git a/pkg/workflow/awf_helpers.go b/pkg/workflow/awf_helpers.go index 0de7049a1d6..fe84d182aba 100644 --- a/pkg/workflow/awf_helpers.go +++ b/pkg/workflow/awf_helpers.go @@ -190,6 +190,21 @@ func BuildAWFArgs(config AWFCommandConfig) []string { awfArgs = append(awfArgs, "--enable-api-proxy") awfHelpersLog.Print("Added --enable-api-proxy for LLM API proxying") + // Add custom API targets if configured in engine.env + // This allows AWF's credential isolation and firewall to work with custom endpoints + // (e.g., corporate LLM routers, Azure OpenAI, self-hosted APIs) + openaiTarget := extractAPITargetHost(config.WorkflowData, "OPENAI_BASE_URL") + if openaiTarget != "" { + awfArgs = append(awfArgs, "--openai-api-target", openaiTarget) + awfHelpersLog.Printf("Added --openai-api-target=%s", openaiTarget) + } + + anthropicTarget := extractAPITargetHost(config.WorkflowData, "ANTHROPIC_BASE_URL") + if anthropicTarget != "" { + awfArgs = append(awfArgs, "--anthropic-api-target", anthropicTarget) + awfHelpersLog.Printf("Added --anthropic-api-target=%s", anthropicTarget) + } + // Add SSL Bump support for HTTPS content inspection (v0.9.0+) sslBumpArgs := getSSLBumpArgs(firewallConfig) awfArgs = append(awfArgs, sslBumpArgs...) @@ -245,3 +260,68 @@ func WrapCommandInShell(command string) string { // Wrap in shell invocation return fmt.Sprintf("/bin/bash -c '%s'", escapedCommand) } + +// extractAPITargetHost extracts the hostname from a custom API base URL in engine.env. +// This supports custom OpenAI-compatible or Anthropic-compatible endpoints (e.g., internal +// LLM routers, Azure OpenAI) while preserving AWF's credential isolation and firewall features. +// +// The function: +// 1. Checks if the specified env var (e.g., "OPENAI_BASE_URL") exists in engine.env +// 2. Extracts the hostname from the URL (e.g., "https://llm-router.internal.example.com/v1" → "llm-router.internal.example.com") +// 3. Returns empty string if no custom URL is configured or if the URL is invalid +// +// Parameters: +// - workflowData: The workflow data containing engine configuration +// - envVar: The environment variable name (e.g., "OPENAI_BASE_URL", "ANTHROPIC_BASE_URL") +// +// Returns: +// - string: The hostname to use as --openai-api-target or --anthropic-api-target, or empty string if not configured +// +// Example: +// +// engine: +// id: codex +// env: +// OPENAI_BASE_URL: "https://llm-router.internal.example.com/v1" +// OPENAI_API_KEY: ${{ secrets.LLM_ROUTER_KEY }} +// +// extractAPITargetHost(workflowData, "OPENAI_BASE_URL") +// // Returns: "llm-router.internal.example.com" +func extractAPITargetHost(workflowData *WorkflowData, envVar string) string { + // Check if engine config and env are available + if workflowData == nil || workflowData.EngineConfig == nil || workflowData.EngineConfig.Env == nil { + return "" + } + + // Get the custom base URL from engine.env + baseURL, exists := workflowData.EngineConfig.Env[envVar] + if !exists || baseURL == "" { + return "" + } + + // Extract hostname from URL + // URLs can be: + // - "https://llm-router.internal.example.com/v1" → "llm-router.internal.example.com" + // - "http://localhost:8080/v1" → "localhost:8080" + // - "api.openai.com" → "api.openai.com" (treated as hostname) + + // Remove protocol prefix if present + host := baseURL + if idx := strings.Index(host, "://"); idx != -1 { + host = host[idx+3:] + } + + // Remove path suffix if present (everything after first /) + if idx := strings.Index(host, "/"); idx != -1 { + host = host[:idx] + } + + // Validate that we have a non-empty hostname + if host == "" { + awfHelpersLog.Printf("Invalid %s URL (no hostname): %s", envVar, baseURL) + return "" + } + + awfHelpersLog.Printf("Extracted API target host from %s: %s", envVar, host) + return host +} diff --git a/pkg/workflow/awf_helpers_test.go b/pkg/workflow/awf_helpers_test.go new file mode 100644 index 00000000000..5547c818386 --- /dev/null +++ b/pkg/workflow/awf_helpers_test.go @@ -0,0 +1,310 @@ +//go:build !integration + +package workflow + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestExtractAPITargetHost tests the extractAPITargetHost function that extracts +// hostnames from custom API base URLs in engine.env +func TestExtractAPITargetHost(t *testing.T) { + tests := []struct { + name string + workflowData *WorkflowData + envVar string + expected string + }{ + { + name: "extracts hostname from HTTPS URL with path", + workflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ + Env: map[string]string{ + "OPENAI_BASE_URL": "https://llm-router.internal.example.com/v1", + }, + }, + }, + envVar: "OPENAI_BASE_URL", + expected: "llm-router.internal.example.com", + }, + { + name: "extracts hostname from HTTP URL with port and path", + workflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ + Env: map[string]string{ + "ANTHROPIC_BASE_URL": "http://localhost:8080/v1", + }, + }, + }, + envVar: "ANTHROPIC_BASE_URL", + expected: "localhost:8080", + }, + { + name: "handles hostname without protocol or path", + workflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ + Env: map[string]string{ + "OPENAI_BASE_URL": "api.openai.com", + }, + }, + }, + envVar: "OPENAI_BASE_URL", + expected: "api.openai.com", + }, + { + name: "handles hostname with port but no protocol", + workflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ + Env: map[string]string{ + "OPENAI_BASE_URL": "localhost:8000", + }, + }, + }, + envVar: "OPENAI_BASE_URL", + expected: "localhost:8000", + }, + { + name: "returns empty string when env var not set", + workflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ + Env: map[string]string{ + "OTHER_VAR": "value", + }, + }, + }, + envVar: "OPENAI_BASE_URL", + expected: "", + }, + { + name: "returns empty string when engine config is nil", + workflowData: &WorkflowData{ + EngineConfig: nil, + }, + envVar: "OPENAI_BASE_URL", + expected: "", + }, + { + name: "returns empty string when workflow data is nil", + workflowData: nil, + envVar: "OPENAI_BASE_URL", + expected: "", + }, + { + name: "returns empty string for empty URL", + workflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ + Env: map[string]string{ + "OPENAI_BASE_URL": "", + }, + }, + }, + envVar: "OPENAI_BASE_URL", + expected: "", + }, + { + name: "extracts Azure OpenAI endpoint hostname", + workflowData: &WorkflowData{ + EngineConfig: &EngineConfig{ + Env: map[string]string{ + "OPENAI_BASE_URL": "https://my-resource.openai.azure.com/openai/deployments/gpt-4", + }, + }, + }, + envVar: "OPENAI_BASE_URL", + expected: "my-resource.openai.azure.com", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := extractAPITargetHost(tt.workflowData, tt.envVar) + assert.Equal(t, tt.expected, result, "Extracted hostname should match expected value") + }) + } +} + +// TestAWFCustomAPITargetFlags tests that BuildAWFArgs includes custom API target flags +// when OPENAI_BASE_URL or ANTHROPIC_BASE_URL are configured in engine.env +func TestAWFCustomAPITargetFlags(t *testing.T) { + t.Run("includes openai-api-target flag when OPENAI_BASE_URL is configured", func(t *testing.T) { + workflowData := &WorkflowData{ + Name: "test-workflow", + EngineConfig: &EngineConfig{ + ID: "codex", + Env: map[string]string{ + "OPENAI_BASE_URL": "https://llm-router.internal.example.com/v1", + "OPENAI_API_KEY": "${{ secrets.LLM_ROUTER_KEY }}", + }, + }, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{ + Enabled: true, + }, + }, + } + + config := AWFCommandConfig{ + EngineName: "codex", + WorkflowData: workflowData, + AllowedDomains: "github.com", + } + + args := BuildAWFArgs(config) + argsStr := strings.Join(args, " ") + + assert.Contains(t, argsStr, "--openai-api-target", "Should include --openai-api-target flag") + assert.Contains(t, argsStr, "llm-router.internal.example.com", "Should include custom hostname") + }) + + t.Run("includes anthropic-api-target flag when ANTHROPIC_BASE_URL is configured", func(t *testing.T) { + workflowData := &WorkflowData{ + Name: "test-workflow", + EngineConfig: &EngineConfig{ + ID: "claude", + Env: map[string]string{ + "ANTHROPIC_BASE_URL": "https://claude-proxy.internal.company.com", + "ANTHROPIC_API_KEY": "${{ secrets.CLAUDE_PROXY_KEY }}", + }, + }, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{ + Enabled: true, + }, + }, + } + + config := AWFCommandConfig{ + EngineName: "claude", + WorkflowData: workflowData, + AllowedDomains: "github.com", + } + + args := BuildAWFArgs(config) + argsStr := strings.Join(args, " ") + + assert.Contains(t, argsStr, "--anthropic-api-target", "Should include --anthropic-api-target flag") + assert.Contains(t, argsStr, "claude-proxy.internal.company.com", "Should include custom hostname") + }) + + t.Run("does not include api-target flags when using default URLs", func(t *testing.T) { + workflowData := &WorkflowData{ + Name: "test-workflow", + EngineConfig: &EngineConfig{ + ID: "codex", + // No custom OPENAI_BASE_URL + }, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{ + Enabled: true, + }, + }, + } + + config := AWFCommandConfig{ + EngineName: "codex", + WorkflowData: workflowData, + AllowedDomains: "github.com", + } + + args := BuildAWFArgs(config) + argsStr := strings.Join(args, " ") + + assert.NotContains(t, argsStr, "--openai-api-target", "Should not include --openai-api-target when not configured") + assert.NotContains(t, argsStr, "--anthropic-api-target", "Should not include --anthropic-api-target when not configured") + }) + + t.Run("includes both api-target flags when both are configured", func(t *testing.T) { + workflowData := &WorkflowData{ + Name: "test-workflow", + EngineConfig: &EngineConfig{ + ID: "custom", + Env: map[string]string{ + "OPENAI_BASE_URL": "https://openai-proxy.company.com/v1", + "ANTHROPIC_BASE_URL": "https://anthropic-proxy.company.com", + }, + }, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{ + Enabled: true, + }, + }, + } + + config := AWFCommandConfig{ + EngineName: "custom", + WorkflowData: workflowData, + AllowedDomains: "github.com", + } + + args := BuildAWFArgs(config) + argsStr := strings.Join(args, " ") + + assert.Contains(t, argsStr, "--openai-api-target", "Should include --openai-api-target flag") + assert.Contains(t, argsStr, "openai-proxy.company.com", "Should include OpenAI custom hostname") + assert.Contains(t, argsStr, "--anthropic-api-target", "Should include --anthropic-api-target flag") + assert.Contains(t, argsStr, "anthropic-proxy.company.com", "Should include Anthropic custom hostname") + }) +} + +// TestEngineExecutionWithCustomAPITarget tests that engine execution steps include +// custom API target flags when configured in engine.env +func TestEngineExecutionWithCustomAPITarget(t *testing.T) { + t.Run("Codex engine includes openai-api-target flag when OPENAI_BASE_URL is configured", func(t *testing.T) { + workflowData := &WorkflowData{ + Name: "test-workflow", + EngineConfig: &EngineConfig{ + ID: "codex", + Env: map[string]string{ + "OPENAI_BASE_URL": "https://llm-router.internal.example.com/v1", + "OPENAI_API_KEY": "${{ secrets.LLM_ROUTER_KEY }}", + }, + }, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{ + Enabled: true, + }, + }, + } + + engine := NewCodexEngine() + steps := engine.GetExecutionSteps(workflowData, "test.log") + + assert.NotEmpty(t, steps, "Should generate execution steps") + + stepContent := strings.Join(steps[0], "\n") + + assert.Contains(t, stepContent, "--openai-api-target", "Should include --openai-api-target flag") + assert.Contains(t, stepContent, "llm-router.internal.example.com", "Should include custom hostname") + }) + + t.Run("Claude engine includes anthropic-api-target flag when ANTHROPIC_BASE_URL is configured", func(t *testing.T) { + workflowData := &WorkflowData{ + Name: "test-workflow", + EngineConfig: &EngineConfig{ + ID: "claude", + Env: map[string]string{ + "ANTHROPIC_BASE_URL": "https://claude-proxy.internal.company.com", + "ANTHROPIC_API_KEY": "${{ secrets.CLAUDE_PROXY_KEY }}", + }, + }, + NetworkPermissions: &NetworkPermissions{ + Firewall: &FirewallConfig{ + Enabled: true, + }, + }, + } + + engine := NewClaudeEngine() + steps := engine.GetExecutionSteps(workflowData, "test.log") + + assert.NotEmpty(t, steps, "Should generate execution steps") + + stepContent := strings.Join(steps[0], "\n") + + assert.Contains(t, stepContent, "--anthropic-api-target", "Should include --anthropic-api-target flag") + assert.Contains(t, stepContent, "claude-proxy.internal.company.com", "Should include custom hostname") + }) +}