diff --git a/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_embedding_cel_test.go b/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_embedding_cel_test.go index 1512d7adf3..c1a5efc278 100644 --- a/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_embedding_cel_test.go +++ b/cmd/thv-operator/test-integration/virtualmcp/virtualmcpserver_embedding_cel_test.go @@ -5,6 +5,9 @@ package controllers import ( + "fmt" + "strings" + . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -53,4 +56,83 @@ var _ = Describe("CEL Validation for embedding provider on VirtualMCPServer", err := k8sClient.Create(ctx, vmcp) Expect(err).NotTo(HaveOccurred()) }) + + It("should reject embeddingHeaders with the tei provider", func() { + vmcp := newVirtualMCPServerWithOptimizer("vmcp-headers-tei", + &vmcpconfig.OptimizerConfig{ + EmbeddingProvider: "tei", + EmbeddingService: "http://embeddings.example:8080", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "toolhive-optimizer"}, + }) + err := k8sClient.Create(ctx, vmcp) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring( + "embeddingHeaders is only supported when embeddingProvider is 'openai'")) + }) + + It("should reject embeddingHeaders when the provider is defaulted to tei", func() { + vmcp := newVirtualMCPServerWithOptimizer("vmcp-headers-default", + &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://embeddings.example:8080", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "toolhive-optimizer"}, + }) + err := k8sClient.Create(ctx, vmcp) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring( + "embeddingHeaders is only supported when embeddingProvider is 'openai'")) + }) + + It("should accept embeddingHeaders with the openai provider", func() { + vmcp := newVirtualMCPServerWithOptimizer("vmcp-headers-openai", + &vmcpconfig.OptimizerConfig{ + EmbeddingProvider: "openai", + EmbeddingService: "http://gateway.example:8080", + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "toolhive-optimizer"}, + }) + err := k8sClient.Create(ctx, vmcp) + Expect(err).NotTo(HaveOccurred()) + }) + + It("should reject reserved names, invalid names, and unsafe values in embeddingHeaders", func() { + for i, tc := range []struct { + headers map[string]vmcpconfig.EmbeddingHeaderValue + want string + }{ + {map[string]vmcpconfig.EmbeddingHeaderValue{"Authorization": "Bearer x"}, "must not include Authorization or Content-Type"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{"authorization": "Bearer x"}, "must not include Authorization or Content-Type"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{"Content-Type": "application/json"}, "must not include Authorization or Content-Type"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{"content-type": "application/json"}, "must not include Authorization or Content-Type"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{"": "value"}, "names must be valid HTTP header names"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{"x bad": "value"}, "names must be valid HTTP header names"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": ""}, "should be at least 1 chars long"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "a\r\nb"}, "should match"}, + {map[string]vmcpconfig.EmbeddingHeaderValue{ + "x-cache-key": vmcpconfig.EmbeddingHeaderValue(strings.Repeat("a", 8193)), + }, "8192"}, + } { + vmcp := newVirtualMCPServerWithOptimizer(fmt.Sprintf("vmcp-headers-reject-%d", i), + &vmcpconfig.OptimizerConfig{ + EmbeddingProvider: "openai", + EmbeddingService: "http://gateway.example:8080", + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: tc.headers, + }) + err := k8sClient.Create(ctx, vmcp) + Expect(err).To(HaveOccurred(), "headers %v should be rejected", tc.headers) + Expect(err.Error()).To(ContainSubstring(tc.want)) + } + }) + + It("should accept uncommon but valid RFC token characters in header names", func() { + vmcp := newVirtualMCPServerWithOptimizer("vmcp-headers-token-chars", + &vmcpconfig.OptimizerConfig{ + EmbeddingProvider: "openai", + EmbeddingService: "http://gateway.example:8080", + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-key'name`x": "value"}, + }) + err := k8sClient.Create(ctx, vmcp) + Expect(err).NotTo(HaveOccurred()) + }) }) diff --git a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml index 82f66b27dd..20c8b0f8ef 100644 --- a/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml +++ b/deploy/charts/operator-crds/files/crds/toolhive.stacklok.dev_virtualmcpservers.yaml @@ -1810,6 +1810,22 @@ spec: instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. properties: + embeddingHeaders: + additionalProperties: + description: |- + EmbeddingHeaderValue is a custom embedding request header value: 1 to 8192 + characters with no control characters other than tab. + maxLength: 8192 + minLength: 1 + pattern: ^[^\x00-\x08\x0A-\x1F\x7F]*$ + type: string + description: |- + EmbeddingHeaders holds additional HTTP headers sent with every embedding + request. Only supported when EmbeddingProvider is "openai". Values are + stored in plain text and must not contain secrets; Authorization + (derived from OPENAI_API_KEY) and Content-Type cannot be set. + maxProperties: 32 + type: object embeddingModel: description: |- EmbeddingModel is the model name requested from the embedding service @@ -1887,6 +1903,16 @@ spec: pattern: ^([0-9]*[.])?[0-9]+$ type: string type: object + x-kubernetes-validations: + - message: embeddingHeaders is only supported when embeddingProvider + is 'openai' + rule: '!has(self.embeddingHeaders) || (has(self.embeddingProvider) + && self.embeddingProvider == ''openai'')' + - message: embeddingHeaders names must be valid HTTP header names + and must not include Authorization or Content-Type + rule: '!has(self.embeddingHeaders) || self.embeddingHeaders.all(k, + k.matches(''^[!#$%&\\x27*+.^_\\x60|~0-9A-Za-z-]+$'') && !(k.lowerAscii() + in [''authorization'', ''content-type'']))' outgoingAuth: description: |- OutgoingAuth configures how the virtual MCP server authenticates to backends. @@ -5216,6 +5242,22 @@ spec: instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. properties: + embeddingHeaders: + additionalProperties: + description: |- + EmbeddingHeaderValue is a custom embedding request header value: 1 to 8192 + characters with no control characters other than tab. + maxLength: 8192 + minLength: 1 + pattern: ^[^\x00-\x08\x0A-\x1F\x7F]*$ + type: string + description: |- + EmbeddingHeaders holds additional HTTP headers sent with every embedding + request. Only supported when EmbeddingProvider is "openai". Values are + stored in plain text and must not contain secrets; Authorization + (derived from OPENAI_API_KEY) and Content-Type cannot be set. + maxProperties: 32 + type: object embeddingModel: description: |- EmbeddingModel is the model name requested from the embedding service @@ -5293,6 +5335,16 @@ spec: pattern: ^([0-9]*[.])?[0-9]+$ type: string type: object + x-kubernetes-validations: + - message: embeddingHeaders is only supported when embeddingProvider + is 'openai' + rule: '!has(self.embeddingHeaders) || (has(self.embeddingProvider) + && self.embeddingProvider == ''openai'')' + - message: embeddingHeaders names must be valid HTTP header names + and must not include Authorization or Content-Type + rule: '!has(self.embeddingHeaders) || self.embeddingHeaders.all(k, + k.matches(''^[!#$%&\\x27*+.^_\\x60|~0-9A-Za-z-]+$'') && !(k.lowerAscii() + in [''authorization'', ''content-type'']))' outgoingAuth: description: |- OutgoingAuth configures how the virtual MCP server authenticates to backends. diff --git a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml index bf51183620..95b6a9d942 100644 --- a/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml +++ b/deploy/charts/operator-crds/templates/toolhive.stacklok.dev_virtualmcpservers.yaml @@ -1813,6 +1813,22 @@ spec: instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. properties: + embeddingHeaders: + additionalProperties: + description: |- + EmbeddingHeaderValue is a custom embedding request header value: 1 to 8192 + characters with no control characters other than tab. + maxLength: 8192 + minLength: 1 + pattern: ^[^\x00-\x08\x0A-\x1F\x7F]*$ + type: string + description: |- + EmbeddingHeaders holds additional HTTP headers sent with every embedding + request. Only supported when EmbeddingProvider is "openai". Values are + stored in plain text and must not contain secrets; Authorization + (derived from OPENAI_API_KEY) and Content-Type cannot be set. + maxProperties: 32 + type: object embeddingModel: description: |- EmbeddingModel is the model name requested from the embedding service @@ -1890,6 +1906,16 @@ spec: pattern: ^([0-9]*[.])?[0-9]+$ type: string type: object + x-kubernetes-validations: + - message: embeddingHeaders is only supported when embeddingProvider + is 'openai' + rule: '!has(self.embeddingHeaders) || (has(self.embeddingProvider) + && self.embeddingProvider == ''openai'')' + - message: embeddingHeaders names must be valid HTTP header names + and must not include Authorization or Content-Type + rule: '!has(self.embeddingHeaders) || self.embeddingHeaders.all(k, + k.matches(''^[!#$%&\\x27*+.^_\\x60|~0-9A-Za-z-]+$'') && !(k.lowerAscii() + in [''authorization'', ''content-type'']))' outgoingAuth: description: |- OutgoingAuth configures how the virtual MCP server authenticates to backends. @@ -5219,6 +5245,22 @@ spec: instead of all backend tools directly. This reduces token usage by allowing LLMs to discover relevant tools on demand rather than receiving all tool definitions. properties: + embeddingHeaders: + additionalProperties: + description: |- + EmbeddingHeaderValue is a custom embedding request header value: 1 to 8192 + characters with no control characters other than tab. + maxLength: 8192 + minLength: 1 + pattern: ^[^\x00-\x08\x0A-\x1F\x7F]*$ + type: string + description: |- + EmbeddingHeaders holds additional HTTP headers sent with every embedding + request. Only supported when EmbeddingProvider is "openai". Values are + stored in plain text and must not contain secrets; Authorization + (derived from OPENAI_API_KEY) and Content-Type cannot be set. + maxProperties: 32 + type: object embeddingModel: description: |- EmbeddingModel is the model name requested from the embedding service @@ -5296,6 +5338,16 @@ spec: pattern: ^([0-9]*[.])?[0-9]+$ type: string type: object + x-kubernetes-validations: + - message: embeddingHeaders is only supported when embeddingProvider + is 'openai' + rule: '!has(self.embeddingHeaders) || (has(self.embeddingProvider) + && self.embeddingProvider == ''openai'')' + - message: embeddingHeaders names must be valid HTTP header names + and must not include Authorization or Content-Type + rule: '!has(self.embeddingHeaders) || self.embeddingHeaders.all(k, + k.matches(''^[!#$%&\\x27*+.^_\\x60|~0-9A-Za-z-]+$'') && !(k.lowerAscii() + in [''authorization'', ''content-type'']))' outgoingAuth: description: |- OutgoingAuth configures how the virtual MCP server authenticates to backends. diff --git a/docs/operator/crd-api.md b/docs/operator/crd-api.md index d783eadb4a..c6372fc93b 100644 --- a/docs/operator/crd-api.md +++ b/docs/operator/crd-api.md @@ -455,6 +455,8 @@ _Appears in:_ | `action` _string_ | Action defines the action to take when the user declines or cancels
- skip_remaining: Skip remaining steps in the workflow
- abort: Abort the entire workflow execution
- continue: Continue to the next step | abort | Enum: [skip_remaining abort continue]
Optional: \{\}
| + + #### vmcp.config.FailureHandlingConfig @@ -567,6 +569,7 @@ _Appears in:_ | `embeddingServiceTimeout` _[vmcp.config.Duration](#vmcpconfigduration)_ | EmbeddingServiceTimeout is the HTTP request timeout for calls to the embedding service.
Defaults to 30s if not specified. | 30s | Pattern: `^([0-9]+(\.[0-9]+)?(ns\|us\|µs\|ms\|s\|m\|h))+$`
Type: string
Optional: \{\}
| | `embeddingProvider` _string_ | EmbeddingProvider selects the wire protocol used to talk to the embedding
service. "tei" speaks the HuggingFace Text Embeddings Inference API;
"openai" speaks the OpenAI-compatible /embeddings API, which lets the
optimizer use OpenAI, Azure OpenAI, or another OpenAI-compatible gateway.
Defaults to "tei" when empty.
The "openai" provider reads EmbeddingService directly and cannot be combined
with EmbeddingServerRef, which provisions a managed TEI server; the operator
rejects that combination at admission. | tei | Enum: [tei openai]
Optional: \{\}
| | `embeddingModel` _string_ | EmbeddingModel is the model name requested from the embedding service
(e.g. "text-embedding-3-small"). Required when EmbeddingProvider is
"openai". Ignored for the "tei" provider, where the model is fixed by the
running TEI container.
The API key for an OpenAI-compatible service is not configured here: it is
read from the OPENAI_API_KEY environment variable so the secret never
lands in a CRD spec or ConfigMap. An empty key omits the Authorization
header, which supports keyless in-cluster gateways. | | Optional: \{\}
| +| `embeddingHeaders` _object (keys:string, values:[vmcp.config.EmbeddingHeaderValue](#vmcpconfigembeddingheadervalue))_ | EmbeddingHeaders holds additional HTTP headers sent with every embedding
request. Only supported when EmbeddingProvider is "openai". Values are
stored in plain text and must not contain secrets; Authorization
(derived from OPENAI_API_KEY) and Content-Type cannot be set. | | MaxProperties: 32
Optional: \{\}
| | `maxToolsToReturn` _integer_ | MaxToolsToReturn is the maximum number of tool results returned by a search query.
Defaults to 8 if not specified or zero. | | Maximum: 50
Minimum: 1
Optional: \{\}
| | `hybridSearchSemanticRatio` _string_ | HybridSearchSemanticRatio controls the balance between semantic (meaning-based)
and keyword search results. 0.0 = all keyword, 1.0 = all semantic.
Defaults to "0.5" if not specified or empty.
Serialized as a string because CRDs do not support float types portably. | | Pattern: `^([0-9]*[.])?[0-9]+$`
Optional: \{\}
| | `semanticDistanceThreshold` _string_ | SemanticDistanceThreshold is the maximum distance for semantic search results.
Results exceeding this threshold are filtered out from semantic search.
This threshold does not apply to keyword search.
Range: 0 = identical, 2 = completely unrelated.
Defaults to "1.0" if not specified or empty.
Serialized as a string because CRDs do not support float types portably. | | Pattern: `^([0-9]*[.])?[0-9]+$`
Optional: \{\}
| diff --git a/examples/operator/virtual-mcps/vmcp_optimizer_openai.yaml b/examples/operator/virtual-mcps/vmcp_optimizer_openai.yaml index b51a740c06..5ef2047f6a 100644 --- a/examples/operator/virtual-mcps/vmcp_optimizer_openai.yaml +++ b/examples/operator/virtual-mcps/vmcp_optimizer_openai.yaml @@ -68,6 +68,10 @@ spec: # Model requested from the service (required for the openai provider). embeddingModel: text-embedding-3-small embeddingServiceTimeout: 15s + # Optional extra headers sent with every embedding request (e.g. a + # gateway cache-scoping key). Plain text — never put secrets here. + embeddingHeaders: + x-cache-key: toolhive-optimizer incomingAuth: type: anonymous diff --git a/pkg/vmcp/config/config.go b/pkg/vmcp/config/config.go index 669e3665ed..32e4b107e5 100644 --- a/pkg/vmcp/config/config.go +++ b/pkg/vmcp/config/config.go @@ -908,8 +908,13 @@ type OutputProperty struct { // OptimizerConfig configures the MCP optimizer. // When enabled, vMCP exposes only find_tool and call_tool operations to clients // instead of all backend tools directly. +// // +kubebuilder:object:generate=true +// +kubebuilder:validation:XValidation:rule="!has(self.embeddingHeaders) || (has(self.embeddingProvider) && self.embeddingProvider == 'openai')",message="embeddingHeaders is only supported when embeddingProvider is 'openai'" +// +kubebuilder:validation:XValidation:rule=`!has(self.embeddingHeaders) || self.embeddingHeaders.all(k, k.matches('^[!#$%&\\x27*+.^_\\x60|~0-9A-Za-z-]+$') && !(k.lowerAscii() in ['authorization', 'content-type']))`,message="embeddingHeaders names must be valid HTTP header names and must not include Authorization or Content-Type" // +gendoc +// +//nolint:lll // CEL validation rules exceed line length limit type OptimizerConfig struct { // EmbeddingService is the full base URL of the embedding service endpoint // (e.g., http://my-embedding.default.svc.cluster.local:8080) for semantic @@ -959,6 +964,14 @@ type OptimizerConfig struct { // +optional EmbeddingModel string `json:"embeddingModel,omitempty" yaml:"embeddingModel,omitempty"` + // EmbeddingHeaders holds additional HTTP headers sent with every embedding + // request. Only supported when EmbeddingProvider is "openai". Values are + // stored in plain text and must not contain secrets; Authorization + // (derived from OPENAI_API_KEY) and Content-Type cannot be set. + // +kubebuilder:validation:MaxProperties=32 + // +optional + EmbeddingHeaders map[string]EmbeddingHeaderValue `json:"embeddingHeaders,omitempty" yaml:"embeddingHeaders,omitempty"` + // MaxToolsToReturn is the maximum number of tool results returned by a search query. // Defaults to 8 if not specified or zero. // +kubebuilder:validation:Minimum=1 @@ -985,6 +998,13 @@ type OptimizerConfig struct { SemanticDistanceThreshold string `json:"semanticDistanceThreshold,omitempty" yaml:"semanticDistanceThreshold,omitempty"` } +// EmbeddingHeaderValue is a custom embedding request header value: 1 to 8192 +// characters with no control characters other than tab. +// +kubebuilder:validation:MinLength=1 +// +kubebuilder:validation:MaxLength=8192 +// +kubebuilder:validation:Pattern=`^[^\x00-\x08\x0A-\x1F\x7F]*$` +type EmbeddingHeaderValue string + // CodeModeConfig configures vMCP code mode (the execute_tool_script virtual tool). // When enabled, agents can submit a Starlark script that calls multiple backend tools // server-side — with loops, conditionals, and parallel() fan-out — and receive a single diff --git a/pkg/vmcp/config/zz_generated.deepcopy.go b/pkg/vmcp/config/zz_generated.deepcopy.go index 2ec5928a5c..e460451b35 100644 --- a/pkg/vmcp/config/zz_generated.deepcopy.go +++ b/pkg/vmcp/config/zz_generated.deepcopy.go @@ -213,7 +213,7 @@ func (in *Config) DeepCopyInto(out *Config) { if in.Optimizer != nil { in, out := &in.Optimizer, &out.Optimizer *out = new(OptimizerConfig) - **out = **in + (*in).DeepCopyInto(*out) } if in.CodeMode != nil { in, out := &in.CodeMode, &out.CodeMode @@ -375,6 +375,13 @@ func (in *OperationalConfig) DeepCopy() *OperationalConfig { // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *OptimizerConfig) DeepCopyInto(out *OptimizerConfig) { *out = *in + if in.EmbeddingHeaders != nil { + in, out := &in.EmbeddingHeaders, &out.EmbeddingHeaders + *out = make(map[string]EmbeddingHeaderValue, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new OptimizerConfig. diff --git a/pkg/vmcp/optimizer/internal/similarity/embedding_client.go b/pkg/vmcp/optimizer/internal/similarity/embedding_client.go index 1d94d206de..cde96f72ef 100644 --- a/pkg/vmcp/optimizer/internal/similarity/embedding_client.go +++ b/pkg/vmcp/optimizer/internal/similarity/embedding_client.go @@ -22,7 +22,8 @@ func NewEmbeddingClient(cfg *types.OptimizerConfig) (types.EmbeddingClient, erro case "", types.EmbeddingProviderTEI: return newTEIClient(cfg.EmbeddingService, cfg.EmbeddingServiceTimeout) case types.EmbeddingProviderOpenAI: - return newOpenAIClient(cfg.EmbeddingService, cfg.EmbeddingModel, cfg.EmbeddingAPIKey, cfg.EmbeddingServiceTimeout) + return newOpenAIClient(cfg.EmbeddingService, cfg.EmbeddingModel, cfg.EmbeddingAPIKey, + cfg.EmbeddingHeaders, cfg.EmbeddingServiceTimeout) default: return nil, fmt.Errorf("unsupported embedding provider %q (supported: %q, %q)", cfg.EmbeddingProvider, types.EmbeddingProviderTEI, types.EmbeddingProviderOpenAI) diff --git a/pkg/vmcp/optimizer/internal/similarity/openai_client.go b/pkg/vmcp/optimizer/internal/similarity/openai_client.go index c9f3d1fdec..907164d16c 100644 --- a/pkg/vmcp/optimizer/internal/similarity/openai_client.go +++ b/pkg/vmcp/optimizer/internal/similarity/openai_client.go @@ -10,6 +10,7 @@ import ( "fmt" "io" "log/slog" + "maps" "net/http" "strings" "time" @@ -29,15 +30,17 @@ type openAIClient struct { baseURL string apiKey string model string + headers map[string]string httpClient *http.Client maxBatchSize int } // newOpenAIClient creates a client that POSTs to baseURL+"/embeddings" using the // given model. A non-empty apiKey is sent as a Bearer token; an empty apiKey -// omits the Authorization header so keyless endpoints work. Zero timeout uses -// defaultTimeout. -func newOpenAIClient(baseURL, model, apiKey string, timeout time.Duration) (*openAIClient, error) { +// omits the Authorization header so keyless endpoints work. headers are set on +// every request but cannot override Content-Type or Authorization. Zero timeout +// uses defaultTimeout. +func newOpenAIClient(baseURL, model, apiKey string, headers map[string]string, timeout time.Duration) (*openAIClient, error) { if baseURL == "" { return nil, fmt.Errorf("OpenAI embedding base URL is required") } @@ -51,12 +54,13 @@ func newOpenAIClient(baseURL, model, apiKey string, timeout time.Duration) (*ope } slog.Debug("OpenAI embedding client created", - "base_url", baseURL, "model", model, "timeout", timeout) + "base_url", baseURL, "model", model, "timeout", timeout, "custom_headers", len(headers)) return &openAIClient{ baseURL: baseURL, apiKey: apiKey, model: model, + headers: maps.Clone(headers), httpClient: &http.Client{Timeout: timeout}, maxBatchSize: openAIMaxBatchSize, }, nil @@ -129,6 +133,10 @@ func (c *openAIClient) embedChunk(ctx context.Context, texts []string) ([][]floa if err != nil { return nil, fmt.Errorf("failed to create OpenAI request: %w", err) } + for name, value := range c.headers { + req.Header.Set(name, value) + } + // Set after the custom headers so they can never be overridden. req.Header.Set("Content-Type", "application/json") if c.apiKey != "" { req.Header.Set("Authorization", "Bearer "+c.apiKey) diff --git a/pkg/vmcp/optimizer/internal/similarity/openai_client_integration_test.go b/pkg/vmcp/optimizer/internal/similarity/openai_client_integration_test.go index 59efacaad2..0812d01791 100644 --- a/pkg/vmcp/optimizer/internal/similarity/openai_client_integration_test.go +++ b/pkg/vmcp/optimizer/internal/similarity/openai_client_integration_test.go @@ -27,7 +27,7 @@ func TestOpenAIClient_Live(t *testing.T) { baseURL := cmp.Or(os.Getenv("OPENAI_EMBEDDING_BASE_URL"), "https://api.openai.com/v1") model := cmp.Or(os.Getenv("OPENAI_EMBEDDING_MODEL"), "text-embedding-3-small") - client, err := newOpenAIClient(baseURL, model, apiKey, 0) + client, err := newOpenAIClient(baseURL, model, apiKey, nil, 0) require.NoError(t, err) t.Cleanup(func() { _ = client.Close() }) diff --git a/pkg/vmcp/optimizer/internal/similarity/openai_client_test.go b/pkg/vmcp/optimizer/internal/similarity/openai_client_test.go index af7da66bc1..dd8a2b6879 100644 --- a/pkg/vmcp/optimizer/internal/similarity/openai_client_test.go +++ b/pkg/vmcp/optimizer/internal/similarity/openai_client_test.go @@ -20,21 +20,21 @@ func Test_newOpenAIClient(t *testing.T) { t.Run("empty URL returns error", func(t *testing.T) { t.Parallel() - client, err := newOpenAIClient("", "text-embedding-3-small", "key", 0) + client, err := newOpenAIClient("", "text-embedding-3-small", "key", nil, 0) require.ErrorContains(t, err, "OpenAI embedding base URL is required") require.Nil(t, client) }) t.Run("empty model returns error", func(t *testing.T) { t.Parallel() - client, err := newOpenAIClient("http://embeddings:8080/v1", "", "key", 0) + client, err := newOpenAIClient("http://embeddings:8080/v1", "", "key", nil, 0) require.ErrorContains(t, err, "OpenAI embedding model is required") require.Nil(t, client) }) t.Run("valid args create client with default batch size", func(t *testing.T) { t.Parallel() - client, err := newOpenAIClient("http://embeddings:8080/v1", "text-embedding-3-small", "key", 0) + client, err := newOpenAIClient("http://embeddings:8080/v1", "text-embedding-3-small", "key", nil, 0) require.NoError(t, err) require.NotNil(t, client) require.Equal(t, openAIMaxBatchSize, client.maxBatchSize) @@ -43,11 +43,20 @@ func Test_newOpenAIClient(t *testing.T) { t.Run("custom timeout", func(t *testing.T) { t.Parallel() - client, err := newOpenAIClient("http://embeddings:8080/v1", "text-embedding-3-small", "key", 5*time.Second) + client, err := newOpenAIClient("http://embeddings:8080/v1", "text-embedding-3-small", "key", nil, 5*time.Second) require.NoError(t, err) require.NotNil(t, client) require.Equal(t, 5*time.Second, client.httpClient.Timeout) }) + + t.Run("headers are cloned at construction", func(t *testing.T) { + t.Parallel() + headers := map[string]string{"x-cache-key": "toolhive"} + client, err := newOpenAIClient("http://embeddings:8080/v1", "text-embedding-3-small", "key", headers, 0) + require.NoError(t, err) + headers["x-cache-key"] = "mutated" + require.Equal(t, "toolhive", client.headers["x-cache-key"]) + }) } func TestOpenAIClient_Embed(t *testing.T) { @@ -278,6 +287,48 @@ func TestOpenAIClient_OmitsAuthHeaderWhenKeyless(t *testing.T) { require.NoError(t, err) } +func TestOpenAIClient_SendsCustomHeaders(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "toolhive-optimizer", r.Header.Get("x-cache-key")) + require.Equal(t, "eu-west", r.Header.Get("X-Gateway-Region")) + require.Equal(t, "application/json", r.Header.Get("Content-Type")) + require.Equal(t, "Bearer test-key", r.Header.Get("Authorization")) + writeOpenAIEmbeddings(t, w, [][]float32{{0.1}}) + })) + t.Cleanup(srv.Close) + + client, err := newOpenAIClient(srv.URL, "text-embedding-3-small", "test-key", map[string]string{ + "x-cache-key": "toolhive-optimizer", + "X-Gateway-Region": "eu-west", + }, 0) + require.NoError(t, err) + + _, err = client.Embed(context.Background(), "hello") + require.NoError(t, err) +} + +func TestOpenAIClient_ProtocolHeadersWinOverCustomHeaders(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "application/json", r.Header.Get("Content-Type")) + require.Equal(t, "Bearer test-key", r.Header.Get("Authorization")) + writeOpenAIEmbeddings(t, w, [][]float32{{0.1}}) + })) + t.Cleanup(srv.Close) + + client, err := newOpenAIClient(srv.URL, "text-embedding-3-small", "test-key", map[string]string{ + "authorization": "Bearer spoofed", + "content-type": "text/plain", + }, 0) + require.NoError(t, err) + + _, err = client.Embed(context.Background(), "hello") + require.NoError(t, err) +} + func TestOpenAIClient_Close(t *testing.T) { t.Parallel() diff --git a/pkg/vmcp/optimizer/internal/types/types.go b/pkg/vmcp/optimizer/internal/types/types.go index 1eb90f424a..9c54f5a92c 100644 --- a/pkg/vmcp/optimizer/internal/types/types.go +++ b/pkg/vmcp/optimizer/internal/types/types.go @@ -94,6 +94,11 @@ type OptimizerConfig struct { // keyless in-cluster gateways. Never populated for the TEI provider. EmbeddingAPIKey string + // EmbeddingHeaders holds additional HTTP headers sent with every request + // to an OpenAI-compatible embedding service. Never populated for the TEI + // provider. + EmbeddingHeaders map[string]string + // MaxToolsToReturn limits the number of tools returned by FindTool. MaxToolsToReturn *int diff --git a/pkg/vmcp/optimizer/optimizer.go b/pkg/vmcp/optimizer/optimizer.go index 7e41d0d0ed..35389cf0f2 100644 --- a/pkg/vmcp/optimizer/optimizer.go +++ b/pkg/vmcp/optimizer/optimizer.go @@ -18,11 +18,13 @@ import ( "log/slog" "os" "strconv" + "strings" "time" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" + httpval "github.com/stacklok/toolhive-core/validation/http" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/tokencounter" @@ -54,6 +56,7 @@ func GetAndValidateConfig(cfg *vmcpconfig.OptimizerConfig) (*Config, error) { EmbeddingServiceTimeout: time.Duration(cfg.EmbeddingServiceTimeout), EmbeddingProvider: cfg.EmbeddingProvider, EmbeddingModel: cfg.EmbeddingModel, + EmbeddingHeaders: convertEmbeddingHeaders(cfg.EmbeddingHeaders), } if err := resolveEmbeddingProvider(optCfg); err != nil { @@ -100,8 +103,9 @@ func GetAndValidateConfig(cfg *vmcpconfig.OptimizerConfig) (*Config, error) { // resolveEmbeddingProvider normalizes and validates the embedding provider on // optCfg in place. An empty provider defaults to TEI so existing configs keep -// working; the OpenAI provider requires a service and model and reads its API -// key from the environment. +// working; the OpenAI provider requires a service and model, reads its API +// key from the environment, and is the only provider that accepts custom +// embedding headers. func resolveEmbeddingProvider(optCfg *Config) error { switch optCfg.EmbeddingProvider { case "": @@ -116,11 +120,61 @@ func resolveEmbeddingProvider(optCfg *Config) error { return fmt.Errorf("optimizer.embeddingModel is required when optimizer.embeddingProvider is %q", types.EmbeddingProviderOpenAI) } + if err := validateEmbeddingHeaders(optCfg.EmbeddingHeaders); err != nil { + return err + } optCfg.EmbeddingAPIKey = os.Getenv(embeddingAPIKeyEnvVar) default: return fmt.Errorf("optimizer.embeddingProvider must be %q or %q, got %q", types.EmbeddingProviderTEI, types.EmbeddingProviderOpenAI, optCfg.EmbeddingProvider) } + + // Defense in depth: mirrors the CEL rule on config.OptimizerConfig, + // covering config sources with no admission validation. + if optCfg.EmbeddingProvider != types.EmbeddingProviderOpenAI && len(optCfg.EmbeddingHeaders) > 0 { + return fmt.Errorf("optimizer.embeddingHeaders is only supported when optimizer.embeddingProvider is %q", + types.EmbeddingProviderOpenAI) + } + + return nil +} + +// convertEmbeddingHeaders converts the config header map to the internal +// plain-string representation, returning nil for an empty map. +func convertEmbeddingHeaders(headers map[string]vmcpconfig.EmbeddingHeaderValue) map[string]string { + if len(headers) == 0 { + return nil + } + out := make(map[string]string, len(headers)) + for name, value := range headers { + out[name] = string(value) + } + return out +} + +// validateEmbeddingHeaders rejects custom embedding headers with invalid +// RFC 7230 names or values, and headers the OpenAI client sets itself: +// Content-Type is always application/json, and Authorization is derived from +// the OPENAI_API_KEY environment variable so the token never lands in config. +// Reserved names are compared case-insensitively, matching HTTP semantics. +// Mirrors the CEL rule on config.OptimizerConfig as defense in depth. +func validateEmbeddingHeaders(headers map[string]string) error { + for name, value := range headers { + if err := httpval.ValidateHeaderName(name); err != nil { + return fmt.Errorf("optimizer.embeddingHeaders: invalid header name %q: %w", name, err) + } + if err := httpval.ValidateHeaderValue(value); err != nil { + return fmt.Errorf("optimizer.embeddingHeaders[%q]: %w", name, err) + } + switch strings.ToLower(name) { + case "authorization": + return fmt.Errorf("optimizer.embeddingHeaders must not set %q: the Authorization header is derived "+ + "from the %s environment variable", name, embeddingAPIKeyEnvVar) + case "content-type": + return fmt.Errorf("optimizer.embeddingHeaders must not set %q: the Content-Type header is always "+ + "application/json", name) + } + } return nil } diff --git a/pkg/vmcp/optimizer/optimizer_test.go b/pkg/vmcp/optimizer/optimizer_test.go index 8ccef87b0c..1c59228507 100644 --- a/pkg/vmcp/optimizer/optimizer_test.go +++ b/pkg/vmcp/optimizer/optimizer_test.go @@ -77,6 +77,103 @@ func TestGetAndValidateConfig(t *testing.T) { EmbeddingModel: "text-embedding-3-small", }, }, + { + name: "openai provider with custom headers", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "toolhive-optimizer"}, + }, + expected: &Config{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]string{"x-cache-key": "toolhive-optimizer"}, + }, + }, + { + name: "openai provider with uncommon valid header name", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-key'name`x": "value"}, + }, + expected: &Config{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]string{"x-key'name`x": "value"}, + }, + }, + { + name: "error: headers with tei provider", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://embeddings:8080", + EmbeddingProvider: types.EmbeddingProviderTEI, + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "toolhive-optimizer"}, + }, + errContains: "optimizer.embeddingHeaders is only supported", + }, + { + name: "error: headers with default provider", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://embeddings:8080", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "toolhive-optimizer"}, + }, + errContains: "optimizer.embeddingHeaders is only supported", + }, + { + name: "error: headers set authorization case-insensitively", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"AUTHORIZATION": "Bearer spoofed"}, + }, + errContains: "optimizer.embeddingHeaders must not set \"AUTHORIZATION\"", + }, + { + name: "error: headers set content-type", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"Content-Type": "text/plain"}, + }, + errContains: "optimizer.embeddingHeaders must not set \"Content-Type\"", + }, + { + name: "error: empty header name", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"": "value"}, + }, + errContains: "invalid header name", + }, + { + name: "error: header name with invalid characters", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x bad": "value"}, + }, + errContains: "invalid header name \"x bad\"", + }, + { + name: "error: header value with control characters", + cfg: &vmcpconfig.OptimizerConfig{ + EmbeddingService: "http://gateway:8080/v1", + EmbeddingProvider: types.EmbeddingProviderOpenAI, + EmbeddingModel: "text-embedding-3-small", + EmbeddingHeaders: map[string]vmcpconfig.EmbeddingHeaderValue{"x-cache-key": "a\r\nb"}, + }, + errContains: "invalid HTTP header value", + }, { name: "error: openai provider without service", cfg: &vmcpconfig.OptimizerConfig{ @@ -263,6 +360,7 @@ func TestGetAndValidateConfig(t *testing.T) { } assert.Equal(t, wantProvider, result.EmbeddingProvider) assert.Equal(t, tt.expected.EmbeddingModel, result.EmbeddingModel) + assert.Equal(t, tt.expected.EmbeddingHeaders, result.EmbeddingHeaders) if tt.expected.MaxToolsToReturn != nil { require.NotNil(t, result.MaxToolsToReturn)