Skip to content

feat: add custom headers and base_url env expansion for providers#2108

Open
alindsilva wants to merge 13 commits intodocker:mainfrom
alindsilva:feature/provider-custom-headers-pr
Open

feat: add custom headers and base_url env expansion for providers#2108
alindsilva wants to merge 13 commits intodocker:mainfrom
alindsilva:feature/provider-custom-headers-pr

Conversation

@alindsilva
Copy link
Copy Markdown

Summary

Add support for custom HTTP headers on provider and model configurations, enabling use cases like API gateways (e.g., Cloudflare AI Gateway) that require additional authentication or routing headers.

Changes

1. Custom headers in config (ee35e1a9)

  • Add headers field (map[string]string) to both ProviderConfig and ModelConfig in the config schema
  • Wire headers from config into ProviderOpts["headers"] in applyProviderDefaults()
  • Model-level headers take precedence over provider-level headers
  • Update JSON schema (agent-schema.json)
  • Add tests for headers wiring logic

2. Forward Start/Stop to inner toolsets (19147d98)

  • Fix "toolset not started" errors in multi-agent configurations
  • Add Start()/Stop() forwarding to FilteredToolset, InstructionsToolset, and ToonToolset wrappers
  • Add tests for all three wrapper types

3. Normalize anyOf schemas + API error logging (14eb42df)

  • Add normalizeUnionTypes() to collapse single-element anyOf/oneOf arrays (required by some providers)
  • Wire into ConvertParametersToSchema pipeline
  • Add debug logging of API error response bodies in streaming handler
  • Add schema normalization tests

4. Headers + base_url env expansion on all providers (11167d78)

  • OpenAI: Read custom headers from ProviderOpts, expand ${VAR} env vars, auth middleware for custom providers
  • Gemini: Custom headers via genai.HTTPOptions, base_url ${VAR} expansion
  • Anthropic: Custom headers via option.WithHeader(), base_url ${VAR} expansion
  • Gather env vars from header values and base_url in GatherEnvVarsForModels() so tests auto-set dummy values
  • Fix old module path reference in test file
  • Add example config (examples/custom_provider.yaml)

Example usage

providers:
  cloudflare_gateway:
    api_type: openai_chatcompletions
    base_url: https://gateway.ai.cloudflare.com/v1/acct/gw/compat
    token_key: GOOGLE_API_KEY
    headers:
      cf-aig-authorization: Bearer ${CLOUDFLARE_AI_GATEWAY_TOKEN}

models:
  gemini_via_cloudflare:
    provider: cloudflare_gateway
    model: google-ai-studio/gemini-3-flash-preview

Testing

  • All existing tests pass
  • New tests for: header wiring, schema normalization, Start/Stop forwarding
  • TestLoadExamples passes with the new custom_provider.yaml example

Add Headers map[string]string to ProviderConfig, allowing custom HTTP
headers on provider definitions. Headers flow through ProviderOpts to
the OpenAI client with env var expansion (${VAR_NAME} syntax).

Includes:
- ProviderConfig.Headers field in config schema (v3 and latest)
- Headers wiring in applyProviderDefaults
- OpenAI client: headers parsing, env expansion, auth middleware
  for custom providers without token_key
- Schema normalization (normalizeUnionTypes) for gateway compatibility
- Handle both map[string]string and map[interface{}]interface{} YAML types
The filter, instructions, and toon toolset wrappers were not forwarding
Start() and Stop() calls to their inner toolsets. This caused MCP tools
to fail with 'toolset not started' errors in multi-agent configurations.
- Convert anyOf patterns like {anyOf: [{type:string},{type:null}]} to
  {type:string} for compatibility with AI gateways (e.g. Cloudflare)
  that don't support anyOf in tool parameter schemas.
- Log HTTP response body on non-2xx API errors for easier debugging.
Add custom headers support and ${VAR_NAME} expansion in base_url to the
Gemini and Anthropic provider clients, matching the existing OpenAI
client capability. Also add Headers field directly to ModelConfig for
convenience (no separate providers section needed).

- Gemini: read headers from ProviderOpts, expand env vars, set on
  genai.HTTPOptions; expand env vars in base_url
