feat: add custom headers and base_url env expansion for providers

agent-schema.json

@@ -106,6 +106,19 @@
           "examples": [
             "CUSTOM_PROVIDER_API_KEY"
           ]
+        },
+        "headers": {
+          "type": "object",
+          "description": "Custom HTTP headers to include in requests. Header values can reference environment variables using ${VAR_NAME} syntax.",
+          "additionalProperties": {
+            "type": "string"
+          },
+          "examples": [
+            {
+              "cf-aig-authorization": "Bearer ${CLOUDFLARE_AI_GATEWAY_TOKEN}",
+              "x-custom-header": "value"
+            }
+          ]
         }
       },
       "required": [
@@ -525,6 +538,19 @@
           "type": "string",
           "description": "Token key for authentication"
         },
+        "headers": {
+          "type": "object",
+          "description": "Custom HTTP headers to include in requests to this model's provider. Header values can reference environment variables using ${VAR_NAME} syntax.",
+          "additionalProperties": {
+            "type": "string"
+          },
+          "examples": [
+            {
+              "cf-aig-authorization": "Bearer ${CLOUDFLARE_AI_GATEWAY_TOKEN}",
+              "x-custom-header": "value"
+            }
+          ]
+        },
         "provider_opts": {
           "type": "object",
           "description": "Provider-specific options. dmr: runtime_flags. anthropic/amazon-bedrock (Claude): interleaved_thinking (boolean, default true). openai/anthropic/google: rerank_prompt (string) to fully override the system prompt used for RAG reranking (advanced - prefer using results.reranking.criteria for domain-specific guidance).",

examples/custom_provider.yaml

@@ -6,45 +6,31 @@
 
 # Define custom providers with reusable configuration
 providers:
-  # Example: A custom OpenAI Chat Completions compatible API gateway
-  my_gateway:
-    api_type: openai_chatcompletions  # Use the Chat Completions API schema
-    base_url: https://api.example.com/
-    token_key: API_KEY_ENV_VAR_NAME  # Environment variable containing the API token
 
-  # Example: A custom OpenAI Responses compatible API gateway
-  responses_provider:
-    api_type: openai_responses
-    base_url: https://responses.example.com/
-    token_key: API_KEY_ENV_VAR_NAME
+  # Example: Cloudflare AI Gateway with custom headers
+  cloudflare_gateway:
+    api_type: openai_chatcompletions
+    base_url: https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/compat
+    token_key: GOOGLE_API_KEY  # Standard Authorization header for provider auth
+    headers:
+      # Custom header for gateway authentication with environment variable expansion
+      cf-aig-authorization: Bearer ${CLOUDFLARE_AI_GATEWAY_TOKEN}
 
 # Define models that use the custom providers
 models:
-  # Model using the custom gateway provider
-  gateway_gpt4o:
-    provider: my_gateway
-    model: gpt-4o
-    max_tokens: 32768
-    temperature: 0.7
 
-  # Model using the responses provider
-  responses_model:
-    provider: responses_provider
-    model: gpt-5
-    max_tokens: 16000
+  # Model using Cloudflare AI Gateway with custom headers
+  gemini_via_cloudflare:
+    provider: cloudflare_gateway
+    model: google-ai-studio/gemini-3-flash-preview
+    max_tokens: 8000
+    temperature: 0.7
 
 # Define agents that use the models
 agents:
   root:
-    model: responses_model
+    model: gemini_via_cloudflare
     description: Main assistant using the custom gateway
     instruction: |
       You are a helpful AI assistant. Be concise and helpful in your responses.
 
-  # Example using shorthand syntax: provider_name/model_name
-  # The provider defaults (base_url, token_key, api_type) are automatically applied
-  subagent:
-    model: my_gateway/gpt-4o-mini
-    description: Sub-agent for specialized tasks
-    instruction: |
-      You are a specialized assistant for specific tasks.

pkg/config/gather.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"maps"
 	"os"
+	"regexp"
 	"slices"
 	"strings"
 
@@ -122,6 +123,36 @@ func addEnvVarsForModelConfig(model *latest.ModelConfig, customProviders map[str
 			}
 		}
 	}
