fix: guardrail identifier in runtime headers

core/providers/bedrock/mantle_test.go

@@ -2,6 +2,7 @@ package bedrock
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	schemas "github.com/maximhq/bifrost/core/schemas"
@@ -110,3 +111,115 @@ func TestParseBedrockRegionAndModelStripsForTheWire(t *testing.T) {
 		})
 	}
 }
+
+// Converse takes a guardrail as a guardrailConfig body field; the OpenAI-compatible
+// endpoints take it as headers and ignore the body field, so the two renderings are not
+// interchangeable. Mantle is deliberately not wired: it accepts the headers and enforces
+// nothing, so there is no rendering that works there.
+func TestWithGuardrailHeaders(t *testing.T) {
+	base := map[string]string{"X-Existing": "keep"}
+
+	t.Run("renders the headers and consumes the extra param", func(t *testing.T) {
+		extra := map[string]any{
+			"guardrailConfig": map[string]any{
+				"guardrailIdentifier": "gr-123",
+				"guardrailVersion":    "DRAFT",
+				"trace":               "ENABLED",
+			},
+			"other": "untouched",
+		}
+		got := withGuardrailHeaders(base, extra)
+
+		if got[guardrailIdentifierHeader] != "gr-123" || got[guardrailVersionHeader] != "DRAFT" {
+			t.Errorf("guardrail headers = %v", got)
+		}
+		if got[guardrailTraceHeader] != "ENABLED" {
+			t.Errorf("trace header = %q", got[guardrailTraceHeader])
+		}
+		if got["X-Existing"] != "keep" {
+			t.Error("existing headers must survive")
+		}
+		// Read, never consumed: core reuses one request across retry attempts, so
+		// removing it would drop the guardrail on every attempt after the first.
+		if _, still := extra["guardrailConfig"]; !still {
+			t.Error("guardrailConfig must survive for the next retry attempt")
+		}
+		if extra["other"] != "untouched" {
+			t.Error("unrelated extra params must be left alone")
+		}
+		// The shared networkConfig map must never be written through.
+		if _, leaked := base[guardrailIdentifierHeader]; leaked {
+			t.Error("base header map was mutated")
+		}
+	})
+
+	t.Run("no guardrail config returns base unchanged", func(t *testing.T) {
+		if got := withGuardrailHeaders(base, map[string]any{"other": 1}); len(got) != 1 {
+			t.Errorf("expected base untouched, got %v", got)
+		}
+	})
+
+	// Both fields are required upstream; half a config is left alone rather than sent.
+	t.Run("identifier without version is not sent", func(t *testing.T) {
+		extra := map[string]any{"guardrailConfig": map[string]any{"guardrailIdentifier": "gr-123"}}
+		got := withGuardrailHeaders(base, extra)
+		if _, ok := got[guardrailIdentifierHeader]; ok {
+			t.Error("a half-formed guardrail config must not be sent")
+		}
+		if _, still := extra["guardrailConfig"]; !still {
+			t.Error("an unused config must be left in place")
+		}
+	})
+
+	// SetExtraHeaders canonicalises keys and keeps the first it reaches; Go map order is
+	// random, so a differently-cased static header must be replaced, not merely shadowed.
+	t.Run("a differently-cased static header is replaced, not doubled", func(t *testing.T) {
+		static := map[string]string{
+			"x-amzn-bedrock-guardrailidentifier": "gr-static",
+			"X-Other":                            "keep",
+		}
+		extra := map[string]any{"guardrailConfig": map[string]any{
+			"guardrailIdentifier": "gr-request", "guardrailVersion": "1"}}
+		got := withGuardrailHeaders(static, extra)
+
+		seen := 0
+		for k, v := range got {
+			if strings.EqualFold(k, guardrailIdentifierHeader) {
+				seen++
+				if v != "gr-request" {
+					t.Errorf("per-request value should win, got %q", v)
+				}
+			}
+		}
+		if seen != 1 {
+			t.Errorf("expected exactly one identifier header, found %d in %v", seen, got)
+		}
+		if got["X-Other"] != "keep" {
+			t.Error("unrelated static headers must survive")
+		}
+		if static["x-amzn-bedrock-guardrailidentifier"] != "gr-static" {
+			t.Error("the caller's map was mutated")
+		}
+	})
+
+	// The same request object is handed to every retry attempt, so repeated calls must
+	// render the same headers rather than degrade after the first.
+	t.Run("repeated calls are stable across retry attempts", func(t *testing.T) {
+		extra := map[string]any{"guardrailConfig": map[string]any{
+			"guardrailIdentifier": "gr-1", "guardrailVersion": "DRAFT"}}
+		first := withGuardrailHeaders(base, extra)
+		second := withGuardrailHeaders(base, extra)
+		if second[guardrailIdentifierHeader] != first[guardrailIdentifierHeader] ||
+			second[guardrailVersionHeader] != first[guardrailVersionHeader] {
+			t.Errorf("second attempt lost the guardrail: %v then %v", first, second)
+		}
+	})
+
+	t.Run("nil base map is handled", func(t *testing.T) {
+		extra := map[string]any{"guardrailConfig": map[string]any{
+			"guardrailIdentifier": "gr-1", "guardrailVersion": "1"}}
+		if got := withGuardrailHeaders(nil, extra); got[guardrailIdentifierHeader] != "gr-1" {
+			t.Errorf("got %v", got)
+		}
+	})
+}