- Anthropic: same pattern with option.WithHeader; expand env vars
  in base_url
- ModelConfig.Headers: new field merged into ProviderOpts['headers']
  with model-level taking precedence over provider-level
- Updated JSON schema and config types (v3 + latest)
Copilot AI review requested due to automatic review settings March 15, 2026 19:36
@alindsilva alindsilva requested a review from a team as a code owner March 15, 2026 19:36
…_url env expansion for providers

* feature/provider-custom-headers-pr:
  feat: add custom headers and base_url env expansion to all providers
  fix: normalize anyOf schemas and add API error response body logging
  fix: forward Start/Stop to inner toolsets in teamloader wrappers
  feat: add custom headers support for provider configs
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR expands provider/model configurability and robustness across the runtime by adding configurable HTTP headers, environment-variable expansion for certain endpoints, schema normalization for stricter gateways, and lifecycle forwarding for wrapped toolsets.

Changes:

  • Add headers fields to provider/model config (plus JSON schema + example) and wire them through provider defaults into per-provider HTTP client initialization (with env expansion in header values).
  • Fix multi-agent “toolset not started” issues by forwarding Start()/Stop() through toolset wrapper types and adding tests.
  • Normalize union/anyOf tool schemas for gateway compatibility and add additional API error detail logging for streaming failures.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 6 comments.

Show a summary per file
File Description
pkg/teamloader/toon.go Forwards Start/Stop to inner toolset when supported.
pkg/teamloader/toon_test.go Adds tests validating Start/Stop forwarding behavior for Toon wrapper.
pkg/teamloader/instructions.go Forwards Start/Stop to inner toolset for instructions wrapper.
pkg/teamloader/instructions_test.go Adds tests for Start/Stop forwarding behavior in instructions wrapper.
pkg/teamloader/filter.go Forwards Start/Stop to inner toolset for filtered toolsets.
pkg/teamloader/filter_test.go Adds tests for Start/Stop forwarding and startable integration chaining.
pkg/runtime/streaming.go Adds debug logging for OpenAI SDK API error response dumps in streaming loop.
pkg/model/provider/provider.go Wires provider/model headers into ProviderOpts["headers"] with model-level precedence.
pkg/model/provider/custom_headers_test.go Adds tests covering provider/model header wiring into ProviderOpts.
pkg/model/provider/openai/client.go Adds support for custom headers (with env expansion) and changes auth handling for custom providers without token_key.
pkg/model/provider/openai/api_type_test.go Updates expectation to reflect stripping Authorization when no token_key is set.
pkg/model/provider/gemini/client.go Adds base_url env expansion and custom headers via genai.HTTPOptions.
pkg/model/provider/anthropic/client.go Adds base_url env expansion and custom headers via request options.
pkg/model/provider/openai/schema.go Adds normalizeUnionTypes() into schema conversion pipeline.
pkg/model/provider/openai/schema_test.go Adds tests for schema normalization of anyOf optional patterns.
pkg/model/provider/schema_test.go Updates expected schema output to match normalized union handling.
pkg/config/latest/types.go Adds headers to provider/model config types.
pkg/config/v3/types.go Adds headers to provider/model config types (older version).
pkg/config/gather.go Gathers env vars referenced in headers (and provider base_url) for missing-env checks.
agent-schema.json Extends JSON schema with headers fields for provider/model.
examples/custom_provider.yaml Updates example to demonstrate Cloudflare gateway headers usage.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

… feedback

- Add Headers field to ModelConfig and ProviderConfig in v4, v5, v6 config
  types so headers survive the JSON upgrade chain from v3 to latest. This
  was causing Cloudflare AI Gateway 401 errors because custom headers were
  silently dropped at the v3→v4 boundary.

- Address Copilot review comments on PR docker#2108:
  1. Remove response body from streaming error logs to avoid leaking
     sensitive data (pkg/runtime/streaming.go)
  2. Deep-copy provider headers map before merging to avoid mutating
     shared config across models (pkg/model/provider/provider.go)
  3. Gather env vars from model-level base_url in addition to provider
     base_url (pkg/config/gather.go)
  4. Expand env vars in OpenAI/Azure base_url consistently with
     Anthropic/Gemini (pkg/model/provider/openai/client.go)
  5. Redact header values from error logs to prevent credential leaks
     (pkg/model/provider/openai/client.go)
  6. Tighten union type normalization to only collapse nullable patterns
     (exactly 2 options with one being null), preserving non-nullable
     unions (pkg/model/provider/openai/schema.go)