+
+	// Gather env vars from headers (model-level and provider-level)
+	gatherEnvVarsFromHeaders(model.Headers, requiredEnv)
+	if customProviders != nil {
+		if provCfg, exists := customProviders[model.Provider]; exists {
+			gatherEnvVarsFromHeaders(provCfg.Headers, requiredEnv)
+			gatherEnvVarsFromString(provCfg.BaseURL, requiredEnv)
+		}
+	}
+}
+
+// envVarPattern matches ${VAR} and $VAR references in strings.
+var envVarPattern = regexp.MustCompile(`\$\{([^}]+)\}|\$([A-Za-z_][A-Za-z0-9_]*)`)
+
+// gatherEnvVarsFromHeaders extracts environment variable names referenced in header values.
+func gatherEnvVarsFromHeaders(headers map[string]string, requiredEnv map[string]bool) {
+	for _, value := range headers {
+		gatherEnvVarsFromString(value, requiredEnv)
+	}
+}
+
+// gatherEnvVarsFromString extracts environment variable names from a string containing $VAR or ${VAR}.
+func gatherEnvVarsFromString(s string, requiredEnv map[string]bool) {
+	for _, match := range envVarPattern.FindAllStringSubmatch(s, -1) {
+		if match[1] != "" {
+			requiredEnv[match[1]] = true
+		} else if match[2] != "" {
+			requiredEnv[match[2]] = true
+		}
+	}
 }
 
 func GatherEnvVarsForTools(ctx context.Context, cfg *latest.Config) ([]string, error) {

pkg/config/latest/types.go

@@ -135,6 +135,9 @@ type ProviderConfig struct {
 	BaseURL string `json:"base_url"`
 	// TokenKey is the environment variable name containing the API token
 	TokenKey string `json:"token_key,omitempty"`
+	// Headers allows custom HTTP headers to be included in requests.
+	// Header values can reference environment variables using ${VAR_NAME} syntax.
+	Headers map[string]string `json:"headers,omitempty"`
 }
 
 // FallbackConfig represents fallback model configuration for an agent.
@@ -392,6 +395,9 @@ type ModelConfig struct {
 	BaseURL           string   `json:"base_url,omitempty"`
 	ParallelToolCalls *bool    `json:"parallel_tool_calls,omitempty"`
 	TokenKey          string   `json:"token_key,omitempty"`
+	// Headers allows custom HTTP headers to be included in requests to this model's provider.
+	// Header values can reference environment variables using ${VAR_NAME} syntax.
+	Headers map[string]string `json:"headers,omitempty"`
 	// ProviderOpts allows provider-specific options.
 	ProviderOpts map[string]any `json:"provider_opts,omitempty"`
 	TrackUsage   *bool          `json:"track_usage,omitempty"`

pkg/config/v3/types.go

@@ -34,6 +34,9 @@ type ProviderConfig struct {
 	BaseURL string `json:"base_url"`
 	// TokenKey is the environment variable name containing the API token
 	TokenKey string `json:"token_key,omitempty"`
+	// Headers allows custom HTTP headers to be included in requests.
+	// Header values can reference environment variables using ${VAR_NAME} syntax.
+	Headers map[string]string `json:"headers,omitempty"`
 }
 
 // AgentConfig represents a single agent configuration
@@ -70,6 +73,8 @@ type ModelConfig struct {
 	BaseURL           string   `json:"base_url,omitempty"`
 	ParallelToolCalls *bool    `json:"parallel_tool_calls,omitempty"`
 	TokenKey          string   `json:"token_key,omitempty"`
+	// Headers allows custom HTTP headers to be included in requests.
+	Headers map[string]string `json:"headers,omitempty"`
 	// ProviderOpts allows provider-specific options. Currently used for "dmr" provider only.
 	ProviderOpts map[string]any `json:"provider_opts,omitempty"`
 	TrackUsage   *bool          `json:"track_usage,omitempty"`

pkg/model/provider/anthropic/client.go

@@ -153,8 +153,43 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
 			option.WithHTTPClient(httpclient.NewHTTPClient()),
 		}
 		if cfg.BaseURL != "" {
-			requestOptions = append(requestOptions, option.WithBaseURL(cfg.BaseURL))
+			expandedBaseURL, err := environment.Expand(ctx, cfg.BaseURL, env)
+			if err != nil {
+				return nil, fmt.Errorf("expanding base_url: %w", err)
+			}
+			requestOptions = append(requestOptions, option.WithBaseURL(expandedBaseURL))
+		}
+
+		// Apply custom headers from provider config if present
+		if cfg.ProviderOpts != nil {
+			if headers, exists := cfg.ProviderOpts["headers"]; exists {
+				headersMap := make(map[string]string)
+				switch h := headers.(type) {
+				case map[string]string:
+					headersMap = h
+				case map[interface{}]interface{}:
+					for k, v := range h {
+						keyStr, okKey := k.(string)
+						valStr, okVal := v.(string)
+						if !okKey || !okVal {
+							return nil, fmt.Errorf("invalid header key/value type: key=%T, value=%T", k, v)
+						}
+						headersMap[keyStr] = valStr
+					}
+				default:
+					return nil, fmt.Errorf("invalid headers configuration: expected map[string]string, got %T", headers)
+				}
+				for key, value := range headersMap {
+					expandedValue, err := environment.Expand(ctx, value, env)
+					if err != nil {
+						return nil, fmt.Errorf("expanding header %s: %w", key, err)
+					}
+					requestOptions = append(requestOptions, option.WithHeader(key, expandedValue))
+					slog.Debug("Applied custom header", "header", key, "provider", cfg.Provider)
+				}
+			}
 		}
+
 		client := anthropic.NewClient(requestOptions...)
 		anthropicClient.clientFn = func(context.Context) (anthropic.Client, error) {
 			return client, nil

pkg/model/provider/custom_headers_test.go

@@ -0,0 +1,135 @@
+package provider
+
+import (
+	"testing"
+
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestApplyProviderDefaults_WithHeaders(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name             string
+		providerName     string
+		providerCfg      latest.ProviderConfig
+		modelCfg         latest.ModelConfig
+		expectedHeaders  map[string]string
+		headersInOpts    bool
+	}{
+		{
+			name:         "custom provider with headers",
+			providerName: "custom",
+			providerCfg: latest.ProviderConfig{
+				BaseURL: "https://gateway.example.com/v1",
+				Headers: map[string]string{
+					"cf-aig-authorization": "Bearer token123",
+					"x-custom-header":      "value",
+				},
+			},
+			modelCfg: latest.ModelConfig{
+				Provider: "custom",
+				Model:    "gpt-4o",
+			},
+			expectedHeaders: map[string]string{
+				"cf-aig-authorization": "Bearer token123",
+				"x-custom-header":      "value",
+			},
+			headersInOpts: true,
+		},
+		{
+			name:         "custom provider without headers",
+			providerName: "custom",
+			providerCfg: latest.ProviderConfig{
+				BaseURL: "https://api.example.com/v1",
+			},
+			modelCfg: latest.ModelConfig{
+				Provider: "custom",
+				Model:    "gpt-4o",
+			},
+			headersInOpts: false,
+		},
+		{
+			name:         "custom provider with empty headers",
+			providerName: "custom",
+			providerCfg: latest.ProviderConfig{
+				BaseURL: "https://api.example.com/v1",
+				Headers: map[string]string{},
+			},
+			modelCfg: latest.ModelConfig{
+				Provider: "custom",
+				Model:    "gpt-4o",
+			},
+			headersInOpts: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			providers := map[string]latest.ProviderConfig{
+				tt.providerName: tt.providerCfg,
+			}
+
+			result := applyProviderDefaults(&tt.modelCfg, providers)
+			require.NotNil(t, result)
+
+			if tt.headersInOpts {
+				require.NotNil(t, result.ProviderOpts, "ProviderOpts should not be nil")
+				headers, ok := result.ProviderOpts["headers"]
+				require.True(t, ok, "headers should be in ProviderOpts")
+
+				headerMap, ok := headers.(map[string]string)
+				require.True(t, ok, "headers should be map[string]string")
+				assert.Equal(t, tt.expectedHeaders, headerMap, "headers should match")
+			} else {
+				if result.ProviderOpts != nil {
+					_, hasHeaders := result.ProviderOpts["headers"]
+					assert.False(t, hasHeaders, "headers should not be in ProviderOpts")
+				}
+			}
+		})
+	}
+}
+
+func TestApplyProviderDefaults_HeadersDoNotOverrideExisting(t *testing.T) {
+	t.Parallel()
+
+	providerCfg := latest.ProviderConfig{
+		BaseURL: "https://gateway.example.com/v1",
+		Headers: map[string]string{
+			"x-provider-header": "from-provider",
+		},
+	}
+
+	modelCfg := latest.ModelConfig{
+		Provider: "custom",
+		Model:    "gpt-4o",
+		ProviderOpts: map[string]any{
+			"headers": map[string]string{
+				"x-model-header": "from-model",
+			},
+		},
+	}
+
+	providers := map[string]latest.ProviderConfig{
+		"custom": providerCfg,
+	}
+
+	result := applyProviderDefaults(&modelCfg, providers)
+	require.NotNil(t, result)
+
+	// Model config's headers should take precedence (not be overwritten)
+	require.NotNil(t, result.ProviderOpts)
+	headers, ok := result.ProviderOpts["headers"]
+	require.True(t, ok)
+
+	headerMap, ok := headers.(map[string]string)
+	require.True(t, ok)
+
+	// Should have model's header, not provider's header
+	assert.Equal(t, map[string]string{"x-model-header": "from-model"}, headerMap)
+}