Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion core/providers/bedrock/bedrock.go
Original file line number Diff line number Diff line change
Expand Up @@ -723,7 +723,7 @@ func (provider *BedrockProvider) listMantleModels(ctx *schemas.BifrostContext, k
provider.logger.Warn("failed to build mantle list-models request: %v", err)
return nil
}
providerUtils.SetExtraHeadersHTTP(ctx, req, provider.networkConfig.ExtraHeaders, nil)
providerUtils.SetExtraHeadersHTTP(ctx, req, WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)), nil)
if key.Value.GetValue() != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key.Value.GetValue()))
} else if bifrostErr := signAWSRequest(ctx, req, key.BedrockKeyConfig, region, bedrockMantleSigningService); bifrostErr != nil {
Expand Down
43 changes: 39 additions & 4 deletions core/providers/bedrock/mantle.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"context"
"fmt"
"io"
"maps"
"net/http"
"strings"

Expand All @@ -13,6 +14,40 @@ import (
schemas "github.com/maximhq/bifrost/core/schemas"
)

const (
// MantleOpenAIProjectHeader selects a Bedrock Mantle project on the OpenAI-compatible surface
// (chat/completions, responses, /models). AWS routes to the account's default project when absent.
MantleOpenAIProjectHeader = "OpenAI-Project"
// MantleAnthropicProjectHeader selects a Bedrock Mantle project on the native-Anthropic surface
// (/anthropic/v1/messages).
MantleAnthropicProjectHeader = "anthropic-workspace-id"
)

// WithMantleProject returns headers with the given Mantle project header set when projectID is
// non-empty, letting AWS fall back to the account's default project when it is empty. It never
// mutates base (which may be the shared networkConfig.ExtraHeaders map). The project header is a
// plain (non x-amz-*) header, so it does not need to be part of the SigV4 SignedHeaders.
func WithMantleProject(base map[string]string, headerName, projectID string) map[string]string {
if projectID == "" {
return base
}
out := maps.Clone(base)
if out == nil {
Comment thread
greptile-apps[bot] marked this conversation as resolved.
out = make(map[string]string, 1)
}
out[headerName] = projectID
return out
}

// resolveMantleProjectID returns the Bedrock project configured for the mantle sub-surface of the
// Bedrock provider, or "" when none is set (AWS then routes to the account's default project).
func resolveMantleProjectID(key schemas.Key) string {
Comment thread
impoiler marked this conversation as resolved.
if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ProjectID != nil {
return key.BedrockKeyConfig.ProjectID.GetValue()
}
return ""
}
Comment on lines +44 to +49

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P0 Duplicate resolver declaration This declares resolveMantleProjectID in package bedrock, but core/providers/bedrock/utils.go also declares a package-level resolveMantleProjectID with a different signature. Go does not support function overloading, so any build or test target that includes this package fails before the new project-scoping code can run. Keep one resolver and update the call sites to that signature.


// isMantleModel reports whether a model should be routed via the Bedrock Mantle
// OpenAI-compatible endpoint. OpenAI-family (gpt-*) and Gemma 4 models are mantle-only
// (they have no Converse equivalent). Gemma 3 is intentionally excluded: it only supports
Expand Down Expand Up @@ -131,7 +166,7 @@ func (provider *BedrockProvider) mantleChatCompletions(
url,
request,
openai.BearerAuthHeader(key),
provider.networkConfig.ExtraHeaders,
WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)),
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
provider.GetProviderKey(),
Expand Down Expand Up @@ -165,7 +200,7 @@ func (provider *BedrockProvider) mantleChatCompletionsStream(

return openai.HandleOpenAIChatCompletionStreaming(
ctx, provider.mantleStreamingClient, url, request,
openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders,
openai.BearerAuthHeader(key), WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)),
provider.networkConfig.StreamIdleTimeoutInSeconds,
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
Expand Down Expand Up @@ -206,7 +241,7 @@ func (provider *BedrockProvider) mantleResponses(
url,
request,
openai.BearerAuthHeader(key),
provider.networkConfig.ExtraHeaders,
WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)),
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
provider.GetProviderKey(),
Expand Down Expand Up @@ -240,7 +275,7 @@ func (provider *BedrockProvider) mantleResponsesStream(

return openai.HandleOpenAIResponsesStreaming(
ctx, provider.mantleStreamingClient, url, request,
openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders,
openai.BearerAuthHeader(key), WithMantleProject(provider.networkConfig.ExtraHeaders, MantleOpenAIProjectHeader, resolveMantleProjectID(key)),
provider.networkConfig.StreamIdleTimeoutInSeconds,
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
Expand Down
75 changes: 75 additions & 0 deletions core/providers/bedrock/mantle_project_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
package bedrock

import (
"testing"

schemas "github.com/maximhq/bifrost/core/schemas"
)

// TestWithMantleProject verifies the project header is added only when a project ID is present,
// the target header name is honoured, and the base map is never mutated.
func TestWithMantleProject(t *testing.T) {
t.Run("empty project returns base unchanged", func(t *testing.T) {
base := map[string]string{"X-Custom": "v"}
got := WithMantleProject(base, MantleOpenAIProjectHeader, "")
if _, ok := got[MantleOpenAIProjectHeader]; ok {
t.Fatalf("expected no project header when project ID is empty, got %v", got)
}
// Empty project must return the base map as-is (default-project behaviour).
if len(got) != 1 || got["X-Custom"] != "v" {
t.Fatalf("expected base returned unchanged, got %v", got)
}
})

t.Run("OpenAI project header set", func(t *testing.T) {
base := map[string]string{"X-Custom": "v"}
got := WithMantleProject(base, MantleOpenAIProjectHeader, "proj_abc")
if got[MantleOpenAIProjectHeader] != "proj_abc" {
t.Fatalf("expected %s=proj_abc, got %v", MantleOpenAIProjectHeader, got)
}
if got["X-Custom"] != "v" {
t.Fatalf("existing headers must be preserved, got %v", got)
}
// base must not be mutated.
if _, ok := base[MantleOpenAIProjectHeader]; ok {
t.Fatalf("base map was mutated: %v", base)
}
})

t.Run("Anthropic workspace header set", func(t *testing.T) {
got := WithMantleProject(nil, MantleAnthropicProjectHeader, "proj_xyz")
if got[MantleAnthropicProjectHeader] != "proj_xyz" {
t.Fatalf("expected %s=proj_xyz, got %v", MantleAnthropicProjectHeader, got)
}
})

t.Run("nil base with empty project stays nil", func(t *testing.T) {
if got := WithMantleProject(nil, MantleOpenAIProjectHeader, ""); got != nil {
t.Fatalf("expected nil when base is nil and project is empty, got %v", got)
}
})
}

// TestResolveMantleProjectID verifies precedence of the BedrockKeyConfig.ProjectID field.
func TestResolveMantleProjectID(t *testing.T) {
tests := []struct {
name string
key schemas.Key
want string
}{
{name: "no bedrock config", key: schemas.Key{}, want: ""},
{name: "config without project", key: schemas.Key{BedrockKeyConfig: &schemas.BedrockKeyConfig{}}, want: ""},
{
name: "config with project",
key: schemas.Key{BedrockKeyConfig: &schemas.BedrockKeyConfig{ProjectID: schemas.NewSecretVar("proj_abc")}},
want: "proj_abc",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := resolveMantleProjectID(tt.key); got != tt.want {
t.Fatalf("resolveMantleProjectID = %q, want %q", got, tt.want)
}
})
}
}
16 changes: 16 additions & 0 deletions core/providers/bedrock/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,22 @@ var bedrockUnsafeToolNameCharRegex = regexp.MustCompile(`[^A-Za-z0-9_-]+`)
// bedrockToolNameAliasKey stores Bedrock wire-name aliases on the request context.
type bedrockToolNameAliasKey struct{}

// resolveMantleProjectID returns the Bedrock project configured for the mantle sub-surface of the
// Bedrock provider, or "" when none is set (AWS then routes to the account's default project).
// Priority: per-alias AliasConfig.ProjectID > key-level BedrockKeyConfig.ProjectID. The per-alias
// override lets one Bedrock credential scope different aliased models to different projects.
func resolveMantleProjectID(ctx *schemas.BifrostContext, key schemas.Key) string {
Comment thread
impoiler marked this conversation as resolved.
if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.ProjectID != nil {
if v := ra.Config.ProjectID.GetValue(); v != "" {
return v
}
}
if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ProjectID != nil {
return key.BedrockKeyConfig.ProjectID.GetValue()
}
return ""
}

// parseBedrockRegionAndModel splits a model string that optionally carries an AWS region prefix
// into its region and bare model ID components.
// If no region prefix is present the returned region is empty and bareModel equals model.
Expand Down
24 changes: 13 additions & 11 deletions core/providers/bedrockmantle/bedrockmantle.go
Original file line number Diff line number Diff line change
Expand Up @@ -122,16 +122,18 @@ func (provider *BedrockMantleProvider) listModelsByKey(ctx *schemas.BifrostConte
region := provider.resolveRegion(ctx, key, "")
mURL := mantleOpenAIURL(region, "", "models")

extraHeaders := provider.networkConfig.ExtraHeaders
// Scope the catalog to the configured project via the OpenAI-Project header (default project
// when unset). It is a plain header, so it does not need to be part of the SigV4 SignedHeaders.
extraHeaders := bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleOpenAIProjectHeader, resolveProjectID(key))
if key.Value.GetValue() == "" {
// SigV4: sign the GET and overlay the signed headers; OpenAI's ListModelsByKey only sets
// a Bearer header when the key carries a value, so the SigV4 Authorization wins here.
sigHeaders, bifrostErr := bedrock.SignMantleV4Headers(ctx, nil, mURL, "", key, region, provider.networkConfig.ExtraHeaders)
if bifrostErr != nil {
return nil, bifrostErr
}
merged := make(map[string]string, len(provider.networkConfig.ExtraHeaders)+len(sigHeaders))
maps.Copy(merged, provider.networkConfig.ExtraHeaders)
merged := make(map[string]string, len(extraHeaders)+len(sigHeaders))
maps.Copy(merged, extraHeaders)
maps.Copy(merged, sigHeaders)
extraHeaders = merged
}
Expand Down Expand Up @@ -184,7 +186,7 @@ func (provider *BedrockMantleProvider) ChatCompletion(ctx *schemas.BifrostContex
ShouldSendBackRawResponse: provider.sendBackRawResponse,
},
openai.BearerAuthHeader(key),
addAnthropicHeaders(provider.networkConfig.ExtraHeaders),
addAnthropicHeaders(bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleAnthropicProjectHeader, resolveProjectID(key))),
provider.mantleSigner(ctx, key, url, "application/json", region),
provider.logger,
)
Expand All @@ -197,7 +199,7 @@ func (provider *BedrockMantleProvider) ChatCompletion(ctx *schemas.BifrostContex
url,
request,
openai.BearerAuthHeader(key),
provider.networkConfig.ExtraHeaders,
bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleOpenAIProjectHeader, resolveProjectID(key)),
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
provider.GetProviderKey(),
Expand Down Expand Up @@ -237,7 +239,7 @@ func (provider *BedrockMantleProvider) ChatCompletionStream(ctx *schemas.Bifrost
url,
jsonData,
openai.BearerAuthHeader(key),
addAnthropicHeaders(provider.networkConfig.ExtraHeaders),
addAnthropicHeaders(bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleAnthropicProjectHeader, resolveProjectID(key))),
provider.networkConfig.StreamIdleTimeoutInSeconds,
provider.networkConfig.BetaHeaderOverrides,
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
Expand All @@ -254,7 +256,7 @@ func (provider *BedrockMantleProvider) ChatCompletionStream(ctx *schemas.Bifrost
url := mantleOpenAIURL(region, schemas.ResolveCanonicalModel(ctx, request.Model), "chat/completions")
return openai.HandleOpenAIChatCompletionStreaming(
ctx, provider.mantleStreamingClient, url, request,
openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders,
openai.BearerAuthHeader(key), bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleOpenAIProjectHeader, resolveProjectID(key)),
provider.networkConfig.StreamIdleTimeoutInSeconds,
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
Expand Down Expand Up @@ -290,7 +292,7 @@ func (provider *BedrockMantleProvider) Responses(ctx *schemas.BifrostContext, ke
ShouldSendBackRawResponse: provider.sendBackRawResponse,
},
openai.BearerAuthHeader(key),
addAnthropicHeaders(provider.networkConfig.ExtraHeaders),
addAnthropicHeaders(bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleAnthropicProjectHeader, resolveProjectID(key))),
provider.mantleSigner(ctx, key, url, "application/json", region),
provider.logger,
)
Expand All @@ -303,7 +305,7 @@ func (provider *BedrockMantleProvider) Responses(ctx *schemas.BifrostContext, ke
url,
request,
openai.BearerAuthHeader(key),
provider.networkConfig.ExtraHeaders,
bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleOpenAIProjectHeader, resolveProjectID(key)),
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
provider.GetProviderKey(),
Expand Down Expand Up @@ -343,7 +345,7 @@ func (provider *BedrockMantleProvider) ResponsesStream(ctx *schemas.BifrostConte
url,
jsonData,
openai.BearerAuthHeader(key),
addAnthropicHeaders(provider.networkConfig.ExtraHeaders),
addAnthropicHeaders(bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleAnthropicProjectHeader, resolveProjectID(key))),
provider.networkConfig.StreamIdleTimeoutInSeconds,
provider.networkConfig.BetaHeaderOverrides,
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
Expand All @@ -360,7 +362,7 @@ func (provider *BedrockMantleProvider) ResponsesStream(ctx *schemas.BifrostConte
url := mantleOpenAIURL(region, schemas.ResolveCanonicalModel(ctx, request.Model), "responses")
return openai.HandleOpenAIResponsesStreaming(
ctx, provider.mantleStreamingClient, url, request,
openai.BearerAuthHeader(key), provider.networkConfig.ExtraHeaders,
openai.BearerAuthHeader(key), bedrock.WithMantleProject(provider.networkConfig.ExtraHeaders, bedrock.MantleOpenAIProjectHeader, resolveProjectID(key)),
provider.networkConfig.StreamIdleTimeoutInSeconds,
providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest),
providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse),
Expand Down
32 changes: 32 additions & 0 deletions core/providers/bedrockmantle/project_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
package bedrockmantle