core/providers/bedrock/runtimeopenai.go

@@ -3,6 +3,8 @@ package bedrock
 import (
 	"context"
 	"fmt"
+	"maps"
+	"strings"
 
 	openai "github.com/maximhq/bifrost/core/providers/openai"
 	providerUtils "github.com/maximhq/bifrost/core/providers/utils"
@@ -17,6 +19,60 @@ func runtimeOpenAIURL(endpoints *schemas.BedrockEndpoints, region, path string)
 	return fmt.Sprintf("https://%s/openai/v1/%s", resolveBedrockHost(endpoints, bedrockServiceRuntime, region), path)
 }
 
+// Guardrail headers for the OpenAI-compatible surfaces. Converse takes the same intent
+// as a guardrailConfig body field; these endpoints take it as headers and ignore the body
+// field entirely, so the two renderings are not interchangeable.
+const (
+	guardrailIdentifierHeader = "X-Amzn-Bedrock-GuardrailIdentifier"
+	guardrailVersionHeader    = "X-Amzn-Bedrock-GuardrailVersion"
+	guardrailTraceHeader      = "X-Amzn-Bedrock-Trace"
+)
+
+// withGuardrailHeaders returns headers naming the guardrail in the request's
+// guardrailConfig extra param, which is how bedrock-runtime's OpenAI-compatible
+// endpoints take it. Converse expresses the same intent as a body field and these
+// endpoints ignore that field, so the two renderings are not interchangeable.
+//
+// Identifier and version are both required upstream, so a config carrying one is left
+// alone rather than half-sent. streamProcessingMode has no header equivalent.
+//
+// The headers are not signed: AWS requires only x-amz-* in SignedHeaders and these are
+// x-amzn-*. Verified live — a guardrail applies identically either way.
+//
+// Mantle is deliberately not wired: it accepts these headers and enforces nothing.
+func withGuardrailHeaders(base map[string]string, extraParams map[string]any) map[string]string {
+	config, _ := extraParams["guardrailConfig"].(map[string]any)
+	identifier, _ := config["guardrailIdentifier"].(string)
+	version, _ := config["guardrailVersion"].(string)
+	if identifier == "" || version == "" {
+		return base
+	}
+
+	out := maps.Clone(base)
+	if out == nil {
+		out = make(map[string]string, 3)
+	}
+	setHeader(out, guardrailIdentifierHeader, identifier)
+	setHeader(out, guardrailVersionHeader, version)
+	if trace, _ := config["trace"].(string); trace != "" {
+		setHeader(out, guardrailTraceHeader, trace)
+	}
+	return out
+}
+
+// setHeader assigns name, first dropping any key that differs from it only in case.
+// SetExtraHeaders canonicalises every key and keeps whichever it reaches first, and Go
+// map order is random, so a differently-cased entry would race this one rather than
+// lose to it.
+func setHeader(headers map[string]string, name, value string) {
+	for existing := range headers {
+		if existing != name && strings.EqualFold(existing, name) {
+			delete(headers, existing)
+		}
+	}
+	headers[name] = value
+}
+
 // runtimeResponses handles non-streaming Responses requests on bedrock-runtime's
 // OpenAI-compatible surface. Payloads and SSE follow the OpenAI Responses spec, so the shared
 // OpenAI handler does the work and only the URL and SigV4 scope differ from mantle.