@alindsilva
Copy link
Copy Markdown
Author

Bug fix: Headers dropped during config upgrade chain

The original PR added Headers fields to ModelConfig and ProviderConfig in v3 and latest config types, but missed the intermediate versions (v4, v5, v6) that sit in the upgrade chain.

Since configs are migrated through the chain v3 → v4 → v5 → v6 → latest via CloneThroughJSON (JSON serialize → deserialize), the missing Headers field in v4 caused all custom headers to be silently dropped at the v3→v4 boundary. By the time the config reached the provider clients, ProviderOpts["headers"] was empty — so gateway auth headers like cf-aig-authorization were never sent, resulting in 401 Unauthorized errors from Cloudflare AI Gateway.

Fix (8d15048): Added Headers map[string]string to both ModelConfig and ProviderConfig in v4, v5, and v6 config types. This same commit also addresses all 6 Copilot review comments (see inline replies).

- Fix import ordering (gci) in custom_headers_test.go and openai/client.go
- Replace if-else-if chains with switch statements (gocritic)
- Use maps.Copy instead of manual loops (modernize)
- Fix nil return in normalizeUnionTypes (gocritic)
- Replace else-if with else if in custom_headers_test.go (gocritic)
- Use assert.Empty instead of assert.Equal with empty string (testifylint)
- Remove extra blank line (gofmt)
Sync PR branch with 339 commits from upstream main.
Resolved conflict in provider.go by keeping headers logic and adapting to
new function signature (changed from pointer to value return).
After merge with main, the built-in alias handling section was incomplete.
Fixed by properly structuring the function:
- Custom provider handling (with headers) → early return
- Built-in alias handling (Aliases map lookup)
- General header merging for non-custom providers
- Model defaults and return
The project migrated from Taskfile to mise in commit 2bf2683 (March 2026),
but the devcontainer.json was not updated. This commit:

- Replaces go-task feature with mise-en-dev feature
- Adds custom Dockerfile to fix yarn repository issue
- Enables 'mise cross' and other mise commands in devcontainer

This aligns the development environment with the current build system.
The merge conflict resolution introduced an extra closing brace that
terminated applyProviderDefaults() prematurely. This caused all code
after the custom provider block (alias handling and header merging for
non-custom providers) to become orphaned code outside the function.

This broke configurations using built-in providers (like 'google') with
custom base_url and headers (e.g., Cloudflare AI Gateway), since the
header merging code was unreachable.

Fixed by removing the erroneous closing brace and ensuring proper
nesting of the conditional blocks.
TROUBLESHOOTING NOTES FOR FUTURE MERGES WITH UPSTREAM
======================================================

## Critical Merge Pitfall: Function Scope Bug (April 4, 2026)

**Issue**: Merge conflict in pkg/model/provider/provider.go introduced an extra
closing brace that terminated applyProviderDefaults() prematurely, making all
code after line 351 unreachable (alias handling, header merging for non-custom
providers).

**Symptoms**:
- Headers don't work even though code looks correct
- 401 Unauthorized errors from Cloudflare AI Gateway
- Custom provider headers work, built-in provider headers fail

**Root Cause**:
Incorrect closing brace placement after custom provider block:
  applyModelDefaults(enhancedCfg)
  return enhancedCfg
  }
}  # ← WRONG: This extra } closes the entire function!