import (
"testing"

schemas "github.com/maximhq/bifrost/core/schemas"
)

// TestResolveProjectID verifies the BedrockMantleKeyConfig.ProjectID field is honoured and that an
// absent project resolves to "" (AWS default project).
func TestResolveProjectID(t *testing.T) {
tests := []struct {
name string
key schemas.Key
want string
}{
{name: "no mantle config", key: schemas.Key{}, want: ""},
{name: "config without project", key: schemas.Key{BedrockMantleKeyConfig: &schemas.BedrockMantleKeyConfig{}}, want: ""},
{
name: "config with project",
key: schemas.Key{BedrockMantleKeyConfig: &schemas.BedrockMantleKeyConfig{ProjectID: schemas.NewSecretVar("proj_xyz")}},
want: "proj_xyz",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := resolveProjectID(tt.key); got != tt.want {
t.Fatalf("resolveProjectID = %q, want %q", got, tt.want)
}
})
}
}
10 changes: 10 additions & 0 deletions core/providers/bedrockmantle/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,16 @@ func addAnthropicHeaders(headers map[string]string) map[string]string {
return out
}

// resolveProjectID returns the Bedrock project configured for this key, or "" when none is set
// (AWS then routes to the account's default project). The value is sent as the OpenAI-Project or
// anthropic-workspace-id header depending on the request surface.
func resolveProjectID(key schemas.Key) string {
if key.BedrockMantleKeyConfig != nil && key.BedrockMantleKeyConfig.ProjectID != nil {
return key.BedrockMantleKeyConfig.ProjectID.GetValue()
}
return ""
}