@@ -28,13 +84,17 @@ func (provider *BedrockProvider) runtimeResponses(
 	region := resolveBedrockRegion(ctx, key, request.Model)
 	url := runtimeOpenAIURL(bedrockEndpoints(key.BedrockKeyConfig), region, "responses")
 	_, request.Model = parseBedrockRegionAndModel(request.Model)
+	extraHeaders := provider.networkConfig.ExtraHeaders
+	if request.Params != nil {
+		extraHeaders = withGuardrailHeaders(extraHeaders, request.Params.ExtraParams)
+	}
 
 	// SigV4 (empty key value): sign the exact body the handler builds via a signer closure.
 	// Bearer (key has a value): no signer; auth flows through the Authorization header.
 	var signer providerUtils.BodySigner
 	if key.Value.GetValue() == "" {
 		signer = func(body []byte) (map[string]string, *schemas.BifrostError) {
-			return signOpenAIV4Headers(ctx, body, url, "application/json", key, region, provider.networkConfig.ExtraHeaders, bedrockSigningService)
+			return signOpenAIV4Headers(ctx, body, url, "application/json", key, region, extraHeaders, bedrockSigningService)
 		}
 	}
 
@@ -45,7 +105,7 @@ func (provider *BedrockProvider) runtimeResponses(
 		url,
 		request,
 		openai.BearerAuthHeader(key),
-		provider.networkConfig.ExtraHeaders,
+		extraHeaders,
 		providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
 		providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
 		provider.GetProviderKey(),
@@ -68,17 +128,21 @@ func (provider *BedrockProvider) runtimeResponsesStream(
 	region := resolveBedrockRegion(ctx, key, request.Model)
 	url := runtimeOpenAIURL(bedrockEndpoints(key.BedrockKeyConfig), region, "responses")
 	_, request.Model = parseBedrockRegionAndModel(request.Model)
+	extraHeaders := provider.networkConfig.ExtraHeaders
+	if request.Params != nil {
+		extraHeaders = withGuardrailHeaders(extraHeaders, request.Params.ExtraParams)
+	}
 
 	var signer providerUtils.BodySigner
 	if key.Value.GetValue() == "" {
 		signer = func(body []byte) (map[string]string, *schemas.BifrostError) {
-			return signOpenAIV4Headers(ctx, body, url, "text/event-stream", key, region, provider.networkConfig.ExtraHeaders, bedrockSigningService)
+			return signOpenAIV4Headers(ctx, body, url, "text/event-stream", key, region, extraHeaders, bedrockSigningService)
 		}
 	}
 
 	return openai.HandleOpenAIResponsesStreaming(
 		ctx, provider.mantleStreamingClient, url, request,
-		openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders,
+		openai.BearerAuthHeader(key), extraHeaders,
 		provider.networkConfig.StreamIdleTimeoutInSeconds,
 		providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
 		providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
@@ -103,11 +167,15 @@ func (provider *BedrockProvider) runtimeChatCompletions(
 	region := resolveBedrockRegion(ctx, key, request.Model)
 	url := runtimeOpenAIURL(bedrockEndpoints(key.BedrockKeyConfig), region, "chat/completions")
 	_, request.Model = parseBedrockRegionAndModel(request.Model)
+	extraHeaders := provider.networkConfig.ExtraHeaders
+	if request.Params != nil {
+		extraHeaders = withGuardrailHeaders(extraHeaders, request.Params.ExtraParams)
+	}
 
 	var signer providerUtils.BodySigner
 	if key.Value.GetValue() == "" {
 		signer = func(body []byte) (map[string]string, *schemas.BifrostError) {
-			return signOpenAIV4Headers(ctx, body, url, "application/json", key, region, provider.networkConfig.ExtraHeaders, bedrockSigningService)
+			return signOpenAIV4Headers(ctx, body, url, "application/json", key, region, extraHeaders, bedrockSigningService)
 		}
 	}
 
@@ -117,7 +185,7 @@ func (provider *BedrockProvider) runtimeChatCompletions(
 		url,
 		request,
 		openai.BearerAuthHeader(key),
-		provider.networkConfig.ExtraHeaders,
+		extraHeaders,
 		providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
 		providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
 		provider.GetProviderKey(),
@@ -140,17 +208,21 @@ func (provider *BedrockProvider) runtimeChatCompletionsStream(
 	region := resolveBedrockRegion(ctx, key, request.Model)
 	url := runtimeOpenAIURL(bedrockEndpoints(key.BedrockKeyConfig), region, "chat/completions")
 	_, request.Model = parseBedrockRegionAndModel(request.Model)
+	extraHeaders := provider.networkConfig.ExtraHeaders
+	if request.Params != nil {
+		extraHeaders = withGuardrailHeaders(extraHeaders, request.Params.ExtraParams)
+	}
 
 	var signer providerUtils.BodySigner
 	if key.Value.GetValue() == "" {
 		signer = func(body []byte) (map[string]string, *schemas.BifrostError) {
-			return signOpenAIV4Headers(ctx, body, url, "text/event-stream", key, region, provider.networkConfig.ExtraHeaders, bedrockSigningService)
+			return signOpenAIV4Headers(ctx, body, url, "text/event-stream", key, region, extraHeaders, bedrockSigningService)
 		}
 	}
 
 	return openai.HandleOpenAIChatCompletionStreaming(
 		ctx, provider.mantleStreamingClient, url, request,
-		openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders,
+		openai.BearerAuthHeader(key), extraHeaders,
 		provider.networkConfig.StreamIdleTimeoutInSeconds,
 		providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
 		providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),