diff --git a/core/providers/bedrock/bedrock.go b/core/providers/bedrock/bedrock.go index 094d920fbe..2dd7e12e63 100644 --- a/core/providers/bedrock/bedrock.go +++ b/core/providers/bedrock/bedrock.go @@ -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 { diff --git a/core/providers/bedrock/mantle.go b/core/providers/bedrock/mantle.go index 30edbbcf04..587cf06d39 100644 --- a/core/providers/bedrock/mantle.go +++ b/core/providers/bedrock/mantle.go @@ -5,6 +5,7 @@ import ( "context" "fmt" "io" + "maps" "net/http" "strings" @@ -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 { + 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 { + if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ProjectID != nil { + return key.BedrockKeyConfig.ProjectID.GetValue() + } + return "" +} + // 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 @@ -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(), @@ -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), @@ -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(), @@ -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), diff --git a/core/providers/bedrock/mantle_project_test.go b/core/providers/bedrock/mantle_project_test.go new file mode 100644 index 0000000000..6e2ef7ef25 --- /dev/null +++ b/core/providers/bedrock/mantle_project_test.go @@ -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) + } + }) + } +} diff --git a/core/providers/bedrock/utils.go b/core/providers/bedrock/utils.go index c615ceed12..657885df41 100644 --- a/core/providers/bedrock/utils.go +++ b/core/providers/bedrock/utils.go @@ -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 { + 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. diff --git a/core/providers/bedrockmantle/bedrockmantle.go b/core/providers/bedrockmantle/bedrockmantle.go index 78ab092ba5..c1bb90bd93 100644 --- a/core/providers/bedrockmantle/bedrockmantle.go +++ b/core/providers/bedrockmantle/bedrockmantle.go @@ -122,7 +122,9 @@ 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. @@ -130,8 +132,8 @@ func (provider *BedrockMantleProvider) listModelsByKey(ctx *schemas.BifrostConte 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 } @@ -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, ) @@ -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(), @@ -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), @@ -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), @@ -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, ) @@ -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(), @@ -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), @@ -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), diff --git a/core/providers/bedrockmantle/project_test.go b/core/providers/bedrockmantle/project_test.go new file mode 100644 index 0000000000..5f674efe79 --- /dev/null +++ b/core/providers/bedrockmantle/project_test.go @@ -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) + } + }) + } +} diff --git a/core/providers/bedrockmantle/utils.go b/core/providers/bedrockmantle/utils.go index d49998cdbc..4f895e17cc 100644 --- a/core/providers/bedrockmantle/utils.go +++ b/core/providers/bedrockmantle/utils.go @@ -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. diff --git a/core/schemas/account.go b/core/schemas/account.go index 0ab3200bbf..bbdb044aa5 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -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 + // Converse/bedrock-runtime paths, which are not project-scoped. + ProjectID *SecretVar `json:"project_id,omitempty"` + BatchS3Config *BatchS3Config `json:"batch_s3_config,omitempty"` // S3 bucket configuration for batch operations } @@ -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"` } // NOTE: To use Bedrock Mantle IAM role authentication, set both AccessKey and SecretKey to empty diff --git a/framework/configstore/clientconfig.go b/framework/configstore/clientconfig.go index ccc16938dc..9505d864d7 100644 --- a/framework/configstore/clientconfig.go +++ b/framework/configstore/clientconfig.go @@ -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 @@ -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 } diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 5fe5c07e7c..fbc9826f17 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -442,6 +442,7 @@ var configstoreMigrationSteps = []migrationStep{ {IDs: []string{"add_fast_mode_cache_pricing_columns"}, run: migrationAddFastModeCachePricingColumns}, {IDs: []string{"add_inference_geo_multiplier_column"}, run: migrationAddInferenceGeoMultiplierColumn}, {IDs: []string{"repair_bare_wildcard_allowed_models"}, run: migrationRepairBareWildcardAllowedModels}, + {IDs: []string{"add_bedrock_project_id_columns"}, run: migrationAddBedrockProjectIDColumns}, } // quoteSQLiteIdentifier quotes a SQLite identifier, escaping any double quotes. @@ -1164,6 +1165,44 @@ func migrationAddBedrockMantleKeyColumns(ctx context.Context, db *gorm.DB, logge return nil } +// migrationAddBedrockProjectIDColumns adds the bedrock_project_id and bedrock_mantle_project_id +// columns to the config_keys table. These scope Bedrock Mantle inference / model listing to a +// specific Bedrock project via the OpenAI-Project / anthropic-workspace-id header. +func migrationAddBedrockProjectIDColumns(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { + migrationName := "add_bedrock_project_id_columns" + logger.Info("[configstore] starting migration %s", migrationName) + defer logger.Info("[configstore] finished migration %s", migrationName) + cols := []string{ + "bedrock_project_id", + "bedrock_mantle_project_id", + } + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: migrationName, + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + for _, col := range cols { + if err := addColumnIfNotExists(tx, logger, &tables.TableKey{}, col); err != nil { + return err + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + for _, col := range cols { + if err := dropColumnIfExists(tx, logger, &tables.TableKey{}, col); err != nil { + return err + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error while running db migration: %s", err.Error()) + } + return nil +} + func migrationAddAllowedOriginsJSONColumn(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { migrationName := "add_allowed_origins_json_column" logger.Info("[configstore] starting migration %s", migrationName) diff --git a/framework/configstore/tables/encryption_test.go b/framework/configstore/tables/encryption_test.go index d243d3a5e1..83f280e9b9 100644 --- a/framework/configstore/tables/encryption_test.go +++ b/framework/configstore/tables/encryption_test.go @@ -184,6 +184,7 @@ func TestTableKey_BedrockFieldsEncryptDecrypt(t *testing.T) { SecretKey: *schemas.NewSecretVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), Region: schemas.NewSecretVar("us-west-2"), ARN: schemas.NewSecretVar("arn:aws:iam::123456789:role/test"), + ProjectID: schemas.NewSecretVar("proj_bedrock123"), BatchS3Config: &schemas.BatchS3Config{ Buckets: []schemas.S3BucketConfig{ {BucketName: "my-batch-bucket", Prefix: "jobs/", IsDefault: true}, @@ -200,6 +201,7 @@ func TestTableKey_BedrockFieldsEncryptDecrypt(t *testing.T) { assert.NotEqual(t, "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", raw["bedrock_secret_key"]) assert.NotEqual(t, "us-west-2", raw["bedrock_region"]) assert.NotEqual(t, "arn:aws:iam::123456789:role/test", raw["bedrock_arn"]) + assert.NotEqual(t, "proj_bedrock123", raw["bedrock_project_id"]) rawAliasesVal := raw["aliases_json"] require.NotNil(t, rawAliasesVal, "aliases_json should be present in raw row") var rawAliasesStr string @@ -224,6 +226,8 @@ func TestTableKey_BedrockFieldsEncryptDecrypt(t *testing.T) { assert.Equal(t, "us-west-2", found.BedrockKeyConfig.Region.GetValue()) require.NotNil(t, found.BedrockKeyConfig.ARN) assert.Equal(t, "arn:aws:iam::123456789:role/test", found.BedrockKeyConfig.ARN.GetValue()) + require.NotNil(t, found.BedrockKeyConfig.ProjectID) + assert.Equal(t, "proj_bedrock123", found.BedrockKeyConfig.ProjectID.GetValue()) assert.Equal(t, "profile-a", found.Aliases["model-a"].ModelID) require.NotNil(t, found.BedrockKeyConfig.BatchS3Config) require.Len(t, found.BedrockKeyConfig.BatchS3Config.Buckets, 1) @@ -832,6 +836,39 @@ func TestTableKey_BedrockSessionTokenEncryptDecrypt(t *testing.T) { assert.Equal(t, "us-east-1", found.BedrockKeyConfig.Region.GetValue()) } +// TestTableKey_BedrockMantleProjectID_RoundTrip verifies the bedrock_mantle_project_id column +// encrypts, persists, and reconstructs onto BedrockMantleKeyConfig.ProjectID. +func TestTableKey_BedrockMantleProjectID_RoundTrip(t *testing.T) { + db := setupTestDB(t) + + key := &TableKey{ + Name: "bedrock-mantle-proj-key", + ProviderID: 1, + Provider: "bedrock_mantle", + KeyID: "bedrock-mantle-proj-uuid", + Value: *schemas.NewSecretVar("mantle-val"), + BedrockMantleKeyConfig: &schemas.BedrockMantleKeyConfig{ + AccessKey: *schemas.NewSecretVar("AKIA-MANTLE-PROJ"), + SecretKey: *schemas.NewSecretVar("wJalr-MANTLE-PROJ"), + Region: schemas.NewSecretVar("us-east-1"), + ProjectID: schemas.NewSecretVar("proj_elvsngya7ixv4dkb26xe"), + }, + } + + require.NoError(t, db.Create(key).Error) + + raw := rawRow(t, db, "config_keys", key.ID) + assert.Equal(t, "encrypted", raw["encryption_status"]) + assert.NotEqual(t, "proj_elvsngya7ixv4dkb26xe", raw["bedrock_mantle_project_id"]) + + var found TableKey + require.NoError(t, db.First(&found, key.ID).Error) + require.NotNil(t, found.BedrockMantleKeyConfig) + require.NotNil(t, found.BedrockMantleKeyConfig.ProjectID) + assert.Equal(t, "proj_elvsngya7ixv4dkb26xe", found.BedrockMantleKeyConfig.ProjectID.GetValue()) + assert.Equal(t, "us-east-1", found.BedrockMantleKeyConfig.Region.GetValue()) +} + // ============================================================================ // MCP — edge cases for connection string / headers combinations // ============================================================================ @@ -1176,6 +1213,12 @@ func TestTableKey_AllProviderConfigs_EncryptDecrypt(t *testing.T) { Region: schemas.NewSecretVar("eu-west-1"), ARN: schemas.NewSecretVar("arn:aws:bedrock:eu-west-1:123:role"), }, + BedrockMantleKeyConfig: &schemas.BedrockMantleKeyConfig{ + AccessKey: *schemas.NewSecretVar("AKIA-MANTLE"), + SecretKey: *schemas.NewSecretVar("wJalr-MANTLE"), + Region: schemas.NewSecretVar("us-east-1"), + ProjectID: schemas.NewSecretVar("proj_mantle456"), + }, } require.NoError(t, db.Create(key).Error) @@ -1190,6 +1233,7 @@ func TestTableKey_AllProviderConfigs_EncryptDecrypt(t *testing.T) { assert.NotEqual(t, "us-central1", raw["vertex_region"]) assert.NotEqual(t, "eu-west-1", raw["bedrock_region"]) assert.NotEqual(t, "arn:aws:bedrock:eu-west-1:123:role", raw["bedrock_arn"]) + assert.NotEqual(t, "proj_mantle456", raw["bedrock_mantle_project_id"]) rawAliasesVal2 := raw["aliases_json"] require.NotNil(t, rawAliasesVal2, "aliases_json should be present in raw row") var rawAliasesStr2 string @@ -1230,6 +1274,13 @@ func TestTableKey_AllProviderConfigs_EncryptDecrypt(t *testing.T) { assert.Equal(t, "eu-west-1", found.BedrockKeyConfig.Region.GetValue()) require.NotNil(t, found.BedrockKeyConfig.ARN) assert.Equal(t, "arn:aws:bedrock:eu-west-1:123:role", found.BedrockKeyConfig.ARN.GetValue()) + + require.NotNil(t, found.BedrockMantleKeyConfig) + assert.Equal(t, "AKIA-MANTLE", found.BedrockMantleKeyConfig.AccessKey.GetValue()) + assert.Equal(t, "wJalr-MANTLE", found.BedrockMantleKeyConfig.SecretKey.GetValue()) + require.NotNil(t, found.BedrockMantleKeyConfig.ProjectID) + assert.Equal(t, "proj_mantle456", found.BedrockMantleKeyConfig.ProjectID.GetValue()) + assert.Equal(t, "profile-claude", found.Aliases["claude-3"].ModelID) } diff --git a/framework/configstore/tables/key.go b/framework/configstore/tables/key.go index d80d38636b..e8986b46d3 100644 --- a/framework/configstore/tables/key.go +++ b/framework/configstore/tables/key.go @@ -55,6 +55,7 @@ type TableKey struct { BedrockRoleARN *schemas.SecretVar `gorm:"type:text" json:"bedrock_role_arn,omitempty"` BedrockExternalID *schemas.SecretVar `gorm:"type:text" json:"bedrock_external_id,omitempty"` BedrockRoleSessionName *schemas.SecretVar `gorm:"type:text" json:"bedrock_role_session_name,omitempty"` + BedrockProjectID *schemas.SecretVar `gorm:"type:text" json:"bedrock_project_id,omitempty"` BedrockBatchS3ConfigJSON *string `gorm:"type:text" json:"-"` // JSON serialized schemas.BatchS3Config // Bedrock Mantle config fields (embedded) @@ -65,6 +66,7 @@ type TableKey struct { BedrockMantleRoleARN *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_role_arn,omitempty"` BedrockMantleExternalID *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_external_id,omitempty"` BedrockMantleRoleSessionName *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_role_session_name,omitempty"` + BedrockMantleProjectID *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_project_id,omitempty"` // VLLM config fields (embedded) VLLMUrl *schemas.SecretVar `gorm:"type:text" json:"vllm_url,omitempty"` @@ -267,6 +269,12 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error { } else { k.BedrockRoleSessionName = nil } + if k.BedrockKeyConfig.ProjectID != nil { + pid := *k.BedrockKeyConfig.ProjectID + k.BedrockProjectID = &pid + } else { + k.BedrockProjectID = nil + } if k.BedrockKeyConfig.BatchS3Config != nil { data, err := sonic.Marshal(k.BedrockKeyConfig.BatchS3Config) if err != nil { @@ -286,6 +294,7 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error { k.BedrockRoleARN = nil k.BedrockExternalID = nil k.BedrockRoleSessionName = nil + k.BedrockProjectID = nil k.BedrockBatchS3ConfigJSON = nil } @@ -333,6 +342,12 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error { } else { k.BedrockMantleRoleSessionName = nil } + if k.BedrockMantleKeyConfig.ProjectID != nil { + pid := *k.BedrockMantleKeyConfig.ProjectID + k.BedrockMantleProjectID = &pid + } else { + k.BedrockMantleProjectID = nil + } } else { k.BedrockMantleAccessKey = nil k.BedrockMantleSecretKey = nil @@ -341,6 +356,7 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error { k.BedrockMantleRoleARN = nil k.BedrockMantleExternalID = nil k.BedrockMantleRoleSessionName = nil + k.BedrockMantleProjectID = nil } if k.Aliases != nil { @@ -461,6 +477,9 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error { if err := encryptSecretVarPtr(&k.BedrockRoleSessionName); err != nil { return fmt.Errorf("failed to encrypt bedrock role session name: %w", err) } + if err := encryptSecretVarPtr(&k.BedrockProjectID); err != nil { + return fmt.Errorf("failed to encrypt bedrock project id: %w", err) + } if err := encryptString(k.BedrockBatchS3ConfigJSON); err != nil { return fmt.Errorf("failed to encrypt bedrock batch s3 config: %w", err) } @@ -486,6 +505,9 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error { if err := encryptSecretVarPtr(&k.BedrockMantleRoleSessionName); err != nil { return fmt.Errorf("failed to encrypt bedrock mantle role session name: %w", err) } + if err := encryptSecretVarPtr(&k.BedrockMantleProjectID); err != nil { + return fmt.Errorf("failed to encrypt bedrock mantle project id: %w", err) + } // Aliases if err := encryptString(k.AliasesJSON); err != nil { return fmt.Errorf("failed to encrypt aliases: %w", err) @@ -567,6 +589,9 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error { if err := decryptSecretVarPtr(&k.BedrockRoleSessionName); err != nil { return fmt.Errorf("failed to decrypt bedrock role session name: %w", err) } + if err := decryptSecretVarPtr(&k.BedrockProjectID); err != nil { + return fmt.Errorf("failed to decrypt bedrock project id: %w", err) + } if err := decryptString(k.BedrockBatchS3ConfigJSON); err != nil { return fmt.Errorf("failed to decrypt bedrock batch s3 config: %w", err) } @@ -592,6 +617,9 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error { if err := decryptSecretVarPtr(&k.BedrockMantleRoleSessionName); err != nil { return fmt.Errorf("failed to decrypt bedrock mantle role session name: %w", err) } + if err := decryptSecretVarPtr(&k.BedrockMantleProjectID); err != nil { + return fmt.Errorf("failed to decrypt bedrock mantle project id: %w", err) + } // Aliases if err := decryptString(k.AliasesJSON); err != nil { return fmt.Errorf("failed to decrypt aliases: %w", err) @@ -674,7 +702,7 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error { k.VertexKeyConfig = config } // Reconstruct Bedrock config if fields are present - if k.BedrockAccessKey != nil || k.BedrockSecretKey != nil || k.BedrockSessionToken != nil || k.BedrockRegion != nil || k.BedrockARN != nil || k.BedrockRoleARN != nil || k.BedrockExternalID != nil || k.BedrockRoleSessionName != nil || (k.BedrockBatchS3ConfigJSON != nil && *k.BedrockBatchS3ConfigJSON != "") { + if k.BedrockAccessKey != nil || k.BedrockSecretKey != nil || k.BedrockSessionToken != nil || k.BedrockRegion != nil || k.BedrockARN != nil || k.BedrockRoleARN != nil || k.BedrockExternalID != nil || k.BedrockRoleSessionName != nil || k.BedrockProjectID != nil || (k.BedrockBatchS3ConfigJSON != nil && *k.BedrockBatchS3ConfigJSON != "") { bedrockConfig := &schemas.BedrockKeyConfig{} if k.BedrockAccessKey != nil { @@ -687,6 +715,7 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error { bedrockConfig.RoleARN = k.BedrockRoleARN bedrockConfig.ExternalID = k.BedrockExternalID bedrockConfig.RoleSessionName = k.BedrockRoleSessionName + bedrockConfig.ProjectID = k.BedrockProjectID if k.BedrockSecretKey != nil { bedrockConfig.SecretKey = *k.BedrockSecretKey @@ -703,7 +732,7 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error { k.BedrockKeyConfig = bedrockConfig } // Reconstruct Bedrock Mantle config if fields are present - if k.BedrockMantleAccessKey != nil || k.BedrockMantleSecretKey != nil || k.BedrockMantleSessionToken != nil || k.BedrockMantleRegion != nil || k.BedrockMantleRoleARN != nil || k.BedrockMantleExternalID != nil || k.BedrockMantleRoleSessionName != nil { + if k.BedrockMantleAccessKey != nil || k.BedrockMantleSecretKey != nil || k.BedrockMantleSessionToken != nil || k.BedrockMantleRegion != nil || k.BedrockMantleRoleARN != nil || k.BedrockMantleExternalID != nil || k.BedrockMantleRoleSessionName != nil || k.BedrockMantleProjectID != nil { mantleConfig := &schemas.BedrockMantleKeyConfig{} if k.BedrockMantleAccessKey != nil { mantleConfig.AccessKey = *k.BedrockMantleAccessKey @@ -716,6 +745,7 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error { mantleConfig.RoleARN = k.BedrockMantleRoleARN mantleConfig.ExternalID = k.BedrockMantleExternalID mantleConfig.RoleSessionName = k.BedrockMantleRoleSessionName + mantleConfig.ProjectID = k.BedrockMantleProjectID k.BedrockMantleKeyConfig = mantleConfig } // Reconstruct Aliases diff --git a/transports/config.schema.json b/transports/config.schema.json index 6550d5e146..7f5e30dea3 100644 --- a/transports/config.schema.json +++ b/transports/config.schema.json @@ -3729,6 +3729,10 @@ "type": "string", "description": "Role session name for AssumeRole (can use env. prefix)" }, + "project_id": { + "type": "string", + "description": "Bedrock project ID scoping the Mantle sub-surface (OpenAI-compatible gpt-*/Gemma routing) via the OpenAI-Project header. When empty, AWS routes to the account's default project. No effect on the Converse/bedrock-runtime paths (can use env. prefix)." + }, "deployments": { "type": "object", "additionalProperties": { @@ -3812,6 +3816,10 @@ "session_name": { "type": "string", "description": "Role session name for AssumeRole (can use env. prefix)" + }, + "project_id": { + "type": "string", + "description": "Bedrock project ID scoping inference and model listing. 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 (can use env. prefix)." } }, "required": ["region"],