// parseBedrockRegionAndModel splits a model string that optionally carries an AWS region prefix
// into its region and bare model ID components.
// If no region prefix is present the returned region is empty and bareModel equals model.
Expand Down
12 changes: 12 additions & 0 deletions core/schemas/account.go
Original file line number Diff line number Diff line change
Expand Up @@ -671,6 +671,12 @@ type BedrockKeyConfig struct {
ExternalID *SecretVar `json:"external_id,omitempty"`
RoleSessionName *SecretVar `json:"session_name,omitempty"`

// ProjectID scopes the Bedrock Mantle sub-surface (OpenAI-compatible gpt-*/Gemma routing and the
// mantle catalog merge in ListModels) to a specific Bedrock project via the "OpenAI-Project"
// header. When empty, AWS routes to the account's default project. It has no effect on the
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Comment thread
greptile-apps[bot] marked this conversation as resolved.
// Converse/bedrock-runtime paths, which are not project-scoped.
ProjectID *SecretVar `json:"project_id,omitempty"`
Comment thread
impoiler marked this conversation as resolved.

BatchS3Config *BatchS3Config `json:"batch_s3_config,omitempty"` // S3 bucket configuration for batch operations
}

Expand All @@ -691,6 +697,12 @@ type BedrockMantleKeyConfig struct {
RoleARN *SecretVar `json:"role_arn,omitempty"`
ExternalID *SecretVar `json:"external_id,omitempty"`
RoleSessionName *SecretVar `json:"session_name,omitempty"`

// ProjectID scopes inference and model listing to a specific Bedrock project. It is sent as the
// "OpenAI-Project" header on the OpenAI-compatible surface and the "anthropic-workspace-id"
// header on the native-Anthropic (Claude) surface. When empty, AWS routes to the account's
// default project.
ProjectID *SecretVar `json:"project_id,omitempty"`
Comment thread
impoiler marked this conversation as resolved.
}

// NOTE: To use Bedrock Mantle IAM role authentication, set both AccessKey and SecretKey to empty
Expand Down
8 changes: 8 additions & 0 deletions framework/configstore/clientconfig.go
Original file line number Diff line number Diff line change
Expand Up @@ -571,6 +571,10 @@ func (p *ProviderConfig) Redacted() *ProviderConfig {
if key.BedrockKeyConfig.RoleSessionName != nil {
bedrockConfig.RoleSessionName = key.BedrockKeyConfig.RoleSessionName.Redacted()
}
// Mantle project ID is an identifier, not a credential — surface it in plaintext.
if key.BedrockKeyConfig.ProjectID != nil {
bedrockConfig.ProjectID = key.BedrockKeyConfig.ProjectID
}
// Add back s3 config
if key.BedrockKeyConfig.BatchS3Config != nil {
bedrockConfig.BatchS3Config = key.BedrockKeyConfig.BatchS3Config
Expand Down Expand Up @@ -598,6 +602,10 @@ func (p *ProviderConfig) Redacted() *ProviderConfig {
if key.BedrockMantleKeyConfig.RoleSessionName != nil {
mantleConfig.RoleSessionName = key.BedrockMantleKeyConfig.RoleSessionName.Redacted()
}
// Project ID is an identifier, not a credential — surface it in plaintext.
if key.BedrockMantleKeyConfig.ProjectID != nil {
mantleConfig.ProjectID = key.BedrockMantleKeyConfig.ProjectID
}
redactedConfig.Keys[i].BedrockMantleKeyConfig = mantleConfig
}

Expand Down
Loading
Loading