if alias, exists := Aliases[cfg.Provider]; exists {
  // This code becomes unreachable!

**Fix**: Remove the extra closing brace - the function should have only ONE
final closing brace at the end.

**Detection**:
- Run: sed -n '290,376p' pkg/model/provider/provider.go | cat -A
- Check indentation: Headers code needs 2-3 tabs, not at column 0
- Test: Custom provider should work, model-level headers should also work

## Custom Headers Feature Architecture

**Config Schema Versions**:
- v7 and earlier: NO Headers field in ModelConfig or ProviderConfig
- latest (unreleased): HAS Headers in both ModelConfig and ProviderConfig
- User configs should OMIT schema_version to use 'latest' automatically
- Setting schema_version: "7" will break headers with 'unknown field' error

**Code Flow**:
1. Config parsing → applyProviderDefaults() in pkg/model/provider/provider.go
2. Custom providers (from providers: section) → headers copied to ProviderOpts
3. Model-level headers merged second (takes precedence over provider headers)
4. Headers stored in enhancedCfg.ProviderOpts["headers"] as map[string]string
5. Each client (openai, anthropic, gemini) reads from ProviderOpts["headers"]
6. Clients expand environment variables () before sending requests

**Key Functions**:
- applyProviderDefaults(): Merges custom provider defaults and model headers
- cloneModelConfig(): Deep copies ModelConfig including ProviderOpts map
- OpenAI/Anthropic/Gemini clients: Read headers from cfg.ProviderOpts["headers"]

## Cloudflare AI Gateway Configuration

**Correct Setup** (using custom provider):
providers:
  google-ai-studio:
    api_type: openai_chatcompletions  # Required for /compat endpoint
    base_url: https://gateway.ai.cloudflare.com/v1/${ACCOUNT}/${GATEWAY}/compat
    headers:
      cf-aig-authorization: Bearer ${CLOUDFLARE_AI_GATEWAY_TOKEN}
      x-goog-api-key: ${GOOGLE_API_KEY}

models:
  gemini_flash:
    provider: google-ai-studio  # References custom provider
    model: google-ai-studio/gemini-3-flash-preview

**Why This Works**:
- Cloudflare /compat endpoint expects OpenAI-compatible API format
- provider: google would use Gemini SDK (native format) - won't work!
- Custom provider with api_type: openai_chatcompletions routes to OpenAI client
- OpenAI client applies headers from ProviderOpts before making requests

**Common Mistakes**:
- Using provider: google instead of custom provider → Gemini SDK used
- Missing api_type: openai_chatcompletions → wrong API format
- Setting schema_version: "7" → headers field rejected
- Undefined environment variables → 401 Unauthorized (empty header values)

## Testing After Merge

1. Verify function structure:
   grep -A60 'func applyProviderDefaults' pkg/model/provider/provider.go

2. Count closing braces (should match):
   sed -n '290,380p' pkg/model/provider/provider.go | grep -c '^[[:space:]]*}$'

3. Test header merging (create temporary test):
   - Custom provider headers should be in ProviderOpts
   - Model-level headers should merge and override provider headers
   - Environment variables should remain as ${VAR} (expanded by clients)

4. Build and run integration test:
   mise build
   ./bin/docker-agent run examples/custom_provider.yaml --debug

5. Check debug logs for:
   - "Applying custom headers"
   - "Applied custom header"
   - Header count should match config

## Build System Notes

- Project migrated from Taskfile to mise on March 28, 2026 (commit 2bf2683)
- DevContainer updated from go-task to mise-en-dev feature
- Use 'mise build' not 'task build'
- Go 1.26.1+ required (use GOTOOLCHAIN=auto if local version too old)

## Upstream Status (as of April 4, 2026)

- docker/docker-agent:main does NOT have Headers field yet
- PR docker#2108 (feature/provider-custom-headers-pr) adds this capability
- Until merge, only this fork supports custom headers
- After merge, remove this commit or keep for historical reference

## Related Commits

- cce9e24: Initial lint fixes for PR
- 60aa6ae: Merged 339 commits from main (March 28 → April 4, 2026)
- 2b0d2fa: First merge conflict resolution (introduced the bug)
- 6d6d986: Fixed function scope bug (removed extra closing brace)
- 8d15048: Added Headers to v4/v5/v6 config types (config upgrade support)

## References

- PR: docker#2108
- Schema: ./agent-schema.json
- Config types: ./pkg/config/latest/types.go
- Provider logic: ./pkg/model/provider/provider.go
- Example config: ./examples/custom_provider.yaml
- Fix gci formatting in provider.go (proper indentation of headers code)
- Fix gocritic else-if in custom_headers_test.go

Note: Upstream lint errors in pkg/teamloader/* are not addressed as
they are not part of the custom headers feature.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants