From 0caebf589e0142a47637d73ba247f0d8daf00fb3 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 8 Sep 2026 18:09:47 +0000 Subject: [PATCH 01/25] feat(coderd): resolve bedrock inference profiles when a provider is written --- aibridge/config/config.go | 6 + aibridge/provider/anthropic.go | 12 +- .../provider/bedrock_inference_profile.go | 53 +++- ...bedrock_inference_profile_internal_test.go | 115 ++++++-- cli/aibridged.go | 22 +- coderd/ai_providers.go | 37 ++- coderd/ai_providers_bedrock.go | 139 ++++++++++ coderd/ai_providers_bedrock_test.go | 255 ++++++++++++++++++ coderd/aibridged/proto/aibridged.pb.go | 179 ++++++------ coderd/aibridged/proto/aibridged.proto | 7 + coderd/aibridgedserver/aibridgedserver.go | 18 +- coderd/coderd.go | 6 +- coderd/coderdtest/coderdtest.go | 10 +- codersdk/aiproviders.go | 21 ++ codersdk/aiproviders_bedrock.go | 27 ++ docs/ai-coder/ai-gateway/providers.md | 11 +- site/src/api/typesGenerated.ts | 12 + 17 files changed, 790 insertions(+), 140 deletions(-) create mode 100644 coderd/ai_providers_bedrock.go create mode 100644 coderd/ai_providers_bedrock_test.go diff --git a/aibridge/config/config.go b/aibridge/config/config.go index ee0c5fec8a6..44ba4ef2a4d 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -68,6 +68,12 @@ type AWSBedrock struct { // Protocol selects the Bedrock wire protocol. The zero value behaves as // BedrockProtocolInvokeModel. Protocol BedrockProtocol + // ResolvedModel is the model ID behind Model, which differs from it only + // when Model is an application inference profile ARN. coderd resolves it + // when the provider is written, so the gateway never calls AWS for it. + ResolvedModel string + // ResolvedSmallFastModel is ResolvedModel for SmallFastModel. + ResolvedSmallFastModel string } // ResolvedProtocol returns the configured protocol, mapping the empty value to diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index e8796fde256..68a349e3f12 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -83,13 +83,13 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. return nil, xerrors.Errorf("bedrock config: %w", err) } - // Resolution only calls AWS for application inference profile ARNs, so - // deployments configured with plain model IDs need no extra permission. - resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) - defer cancel() - model, smallFastModel, err := resolveBedrockModels(resolveCtx, runtimeCfg, awsCfg) + // coderd resolves application inference profile ARNs when the provider + // is written, so construction never calls AWS. A missing resolution + // means the stored provider predates that step or was edited around it; + // serving it would silently misshape every request. + model, smallFastModel, err := resolvedBedrockModels(runtimeCfg) if err != nil { - return nil, xerrors.Errorf("resolve bedrock models: %w", err) + return nil, err } bedrock = messages.NewBedrockRuntime(runtimeCfg, awsCfg.Credentials, model, smallFastModel) diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index b0ed2f55664..6b9bc88762c 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -22,8 +22,8 @@ const bedrockService = "bedrock" const applicationInferenceProfileResourceType = "application-inference-profile" // inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls made -// while constructing a provider, which also cover the first credential -// resolution (STS/IRSA). +// while writing a provider, which also cover the first credential resolution +// (STS/IRSA). const inferenceProfileResolutionTimeout = 30 * time.Second // isApplicationInferenceProfileARN reports whether model is an application @@ -88,16 +88,33 @@ func modelIDFromARN(modelARN string) (string, error) { return model, nil } -// resolveBedrockModels resolves the configured model identifiers to the model +// ResolveBedrockModels resolves the configured model identifiers to the model // IDs used for capability detection, usage recording, and pricing. Identifiers // that are not application inference profile ARNs are returned unchanged and // cost no AWS call. -func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, awsCfg aws.Config) (model, smallFastModel string, err error) { +// +// It runs where a Bedrock provider is written rather than where it is served, +// so the gateway never calls the Bedrock control plane. The identity comes from +// cfg, including any role assumed via config.AWSBedrock.RoleARN, so the +// required bedrock:GetInferenceProfile permission belongs to that identity. +func ResolveBedrockModels(ctx context.Context, cfg config.AWSBedrock) (model, smallFastModel string, err error) { + if !isApplicationInferenceProfileARN(cfg.Model) && !isApplicationInferenceProfileARN(cfg.SmallFastModel) { + return cfg.Model, cfg.SmallFastModel, nil + } + + awsCfg, err := buildBedrockCredentials(ctx, cfg) + if err != nil { + return "", "", xerrors.Errorf("build bedrock credentials: %w", err) + } + + resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) + defer cancel() + resolveOne := func(configured string) (string, error) { if !isApplicationInferenceProfileARN(configured) { return configured, nil } - return resolveInferenceProfile(ctx, awsCfg, configured) + return resolveInferenceProfile(resolveCtx, awsCfg, configured) } model, err = resolveOne(cfg.Model) @@ -110,3 +127,29 @@ func resolveBedrockModels(ctx context.Context, cfg config.AWSBedrock, awsCfg aws } return model, smallFastModel, nil } + +// resolvedBedrockModels returns the model identities to serve with. A +// configured identifier that needs no resolution is its own identity; an +// application inference profile ARN requires the resolution stored with the +// provider. +func resolvedBedrockModels(cfg config.AWSBedrock) (model, smallFastModel string, err error) { + identity := func(configured, resolved, field string) (string, error) { + if !isApplicationInferenceProfileARN(configured) { + return configured, nil + } + if resolved == "" { + return "", xerrors.Errorf("%s %q is an application inference profile with no resolved model; re-save the provider to resolve it", field, configured) + } + return resolved, nil + } + + model, err = identity(cfg.Model, cfg.ResolvedModel, "model") + if err != nil { + return "", "", err + } + smallFastModel, err = identity(cfg.SmallFastModel, cfg.ResolvedSmallFastModel, "small fast model") + if err != nil { + return "", "", err + } + return model, smallFastModel, nil +} diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 2ac34857e97..1e551b8c468 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -124,15 +124,19 @@ func TestModelIDFromARN(t *testing.T) { } } -// TestNewAnthropic_InferenceProfileResolution drives the Bedrock -// GetInferenceProfile path against a mock endpoint. +// TestResolveBedrockModels drives the Bedrock GetInferenceProfile path against +// a mock endpoint. Resolution runs where a provider is written, so this covers +// what coderd calls, not what the gateway does when serving. // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetInferenceProfile.html // NOTE: no t.Parallel() because the subtests use t.Setenv. -func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { - const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" +func TestResolveBedrockModels(t *testing.T) { + const ( + profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" + smallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" + ) - bedrockCfg := func(model, smallFastModel string) *config.AWSBedrock { - return &config.AWSBedrock{ + bedrockCfg := func(model, smallFastModel string) config.AWSBedrock { + return config.AWSBedrock{ Region: "us-east-1", AccessKey: "test-key", AccessKeySecret: "test-secret", @@ -155,24 +159,22 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { return srv.URL, &got } - t.Run("resolved profile drives the model id", func(t *testing.T) { + t.Run("profile resolves to its model", func(t *testing.T) { url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8"}]}`)) }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) - // The profile stays the configured identifier so AWS attributes spend to it. - require.Equal(t, profileARN, p.bedrock.ConfiguredModel()) - require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) + require.Equal(t, "anthropic.claude-opus-4-8", model) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) require.Len(t, *paths, 1, "only the profile ARN is resolved") require.Contains(t, (*paths)[0], profileARN) }) - t.Run("failed resolution fails construction", func(t *testing.T) { + t.Run("failed resolution is an error", func(t *testing.T) { url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") @@ -181,50 +183,109 @@ func TestNewAnthropic_InferenceProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) - require.ErrorContains(t, err, "resolve bedrock models") + _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + require.ErrorContains(t, err, "resolve model") require.ErrorContains(t, err, "GetInferenceProfile") }) - t.Run("profile without a model fails construction", func(t *testing.T) { + t.Run("profile without a model is an error", func(t *testing.T) { url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"models":[]}`)) }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.ErrorContains(t, err, "references no model") }) t.Run("small fast profile resolves independently", func(t *testing.T) { - const smallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" - url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-haiku-4-5"}]}`)) }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) + model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) - require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) - require.Equal(t, smallFastProfileARN, p.bedrock.ConfiguredSmallFastModel()) + require.Equal(t, "eu.anthropic.claude-opus-4-8", model) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) require.Len(t, *paths, 1, "only the small fast profile ARN is resolved") require.Contains(t, (*paths)[0], smallFastProfileARN) }) - t.Run("plain model id needs no resolution", func(t *testing.T) { + t.Run("plain model ids need no resolution", func(t *testing.T) { url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { - t.Error("Bedrock called for a plain model id") + t.Error("Bedrock called for plain model ids") }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) + model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) + require.NoError(t, err) + require.Equal(t, "eu.anthropic.claude-opus-4-8", model) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + require.Empty(t, *paths) + }) +} + +// TestNewAnthropic_ServesStoredResolution covers what the gateway does with the +// resolution coderd stored: it serves it, and refuses to serve an opaque +// profile ARN that has none. +func TestNewAnthropic_ServesStoredResolution(t *testing.T) { + t.Parallel() + + const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" + + bedrockCfg := func(mutate func(*config.AWSBedrock)) *config.AWSBedrock { + cfg := &config.AWSBedrock{ + Region: "us-east-1", + AccessKey: "test-key", + AccessKeySecret: "test-secret", + Model: profileARN, + SmallFastModel: "anthropic.claude-haiku-4-5", + } + mutate(cfg) + return cfg + } + + t.Run("stored resolution drives the model id", func(t *testing.T) { + t.Parallel() + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(func(cfg *config.AWSBedrock) { + cfg.ResolvedModel = "anthropic.claude-opus-4-8" + })) + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) + // The profile stays the configured identifier so AWS attributes spend to it. + require.Equal(t, profileARN, p.bedrock.ConfiguredModel()) + require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) + }) + + t.Run("unresolved profile fails construction", func(t *testing.T) { + t.Parallel() + + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(func(*config.AWSBedrock) {})) + require.ErrorContains(t, err, "no resolved model") + }) + + t.Run("unresolved small fast profile fails construction", func(t *testing.T) { + t.Parallel() + + _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(func(cfg *config.AWSBedrock) { + cfg.Model = "eu.anthropic.claude-opus-4-8" + cfg.SmallFastModel = profileARN + })) + require.ErrorContains(t, err, "small fast model") + }) + + t.Run("plain model ids serve themselves", func(t *testing.T) { + t.Parallel() + + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(func(cfg *config.AWSBedrock) { + cfg.Model = "eu.anthropic.claude-opus-4-8" + })) require.NoError(t, err) require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ResolvedModel()) require.Equal(t, "eu.anthropic.claude-opus-4-8", p.bedrock.ConfiguredModel()) - require.Empty(t, *paths) }) } diff --git a/cli/aibridged.go b/cli/aibridged.go index 1dbc443c426..f5f9c4fa2e3 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -216,6 +216,8 @@ func protoToProviderSpec(pp *proto.AIProvider) aiProviderSpec { bedrock.RoleARN = b.GetRoleArn() bedrock.ExternalID = b.GetExternalId() bedrock.Protocol = codersdk.AIProviderBedrockProtocol(b.GetProtocol()) + bedrock.ResolvedModel = b.GetResolvedModel() + bedrock.ResolvedSmallFastModel = b.GetResolvedSmallFastModel() spec.Bedrock = new(bedrock) } return spec @@ -350,15 +352,17 @@ func bedrockConfig(baseURL string, bedrock *codersdk.AIProviderBedrockSettings) accessKey := ptr.NilToEmpty(bedrockSettings.AccessKey) accessKeySecret := ptr.NilToEmpty(bedrockSettings.AccessKeySecret) return &aibridge.AWSBedrockConfig{ - BaseURL: baseURL, - Region: bedrockSettings.Region, - AccessKey: accessKey, - AccessKeySecret: accessKeySecret, - Model: bedrockSettings.Model, - SmallFastModel: bedrockSettings.SmallFastModel, - RoleARN: bedrockSettings.RoleARN, - ExternalID: bedrockSettings.ExternalID, - Protocol: config.BedrockProtocol(bedrockSettings.ResolvedProtocol()), + BaseURL: baseURL, + Region: bedrockSettings.Region, + AccessKey: accessKey, + AccessKeySecret: accessKeySecret, + Model: bedrockSettings.Model, + SmallFastModel: bedrockSettings.SmallFastModel, + RoleARN: bedrockSettings.RoleARN, + ExternalID: bedrockSettings.ExternalID, + Protocol: config.BedrockProtocol(bedrockSettings.ResolvedProtocol()), + ResolvedModel: bedrockSettings.ResolvedModel, + ResolvedSmallFastModel: bedrockSettings.ResolvedSmallFastModel, } } diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index 71b40597751..e5b897e37fb 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -186,6 +186,14 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { // Generate the server-owned external ID when the provider assumes a role. ensureBedrockExternalID(&req.Settings) + // Resolve application inference profile ARNs before storing them. Doing it + // here means the operator learns immediately that a profile is wrong or + // unreachable, and the gateway never calls AWS to find out. + if err := api.resolveBedrockModels(ctx, &req.Settings); err != nil { + api.writeAIProviderResolutionError(ctx, rw, err) + return + } + settings, err := encodeAIProviderSettings(req.Settings) if err != nil { api.Logger.Error(ctx, "encode AI provider settings", slog.Error(err)) @@ -309,12 +317,21 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { idOrName := chi.URLParam(r, "idOrName") + // Resolve outside the transaction: it is a network call. The merge is redone + // inside the transaction against the row that gets written, and the + // resolution is applied only when the model identifiers still match. + resolvedPreview, hasResolvedPreview, err := api.previewResolvedBedrockSettings(ctx, idOrName, req.Settings) + if err != nil { + api.writeAIProviderResolutionError(ctx, rw, err) + return + } + var ( updated database.AIProvider keys []database.AIProviderKey keyChanges aiProviderKeyChanges ) - err := api.Database.InTx(func(tx database.Store) error { + err = api.Database.InTx(func(tx database.Store) error { old, err := lookupAIProvider(ctx, tx, idOrName) if err != nil { return err @@ -345,6 +362,13 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { // Generate the server-owned external ID when the provider assumes a role // and lacks one. ensureBedrockExternalID(&existing) + if req.Settings != nil && existing.Bedrock != nil { + if !hasResolvedPreview || !bedrockModelsMatch(existing, resolvedPreview) { + return errAIProviderChangedDuringUpdate + } + existing.Bedrock.ResolvedModel = resolvedPreview.Bedrock.ResolvedModel + existing.Bedrock.ResolvedSmallFastModel = resolvedPreview.Bedrock.ResolvedSmallFastModel + } settings, err := encodeAIProviderSettings(existing) if err != nil { return xerrors.Errorf("encode settings: %w", err) @@ -423,6 +447,12 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { }) return } + if errors.Is(err, errAIProviderChangedDuringUpdate) { + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "The AI provider changed while it was being updated. Retry the request.", + }) + return + } if errors.Is(err, errAIProviderKeyUnknown) { // Use the sentinel directly so the response message does not // leak the "execute transaction:" wrapper xerrors added on the @@ -536,6 +566,11 @@ var errAIProviderBedrockTypeMismatch = xerrors.New("bedrock settings are only va // patch may echo the stored value but not set a different one. var errAIProviderExternalIDReadOnly = xerrors.New("external_id is server-generated and cannot be changed") +// errAIProviderChangedDuringUpdate is the sentinel returned from inside the +// update transaction when the provider's model identifiers changed after they +// were resolved, so the resolution no longer describes what would be stored. +var errAIProviderChangedDuringUpdate = xerrors.New("provider changed while it was being updated, retry the request") + // errAIProviderInvalidName is returned from lookupAIProvider when the // idOrName parameter is neither a UUID nor a syntactically-valid name. // The handler translates this into a 400 so an integrator gets a hint diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go new file mode 100644 index 00000000000..8a183797f98 --- /dev/null +++ b/coderd/ai_providers_bedrock.go @@ -0,0 +1,139 @@ +package coderd + +import ( + "context" + "net/http" + + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/provider" + "github.com/coder/coder/v2/coderd/database/db2sdk" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/codersdk" +) + +// errAIProviderProfileUnresolvable wraps a failed Bedrock inference profile +// lookup so the write path reports it as a client-visible validation failure +// rather than an internal error. +var errAIProviderProfileUnresolvable = xerrors.New("resolve bedrock inference profile") + +// BedrockModelResolver resolves the configured Bedrock model identifiers of a +// provider to the model IDs the gateway records for capability detection, +// usage, and pricing. Identifiers that are not application inference profile +// ARNs resolve to themselves without calling AWS. +// +// Resolution runs here, where the provider is written, so the gateway never +// calls the Bedrock control plane: not at startup, not on reload, and not on a +// request. It is an interface so tests can supply results without AWS. +type BedrockModelResolver interface { + ResolveModels(ctx context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) +} + +// awsBedrockModelResolver resolves through the AWS Bedrock control plane using +// the provider's own credentials, including any assumed role. +type awsBedrockModelResolver struct{} + +func (awsBedrockModelResolver) ResolveModels(ctx context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) { + cfg := config.AWSBedrock{ + Region: settings.Region, + Model: settings.Model, + SmallFastModel: settings.SmallFastModel, + RoleARN: settings.RoleARN, + ExternalID: settings.ExternalID, + Protocol: config.BedrockProtocol(settings.ResolvedProtocol()), + } + if settings.AccessKey != nil { + cfg.AccessKey = *settings.AccessKey + } + if settings.AccessKeySecret != nil { + cfg.AccessKeySecret = *settings.AccessKeySecret + } + return provider.ResolveBedrockModels(ctx, cfg) +} + +func (api *API) bedrockModelResolver() BedrockModelResolver { + if api.AIProviderBedrockResolver != nil { + return api.AIProviderBedrockResolver + } + return awsBedrockModelResolver{} +} + +// resolveBedrockModels fills in the server-owned resolved identifiers on +// settings. A resolved value is stored only when it differs from the configured +// one, so plain model IDs stay unresolved and keep serving themselves. +// +// A failure is returned to the caller: an unresolvable profile must not be +// stored, because the gateway cannot tell what an opaque ARN refers to and +// would misshape every request made through it. +func (api *API) resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSettings) error { + if settings.Bedrock == nil { + return nil + } + + model, smallFastModel, err := api.bedrockModelResolver().ResolveModels(ctx, *settings.Bedrock) + if err != nil { + return xerrors.Errorf("%w: %w", errAIProviderProfileUnresolvable, err) + } + + settings.Bedrock.ResolvedModel = "" + if model != settings.Bedrock.Model { + settings.Bedrock.ResolvedModel = model + } + settings.Bedrock.ResolvedSmallFastModel = "" + if smallFastModel != settings.Bedrock.SmallFastModel { + settings.Bedrock.ResolvedSmallFastModel = smallFastModel + } + return nil +} + +// bedrockModelsMatch reports whether two settings configure the same model +// identifiers. The update path resolves against a snapshot taken outside the +// transaction, so it re-checks the merged settings before storing the result. +func bedrockModelsMatch(a, b codersdk.AIProviderSettings) bool { + if a.Bedrock == nil || b.Bedrock == nil { + return a.Bedrock == b.Bedrock + } + return a.Bedrock.Model == b.Bedrock.Model && a.Bedrock.SmallFastModel == b.Bedrock.SmallFastModel +} + +// previewResolvedBedrockSettings merges patch onto the stored settings of the +// named provider and resolves the result, so the write path can perform the +// AWS lookup outside its transaction. The boolean reports whether a preview was +// produced: there is nothing to resolve, or the provider cannot be read, in +// which case the transaction reports the failure with its own error handling. +func (api *API) previewResolvedBedrockSettings(ctx context.Context, idOrName string, patch *codersdk.AIProviderSettings) (codersdk.AIProviderSettings, bool, error) { + if patch == nil || patch.Bedrock == nil { + return codersdk.AIProviderSettings{}, false, nil + } + old, err := lookupAIProvider(ctx, api.Database, idOrName) + if err != nil { + //nolint:nilerr // The transaction reports lookup failures. + return codersdk.AIProviderSettings{}, false, nil + } + existing, err := db2sdk.AIProviderSettings(old.Settings) + if err != nil { + //nolint:nilerr // The transaction reports decode failures. + return codersdk.AIProviderSettings{}, false, nil + } + + preview := mergeAIProviderSettings(existing, *patch) + ensureBedrockExternalID(&preview) + if err := api.resolveBedrockModels(ctx, &preview); err != nil { + return codersdk.AIProviderSettings{}, false, err + } + return preview, true, nil +} + +// writeAIProviderResolutionError reports a failed Bedrock model resolution. The +// write is rejected rather than stored unresolved: the gateway cannot serve an +// opaque profile ARN, so accepting it would produce a provider that fails every +// request. +func (api *API) writeAIProviderResolutionError(ctx context.Context, rw http.ResponseWriter, err error) { + api.Logger.Warn(ctx, "resolve bedrock inference profile", slog.Error(err)) + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Could not resolve the Bedrock application inference profile. Check that the ARN is correct and that the AWS identity used by Coder is allowed bedrock:GetInferenceProfile.", + Detail: err.Error(), + }) +} diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go new file mode 100644 index 00000000000..174d1165314 --- /dev/null +++ b/coderd/ai_providers_bedrock_test.go @@ -0,0 +1,255 @@ +package coderd_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +const ( + testProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" + testSmallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" +) + +// stubBedrockResolver stands in for the AWS Bedrock control plane. It records +// the settings it was asked to resolve so tests can assert whether a write +// consulted AWS at all. +type stubBedrockResolver struct { + models map[string]string + err error + calls []codersdk.AIProviderBedrockSettings +} + +func (s *stubBedrockResolver) ResolveModels(_ context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) { + s.calls = append(s.calls, settings) + if s.err != nil { + return "", "", s.err + } + resolve := func(configured string) string { + if model, ok := s.models[configured]; ok { + return model + } + return configured + } + return resolve(settings.Model), resolve(settings.SmallFastModel), nil +} + +func bedrockSettings(model, smallFastModel string) *codersdk.AIProviderSettings { + accessKey := "test-key" + accessKeySecret := "test-secret" + return &codersdk.AIProviderSettings{ + Bedrock: &codersdk.AIProviderBedrockSettings{ + Region: "us-east-1", + AccessKey: &accessKey, + AccessKeySecret: &accessKeySecret, + Model: model, + SmallFastModel: smallFastModel, + }, + } +} + +func TestAIProvidersBedrockProfileResolution(t *testing.T) { + t.Parallel() + + newClient := func(t *testing.T, resolver coderd.BedrockModelResolver) *codersdk.Client { + t.Helper() + + client := coderdtest.New(t, &coderdtest.Options{AIProviderBedrockResolver: resolver}) + _ = coderdtest.CreateFirstUser(t, client) + return client + } + + t.Run("CreateStoresResolvedModels", func(t *testing.T) { + t.Parallel() + + resolver := &stubBedrockResolver{models: map[string]string{ + testProfileARN: "anthropic.claude-opus-4-8", + testSmallFastProfileARN: "anthropic.claude-haiku-4-5", + }} + client := newClient(t, resolver) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-profiles", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, testSmallFastProfileARN), + }) + require.NoError(t, err) + require.NotNil(t, created.Settings.Bedrock) + // The configured identifiers stay untouched: they remain the Bedrock + // invocation target, and AWS attributes spend to them. + require.Equal(t, testProfileARN, created.Settings.Bedrock.Model) + require.Equal(t, testSmallFastProfileARN, created.Settings.Bedrock.SmallFastModel) + require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) + require.Equal(t, "anthropic.claude-haiku-4-5", created.Settings.Bedrock.ResolvedSmallFastModel) + }) + + t.Run("CreateLeavesPlainModelIDsUnresolved", func(t *testing.T) { + t.Parallel() + + resolver := &stubBedrockResolver{} + client := newClient(t, resolver) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-plain", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + require.NotNil(t, created.Settings.Bedrock) + require.Empty(t, created.Settings.Bedrock.ResolvedModel) + require.Empty(t, created.Settings.Bedrock.ResolvedSmallFastModel) + }) + + t.Run("CreateRejectsUnresolvableProfile", func(t *testing.T) { + t.Parallel() + + resolver := &stubBedrockResolver{err: xerrors.New("AccessDeniedException: not authorized to perform bedrock:GetInferenceProfile")} + client := newClient(t, resolver) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + _, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-denied", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, 400, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Detail, "bedrock:GetInferenceProfile") + + //nolint:gocritic // Owner role is the audience for this endpoint. + providers, err := client.AIProviders(ctx) + require.NoError(t, err) + require.Empty(t, providers, "an unresolvable provider is not stored") + }) + + t.Run("CreateRejectsClientSuppliedResolution", func(t *testing.T) { + t.Parallel() + + client := newClient(t, &stubBedrockResolver{}) + ctx := testutil.Context(t, testutil.WaitLong) + + settings := bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") + settings.Bedrock.ResolvedModel = "anthropic.claude-opus-4-8" + + //nolint:gocritic // Owner role is the audience for this endpoint. + _, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-spoofed", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *settings, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, 400, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Error(), "resolved_model") + }) + + t.Run("UpdateReresolvesChangedProfile", func(t *testing.T) { + t.Parallel() + + resolver := &stubBedrockResolver{models: map[string]string{ + testProfileARN: "anthropic.claude-opus-4-8", + testSmallFastProfileARN: "anthropic.claude-haiku-4-5", + }} + client := newClient(t, resolver) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-update", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + require.Empty(t, created.Settings.Bedrock.ResolvedModel) + + //nolint:gocritic // Owner role is the audience for this endpoint. + updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + Settings: bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + require.Equal(t, testProfileARN, updated.Settings.Bedrock.Model) + require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) + require.Empty(t, updated.Settings.Bedrock.ResolvedSmallFastModel) + }) + + t.Run("UpdateClearsResolutionWhenProfileReplacedByModelID", func(t *testing.T) { + t.Parallel() + + resolver := &stubBedrockResolver{models: map[string]string{ + testProfileARN: "anthropic.claude-opus-4-8", + }} + client := newClient(t, resolver) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-replace", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) + + //nolint:gocritic // Owner role is the audience for this endpoint. + updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + Settings: bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + require.Empty(t, updated.Settings.Bedrock.ResolvedModel, "a plain model id resolves to itself") + }) + + t.Run("UpdateWithoutSettingsSkipsResolution", func(t *testing.T) { + t.Parallel() + + resolver := &stubBedrockResolver{models: map[string]string{ + testProfileARN: "anthropic.claude-opus-4-8", + }} + client := newClient(t, resolver) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-keep", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + callsAfterCreate := len(resolver.calls) + + enabled := false + //nolint:gocritic // Owner role is the audience for this endpoint. + updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + Enabled: &enabled, + }) + require.NoError(t, err) + require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) + require.Len(t, resolver.calls, callsAfterCreate, "an unrelated update does not call AWS") + }) +} diff --git a/coderd/aibridged/proto/aibridged.pb.go b/coderd/aibridged/proto/aibridged.pb.go index 4d099371e7d..c6024bdf13d 100644 --- a/coderd/aibridged/proto/aibridged.pb.go +++ b/coderd/aibridged/proto/aibridged.pb.go @@ -1674,6 +1674,13 @@ type AIProviderKindBedrock struct { // protocol selects the Bedrock wire protocol ("invoke-model" or "mantle"). // Empty falls back to invoke-model. Protocol string `protobuf:"bytes,8,opt,name=protocol,proto3" json:"protocol,omitempty"` + // resolved_model is the model ID behind model, which differs from it only + // when model is an application inference profile ARN. coderd resolves it when + // the provider is written, so the gateway never calls the Bedrock control + // plane. + ResolvedModel string `protobuf:"bytes,9,opt,name=resolved_model,json=resolvedModel,proto3" json:"resolved_model,omitempty"` + // resolved_small_fast_model is resolved_model for small_fast_model. + ResolvedSmallFastModel string `protobuf:"bytes,10,opt,name=resolved_small_fast_model,json=resolvedSmallFastModel,proto3" json:"resolved_small_fast_model,omitempty"` } func (x *AIProviderKindBedrock) Reset() { @@ -1764,6 +1771,20 @@ func (x *AIProviderKindBedrock) GetProtocol() string { return "" } +func (x *AIProviderKindBedrock) GetResolvedModel() string { + if x != nil { + return x.ResolvedModel + } + return "" +} + +func (x *AIProviderKindBedrock) GetResolvedSmallFastModel() string { + if x != nil { + return x.ResolvedSmallFastModel + } + return "" +} + var File_coderd_aibridged_proto_aibridged_proto protoreflect.FileDescriptor var file_coderd_aibridged_proto_aibridged_proto_rawDesc = []byte{ @@ -2060,7 +2081,7 @@ var file_coderd_aibridged_proto_aibridged_proto_rawDesc = []byte{ 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x4b, 0x69, 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x52, 0x07, 0x62, 0x65, 0x64, 0x72, 0x6f, - 0x63, 0x6b, 0x22, 0x92, 0x02, 0x0a, 0x15, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x63, 0x6b, 0x22, 0xf4, 0x02, 0x0a, 0x15, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x4b, 0x69, 0x6e, 0x64, 0x42, 0x65, 0x64, 0x72, 0x6f, 0x63, 0x6b, 0x12, 0x16, 0x0a, 0x06, 0x72, 0x65, 0x67, 0x69, 0x6f, 0x6e, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x72, 0x65, 0x67, 0x69, 0x6f, 0x6e, 0x12, 0x1d, 0x0a, 0x0a, 0x61, 0x63, 0x63, 0x65, 0x73, 0x73, 0x5f, 0x6b, @@ -2077,83 +2098,89 @@ var file_coderd_aibridged_proto_aibridged_proto_rawDesc = []byte{ 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x5f, 0x69, 0x64, 0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x65, 0x78, 0x74, 0x65, 0x72, 0x6e, 0x61, 0x6c, 0x49, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x18, 0x08, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x70, - 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x32, 0xa9, 0x04, 0x0a, 0x08, 0x52, 0x65, 0x63, 0x6f, - 0x72, 0x64, 0x65, 0x72, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, - 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, - 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x21, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x63, 0x6f, 0x6c, 0x12, 0x25, 0x0a, 0x0e, 0x72, 0x65, 0x73, 0x6f, 0x6c, + 0x76, 0x65, 0x64, 0x5f, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x09, 0x20, 0x01, 0x28, 0x09, 0x52, + 0x0d, 0x72, 0x65, 0x73, 0x6f, 0x6c, 0x76, 0x65, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x12, 0x39, + 0x0a, 0x19, 0x72, 0x65, 0x73, 0x6f, 0x6c, 0x76, 0x65, 0x64, 0x5f, 0x73, 0x6d, 0x61, 0x6c, 0x6c, + 0x5f, 0x66, 0x61, 0x73, 0x74, 0x5f, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x0a, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x16, 0x72, 0x65, 0x73, 0x6f, 0x6c, 0x76, 0x65, 0x64, 0x53, 0x6d, 0x61, 0x6c, 0x6c, + 0x46, 0x61, 0x73, 0x74, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x32, 0xa9, 0x04, 0x0a, 0x08, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x65, 0x72, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, + 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, - 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, - 0x68, 0x0a, 0x17, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, - 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x12, 0x25, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, - 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x1a, 0x26, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, - 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, - 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, 0x10, 0x52, 0x65, 0x63, - 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1e, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, - 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, - 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x56, - 0x0a, 0x11, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, - 0x61, 0x67, 0x65, 0x12, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, - 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, - 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, - 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1d, 0x2e, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, - 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, - 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, - 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, 0x63, 0x6f, - 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x12, 0x20, - 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, - 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, - 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, - 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, - 0x6e, 0x73, 0x65, 0x32, 0xeb, 0x01, 0x0a, 0x0f, 0x4d, 0x43, 0x50, 0x43, 0x6f, 0x6e, 0x66, 0x69, - 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x5c, 0x0a, 0x13, 0x47, 0x65, 0x74, 0x4d, 0x43, - 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x12, 0x21, - 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, - 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x1a, 0x22, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, - 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, 0x73, - 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x7a, 0x0a, 0x1d, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, - 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, - 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x12, 0x2b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, - 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, - 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x71, 0x75, - 0x65, 0x73, 0x74, 0x1a, 0x2c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, - 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, - 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, - 0x65, 0x32, 0xaa, 0x01, 0x0a, 0x0a, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x72, - 0x12, 0x47, 0x0a, 0x0c, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, - 0x12, 0x1a, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, - 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, 0x2e, 0x70, - 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, - 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, 0x10, 0x49, 0x73, 0x42, - 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x12, 0x1e, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, - 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, - 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xbc, - 0x01, 0x0a, 0x14, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, - 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x4d, 0x0a, 0x0e, 0x47, 0x65, 0x74, 0x41, 0x49, - 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1c, 0x2e, 0x70, 0x72, 0x6f, 0x74, + 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x21, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, + 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, + 0x65, 0x12, 0x68, 0x0a, 0x17, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, + 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x12, 0x25, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, + 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, + 0x65, 0x73, 0x74, 0x1a, 0x26, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, + 0x72, 0x64, 0x49, 0x6e, 0x74, 0x65, 0x72, 0x63, 0x65, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x45, 0x6e, + 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, 0x10, 0x52, + 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, + 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, + 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, + 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, + 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, + 0x12, 0x56, 0x0a, 0x11, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, + 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, + 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, + 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, + 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x52, 0x65, 0x63, 0x6f, + 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1d, 0x2e, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, + 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, + 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x59, 0x0a, 0x12, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, + 0x12, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x4d, + 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x1a, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, + 0x64, 0x4d, 0x6f, 0x64, 0x65, 0x6c, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x74, 0x52, 0x65, 0x73, + 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xeb, 0x01, 0x0a, 0x0f, 0x4d, 0x43, 0x50, 0x43, 0x6f, 0x6e, + 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x5c, 0x0a, 0x13, 0x47, 0x65, 0x74, + 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, + 0x12, 0x21, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, + 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, 0x65, 0x71, 0x75, + 0x65, 0x73, 0x74, 0x1a, 0x22, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x4d, + 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x73, 0x52, + 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x7a, 0x0a, 0x1d, 0x47, 0x65, 0x74, 0x4d, 0x43, + 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, + 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x12, 0x2b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x2e, 0x47, 0x65, 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, + 0x65, 0x73, 0x73, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, + 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x2c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, + 0x74, 0x4d, 0x43, 0x50, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x41, 0x63, 0x63, 0x65, 0x73, 0x73, + 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x42, 0x61, 0x74, 0x63, 0x68, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x32, 0xaa, 0x01, 0x0a, 0x0a, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, + 0x65, 0x72, 0x12, 0x47, 0x0a, 0x0c, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, 0x7a, + 0x65, 0x64, 0x12, 0x1a, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, + 0x68, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1b, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x41, 0x75, 0x74, 0x68, 0x6f, 0x72, 0x69, + 0x7a, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, 0x0a, 0x10, 0x49, + 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x12, + 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, + 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, + 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x49, 0x73, 0x42, 0x75, 0x64, 0x67, 0x65, 0x74, + 0x45, 0x78, 0x63, 0x65, 0x65, 0x64, 0x65, 0x64, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, + 0x32, 0xbc, 0x01, 0x0a, 0x14, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x43, 0x6f, 0x6e, + 0x66, 0x69, 0x67, 0x75, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x4d, 0x0a, 0x0e, 0x47, 0x65, 0x74, + 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1c, 0x2e, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, + 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, - 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1d, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, - 0x47, 0x65, 0x74, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, - 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x55, 0x0a, 0x10, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, - 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, - 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, - 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, - 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x30, 0x01, 0x42, 0x32, 0x5a, - 0x30, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x6f, 0x64, 0x65, - 0x72, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x76, 0x32, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, - 0x64, 0x2f, 0x61, 0x69, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x64, 0x2f, 0x70, 0x72, 0x6f, 0x74, - 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, + 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x55, 0x0a, 0x10, 0x57, 0x61, 0x74, 0x63, + 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x73, 0x12, 0x1e, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, + 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x70, + 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x57, 0x61, 0x74, 0x63, 0x68, 0x41, 0x49, 0x50, 0x72, 0x6f, 0x76, + 0x69, 0x64, 0x65, 0x72, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x30, 0x01, 0x42, + 0x32, 0x5a, 0x30, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x6f, + 0x64, 0x65, 0x72, 0x2f, 0x63, 0x6f, 0x64, 0x65, 0x72, 0x2f, 0x76, 0x32, 0x2f, 0x63, 0x6f, 0x64, + 0x65, 0x72, 0x64, 0x2f, 0x61, 0x69, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x64, 0x2f, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, } var ( diff --git a/coderd/aibridged/proto/aibridged.proto b/coderd/aibridged/proto/aibridged.proto index c5880ef594e..dcfc1db46d2 100644 --- a/coderd/aibridged/proto/aibridged.proto +++ b/coderd/aibridged/proto/aibridged.proto @@ -236,4 +236,11 @@ message AIProviderKindBedrock { // protocol selects the Bedrock wire protocol ("invoke-model" or "mantle"). // Empty falls back to invoke-model. string protocol = 8; + // resolved_model is the model ID behind model, which differs from it only + // when model is an application inference profile ARN. coderd resolves it when + // the provider is written, so the gateway never calls the Bedrock control + // plane. + string resolved_model = 9; + // resolved_small_fast_model is resolved_model for small_fast_model. + string resolved_small_fast_model = 10; } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 9f4f039ee75..cf76f9e27c0 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -1205,14 +1205,16 @@ func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey) ( } if settings.Bedrock != nil { p.Bedrock = &proto.AIProviderKindBedrock{ - Region: settings.Bedrock.Region, - AccessKey: ptr.NilToEmpty(settings.Bedrock.AccessKey), - AccessKeySecret: ptr.NilToEmpty(settings.Bedrock.AccessKeySecret), - Model: settings.Bedrock.Model, - SmallFastModel: settings.Bedrock.SmallFastModel, - RoleArn: settings.Bedrock.RoleARN, - ExternalId: settings.Bedrock.ExternalID, - Protocol: string(settings.Bedrock.Protocol), + Region: settings.Bedrock.Region, + AccessKey: ptr.NilToEmpty(settings.Bedrock.AccessKey), + AccessKeySecret: ptr.NilToEmpty(settings.Bedrock.AccessKeySecret), + Model: settings.Bedrock.Model, + SmallFastModel: settings.Bedrock.SmallFastModel, + RoleArn: settings.Bedrock.RoleARN, + ExternalId: settings.Bedrock.ExternalID, + Protocol: string(settings.Bedrock.Protocol), + ResolvedModel: settings.Bedrock.ResolvedModel, + ResolvedSmallFastModel: settings.Bedrock.ResolvedSmallFastModel, } } diff --git a/coderd/coderd.go b/coderd/coderd.go index 5051ac7dc0c..dc9a501b8dd 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -176,7 +176,11 @@ type Options struct { // CacheDir is used for caching files served by the API. CacheDir string - Auditor audit.Auditor + Auditor audit.Auditor + // AIProviderBedrockResolver resolves Bedrock application inference profile + // ARNs when an AI provider is written. Defaults to resolving through the AWS + // Bedrock control plane; tests substitute their own. + AIProviderBedrockResolver BedrockModelResolver ConnectionLogger connectionlog.ConnectionLogger AgentConnectionUpdateFrequency time.Duration AgentInactiveDisconnectTimeout time.Duration diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 8179431c2c6..6723215ea2d 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -131,9 +131,12 @@ type Options struct { AutobuildTicker <-chan time.Time AutobuildStats chan<- autobuild.Stats Auditor audit.Auditor - TLSCertificates []tls.Certificate - ExternalAuthConfigs []*externalauth.Config - TrialGenerator func(ctx context.Context, body codersdk.LicensorTrialRequest) error + // AIProviderBedrockResolver resolves Bedrock application inference profile + // ARNs when an AI provider is written. Tests set it to avoid calling AWS. + AIProviderBedrockResolver coderd.BedrockModelResolver + TLSCertificates []tls.Certificate + ExternalAuthConfigs []*externalauth.Config + TrialGenerator func(ctx context.Context, body codersdk.LicensorTrialRequest) error // MCPAllowedPrivateCIDRs exempts IP ranges from the MCP // SSRF guard for MCP server and OAuth2 traffic. Defaults to loopback so // tests can serve mock MCP and authorization servers via httptest. @@ -642,6 +645,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can UsageInserter: usageInserter, Auditor: options.Auditor, + AIProviderBedrockResolver: options.AIProviderBedrockResolver, ConnectionLogger: options.ConnectionLogger, AWSCertificates: options.AWSCertificates, AzureCertificates: options.AzureCertificates, diff --git a/codersdk/aiproviders.go b/codersdk/aiproviders.go index 6e81e3616a1..058351d782e 100644 --- a/codersdk/aiproviders.go +++ b/codersdk/aiproviders.go @@ -293,6 +293,7 @@ func (req CreateAIProviderRequest) Validate() []ValidationError { Detail: "external_id is server-generated and cannot be set", }) } + validations = append(validations, validateAIProviderBedrockResolvedModelsUnset(*req.Settings.Bedrock)...) validations = append(validations, validateAIProviderBedrockMantleRegion(*req.Settings.Bedrock)...) validations = append(validations, validateAIProviderBedrockModels(*req.Settings.Bedrock)...) } @@ -428,6 +429,26 @@ func validateAIProviderBedrockModels(b AIProviderBedrockSettings) []ValidationEr return validations } +// validateAIProviderBedrockResolvedModelsUnset rejects client-supplied resolved +// model identifiers on create. The server resolves them through AWS and owns +// the values, the same way it owns the STS external ID. +func validateAIProviderBedrockResolvedModelsUnset(b AIProviderBedrockSettings) []ValidationError { + var validations []ValidationError + if b.ResolvedModel != "" { + validations = append(validations, ValidationError{ + Field: "settings.resolved_model", + Detail: "resolved_model is server-resolved and cannot be set", + }) + } + if b.ResolvedSmallFastModel != "" { + validations = append(validations, ValidationError{ + Field: "settings.resolved_small_fast_model", + Detail: "resolved_small_fast_model is server-resolved and cannot be set", + }) + } + return validations +} + func validateAIProviderRoleARN(roleARN string) []ValidationError { if roleARN == "" { return nil diff --git a/codersdk/aiproviders_bedrock.go b/codersdk/aiproviders_bedrock.go index b3bc94e9e4b..d9c961cbb97 100644 --- a/codersdk/aiproviders_bedrock.go +++ b/codersdk/aiproviders_bedrock.go @@ -61,6 +61,33 @@ type AIProviderBedrockSettings struct { // AIProviderBedrockProtocolInvokeModel, so existing rows keep the legacy // behavior. Protocol AIProviderBedrockProtocol `json:"protocol,omitempty"` + // ResolvedModel is the Bedrock model ID behind Model. It differs from Model + // only when Model is an application inference profile ARN, whose identifier + // is opaque. The server resolves it through AWS when the provider is + // written and owns the value: create and update reject a client-supplied + // one that differs from the stored value. + ResolvedModel string `json:"resolved_model,omitempty"` + // ResolvedSmallFastModel is ResolvedModel for SmallFastModel. + ResolvedSmallFastModel string `json:"resolved_small_fast_model,omitempty"` +} + +// ModelIdentity returns the model ID the gateway records for usage, pricing, +// and capability detection. It is the resolved value when the configured +// identifier needed resolution, and the configured identifier otherwise. +func (b AIProviderBedrockSettings) ModelIdentity() string { + if b.ResolvedModel != "" { + return b.ResolvedModel + } + return b.Model +} + +// SmallFastModelIdentity is [AIProviderBedrockSettings.ModelIdentity] for the +// small/fast model. +func (b AIProviderBedrockSettings) SmallFastModelIdentity() string { + if b.ResolvedSmallFastModel != "" { + return b.ResolvedSmallFastModel + } + return b.SmallFastModel } // ResolvedProtocol returns the configured protocol, mapping the empty value to diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index 768d2e047cd..89522b509dd 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -251,10 +251,13 @@ AI Gateway passes the profile upstream so AWS records the attribution, while internally resolving and using the underlying model identity, including for usage pricing. -Resolution requires a `GetInferenceProfile` call, so the AWS identity used by -the gateway must have `bedrock:GetInferenceProfile` permission for the -profile. Providers configured with plain model identifiers do not need this -permission. If resolution fails, the provider is skipped. +Resolution requires a `GetInferenceProfile` call, which Coder makes when the +provider is saved, not when a request is served. The AWS identity Coder uses, +which is the provider's access keys when configured and otherwise the identity +of the Coder deployment, must have `bedrock:GetInferenceProfile` permission for +the profile. Saving fails when the lookup fails, so a profile that cannot be +resolved is never stored. Providers configured with plain model identifiers do +not need this permission. ### GitHub Copilot diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 1624e90af5b..ac2bc17325b 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -482,6 +482,18 @@ export interface AIProviderBedrockSettings { * behavior. */ readonly protocol?: AIProviderBedrockProtocol; + /** + * ResolvedModel is the Bedrock model ID behind Model. It differs from Model + * only when Model is an application inference profile ARN, whose identifier + * is opaque. The server resolves it through AWS when the provider is + * written and owns the value: create and update reject a client-supplied + * one that differs from the stored value. + */ + readonly resolved_model?: string; + /** + * ResolvedSmallFastModel is ResolvedModel for SmallFastModel. + */ + readonly resolved_small_fast_model?: string; } // From codersdk/aiproviders_bedrock.go From 7d51f6d33ec3ed830be8db32b7726d41f7c1418a Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 8 Sep 2026 19:38:54 +0000 Subject: [PATCH 02/25] refactor(aibridge/provider): resolve profiles from an explicit bedrock identity --- aibridge/provider/anthropic.go | 2 +- aibridge/provider/bedrock.go | 69 +++++++++++----- .../provider/bedrock_inference_profile.go | 56 ++----------- ...bedrock_inference_profile_internal_test.go | 78 ++++++------------- aibridge/provider/bedrock_internal_test.go | 22 +++--- coderd/ai_providers_bedrock.go | 56 +++++++++---- 6 files changed, 135 insertions(+), 148 deletions(-) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 68a349e3f12..72776021037 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -69,7 +69,7 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // so it is cheap to run at construction. var bedrock *messages.BedrockRuntime if bedrockCfg != nil { - awsCfg, err := buildBedrockCredentials(ctx, *bedrockCfg) + awsCfg, err := bedrockRuntimeCredentials(ctx, *bedrockCfg) if err != nil { return nil, xerrors.Errorf("build bedrock credentials: %w", err) } diff --git a/aibridge/provider/bedrock.go b/aibridge/provider/bedrock.go index 3ef73ae522c..3faa0b97fbd 100644 --- a/aibridge/provider/bedrock.go +++ b/aibridge/provider/bedrock.go @@ -19,7 +19,35 @@ import ( // A stable value keeps them identifiable in CloudTrail. const bedrockSessionName = "coder-aigateway" -// buildBedrockCredentials resolves the base identity and, when a role ARN +// BedrockIdentity is the subset of Bedrock provider settings that determines +// which AWS identity signs a request. It carries no endpoint or protocol +// choice, so it serves the data plane and the control plane equally. +type BedrockIdentity struct { + // Region resolves the endpoint and signs requests. It may be empty only + // when the caller supplies an endpoint of its own, in which case the AWS + // environment must supply the region. + Region string + // AccessKey and AccessKeySecret select static credentials. When either is + // empty the AWS SDK default credential chain resolves the base identity. + AccessKey string + AccessKeySecret string + // RoleARN, when set, is assumed via STS on top of the base identity. + RoleARN string + // ExternalID is sent as the STS external ID on the AssumeRole call. + ExternalID string +} + +func bedrockIdentity(cfg config.AWSBedrock) BedrockIdentity { + return BedrockIdentity{ + Region: cfg.Region, + AccessKey: cfg.AccessKey, + AccessKeySecret: cfg.AccessKeySecret, + RoleARN: cfg.RoleARN, + ExternalID: cfg.ExternalID, + } +} + +// BuildBedrockCredentials resolves the base identity and, when a role ARN // is configured, assumes that role via STS. The base identity is either // static keys or the AWS SDK default credential chain, which covers IRSA, // EKS Pod Identity, EC2 Instance Profile, and more. @@ -29,37 +57,33 @@ const bedrockSessionName = "coder-aigateway" // and identity rather than assembling their own. // // The credentials are wrapped in aws.NewCredentialsCache, which caches and -// rotates the resolved temporary credentials. buildBedrockCredentials should be +// rotates the resolved temporary credentials. BuildBedrockCredentials should be // called once when the Bedrock provider is constructed, and the returned // Credential Provider should be shared across all LLM requests to the Bedrock // Provider, so per-request credential retrieval is served from this cache // rather than re-resolving (and re-assuming) on every request. No network call // is made here: the base identity and any AssumeRole are resolved lazily on // first retrieval. -func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Config, error) { - if cfg.Region == "" && cfg.BaseURL == "" { - return aws.Config{}, xerrors.New("region or base url required") - } - +func BuildBedrockCredentials(ctx context.Context, id BedrockIdentity) (aws.Config, error) { var loadOpts []func(*awsconfig.LoadOptions) error - if cfg.Region != "" { - loadOpts = append(loadOpts, awsconfig.WithRegion(cfg.Region)) + if id.Region != "" { + loadOpts = append(loadOpts, awsconfig.WithRegion(id.Region)) } // Use static credentials when explicitly provided, otherwise fall back to // the SDK default credential chain. switch { // Both set: use static credentials directly. - case cfg.AccessKey != "" && cfg.AccessKeySecret != "": + case id.AccessKey != "" && id.AccessKeySecret != "": loadOpts = append(loadOpts, awsconfig.WithCredentialsProvider( credentials.NewStaticCredentialsProvider( - cfg.AccessKey, - cfg.AccessKeySecret, + id.AccessKey, + id.AccessKeySecret, "", ), )) // Only one set: misconfiguration. - case cfg.AccessKey != "" || cfg.AccessKeySecret != "": + case id.AccessKey != "" || id.AccessKeySecret != "": return aws.Config{}, xerrors.New("both access key and access key secret must be provided together") // Neither set: SDK default credential chain resolves the base identity. default: @@ -73,7 +97,7 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Co // Assuming a role calls STS, which needs a region to resolve its endpoint. // The region may come from the config or the AWS environment; if neither // supplies one, fail here. - if cfg.RoleARN != "" && base.Region == "" { + if id.RoleARN != "" && base.Region == "" { return aws.Config{}, xerrors.New("region is required to assume a role, but was not specified") } @@ -83,7 +107,7 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Co // is already cache-wrapped, so only the AssumeRoleProvider is wrapped with a // cache to avoid re-assuming the role on every request. credsProvider := base.Credentials - if cfg.RoleARN != "" { + if id.RoleARN != "" { // Disable keep-alive on the STS client so each AssumeRole opens a // fresh connection. Observed: with keep-alive, AssumeRole calls reuse // one connection pinned to a single STS endpoint, and after a @@ -101,10 +125,10 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Co t.DisableKeepAlives = true }) }) - credsProvider = stscreds.NewAssumeRoleProvider(stsClient, cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) { + credsProvider = stscreds.NewAssumeRoleProvider(stsClient, id.RoleARN, func(o *stscreds.AssumeRoleOptions) { o.RoleSessionName = bedrockSessionName - if cfg.ExternalID != "" { - o.ExternalID = aws.String(cfg.ExternalID) + if id.ExternalID != "" { + o.ExternalID = aws.String(id.ExternalID) } }) credsProvider = aws.NewCredentialsCache(credsProvider) @@ -115,3 +139,12 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Co base.Credentials = credsProvider return base, nil } + +// bedrockRuntimeCredentials is [BuildBedrockCredentials] for a full provider +// configuration, rejecting a config that gives the runtime no endpoint at all. +func bedrockRuntimeCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Config, error) { + if cfg.Region == "" && cfg.BaseURL == "" { + return aws.Config{}, xerrors.New("region or base url required") + } + return BuildBedrockCredentials(ctx, bedrockIdentity(cfg)) +} diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 6b9bc88762c..7c73af248e7 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -3,7 +3,6 @@ package provider import ( "context" "strings" - "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/aws/arn" @@ -21,17 +20,12 @@ const bedrockService = "bedrock" // Bedrock spend to a team or workload via cost allocation tags. const applicationInferenceProfileResourceType = "application-inference-profile" -// inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls made -// while writing a provider, which also cover the first credential resolution -// (STS/IRSA). -const inferenceProfileResolutionTimeout = 30 * time.Second - -// isApplicationInferenceProfileARN reports whether model is an application +// IsApplicationInferenceProfileARN reports whether model is an application // inference profile ARN, whose identifier is opaque and must be resolved // through AWS. Plain model IDs and system-defined inference profile ARNs, which // AWS documents as {geoRegion}.{modelId}, embed the model ID and need no // lookup. -func isApplicationInferenceProfileARN(model string) bool { +func IsApplicationInferenceProfileARN(model string) bool { parsed, err := arn.Parse(model) if err != nil || parsed.Service != bedrockService { return false @@ -40,7 +34,7 @@ func isApplicationInferenceProfileARN(model string) bool { return ok && resourceType == applicationInferenceProfileResourceType } -// resolveInferenceProfile returns the Bedrock model ID behind an application +// ResolveInferenceProfile returns the Bedrock model ID behind an application // inference profile ARN. // // awsCfg carries the identity that invokes Bedrock, including any role assumed @@ -50,7 +44,7 @@ func isApplicationInferenceProfileARN(model string) bool { // A profile that wraps a cross-region system-defined profile lists one model // per region. Those entries differ only in the ARN region, which the model ID // does not carry, so any entry resolves to the same model. -func resolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error) { +func ResolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error) { client := bedrock.NewFromConfig(awsCfg) out, err := client.GetInferenceProfile(ctx, &bedrock.GetInferenceProfileInput{ @@ -88,53 +82,13 @@ func modelIDFromARN(modelARN string) (string, error) { return model, nil } -// ResolveBedrockModels resolves the configured model identifiers to the model -// IDs used for capability detection, usage recording, and pricing. Identifiers -// that are not application inference profile ARNs are returned unchanged and -// cost no AWS call. -// -// It runs where a Bedrock provider is written rather than where it is served, -// so the gateway never calls the Bedrock control plane. The identity comes from -// cfg, including any role assumed via config.AWSBedrock.RoleARN, so the -// required bedrock:GetInferenceProfile permission belongs to that identity. -func ResolveBedrockModels(ctx context.Context, cfg config.AWSBedrock) (model, smallFastModel string, err error) { - if !isApplicationInferenceProfileARN(cfg.Model) && !isApplicationInferenceProfileARN(cfg.SmallFastModel) { - return cfg.Model, cfg.SmallFastModel, nil - } - - awsCfg, err := buildBedrockCredentials(ctx, cfg) - if err != nil { - return "", "", xerrors.Errorf("build bedrock credentials: %w", err) - } - - resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) - defer cancel() - - resolveOne := func(configured string) (string, error) { - if !isApplicationInferenceProfileARN(configured) { - return configured, nil - } - return resolveInferenceProfile(resolveCtx, awsCfg, configured) - } - - model, err = resolveOne(cfg.Model) - if err != nil { - return "", "", xerrors.Errorf("resolve model: %w", err) - } - smallFastModel, err = resolveOne(cfg.SmallFastModel) - if err != nil { - return "", "", xerrors.Errorf("resolve small fast model: %w", err) - } - return model, smallFastModel, nil -} - // resolvedBedrockModels returns the model identities to serve with. A // configured identifier that needs no resolution is its own identity; an // application inference profile ARN requires the resolution stored with the // provider. func resolvedBedrockModels(cfg config.AWSBedrock) (model, smallFastModel string, err error) { identity := func(configured, resolved, field string) (string, error) { - if !isApplicationInferenceProfileARN(configured) { + if !IsApplicationInferenceProfileARN(configured) { return configured, nil } if resolved == "" { diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 1e551b8c468..f7942638bd6 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -6,6 +6,7 @@ import ( "net/http/httptest" "testing" + "github.com/aws/aws-sdk-go-v2/aws" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/aibridge/config" @@ -72,7 +73,7 @@ func TestIsApplicationInferenceProfileARN(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - require.Equal(t, tt.want, isApplicationInferenceProfileARN(tt.model)) + require.Equal(t, tt.want, IsApplicationInferenceProfileARN(tt.model)) }) } } @@ -124,26 +125,13 @@ func TestModelIDFromARN(t *testing.T) { } } -// TestResolveBedrockModels drives the Bedrock GetInferenceProfile path against -// a mock endpoint. Resolution runs where a provider is written, so this covers -// what coderd calls, not what the gateway does when serving. +// TestResolveInferenceProfile drives the Bedrock GetInferenceProfile path +// against a mock endpoint. Resolution runs where a provider is written, so this +// covers the primitive coderd calls, not what the gateway does when serving. // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetInferenceProfile.html // NOTE: no t.Parallel() because the subtests use t.Setenv. -func TestResolveBedrockModels(t *testing.T) { - const ( - profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" - smallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" - ) - - bedrockCfg := func(model, smallFastModel string) config.AWSBedrock { - return config.AWSBedrock{ - Region: "us-east-1", - AccessKey: "test-key", - AccessKeySecret: "test-secret", - Model: model, - SmallFastModel: smallFastModel, - } - } +func TestResolveInferenceProfile(t *testing.T) { + const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" // mockBedrock serves the Bedrock control-plane API and records the paths it // receives. Callers point the SDK at the returned URL. @@ -159,6 +147,18 @@ func TestResolveBedrockModels(t *testing.T) { return srv.URL, &got } + credentials := func(t *testing.T) aws.Config { + t.Helper() + + awsCfg, err := BuildBedrockCredentials(context.Background(), BedrockIdentity{ + Region: "us-east-1", + AccessKey: "test-key", + AccessKeySecret: "test-secret", + }) + require.NoError(t, err) + return awsCfg + } + t.Run("profile resolves to its model", func(t *testing.T) { url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") @@ -166,15 +166,14 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + model, err := ResolveInferenceProfile(context.Background(), credentials(t), profileARN) require.NoError(t, err) require.Equal(t, "anthropic.claude-opus-4-8", model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) - require.Len(t, *paths, 1, "only the profile ARN is resolved") + require.Len(t, *paths, 1) require.Contains(t, (*paths)[0], profileARN) }) - t.Run("failed resolution is an error", func(t *testing.T) { + t.Run("failed lookup is an error", func(t *testing.T) { url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") @@ -183,8 +182,7 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) - require.ErrorContains(t, err, "resolve model") + _, err := ResolveInferenceProfile(context.Background(), credentials(t), profileARN) require.ErrorContains(t, err, "GetInferenceProfile") }) @@ -195,37 +193,9 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + _, err := ResolveInferenceProfile(context.Background(), credentials(t), profileARN) require.ErrorContains(t, err, "references no model") }) - - t.Run("small fast profile resolves independently", func(t *testing.T) { - url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-haiku-4-5"}]}`)) - }) - t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - - model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) - require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) - require.Len(t, *paths, 1, "only the small fast profile ARN is resolved") - require.Contains(t, (*paths)[0], smallFastProfileARN) - }) - - t.Run("plain model ids need no resolution", func(t *testing.T) { - url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { - t.Error("Bedrock called for plain model ids") - }) - t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - - model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) - require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) - require.Empty(t, *paths) - }) } // TestNewAnthropic_ServesStoredResolution covers what the gateway does with the diff --git a/aibridge/provider/bedrock_internal_test.go b/aibridge/provider/bedrock_internal_test.go index 1577f3b5262..6c542bff4e2 100644 --- a/aibridge/provider/bedrock_internal_test.go +++ b/aibridge/provider/bedrock_internal_test.go @@ -49,7 +49,7 @@ func TestBuildBedrockCredentialsValidation(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - _, err := buildBedrockCredentials(context.Background(), tt.cfg) + _, err := bedrockRuntimeCredentials(context.Background(), tt.cfg) require.Error(t, err) require.Contains(t, err.Error(), tt.errorMsg) }) @@ -60,7 +60,7 @@ func TestBuildBedrockCredentialsValidation(t *testing.T) { func TestBuildBedrockCredentialsStatic(t *testing.T) { t.Parallel() - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", AccessKey: "test-key", AccessKeySecret: "test-secret", @@ -132,10 +132,10 @@ func TestBuildBedrockCredentialsDefaultChain(t *testing.T) { t.Setenv(key, val) } - // buildBedrockCredentials only wires up the provider chain; it + // bedrockRuntimeCredentials only wires up the provider chain; it // does not resolve credentials, so it succeeds regardless of // credential availability. Resolution failures surface on Retrieve. - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", }) require.NoError(t, err) @@ -193,7 +193,7 @@ func TestBuildBedrockCredentialsAssumeRole(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -254,7 +254,7 @@ func TestBuildBedrockCredentialsAssumeRoleExternalID(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", ExternalID: tt.externalID, @@ -293,7 +293,7 @@ func TestBuildBedrockCredentialsAssumeRoleError(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -337,7 +337,7 @@ func TestBuildBedrockCredentialsAssumeRoleCaches(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -386,7 +386,7 @@ func TestBuildBedrockCredentialsAssumeRoleRefreshesOnExpiry(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -415,7 +415,7 @@ func TestBuildBedrockCredentialsAssumeRoleRequiresRegion(t *testing.T) { t.Setenv("AWS_SHARED_CREDENTIALS_FILE", "/dev/null") t.Setenv("AWS_EC2_METADATA_DISABLED", "true") - _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + _, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ BaseURL: "https://bedrock-runtime.example.com", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -430,7 +430,7 @@ func TestBuildBedrockCredentialsAssumeRoleRegionFromEnv(t *testing.T) { t.Setenv("AWS_REGION", "us-west-2") // BaseURL set with no explicit region: the region comes from AWS_REGION. - awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ BaseURL: "https://bedrock-runtime.example.com", RoleARN: "arn:aws:iam::123456789012:role/target", }) diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index 8a183797f98..e275d7f10b9 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -3,17 +3,23 @@ package coderd import ( "context" "net/http" + "time" "golang.org/x/xerrors" "cdr.dev/slog/v3" - "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/provider" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" ) +// inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls a +// single write makes, which also cover the first credential resolution +// (STS/IRSA). +const inferenceProfileResolutionTimeout = 30 * time.Second + // errAIProviderProfileUnresolvable wraps a failed Bedrock inference profile // lookup so the write path reports it as a client-visible validation failure // rather than an internal error. @@ -35,22 +41,46 @@ type BedrockModelResolver interface { // the provider's own credentials, including any assumed role. type awsBedrockModelResolver struct{} +// ResolveModels resolves the configured identifiers to the model IDs the +// gateway records for capability detection, usage, and pricing. Only +// application inference profile ARNs are opaque, so only they cost an AWS call; +// every other identifier resolves to itself. func (awsBedrockModelResolver) ResolveModels(ctx context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) { - cfg := config.AWSBedrock{ - Region: settings.Region, - Model: settings.Model, - SmallFastModel: settings.SmallFastModel, - RoleARN: settings.RoleARN, - ExternalID: settings.ExternalID, - Protocol: config.BedrockProtocol(settings.ResolvedProtocol()), + if !provider.IsApplicationInferenceProfileARN(settings.Model) && + !provider.IsApplicationInferenceProfileARN(settings.SmallFastModel) { + return settings.Model, settings.SmallFastModel, nil } - if settings.AccessKey != nil { - cfg.AccessKey = *settings.AccessKey + + ctx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) + defer cancel() + + awsCfg, err := provider.BuildBedrockCredentials(ctx, provider.BedrockIdentity{ + Region: settings.Region, + AccessKey: ptr.NilToEmpty(settings.AccessKey), + AccessKeySecret: ptr.NilToEmpty(settings.AccessKeySecret), + RoleARN: settings.RoleARN, + ExternalID: settings.ExternalID, + }) + if err != nil { + return "", "", xerrors.Errorf("build bedrock credentials: %w", err) } - if settings.AccessKeySecret != nil { - cfg.AccessKeySecret = *settings.AccessKeySecret + + resolve := func(configured string) (string, error) { + if !provider.IsApplicationInferenceProfileARN(configured) { + return configured, nil + } + return provider.ResolveInferenceProfile(ctx, awsCfg, configured) + } + + model, err = resolve(settings.Model) + if err != nil { + return "", "", xerrors.Errorf("resolve model: %w", err) + } + smallFastModel, err = resolve(settings.SmallFastModel) + if err != nil { + return "", "", xerrors.Errorf("resolve small fast model: %w", err) } - return provider.ResolveBedrockModels(ctx, cfg) + return model, smallFastModel, nil } func (api *API) bedrockModelResolver() BedrockModelResolver { From 2fde4d3fca79b8499cdc75b34536f66762adee14 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 8 Sep 2026 19:57:27 +0000 Subject: [PATCH 03/25] Revert "refactor(aibridge/provider): resolve profiles from an explicit bedrock identity" This reverts commit 7d51f6d33ec3ed830be8db32b7726d41f7c1418a. --- aibridge/provider/anthropic.go | 2 +- aibridge/provider/bedrock.go | 69 +++++----------- .../provider/bedrock_inference_profile.go | 56 +++++++++++-- ...bedrock_inference_profile_internal_test.go | 78 +++++++++++++------ aibridge/provider/bedrock_internal_test.go | 22 +++--- coderd/ai_providers_bedrock.go | 56 ++++--------- 6 files changed, 148 insertions(+), 135 deletions(-) diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 72776021037..68a349e3f12 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -69,7 +69,7 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. // so it is cheap to run at construction. var bedrock *messages.BedrockRuntime if bedrockCfg != nil { - awsCfg, err := bedrockRuntimeCredentials(ctx, *bedrockCfg) + awsCfg, err := buildBedrockCredentials(ctx, *bedrockCfg) if err != nil { return nil, xerrors.Errorf("build bedrock credentials: %w", err) } diff --git a/aibridge/provider/bedrock.go b/aibridge/provider/bedrock.go index 3faa0b97fbd..3ef73ae522c 100644 --- a/aibridge/provider/bedrock.go +++ b/aibridge/provider/bedrock.go @@ -19,35 +19,7 @@ import ( // A stable value keeps them identifiable in CloudTrail. const bedrockSessionName = "coder-aigateway" -// BedrockIdentity is the subset of Bedrock provider settings that determines -// which AWS identity signs a request. It carries no endpoint or protocol -// choice, so it serves the data plane and the control plane equally. -type BedrockIdentity struct { - // Region resolves the endpoint and signs requests. It may be empty only - // when the caller supplies an endpoint of its own, in which case the AWS - // environment must supply the region. - Region string - // AccessKey and AccessKeySecret select static credentials. When either is - // empty the AWS SDK default credential chain resolves the base identity. - AccessKey string - AccessKeySecret string - // RoleARN, when set, is assumed via STS on top of the base identity. - RoleARN string - // ExternalID is sent as the STS external ID on the AssumeRole call. - ExternalID string -} - -func bedrockIdentity(cfg config.AWSBedrock) BedrockIdentity { - return BedrockIdentity{ - Region: cfg.Region, - AccessKey: cfg.AccessKey, - AccessKeySecret: cfg.AccessKeySecret, - RoleARN: cfg.RoleARN, - ExternalID: cfg.ExternalID, - } -} - -// BuildBedrockCredentials resolves the base identity and, when a role ARN +// buildBedrockCredentials resolves the base identity and, when a role ARN // is configured, assumes that role via STS. The base identity is either // static keys or the AWS SDK default credential chain, which covers IRSA, // EKS Pod Identity, EC2 Instance Profile, and more. @@ -57,33 +29,37 @@ func bedrockIdentity(cfg config.AWSBedrock) BedrockIdentity { // and identity rather than assembling their own. // // The credentials are wrapped in aws.NewCredentialsCache, which caches and -// rotates the resolved temporary credentials. BuildBedrockCredentials should be +// rotates the resolved temporary credentials. buildBedrockCredentials should be // called once when the Bedrock provider is constructed, and the returned // Credential Provider should be shared across all LLM requests to the Bedrock // Provider, so per-request credential retrieval is served from this cache // rather than re-resolving (and re-assuming) on every request. No network call // is made here: the base identity and any AssumeRole are resolved lazily on // first retrieval. -func BuildBedrockCredentials(ctx context.Context, id BedrockIdentity) (aws.Config, error) { +func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Config, error) { + if cfg.Region == "" && cfg.BaseURL == "" { + return aws.Config{}, xerrors.New("region or base url required") + } + var loadOpts []func(*awsconfig.LoadOptions) error - if id.Region != "" { - loadOpts = append(loadOpts, awsconfig.WithRegion(id.Region)) + if cfg.Region != "" { + loadOpts = append(loadOpts, awsconfig.WithRegion(cfg.Region)) } // Use static credentials when explicitly provided, otherwise fall back to // the SDK default credential chain. switch { // Both set: use static credentials directly. - case id.AccessKey != "" && id.AccessKeySecret != "": + case cfg.AccessKey != "" && cfg.AccessKeySecret != "": loadOpts = append(loadOpts, awsconfig.WithCredentialsProvider( credentials.NewStaticCredentialsProvider( - id.AccessKey, - id.AccessKeySecret, + cfg.AccessKey, + cfg.AccessKeySecret, "", ), )) // Only one set: misconfiguration. - case id.AccessKey != "" || id.AccessKeySecret != "": + case cfg.AccessKey != "" || cfg.AccessKeySecret != "": return aws.Config{}, xerrors.New("both access key and access key secret must be provided together") // Neither set: SDK default credential chain resolves the base identity. default: @@ -97,7 +73,7 @@ func BuildBedrockCredentials(ctx context.Context, id BedrockIdentity) (aws.Confi // Assuming a role calls STS, which needs a region to resolve its endpoint. // The region may come from the config or the AWS environment; if neither // supplies one, fail here. - if id.RoleARN != "" && base.Region == "" { + if cfg.RoleARN != "" && base.Region == "" { return aws.Config{}, xerrors.New("region is required to assume a role, but was not specified") } @@ -107,7 +83,7 @@ func BuildBedrockCredentials(ctx context.Context, id BedrockIdentity) (aws.Confi // is already cache-wrapped, so only the AssumeRoleProvider is wrapped with a // cache to avoid re-assuming the role on every request. credsProvider := base.Credentials - if id.RoleARN != "" { + if cfg.RoleARN != "" { // Disable keep-alive on the STS client so each AssumeRole opens a // fresh connection. Observed: with keep-alive, AssumeRole calls reuse // one connection pinned to a single STS endpoint, and after a @@ -125,10 +101,10 @@ func BuildBedrockCredentials(ctx context.Context, id BedrockIdentity) (aws.Confi t.DisableKeepAlives = true }) }) - credsProvider = stscreds.NewAssumeRoleProvider(stsClient, id.RoleARN, func(o *stscreds.AssumeRoleOptions) { + credsProvider = stscreds.NewAssumeRoleProvider(stsClient, cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) { o.RoleSessionName = bedrockSessionName - if id.ExternalID != "" { - o.ExternalID = aws.String(id.ExternalID) + if cfg.ExternalID != "" { + o.ExternalID = aws.String(cfg.ExternalID) } }) credsProvider = aws.NewCredentialsCache(credsProvider) @@ -139,12 +115,3 @@ func BuildBedrockCredentials(ctx context.Context, id BedrockIdentity) (aws.Confi base.Credentials = credsProvider return base, nil } - -// bedrockRuntimeCredentials is [BuildBedrockCredentials] for a full provider -// configuration, rejecting a config that gives the runtime no endpoint at all. -func bedrockRuntimeCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Config, error) { - if cfg.Region == "" && cfg.BaseURL == "" { - return aws.Config{}, xerrors.New("region or base url required") - } - return BuildBedrockCredentials(ctx, bedrockIdentity(cfg)) -} diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 7c73af248e7..6b9bc88762c 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -3,6 +3,7 @@ package provider import ( "context" "strings" + "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/aws/arn" @@ -20,12 +21,17 @@ const bedrockService = "bedrock" // Bedrock spend to a team or workload via cost allocation tags. const applicationInferenceProfileResourceType = "application-inference-profile" -// IsApplicationInferenceProfileARN reports whether model is an application +// inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls made +// while writing a provider, which also cover the first credential resolution +// (STS/IRSA). +const inferenceProfileResolutionTimeout = 30 * time.Second + +// isApplicationInferenceProfileARN reports whether model is an application // inference profile ARN, whose identifier is opaque and must be resolved // through AWS. Plain model IDs and system-defined inference profile ARNs, which // AWS documents as {geoRegion}.{modelId}, embed the model ID and need no // lookup. -func IsApplicationInferenceProfileARN(model string) bool { +func isApplicationInferenceProfileARN(model string) bool { parsed, err := arn.Parse(model) if err != nil || parsed.Service != bedrockService { return false @@ -34,7 +40,7 @@ func IsApplicationInferenceProfileARN(model string) bool { return ok && resourceType == applicationInferenceProfileResourceType } -// ResolveInferenceProfile returns the Bedrock model ID behind an application +// resolveInferenceProfile returns the Bedrock model ID behind an application // inference profile ARN. // // awsCfg carries the identity that invokes Bedrock, including any role assumed @@ -44,7 +50,7 @@ func IsApplicationInferenceProfileARN(model string) bool { // A profile that wraps a cross-region system-defined profile lists one model // per region. Those entries differ only in the ARN region, which the model ID // does not carry, so any entry resolves to the same model. -func ResolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error) { +func resolveInferenceProfile(ctx context.Context, awsCfg aws.Config, profileARN string) (string, error) { client := bedrock.NewFromConfig(awsCfg) out, err := client.GetInferenceProfile(ctx, &bedrock.GetInferenceProfileInput{ @@ -82,13 +88,53 @@ func modelIDFromARN(modelARN string) (string, error) { return model, nil } +// ResolveBedrockModels resolves the configured model identifiers to the model +// IDs used for capability detection, usage recording, and pricing. Identifiers +// that are not application inference profile ARNs are returned unchanged and +// cost no AWS call. +// +// It runs where a Bedrock provider is written rather than where it is served, +// so the gateway never calls the Bedrock control plane. The identity comes from +// cfg, including any role assumed via config.AWSBedrock.RoleARN, so the +// required bedrock:GetInferenceProfile permission belongs to that identity. +func ResolveBedrockModels(ctx context.Context, cfg config.AWSBedrock) (model, smallFastModel string, err error) { + if !isApplicationInferenceProfileARN(cfg.Model) && !isApplicationInferenceProfileARN(cfg.SmallFastModel) { + return cfg.Model, cfg.SmallFastModel, nil + } + + awsCfg, err := buildBedrockCredentials(ctx, cfg) + if err != nil { + return "", "", xerrors.Errorf("build bedrock credentials: %w", err) + } + + resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) + defer cancel() + + resolveOne := func(configured string) (string, error) { + if !isApplicationInferenceProfileARN(configured) { + return configured, nil + } + return resolveInferenceProfile(resolveCtx, awsCfg, configured) + } + + model, err = resolveOne(cfg.Model) + if err != nil { + return "", "", xerrors.Errorf("resolve model: %w", err) + } + smallFastModel, err = resolveOne(cfg.SmallFastModel) + if err != nil { + return "", "", xerrors.Errorf("resolve small fast model: %w", err) + } + return model, smallFastModel, nil +} + // resolvedBedrockModels returns the model identities to serve with. A // configured identifier that needs no resolution is its own identity; an // application inference profile ARN requires the resolution stored with the // provider. func resolvedBedrockModels(cfg config.AWSBedrock) (model, smallFastModel string, err error) { identity := func(configured, resolved, field string) (string, error) { - if !IsApplicationInferenceProfileARN(configured) { + if !isApplicationInferenceProfileARN(configured) { return configured, nil } if resolved == "" { diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index f7942638bd6..1e551b8c468 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -6,7 +6,6 @@ import ( "net/http/httptest" "testing" - "github.com/aws/aws-sdk-go-v2/aws" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/aibridge/config" @@ -73,7 +72,7 @@ func TestIsApplicationInferenceProfileARN(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - require.Equal(t, tt.want, IsApplicationInferenceProfileARN(tt.model)) + require.Equal(t, tt.want, isApplicationInferenceProfileARN(tt.model)) }) } } @@ -125,13 +124,26 @@ func TestModelIDFromARN(t *testing.T) { } } -// TestResolveInferenceProfile drives the Bedrock GetInferenceProfile path -// against a mock endpoint. Resolution runs where a provider is written, so this -// covers the primitive coderd calls, not what the gateway does when serving. +// TestResolveBedrockModels drives the Bedrock GetInferenceProfile path against +// a mock endpoint. Resolution runs where a provider is written, so this covers +// what coderd calls, not what the gateway does when serving. // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetInferenceProfile.html // NOTE: no t.Parallel() because the subtests use t.Setenv. -func TestResolveInferenceProfile(t *testing.T) { - const profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" +func TestResolveBedrockModels(t *testing.T) { + const ( + profileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5" + smallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" + ) + + bedrockCfg := func(model, smallFastModel string) config.AWSBedrock { + return config.AWSBedrock{ + Region: "us-east-1", + AccessKey: "test-key", + AccessKeySecret: "test-secret", + Model: model, + SmallFastModel: smallFastModel, + } + } // mockBedrock serves the Bedrock control-plane API and records the paths it // receives. Callers point the SDK at the returned URL. @@ -147,18 +159,6 @@ func TestResolveInferenceProfile(t *testing.T) { return srv.URL, &got } - credentials := func(t *testing.T) aws.Config { - t.Helper() - - awsCfg, err := BuildBedrockCredentials(context.Background(), BedrockIdentity{ - Region: "us-east-1", - AccessKey: "test-key", - AccessKeySecret: "test-secret", - }) - require.NoError(t, err) - return awsCfg - } - t.Run("profile resolves to its model", func(t *testing.T) { url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") @@ -166,14 +166,15 @@ func TestResolveInferenceProfile(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - model, err := ResolveInferenceProfile(context.Background(), credentials(t), profileARN) + model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.NoError(t, err) require.Equal(t, "anthropic.claude-opus-4-8", model) - require.Len(t, *paths, 1) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + require.Len(t, *paths, 1, "only the profile ARN is resolved") require.Contains(t, (*paths)[0], profileARN) }) - t.Run("failed lookup is an error", func(t *testing.T) { + t.Run("failed resolution is an error", func(t *testing.T) { url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") @@ -182,7 +183,8 @@ func TestResolveInferenceProfile(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := ResolveInferenceProfile(context.Background(), credentials(t), profileARN) + _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + require.ErrorContains(t, err, "resolve model") require.ErrorContains(t, err, "GetInferenceProfile") }) @@ -193,9 +195,37 @@ func TestResolveInferenceProfile(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, err := ResolveInferenceProfile(context.Background(), credentials(t), profileARN) + _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.ErrorContains(t, err, "references no model") }) + + t.Run("small fast profile resolves independently", func(t *testing.T) { + url, paths := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"models":[{"modelArn":"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-haiku-4-5"}]}`)) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) + require.NoError(t, err) + require.Equal(t, "eu.anthropic.claude-opus-4-8", model) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + require.Len(t, *paths, 1, "only the small fast profile ARN is resolved") + require.Contains(t, (*paths)[0], smallFastProfileARN) + }) + + t.Run("plain model ids need no resolution", func(t *testing.T) { + url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { + t.Error("Bedrock called for plain model ids") + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) + require.NoError(t, err) + require.Equal(t, "eu.anthropic.claude-opus-4-8", model) + require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + require.Empty(t, *paths) + }) } // TestNewAnthropic_ServesStoredResolution covers what the gateway does with the diff --git a/aibridge/provider/bedrock_internal_test.go b/aibridge/provider/bedrock_internal_test.go index 6c542bff4e2..1577f3b5262 100644 --- a/aibridge/provider/bedrock_internal_test.go +++ b/aibridge/provider/bedrock_internal_test.go @@ -49,7 +49,7 @@ func TestBuildBedrockCredentialsValidation(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - _, err := bedrockRuntimeCredentials(context.Background(), tt.cfg) + _, err := buildBedrockCredentials(context.Background(), tt.cfg) require.Error(t, err) require.Contains(t, err.Error(), tt.errorMsg) }) @@ -60,7 +60,7 @@ func TestBuildBedrockCredentialsValidation(t *testing.T) { func TestBuildBedrockCredentialsStatic(t *testing.T) { t.Parallel() - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", AccessKey: "test-key", AccessKeySecret: "test-secret", @@ -132,10 +132,10 @@ func TestBuildBedrockCredentialsDefaultChain(t *testing.T) { t.Setenv(key, val) } - // bedrockRuntimeCredentials only wires up the provider chain; it + // buildBedrockCredentials only wires up the provider chain; it // does not resolve credentials, so it succeeds regardless of // credential availability. Resolution failures surface on Retrieve. - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", }) require.NoError(t, err) @@ -193,7 +193,7 @@ func TestBuildBedrockCredentialsAssumeRole(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -254,7 +254,7 @@ func TestBuildBedrockCredentialsAssumeRoleExternalID(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", ExternalID: tt.externalID, @@ -293,7 +293,7 @@ func TestBuildBedrockCredentialsAssumeRoleError(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -337,7 +337,7 @@ func TestBuildBedrockCredentialsAssumeRoleCaches(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -386,7 +386,7 @@ func TestBuildBedrockCredentialsAssumeRoleRefreshesOnExpiry(t *testing.T) { t.Setenv("AWS_ACCESS_KEY_ID", "base-key") t.Setenv("AWS_SECRET_ACCESS_KEY", "base-secret") - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ Region: "us-east-1", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -415,7 +415,7 @@ func TestBuildBedrockCredentialsAssumeRoleRequiresRegion(t *testing.T) { t.Setenv("AWS_SHARED_CREDENTIALS_FILE", "/dev/null") t.Setenv("AWS_EC2_METADATA_DISABLED", "true") - _, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + _, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ BaseURL: "https://bedrock-runtime.example.com", RoleARN: "arn:aws:iam::123456789012:role/target", }) @@ -430,7 +430,7 @@ func TestBuildBedrockCredentialsAssumeRoleRegionFromEnv(t *testing.T) { t.Setenv("AWS_REGION", "us-west-2") // BaseURL set with no explicit region: the region comes from AWS_REGION. - awsCfg, err := bedrockRuntimeCredentials(context.Background(), config.AWSBedrock{ + awsCfg, err := buildBedrockCredentials(context.Background(), config.AWSBedrock{ BaseURL: "https://bedrock-runtime.example.com", RoleARN: "arn:aws:iam::123456789012:role/target", }) diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index e275d7f10b9..8a183797f98 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -3,23 +3,17 @@ package coderd import ( "context" "net/http" - "time" "golang.org/x/xerrors" "cdr.dev/slog/v3" + "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/provider" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/httpapi" - "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" ) -// inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls a -// single write makes, which also cover the first credential resolution -// (STS/IRSA). -const inferenceProfileResolutionTimeout = 30 * time.Second - // errAIProviderProfileUnresolvable wraps a failed Bedrock inference profile // lookup so the write path reports it as a client-visible validation failure // rather than an internal error. @@ -41,46 +35,22 @@ type BedrockModelResolver interface { // the provider's own credentials, including any assumed role. type awsBedrockModelResolver struct{} -// ResolveModels resolves the configured identifiers to the model IDs the -// gateway records for capability detection, usage, and pricing. Only -// application inference profile ARNs are opaque, so only they cost an AWS call; -// every other identifier resolves to itself. func (awsBedrockModelResolver) ResolveModels(ctx context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) { - if !provider.IsApplicationInferenceProfileARN(settings.Model) && - !provider.IsApplicationInferenceProfileARN(settings.SmallFastModel) { - return settings.Model, settings.SmallFastModel, nil + cfg := config.AWSBedrock{ + Region: settings.Region, + Model: settings.Model, + SmallFastModel: settings.SmallFastModel, + RoleARN: settings.RoleARN, + ExternalID: settings.ExternalID, + Protocol: config.BedrockProtocol(settings.ResolvedProtocol()), } - - ctx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) - defer cancel() - - awsCfg, err := provider.BuildBedrockCredentials(ctx, provider.BedrockIdentity{ - Region: settings.Region, - AccessKey: ptr.NilToEmpty(settings.AccessKey), - AccessKeySecret: ptr.NilToEmpty(settings.AccessKeySecret), - RoleARN: settings.RoleARN, - ExternalID: settings.ExternalID, - }) - if err != nil { - return "", "", xerrors.Errorf("build bedrock credentials: %w", err) + if settings.AccessKey != nil { + cfg.AccessKey = *settings.AccessKey } - - resolve := func(configured string) (string, error) { - if !provider.IsApplicationInferenceProfileARN(configured) { - return configured, nil - } - return provider.ResolveInferenceProfile(ctx, awsCfg, configured) - } - - model, err = resolve(settings.Model) - if err != nil { - return "", "", xerrors.Errorf("resolve model: %w", err) - } - smallFastModel, err = resolve(settings.SmallFastModel) - if err != nil { - return "", "", xerrors.Errorf("resolve small fast model: %w", err) + if settings.AccessKeySecret != nil { + cfg.AccessKeySecret = *settings.AccessKeySecret } - return model, smallFastModel, nil + return provider.ResolveBedrockModels(ctx, cfg) } func (api *API) bedrockModelResolver() BedrockModelResolver { From 493b668b8bb74d7fba0ebf4b1dc9f449a7be6625 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 8 Sep 2026 20:00:53 +0000 Subject: [PATCH 04/25] refactor(coderd): share the bedrock settings converter with the write path --- cli/aibridged.go | 32 ++++----------------------- coderd/ai_providers_bedrock.go | 24 ++++++++------------ coderd/aibridge/bedrock.go | 40 ++++++++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 43 deletions(-) create mode 100644 coderd/aibridge/bedrock.go diff --git a/cli/aibridged.go b/cli/aibridged.go index f5f9c4fa2e3..030fbc00e65 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -14,11 +14,11 @@ import ( "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/coderd" + agplaibridge "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/aibridged" "github.com/coder/coder/v2/coderd/aibridged/proto" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/tracing" - "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" "github.com/coder/quartz" ) @@ -336,34 +336,10 @@ func buildAIProviderKeyPool(providerName string, keys []string, metrics *aibridg return keypool.New(providerName, keys, quartz.NewReal(), metrics) } -// bedrockConfig returns nil when the settings are absent or when the -// Bedrock fields are not actually configured. The provider's BaseURL is -// the generic upstream endpoint and is always non-empty, so it cannot -// serve as a Bedrock detection signal; gate on the settings alone via -// [codersdk.AIProviderBedrockSettings.IsConfigured]. +// bedrockConfig is [agplaibridge.BedrockConfig], shared with the provider +// write path so both map stored settings the same way. func bedrockConfig(baseURL string, bedrock *codersdk.AIProviderBedrockSettings) *aibridge.AWSBedrockConfig { - if bedrock == nil { - return nil - } - bedrockSettings := *bedrock - if !bedrockSettings.IsConfigured() { - return nil - } - accessKey := ptr.NilToEmpty(bedrockSettings.AccessKey) - accessKeySecret := ptr.NilToEmpty(bedrockSettings.AccessKeySecret) - return &aibridge.AWSBedrockConfig{ - BaseURL: baseURL, - Region: bedrockSettings.Region, - AccessKey: accessKey, - AccessKeySecret: accessKeySecret, - Model: bedrockSettings.Model, - SmallFastModel: bedrockSettings.SmallFastModel, - RoleARN: bedrockSettings.RoleARN, - ExternalID: bedrockSettings.ExternalID, - Protocol: config.BedrockProtocol(bedrockSettings.ResolvedProtocol()), - ResolvedModel: bedrockSettings.ResolvedModel, - ResolvedSmallFastModel: bedrockSettings.ResolvedSmallFastModel, - } + return agplaibridge.BedrockConfig(baseURL, bedrock) } // circuitBreakerConfig returns nil when the breaker is disabled. diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index 8a183797f98..f71a6a36674 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -7,8 +7,8 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog/v3" - "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/provider" + agplaibridge "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/codersdk" @@ -35,22 +35,16 @@ type BedrockModelResolver interface { // the provider's own credentials, including any assumed role. type awsBedrockModelResolver struct{} +// ResolveModels resolves the configured identifiers to the model IDs the +// gateway records for capability detection, usage, and pricing. Only +// application inference profile ARNs are opaque, so only they cost an AWS call; +// every other identifier resolves to itself. func (awsBedrockModelResolver) ResolveModels(ctx context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) { - cfg := config.AWSBedrock{ - Region: settings.Region, - Model: settings.Model, - SmallFastModel: settings.SmallFastModel, - RoleARN: settings.RoleARN, - ExternalID: settings.ExternalID, - Protocol: config.BedrockProtocol(settings.ResolvedProtocol()), + cfg := agplaibridge.BedrockConfig("", &settings) + if cfg == nil { + return settings.Model, settings.SmallFastModel, nil } - if settings.AccessKey != nil { - cfg.AccessKey = *settings.AccessKey - } - if settings.AccessKeySecret != nil { - cfg.AccessKeySecret = *settings.AccessKeySecret - } - return provider.ResolveBedrockModels(ctx, cfg) + return provider.ResolveBedrockModels(ctx, *cfg) } func (api *API) bedrockModelResolver() BedrockModelResolver { diff --git a/coderd/aibridge/bedrock.go b/coderd/aibridge/bedrock.go new file mode 100644 index 00000000000..49182d484be --- /dev/null +++ b/coderd/aibridge/bedrock.go @@ -0,0 +1,40 @@ +package aibridge + +import ( + aibridgeconfig "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/coderd/util/ptr" + "github.com/coder/coder/v2/codersdk" +) + +// BedrockConfig maps stored provider settings onto the runtime Bedrock +// configuration. It is shared by the gateway, which serves requests with it, +// and by the provider write path, which resolves application inference profile +// ARNs with it. +// +// It returns nil when the settings are absent or when the Bedrock fields are +// not actually configured. The provider's BaseURL is the generic upstream +// endpoint and is always non-empty, so it cannot serve as a Bedrock detection +// signal; gate on the settings alone via +// [codersdk.AIProviderBedrockSettings.IsConfigured]. +func BedrockConfig(baseURL string, bedrock *codersdk.AIProviderBedrockSettings) *aibridgeconfig.AWSBedrock { + if bedrock == nil { + return nil + } + settings := *bedrock + if !settings.IsConfigured() { + return nil + } + return &aibridgeconfig.AWSBedrock{ + BaseURL: baseURL, + Region: settings.Region, + AccessKey: ptr.NilToEmpty(settings.AccessKey), + AccessKeySecret: ptr.NilToEmpty(settings.AccessKeySecret), + Model: settings.Model, + SmallFastModel: settings.SmallFastModel, + RoleARN: settings.RoleARN, + ExternalID: settings.ExternalID, + Protocol: aibridgeconfig.BedrockProtocol(settings.ResolvedProtocol()), + ResolvedModel: settings.ResolvedModel, + ResolvedSmallFastModel: settings.ResolvedSmallFastModel, + } +} From 0774d48c85a2f071aaab1775087202e0f71c88a1 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 8 Sep 2026 20:24:33 +0000 Subject: [PATCH 05/25] refactor(coderd): resolve Bedrock profiles directly in provider writes --- coderd/ai_providers.go | 2 +- coderd/ai_providers_bedrock.go | 50 +++------ coderd/ai_providers_bedrock_test.go | 158 ++++++++++++++++------------ coderd/coderd.go | 6 +- coderd/coderdtest/coderdtest.go | 10 +- 5 files changed, 106 insertions(+), 120 deletions(-) diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index e5b897e37fb..b781d3f614d 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -189,7 +189,7 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { // Resolve application inference profile ARNs before storing them. Doing it // here means the operator learns immediately that a profile is wrong or // unreachable, and the gateway never calls AWS to find out. - if err := api.resolveBedrockModels(ctx, &req.Settings); err != nil { + if err := resolveBedrockModels(ctx, &req.Settings); err != nil { api.writeAIProviderResolutionError(ctx, rw, err) return } diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index f71a6a36674..a355f160d31 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -19,54 +19,28 @@ import ( // rather than an internal error. var errAIProviderProfileUnresolvable = xerrors.New("resolve bedrock inference profile") -// BedrockModelResolver resolves the configured Bedrock model identifiers of a -// provider to the model IDs the gateway records for capability detection, -// usage, and pricing. Identifiers that are not application inference profile -// ARNs resolve to themselves without calling AWS. +// resolveBedrockModels fills in the server-owned resolved identifiers on +// settings. Only application inference profile ARNs are opaque, so only they +// cost an AWS call; every other identifier resolves to itself and is stored +// unresolved. // // Resolution runs here, where the provider is written, so the gateway never // calls the Bedrock control plane: not at startup, not on reload, and not on a -// request. It is an interface so tests can supply results without AWS. -type BedrockModelResolver interface { - ResolveModels(ctx context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) -} - -// awsBedrockModelResolver resolves through the AWS Bedrock control plane using -// the provider's own credentials, including any assumed role. -type awsBedrockModelResolver struct{} - -// ResolveModels resolves the configured identifiers to the model IDs the -// gateway records for capability detection, usage, and pricing. Only -// application inference profile ARNs are opaque, so only they cost an AWS call; -// every other identifier resolves to itself. -func (awsBedrockModelResolver) ResolveModels(ctx context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) { - cfg := agplaibridge.BedrockConfig("", &settings) - if cfg == nil { - return settings.Model, settings.SmallFastModel, nil - } - return provider.ResolveBedrockModels(ctx, *cfg) -} - -func (api *API) bedrockModelResolver() BedrockModelResolver { - if api.AIProviderBedrockResolver != nil { - return api.AIProviderBedrockResolver - } - return awsBedrockModelResolver{} -} - -// resolveBedrockModels fills in the server-owned resolved identifiers on -// settings. A resolved value is stored only when it differs from the configured -// one, so plain model IDs stay unresolved and keep serving themselves. +// request. // // A failure is returned to the caller: an unresolvable profile must not be // stored, because the gateway cannot tell what an opaque ARN refers to and // would misshape every request made through it. -func (api *API) resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSettings) error { +func resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSettings) error { if settings.Bedrock == nil { return nil } + cfg := agplaibridge.BedrockConfig("", settings.Bedrock) + if cfg == nil { + return nil + } - model, smallFastModel, err := api.bedrockModelResolver().ResolveModels(ctx, *settings.Bedrock) + model, smallFastModel, err := provider.ResolveBedrockModels(ctx, *cfg) if err != nil { return xerrors.Errorf("%w: %w", errAIProviderProfileUnresolvable, err) } @@ -114,7 +88,7 @@ func (api *API) previewResolvedBedrockSettings(ctx context.Context, idOrName str preview := mergeAIProviderSettings(existing, *patch) ensureBedrockExternalID(&preview) - if err := api.resolveBedrockModels(ctx, &preview); err != nil { + if err := resolveBedrockModels(ctx, &preview); err != nil { return codersdk.AIProviderSettings{}, false, err } return preview, true, nil diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go index 174d1165314..a54c7d36692 100644 --- a/coderd/ai_providers_bedrock_test.go +++ b/coderd/ai_providers_bedrock_test.go @@ -1,13 +1,15 @@ package coderd_test import ( - "context" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" "testing" "github.com/stretchr/testify/require" - "golang.org/x/xerrors" - "github.com/coder/coder/v2/coderd" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" @@ -18,29 +20,6 @@ const ( testSmallFastProfileARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/8x1qk20fzp3r" ) -// stubBedrockResolver stands in for the AWS Bedrock control plane. It records -// the settings it was asked to resolve so tests can assert whether a write -// consulted AWS at all. -type stubBedrockResolver struct { - models map[string]string - err error - calls []codersdk.AIProviderBedrockSettings -} - -func (s *stubBedrockResolver) ResolveModels(_ context.Context, settings codersdk.AIProviderBedrockSettings) (model, smallFastModel string, err error) { - s.calls = append(s.calls, settings) - if s.err != nil { - return "", "", s.err - } - resolve := func(configured string) string { - if model, ok := s.models[configured]; ok { - return model - } - return configured - } - return resolve(settings.Model), resolve(settings.SmallFastModel), nil -} - func bedrockSettings(model, smallFastModel string) *codersdk.AIProviderSettings { accessKey := "test-key" accessKeySecret := "test-secret" @@ -55,25 +34,52 @@ func bedrockSettings(model, smallFastModel string) *codersdk.AIProviderSettings } } -func TestAIProvidersBedrockProfileResolution(t *testing.T) { - t.Parallel() - - newClient := func(t *testing.T, resolver coderd.BedrockModelResolver) *codersdk.Client { - t.Helper() +// mockBedrock serves the Bedrock control-plane API and records the profile +// lookups it receives. Callers point the AWS SDK at the returned URL. +func mockBedrock(t *testing.T, handler http.HandlerFunc) (url string, paths func() []string) { + t.Helper() + + var ( + mu sync.Mutex + got []string + ) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + got = append(got, r.URL.Path) + mu.Unlock() + handler(w, r) + })) + t.Cleanup(srv.Close) + return srv.URL, func() []string { + mu.Lock() + defer mu.Unlock() + return slices.Clone(got) + } +} - client := coderdtest.New(t, &coderdtest.Options{AIProviderBedrockResolver: resolver}) - _ = coderdtest.CreateFirstUser(t, client) - return client +func respondWithModel(modelARN string) http.HandlerFunc { + return func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"models":[{"modelArn":"` + modelARN + `"}]}`)) } +} +// TestAIProvidersBedrockProfileResolution drives provider writes against a mock +// Bedrock control plane, so the AWS SDK path runs for real. +// NOTE: no t.Parallel() because the subtests use t.Setenv. +func TestAIProvidersBedrockProfileResolution(t *testing.T) { t.Run("CreateStoresResolvedModels", func(t *testing.T) { - t.Parallel() + url, paths := mockBedrock(t, func(w http.ResponseWriter, r *http.Request) { + modelARN := "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8" + if !strings.Contains(r.URL.Path, "46u2vhiyo6z5") { + modelARN = "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-haiku-4-5" + } + respondWithModel(modelARN)(w, r) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - resolver := &stubBedrockResolver{models: map[string]string{ - testProfileARN: "anthropic.claude-opus-4-8", - testSmallFastProfileARN: "anthropic.claude-haiku-4-5", - }} - client := newClient(t, resolver) + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) //nolint:gocritic // Owner role is the audience for this endpoint. @@ -92,13 +98,17 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Equal(t, testSmallFastProfileARN, created.Settings.Bedrock.SmallFastModel) require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) require.Equal(t, "anthropic.claude-haiku-4-5", created.Settings.Bedrock.ResolvedSmallFastModel) + require.Len(t, paths(), 2, "each profile is resolved once") }) t.Run("CreateLeavesPlainModelIDsUnresolved", func(t *testing.T) { - t.Parallel() + url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { + t.Error("Bedrock called for plain model ids") + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - resolver := &stubBedrockResolver{} - client := newClient(t, resolver) + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) //nolint:gocritic // Owner role is the audience for this endpoint. @@ -113,13 +123,20 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.NotNil(t, created.Settings.Bedrock) require.Empty(t, created.Settings.Bedrock.ResolvedModel) require.Empty(t, created.Settings.Bedrock.ResolvedSmallFastModel) + require.Empty(t, paths()) }) t.Run("CreateRejectsUnresolvableProfile", func(t *testing.T) { - t.Parallel() + url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"message":"not authorized to perform bedrock:GetInferenceProfile"}`)) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - resolver := &stubBedrockResolver{err: xerrors.New("AccessDeniedException: not authorized to perform bedrock:GetInferenceProfile")} - client := newClient(t, resolver) + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) //nolint:gocritic // Owner role is the audience for this endpoint. @@ -132,8 +149,8 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) var sdkErr *codersdk.Error require.ErrorAs(t, err, &sdkErr) - require.Equal(t, 400, sdkErr.StatusCode()) - require.Contains(t, sdkErr.Detail, "bedrock:GetInferenceProfile") + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + require.Contains(t, sdkErr.Detail, "GetInferenceProfile") //nolint:gocritic // Owner role is the audience for this endpoint. providers, err := client.AIProviders(ctx) @@ -142,9 +159,13 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Run("CreateRejectsClientSuppliedResolution", func(t *testing.T) { - t.Parallel() + url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { + t.Error("Bedrock called for a rejected request") + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - client := newClient(t, &stubBedrockResolver{}) + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) settings := bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") @@ -160,18 +181,17 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) var sdkErr *codersdk.Error require.ErrorAs(t, err, &sdkErr) - require.Equal(t, 400, sdkErr.StatusCode()) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) require.Contains(t, sdkErr.Error(), "resolved_model") + require.Empty(t, paths(), "validation rejects the write before any AWS call") }) t.Run("UpdateReresolvesChangedProfile", func(t *testing.T) { - t.Parallel() + url, _ := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - resolver := &stubBedrockResolver{models: map[string]string{ - testProfileARN: "anthropic.claude-opus-4-8", - testSmallFastProfileARN: "anthropic.claude-haiku-4-5", - }} - client := newClient(t, resolver) + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) //nolint:gocritic // Owner role is the audience for this endpoint. @@ -196,12 +216,11 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Run("UpdateClearsResolutionWhenProfileReplacedByModelID", func(t *testing.T) { - t.Parallel() + url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - resolver := &stubBedrockResolver{models: map[string]string{ - testProfileARN: "anthropic.claude-opus-4-8", - }} - client := newClient(t, resolver) + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) //nolint:gocritic // Owner role is the audience for this endpoint. @@ -214,6 +233,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) require.NoError(t, err) require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) + callsAfterCreate := len(paths()) //nolint:gocritic // Owner role is the audience for this endpoint. updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ @@ -221,15 +241,15 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) require.NoError(t, err) require.Empty(t, updated.Settings.Bedrock.ResolvedModel, "a plain model id resolves to itself") + require.Len(t, paths(), callsAfterCreate, "no profile is left to resolve") }) t.Run("UpdateWithoutSettingsSkipsResolution", func(t *testing.T) { - t.Parallel() + url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - resolver := &stubBedrockResolver{models: map[string]string{ - testProfileARN: "anthropic.claude-opus-4-8", - }} - client := newClient(t, resolver) + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) //nolint:gocritic // Owner role is the audience for this endpoint. @@ -241,7 +261,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - callsAfterCreate := len(resolver.calls) + callsAfterCreate := len(paths()) enabled := false //nolint:gocritic // Owner role is the audience for this endpoint. @@ -250,6 +270,6 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) require.NoError(t, err) require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) - require.Len(t, resolver.calls, callsAfterCreate, "an unrelated update does not call AWS") + require.Len(t, paths(), callsAfterCreate, "an unrelated update does not call AWS") }) } diff --git a/coderd/coderd.go b/coderd/coderd.go index dc9a501b8dd..5051ac7dc0c 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -176,11 +176,7 @@ type Options struct { // CacheDir is used for caching files served by the API. CacheDir string - Auditor audit.Auditor - // AIProviderBedrockResolver resolves Bedrock application inference profile - // ARNs when an AI provider is written. Defaults to resolving through the AWS - // Bedrock control plane; tests substitute their own. - AIProviderBedrockResolver BedrockModelResolver + Auditor audit.Auditor ConnectionLogger connectionlog.ConnectionLogger AgentConnectionUpdateFrequency time.Duration AgentInactiveDisconnectTimeout time.Duration diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 6723215ea2d..8179431c2c6 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -131,12 +131,9 @@ type Options struct { AutobuildTicker <-chan time.Time AutobuildStats chan<- autobuild.Stats Auditor audit.Auditor - // AIProviderBedrockResolver resolves Bedrock application inference profile - // ARNs when an AI provider is written. Tests set it to avoid calling AWS. - AIProviderBedrockResolver coderd.BedrockModelResolver - TLSCertificates []tls.Certificate - ExternalAuthConfigs []*externalauth.Config - TrialGenerator func(ctx context.Context, body codersdk.LicensorTrialRequest) error + TLSCertificates []tls.Certificate + ExternalAuthConfigs []*externalauth.Config + TrialGenerator func(ctx context.Context, body codersdk.LicensorTrialRequest) error // MCPAllowedPrivateCIDRs exempts IP ranges from the MCP // SSRF guard for MCP server and OAuth2 traffic. Defaults to loopback so // tests can serve mock MCP and authorization servers via httptest. @@ -645,7 +642,6 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can UsageInserter: usageInserter, Auditor: options.Auditor, - AIProviderBedrockResolver: options.AIProviderBedrockResolver, ConnectionLogger: options.ConnectionLogger, AWSCertificates: options.AWSCertificates, AzureCertificates: options.AzureCertificates, From 96adb4ed7c3e8830df511b2413cee6495234bb5e Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Tue, 8 Sep 2026 21:43:14 +0000 Subject: [PATCH 06/25] refactor(coderd): store bedrock profile resolution after the provider write --- coderd/ai_providers.go | 59 ++++++--------- coderd/ai_providers_bedrock.go | 103 +++++++++++++++----------- coderd/ai_providers_bedrock_test.go | 5 +- docs/ai-coder/ai-gateway/providers.md | 6 +- 4 files changed, 89 insertions(+), 84 deletions(-) diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index b781d3f614d..c5fad09933d 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -186,14 +186,6 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { // Generate the server-owned external ID when the provider assumes a role. ensureBedrockExternalID(&req.Settings) - // Resolve application inference profile ARNs before storing them. Doing it - // here means the operator learns immediately that a profile is wrong or - // unreachable, and the gateway never calls AWS to find out. - if err := resolveBedrockModels(ctx, &req.Settings); err != nil { - api.writeAIProviderResolutionError(ctx, rw, err) - return - } - settings, err := encodeAIProviderSettings(req.Settings) if err != nil { api.Logger.Error(ctx, "encode AI provider settings", slog.Error(err)) @@ -252,6 +244,16 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { } aReq.New = row + // Resolve inference profile ARNs once the provider is stored, then announce + // it. The gateway never sees the unresolved provider, and never calls the + // Bedrock control plane itself. + row, err = api.applyBedrockResolution(ctx, row) + if err != nil { + api.writeAIProviderResolutionError(ctx, rw, err) + return + } + aReq.New = row + auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, aiProviderKeyChanges{Added: keys}) api.publishAIProvidersChanged(ctx) @@ -317,21 +319,12 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { idOrName := chi.URLParam(r, "idOrName") - // Resolve outside the transaction: it is a network call. The merge is redone - // inside the transaction against the row that gets written, and the - // resolution is applied only when the model identifiers still match. - resolvedPreview, hasResolvedPreview, err := api.previewResolvedBedrockSettings(ctx, idOrName, req.Settings) - if err != nil { - api.writeAIProviderResolutionError(ctx, rw, err) - return - } - var ( updated database.AIProvider keys []database.AIProviderKey keyChanges aiProviderKeyChanges ) - err = api.Database.InTx(func(tx database.Store) error { + err := api.Database.InTx(func(tx database.Store) error { old, err := lookupAIProvider(ctx, tx, idOrName) if err != nil { return err @@ -362,13 +355,6 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { // Generate the server-owned external ID when the provider assumes a role // and lacks one. ensureBedrockExternalID(&existing) - if req.Settings != nil && existing.Bedrock != nil { - if !hasResolvedPreview || !bedrockModelsMatch(existing, resolvedPreview) { - return errAIProviderChangedDuringUpdate - } - existing.Bedrock.ResolvedModel = resolvedPreview.Bedrock.ResolvedModel - existing.Bedrock.ResolvedSmallFastModel = resolvedPreview.Bedrock.ResolvedSmallFastModel - } settings, err := encodeAIProviderSettings(existing) if err != nil { return xerrors.Errorf("encode settings: %w", err) @@ -447,12 +433,6 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { }) return } - if errors.Is(err, errAIProviderChangedDuringUpdate) { - httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ - Message: "The AI provider changed while it was being updated. Retry the request.", - }) - return - } if errors.Is(err, errAIProviderKeyUnknown) { // Use the sentinel directly so the response message does not // leak the "execute transaction:" wrapper xerrors added on the @@ -468,6 +448,18 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { return } + // An update that carries no settings cannot change the configured + // identifiers or the credentials they resolve under, so the stored + // resolution still holds. + if req.Settings != nil { + updated, err = api.applyBedrockResolution(ctx, updated) + if err != nil { + api.writeAIProviderResolutionError(ctx, rw, err) + return + } + aReq.New = updated + } + auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, keyChanges) api.publishAIProvidersChanged(ctx) @@ -566,11 +558,6 @@ var errAIProviderBedrockTypeMismatch = xerrors.New("bedrock settings are only va // patch may echo the stored value but not set a different one. var errAIProviderExternalIDReadOnly = xerrors.New("external_id is server-generated and cannot be changed") -// errAIProviderChangedDuringUpdate is the sentinel returned from inside the -// update transaction when the provider's model identifiers changed after they -// were resolved, so the resolution no longer describes what would be stored. -var errAIProviderChangedDuringUpdate = xerrors.New("provider changed while it was being updated, retry the request") - // errAIProviderInvalidName is returned from lookupAIProvider when the // idOrName parameter is neither a UUID nor a syntactically-valid name. // The handler translates this into a 400 so an integrator gets a hint diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index a355f160d31..3bf234587c6 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -2,6 +2,8 @@ package coderd import ( "context" + "database/sql" + "errors" "net/http" "golang.org/x/xerrors" @@ -9,15 +11,22 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/provider" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" + "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/codersdk" ) -// errAIProviderProfileUnresolvable wraps a failed Bedrock inference profile -// lookup so the write path reports it as a client-visible validation failure -// rather than an internal error. -var errAIProviderProfileUnresolvable = xerrors.New("resolve bedrock inference profile") +// bedrockProfileUnresolvableError marks a failed Bedrock inference profile lookup +// so the write path reports it as a client-visible validation failure rather +// than an internal error. +type bedrockProfileUnresolvableError struct{ err error } + +func (e bedrockProfileUnresolvableError) Error() string { + return "resolve bedrock inference profile: " + e.err.Error() +} + +func (e bedrockProfileUnresolvableError) Unwrap() error { return e.err } // resolveBedrockModels fills in the server-owned resolved identifiers on // settings. Only application inference profile ARNs are opaque, so only they @@ -27,14 +36,7 @@ var errAIProviderProfileUnresolvable = xerrors.New("resolve bedrock inference pr // Resolution runs here, where the provider is written, so the gateway never // calls the Bedrock control plane: not at startup, not on reload, and not on a // request. -// -// A failure is returned to the caller: an unresolvable profile must not be -// stored, because the gateway cannot tell what an opaque ARN refers to and -// would misshape every request made through it. func resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSettings) error { - if settings.Bedrock == nil { - return nil - } cfg := agplaibridge.BedrockConfig("", settings.Bedrock) if cfg == nil { return nil @@ -42,7 +44,7 @@ func resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSett model, smallFastModel, err := provider.ResolveBedrockModels(ctx, *cfg) if err != nil { - return xerrors.Errorf("%w: %w", errAIProviderProfileUnresolvable, err) + return bedrockProfileUnresolvableError{err: err} } settings.Bedrock.ResolvedModel = "" @@ -56,49 +58,62 @@ func resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSett return nil } -// bedrockModelsMatch reports whether two settings configure the same model -// identifiers. The update path resolves against a snapshot taken outside the -// transaction, so it re-checks the merged settings before storing the result. -func bedrockModelsMatch(a, b codersdk.AIProviderSettings) bool { - if a.Bedrock == nil || b.Bedrock == nil { - return a.Bedrock == b.Bedrock +// applyBedrockResolution resolves the inference profile ARNs of a provider that +// was just written and stores the result on it. Resolution is an AWS call, so +// it runs after the write transaction has committed rather than holding a +// database connection open across the network. +// +// The provider row is the merged settings the operator will actually use, which +// is why resolution reads it back instead of the request. +func (api *API) applyBedrockResolution(ctx context.Context, row database.AIProvider) (database.AIProvider, error) { + settings, err := db2sdk.AIProviderSettings(row.Settings) + if err != nil { + return row, xerrors.Errorf("decode settings: %w", err) + } + if settings.Bedrock == nil { + return row, nil } - return a.Bedrock.Model == b.Bedrock.Model && a.Bedrock.SmallFastModel == b.Bedrock.SmallFastModel -} -// previewResolvedBedrockSettings merges patch onto the stored settings of the -// named provider and resolves the result, so the write path can perform the -// AWS lookup outside its transaction. The boolean reports whether a preview was -// produced: there is nothing to resolve, or the provider cannot be read, in -// which case the transaction reports the failure with its own error handling. -func (api *API) previewResolvedBedrockSettings(ctx context.Context, idOrName string, patch *codersdk.AIProviderSettings) (codersdk.AIProviderSettings, bool, error) { - if patch == nil || patch.Bedrock == nil { - return codersdk.AIProviderSettings{}, false, nil + stored := *settings.Bedrock + if err := resolveBedrockModels(ctx, &settings); err != nil { + return row, err } - old, err := lookupAIProvider(ctx, api.Database, idOrName) - if err != nil { - //nolint:nilerr // The transaction reports lookup failures. - return codersdk.AIProviderSettings{}, false, nil + if settings.Bedrock.ResolvedModel == stored.ResolvedModel && + settings.Bedrock.ResolvedSmallFastModel == stored.ResolvedSmallFastModel { + return row, nil } - existing, err := db2sdk.AIProviderSettings(old.Settings) + + encoded, err := encodeAIProviderSettings(settings) if err != nil { - //nolint:nilerr // The transaction reports decode failures. - return codersdk.AIProviderSettings{}, false, nil + return row, xerrors.Errorf("encode settings: %w", err) } - - preview := mergeAIProviderSettings(existing, *patch) - ensureBedrockExternalID(&preview) - if err := resolveBedrockModels(ctx, &preview); err != nil { - return codersdk.AIProviderSettings{}, false, err + updated, err := api.Database.UpdateAIProvider(ctx, database.UpdateAIProviderParams{ + ID: row.ID, + Type: row.Type, + DisplayName: row.DisplayName, + Icon: row.Icon, + Enabled: row.Enabled, + BaseUrl: row.BaseUrl, + Settings: encoded, + // SettingsKeyID is set by the dbcrypt wrapper. + SettingsKeyID: sql.NullString{}, + }) + if err != nil { + return row, xerrors.Errorf("store resolved models: %w", err) } - return preview, true, nil + return updated, nil } // writeAIProviderResolutionError reports a failed Bedrock model resolution. The -// write is rejected rather than stored unresolved: the gateway cannot serve an -// opaque profile ARN, so accepting it would produce a provider that fails every -// request. +// provider keeps the identifiers the operator asked for, but without a +// resolution the gateway cannot tell what an opaque profile ARN refers to, so +// it refuses to serve the provider until the write succeeds. func (api *API) writeAIProviderResolutionError(ctx context.Context, rw http.ResponseWriter, err error) { + var unresolvable bedrockProfileUnresolvableError + if !errors.As(err, &unresolvable) { + writeAIProviderError(ctx, api.Logger, rw, err, "resolve bedrock inference profile", "Internal error resolving the Bedrock application inference profile.") + return + } api.Logger.Warn(ctx, "resolve bedrock inference profile", slog.Error(err)) httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Could not resolve the Bedrock application inference profile. Check that the ARN is correct and that the AWS identity used by Coder is allowed bedrock:GetInferenceProfile.", diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go index a54c7d36692..8d964596e4c 100644 --- a/coderd/ai_providers_bedrock_test.go +++ b/coderd/ai_providers_bedrock_test.go @@ -152,10 +152,13 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) require.Contains(t, sdkErr.Detail, "GetInferenceProfile") + // The provider is stored with the ARN the operator asked for, but + // without a resolution the gateway refuses to serve it. //nolint:gocritic // Owner role is the audience for this endpoint. providers, err := client.AIProviders(ctx) require.NoError(t, err) - require.Empty(t, providers, "an unresolvable provider is not stored") + require.Len(t, providers, 1) + require.Empty(t, providers[0].Settings.Bedrock.ResolvedModel) }) t.Run("CreateRejectsClientSuppliedResolution", func(t *testing.T) { diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index 89522b509dd..670b8568c68 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -255,9 +255,9 @@ Resolution requires a `GetInferenceProfile` call, which Coder makes when the provider is saved, not when a request is served. The AWS identity Coder uses, which is the provider's access keys when configured and otherwise the identity of the Coder deployment, must have `bedrock:GetInferenceProfile` permission for -the profile. Saving fails when the lookup fails, so a profile that cannot be -resolved is never stored. Providers configured with plain model identifiers do -not need this permission. +the profile. Saving reports an error when the lookup fails, and the provider +cannot serve requests until a later save resolves it. Providers configured with +plain model identifiers do not need this permission. ### GitHub Copilot From 0ba3ad9f4099c9b0f859ff1677a2ad9da7877fb3 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 14:50:54 +0000 Subject: [PATCH 07/25] refactor(coderd): store bedrock profile resolution in its own table --- cli/aibridged.go | 14 +- coderd/ai_providers.go | 20 ++- coderd/ai_providers_bedrock.go | 105 ++++++--------- coderd/ai_providers_bedrock_test.go | 123 +++++++++++------- coderd/aibridge/bedrock.go | 20 ++- coderd/aibridgedserver/aibridgedserver.go | 31 ++++- coderd/database/dbauthz/dbauthz.go | 21 +++ coderd/database/dbauthz/dbauthz_test.go | 23 ++++ coderd/database/dbmetrics/querymetrics.go | 24 ++++ coderd/database/dbmock/dbmock.go | 43 ++++++ coderd/database/dump.sql | 16 +++ coderd/database/foreign_key_constraint.go | 1 + ..._provider_bedrock_resolved_models.down.sql | 1 + ...ai_provider_bedrock_resolved_models.up.sql | 17 +++ ...ai_provider_bedrock_resolved_models.up.sql | 31 +++++ coderd/database/models.go | 8 ++ coderd/database/querier.go | 6 + coderd/database/queries.sql.go | 68 ++++++++++ coderd/database/queries/ai_providers.sql | 26 ++++ coderd/database/unique_constraint.go | 1 + codersdk/aiproviders.go | 21 --- codersdk/aiproviders_bedrock.go | 27 ---- site/src/api/typesGenerated.ts | 12 -- 23 files changed, 459 insertions(+), 200 deletions(-) create mode 100644 coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql create mode 100644 coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql diff --git a/cli/aibridged.go b/cli/aibridged.go index 030fbc00e65..4aee68221a1 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -216,9 +216,9 @@ func protoToProviderSpec(pp *proto.AIProvider) aiProviderSpec { bedrock.RoleARN = b.GetRoleArn() bedrock.ExternalID = b.GetExternalId() bedrock.Protocol = codersdk.AIProviderBedrockProtocol(b.GetProtocol()) - bedrock.ResolvedModel = b.GetResolvedModel() - bedrock.ResolvedSmallFastModel = b.GetResolvedSmallFastModel() spec.Bedrock = new(bedrock) + spec.BedrockResolvedModel = b.GetResolvedModel() + spec.BedrockResolvedSmallFastModel = b.GetResolvedSmallFastModel() } return spec } @@ -237,6 +237,12 @@ type aiProviderSpec struct { // Bedrock holds Bedrock-specific settings when the provider targets // AWS Bedrock; nil otherwise. Bedrock *codersdk.AIProviderBedrockSettings + // BedrockResolvedModel and BedrockResolvedSmallFastModel are the models the + // configured identifiers refer to. They are set only when an identifier is + // an application inference profile ARN, which coderd resolved when the + // provider was written. + BedrockResolvedModel string + BedrockResolvedSmallFastModel string } // buildProvider constructs the appropriate [aibridge.Provider] for a @@ -285,6 +291,10 @@ func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBrid case database.AIProviderTypeAnthropic, database.AIProviderTypeBedrock: bedrock := bedrockConfig(spec.BaseURL, spec.Bedrock) + if bedrock != nil { + bedrock.ResolvedModel = spec.BedrockResolvedModel + bedrock.ResolvedSmallFastModel = spec.BedrockResolvedSmallFastModel + } // A spec typed 'bedrock' authenticates exclusively via settings; // without populated Bedrock credentials it cannot make upstream // calls, so refuse rather than falling back to an unsigned diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index c5fad09933d..ff7ca6e6231 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -247,12 +247,10 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { // Resolve inference profile ARNs once the provider is stored, then announce // it. The gateway never sees the unresolved provider, and never calls the // Bedrock control plane itself. - row, err = api.applyBedrockResolution(ctx, row) - if err != nil { + if err := api.resolveBedrockModels(ctx, row); err != nil { api.writeAIProviderResolutionError(ctx, rw, err) return } - aReq.New = row auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, aiProviderKeyChanges{Added: keys}) api.publishAIProvidersChanged(ctx) @@ -370,6 +368,15 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { return errCopilotRejectsAPIKeys } + // The patch may point the provider at different identifiers, so the + // stored resolution no longer describes it. Resolution runs after the + // transaction, because it is an AWS call. + if req.Settings != nil { + if err := clearBedrockModelResolution(ctx, tx, old.ID); err != nil { + return err + } + } + displayName := old.DisplayName if req.DisplayName != nil { // Empty string clears the column. @@ -452,12 +459,10 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { // identifiers or the credentials they resolve under, so the stored // resolution still holds. if req.Settings != nil { - updated, err = api.applyBedrockResolution(ctx, updated) - if err != nil { + if err := api.resolveBedrockModels(ctx, updated); err != nil { api.writeAIProviderResolutionError(ctx, rw, err) return } - aReq.New = updated } auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, keyChanges) @@ -509,7 +514,8 @@ func (api *API) aiProvidersDelete(rw http.ResponseWriter, r *http.Request) { if err := tx.DeleteAIProviderByID(ctx, row.ID); err != nil { return xerrors.Errorf("delete ai provider: %w", err) } - return nil + // Providers are soft-deleted, so the foreign key never cascades. + return clearBedrockModelResolution(ctx, tx, row.ID) }, &database.TxOptions{TxIdentifier: "delete_ai_provider"}) if err != nil { writeAIProviderError(ctx, api.Logger, rw, err, "delete AI provider", "Internal error deleting AI provider.") diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index 3bf234587c6..75bbbe26640 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -2,10 +2,9 @@ package coderd import ( "context" - "database/sql" - "errors" "net/http" + "github.com/google/uuid" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -17,9 +16,9 @@ import ( "github.com/coder/coder/v2/codersdk" ) -// bedrockProfileUnresolvableError marks a failed Bedrock inference profile lookup -// so the write path reports it as a client-visible validation failure rather -// than an internal error. +// bedrockProfileUnresolvableError marks a failed Bedrock inference profile +// lookup so the write path reports it as a client-visible validation failure +// rather than an internal error. type bedrockProfileUnresolvableError struct{ err error } func (e bedrockProfileUnresolvableError) Error() string { @@ -28,15 +27,23 @@ func (e bedrockProfileUnresolvableError) Error() string { func (e bedrockProfileUnresolvableError) Unwrap() error { return e.err } -// resolveBedrockModels fills in the server-owned resolved identifiers on -// settings. Only application inference profile ARNs are opaque, so only they -// cost an AWS call; every other identifier resolves to itself and is stored -// unresolved. +// resolveBedrockModels records which models the provider's application +// inference profile ARNs refer to. An ARN identifies a billing wrapper rather +// than a model, so the gateway needs the mapping to detect capabilities, price +// usage, and record interceptions. // -// Resolution runs here, where the provider is written, so the gateway never -// calls the Bedrock control plane: not at startup, not on reload, and not on a -// request. -func resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSettings) error { +// Resolution is an AWS call, so it runs after the provider write has committed +// rather than holding a database transaction open across the network. It reads +// the stored row because that is the merged configuration the provider will +// actually use. +// +// Nothing is stored for a provider configured with plain model IDs: they are +// already model identities and need no mapping. +func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvider) error { + settings, err := db2sdk.AIProviderSettings(row.Settings) + if err != nil { + return xerrors.Errorf("decode settings: %w", err) + } cfg := agplaibridge.BedrockConfig("", settings.Bedrock) if cfg == nil { return nil @@ -46,71 +53,39 @@ func resolveBedrockModels(ctx context.Context, settings *codersdk.AIProviderSett if err != nil { return bedrockProfileUnresolvableError{err: err} } - - settings.Bedrock.ResolvedModel = "" - if model != settings.Bedrock.Model { - settings.Bedrock.ResolvedModel = model - } - settings.Bedrock.ResolvedSmallFastModel = "" - if smallFastModel != settings.Bedrock.SmallFastModel { - settings.Bedrock.ResolvedSmallFastModel = smallFastModel + if model == cfg.Model && smallFastModel == cfg.SmallFastModel { + return nil } - return nil -} -// applyBedrockResolution resolves the inference profile ARNs of a provider that -// was just written and stores the result on it. Resolution is an AWS call, so -// it runs after the write transaction has committed rather than holding a -// database connection open across the network. -// -// The provider row is the merged settings the operator will actually use, which -// is why resolution reads it back instead of the request. -func (api *API) applyBedrockResolution(ctx context.Context, row database.AIProvider) (database.AIProvider, error) { - settings, err := db2sdk.AIProviderSettings(row.Settings) + err = api.Database.UpsertAIProviderBedrockResolvedModels(ctx, database.UpsertAIProviderBedrockResolvedModelsParams{ + AIProviderID: row.ID, + ResolvedModel: model, + ResolvedSmallFastModel: smallFastModel, + }) if err != nil { - return row, xerrors.Errorf("decode settings: %w", err) - } - if settings.Bedrock == nil { - return row, nil + return xerrors.Errorf("store resolved models: %w", err) } + return nil +} - stored := *settings.Bedrock - if err := resolveBedrockModels(ctx, &settings); err != nil { - return row, err +// clearBedrockModelResolution drops a provider's stored resolution. The write +// path calls it whenever the configured identifiers may have changed, so a +// stale mapping never outlives the ARN it describes. The provider is then +// unresolved until resolution succeeds, and the gateway will not serve it. +func clearBedrockModelResolution(ctx context.Context, db database.Store, providerID uuid.UUID) error { + if err := db.DeleteAIProviderBedrockResolvedModels(ctx, providerID); err != nil { + return xerrors.Errorf("clear resolved models: %w", err) } - if settings.Bedrock.ResolvedModel == stored.ResolvedModel && - settings.Bedrock.ResolvedSmallFastModel == stored.ResolvedSmallFastModel { - return row, nil - } - - encoded, err := encodeAIProviderSettings(settings) - if err != nil { - return row, xerrors.Errorf("encode settings: %w", err) - } - updated, err := api.Database.UpdateAIProvider(ctx, database.UpdateAIProviderParams{ - ID: row.ID, - Type: row.Type, - DisplayName: row.DisplayName, - Icon: row.Icon, - Enabled: row.Enabled, - BaseUrl: row.BaseUrl, - Settings: encoded, - // SettingsKeyID is set by the dbcrypt wrapper. - SettingsKeyID: sql.NullString{}, - }) - if err != nil { - return row, xerrors.Errorf("store resolved models: %w", err) - } - return updated, nil + return nil } // writeAIProviderResolutionError reports a failed Bedrock model resolution. The // provider keeps the identifiers the operator asked for, but without a // resolution the gateway cannot tell what an opaque profile ARN refers to, so -// it refuses to serve the provider until the write succeeds. +// it refuses to serve the provider until a later save resolves it. func (api *API) writeAIProviderResolutionError(ctx context.Context, rw http.ResponseWriter, err error) { var unresolvable bedrockProfileUnresolvableError - if !errors.As(err, &unresolvable) { + if !xerrors.As(err, &unresolvable) { writeAIProviderError(ctx, api.Logger, rw, err, "resolve bedrock inference profile", "Internal error resolving the Bedrock application inference profile.") return } diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go index 8d964596e4c..fd5830799a6 100644 --- a/coderd/ai_providers_bedrock_test.go +++ b/coderd/ai_providers_bedrock_test.go @@ -1,6 +1,7 @@ package coderd_test import ( + "context" "net/http" "net/http/httptest" "slices" @@ -8,9 +9,12 @@ import ( "sync" "testing" + "github.com/google/uuid" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) @@ -64,6 +68,14 @@ func respondWithModel(modelARN string) http.HandlerFunc { } } +func resolvedModels(ctx context.Context, t *testing.T, db database.Store, providerID uuid.UUID) []database.AIProviderBedrockResolvedModel { + t.Helper() + + rows, err := db.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, []uuid.UUID{providerID}) + require.NoError(t, err) + return rows +} + // TestAIProvidersBedrockProfileResolution drives provider writes against a mock // Bedrock control plane, so the AWS SDK path runs for real. // NOTE: no t.Parallel() because the subtests use t.Setenv. @@ -78,7 +90,8 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - client := coderdtest.New(t, nil) + db, ps := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -96,9 +109,12 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { // invocation target, and AWS attributes spend to them. require.Equal(t, testProfileARN, created.Settings.Bedrock.Model) require.Equal(t, testSmallFastProfileARN, created.Settings.Bedrock.SmallFastModel) - require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) - require.Equal(t, "anthropic.claude-haiku-4-5", created.Settings.Bedrock.ResolvedSmallFastModel) require.Len(t, paths(), 2, "each profile is resolved once") + + rows := resolvedModels(ctx, t, db, created.ID) + require.Len(t, rows, 1) + require.Equal(t, "anthropic.claude-opus-4-8", rows[0].ResolvedModel) + require.Equal(t, "anthropic.claude-haiku-4-5", rows[0].ResolvedSmallFastModel) }) t.Run("CreateLeavesPlainModelIDsUnresolved", func(t *testing.T) { @@ -107,7 +123,8 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - client := coderdtest.New(t, nil) + db, ps := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -120,10 +137,8 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.NotNil(t, created.Settings.Bedrock) - require.Empty(t, created.Settings.Bedrock.ResolvedModel) - require.Empty(t, created.Settings.Bedrock.ResolvedSmallFastModel) require.Empty(t, paths()) + require.Empty(t, resolvedModels(ctx, t, db, created.ID), "plain model ids are already model identities") }) t.Run("CreateRejectsUnresolvableProfile", func(t *testing.T) { @@ -135,7 +150,8 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - client := coderdtest.New(t, nil) + db, ps := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -158,42 +174,15 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { providers, err := client.AIProviders(ctx) require.NoError(t, err) require.Len(t, providers, 1) - require.Empty(t, providers[0].Settings.Bedrock.ResolvedModel) - }) - - t.Run("CreateRejectsClientSuppliedResolution", func(t *testing.T) { - url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { - t.Error("Bedrock called for a rejected request") - }) - t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - - client := coderdtest.New(t, nil) - _ = coderdtest.CreateFirstUser(t, client) - ctx := testutil.Context(t, testutil.WaitLong) - - settings := bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") - settings.Bedrock.ResolvedModel = "anthropic.claude-opus-4-8" - - //nolint:gocritic // Owner role is the audience for this endpoint. - _, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ - Name: "bedrock-spoofed", - Type: codersdk.AIProviderTypeBedrock, - BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", - Enabled: true, - Settings: *settings, - }) - var sdkErr *codersdk.Error - require.ErrorAs(t, err, &sdkErr) - require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) - require.Contains(t, sdkErr.Error(), "resolved_model") - require.Empty(t, paths(), "validation rejects the write before any AWS call") + require.Empty(t, resolvedModels(ctx, t, db, providers[0].ID)) }) t.Run("UpdateReresolvesChangedProfile", func(t *testing.T) { url, _ := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - client := coderdtest.New(t, nil) + db, ps := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -206,7 +195,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Empty(t, created.Settings.Bedrock.ResolvedModel) + require.Empty(t, resolvedModels(ctx, t, db, created.ID)) //nolint:gocritic // Owner role is the audience for this endpoint. updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ @@ -214,15 +203,20 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) require.NoError(t, err) require.Equal(t, testProfileARN, updated.Settings.Bedrock.Model) - require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) - require.Empty(t, updated.Settings.Bedrock.ResolvedSmallFastModel) + + rows := resolvedModels(ctx, t, db, created.ID) + require.Len(t, rows, 1) + require.Equal(t, "anthropic.claude-opus-4-8", rows[0].ResolvedModel) + // The small/fast model is a plain ID, so it resolves to itself. + require.Equal(t, "anthropic.claude-haiku-4-5", rows[0].ResolvedSmallFastModel) }) t.Run("UpdateClearsResolutionWhenProfileReplacedByModelID", func(t *testing.T) { url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - client := coderdtest.New(t, nil) + db, ps := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -235,23 +229,24 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) + require.Len(t, resolvedModels(ctx, t, db, created.ID), 1) callsAfterCreate := len(paths()) //nolint:gocritic // Owner role is the audience for this endpoint. - updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + _, err = client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ Settings: bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Empty(t, updated.Settings.Bedrock.ResolvedModel, "a plain model id resolves to itself") + require.Empty(t, resolvedModels(ctx, t, db, created.ID), "a plain model id needs no mapping") require.Len(t, paths(), callsAfterCreate, "no profile is left to resolve") }) - t.Run("UpdateWithoutSettingsSkipsResolution", func(t *testing.T) { + t.Run("UpdateWithoutSettingsKeepsResolution", func(t *testing.T) { url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - client := coderdtest.New(t, nil) + db, ps := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -268,11 +263,41 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { enabled := false //nolint:gocritic // Owner role is the audience for this endpoint. - updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + _, err = client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ Enabled: &enabled, }) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) + + rows := resolvedModels(ctx, t, db, created.ID) + require.Len(t, rows, 1) + require.Equal(t, "anthropic.claude-opus-4-8", rows[0].ResolvedModel) require.Len(t, paths(), callsAfterCreate, "an unrelated update does not call AWS") }) + + t.Run("DeleteClearsResolution", func(t *testing.T) { + url, _ := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + db, ps := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) + _ = coderdtest.CreateFirstUser(t, client) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-delete", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + require.Len(t, resolvedModels(ctx, t, db, created.ID), 1) + + // Providers are soft-deleted, so the mapping has to be removed + // explicitly rather than by the foreign key. + //nolint:gocritic // Owner role is the audience for this endpoint. + require.NoError(t, client.DeleteAIProvider(ctx, created.ID.String())) + require.Empty(t, resolvedModels(ctx, t, db, created.ID)) + }) } diff --git a/coderd/aibridge/bedrock.go b/coderd/aibridge/bedrock.go index 49182d484be..16fa9f2f57b 100644 --- a/coderd/aibridge/bedrock.go +++ b/coderd/aibridge/bedrock.go @@ -25,16 +25,14 @@ func BedrockConfig(baseURL string, bedrock *codersdk.AIProviderBedrockSettings) return nil } return &aibridgeconfig.AWSBedrock{ - BaseURL: baseURL, - Region: settings.Region, - AccessKey: ptr.NilToEmpty(settings.AccessKey), - AccessKeySecret: ptr.NilToEmpty(settings.AccessKeySecret), - Model: settings.Model, - SmallFastModel: settings.SmallFastModel, - RoleARN: settings.RoleARN, - ExternalID: settings.ExternalID, - Protocol: aibridgeconfig.BedrockProtocol(settings.ResolvedProtocol()), - ResolvedModel: settings.ResolvedModel, - ResolvedSmallFastModel: settings.ResolvedSmallFastModel, + BaseURL: baseURL, + Region: settings.Region, + AccessKey: ptr.NilToEmpty(settings.AccessKey), + AccessKeySecret: ptr.NilToEmpty(settings.AccessKeySecret), + Model: settings.Model, + SmallFastModel: settings.SmallFastModel, + RoleARN: settings.RoleARN, + ExternalID: settings.ExternalID, + Protocol: aibridgeconfig.BedrockProtocol(settings.ResolvedProtocol()), } } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index cf76f9e27c0..db544ce69cd 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -100,6 +100,7 @@ type store interface { // any in-flight env seed holding LockIDAIProvidersEnvSeed. GetAIProviders(ctx context.Context, arg database.GetAIProvidersParams) ([]database.AIProvider, error) GetAIProviderKeysByProviderIDs(ctx context.Context, providerIDs []uuid.UUID) ([]database.AIProviderKey, error) + GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIDs []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) InTx(func(database.Store) error, *database.TxOptions) error } @@ -955,8 +956,9 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ ctx = dbauthz.AsAIBridged(ctx) var ( - rows []database.AIProvider - keysByProvider map[uuid.UUID][]database.AIProviderKey + rows []database.AIProvider + keysByProvider map[uuid.UUID][]database.AIProviderKey + resolvedByProvider map[uuid.UUID]database.AIProviderBedrockResolvedModel ) // Wrap both reads in a read-only transaction so the provider list and the // key list are consistent with each other, and so the seed lock is held @@ -994,6 +996,19 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ for _, k := range keyRows { keysByProvider[k.ProviderID] = append(keysByProvider[k.ProviderID], k) } + + // Bedrock application inference profile ARNs are opaque, so the models + // they refer to are resolved when the provider is written and read back + // here. A provider without a mapping either configures plain model IDs + // or could not be resolved. + resolvedRows, err := tx.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, ids) + if err != nil { + return xerrors.Errorf("get ai provider resolved models: %w", err) + } + resolvedByProvider = make(map[uuid.UUID]database.AIProviderBedrockResolvedModel, len(resolvedRows)) + for _, r := range resolvedRows { + resolvedByProvider[r.AIProviderID] = r + } return nil }, &database.TxOptions{ReadOnly: true, TxIdentifier: "get_ai_providers"}) if err != nil { @@ -1002,7 +1017,7 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ providers := make([]*proto.AIProvider, 0, len(rows)) for _, row := range rows { - p, err := aiProviderToProto(row, keysByProvider[row.ID]) + p, err := aiProviderToProto(row, keysByProvider[row.ID], resolvedByProvider[row.ID]) if err != nil { // Skip the offending row rather than failing the whole fetch: // one row with a corrupt settings blob must not break provider @@ -1180,7 +1195,11 @@ func parseOptionalInt32(n *int32) sql.NullInt32 { // providers) to the proto representation served to AI Gateway daemons. Keys and // Bedrock settings are only attached for enabled providers; disabled providers // never call upstream so their secrets are withheld. -func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey) (*proto.AIProvider, error) { +// +// resolved carries the models the provider's application inference profile ARNs +// refer to, and is the zero value when the provider configures plain model IDs +// or when its profiles have not been resolved. +func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey, resolved database.AIProviderBedrockResolvedModel) (*proto.AIProvider, error) { p := &proto.AIProvider{ Name: row.Name, Type: string(row.Type), @@ -1213,8 +1232,8 @@ func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey) ( RoleArn: settings.Bedrock.RoleARN, ExternalId: settings.Bedrock.ExternalID, Protocol: string(settings.Bedrock.Protocol), - ResolvedModel: settings.Bedrock.ResolvedModel, - ResolvedSmallFastModel: settings.Bedrock.ResolvedSmallFastModel, + ResolvedModel: resolved.ResolvedModel, + ResolvedSmallFastModel: resolved.ResolvedSmallFastModel, } } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index e27731ab0fe..e435a9064e0 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2118,6 +2118,13 @@ func (q *querier) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) (databas return q.db.DeleteAIGatewayKey(ctx, id) } +func (q *querier) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil { + return err + } + return q.db.DeleteAIProviderBedrockResolvedModels(ctx, aiProviderID) +} + func (q *querier) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAIProvider); err != nil { return err @@ -2953,6 +2960,13 @@ func (q *querier) GetAIModelPrices(ctx context.Context, arg database.GetAIModelP return q.db.GetAIModelPrices(ctx, arg) } +func (q *querier) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIDs []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIProvider); err != nil { + return nil, err + } + return q.db.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, aiProviderIDs) +} + func (q *querier) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIProvider); err != nil { return database.AIProvider{}, err @@ -9061,6 +9075,13 @@ func (q *querier) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAI return q.db.UpsertAIModelPrices(ctx, arg) } +func (q *querier) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg database.UpsertAIProviderBedrockResolvedModelsParams) error { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil { + return err + } + return q.db.UpsertAIProviderBedrockResolvedModels(ctx, arg) +} + func (q *querier) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) { if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceAiSeat); err != nil { return false, err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 5dd49f43068..1ae8ea99916 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -7389,6 +7389,29 @@ func (s *MethodTestSuite) TestAIBridge() { dbm.EXPECT().UpdateEncryptedAIProviderSettings(gomock.Any(), arg).Return(provider, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns(provider) })) + s.Run("UpsertAIProviderBedrockResolvedModels", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + provider := testutil.Fake(s.T(), faker, database.AIProvider{}) + arg := database.UpsertAIProviderBedrockResolvedModelsParams{ + AIProviderID: provider.ID, + ResolvedModel: "anthropic.claude-opus-4-8", + ResolvedSmallFastModel: "anthropic.claude-haiku-4-5", + } + dbm.EXPECT().UpsertAIProviderBedrockResolvedModels(gomock.Any(), arg).Return(nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns() + })) + s.Run("DeleteAIProviderBedrockResolvedModels", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + provider := testutil.Fake(s.T(), faker, database.AIProvider{}) + dbm.EXPECT().DeleteAIProviderBedrockResolvedModels(gomock.Any(), provider.ID).Return(nil).AnyTimes() + check.Args(provider.ID).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns() + })) + s.Run("GetAIProviderBedrockResolvedModelsByProviderIDs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + providerA := testutil.Fake(s.T(), faker, database.AIProvider{}) + providerB := testutil.Fake(s.T(), faker, database.AIProvider{}) + providerIDs := []uuid.UUID{providerA.ID, providerB.ID} + resolved := testutil.Fake(s.T(), faker, database.AIProviderBedrockResolvedModel{AIProviderID: providerA.ID}) + dbm.EXPECT().GetAIProviderBedrockResolvedModelsByProviderIDs(gomock.Any(), providerIDs).Return([]database.AIProviderBedrockResolvedModel{resolved}, nil).AnyTimes() + check.Args(providerIDs).Asserts(rbac.ResourceAIProvider, policy.ActionRead).Returns([]database.AIProviderBedrockResolvedModel{resolved}) + })) s.Run("GetAIProviderKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { key := testutil.Fake(s.T(), faker, database.AIProviderKey{}) dbm.EXPECT().GetAIProviderKeyByID(gomock.Any(), key.ID).Return(key, nil).AnyTimes() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 365899afbd2..356e8e688d5 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -424,6 +424,14 @@ func (m queryMetricsStore) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) return r0, r1 } +func (m queryMetricsStore) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { + start := time.Now() + r0 := m.s.DeleteAIProviderBedrockResolvedModels(ctx, aiProviderID) + m.queryLatencies.WithLabelValues("DeleteAIProviderBedrockResolvedModels").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteAIProviderBedrockResolvedModels").Inc() + return r0 +} + func (m queryMetricsStore) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error { start := time.Now() r0 := m.s.DeleteAIProviderByID(ctx, id) @@ -1184,6 +1192,14 @@ func (m queryMetricsStore) GetAIModelPrices(ctx context.Context, arg database.Ge return r0, r1 } +func (m queryMetricsStore) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) { + start := time.Now() + r0, r1 := m.s.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, aiProviderIds) + m.queryLatencies.WithLabelValues("GetAIProviderBedrockResolvedModelsByProviderIDs").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIProviderBedrockResolvedModelsByProviderIDs").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { start := time.Now() r0, r1 := m.s.GetAIProviderByID(ctx, id) @@ -6320,6 +6336,14 @@ func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, arg database return r0 } +func (m queryMetricsStore) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg database.UpsertAIProviderBedrockResolvedModelsParams) error { + start := time.Now() + r0 := m.s.UpsertAIProviderBedrockResolvedModels(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertAIProviderBedrockResolvedModels").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIProviderBedrockResolvedModels").Inc() + return r0 +} + func (m queryMetricsStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) { start := time.Now() r0, r1 := m.s.UpsertAISeatState(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index f52286cf511..34e3eb6d98f 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -676,6 +676,20 @@ func (mr *MockStoreMockRecorder) DeleteAIGatewayKey(ctx, id any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAIGatewayKey", reflect.TypeOf((*MockStore)(nil).DeleteAIGatewayKey), ctx, id) } +// DeleteAIProviderBedrockResolvedModels mocks base method. +func (m *MockStore) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAIProviderBedrockResolvedModels", ctx, aiProviderID) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteAIProviderBedrockResolvedModels indicates an expected call of DeleteAIProviderBedrockResolvedModels. +func (mr *MockStoreMockRecorder) DeleteAIProviderBedrockResolvedModels(ctx, aiProviderID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAIProviderBedrockResolvedModels", reflect.TypeOf((*MockStore)(nil).DeleteAIProviderBedrockResolvedModels), ctx, aiProviderID) +} + // DeleteAIProviderByID mocks base method. func (m *MockStore) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error { m.ctrl.T.Helper() @@ -2068,6 +2082,21 @@ func (mr *MockStoreMockRecorder) GetAIModelPrices(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIModelPrices", reflect.TypeOf((*MockStore)(nil).GetAIModelPrices), ctx, arg) } +// GetAIProviderBedrockResolvedModelsByProviderIDs mocks base method. +func (m *MockStore) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAIProviderBedrockResolvedModelsByProviderIDs", ctx, aiProviderIds) + ret0, _ := ret[0].([]database.AIProviderBedrockResolvedModel) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAIProviderBedrockResolvedModelsByProviderIDs indicates an expected call of GetAIProviderBedrockResolvedModelsByProviderIDs. +func (mr *MockStoreMockRecorder) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, aiProviderIds any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIProviderBedrockResolvedModelsByProviderIDs", reflect.TypeOf((*MockStore)(nil).GetAIProviderBedrockResolvedModelsByProviderIDs), ctx, aiProviderIds) +} + // GetAIProviderByID mocks base method. func (m *MockStore) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { m.ctrl.T.Helper() @@ -11908,6 +11937,20 @@ func (mr *MockStoreMockRecorder) UpsertAIModelPrices(ctx, arg any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIModelPrices", reflect.TypeOf((*MockStore)(nil).UpsertAIModelPrices), ctx, arg) } +// UpsertAIProviderBedrockResolvedModels mocks base method. +func (m *MockStore) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg database.UpsertAIProviderBedrockResolvedModelsParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpsertAIProviderBedrockResolvedModels", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// UpsertAIProviderBedrockResolvedModels indicates an expected call of UpsertAIProviderBedrockResolvedModels. +func (mr *MockStoreMockRecorder) UpsertAIProviderBedrockResolvedModels(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIProviderBedrockResolvedModels", reflect.TypeOf((*MockStore)(nil).UpsertAIProviderBedrockResolvedModels), ctx, arg) +} + // UpsertAISeatState mocks base method. func (m *MockStore) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 6e93c540739..8bda5f9a52e 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1603,6 +1603,16 @@ COMMENT ON TABLE ai_model_prices IS 'Per-model token prices used by AI Bridge to COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.'; +CREATE TABLE ai_provider_bedrock_resolved_models ( + ai_provider_id uuid NOT NULL, + resolved_model text NOT NULL, + resolved_small_fast_model text NOT NULL +); + +COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_model IS 'The model ID behind the provider''s configured model identifier. Equal to the configured value when that value is already a model ID.'; + +COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_small_fast_model IS 'resolved_model for the provider''s configured small/fast model identifier.'; + CREATE TABLE ai_provider_keys ( id uuid DEFAULT gen_random_uuid() NOT NULL, provider_id uuid NOT NULL, @@ -4372,6 +4382,9 @@ ALTER TABLE ONLY ai_gateway_keys ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source); +ALTER TABLE ONLY ai_provider_bedrock_resolved_models + ADD CONSTRAINT ai_provider_bedrock_resolved_models_pkey PRIMARY KEY (ai_provider_id); + ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id); @@ -5254,6 +5267,9 @@ COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS 'U the uniqueness requirement. A trigger allows us to enforce uniqueness going forward without requiring a migration to clean up historical data.'; +ALTER TABLE ONLY ai_provider_bedrock_resolved_models + ADD CONSTRAINT ai_provider_bedrock_resolved_models_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; + ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_api_key_key_id_fkey FOREIGN KEY (api_key_key_id) REFERENCES dbcrypt_keys(active_key_digest); diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index 251ce1aec5c..d23b2fe5247 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -6,6 +6,7 @@ type ForeignKeyConstraint string // ForeignKeyConstraint enums. const ( + ForeignKeyAIProviderBedrockResolvedModelsAIProviderID ForeignKeyConstraint = "ai_provider_bedrock_resolved_models_ai_provider_id_fkey" // ALTER TABLE ONLY ai_provider_bedrock_resolved_models ADD CONSTRAINT ai_provider_bedrock_resolved_models_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; ForeignKeyAIProviderKeysAPIKeyKeyID ForeignKeyConstraint = "ai_provider_keys_api_key_key_id_fkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_api_key_key_id_fkey FOREIGN KEY (api_key_key_id) REFERENCES dbcrypt_keys(active_key_digest); ForeignKeyAIProviderKeysProviderID ForeignKeyConstraint = "ai_provider_keys_provider_id_fkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; ForeignKeyAIProvidersSettingsKeyID ForeignKeyConstraint = "ai_providers_settings_key_id_fkey" // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_settings_key_id_fkey FOREIGN KEY (settings_key_id) REFERENCES dbcrypt_keys(active_key_digest); diff --git a/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql b/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql new file mode 100644 index 00000000000..6e2f68cf04f --- /dev/null +++ b/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql @@ -0,0 +1 @@ +DROP TABLE ai_provider_bedrock_resolved_models; diff --git a/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql b/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql new file mode 100644 index 00000000000..9dd63146580 --- /dev/null +++ b/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql @@ -0,0 +1,17 @@ +-- Application inference profile ARNs are opaque: they identify a Bedrock +-- billing wrapper, not a model. This table records the model each configured +-- ARN resolves to, so the gateway can detect capabilities, price usage, and +-- record interceptions without calling the Bedrock control plane itself. +-- +-- A provider has a row only while one of its identifiers is an ARN. Providers +-- configured with plain model IDs need no mapping, and neither does a provider +-- whose profile could not be resolved: the gateway refuses to serve it. +CREATE TABLE ai_provider_bedrock_resolved_models ( + ai_provider_id uuid PRIMARY KEY REFERENCES ai_providers (id) ON DELETE CASCADE, + resolved_model text NOT NULL, + resolved_small_fast_model text NOT NULL +); + +COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_model IS 'The model ID behind the provider''s configured model identifier. Equal to the configured value when that value is already a model ID.'; + +COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_small_fast_model IS 'resolved_model for the provider''s configured small/fast model identifier.'; diff --git a/coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql b/coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql new file mode 100644 index 00000000000..9b1fd9996f7 --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql @@ -0,0 +1,31 @@ +INSERT INTO ai_providers ( + id, + type, + name, + display_name, + enabled, + deleted, + base_url, + settings +) VALUES + ( + '8e3c6e18-2b75-4c3f-9b35-9d1c6f4e1a04', + 'bedrock', + 'bedrock-inference-profile', + 'Bedrock via Application Inference Profile (Fixture)', + TRUE, + FALSE, + 'https://bedrock-runtime.us-west-2.amazonaws.com/', + '{"_type":"bedrock","_version":1,"region":"us-west-2","model":"arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/fixtureprofile","small_fast_model":"anthropic.claude-3-5-haiku-20241022-v1:0","access_key":"fixture-bedrock-access-key","access_key_secret":"fixture-bedrock-access-key-secret"}' + ); + +INSERT INTO ai_provider_bedrock_resolved_models ( + ai_provider_id, + resolved_model, + resolved_small_fast_model +) VALUES + ( + '8e3c6e18-2b75-4c3f-9b35-9d1c6f4e1a04', + 'anthropic.claude-sonnet-4-5-20250929-v1:0', + 'anthropic.claude-3-5-haiku-20241022-v1:0' + ); diff --git a/coderd/database/models.go b/coderd/database/models.go index b5a1340b353..67df38495eb 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4991,6 +4991,14 @@ type AIProvider struct { Icon string `db:"icon" json:"icon"` } +type AIProviderBedrockResolvedModel struct { + AIProviderID uuid.UUID `db:"ai_provider_id" json:"ai_provider_id"` + // The model ID behind the provider's configured model identifier. Equal to the configured value when that value is already a model ID. + ResolvedModel string `db:"resolved_model" json:"resolved_model"` + // resolved_model for the provider's configured small/fast model identifier. + ResolvedSmallFastModel string `db:"resolved_small_fast_model" json:"resolved_small_fast_model"` +} + // API keys associated with AI providers. Bedrock providers have zero keys (they authenticate via settings). OpenAI and Anthropic providers have one or more keys for failover. type AIProviderKey struct { ID uuid.UUID `db:"id" json:"id"` diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 88e86e97304..d18293895c1 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -117,6 +117,7 @@ type sqlcQuerier interface { CreateUserSecret(ctx context.Context, arg CreateUserSecretParams) (UserSecret, error) CustomRoles(ctx context.Context, arg CustomRolesParams) ([]CustomRole, error) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) (DeleteAIGatewayKeyRow, error) + DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error DeleteAIProviderKey(ctx context.Context, id uuid.UUID) error DeleteAPIKeyByID(ctx context.Context, id string) error @@ -340,6 +341,7 @@ type sqlcQuerier interface { // each source forms its own group and nothing collapses. Every other source // contributes the same constant, leaving the key as (provider, model). GetAIModelPrices(ctx context.Context, arg GetAIModelPricesParams) ([]AIModelPrice, error) + GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]AIProviderBedrockResolvedModel, error) GetAIProviderByID(ctx context.Context, id uuid.UUID) (AIProvider, error) // Lock the provider row until the model-config write completes. The // transaction alone does not stop a concurrent soft-delete or disable @@ -1667,6 +1669,10 @@ type sqlcQuerier interface { // differs, so updated_at records when a price last changed. Prices are // nullable and a NULL on either side counts as a difference. UpsertAIModelPrices(ctx context.Context, arg UpsertAIModelPricesParams) error + // Records the models an application inference profile ARN resolves to. The + // provider write path resolves the ARN through the Bedrock control plane and + // stores the answer here, so the gateway never has to. + UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg UpsertAIProviderBedrockResolvedModelsParams) error // Returns true if a new rows was inserted, false otherwise. UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, error) UpsertAnnouncementBanners(ctx context.Context, value string) error diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index ccbd4a5af2a..f017339ddee 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -553,6 +553,18 @@ func (q *sqlQuerier) UpdateEncryptedAIProviderKey(ctx context.Context, arg Updat return i, err } +const deleteAIProviderBedrockResolvedModels = `-- name: DeleteAIProviderBedrockResolvedModels :exec +DELETE FROM + ai_provider_bedrock_resolved_models +WHERE + ai_provider_id = $1::uuid +` + +func (q *sqlQuerier) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { + _, err := q.db.ExecContext(ctx, deleteAIProviderBedrockResolvedModels, aiProviderID) + return err +} + const deleteAIProviderByID = `-- name: DeleteAIProviderByID :exec UPDATE ai_providers @@ -569,6 +581,38 @@ func (q *sqlQuerier) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) err return err } +const getAIProviderBedrockResolvedModelsByProviderIDs = `-- name: GetAIProviderBedrockResolvedModelsByProviderIDs :many +SELECT + ai_provider_id, resolved_model, resolved_small_fast_model +FROM + ai_provider_bedrock_resolved_models +WHERE + ai_provider_id = ANY($1::uuid[]) +` + +func (q *sqlQuerier) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]AIProviderBedrockResolvedModel, error) { + rows, err := q.db.QueryContext(ctx, getAIProviderBedrockResolvedModelsByProviderIDs, pq.Array(aiProviderIds)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []AIProviderBedrockResolvedModel + for rows.Next() { + var i AIProviderBedrockResolvedModel + if err := rows.Scan(&i.AIProviderID, &i.ResolvedModel, &i.ResolvedSmallFastModel); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getAIProviderByID = `-- name: GetAIProviderByID :one SELECT id, type, name, display_name, enabled, deleted, base_url, settings, settings_key_id, created_at, updated_at, icon @@ -884,6 +928,30 @@ func (q *sqlQuerier) UpdateEncryptedAIProviderSettings(ctx context.Context, arg return i, err } +const upsertAIProviderBedrockResolvedModels = `-- name: UpsertAIProviderBedrockResolvedModels :exec +INSERT INTO + ai_provider_bedrock_resolved_models (ai_provider_id, resolved_model, resolved_small_fast_model) +VALUES + ($1::uuid, $2::text, $3::text) +ON CONFLICT (ai_provider_id) DO UPDATE SET + resolved_model = $2::text, + resolved_small_fast_model = $3::text +` + +type UpsertAIProviderBedrockResolvedModelsParams struct { + AIProviderID uuid.UUID `db:"ai_provider_id" json:"ai_provider_id"` + ResolvedModel string `db:"resolved_model" json:"resolved_model"` + ResolvedSmallFastModel string `db:"resolved_small_fast_model" json:"resolved_small_fast_model"` +} + +// Records the models an application inference profile ARN resolves to. The +// provider write path resolves the ARN through the Bedrock control plane and +// stores the answer here, so the gateway never has to. +func (q *sqlQuerier) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg UpsertAIProviderBedrockResolvedModelsParams) error { + _, err := q.db.ExecContext(ctx, upsertAIProviderBedrockResolvedModels, arg.AIProviderID, arg.ResolvedModel, arg.ResolvedSmallFastModel) + return err +} + const calculateAIBridgeInterceptionsTelemetrySummary = `-- name: CalculateAIBridgeInterceptionsTelemetrySummary :one WITH interceptions_in_range AS ( -- Get all matching interceptions in the given timeframe. diff --git a/coderd/database/queries/ai_providers.sql b/coderd/database/queries/ai_providers.sql index 2971918e46f..6e003ce7c20 100644 --- a/coderd/database/queries/ai_providers.sql +++ b/coderd/database/queries/ai_providers.sql @@ -106,3 +106,29 @@ WHERE id = @id::uuid RETURNING *; + +-- name: UpsertAIProviderBedrockResolvedModels :exec +-- Records the models an application inference profile ARN resolves to. The +-- provider write path resolves the ARN through the Bedrock control plane and +-- stores the answer here, so the gateway never has to. +INSERT INTO + ai_provider_bedrock_resolved_models (ai_provider_id, resolved_model, resolved_small_fast_model) +VALUES + (@ai_provider_id::uuid, @resolved_model::text, @resolved_small_fast_model::text) +ON CONFLICT (ai_provider_id) DO UPDATE SET + resolved_model = @resolved_model::text, + resolved_small_fast_model = @resolved_small_fast_model::text; + +-- name: DeleteAIProviderBedrockResolvedModels :exec +DELETE FROM + ai_provider_bedrock_resolved_models +WHERE + ai_provider_id = @ai_provider_id::uuid; + +-- name: GetAIProviderBedrockResolvedModelsByProviderIDs :many +SELECT + * +FROM + ai_provider_bedrock_resolved_models +WHERE + ai_provider_id = ANY(@ai_provider_ids::uuid[]); diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 9d2390a6468..332ae755c0a 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -9,6 +9,7 @@ const ( UniqueAgentStatsPkey UniqueConstraint = "agent_stats_pkey" // ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); UniqueAIGatewayKeysPkey UniqueConstraint = "ai_gateway_keys_pkey" // ALTER TABLE ONLY ai_gateway_keys ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id); UniqueAIModelPricesPkey UniqueConstraint = "ai_model_prices_pkey" // ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source); + UniqueAIProviderBedrockResolvedModelsPkey UniqueConstraint = "ai_provider_bedrock_resolved_models_pkey" // ALTER TABLE ONLY ai_provider_bedrock_resolved_models ADD CONSTRAINT ai_provider_bedrock_resolved_models_pkey PRIMARY KEY (ai_provider_id); UniqueAIProviderKeysPkey UniqueConstraint = "ai_provider_keys_pkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id); UniqueAIProvidersPkey UniqueConstraint = "ai_providers_pkey" // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_pkey PRIMARY KEY (id); UniqueAISeatStatePkey UniqueConstraint = "ai_seat_state_pkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id); diff --git a/codersdk/aiproviders.go b/codersdk/aiproviders.go index 058351d782e..6e81e3616a1 100644 --- a/codersdk/aiproviders.go +++ b/codersdk/aiproviders.go @@ -293,7 +293,6 @@ func (req CreateAIProviderRequest) Validate() []ValidationError { Detail: "external_id is server-generated and cannot be set", }) } - validations = append(validations, validateAIProviderBedrockResolvedModelsUnset(*req.Settings.Bedrock)...) validations = append(validations, validateAIProviderBedrockMantleRegion(*req.Settings.Bedrock)...) validations = append(validations, validateAIProviderBedrockModels(*req.Settings.Bedrock)...) } @@ -429,26 +428,6 @@ func validateAIProviderBedrockModels(b AIProviderBedrockSettings) []ValidationEr return validations } -// validateAIProviderBedrockResolvedModelsUnset rejects client-supplied resolved -// model identifiers on create. The server resolves them through AWS and owns -// the values, the same way it owns the STS external ID. -func validateAIProviderBedrockResolvedModelsUnset(b AIProviderBedrockSettings) []ValidationError { - var validations []ValidationError - if b.ResolvedModel != "" { - validations = append(validations, ValidationError{ - Field: "settings.resolved_model", - Detail: "resolved_model is server-resolved and cannot be set", - }) - } - if b.ResolvedSmallFastModel != "" { - validations = append(validations, ValidationError{ - Field: "settings.resolved_small_fast_model", - Detail: "resolved_small_fast_model is server-resolved and cannot be set", - }) - } - return validations -} - func validateAIProviderRoleARN(roleARN string) []ValidationError { if roleARN == "" { return nil diff --git a/codersdk/aiproviders_bedrock.go b/codersdk/aiproviders_bedrock.go index d9c961cbb97..b3bc94e9e4b 100644 --- a/codersdk/aiproviders_bedrock.go +++ b/codersdk/aiproviders_bedrock.go @@ -61,33 +61,6 @@ type AIProviderBedrockSettings struct { // AIProviderBedrockProtocolInvokeModel, so existing rows keep the legacy // behavior. Protocol AIProviderBedrockProtocol `json:"protocol,omitempty"` - // ResolvedModel is the Bedrock model ID behind Model. It differs from Model - // only when Model is an application inference profile ARN, whose identifier - // is opaque. The server resolves it through AWS when the provider is - // written and owns the value: create and update reject a client-supplied - // one that differs from the stored value. - ResolvedModel string `json:"resolved_model,omitempty"` - // ResolvedSmallFastModel is ResolvedModel for SmallFastModel. - ResolvedSmallFastModel string `json:"resolved_small_fast_model,omitempty"` -} - -// ModelIdentity returns the model ID the gateway records for usage, pricing, -// and capability detection. It is the resolved value when the configured -// identifier needed resolution, and the configured identifier otherwise. -func (b AIProviderBedrockSettings) ModelIdentity() string { - if b.ResolvedModel != "" { - return b.ResolvedModel - } - return b.Model -} - -// SmallFastModelIdentity is [AIProviderBedrockSettings.ModelIdentity] for the -// small/fast model. -func (b AIProviderBedrockSettings) SmallFastModelIdentity() string { - if b.ResolvedSmallFastModel != "" { - return b.ResolvedSmallFastModel - } - return b.SmallFastModel } // ResolvedProtocol returns the configured protocol, mapping the empty value to diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index ac2bc17325b..1624e90af5b 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -482,18 +482,6 @@ export interface AIProviderBedrockSettings { * behavior. */ readonly protocol?: AIProviderBedrockProtocol; - /** - * ResolvedModel is the Bedrock model ID behind Model. It differs from Model - * only when Model is an application inference profile ARN, whose identifier - * is opaque. The server resolves it through AWS when the provider is - * written and owns the value: create and update reject a client-supplied - * one that differs from the stored value. - */ - readonly resolved_model?: string; - /** - * ResolvedSmallFastModel is ResolvedModel for SmallFastModel. - */ - readonly resolved_small_fast_model?: string; } // From codersdk/aiproviders_bedrock.go From 5f4d2a3b7ae85b5ec6639d150922d3dd910cbb4b Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 16:04:50 +0000 Subject: [PATCH 08/25] refactor(aibridge): derive bedrock model identity from the provider config --- aibridge/intercept/messages/base.go | 19 +++++++--- .../intercept/messages/base_internal_test.go | 38 +++++++++---------- aibridge/provider/anthropic.go | 11 +----- .../provider/bedrock_inference_profile.go | 26 ------------- ...bedrock_inference_profile_internal_test.go | 24 +++++------- 5 files changed, 43 insertions(+), 75 deletions(-) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 67b19b1bdbc..df9f9ef4fa8 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -78,11 +78,20 @@ type BedrockRuntime struct { } // NewBedrockRuntime bundles the Bedrock config and credentials with the model -// IDs behind the configured identifiers. The resolved IDs differ from the -// configured ones only when those are application inference profile ARNs, which -// are opaque and must be resolved through AWS; every other identifier resolves -// to itself. -func NewBedrockRuntime(cfg aibconfig.AWSBedrock, creds aws.CredentialsProvider, resolvedModel, resolvedSmallFastModel string) *BedrockRuntime { +// IDs behind the configured identifiers. +// +// An identifier resolves to itself unless it is an application inference +// profile ARN, which is opaque. Those are resolved through AWS when the +// provider is written and arrive on cfg, so construction never calls AWS. +func NewBedrockRuntime(cfg aibconfig.AWSBedrock, creds aws.CredentialsProvider) *BedrockRuntime { + resolvedModel := cfg.ResolvedModel + if resolvedModel == "" { + resolvedModel = cfg.Model + } + resolvedSmallFastModel := cfg.ResolvedSmallFastModel + if resolvedSmallFastModel == "" { + resolvedSmallFastModel = cfg.SmallFastModel + } return &BedrockRuntime{ Cfg: cfg, Creds: creds, diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index 38c7d185c00..1bbc1c4b9de 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -179,7 +179,7 @@ func TestAWSBedrockValidation(t *testing.T) { t.Parallel() base := &interceptionBase{ - bedrock: NewBedrockRuntime(tt.cfg, credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), "", ""), + bedrock: NewBedrockRuntime(tt.cfg, credentials.NewStaticCredentialsProvider("test-key", "test-secret", "")), } opts, err := base.withBedrockInvokeModelOptions(context.Background()) @@ -217,9 +217,11 @@ func TestModelForBedrockInvokeModel(t *testing.T) { ) runtime := NewBedrockRuntime(config.AWSBedrock{ - Model: profileARN, - SmallFastModel: smallFastProfileARN, - }, nil, "anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") + Model: profileARN, + SmallFastModel: smallFastProfileARN, + ResolvedModel: "anthropic.claude-opus-4-8", + ResolvedSmallFastModel: "anthropic.claude-haiku-4-5", + }, nil) tests := []struct { name string @@ -274,9 +276,11 @@ func TestSmallFastModelCapturedAtConstruction(t *testing.T) { ) runtime := NewBedrockRuntime(config.AWSBedrock{ - Model: profileARN, - SmallFastModel: smallFastProfileARN, - }, nil, "anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") + Model: profileARN, + SmallFastModel: smallFastProfileARN, + ResolvedModel: "anthropic.claude-opus-4-8", + ResolvedSmallFastModel: "anthropic.claude-haiku-4-5", + }, nil) const haikuPayload = `{"model":"claude-haiku-4-5","max_tokens":10000}` const opusPayload = `{"model":"claude-opus-4-8","max_tokens":10000}` @@ -336,7 +340,7 @@ func TestModelForPlainBedrockModelID(t *testing.T) { bedrock: NewBedrockRuntime(config.AWSBedrock{ Model: "eu.anthropic.claude-opus-4-8", SmallFastModel: "anthropic.claude-haiku-4-5", - }, nil, "eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), + }, nil), logger: slog.Make(), } @@ -958,18 +962,14 @@ func TestAugmentRequestForBedrock_AdaptiveThinking(t *testing.T) { } // Plain model IDs resolve to themselves; an application inference - // profile ARN resolves to the model behind it. - resolvedModel := tc.resolvedModel - if resolvedModel == "" { - resolvedModel = tc.bedrockModel - } - i := &interceptionBase{ reqPayload: mustMessagesPayload(t, tc.requestBody), bedrock: NewBedrockRuntime(config.AWSBedrock{ - Model: tc.bedrockModel, - SmallFastModel: "anthropic.claude-haiku-3-5", - }, nil, resolvedModel, "anthropic.claude-haiku-3-5"), + Model: tc.bedrockModel, + SmallFastModel: "anthropic.claude-haiku-3-5", + ResolvedModel: tc.resolvedModel, + ResolvedSmallFastModel: "anthropic.claude-haiku-3-5", + }, nil), clientHeaders: clientHeaders, logger: slog.Make(), } @@ -1322,7 +1322,7 @@ func TestBedrockMantleIsPassthrough(t *testing.T) { Region: "us-east-1", BaseURL: "https://bedrock-mantle.us-east-1.api.aws/anthropic", Protocol: config.BedrockProtocolMantle, - }, credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), "", ""), + }, credentials.NewStaticCredentialsProvider("test-key", "test-secret", "")), logger: slog.Make(), } @@ -1373,7 +1373,7 @@ func TestAWSMantleOptionsValidation(t *testing.T) { t.Parallel() base := &interceptionBase{ - bedrock: NewBedrockRuntime(tt.cfg, credentials.NewStaticCredentialsProvider("test-key", "test-secret", ""), "", ""), + bedrock: NewBedrockRuntime(tt.cfg, credentials.NewStaticCredentialsProvider("test-key", "test-secret", "")), } opts, err := base.withBedrockMantleOptions(t.Context()) if tt.errorMsg != "" { diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 68a349e3f12..d2dcf39b1d0 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -83,16 +83,7 @@ func NewAnthropic(ctx context.Context, cfg config.Anthropic, bedrockCfg *config. return nil, xerrors.Errorf("bedrock config: %w", err) } - // coderd resolves application inference profile ARNs when the provider - // is written, so construction never calls AWS. A missing resolution - // means the stored provider predates that step or was edited around it; - // serving it would silently misshape every request. - model, smallFastModel, err := resolvedBedrockModels(runtimeCfg) - if err != nil { - return nil, err - } - - bedrock = messages.NewBedrockRuntime(runtimeCfg, awsCfg.Credentials, model, smallFastModel) + bedrock = messages.NewBedrockRuntime(runtimeCfg, awsCfg.Credentials) } return &Anthropic{ diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 6b9bc88762c..034305a4536 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -127,29 +127,3 @@ func ResolveBedrockModels(ctx context.Context, cfg config.AWSBedrock) (model, sm } return model, smallFastModel, nil } - -// resolvedBedrockModels returns the model identities to serve with. A -// configured identifier that needs no resolution is its own identity; an -// application inference profile ARN requires the resolution stored with the -// provider. -func resolvedBedrockModels(cfg config.AWSBedrock) (model, smallFastModel string, err error) { - identity := func(configured, resolved, field string) (string, error) { - if !isApplicationInferenceProfileARN(configured) { - return configured, nil - } - if resolved == "" { - return "", xerrors.Errorf("%s %q is an application inference profile with no resolved model; re-save the provider to resolve it", field, configured) - } - return resolved, nil - } - - model, err = identity(cfg.Model, cfg.ResolvedModel, "model") - if err != nil { - return "", "", err - } - smallFastModel, err = identity(cfg.SmallFastModel, cfg.ResolvedSmallFastModel, "small fast model") - if err != nil { - return "", "", err - } - return model, smallFastModel, nil -} diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 1e551b8c468..67a0bdbe262 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -229,8 +229,7 @@ func TestResolveBedrockModels(t *testing.T) { } // TestNewAnthropic_ServesStoredResolution covers what the gateway does with the -// resolution coderd stored: it serves it, and refuses to serve an opaque -// profile ARN that has none. +// resolution coderd stored. func TestNewAnthropic_ServesStoredResolution(t *testing.T) { t.Parallel() @@ -261,21 +260,16 @@ func TestNewAnthropic_ServesStoredResolution(t *testing.T) { require.Equal(t, "anthropic.claude-haiku-4-5", p.bedrock.ResolvedSmallFastModel()) }) - t.Run("unresolved profile fails construction", func(t *testing.T) { + t.Run("unresolved profile serves the configured identifier", func(t *testing.T) { t.Parallel() - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(func(*config.AWSBedrock) {})) - require.ErrorContains(t, err, "no resolved model") - }) - - t.Run("unresolved small fast profile fails construction", func(t *testing.T) { - t.Parallel() - - _, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(func(cfg *config.AWSBedrock) { - cfg.Model = "eu.anthropic.claude-opus-4-8" - cfg.SmallFastModel = profileARN - })) - require.ErrorContains(t, err, "small fast model") + // A save whose profile lookup failed stores no resolution. The provider + // still serves, with the ARN as its own identity, which is wrong for + // capability detection and pricing but visible to the operator as the + // error their save returned. + p, err := NewAnthropic(context.Background(), config.Anthropic{}, bedrockCfg(func(*config.AWSBedrock) {})) + require.NoError(t, err) + require.Equal(t, profileARN, p.bedrock.ResolvedModel()) }) t.Run("plain model ids serve themselves", func(t *testing.T) { From 5a04d323ee1fdb1c58aacd1806fc9a6ee9c56c28 Mon Sep 17 00:00:00 2001 From: evgeniy-scherbina Date: Wed, 9 Sep 2026 12:07:31 -0400 Subject: [PATCH 09/25] docs: minor changes --- aibridge/provider/bedrock_inference_profile.go | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index 034305a4536..c439728567e 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -22,8 +22,8 @@ const bedrockService = "bedrock" const applicationInferenceProfileResourceType = "application-inference-profile" // inferenceProfileResolutionTimeout bounds the Bedrock control-plane calls made -// while writing a provider, which also cover the first credential resolution -// (STS/IRSA). +// while constructing a provider, which also cover the first credential +// resolution (STS/IRSA). const inferenceProfileResolutionTimeout = 30 * time.Second // isApplicationInferenceProfileARN reports whether model is an application @@ -93,10 +93,8 @@ func modelIDFromARN(modelARN string) (string, error) { // that are not application inference profile ARNs are returned unchanged and // cost no AWS call. // -// It runs where a Bedrock provider is written rather than where it is served, -// so the gateway never calls the Bedrock control plane. The identity comes from -// cfg, including any role assumed via config.AWSBedrock.RoleARN, so the -// required bedrock:GetInferenceProfile permission belongs to that identity. +// The identity comes from cfg, including any role assumed via config.AWSBedrock.RoleARN, +// so the required bedrock:GetInferenceProfile permission belongs to that identity. func ResolveBedrockModels(ctx context.Context, cfg config.AWSBedrock) (model, smallFastModel string, err error) { if !isApplicationInferenceProfileARN(cfg.Model) && !isApplicationInferenceProfileARN(cfg.SmallFastModel) { return cfg.Model, cfg.SmallFastModel, nil From 42f9cf3087941b3a580826f8b870c0c1e6108e53 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 16:24:20 +0000 Subject: [PATCH 10/25] refactor(aibridge/config): add resolved model fallback accessors --- aibridge/config/config.go | 25 ++++++++++++++++++++++++- aibridge/intercept/messages/base.go | 28 +++++----------------------- 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/aibridge/config/config.go b/aibridge/config/config.go index 44ba4ef2a4d..5de91fa8a55 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -70,7 +70,9 @@ type AWSBedrock struct { Protocol BedrockProtocol // ResolvedModel is the model ID behind Model, which differs from it only // when Model is an application inference profile ARN. coderd resolves it - // when the provider is written, so the gateway never calls AWS for it. + // when the provider is written, so the gateway never calls AWS for it. It + // is empty when nothing needed resolving, so read it through + // [AWSBedrock.ResolvedModelWithFallback]. ResolvedModel string // ResolvedSmallFastModel is ResolvedModel for SmallFastModel. ResolvedSmallFastModel string @@ -86,6 +88,27 @@ func (c AWSBedrock) ResolvedProtocol() BedrockProtocol { return c.Protocol } +// ResolvedModelWithFallback returns the model ID to record usage, pricing, and +// capabilities against. It falls back to Model, which is its own identity +// unless it is an application inference profile ARN. An unresolved ARN +// therefore serves as itself, matching the behavior before coderd resolved +// profiles; the operator saw the resolution failure when saving the provider. +func (c AWSBedrock) ResolvedModelWithFallback() string { + if c.ResolvedModel == "" { + return c.Model + } + return c.ResolvedModel +} + +// ResolvedSmallFastModelWithFallback is +// [AWSBedrock.ResolvedModelWithFallback] for the small/fast model. +func (c AWSBedrock) ResolvedSmallFastModelWithFallback() string { + if c.ResolvedSmallFastModel == "" { + return c.SmallFastModel + } + return c.ResolvedSmallFastModel +} + // Validate verifies protocol-specific Bedrock configuration. func (c AWSBedrock) Validate() error { switch c.ResolvedProtocol() { diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index df9f9ef4fa8..f25e5114833 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -72,31 +72,13 @@ var bedrockSupportedBetaFlags = map[string]bool{ type BedrockRuntime struct { Cfg aibconfig.AWSBedrock Creds aws.CredentialsProvider - - resolvedModel string - resolvedSmallFastModel string } -// NewBedrockRuntime bundles the Bedrock config and credentials with the model -// IDs behind the configured identifiers. -// -// An identifier resolves to itself unless it is an application inference -// profile ARN, which is opaque. Those are resolved through AWS when the -// provider is written and arrive on cfg, so construction never calls AWS. +// NewBedrockRuntime bundles the Bedrock config and credentials. func NewBedrockRuntime(cfg aibconfig.AWSBedrock, creds aws.CredentialsProvider) *BedrockRuntime { - resolvedModel := cfg.ResolvedModel - if resolvedModel == "" { - resolvedModel = cfg.Model - } - resolvedSmallFastModel := cfg.ResolvedSmallFastModel - if resolvedSmallFastModel == "" { - resolvedSmallFastModel = cfg.SmallFastModel - } return &BedrockRuntime{ - Cfg: cfg, - Creds: creds, - resolvedModel: resolvedModel, - resolvedSmallFastModel: resolvedSmallFastModel, + Cfg: cfg, + Creds: creds, } } @@ -117,13 +99,13 @@ func (b *BedrockRuntime) ConfiguredSmallFastModel() string { // Model capabilities, usage records, pricing, and metrics all key off this // rather than the configured identifier. func (b *BedrockRuntime) ResolvedModel() string { - return b.resolvedModel + return b.Cfg.ResolvedModelWithFallback() } // ResolvedSmallFastModel is [BedrockRuntime.ResolvedModel] for the small/fast // model. func (b *BedrockRuntime) ResolvedSmallFastModel() string { - return b.resolvedSmallFastModel + return b.Cfg.ResolvedSmallFastModelWithFallback() } type interceptionBase struct { From 051d2b50532d57302277d80af8a62ae3e46783ae Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 16:48:13 +0000 Subject: [PATCH 11/25] refactor(aibridge/config): simplify resolved model fallback accessors --- aibridge/config/config.go | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/aibridge/config/config.go b/aibridge/config/config.go index 5de91fa8a55..58bb4e1e049 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -70,9 +70,7 @@ type AWSBedrock struct { Protocol BedrockProtocol // ResolvedModel is the model ID behind Model, which differs from it only // when Model is an application inference profile ARN. coderd resolves it - // when the provider is written, so the gateway never calls AWS for it. It - // is empty when nothing needed resolving, so read it through - // [AWSBedrock.ResolvedModelWithFallback]. + // when the provider is written, so the gateway never calls AWS for it. ResolvedModel string // ResolvedSmallFastModel is ResolvedModel for SmallFastModel. ResolvedSmallFastModel string @@ -88,25 +86,20 @@ func (c AWSBedrock) ResolvedProtocol() BedrockProtocol { return c.Protocol } -// ResolvedModelWithFallback returns the model ID to record usage, pricing, and -// capabilities against. It falls back to Model, which is its own identity -// unless it is an application inference profile ARN. An unresolved ARN -// therefore serves as itself, matching the behavior before coderd resolved -// profiles; the operator saw the resolution failure when saving the provider. func (c AWSBedrock) ResolvedModelWithFallback() string { - if c.ResolvedModel == "" { - return c.Model + if c.ResolvedModel != "" { + return c.ResolvedModel } - return c.ResolvedModel + return c.Model } // ResolvedSmallFastModelWithFallback is // [AWSBedrock.ResolvedModelWithFallback] for the small/fast model. func (c AWSBedrock) ResolvedSmallFastModelWithFallback() string { - if c.ResolvedSmallFastModel == "" { - return c.SmallFastModel + if c.ResolvedSmallFastModel != "" { + return c.ResolvedSmallFastModel } - return c.ResolvedSmallFastModel + return c.SmallFastModel } // Validate verifies protocol-specific Bedrock configuration. From a56f048f7827caddf55ec20e861ae3166b182914 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 16:54:51 +0000 Subject: [PATCH 12/25] docs(aibridge/config): drop redundant accessor comment --- aibridge/config/config.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/aibridge/config/config.go b/aibridge/config/config.go index 58bb4e1e049..8ec4c9e6072 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -93,8 +93,6 @@ func (c AWSBedrock) ResolvedModelWithFallback() string { return c.Model } -// ResolvedSmallFastModelWithFallback is -// [AWSBedrock.ResolvedModelWithFallback] for the small/fast model. func (c AWSBedrock) ResolvedSmallFastModelWithFallback() string { if c.ResolvedSmallFastModel != "" { return c.ResolvedSmallFastModel From 4e555de78aa73991c10c97bc0210351ef371bde9 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 17:31:19 +0000 Subject: [PATCH 13/25] refactor(codersdk): carry bedrock model resolution on the settings type --- cli/aibridged.go | 14 ++------------ coderd/aibridge/bedrock.go | 20 +++++++++++--------- codersdk/aiproviders_bedrock.go | 10 ++++++++++ 3 files changed, 23 insertions(+), 21 deletions(-) diff --git a/cli/aibridged.go b/cli/aibridged.go index 4aee68221a1..030fbc00e65 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -216,9 +216,9 @@ func protoToProviderSpec(pp *proto.AIProvider) aiProviderSpec { bedrock.RoleARN = b.GetRoleArn() bedrock.ExternalID = b.GetExternalId() bedrock.Protocol = codersdk.AIProviderBedrockProtocol(b.GetProtocol()) + bedrock.ResolvedModel = b.GetResolvedModel() + bedrock.ResolvedSmallFastModel = b.GetResolvedSmallFastModel() spec.Bedrock = new(bedrock) - spec.BedrockResolvedModel = b.GetResolvedModel() - spec.BedrockResolvedSmallFastModel = b.GetResolvedSmallFastModel() } return spec } @@ -237,12 +237,6 @@ type aiProviderSpec struct { // Bedrock holds Bedrock-specific settings when the provider targets // AWS Bedrock; nil otherwise. Bedrock *codersdk.AIProviderBedrockSettings - // BedrockResolvedModel and BedrockResolvedSmallFastModel are the models the - // configured identifiers refer to. They are set only when an identifier is - // an application inference profile ARN, which coderd resolved when the - // provider was written. - BedrockResolvedModel string - BedrockResolvedSmallFastModel string } // buildProvider constructs the appropriate [aibridge.Provider] for a @@ -291,10 +285,6 @@ func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBrid case database.AIProviderTypeAnthropic, database.AIProviderTypeBedrock: bedrock := bedrockConfig(spec.BaseURL, spec.Bedrock) - if bedrock != nil { - bedrock.ResolvedModel = spec.BedrockResolvedModel - bedrock.ResolvedSmallFastModel = spec.BedrockResolvedSmallFastModel - } // A spec typed 'bedrock' authenticates exclusively via settings; // without populated Bedrock credentials it cannot make upstream // calls, so refuse rather than falling back to an unsigned diff --git a/coderd/aibridge/bedrock.go b/coderd/aibridge/bedrock.go index 16fa9f2f57b..49182d484be 100644 --- a/coderd/aibridge/bedrock.go +++ b/coderd/aibridge/bedrock.go @@ -25,14 +25,16 @@ func BedrockConfig(baseURL string, bedrock *codersdk.AIProviderBedrockSettings) return nil } return &aibridgeconfig.AWSBedrock{ - BaseURL: baseURL, - Region: settings.Region, - AccessKey: ptr.NilToEmpty(settings.AccessKey), - AccessKeySecret: ptr.NilToEmpty(settings.AccessKeySecret), - Model: settings.Model, - SmallFastModel: settings.SmallFastModel, - RoleARN: settings.RoleARN, - ExternalID: settings.ExternalID, - Protocol: aibridgeconfig.BedrockProtocol(settings.ResolvedProtocol()), + BaseURL: baseURL, + Region: settings.Region, + AccessKey: ptr.NilToEmpty(settings.AccessKey), + AccessKeySecret: ptr.NilToEmpty(settings.AccessKeySecret), + Model: settings.Model, + SmallFastModel: settings.SmallFastModel, + RoleARN: settings.RoleARN, + ExternalID: settings.ExternalID, + Protocol: aibridgeconfig.BedrockProtocol(settings.ResolvedProtocol()), + ResolvedModel: settings.ResolvedModel, + ResolvedSmallFastModel: settings.ResolvedSmallFastModel, } } diff --git a/codersdk/aiproviders_bedrock.go b/codersdk/aiproviders_bedrock.go index b3bc94e9e4b..2a25093db7c 100644 --- a/codersdk/aiproviders_bedrock.go +++ b/codersdk/aiproviders_bedrock.go @@ -61,6 +61,16 @@ type AIProviderBedrockSettings struct { // AIProviderBedrockProtocolInvokeModel, so existing rows keep the legacy // behavior. Protocol AIProviderBedrockProtocol `json:"protocol,omitempty"` + // ResolvedModel and ResolvedSmallFastModel carry the model IDs behind the + // configured identifiers, which differ from them only for application + // inference profile ARNs. coderd resolves those when the provider is + // written and stores them in ai_provider_bedrock_resolved_models. + // + // They are in-process plumbing for the gateway, not part of this type's + // wire form: the API neither accepts nor returns them, and they are never + // stored in the settings blob. + ResolvedModel string `json:"-"` + ResolvedSmallFastModel string `json:"-"` } // ResolvedProtocol returns the configured protocol, mapping the empty value to From b6c5568f3d30f23021eb3686a9ad08cebf5efd25 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 18:59:30 +0000 Subject: [PATCH 14/25] refactor(coderd): key bedrock model resolution by inference profile arn --- .../provider/bedrock_inference_profile.go | 46 +++++----- ...bedrock_inference_profile_internal_test.go | 20 ++--- coderd/ai_providers.go | 17 +--- coderd/ai_providers_bedrock.go | 44 ++++------ coderd/ai_providers_bedrock_test.go | 83 +++++++++---------- coderd/aibridgedserver/aibridgedserver.go | 55 +++++++----- coderd/database/dbauthz/dbauthz.go | 33 +++----- coderd/database/dbauthz/dbauthz_test.go | 29 +++---- coderd/database/dbmetrics/querymetrics.go | 40 ++++----- coderd/database/dbmock/dbmock.go | 68 ++++++--------- coderd/database/dump.sql | 26 +++--- coderd/database/foreign_key_constraint.go | 1 - ..._bedrock_inference_profile_models.down.sql | 1 + ...ai_bedrock_inference_profile_models.up.sql | 14 ++++ ..._provider_bedrock_resolved_models.down.sql | 1 - ...ai_provider_bedrock_resolved_models.up.sql | 17 ---- ...ai_bedrock_inference_profile_models.up.sql | 8 ++ ...ai_provider_bedrock_resolved_models.up.sql | 31 ------- coderd/database/models.go | 14 ++-- coderd/database/querier.go | 12 +-- coderd/database/queries.sql.go | 57 +++++-------- coderd/database/queries/ai_providers.sql | 28 +++---- coderd/database/unique_constraint.go | 2 +- 23 files changed, 270 insertions(+), 377 deletions(-) create mode 100644 coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql create mode 100644 coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql delete mode 100644 coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql delete mode 100644 coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql delete mode 100644 coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql diff --git a/aibridge/provider/bedrock_inference_profile.go b/aibridge/provider/bedrock_inference_profile.go index c439728567e..694a93fc714 100644 --- a/aibridge/provider/bedrock_inference_profile.go +++ b/aibridge/provider/bedrock_inference_profile.go @@ -88,40 +88,42 @@ func modelIDFromARN(modelARN string) (string, error) { return model, nil } -// ResolveBedrockModels resolves the configured model identifiers to the model -// IDs used for capability detection, usage recording, and pricing. Identifiers -// that are not application inference profile ARNs are returned unchanged and -// cost no AWS call. +// ResolveBedrockModels resolves the application inference profile ARNs among +// the configured model identifiers, returning what each ARN refers to. The +// result is empty when neither identifier is an ARN, which costs no AWS call. // // The identity comes from cfg, including any role assumed via config.AWSBedrock.RoleARN, // so the required bedrock:GetInferenceProfile permission belongs to that identity. -func ResolveBedrockModels(ctx context.Context, cfg config.AWSBedrock) (model, smallFastModel string, err error) { - if !isApplicationInferenceProfileARN(cfg.Model) && !isApplicationInferenceProfileARN(cfg.SmallFastModel) { - return cfg.Model, cfg.SmallFastModel, nil +func ResolveBedrockModels(ctx context.Context, cfg config.AWSBedrock) (map[string]string, error) { + resolved := make(map[string]string, 2) + + var profiles []string + for _, configured := range []string{cfg.Model, cfg.SmallFastModel} { + if isApplicationInferenceProfileARN(configured) { + profiles = append(profiles, configured) + } + } + if len(profiles) == 0 { + return resolved, nil } awsCfg, err := buildBedrockCredentials(ctx, cfg) if err != nil { - return "", "", xerrors.Errorf("build bedrock credentials: %w", err) + return nil, xerrors.Errorf("build bedrock credentials: %w", err) } resolveCtx, cancel := context.WithTimeout(ctx, inferenceProfileResolutionTimeout) defer cancel() - resolveOne := func(configured string) (string, error) { - if !isApplicationInferenceProfileARN(configured) { - return configured, nil + for _, profileARN := range profiles { + if _, ok := resolved[profileARN]; ok { + continue } - return resolveInferenceProfile(resolveCtx, awsCfg, configured) - } - - model, err = resolveOne(cfg.Model) - if err != nil { - return "", "", xerrors.Errorf("resolve model: %w", err) - } - smallFastModel, err = resolveOne(cfg.SmallFastModel) - if err != nil { - return "", "", xerrors.Errorf("resolve small fast model: %w", err) + model, err := resolveInferenceProfile(resolveCtx, awsCfg, profileARN) + if err != nil { + return nil, err + } + resolved[profileARN] = model } - return model, smallFastModel, nil + return resolved, nil } diff --git a/aibridge/provider/bedrock_inference_profile_internal_test.go b/aibridge/provider/bedrock_inference_profile_internal_test.go index 67a0bdbe262..91481d7fbc0 100644 --- a/aibridge/provider/bedrock_inference_profile_internal_test.go +++ b/aibridge/provider/bedrock_inference_profile_internal_test.go @@ -166,10 +166,9 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + resolved, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + require.Equal(t, map[string]string{profileARN: "anthropic.claude-opus-4-8"}, resolved) require.Len(t, *paths, 1, "only the profile ARN is resolved") require.Contains(t, (*paths)[0], profileARN) }) @@ -183,8 +182,7 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) - require.ErrorContains(t, err, "resolve model") + _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.ErrorContains(t, err, "GetInferenceProfile") }) @@ -195,7 +193,7 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - _, _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) + _, err := ResolveBedrockModels(context.Background(), bedrockCfg(profileARN, "anthropic.claude-haiku-4-5")) require.ErrorContains(t, err, "references no model") }) @@ -206,10 +204,9 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) + resolved, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", smallFastProfileARN)) require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + require.Equal(t, map[string]string{smallFastProfileARN: "anthropic.claude-haiku-4-5"}, resolved) require.Len(t, *paths, 1, "only the small fast profile ARN is resolved") require.Contains(t, (*paths)[0], smallFastProfileARN) }) @@ -220,10 +217,9 @@ func TestResolveBedrockModels(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - model, smallFastModel, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) + resolved, err := ResolveBedrockModels(context.Background(), bedrockCfg("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5")) require.NoError(t, err) - require.Equal(t, "eu.anthropic.claude-opus-4-8", model) - require.Equal(t, "anthropic.claude-haiku-4-5", smallFastModel) + require.Empty(t, resolved) require.Empty(t, *paths) }) } diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index ff7ca6e6231..70cd3c30670 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -245,8 +245,7 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { aReq.New = row // Resolve inference profile ARNs once the provider is stored, then announce - // it. The gateway never sees the unresolved provider, and never calls the - // Bedrock control plane itself. + // it. The gateway never calls the Bedrock control plane itself. if err := api.resolveBedrockModels(ctx, row); err != nil { api.writeAIProviderResolutionError(ctx, rw, err) return @@ -368,15 +367,6 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { return errCopilotRejectsAPIKeys } - // The patch may point the provider at different identifiers, so the - // stored resolution no longer describes it. Resolution runs after the - // transaction, because it is an AWS call. - if req.Settings != nil { - if err := clearBedrockModelResolution(ctx, tx, old.ID); err != nil { - return err - } - } - displayName := old.DisplayName if req.DisplayName != nil { // Empty string clears the column. @@ -456,7 +446,7 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { } // An update that carries no settings cannot change the configured - // identifiers or the credentials they resolve under, so the stored + // identifiers or the credentials they resolve under, so any stored // resolution still holds. if req.Settings != nil { if err := api.resolveBedrockModels(ctx, updated); err != nil { @@ -514,8 +504,7 @@ func (api *API) aiProvidersDelete(rw http.ResponseWriter, r *http.Request) { if err := tx.DeleteAIProviderByID(ctx, row.ID); err != nil { return xerrors.Errorf("delete ai provider: %w", err) } - // Providers are soft-deleted, so the foreign key never cascades. - return clearBedrockModelResolution(ctx, tx, row.ID) + return nil }, &database.TxOptions{TxIdentifier: "delete_ai_provider"}) if err != nil { writeAIProviderError(ctx, api.Logger, rw, err, "delete AI provider", "Internal error deleting AI provider.") diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index 75bbbe26640..4645a8b8ea1 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -4,7 +4,6 @@ import ( "context" "net/http" - "github.com/google/uuid" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -27,8 +26,8 @@ func (e bedrockProfileUnresolvableError) Error() string { func (e bedrockProfileUnresolvableError) Unwrap() error { return e.err } -// resolveBedrockModels records which models the provider's application -// inference profile ARNs refer to. An ARN identifies a billing wrapper rather +// resolveBedrockModels records which model each of the provider's application +// inference profile ARNs refers to. An ARN identifies a billing wrapper rather // than a model, so the gateway needs the mapping to detect capabilities, price // usage, and record interceptions. // @@ -37,44 +36,33 @@ func (e bedrockProfileUnresolvableError) Unwrap() error { return e.err } // the stored row because that is the merged configuration the provider will // actually use. // -// Nothing is stored for a provider configured with plain model IDs: they are -// already model identities and need no mapping. +// It runs on every save, even for an ARN another provider already resolved, so +// that saving proves this provider's own identity can read the profile. func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvider) error { settings, err := db2sdk.AIProviderSettings(row.Settings) if err != nil { return xerrors.Errorf("decode settings: %w", err) } + // Resolution is a Bedrock control-plane call, whereas the provider's + // BaseURL is its runtime endpoint, so it is deliberately not carried here. cfg := agplaibridge.BedrockConfig("", settings.Bedrock) if cfg == nil { return nil } - model, smallFastModel, err := provider.ResolveBedrockModels(ctx, *cfg) + resolved, err := provider.ResolveBedrockModels(ctx, *cfg) if err != nil { return bedrockProfileUnresolvableError{err: err} } - if model == cfg.Model && smallFastModel == cfg.SmallFastModel { - return nil - } - - err = api.Database.UpsertAIProviderBedrockResolvedModels(ctx, database.UpsertAIProviderBedrockResolvedModelsParams{ - AIProviderID: row.ID, - ResolvedModel: model, - ResolvedSmallFastModel: smallFastModel, - }) - if err != nil { - return xerrors.Errorf("store resolved models: %w", err) - } - return nil -} -// clearBedrockModelResolution drops a provider's stored resolution. The write -// path calls it whenever the configured identifiers may have changed, so a -// stale mapping never outlives the ARN it describes. The provider is then -// unresolved until resolution succeeds, and the gateway will not serve it. -func clearBedrockModelResolution(ctx context.Context, db database.Store, providerID uuid.UUID) error { - if err := db.DeleteAIProviderBedrockResolvedModels(ctx, providerID); err != nil { - return xerrors.Errorf("clear resolved models: %w", err) + for profileARN, model := range resolved { + err := api.Database.UpsertAIBedrockInferenceProfileModel(ctx, database.UpsertAIBedrockInferenceProfileModelParams{ + InferenceProfileArn: profileARN, + ResolvedModel: model, + }) + if err != nil { + return xerrors.Errorf("store resolved model for %q: %w", profileARN, err) + } } return nil } @@ -82,7 +70,7 @@ func clearBedrockModelResolution(ctx context.Context, db database.Store, provide // writeAIProviderResolutionError reports a failed Bedrock model resolution. The // provider keeps the identifiers the operator asked for, but without a // resolution the gateway cannot tell what an opaque profile ARN refers to, so -// it refuses to serve the provider until a later save resolves it. +// it serves the ARN as its own identity until a later save resolves it. func (api *API) writeAIProviderResolutionError(ctx context.Context, rw http.ResponseWriter, err error) { var unresolvable bedrockProfileUnresolvableError if !xerrors.As(err, &unresolvable) { diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go index fd5830799a6..bbd416511eb 100644 --- a/coderd/ai_providers_bedrock_test.go +++ b/coderd/ai_providers_bedrock_test.go @@ -9,7 +9,6 @@ import ( "sync" "testing" - "github.com/google/uuid" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/coderdtest" @@ -68,12 +67,18 @@ func respondWithModel(modelARN string) http.HandlerFunc { } } -func resolvedModels(ctx context.Context, t *testing.T, db database.Store, providerID uuid.UUID) []database.AIProviderBedrockResolvedModel { +// resolvedModel returns the model stored for an inference profile ARN, or the +// empty string when the ARN has no mapping. +func resolvedModel(ctx context.Context, t *testing.T, db database.Store, profileARN string) string { t.Helper() - rows, err := db.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, []uuid.UUID{providerID}) + rows, err := db.GetAIBedrockInferenceProfileModels(ctx, []string{profileARN}) require.NoError(t, err) - return rows + if len(rows) == 0 { + return "" + } + require.Len(t, rows, 1) + return rows[0].ResolvedModel } // TestAIProvidersBedrockProfileResolution drives provider writes against a mock @@ -110,11 +115,8 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Equal(t, testProfileARN, created.Settings.Bedrock.Model) require.Equal(t, testSmallFastProfileARN, created.Settings.Bedrock.SmallFastModel) require.Len(t, paths(), 2, "each profile is resolved once") - - rows := resolvedModels(ctx, t, db, created.ID) - require.Len(t, rows, 1) - require.Equal(t, "anthropic.claude-opus-4-8", rows[0].ResolvedModel) - require.Equal(t, "anthropic.claude-haiku-4-5", rows[0].ResolvedSmallFastModel) + require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) + require.Equal(t, "anthropic.claude-haiku-4-5", resolvedModel(ctx, t, db, testSmallFastProfileARN)) }) t.Run("CreateLeavesPlainModelIDsUnresolved", func(t *testing.T) { @@ -137,8 +139,8 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Empty(t, paths()) - require.Empty(t, resolvedModels(ctx, t, db, created.ID), "plain model ids are already model identities") + require.NotNil(t, created.Settings.Bedrock) + require.Empty(t, paths(), "plain model ids are already model identities") }) t.Run("CreateRejectsUnresolvableProfile", func(t *testing.T) { @@ -169,12 +171,12 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Contains(t, sdkErr.Detail, "GetInferenceProfile") // The provider is stored with the ARN the operator asked for, but - // without a resolution the gateway refuses to serve it. + // nothing maps that ARN, so the gateway serves it as its own identity. //nolint:gocritic // Owner role is the audience for this endpoint. providers, err := client.AIProviders(ctx) require.NoError(t, err) require.Len(t, providers, 1) - require.Empty(t, resolvedModels(ctx, t, db, providers[0].ID)) + require.Empty(t, resolvedModel(ctx, t, db, testProfileARN)) }) t.Run("UpdateReresolvesChangedProfile", func(t *testing.T) { @@ -195,7 +197,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Empty(t, resolvedModels(ctx, t, db, created.ID)) + require.Empty(t, resolvedModel(ctx, t, db, testProfileARN)) //nolint:gocritic // Owner role is the audience for this endpoint. updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ @@ -203,15 +205,10 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) require.NoError(t, err) require.Equal(t, testProfileARN, updated.Settings.Bedrock.Model) - - rows := resolvedModels(ctx, t, db, created.ID) - require.Len(t, rows, 1) - require.Equal(t, "anthropic.claude-opus-4-8", rows[0].ResolvedModel) - // The small/fast model is a plain ID, so it resolves to itself. - require.Equal(t, "anthropic.claude-haiku-4-5", rows[0].ResolvedSmallFastModel) + require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) }) - t.Run("UpdateClearsResolutionWhenProfileReplacedByModelID", func(t *testing.T) { + t.Run("UpdateToPlainModelIDNeedsNoResolution", func(t *testing.T) { url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) @@ -229,7 +226,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Len(t, resolvedModels(ctx, t, db, created.ID), 1) + require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) callsAfterCreate := len(paths()) //nolint:gocritic // Owner role is the audience for this endpoint. @@ -237,7 +234,6 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Empty(t, resolvedModels(ctx, t, db, created.ID), "a plain model id needs no mapping") require.Len(t, paths(), callsAfterCreate, "no profile is left to resolve") }) @@ -267,15 +263,12 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Enabled: &enabled, }) require.NoError(t, err) - - rows := resolvedModels(ctx, t, db, created.ID) - require.Len(t, rows, 1) - require.Equal(t, "anthropic.claude-opus-4-8", rows[0].ResolvedModel) + require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) require.Len(t, paths(), callsAfterCreate, "an unrelated update does not call AWS") }) - t.Run("DeleteClearsResolution", func(t *testing.T) { - url, _ := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) + t.Run("SavingResolvesEvenWhenTheARNIsAlreadyMapped", func(t *testing.T) { + url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) db, ps := dbtestutil.NewDB(t) @@ -283,21 +276,21 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) - //nolint:gocritic // Owner role is the audience for this endpoint. - created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ - Name: "bedrock-delete", - Type: codersdk.AIProviderTypeBedrock, - BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", - Enabled: true, - Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), - }) - require.NoError(t, err) - require.Len(t, resolvedModels(ctx, t, db, created.ID), 1) - - // Providers are soft-deleted, so the mapping has to be removed - // explicitly rather than by the foreign key. - //nolint:gocritic // Owner role is the audience for this endpoint. - require.NoError(t, client.DeleteAIProvider(ctx, created.ID.String())) - require.Empty(t, resolvedModels(ctx, t, db, created.ID)) + for _, name := range []string{"bedrock-first", "bedrock-second"} { + //nolint:gocritic // Owner role is the audience for this endpoint. + _, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: name, + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + } + + // The mapping is shared, but each save proves that provider's own + // identity can read the profile. + require.Len(t, paths(), 2) + require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) }) } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index db544ce69cd..4f2cdc6a41b 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -100,7 +100,7 @@ type store interface { // any in-flight env seed holding LockIDAIProvidersEnvSeed. GetAIProviders(ctx context.Context, arg database.GetAIProvidersParams) ([]database.AIProvider, error) GetAIProviderKeysByProviderIDs(ctx context.Context, providerIDs []uuid.UUID) ([]database.AIProviderKey, error) - GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIDs []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) + GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) InTx(func(database.Store) error, *database.TxOptions) error } @@ -956,9 +956,9 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ ctx = dbauthz.AsAIBridged(ctx) var ( - rows []database.AIProvider - keysByProvider map[uuid.UUID][]database.AIProviderKey - resolvedByProvider map[uuid.UUID]database.AIProviderBedrockResolvedModel + rows []database.AIProvider + keysByProvider map[uuid.UUID][]database.AIProviderKey + modelByProfile map[string]string ) // Wrap both reads in a read-only transaction so the provider list and the // key list are consistent with each other, and so the seed lock is held @@ -998,16 +998,15 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ } // Bedrock application inference profile ARNs are opaque, so the models - // they refer to are resolved when the provider is written and read back - // here. A provider without a mapping either configures plain model IDs - // or could not be resolved. - resolvedRows, err := tx.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, ids) + // they refer to are resolved when a provider is written and read back + // here. An ARN without a mapping is served as its own identity. + profileRows, err := tx.GetAIBedrockInferenceProfileModels(ctx, bedrockModelIdentifiers(rows)) if err != nil { - return xerrors.Errorf("get ai provider resolved models: %w", err) + return xerrors.Errorf("get bedrock inference profile models: %w", err) } - resolvedByProvider = make(map[uuid.UUID]database.AIProviderBedrockResolvedModel, len(resolvedRows)) - for _, r := range resolvedRows { - resolvedByProvider[r.AIProviderID] = r + modelByProfile = make(map[string]string, len(profileRows)) + for _, r := range profileRows { + modelByProfile[r.InferenceProfileArn] = r.ResolvedModel } return nil }, &database.TxOptions{ReadOnly: true, TxIdentifier: "get_ai_providers"}) @@ -1017,7 +1016,7 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ providers := make([]*proto.AIProvider, 0, len(rows)) for _, row := range rows { - p, err := aiProviderToProto(row, keysByProvider[row.ID], resolvedByProvider[row.ID]) + p, err := aiProviderToProto(row, keysByProvider[row.ID], modelByProfile) if err != nil { // Skip the offending row rather than failing the whole fetch: // one row with a corrupt settings blob must not break provider @@ -1191,15 +1190,33 @@ func parseOptionalInt32(n *int32) sql.NullInt32 { return sql.NullInt32{Int32: *n, Valid: true} } +// bedrockModelIdentifiers returns the configured Bedrock model identifiers of +// every enabled provider, which is the key set for the inference profile +// mapping. Rows whose settings cannot be decoded are skipped; aiProviderToProto +// reports that failure when it builds the payload. +func bedrockModelIdentifiers(rows []database.AIProvider) []string { + var identifiers []string + for _, row := range rows { + if !row.Enabled { + continue + } + settings, err := db2sdk.AIProviderSettings(row.Settings) + if err != nil || settings.Bedrock == nil { + continue + } + identifiers = append(identifiers, settings.Bedrock.Model, settings.Bedrock.SmallFastModel) + } + return identifiers +} + // aiProviderToProto maps a single ai_providers row (and its keys, for enabled // providers) to the proto representation served to AI Gateway daemons. Keys and // Bedrock settings are only attached for enabled providers; disabled providers // never call upstream so their secrets are withheld. // -// resolved carries the models the provider's application inference profile ARNs -// refer to, and is the zero value when the provider configures plain model IDs -// or when its profiles have not been resolved. -func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey, resolved database.AIProviderBedrockResolvedModel) (*proto.AIProvider, error) { +// modelByProfile maps an application inference profile ARN to the model it +// wraps. An identifier absent from it needs no resolution, or has none yet. +func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey, modelByProfile map[string]string) (*proto.AIProvider, error) { p := &proto.AIProvider{ Name: row.Name, Type: string(row.Type), @@ -1232,8 +1249,8 @@ func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey, r RoleArn: settings.Bedrock.RoleARN, ExternalId: settings.Bedrock.ExternalID, Protocol: string(settings.Bedrock.Protocol), - ResolvedModel: resolved.ResolvedModel, - ResolvedSmallFastModel: resolved.ResolvedSmallFastModel, + ResolvedModel: modelByProfile[settings.Bedrock.Model], + ResolvedSmallFastModel: modelByProfile[settings.Bedrock.SmallFastModel], } } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index e435a9064e0..7dbdfdb378f 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2118,13 +2118,6 @@ func (q *querier) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) (databas return q.db.DeleteAIGatewayKey(ctx, id) } -func (q *querier) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { - if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil { - return err - } - return q.db.DeleteAIProviderBedrockResolvedModels(ctx, aiProviderID) -} - func (q *querier) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAIProvider); err != nil { return err @@ -2877,6 +2870,13 @@ func (q *querier) FindMatchingPresetID(ctx context.Context, arg database.FindMat return q.db.FindMatchingPresetID(ctx, arg) } +func (q *querier) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIProvider); err != nil { + return nil, err + } + return q.db.GetAIBedrockInferenceProfileModels(ctx, inferenceProfileArns) +} + func (q *querier) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) { // The aggregate covers one chat tree, so it is authorized through the // root chat. Members cannot read interception rows back, but they can @@ -2960,13 +2960,6 @@ func (q *querier) GetAIModelPrices(ctx context.Context, arg database.GetAIModelP return q.db.GetAIModelPrices(ctx, arg) } -func (q *querier) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIDs []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) { - if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIProvider); err != nil { - return nil, err - } - return q.db.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, aiProviderIDs) -} - func (q *querier) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIProvider); err != nil { return database.AIProvider{}, err @@ -9068,18 +9061,18 @@ func (q *querier) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg datab return q.db.UpdateWorkspacesTTLByTemplateID(ctx, arg) } -func (q *querier) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { - if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAiModelPrice); err != nil { +func (q *querier) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg database.UpsertAIBedrockInferenceProfileModelParams) error { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil { return err } - return q.db.UpsertAIModelPrices(ctx, arg) + return q.db.UpsertAIBedrockInferenceProfileModel(ctx, arg) } -func (q *querier) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg database.UpsertAIProviderBedrockResolvedModelsParams) error { - if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil { +func (q *querier) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAiModelPrice); err != nil { return err } - return q.db.UpsertAIProviderBedrockResolvedModels(ctx, arg) + return q.db.UpsertAIModelPrices(ctx, arg) } func (q *querier) UpsertAISeatState(ctx context.Context, arg database.UpsertAISeatStateParams) (bool, error) { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 1ae8ea99916..9f32ae4fe00 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -7389,28 +7389,19 @@ func (s *MethodTestSuite) TestAIBridge() { dbm.EXPECT().UpdateEncryptedAIProviderSettings(gomock.Any(), arg).Return(provider, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns(provider) })) - s.Run("UpsertAIProviderBedrockResolvedModels", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - provider := testutil.Fake(s.T(), faker, database.AIProvider{}) - arg := database.UpsertAIProviderBedrockResolvedModelsParams{ - AIProviderID: provider.ID, - ResolvedModel: "anthropic.claude-opus-4-8", - ResolvedSmallFastModel: "anthropic.claude-haiku-4-5", + s.Run("UpsertAIBedrockInferenceProfileModel", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.UpsertAIBedrockInferenceProfileModelParams{ + InferenceProfileArn: "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5", + ResolvedModel: "anthropic.claude-opus-4-8", } - dbm.EXPECT().UpsertAIProviderBedrockResolvedModels(gomock.Any(), arg).Return(nil).AnyTimes() + dbm.EXPECT().UpsertAIBedrockInferenceProfileModel(gomock.Any(), arg).Return(nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns() })) - s.Run("DeleteAIProviderBedrockResolvedModels", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - provider := testutil.Fake(s.T(), faker, database.AIProvider{}) - dbm.EXPECT().DeleteAIProviderBedrockResolvedModels(gomock.Any(), provider.ID).Return(nil).AnyTimes() - check.Args(provider.ID).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns() - })) - s.Run("GetAIProviderBedrockResolvedModelsByProviderIDs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - providerA := testutil.Fake(s.T(), faker, database.AIProvider{}) - providerB := testutil.Fake(s.T(), faker, database.AIProvider{}) - providerIDs := []uuid.UUID{providerA.ID, providerB.ID} - resolved := testutil.Fake(s.T(), faker, database.AIProviderBedrockResolvedModel{AIProviderID: providerA.ID}) - dbm.EXPECT().GetAIProviderBedrockResolvedModelsByProviderIDs(gomock.Any(), providerIDs).Return([]database.AIProviderBedrockResolvedModel{resolved}, nil).AnyTimes() - check.Args(providerIDs).Asserts(rbac.ResourceAIProvider, policy.ActionRead).Returns([]database.AIProviderBedrockResolvedModel{resolved}) + s.Run("GetAIBedrockInferenceProfileModels", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + resolved := testutil.Fake(s.T(), faker, database.AIBedrockInferenceProfileModel{}) + arg := []string{resolved.InferenceProfileArn} + dbm.EXPECT().GetAIBedrockInferenceProfileModels(gomock.Any(), arg).Return([]database.AIBedrockInferenceProfileModel{resolved}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionRead).Returns([]database.AIBedrockInferenceProfileModel{resolved}) })) s.Run("GetAIProviderKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { key := testutil.Fake(s.T(), faker, database.AIProviderKey{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 356e8e688d5..243ea4cfeca 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -424,14 +424,6 @@ func (m queryMetricsStore) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) return r0, r1 } -func (m queryMetricsStore) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { - start := time.Now() - r0 := m.s.DeleteAIProviderBedrockResolvedModels(ctx, aiProviderID) - m.queryLatencies.WithLabelValues("DeleteAIProviderBedrockResolvedModels").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteAIProviderBedrockResolvedModels").Inc() - return r0 -} - func (m queryMetricsStore) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error { start := time.Now() r0 := m.s.DeleteAIProviderByID(ctx, id) @@ -1104,6 +1096,14 @@ func (m queryMetricsStore) FindMatchingPresetID(ctx context.Context, arg databas return r0, r1 } +func (m queryMetricsStore) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) { + start := time.Now() + r0, r1 := m.s.GetAIBedrockInferenceProfileModels(ctx, inferenceProfileArns) + m.queryLatencies.WithLabelValues("GetAIBedrockInferenceProfileModels").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIBedrockInferenceProfileModels").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) { start := time.Now() r0, r1 := m.s.GetAIBridgeChatCost(ctx, rootChatID) @@ -1192,14 +1192,6 @@ func (m queryMetricsStore) GetAIModelPrices(ctx context.Context, arg database.Ge return r0, r1 } -func (m queryMetricsStore) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) { - start := time.Now() - r0, r1 := m.s.GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, aiProviderIds) - m.queryLatencies.WithLabelValues("GetAIProviderBedrockResolvedModelsByProviderIDs").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIProviderBedrockResolvedModelsByProviderIDs").Inc() - return r0, r1 -} - func (m queryMetricsStore) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { start := time.Now() r0, r1 := m.s.GetAIProviderByID(ctx, id) @@ -6328,19 +6320,19 @@ func (m queryMetricsStore) UpdateWorkspacesTTLByTemplateID(ctx context.Context, return r0 } -func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { +func (m queryMetricsStore) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg database.UpsertAIBedrockInferenceProfileModelParams) error { start := time.Now() - r0 := m.s.UpsertAIModelPrices(ctx, arg) - m.queryLatencies.WithLabelValues("UpsertAIModelPrices").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIModelPrices").Inc() + r0 := m.s.UpsertAIBedrockInferenceProfileModel(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertAIBedrockInferenceProfileModel").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIBedrockInferenceProfileModel").Inc() return r0 } -func (m queryMetricsStore) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg database.UpsertAIProviderBedrockResolvedModelsParams) error { +func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { start := time.Now() - r0 := m.s.UpsertAIProviderBedrockResolvedModels(ctx, arg) - m.queryLatencies.WithLabelValues("UpsertAIProviderBedrockResolvedModels").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIProviderBedrockResolvedModels").Inc() + r0 := m.s.UpsertAIModelPrices(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertAIModelPrices").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIModelPrices").Inc() return r0 } diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 34e3eb6d98f..5e9c7e19cd2 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -676,20 +676,6 @@ func (mr *MockStoreMockRecorder) DeleteAIGatewayKey(ctx, id any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAIGatewayKey", reflect.TypeOf((*MockStore)(nil).DeleteAIGatewayKey), ctx, id) } -// DeleteAIProviderBedrockResolvedModels mocks base method. -func (m *MockStore) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "DeleteAIProviderBedrockResolvedModels", ctx, aiProviderID) - ret0, _ := ret[0].(error) - return ret0 -} - -// DeleteAIProviderBedrockResolvedModels indicates an expected call of DeleteAIProviderBedrockResolvedModels. -func (mr *MockStoreMockRecorder) DeleteAIProviderBedrockResolvedModels(ctx, aiProviderID any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAIProviderBedrockResolvedModels", reflect.TypeOf((*MockStore)(nil).DeleteAIProviderBedrockResolvedModels), ctx, aiProviderID) -} - // DeleteAIProviderByID mocks base method. func (m *MockStore) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error { m.ctrl.T.Helper() @@ -1917,6 +1903,21 @@ func (mr *MockStoreMockRecorder) FindMatchingPresetID(ctx, arg any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindMatchingPresetID", reflect.TypeOf((*MockStore)(nil).FindMatchingPresetID), ctx, arg) } +// GetAIBedrockInferenceProfileModels mocks base method. +func (m *MockStore) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAIBedrockInferenceProfileModels", ctx, inferenceProfileArns) + ret0, _ := ret[0].([]database.AIBedrockInferenceProfileModel) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAIBedrockInferenceProfileModels indicates an expected call of GetAIBedrockInferenceProfileModels. +func (mr *MockStoreMockRecorder) GetAIBedrockInferenceProfileModels(ctx, inferenceProfileArns any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBedrockInferenceProfileModels", reflect.TypeOf((*MockStore)(nil).GetAIBedrockInferenceProfileModels), ctx, inferenceProfileArns) +} + // GetAIBridgeChatCost mocks base method. func (m *MockStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) { m.ctrl.T.Helper() @@ -2082,21 +2083,6 @@ func (mr *MockStoreMockRecorder) GetAIModelPrices(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIModelPrices", reflect.TypeOf((*MockStore)(nil).GetAIModelPrices), ctx, arg) } -// GetAIProviderBedrockResolvedModelsByProviderIDs mocks base method. -func (m *MockStore) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]database.AIProviderBedrockResolvedModel, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetAIProviderBedrockResolvedModelsByProviderIDs", ctx, aiProviderIds) - ret0, _ := ret[0].([]database.AIProviderBedrockResolvedModel) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetAIProviderBedrockResolvedModelsByProviderIDs indicates an expected call of GetAIProviderBedrockResolvedModelsByProviderIDs. -func (mr *MockStoreMockRecorder) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx, aiProviderIds any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIProviderBedrockResolvedModelsByProviderIDs", reflect.TypeOf((*MockStore)(nil).GetAIProviderBedrockResolvedModelsByProviderIDs), ctx, aiProviderIds) -} - // GetAIProviderByID mocks base method. func (m *MockStore) GetAIProviderByID(ctx context.Context, id uuid.UUID) (database.AIProvider, error) { m.ctrl.T.Helper() @@ -11923,32 +11909,32 @@ func (mr *MockStoreMockRecorder) UpdateWorkspacesTTLByTemplateID(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspacesTTLByTemplateID", reflect.TypeOf((*MockStore)(nil).UpdateWorkspacesTTLByTemplateID), ctx, arg) } -// UpsertAIModelPrices mocks base method. -func (m *MockStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { +// UpsertAIBedrockInferenceProfileModel mocks base method. +func (m *MockStore) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg database.UpsertAIBedrockInferenceProfileModelParams) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpsertAIModelPrices", ctx, arg) + ret := m.ctrl.Call(m, "UpsertAIBedrockInferenceProfileModel", ctx, arg) ret0, _ := ret[0].(error) return ret0 } -// UpsertAIModelPrices indicates an expected call of UpsertAIModelPrices. -func (mr *MockStoreMockRecorder) UpsertAIModelPrices(ctx, arg any) *gomock.Call { +// UpsertAIBedrockInferenceProfileModel indicates an expected call of UpsertAIBedrockInferenceProfileModel. +func (mr *MockStoreMockRecorder) UpsertAIBedrockInferenceProfileModel(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIModelPrices", reflect.TypeOf((*MockStore)(nil).UpsertAIModelPrices), ctx, arg) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIBedrockInferenceProfileModel", reflect.TypeOf((*MockStore)(nil).UpsertAIBedrockInferenceProfileModel), ctx, arg) } -// UpsertAIProviderBedrockResolvedModels mocks base method. -func (m *MockStore) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg database.UpsertAIProviderBedrockResolvedModelsParams) error { +// UpsertAIModelPrices mocks base method. +func (m *MockStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpsertAIProviderBedrockResolvedModels", ctx, arg) + ret := m.ctrl.Call(m, "UpsertAIModelPrices", ctx, arg) ret0, _ := ret[0].(error) return ret0 } -// UpsertAIProviderBedrockResolvedModels indicates an expected call of UpsertAIProviderBedrockResolvedModels. -func (mr *MockStoreMockRecorder) UpsertAIProviderBedrockResolvedModels(ctx, arg any) *gomock.Call { +// UpsertAIModelPrices indicates an expected call of UpsertAIModelPrices. +func (mr *MockStoreMockRecorder) UpsertAIModelPrices(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIProviderBedrockResolvedModels", reflect.TypeOf((*MockStore)(nil).UpsertAIProviderBedrockResolvedModels), ctx, arg) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIModelPrices", reflect.TypeOf((*MockStore)(nil).UpsertAIModelPrices), ctx, arg) } // UpsertAISeatState mocks base method. diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 8bda5f9a52e..03bb1dc1397 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1567,6 +1567,13 @@ $$; COMMENT ON FUNCTION update_chat_history_after_message_update() IS 'Component of chatd. Updates history_version and generation_attempt on chats when chat_messages is updated. Excludes changes to search_tsv and search_tsv_config.'; +CREATE TABLE ai_bedrock_inference_profile_models ( + inference_profile_arn text NOT NULL, + resolved_model text NOT NULL +); + +COMMENT ON COLUMN ai_bedrock_inference_profile_models.resolved_model IS 'The Bedrock model ID the inference profile wraps.'; + CREATE TABLE ai_gateway_keys ( id uuid NOT NULL, created_at timestamp with time zone NOT NULL, @@ -1603,16 +1610,6 @@ COMMENT ON TABLE ai_model_prices IS 'Per-model token prices used by AI Bridge to COMMENT ON COLUMN ai_model_prices.source IS 'Where the price came from: default for the embedded price book, custom for a price set through the API. Both can exist for the same model.'; -CREATE TABLE ai_provider_bedrock_resolved_models ( - ai_provider_id uuid NOT NULL, - resolved_model text NOT NULL, - resolved_small_fast_model text NOT NULL -); - -COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_model IS 'The model ID behind the provider''s configured model identifier. Equal to the configured value when that value is already a model ID.'; - -COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_small_fast_model IS 'resolved_model for the provider''s configured small/fast model identifier.'; - CREATE TABLE ai_provider_keys ( id uuid DEFAULT gen_random_uuid() NOT NULL, provider_id uuid NOT NULL, @@ -4376,15 +4373,15 @@ ALTER TABLE ONLY workspace_resource_metadata ALTER COLUMN id SET DEFAULT nextval ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); +ALTER TABLE ONLY ai_bedrock_inference_profile_models + ADD CONSTRAINT ai_bedrock_inference_profile_models_pkey PRIMARY KEY (inference_profile_arn); + ALTER TABLE ONLY ai_gateway_keys ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id); ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source); -ALTER TABLE ONLY ai_provider_bedrock_resolved_models - ADD CONSTRAINT ai_provider_bedrock_resolved_models_pkey PRIMARY KEY (ai_provider_id); - ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id); @@ -5267,9 +5264,6 @@ COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS 'U the uniqueness requirement. A trigger allows us to enforce uniqueness going forward without requiring a migration to clean up historical data.'; -ALTER TABLE ONLY ai_provider_bedrock_resolved_models - ADD CONSTRAINT ai_provider_bedrock_resolved_models_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; - ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_api_key_key_id_fkey FOREIGN KEY (api_key_key_id) REFERENCES dbcrypt_keys(active_key_digest); diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index d23b2fe5247..251ce1aec5c 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -6,7 +6,6 @@ type ForeignKeyConstraint string // ForeignKeyConstraint enums. const ( - ForeignKeyAIProviderBedrockResolvedModelsAIProviderID ForeignKeyConstraint = "ai_provider_bedrock_resolved_models_ai_provider_id_fkey" // ALTER TABLE ONLY ai_provider_bedrock_resolved_models ADD CONSTRAINT ai_provider_bedrock_resolved_models_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; ForeignKeyAIProviderKeysAPIKeyKeyID ForeignKeyConstraint = "ai_provider_keys_api_key_key_id_fkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_api_key_key_id_fkey FOREIGN KEY (api_key_key_id) REFERENCES dbcrypt_keys(active_key_digest); ForeignKeyAIProviderKeysProviderID ForeignKeyConstraint = "ai_provider_keys_provider_id_fkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE; ForeignKeyAIProvidersSettingsKeyID ForeignKeyConstraint = "ai_providers_settings_key_id_fkey" // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_settings_key_id_fkey FOREIGN KEY (settings_key_id) REFERENCES dbcrypt_keys(active_key_digest); diff --git a/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql b/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql new file mode 100644 index 00000000000..00cf210e554 --- /dev/null +++ b/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql @@ -0,0 +1 @@ +DROP TABLE ai_bedrock_inference_profile_models; diff --git a/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql b/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql new file mode 100644 index 00000000000..c15483964a0 --- /dev/null +++ b/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql @@ -0,0 +1,14 @@ +-- An application inference profile ARN identifies a Bedrock billing wrapper +-- rather than a model, and the model it wraps is fixed: Bedrock offers no way +-- to repoint a profile, so a new target requires a new profile and a new ARN. +-- +-- This table records what each ARN resolves to, so the gateway can detect +-- capabilities, price usage, and record interceptions without calling the +-- Bedrock control plane. Rows are written when a provider is saved and are +-- never invalidated, only corrected by a later save. +CREATE TABLE ai_bedrock_inference_profile_models ( + inference_profile_arn text PRIMARY KEY, + resolved_model text NOT NULL +); + +COMMENT ON COLUMN ai_bedrock_inference_profile_models.resolved_model IS 'The Bedrock model ID the inference profile wraps.'; diff --git a/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql b/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql deleted file mode 100644 index 6e2f68cf04f..00000000000 --- a/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.down.sql +++ /dev/null @@ -1 +0,0 @@ -DROP TABLE ai_provider_bedrock_resolved_models; diff --git a/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql b/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql deleted file mode 100644 index 9dd63146580..00000000000 --- a/coderd/database/migrations/000591_ai_provider_bedrock_resolved_models.up.sql +++ /dev/null @@ -1,17 +0,0 @@ --- Application inference profile ARNs are opaque: they identify a Bedrock --- billing wrapper, not a model. This table records the model each configured --- ARN resolves to, so the gateway can detect capabilities, price usage, and --- record interceptions without calling the Bedrock control plane itself. --- --- A provider has a row only while one of its identifiers is an ARN. Providers --- configured with plain model IDs need no mapping, and neither does a provider --- whose profile could not be resolved: the gateway refuses to serve it. -CREATE TABLE ai_provider_bedrock_resolved_models ( - ai_provider_id uuid PRIMARY KEY REFERENCES ai_providers (id) ON DELETE CASCADE, - resolved_model text NOT NULL, - resolved_small_fast_model text NOT NULL -); - -COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_model IS 'The model ID behind the provider''s configured model identifier. Equal to the configured value when that value is already a model ID.'; - -COMMENT ON COLUMN ai_provider_bedrock_resolved_models.resolved_small_fast_model IS 'resolved_model for the provider''s configured small/fast model identifier.'; diff --git a/coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql b/coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql new file mode 100644 index 00000000000..7b9e9984382 --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql @@ -0,0 +1,8 @@ +INSERT INTO ai_bedrock_inference_profile_models ( + inference_profile_arn, + resolved_model +) VALUES + ( + 'arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/fixtureprofile', + 'anthropic.claude-sonnet-4-5-20250929-v1:0' + ); diff --git a/coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql b/coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql deleted file mode 100644 index 9b1fd9996f7..00000000000 --- a/coderd/database/migrations/testdata/fixtures/000591_ai_provider_bedrock_resolved_models.up.sql +++ /dev/null @@ -1,31 +0,0 @@ -INSERT INTO ai_providers ( - id, - type, - name, - display_name, - enabled, - deleted, - base_url, - settings -) VALUES - ( - '8e3c6e18-2b75-4c3f-9b35-9d1c6f4e1a04', - 'bedrock', - 'bedrock-inference-profile', - 'Bedrock via Application Inference Profile (Fixture)', - TRUE, - FALSE, - 'https://bedrock-runtime.us-west-2.amazonaws.com/', - '{"_type":"bedrock","_version":1,"region":"us-west-2","model":"arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/fixtureprofile","small_fast_model":"anthropic.claude-3-5-haiku-20241022-v1:0","access_key":"fixture-bedrock-access-key","access_key_secret":"fixture-bedrock-access-key-secret"}' - ); - -INSERT INTO ai_provider_bedrock_resolved_models ( - ai_provider_id, - resolved_model, - resolved_small_fast_model -) VALUES - ( - '8e3c6e18-2b75-4c3f-9b35-9d1c6f4e1a04', - 'anthropic.claude-sonnet-4-5-20250929-v1:0', - 'anthropic.claude-3-5-haiku-20241022-v1:0' - ); diff --git a/coderd/database/models.go b/coderd/database/models.go index 67df38495eb..a011a3bd441 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4850,6 +4850,12 @@ func AllWorkspaceTransitionValues() []WorkspaceTransition { } } +type AIBedrockInferenceProfileModel struct { + InferenceProfileArn string `db:"inference_profile_arn" json:"inference_profile_arn"` + // The Bedrock model ID the inference profile wraps. + ResolvedModel string `db:"resolved_model" json:"resolved_model"` +} + // Audit log of requests intercepted by AI Bridge type AIBridgeInterception struct { ID uuid.UUID `db:"id" json:"id"` @@ -4991,14 +4997,6 @@ type AIProvider struct { Icon string `db:"icon" json:"icon"` } -type AIProviderBedrockResolvedModel struct { - AIProviderID uuid.UUID `db:"ai_provider_id" json:"ai_provider_id"` - // The model ID behind the provider's configured model identifier. Equal to the configured value when that value is already a model ID. - ResolvedModel string `db:"resolved_model" json:"resolved_model"` - // resolved_model for the provider's configured small/fast model identifier. - ResolvedSmallFastModel string `db:"resolved_small_fast_model" json:"resolved_small_fast_model"` -} - // API keys associated with AI providers. Bedrock providers have zero keys (they authenticate via settings). OpenAI and Anthropic providers have one or more keys for failover. type AIProviderKey struct { ID uuid.UUID `db:"id" json:"id"` diff --git a/coderd/database/querier.go b/coderd/database/querier.go index d18293895c1..71d32f5f92f 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -117,7 +117,6 @@ type sqlcQuerier interface { CreateUserSecret(ctx context.Context, arg CreateUserSecretParams) (UserSecret, error) CustomRoles(ctx context.Context, arg CustomRolesParams) ([]CustomRole, error) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) (DeleteAIGatewayKeyRow, error) - DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error DeleteAIProviderKey(ctx context.Context, id uuid.UUID) error DeleteAPIKeyByID(ctx context.Context, id string) error @@ -295,6 +294,7 @@ type sqlcQuerier interface { // The query finds presets where all preset parameters are present in the provided parameters, // and returns the preset with the most parameters (largest subset). FindMatchingPresetID(ctx context.Context, arg FindMatchingPresetIDParams) (uuid.UUID, error) + GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]AIBedrockInferenceProfileModel, error) // AI Gateway cost for one chat tree: the root chat plus every subagent // beneath it. The spawning chat's ID is recorded as the interception session // ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed @@ -341,7 +341,6 @@ type sqlcQuerier interface { // each source forms its own group and nothing collapses. Every other source // contributes the same constant, leaving the key as (provider, model). GetAIModelPrices(ctx context.Context, arg GetAIModelPricesParams) ([]AIModelPrice, error) - GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]AIProviderBedrockResolvedModel, error) GetAIProviderByID(ctx context.Context, id uuid.UUID) (AIProvider, error) // Lock the provider row until the model-config write completes. The // transaction alone does not stop a concurrent soft-delete or disable @@ -1661,6 +1660,11 @@ type sqlcQuerier interface { UpdateWorkspaceTTL(ctx context.Context, arg UpdateWorkspaceTTLParams) error UpdateWorkspacesDormantDeletingAtByTemplateID(ctx context.Context, arg UpdateWorkspacesDormantDeletingAtByTemplateIDParams) ([]WorkspaceTable, error) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg UpdateWorkspacesTTLByTemplateIDParams) error + // Records the model an application inference profile ARN resolves to. The + // provider write path resolves the ARN through the Bedrock control plane and + // stores the answer here, so the gateway never has to. An upsert rather than + // an insert so a later save corrects a stored value. + UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg UpsertAIBedrockInferenceProfileModelParams) error // Upsert a batch of model prices from a JSON array, all recorded under the // given source. Each element must have provider, model, and the four price // fields, and null prices are written as SQL NULL. @@ -1669,10 +1673,6 @@ type sqlcQuerier interface { // differs, so updated_at records when a price last changed. Prices are // nullable and a NULL on either side counts as a difference. UpsertAIModelPrices(ctx context.Context, arg UpsertAIModelPricesParams) error - // Records the models an application inference profile ARN resolves to. The - // provider write path resolves the ARN through the Bedrock control plane and - // stores the answer here, so the gateway never has to. - UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg UpsertAIProviderBedrockResolvedModelsParams) error // Returns true if a new rows was inserted, false otherwise. UpsertAISeatState(ctx context.Context, arg UpsertAISeatStateParams) (bool, error) UpsertAnnouncementBanners(ctx context.Context, value string) error diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index f017339ddee..93121e8c4c8 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -553,18 +553,6 @@ func (q *sqlQuerier) UpdateEncryptedAIProviderKey(ctx context.Context, arg Updat return i, err } -const deleteAIProviderBedrockResolvedModels = `-- name: DeleteAIProviderBedrockResolvedModels :exec -DELETE FROM - ai_provider_bedrock_resolved_models -WHERE - ai_provider_id = $1::uuid -` - -func (q *sqlQuerier) DeleteAIProviderBedrockResolvedModels(ctx context.Context, aiProviderID uuid.UUID) error { - _, err := q.db.ExecContext(ctx, deleteAIProviderBedrockResolvedModels, aiProviderID) - return err -} - const deleteAIProviderByID = `-- name: DeleteAIProviderByID :exec UPDATE ai_providers @@ -581,25 +569,25 @@ func (q *sqlQuerier) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) err return err } -const getAIProviderBedrockResolvedModelsByProviderIDs = `-- name: GetAIProviderBedrockResolvedModelsByProviderIDs :many +const getAIBedrockInferenceProfileModels = `-- name: GetAIBedrockInferenceProfileModels :many SELECT - ai_provider_id, resolved_model, resolved_small_fast_model + inference_profile_arn, resolved_model FROM - ai_provider_bedrock_resolved_models + ai_bedrock_inference_profile_models WHERE - ai_provider_id = ANY($1::uuid[]) + inference_profile_arn = ANY($1::text[]) ` -func (q *sqlQuerier) GetAIProviderBedrockResolvedModelsByProviderIDs(ctx context.Context, aiProviderIds []uuid.UUID) ([]AIProviderBedrockResolvedModel, error) { - rows, err := q.db.QueryContext(ctx, getAIProviderBedrockResolvedModelsByProviderIDs, pq.Array(aiProviderIds)) +func (q *sqlQuerier) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]AIBedrockInferenceProfileModel, error) { + rows, err := q.db.QueryContext(ctx, getAIBedrockInferenceProfileModels, pq.Array(inferenceProfileArns)) if err != nil { return nil, err } defer rows.Close() - var items []AIProviderBedrockResolvedModel + var items []AIBedrockInferenceProfileModel for rows.Next() { - var i AIProviderBedrockResolvedModel - if err := rows.Scan(&i.AIProviderID, &i.ResolvedModel, &i.ResolvedSmallFastModel); err != nil { + var i AIBedrockInferenceProfileModel + if err := rows.Scan(&i.InferenceProfileArn, &i.ResolvedModel); err != nil { return nil, err } items = append(items, i) @@ -928,27 +916,26 @@ func (q *sqlQuerier) UpdateEncryptedAIProviderSettings(ctx context.Context, arg return i, err } -const upsertAIProviderBedrockResolvedModels = `-- name: UpsertAIProviderBedrockResolvedModels :exec +const upsertAIBedrockInferenceProfileModel = `-- name: UpsertAIBedrockInferenceProfileModel :exec INSERT INTO - ai_provider_bedrock_resolved_models (ai_provider_id, resolved_model, resolved_small_fast_model) + ai_bedrock_inference_profile_models (inference_profile_arn, resolved_model) VALUES - ($1::uuid, $2::text, $3::text) -ON CONFLICT (ai_provider_id) DO UPDATE SET - resolved_model = $2::text, - resolved_small_fast_model = $3::text + ($1::text, $2::text) +ON CONFLICT (inference_profile_arn) DO UPDATE SET + resolved_model = $2::text ` -type UpsertAIProviderBedrockResolvedModelsParams struct { - AIProviderID uuid.UUID `db:"ai_provider_id" json:"ai_provider_id"` - ResolvedModel string `db:"resolved_model" json:"resolved_model"` - ResolvedSmallFastModel string `db:"resolved_small_fast_model" json:"resolved_small_fast_model"` +type UpsertAIBedrockInferenceProfileModelParams struct { + InferenceProfileArn string `db:"inference_profile_arn" json:"inference_profile_arn"` + ResolvedModel string `db:"resolved_model" json:"resolved_model"` } -// Records the models an application inference profile ARN resolves to. The +// Records the model an application inference profile ARN resolves to. The // provider write path resolves the ARN through the Bedrock control plane and -// stores the answer here, so the gateway never has to. -func (q *sqlQuerier) UpsertAIProviderBedrockResolvedModels(ctx context.Context, arg UpsertAIProviderBedrockResolvedModelsParams) error { - _, err := q.db.ExecContext(ctx, upsertAIProviderBedrockResolvedModels, arg.AIProviderID, arg.ResolvedModel, arg.ResolvedSmallFastModel) +// stores the answer here, so the gateway never has to. An upsert rather than +// an insert so a later save corrects a stored value. +func (q *sqlQuerier) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg UpsertAIBedrockInferenceProfileModelParams) error { + _, err := q.db.ExecContext(ctx, upsertAIBedrockInferenceProfileModel, arg.InferenceProfileArn, arg.ResolvedModel) return err } diff --git a/coderd/database/queries/ai_providers.sql b/coderd/database/queries/ai_providers.sql index 6e003ce7c20..d322087153d 100644 --- a/coderd/database/queries/ai_providers.sql +++ b/coderd/database/queries/ai_providers.sql @@ -107,28 +107,22 @@ WHERE RETURNING *; --- name: UpsertAIProviderBedrockResolvedModels :exec --- Records the models an application inference profile ARN resolves to. The +-- name: UpsertAIBedrockInferenceProfileModel :exec +-- Records the model an application inference profile ARN resolves to. The -- provider write path resolves the ARN through the Bedrock control plane and --- stores the answer here, so the gateway never has to. +-- stores the answer here, so the gateway never has to. An upsert rather than +-- an insert so a later save corrects a stored value. INSERT INTO - ai_provider_bedrock_resolved_models (ai_provider_id, resolved_model, resolved_small_fast_model) + ai_bedrock_inference_profile_models (inference_profile_arn, resolved_model) VALUES - (@ai_provider_id::uuid, @resolved_model::text, @resolved_small_fast_model::text) -ON CONFLICT (ai_provider_id) DO UPDATE SET - resolved_model = @resolved_model::text, - resolved_small_fast_model = @resolved_small_fast_model::text; + (@inference_profile_arn::text, @resolved_model::text) +ON CONFLICT (inference_profile_arn) DO UPDATE SET + resolved_model = @resolved_model::text; --- name: DeleteAIProviderBedrockResolvedModels :exec -DELETE FROM - ai_provider_bedrock_resolved_models -WHERE - ai_provider_id = @ai_provider_id::uuid; - --- name: GetAIProviderBedrockResolvedModelsByProviderIDs :many +-- name: GetAIBedrockInferenceProfileModels :many SELECT * FROM - ai_provider_bedrock_resolved_models + ai_bedrock_inference_profile_models WHERE - ai_provider_id = ANY(@ai_provider_ids::uuid[]); + inference_profile_arn = ANY(@inference_profile_arns::text[]); diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 332ae755c0a..8c849a6877c 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -7,9 +7,9 @@ type UniqueConstraint string // UniqueConstraint enums. const ( UniqueAgentStatsPkey UniqueConstraint = "agent_stats_pkey" // ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); + UniqueAIBedrockInferenceProfileModelsPkey UniqueConstraint = "ai_bedrock_inference_profile_models_pkey" // ALTER TABLE ONLY ai_bedrock_inference_profile_models ADD CONSTRAINT ai_bedrock_inference_profile_models_pkey PRIMARY KEY (inference_profile_arn); UniqueAIGatewayKeysPkey UniqueConstraint = "ai_gateway_keys_pkey" // ALTER TABLE ONLY ai_gateway_keys ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id); UniqueAIModelPricesPkey UniqueConstraint = "ai_model_prices_pkey" // ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source); - UniqueAIProviderBedrockResolvedModelsPkey UniqueConstraint = "ai_provider_bedrock_resolved_models_pkey" // ALTER TABLE ONLY ai_provider_bedrock_resolved_models ADD CONSTRAINT ai_provider_bedrock_resolved_models_pkey PRIMARY KEY (ai_provider_id); UniqueAIProviderKeysPkey UniqueConstraint = "ai_provider_keys_pkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id); UniqueAIProvidersPkey UniqueConstraint = "ai_providers_pkey" // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_pkey PRIMARY KEY (id); UniqueAISeatStatePkey UniqueConstraint = "ai_seat_state_pkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id); From 03bfa2c3642b5d9e9d0148b1bb34eb15c33887cf Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 19:27:28 +0000 Subject: [PATCH 15/25] refactor(coderd): report any bedrock resolution failure as a bad request --- coderd/ai_providers_bedrock.go | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index 4645a8b8ea1..6452af03528 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -15,17 +15,6 @@ import ( "github.com/coder/coder/v2/codersdk" ) -// bedrockProfileUnresolvableError marks a failed Bedrock inference profile -// lookup so the write path reports it as a client-visible validation failure -// rather than an internal error. -type bedrockProfileUnresolvableError struct{ err error } - -func (e bedrockProfileUnresolvableError) Error() string { - return "resolve bedrock inference profile: " + e.err.Error() -} - -func (e bedrockProfileUnresolvableError) Unwrap() error { return e.err } - // resolveBedrockModels records which model each of the provider's application // inference profile ARNs refers to. An ARN identifies a billing wrapper rather // than a model, so the gateway needs the mapping to detect capabilities, price @@ -52,7 +41,7 @@ func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvide resolved, err := provider.ResolveBedrockModels(ctx, *cfg) if err != nil { - return bedrockProfileUnresolvableError{err: err} + return xerrors.Errorf("resolve bedrock inference profile: %w", err) } for profileARN, model := range resolved { @@ -72,11 +61,6 @@ func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvide // resolution the gateway cannot tell what an opaque profile ARN refers to, so // it serves the ARN as its own identity until a later save resolves it. func (api *API) writeAIProviderResolutionError(ctx context.Context, rw http.ResponseWriter, err error) { - var unresolvable bedrockProfileUnresolvableError - if !xerrors.As(err, &unresolvable) { - writeAIProviderError(ctx, api.Logger, rw, err, "resolve bedrock inference profile", "Internal error resolving the Bedrock application inference profile.") - return - } api.Logger.Warn(ctx, "resolve bedrock inference profile", slog.Error(err)) httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Could not resolve the Bedrock application inference profile. Check that the ARN is correct and that the AWS identity used by Coder is allowed bedrock:GetInferenceProfile.", From b1fdf6c8cdc6558f38a34b59c177df3d0149a907 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Wed, 9 Sep 2026 19:35:25 +0000 Subject: [PATCH 16/25] docs(coderd): trim bedrock resolution comments --- coderd/ai_providers_bedrock.go | 25 +++++++------------------ 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index 6452af03528..de4fd2322a3 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -15,25 +15,15 @@ import ( "github.com/coder/coder/v2/codersdk" ) -// resolveBedrockModels records which model each of the provider's application -// inference profile ARNs refers to. An ARN identifies a billing wrapper rather -// than a model, so the gateway needs the mapping to detect capabilities, price -// usage, and record interceptions. -// -// Resolution is an AWS call, so it runs after the provider write has committed -// rather than holding a database transaction open across the network. It reads -// the stored row because that is the merged configuration the provider will -// actually use. -// -// It runs on every save, even for an ARN another provider already resolved, so -// that saving proves this provider's own identity can read the profile. +// resolveBedrockModels stores the model each of the provider's application +// inference profile ARNs refers to. It runs after the write commits, and on +// every save, because it calls AWS. func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvider) error { settings, err := db2sdk.AIProviderSettings(row.Settings) if err != nil { return xerrors.Errorf("decode settings: %w", err) } - // Resolution is a Bedrock control-plane call, whereas the provider's - // BaseURL is its runtime endpoint, so it is deliberately not carried here. + // BaseURL is the runtime endpoint; resolution calls the control plane. cfg := agplaibridge.BedrockConfig("", settings.Bedrock) if cfg == nil { return nil @@ -56,10 +46,9 @@ func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvide return nil } -// writeAIProviderResolutionError reports a failed Bedrock model resolution. The -// provider keeps the identifiers the operator asked for, but without a -// resolution the gateway cannot tell what an opaque profile ARN refers to, so -// it serves the ARN as its own identity until a later save resolves it. +// writeAIProviderResolutionError reports a failed resolution. The provider is +// stored either way, and serves the ARN as its own identity until a later save +// resolves it. func (api *API) writeAIProviderResolutionError(ctx context.Context, rw http.ResponseWriter, err error) { api.Logger.Warn(ctx, "resolve bedrock inference profile", slog.Error(err)) httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ From 886dd2ed1aafe4f7f886ec5ede879c67e80e761d Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 12:27:55 +0000 Subject: [PATCH 17/25] revert(docs): restore the bedrock inference profile section --- docs/ai-coder/ai-gateway/providers.md | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index 670b8568c68..768d2e047cd 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -251,13 +251,10 @@ AI Gateway passes the profile upstream so AWS records the attribution, while internally resolving and using the underlying model identity, including for usage pricing. -Resolution requires a `GetInferenceProfile` call, which Coder makes when the -provider is saved, not when a request is served. The AWS identity Coder uses, -which is the provider's access keys when configured and otherwise the identity -of the Coder deployment, must have `bedrock:GetInferenceProfile` permission for -the profile. Saving reports an error when the lookup fails, and the provider -cannot serve requests until a later save resolves it. Providers configured with -plain model identifiers do not need this permission. +Resolution requires a `GetInferenceProfile` call, so the AWS identity used by +the gateway must have `bedrock:GetInferenceProfile` permission for the +profile. Providers configured with plain model identifiers do not need this +permission. If resolution fails, the provider is skipped. ### GitHub Copilot From e68042a212723feeedad910f2dcc446843d98bf8 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 11:22:02 -0400 Subject: [PATCH 18/25] refactor(coderd): store bedrock model resolution in provider settings (#29175) Follow-up to #29112, which stores the resolved Bedrock model in its own table. Keyed by inference profile ARN, that table needed a migration, three queries, dbauthz wrappers, a store-interface method, and a join in the provider payload. This PR keeps the values in the settings blob instead, where the rest of the provider's Bedrock configuration already lives. The plumbing that disappears is the point: `aibridgedserver` no longer collects identifiers, queries mappings, or threads a map through `aiProviderToProto`; it reads two fields off the settings it already decoded. Net 314 deletions against 119 insertions, and no schema change. Resolution still runs after the write commits, so no AWS call happens inside a transaction, and still runs on every save. Storing the result now means a second `UpdateAIProvider` with the resolved settings, which is the one thing the table did not need. Server ownership is enforced by clearing rather than validating. The write path zeroes `resolved_model` and `resolved_small_fast_model` before storing and rewrites them after resolution, so a client-supplied value is discarded rather than rejected, and a stale value cannot survive a settings change. That is one two-line helper instead of the validation, merge carry-forward, and compare-before-write that an earlier attempt at this needed. What is lost relative to the table: two providers configured with the same ARN each store their own copy, and each resolves it separately. Resolution already ran per save, so this costs storage rather than AWS calls. Behavior is unchanged. An unresolved ARN still serves as its own identity, a failed lookup still reports 400 while leaving the provider stored, and the gateway still never calls the Bedrock control plane. Relates to https://linear.app/codercom/issue/AIGOV-488 Created by Coder Agents on behalf of @evgeniy-scherbina. --- coderd/ai_providers.go | 13 +- coderd/ai_providers_bedrock.go | 56 +++++--- coderd/ai_providers_bedrock_test.go | 121 ++++++++---------- coderd/aibridgedserver/aibridgedserver.go | 44 +------ coderd/database/dbauthz/dbauthz.go | 14 -- coderd/database/dbauthz/dbauthz_test.go | 14 -- coderd/database/dbmetrics/querymetrics.go | 16 --- coderd/database/dbmock/dbmock.go | 29 ----- ..._bedrock_inference_profile_models.down.sql | 1 - ...ai_bedrock_inference_profile_models.up.sql | 14 -- ...ai_bedrock_inference_profile_models.up.sql | 8 -- coderd/database/querier.go | 6 - coderd/database/queries.sql.go | 55 -------- coderd/database/queries/ai_providers.sql | 20 --- codersdk/aiproviders_bedrock.go | 14 +- site/src/api/typesGenerated.ts | 8 ++ 16 files changed, 119 insertions(+), 314 deletions(-) delete mode 100644 coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql delete mode 100644 coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql delete mode 100644 coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index 70cd3c30670..c462a95c133 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -185,6 +185,7 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { // Generate the server-owned external ID when the provider assumes a role. ensureBedrockExternalID(&req.Settings) + clearBedrockModelResolution(&req.Settings) settings, err := encodeAIProviderSettings(req.Settings) if err != nil { @@ -246,10 +247,12 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { // Resolve inference profile ARNs once the provider is stored, then announce // it. The gateway never calls the Bedrock control plane itself. - if err := api.resolveBedrockModels(ctx, row); err != nil { + row, err = api.resolveBedrockModels(ctx, row) + if err != nil { api.writeAIProviderResolutionError(ctx, rw, err) return } + aReq.New = row auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, aiProviderKeyChanges{Added: keys}) api.publishAIProvidersChanged(ctx) @@ -339,6 +342,10 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { return err } existing = mergeAIProviderSettings(existing, *req.Settings) + // The patch may point the provider at different identifiers, and a + // client cannot supply resolutions of its own. Resolution runs again + // after the transaction. + clearBedrockModelResolution(&existing) } // Bedrock settings are only meaningful for anthropic- or // bedrock-typed providers; rejecting the mismatch keeps a @@ -449,10 +456,12 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { // identifiers or the credentials they resolve under, so any stored // resolution still holds. if req.Settings != nil { - if err := api.resolveBedrockModels(ctx, updated); err != nil { + updated, err = api.resolveBedrockModels(ctx, updated) + if err != nil { api.writeAIProviderResolutionError(ctx, rw, err) return } + aReq.New = updated } auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, keyChanges) diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index de4fd2322a3..42db6eab558 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -2,6 +2,7 @@ package coderd import ( "context" + "database/sql" "net/http" "golang.org/x/xerrors" @@ -16,34 +17,59 @@ import ( ) // resolveBedrockModels stores the model each of the provider's application -// inference profile ARNs refers to. It runs after the write commits, and on -// every save, because it calls AWS. -func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvider) error { +// inference profile ARNs refers to, returning the updated provider. It runs +// after the write commits, and on every save, because it calls AWS. +func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvider) (database.AIProvider, error) { settings, err := db2sdk.AIProviderSettings(row.Settings) if err != nil { - return xerrors.Errorf("decode settings: %w", err) + return row, xerrors.Errorf("decode settings: %w", err) } // BaseURL is the runtime endpoint; resolution calls the control plane. cfg := agplaibridge.BedrockConfig("", settings.Bedrock) if cfg == nil { - return nil + return row, nil } resolved, err := provider.ResolveBedrockModels(ctx, *cfg) if err != nil { - return xerrors.Errorf("resolve bedrock inference profile: %w", err) + return row, xerrors.Errorf("resolve bedrock inference profile: %w", err) } + if len(resolved) == 0 { + return row, nil + } + settings.Bedrock.ResolvedModel = resolved[settings.Bedrock.Model] + settings.Bedrock.ResolvedSmallFastModel = resolved[settings.Bedrock.SmallFastModel] + + encoded, err := encodeAIProviderSettings(settings) + if err != nil { + return row, xerrors.Errorf("encode settings: %w", err) + } + updated, err := api.Database.UpdateAIProvider(ctx, database.UpdateAIProviderParams{ + ID: row.ID, + Type: row.Type, + DisplayName: row.DisplayName, + Icon: row.Icon, + Enabled: row.Enabled, + BaseUrl: row.BaseUrl, + Settings: encoded, + // SettingsKeyID is set by the dbcrypt wrapper. + SettingsKeyID: sql.NullString{}, + }) + if err != nil { + return row, xerrors.Errorf("store resolved models: %w", err) + } + return updated, nil +} - for profileARN, model := range resolved { - err := api.Database.UpsertAIBedrockInferenceProfileModel(ctx, database.UpsertAIBedrockInferenceProfileModelParams{ - InferenceProfileArn: profileARN, - ResolvedModel: model, - }) - if err != nil { - return xerrors.Errorf("store resolved model for %q: %w", profileARN, err) - } +// clearBedrockModelResolution drops resolved identifiers a client supplied or +// an earlier save stored. The values are server-owned and rewritten after the +// write, so anything present beforehand is stale or forged. +func clearBedrockModelResolution(settings *codersdk.AIProviderSettings) { + if settings.Bedrock == nil { + return } - return nil + settings.Bedrock.ResolvedModel = "" + settings.Bedrock.ResolvedSmallFastModel = "" } // writeAIProviderResolutionError reports a failed resolution. The provider is diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go index bbd416511eb..7af50bb9f7a 100644 --- a/coderd/ai_providers_bedrock_test.go +++ b/coderd/ai_providers_bedrock_test.go @@ -1,7 +1,6 @@ package coderd_test import ( - "context" "net/http" "net/http/httptest" "slices" @@ -12,8 +11,6 @@ import ( "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/coderdtest" - "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) @@ -67,20 +64,6 @@ func respondWithModel(modelARN string) http.HandlerFunc { } } -// resolvedModel returns the model stored for an inference profile ARN, or the -// empty string when the ARN has no mapping. -func resolvedModel(ctx context.Context, t *testing.T, db database.Store, profileARN string) string { - t.Helper() - - rows, err := db.GetAIBedrockInferenceProfileModels(ctx, []string{profileARN}) - require.NoError(t, err) - if len(rows) == 0 { - return "" - } - require.Len(t, rows, 1) - return rows[0].ResolvedModel -} - // TestAIProvidersBedrockProfileResolution drives provider writes against a mock // Bedrock control plane, so the AWS SDK path runs for real. // NOTE: no t.Parallel() because the subtests use t.Setenv. @@ -95,8 +78,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - db, ps := dbtestutil.NewDB(t) - client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) + client := coderdtest.New(t, nil) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -114,9 +96,9 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { // invocation target, and AWS attributes spend to them. require.Equal(t, testProfileARN, created.Settings.Bedrock.Model) require.Equal(t, testSmallFastProfileARN, created.Settings.Bedrock.SmallFastModel) + require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) + require.Equal(t, "anthropic.claude-haiku-4-5", created.Settings.Bedrock.ResolvedSmallFastModel) require.Len(t, paths(), 2, "each profile is resolved once") - require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) - require.Equal(t, "anthropic.claude-haiku-4-5", resolvedModel(ctx, t, db, testSmallFastProfileARN)) }) t.Run("CreateLeavesPlainModelIDsUnresolved", func(t *testing.T) { @@ -125,8 +107,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - db, ps := dbtestutil.NewDB(t) - client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) + client := coderdtest.New(t, nil) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -139,11 +120,38 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.NotNil(t, created.Settings.Bedrock) - require.Empty(t, paths(), "plain model ids are already model identities") + require.Empty(t, created.Settings.Bedrock.ResolvedModel) + require.Empty(t, created.Settings.Bedrock.ResolvedSmallFastModel) + require.Empty(t, paths()) + }) + + t.Run("CreateIgnoresClientSuppliedResolution", func(t *testing.T) { + url, paths := mockBedrock(t, func(http.ResponseWriter, *http.Request) { + t.Error("Bedrock called for plain model ids") + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) + ctx := testutil.Context(t, testutil.WaitLong) + + settings := bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5") + settings.Bedrock.ResolvedModel = "anthropic.claude-opus-4-8" + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-spoofed", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *settings, + }) + require.NoError(t, err) + require.Empty(t, created.Settings.Bedrock.ResolvedModel, "the server owns the resolution") + require.Empty(t, paths()) }) - t.Run("CreateRejectsUnresolvableProfile", func(t *testing.T) { + t.Run("CreateReportsUnresolvableProfile", func(t *testing.T) { url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") @@ -152,8 +160,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - db, ps := dbtestutil.NewDB(t) - client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) + client := coderdtest.New(t, nil) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -170,21 +177,20 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) require.Contains(t, sdkErr.Detail, "GetInferenceProfile") - // The provider is stored with the ARN the operator asked for, but - // nothing maps that ARN, so the gateway serves it as its own identity. + // The provider is stored with the ARN the operator asked for, and + // serves it as its own identity until a later save resolves it. //nolint:gocritic // Owner role is the audience for this endpoint. providers, err := client.AIProviders(ctx) require.NoError(t, err) require.Len(t, providers, 1) - require.Empty(t, resolvedModel(ctx, t, db, testProfileARN)) + require.Empty(t, providers[0].Settings.Bedrock.ResolvedModel) }) t.Run("UpdateReresolvesChangedProfile", func(t *testing.T) { url, _ := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - db, ps := dbtestutil.NewDB(t) - client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) + client := coderdtest.New(t, nil) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -197,7 +203,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Empty(t, resolvedModel(ctx, t, db, testProfileARN)) + require.Empty(t, created.Settings.Bedrock.ResolvedModel) //nolint:gocritic // Owner role is the audience for this endpoint. updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ @@ -205,15 +211,15 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { }) require.NoError(t, err) require.Equal(t, testProfileARN, updated.Settings.Bedrock.Model) - require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) + require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) + require.Empty(t, updated.Settings.Bedrock.ResolvedSmallFastModel, "a plain model id is its own identity") }) - t.Run("UpdateToPlainModelIDNeedsNoResolution", func(t *testing.T) { + t.Run("UpdateToPlainModelIDClearsResolution", func(t *testing.T) { url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - db, ps := dbtestutil.NewDB(t) - client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) + client := coderdtest.New(t, nil) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -226,14 +232,15 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) + require.Equal(t, "anthropic.claude-opus-4-8", created.Settings.Bedrock.ResolvedModel) callsAfterCreate := len(paths()) //nolint:gocritic // Owner role is the audience for this endpoint. - _, err = client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ Settings: bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), }) require.NoError(t, err) + require.Empty(t, updated.Settings.Bedrock.ResolvedModel) require.Len(t, paths(), callsAfterCreate, "no profile is left to resolve") }) @@ -241,8 +248,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - db, ps := dbtestutil.NewDB(t) - client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) + client := coderdtest.New(t, nil) _ = coderdtest.CreateFirstUser(t, client) ctx := testutil.Context(t, testutil.WaitLong) @@ -259,38 +265,11 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { enabled := false //nolint:gocritic // Owner role is the audience for this endpoint. - _, err = client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + updated, err := client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ Enabled: &enabled, }) require.NoError(t, err) - require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) + require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) require.Len(t, paths(), callsAfterCreate, "an unrelated update does not call AWS") }) - - t.Run("SavingResolvesEvenWhenTheARNIsAlreadyMapped", func(t *testing.T) { - url, paths := mockBedrock(t, respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")) - t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) - - db, ps := dbtestutil.NewDB(t) - client := coderdtest.New(t, &coderdtest.Options{Database: db, Pubsub: ps}) - _ = coderdtest.CreateFirstUser(t, client) - ctx := testutil.Context(t, testutil.WaitLong) - - for _, name := range []string{"bedrock-first", "bedrock-second"} { - //nolint:gocritic // Owner role is the audience for this endpoint. - _, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ - Name: name, - Type: codersdk.AIProviderTypeBedrock, - BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", - Enabled: true, - Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), - }) - require.NoError(t, err) - } - - // The mapping is shared, but each save proves that provider's own - // identity can read the profile. - require.Len(t, paths(), 2) - require.Equal(t, "anthropic.claude-opus-4-8", resolvedModel(ctx, t, db, testProfileARN)) - }) } diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 4f2cdc6a41b..cf76f9e27c0 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -100,7 +100,6 @@ type store interface { // any in-flight env seed holding LockIDAIProvidersEnvSeed. GetAIProviders(ctx context.Context, arg database.GetAIProvidersParams) ([]database.AIProvider, error) GetAIProviderKeysByProviderIDs(ctx context.Context, providerIDs []uuid.UUID) ([]database.AIProviderKey, error) - GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) InTx(func(database.Store) error, *database.TxOptions) error } @@ -958,7 +957,6 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ var ( rows []database.AIProvider keysByProvider map[uuid.UUID][]database.AIProviderKey - modelByProfile map[string]string ) // Wrap both reads in a read-only transaction so the provider list and the // key list are consistent with each other, and so the seed lock is held @@ -996,18 +994,6 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ for _, k := range keyRows { keysByProvider[k.ProviderID] = append(keysByProvider[k.ProviderID], k) } - - // Bedrock application inference profile ARNs are opaque, so the models - // they refer to are resolved when a provider is written and read back - // here. An ARN without a mapping is served as its own identity. - profileRows, err := tx.GetAIBedrockInferenceProfileModels(ctx, bedrockModelIdentifiers(rows)) - if err != nil { - return xerrors.Errorf("get bedrock inference profile models: %w", err) - } - modelByProfile = make(map[string]string, len(profileRows)) - for _, r := range profileRows { - modelByProfile[r.InferenceProfileArn] = r.ResolvedModel - } return nil }, &database.TxOptions{ReadOnly: true, TxIdentifier: "get_ai_providers"}) if err != nil { @@ -1016,7 +1002,7 @@ func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequ providers := make([]*proto.AIProvider, 0, len(rows)) for _, row := range rows { - p, err := aiProviderToProto(row, keysByProvider[row.ID], modelByProfile) + p, err := aiProviderToProto(row, keysByProvider[row.ID]) if err != nil { // Skip the offending row rather than failing the whole fetch: // one row with a corrupt settings blob must not break provider @@ -1190,33 +1176,11 @@ func parseOptionalInt32(n *int32) sql.NullInt32 { return sql.NullInt32{Int32: *n, Valid: true} } -// bedrockModelIdentifiers returns the configured Bedrock model identifiers of -// every enabled provider, which is the key set for the inference profile -// mapping. Rows whose settings cannot be decoded are skipped; aiProviderToProto -// reports that failure when it builds the payload. -func bedrockModelIdentifiers(rows []database.AIProvider) []string { - var identifiers []string - for _, row := range rows { - if !row.Enabled { - continue - } - settings, err := db2sdk.AIProviderSettings(row.Settings) - if err != nil || settings.Bedrock == nil { - continue - } - identifiers = append(identifiers, settings.Bedrock.Model, settings.Bedrock.SmallFastModel) - } - return identifiers -} - // aiProviderToProto maps a single ai_providers row (and its keys, for enabled // providers) to the proto representation served to AI Gateway daemons. Keys and // Bedrock settings are only attached for enabled providers; disabled providers // never call upstream so their secrets are withheld. -// -// modelByProfile maps an application inference profile ARN to the model it -// wraps. An identifier absent from it needs no resolution, or has none yet. -func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey, modelByProfile map[string]string) (*proto.AIProvider, error) { +func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey) (*proto.AIProvider, error) { p := &proto.AIProvider{ Name: row.Name, Type: string(row.Type), @@ -1249,8 +1213,8 @@ func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey, m RoleArn: settings.Bedrock.RoleARN, ExternalId: settings.Bedrock.ExternalID, Protocol: string(settings.Bedrock.Protocol), - ResolvedModel: modelByProfile[settings.Bedrock.Model], - ResolvedSmallFastModel: modelByProfile[settings.Bedrock.SmallFastModel], + ResolvedModel: settings.Bedrock.ResolvedModel, + ResolvedSmallFastModel: settings.Bedrock.ResolvedSmallFastModel, } } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 7dbdfdb378f..e27731ab0fe 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2870,13 +2870,6 @@ func (q *querier) FindMatchingPresetID(ctx context.Context, arg database.FindMat return q.db.FindMatchingPresetID(ctx, arg) } -func (q *querier) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) { - if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIProvider); err != nil { - return nil, err - } - return q.db.GetAIBedrockInferenceProfileModels(ctx, inferenceProfileArns) -} - func (q *querier) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) { // The aggregate covers one chat tree, so it is authorized through the // root chat. Members cannot read interception rows back, but they can @@ -9061,13 +9054,6 @@ func (q *querier) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg datab return q.db.UpdateWorkspacesTTLByTemplateID(ctx, arg) } -func (q *querier) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg database.UpsertAIBedrockInferenceProfileModelParams) error { - if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil { - return err - } - return q.db.UpsertAIBedrockInferenceProfileModel(ctx, arg) -} - func (q *querier) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAiModelPrice); err != nil { return err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 9f32ae4fe00..5dd49f43068 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -7389,20 +7389,6 @@ func (s *MethodTestSuite) TestAIBridge() { dbm.EXPECT().UpdateEncryptedAIProviderSettings(gomock.Any(), arg).Return(provider, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns(provider) })) - s.Run("UpsertAIBedrockInferenceProfileModel", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - arg := database.UpsertAIBedrockInferenceProfileModelParams{ - InferenceProfileArn: "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/46u2vhiyo6z5", - ResolvedModel: "anthropic.claude-opus-4-8", - } - dbm.EXPECT().UpsertAIBedrockInferenceProfileModel(gomock.Any(), arg).Return(nil).AnyTimes() - check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns() - })) - s.Run("GetAIBedrockInferenceProfileModels", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - resolved := testutil.Fake(s.T(), faker, database.AIBedrockInferenceProfileModel{}) - arg := []string{resolved.InferenceProfileArn} - dbm.EXPECT().GetAIBedrockInferenceProfileModels(gomock.Any(), arg).Return([]database.AIBedrockInferenceProfileModel{resolved}, nil).AnyTimes() - check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionRead).Returns([]database.AIBedrockInferenceProfileModel{resolved}) - })) s.Run("GetAIProviderKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { key := testutil.Fake(s.T(), faker, database.AIProviderKey{}) dbm.EXPECT().GetAIProviderKeyByID(gomock.Any(), key.ID).Return(key, nil).AnyTimes() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 243ea4cfeca..365899afbd2 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1096,14 +1096,6 @@ func (m queryMetricsStore) FindMatchingPresetID(ctx context.Context, arg databas return r0, r1 } -func (m queryMetricsStore) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) { - start := time.Now() - r0, r1 := m.s.GetAIBedrockInferenceProfileModels(ctx, inferenceProfileArns) - m.queryLatencies.WithLabelValues("GetAIBedrockInferenceProfileModels").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIBedrockInferenceProfileModels").Inc() - return r0, r1 -} - func (m queryMetricsStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) { start := time.Now() r0, r1 := m.s.GetAIBridgeChatCost(ctx, rootChatID) @@ -6320,14 +6312,6 @@ func (m queryMetricsStore) UpdateWorkspacesTTLByTemplateID(ctx context.Context, return r0 } -func (m queryMetricsStore) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg database.UpsertAIBedrockInferenceProfileModelParams) error { - start := time.Now() - r0 := m.s.UpsertAIBedrockInferenceProfileModel(ctx, arg) - m.queryLatencies.WithLabelValues("UpsertAIBedrockInferenceProfileModel").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertAIBedrockInferenceProfileModel").Inc() - return r0 -} - func (m queryMetricsStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { start := time.Now() r0 := m.s.UpsertAIModelPrices(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 5e9c7e19cd2..f52286cf511 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1903,21 +1903,6 @@ func (mr *MockStoreMockRecorder) FindMatchingPresetID(ctx, arg any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FindMatchingPresetID", reflect.TypeOf((*MockStore)(nil).FindMatchingPresetID), ctx, arg) } -// GetAIBedrockInferenceProfileModels mocks base method. -func (m *MockStore) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]database.AIBedrockInferenceProfileModel, error) { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetAIBedrockInferenceProfileModels", ctx, inferenceProfileArns) - ret0, _ := ret[0].([]database.AIBedrockInferenceProfileModel) - ret1, _ := ret[1].(error) - return ret0, ret1 -} - -// GetAIBedrockInferenceProfileModels indicates an expected call of GetAIBedrockInferenceProfileModels. -func (mr *MockStoreMockRecorder) GetAIBedrockInferenceProfileModels(ctx, inferenceProfileArns any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBedrockInferenceProfileModels", reflect.TypeOf((*MockStore)(nil).GetAIBedrockInferenceProfileModels), ctx, inferenceProfileArns) -} - // GetAIBridgeChatCost mocks base method. func (m *MockStore) GetAIBridgeChatCost(ctx context.Context, rootChatID uuid.UUID) (database.GetAIBridgeChatCostRow, error) { m.ctrl.T.Helper() @@ -11909,20 +11894,6 @@ func (mr *MockStoreMockRecorder) UpdateWorkspacesTTLByTemplateID(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateWorkspacesTTLByTemplateID", reflect.TypeOf((*MockStore)(nil).UpdateWorkspacesTTLByTemplateID), ctx, arg) } -// UpsertAIBedrockInferenceProfileModel mocks base method. -func (m *MockStore) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg database.UpsertAIBedrockInferenceProfileModelParams) error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpsertAIBedrockInferenceProfileModel", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 -} - -// UpsertAIBedrockInferenceProfileModel indicates an expected call of UpsertAIBedrockInferenceProfileModel. -func (mr *MockStoreMockRecorder) UpsertAIBedrockInferenceProfileModel(ctx, arg any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertAIBedrockInferenceProfileModel", reflect.TypeOf((*MockStore)(nil).UpsertAIBedrockInferenceProfileModel), ctx, arg) -} - // UpsertAIModelPrices mocks base method. func (m *MockStore) UpsertAIModelPrices(ctx context.Context, arg database.UpsertAIModelPricesParams) error { m.ctrl.T.Helper() diff --git a/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql b/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql deleted file mode 100644 index 00cf210e554..00000000000 --- a/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.down.sql +++ /dev/null @@ -1 +0,0 @@ -DROP TABLE ai_bedrock_inference_profile_models; diff --git a/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql b/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql deleted file mode 100644 index c15483964a0..00000000000 --- a/coderd/database/migrations/000591_ai_bedrock_inference_profile_models.up.sql +++ /dev/null @@ -1,14 +0,0 @@ --- An application inference profile ARN identifies a Bedrock billing wrapper --- rather than a model, and the model it wraps is fixed: Bedrock offers no way --- to repoint a profile, so a new target requires a new profile and a new ARN. --- --- This table records what each ARN resolves to, so the gateway can detect --- capabilities, price usage, and record interceptions without calling the --- Bedrock control plane. Rows are written when a provider is saved and are --- never invalidated, only corrected by a later save. -CREATE TABLE ai_bedrock_inference_profile_models ( - inference_profile_arn text PRIMARY KEY, - resolved_model text NOT NULL -); - -COMMENT ON COLUMN ai_bedrock_inference_profile_models.resolved_model IS 'The Bedrock model ID the inference profile wraps.'; diff --git a/coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql b/coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql deleted file mode 100644 index 7b9e9984382..00000000000 --- a/coderd/database/migrations/testdata/fixtures/000591_ai_bedrock_inference_profile_models.up.sql +++ /dev/null @@ -1,8 +0,0 @@ -INSERT INTO ai_bedrock_inference_profile_models ( - inference_profile_arn, - resolved_model -) VALUES - ( - 'arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/fixtureprofile', - 'anthropic.claude-sonnet-4-5-20250929-v1:0' - ); diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 71d32f5f92f..88e86e97304 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -294,7 +294,6 @@ type sqlcQuerier interface { // The query finds presets where all preset parameters are present in the provided parameters, // and returns the preset with the most parameters (largest subset). FindMatchingPresetID(ctx context.Context, arg FindMatchingPresetIDParams) (uuid.UUID, error) - GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]AIBedrockInferenceProfileModel, error) // AI Gateway cost for one chat tree: the root chat plus every subagent // beneath it. The spawning chat's ID is recorded as the interception session // ID (see chatprovider.CoderHeaders), so a subagent's requests are attributed @@ -1660,11 +1659,6 @@ type sqlcQuerier interface { UpdateWorkspaceTTL(ctx context.Context, arg UpdateWorkspaceTTLParams) error UpdateWorkspacesDormantDeletingAtByTemplateID(ctx context.Context, arg UpdateWorkspacesDormantDeletingAtByTemplateIDParams) ([]WorkspaceTable, error) UpdateWorkspacesTTLByTemplateID(ctx context.Context, arg UpdateWorkspacesTTLByTemplateIDParams) error - // Records the model an application inference profile ARN resolves to. The - // provider write path resolves the ARN through the Bedrock control plane and - // stores the answer here, so the gateway never has to. An upsert rather than - // an insert so a later save corrects a stored value. - UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg UpsertAIBedrockInferenceProfileModelParams) error // Upsert a batch of model prices from a JSON array, all recorded under the // given source. Each element must have provider, model, and the four price // fields, and null prices are written as SQL NULL. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 93121e8c4c8..ccbd4a5af2a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -569,38 +569,6 @@ func (q *sqlQuerier) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) err return err } -const getAIBedrockInferenceProfileModels = `-- name: GetAIBedrockInferenceProfileModels :many -SELECT - inference_profile_arn, resolved_model -FROM - ai_bedrock_inference_profile_models -WHERE - inference_profile_arn = ANY($1::text[]) -` - -func (q *sqlQuerier) GetAIBedrockInferenceProfileModels(ctx context.Context, inferenceProfileArns []string) ([]AIBedrockInferenceProfileModel, error) { - rows, err := q.db.QueryContext(ctx, getAIBedrockInferenceProfileModels, pq.Array(inferenceProfileArns)) - if err != nil { - return nil, err - } - defer rows.Close() - var items []AIBedrockInferenceProfileModel - for rows.Next() { - var i AIBedrockInferenceProfileModel - if err := rows.Scan(&i.InferenceProfileArn, &i.ResolvedModel); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Close(); err != nil { - return nil, err - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} - const getAIProviderByID = `-- name: GetAIProviderByID :one SELECT id, type, name, display_name, enabled, deleted, base_url, settings, settings_key_id, created_at, updated_at, icon @@ -916,29 +884,6 @@ func (q *sqlQuerier) UpdateEncryptedAIProviderSettings(ctx context.Context, arg return i, err } -const upsertAIBedrockInferenceProfileModel = `-- name: UpsertAIBedrockInferenceProfileModel :exec -INSERT INTO - ai_bedrock_inference_profile_models (inference_profile_arn, resolved_model) -VALUES - ($1::text, $2::text) -ON CONFLICT (inference_profile_arn) DO UPDATE SET - resolved_model = $2::text -` - -type UpsertAIBedrockInferenceProfileModelParams struct { - InferenceProfileArn string `db:"inference_profile_arn" json:"inference_profile_arn"` - ResolvedModel string `db:"resolved_model" json:"resolved_model"` -} - -// Records the model an application inference profile ARN resolves to. The -// provider write path resolves the ARN through the Bedrock control plane and -// stores the answer here, so the gateway never has to. An upsert rather than -// an insert so a later save corrects a stored value. -func (q *sqlQuerier) UpsertAIBedrockInferenceProfileModel(ctx context.Context, arg UpsertAIBedrockInferenceProfileModelParams) error { - _, err := q.db.ExecContext(ctx, upsertAIBedrockInferenceProfileModel, arg.InferenceProfileArn, arg.ResolvedModel) - return err -} - const calculateAIBridgeInterceptionsTelemetrySummary = `-- name: CalculateAIBridgeInterceptionsTelemetrySummary :one WITH interceptions_in_range AS ( -- Get all matching interceptions in the given timeframe. diff --git a/coderd/database/queries/ai_providers.sql b/coderd/database/queries/ai_providers.sql index d322087153d..2971918e46f 100644 --- a/coderd/database/queries/ai_providers.sql +++ b/coderd/database/queries/ai_providers.sql @@ -106,23 +106,3 @@ WHERE id = @id::uuid RETURNING *; - --- name: UpsertAIBedrockInferenceProfileModel :exec --- Records the model an application inference profile ARN resolves to. The --- provider write path resolves the ARN through the Bedrock control plane and --- stores the answer here, so the gateway never has to. An upsert rather than --- an insert so a later save corrects a stored value. -INSERT INTO - ai_bedrock_inference_profile_models (inference_profile_arn, resolved_model) -VALUES - (@inference_profile_arn::text, @resolved_model::text) -ON CONFLICT (inference_profile_arn) DO UPDATE SET - resolved_model = @resolved_model::text; - --- name: GetAIBedrockInferenceProfileModels :many -SELECT - * -FROM - ai_bedrock_inference_profile_models -WHERE - inference_profile_arn = ANY(@inference_profile_arns::text[]); diff --git a/codersdk/aiproviders_bedrock.go b/codersdk/aiproviders_bedrock.go index 2a25093db7c..a4d87d9f53f 100644 --- a/codersdk/aiproviders_bedrock.go +++ b/codersdk/aiproviders_bedrock.go @@ -61,16 +61,12 @@ type AIProviderBedrockSettings struct { // AIProviderBedrockProtocolInvokeModel, so existing rows keep the legacy // behavior. Protocol AIProviderBedrockProtocol `json:"protocol,omitempty"` - // ResolvedModel and ResolvedSmallFastModel carry the model IDs behind the + // ResolvedModel and ResolvedSmallFastModel are the model IDs behind the // configured identifiers, which differ from them only for application - // inference profile ARNs. coderd resolves those when the provider is - // written and stores them in ai_provider_bedrock_resolved_models. - // - // They are in-process plumbing for the gateway, not part of this type's - // wire form: the API neither accepts nor returns them, and they are never - // stored in the settings blob. - ResolvedModel string `json:"-"` - ResolvedSmallFastModel string `json:"-"` + // inference profile ARNs. The server resolves those through AWS when the + // provider is written and owns the values; a client cannot set them. + ResolvedModel string `json:"resolved_model,omitempty"` + ResolvedSmallFastModel string `json:"resolved_small_fast_model,omitempty"` } // ResolvedProtocol returns the configured protocol, mapping the empty value to diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 1624e90af5b..399e485be0d 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -482,6 +482,14 @@ export interface AIProviderBedrockSettings { * behavior. */ readonly protocol?: AIProviderBedrockProtocol; + /** + * ResolvedModel and ResolvedSmallFastModel are the model IDs behind the + * configured identifiers, which differ from them only for application + * inference profile ARNs. The server resolves those through AWS when the + * provider is written and owns the values; a client cannot set them. + */ + readonly resolved_model?: string; + readonly resolved_small_fast_model?: string; } // From codersdk/aiproviders_bedrock.go From 7a5d01bcc3195d9defd58239fc90fe150af3406f Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 15:28:06 +0000 Subject: [PATCH 19/25] chore(coderd/database): drop generated remnants of the resolved models table --- coderd/database/dump.sql | 10 ---------- coderd/database/models.go | 6 ------ coderd/database/unique_constraint.go | 1 - 3 files changed, 17 deletions(-) diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 03bb1dc1397..6e93c540739 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1567,13 +1567,6 @@ $$; COMMENT ON FUNCTION update_chat_history_after_message_update() IS 'Component of chatd. Updates history_version and generation_attempt on chats when chat_messages is updated. Excludes changes to search_tsv and search_tsv_config.'; -CREATE TABLE ai_bedrock_inference_profile_models ( - inference_profile_arn text NOT NULL, - resolved_model text NOT NULL -); - -COMMENT ON COLUMN ai_bedrock_inference_profile_models.resolved_model IS 'The Bedrock model ID the inference profile wraps.'; - CREATE TABLE ai_gateway_keys ( id uuid NOT NULL, created_at timestamp with time zone NOT NULL, @@ -4373,9 +4366,6 @@ ALTER TABLE ONLY workspace_resource_metadata ALTER COLUMN id SET DEFAULT nextval ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); -ALTER TABLE ONLY ai_bedrock_inference_profile_models - ADD CONSTRAINT ai_bedrock_inference_profile_models_pkey PRIMARY KEY (inference_profile_arn); - ALTER TABLE ONLY ai_gateway_keys ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id); diff --git a/coderd/database/models.go b/coderd/database/models.go index a011a3bd441..b5a1340b353 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4850,12 +4850,6 @@ func AllWorkspaceTransitionValues() []WorkspaceTransition { } } -type AIBedrockInferenceProfileModel struct { - InferenceProfileArn string `db:"inference_profile_arn" json:"inference_profile_arn"` - // The Bedrock model ID the inference profile wraps. - ResolvedModel string `db:"resolved_model" json:"resolved_model"` -} - // Audit log of requests intercepted by AI Bridge type AIBridgeInterception struct { ID uuid.UUID `db:"id" json:"id"` diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 8c849a6877c..9d2390a6468 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -7,7 +7,6 @@ type UniqueConstraint string // UniqueConstraint enums. const ( UniqueAgentStatsPkey UniqueConstraint = "agent_stats_pkey" // ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); - UniqueAIBedrockInferenceProfileModelsPkey UniqueConstraint = "ai_bedrock_inference_profile_models_pkey" // ALTER TABLE ONLY ai_bedrock_inference_profile_models ADD CONSTRAINT ai_bedrock_inference_profile_models_pkey PRIMARY KEY (inference_profile_arn); UniqueAIGatewayKeysPkey UniqueConstraint = "ai_gateway_keys_pkey" // ALTER TABLE ONLY ai_gateway_keys ADD CONSTRAINT ai_gateway_keys_pkey PRIMARY KEY (id); UniqueAIModelPricesPkey UniqueConstraint = "ai_model_prices_pkey" // ALTER TABLE ONLY ai_model_prices ADD CONSTRAINT ai_model_prices_pkey PRIMARY KEY (provider, model, source); UniqueAIProviderKeysPkey UniqueConstraint = "ai_provider_keys_pkey" // ALTER TABLE ONLY ai_provider_keys ADD CONSTRAINT ai_provider_keys_pkey PRIMARY KEY (id); From 11ba4738d1e0f40087274aa156c69362da6384bc Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 12:51:46 -0400 Subject: [PATCH 20/25] refactor(coderd): resolve bedrock profiles before the provider write (#29179) Follow-up to #29112, which writes the provider first and stores the resolution in a second update. That leaves a window where a provider exists with an unresolved profile ARN, which then has to be explained everywhere: the gateway serves the ARN as its own identity, the audit entry records a pre-resolution row, and the publish is skipped so the provider only reaches the gateways after a restart. Resolving before the write removes the state instead of describing it. A failed lookup now rejects the save and stores nothing, on create and on update alike. Create resolves its own request, since a create carries the complete configuration. Update cannot: a `PATCH` supplies the model identifiers but inherits credentials and the external ID from the stored row, so the config to resolve is stored + patch. That merge now lives in `lookupAndMergeSettings` and is called twice, once before the transaction to resolve against and once inside it against the row being written. The two cannot disagree on the model identifiers, because both take them from the patch, so no reconciliation or conflict handling is needed. The server-owned STS external ID is still generated inside the transaction, as before. A value generated for the preview would be one the operator's role trust policy could not reference yet either way, so resolution assumes the role with the stored value. Net effect against the base branch: the second `UpdateAIProvider` is gone, `clearBedrockModelResolution` is gone (the resolution result is always assigned, which discards anything a client sent), and the update transaction is shorter than before this PR series started. Relates to https://linear.app/codercom/issue/AIGOV-488 Created by Coder Agents on behalf of @evgeniy-scherbina. --- coderd/ai_providers.go | 79 +++++++++++++++++------------ coderd/ai_providers_bedrock.go | 65 +++++++----------------- coderd/ai_providers_bedrock_test.go | 55 ++++++++++++++++++-- 3 files changed, 114 insertions(+), 85 deletions(-) diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index c462a95c133..04c2cd4f5fe 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -185,7 +185,16 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { // Generate the server-owned external ID when the provider assumes a role. ensureBedrockExternalID(&req.Settings) - clearBedrockModelResolution(&req.Settings) + + // Resolve application inference profile ARNs before storing them, so an + // unresolvable profile is never written and the gateway never calls the + // Bedrock control plane. + resolved, err := resolveBedrockProfiles(ctx, req.Settings) + if err != nil { + api.writeAIProviderResolutionError(ctx, rw, err) + return + } + applyBedrockResolution(&req.Settings, resolved) settings, err := encodeAIProviderSettings(req.Settings) if err != nil { @@ -245,15 +254,6 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { } aReq.New = row - // Resolve inference profile ARNs once the provider is stored, then announce - // it. The gateway never calls the Bedrock control plane itself. - row, err = api.resolveBedrockModels(ctx, row) - if err != nil { - api.writeAIProviderResolutionError(ctx, rw, err) - return - } - aReq.New = row - auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, aiProviderKeyChanges{Added: keys}) api.publishAIProvidersChanged(ctx) @@ -319,33 +319,40 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { idOrName := chi.URLParam(r, "idOrName") + // Resolve outside the transaction, because it calls AWS. The merge is + // redone inside against the row that gets written; both merges take the + // model identifiers from the patch, so they cannot disagree on them. + var resolved map[string]string + if req.Settings != nil { + _, preview, err := lookupAndMergeSettings(ctx, api.Database, idOrName, req.Settings) + if err != nil { + writeAIProviderError(ctx, api.Logger, rw, err, "update AI provider", "Internal error updating AI provider.") + return + } + resolved, err = resolveBedrockProfiles(ctx, preview) + if err != nil { + api.writeAIProviderResolutionError(ctx, rw, err) + return + } + } + var ( updated database.AIProvider keys []database.AIProviderKey keyChanges aiProviderKeyChanges ) err := api.Database.InTx(func(tx database.Store) error { - old, err := lookupAIProvider(ctx, tx, idOrName) + old, existing, err := lookupAndMergeSettings(ctx, tx, idOrName, req.Settings) if err != nil { return err } aReq.Old = old - // Decode the existing settings to merge with the patch. The dbcrypt - // wrapper has already decrypted the blob for us. - existing, err := db2sdk.AIProviderSettings(old.Settings) - if err != nil { - return xerrors.Errorf("decode existing settings: %w", err) - } if req.Settings != nil { if err := validateBedrockExternalIDUnchanged(existing, *req.Settings); err != nil { return err } - existing = mergeAIProviderSettings(existing, *req.Settings) - // The patch may point the provider at different identifiers, and a - // client cannot supply resolutions of its own. Resolution runs again - // after the transaction. - clearBedrockModelResolution(&existing) + applyBedrockResolution(&existing, resolved) } // Bedrock settings are only meaningful for anthropic- or // bedrock-typed providers; rejecting the mismatch keeps a @@ -356,8 +363,6 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { old.Type != database.AIProviderTypeBedrock { return errAIProviderBedrockTypeMismatch } - // Generate the server-owned external ID when the provider assumes a role - // and lacks one. ensureBedrockExternalID(&existing) settings, err := encodeAIProviderSettings(existing) if err != nil { @@ -455,14 +460,6 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { // An update that carries no settings cannot change the configured // identifiers or the credentials they resolve under, so any stored // resolution still holds. - if req.Settings != nil { - updated, err = api.resolveBedrockModels(ctx, updated) - if err != nil { - api.writeAIProviderResolutionError(ctx, rw, err) - return - } - aReq.New = updated - } auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, keyChanges) api.publishAIProvidersChanged(ctx) @@ -868,6 +865,24 @@ func encodeAIProviderSettings(s codersdk.AIProviderSettings) (sql.NullString, er return sql.NullString{String: string(out), Valid: true}, nil } +// lookupAndMergeSettings loads a provider and merges patch onto its stored +// settings. +func lookupAndMergeSettings(ctx context.Context, db database.Store, idOrName string, patch *codersdk.AIProviderSettings) (database.AIProvider, codersdk.AIProviderSettings, error) { + old, err := lookupAIProvider(ctx, db, idOrName) + if err != nil { + return database.AIProvider{}, codersdk.AIProviderSettings{}, err + } + // The dbcrypt wrapper has already decrypted the blob for us. + settings, err := db2sdk.AIProviderSettings(old.Settings) + if err != nil { + return database.AIProvider{}, codersdk.AIProviderSettings{}, xerrors.Errorf("decode existing settings: %w", err) + } + if patch != nil { + settings = mergeAIProviderSettings(settings, *patch) + } + return old, settings, nil +} + // mergeAIProviderSettings overlays a patch onto an existing settings // value. Write-only fields (Bedrock AccessKey and AccessKeySecret) use // pointers so the patch can distinguish "omitted, keep existing" (nil) diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index 42db6eab558..f57770d683e 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -2,7 +2,6 @@ package coderd import ( "context" - "database/sql" "net/http" "golang.org/x/xerrors" @@ -10,71 +9,41 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/provider" agplaibridge "github.com/coder/coder/v2/coderd/aibridge" - "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/codersdk" ) -// resolveBedrockModels stores the model each of the provider's application -// inference profile ARNs refers to, returning the updated provider. It runs -// after the write commits, and on every save, because it calls AWS. -func (api *API) resolveBedrockModels(ctx context.Context, row database.AIProvider) (database.AIProvider, error) { - settings, err := db2sdk.AIProviderSettings(row.Settings) - if err != nil { - return row, xerrors.Errorf("decode settings: %w", err) - } +// resolveBedrockProfiles asks AWS which model each application inference +// profile ARN in settings refers to. The result is empty when no identifier is +// an ARN, which costs no AWS call. +func resolveBedrockProfiles(ctx context.Context, settings codersdk.AIProviderSettings) (map[string]string, error) { + resolved := map[string]string{} // BaseURL is the runtime endpoint; resolution calls the control plane. cfg := agplaibridge.BedrockConfig("", settings.Bedrock) if cfg == nil { - return row, nil + return resolved, nil } - resolved, err := provider.ResolveBedrockModels(ctx, *cfg) if err != nil { - return row, xerrors.Errorf("resolve bedrock inference profile: %w", err) - } - if len(resolved) == 0 { - return row, nil - } - settings.Bedrock.ResolvedModel = resolved[settings.Bedrock.Model] - settings.Bedrock.ResolvedSmallFastModel = resolved[settings.Bedrock.SmallFastModel] - - encoded, err := encodeAIProviderSettings(settings) - if err != nil { - return row, xerrors.Errorf("encode settings: %w", err) + return nil, xerrors.Errorf("resolve bedrock inference profile: %w", err) } - updated, err := api.Database.UpdateAIProvider(ctx, database.UpdateAIProviderParams{ - ID: row.ID, - Type: row.Type, - DisplayName: row.DisplayName, - Icon: row.Icon, - Enabled: row.Enabled, - BaseUrl: row.BaseUrl, - Settings: encoded, - // SettingsKeyID is set by the dbcrypt wrapper. - SettingsKeyID: sql.NullString{}, - }) - if err != nil { - return row, xerrors.Errorf("store resolved models: %w", err) - } - return updated, nil + return resolved, nil } -// clearBedrockModelResolution drops resolved identifiers a client supplied or -// an earlier save stored. The values are server-owned and rewritten after the -// write, so anything present beforehand is stale or forged. -func clearBedrockModelResolution(settings *codersdk.AIProviderSettings) { +// applyBedrockResolution records what the configured identifiers refer to. An +// identifier that is not an application inference profile ARN is its own +// identity and stores nothing, which also discards any value a client supplied. +func applyBedrockResolution(settings *codersdk.AIProviderSettings, resolved map[string]string) { if settings.Bedrock == nil { return } - settings.Bedrock.ResolvedModel = "" - settings.Bedrock.ResolvedSmallFastModel = "" + settings.Bedrock.ResolvedModel = resolved[settings.Bedrock.Model] + settings.Bedrock.ResolvedSmallFastModel = resolved[settings.Bedrock.SmallFastModel] } -// writeAIProviderResolutionError reports a failed resolution. The provider is -// stored either way, and serves the ARN as its own identity until a later save -// resolves it. +// writeAIProviderResolutionError reports a failed resolution. The write is +// rejected, because a stored ARN with no resolution would be served as its own +// identity and misshape every request made through it. func (api *API) writeAIProviderResolutionError(ctx context.Context, rw http.ResponseWriter, err error) { api.Logger.Warn(ctx, "resolve bedrock inference profile", slog.Error(err)) httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go index 7af50bb9f7a..8af28f74bc9 100644 --- a/coderd/ai_providers_bedrock_test.go +++ b/coderd/ai_providers_bedrock_test.go @@ -151,7 +151,7 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Empty(t, paths()) }) - t.Run("CreateReportsUnresolvableProfile", func(t *testing.T) { + t.Run("CreateRejectsUnresolvableProfile", func(t *testing.T) { url, _ := mockBedrock(t, func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") @@ -177,13 +177,58 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) require.Contains(t, sdkErr.Detail, "GetInferenceProfile") - // The provider is stored with the ARN the operator asked for, and - // serves it as its own identity until a later save resolves it. + // The write is rejected: a stored ARN with no resolution would be + // served as its own identity. //nolint:gocritic // Owner role is the audience for this endpoint. providers, err := client.AIProviders(ctx) require.NoError(t, err) - require.Len(t, providers, 1) - require.Empty(t, providers[0].Settings.Bedrock.ResolvedModel) + require.Empty(t, providers) + }) + + t.Run("UpdateRejectsUnresolvableProfile", func(t *testing.T) { + var deny bool + url, _ := mockBedrock(t, func(w http.ResponseWriter, r *http.Request) { + if !deny { + respondWithModel("arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-opus-4-8")(w, r) + return + } + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Amzn-Errortype", "AccessDeniedException") + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"message":"not authorized to perform bedrock:GetInferenceProfile"}`)) + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + client := coderdtest.New(t, nil) + _ = coderdtest.CreateFirstUser(t, client) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := client.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-update-denied", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + + deny = true + //nolint:gocritic // Owner role is the audience for this endpoint. + _, err = client.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + Settings: bedrockSettings(testSmallFastProfileARN, "anthropic.claude-haiku-4-5"), + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode()) + + // The stored provider still describes what it did before the failed + // update. + //nolint:gocritic // Owner role is the audience for this endpoint. + current, err := client.AIProvider(ctx, created.ID.String()) + require.NoError(t, err) + require.Equal(t, testProfileARN, current.Settings.Bedrock.Model) + require.Equal(t, "anthropic.claude-opus-4-8", current.Settings.Bedrock.ResolvedModel) }) t.Run("UpdateReresolvesChangedProfile", func(t *testing.T) { From 1e1a54d9c9d00829659d38ecdc9210cfbc874109 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 17:14:43 +0000 Subject: [PATCH 21/25] docs(coderd): drop stale comment in the provider update handler --- coderd/ai_providers.go | 4 ---- 1 file changed, 4 deletions(-) diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index 04c2cd4f5fe..89dc2bfa8aa 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -457,10 +457,6 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { return } - // An update that carries no settings cannot change the configured - // identifiers or the credentials they resolve under, so any stored - // resolution still holds. - auditAIProviderKeyChanges(ctx, r, *auditor, api.Logger, keyChanges) api.publishAIProvidersChanged(ctx) From a417d6a534412393aac6f52a3a69e32473728c27 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 18:28:50 +0000 Subject: [PATCH 22/25] fix(coderd): authorize AI provider writes before resolving bedrock profiles --- coderd/ai_providers.go | 17 ++++++++++++ coderd/ai_providers_bedrock_test.go | 40 +++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index 89dc2bfa8aa..f83f7ef69bf 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -24,6 +24,8 @@ import ( "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/coderd/httpmw" coderpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/codersdk" ) @@ -160,6 +162,14 @@ func (api *API) aiProvidersCreate(rw http.ResponseWriter, r *http.Request) { ) defer commitAudit() + // Provider configuration has side effects outside the database, notably + // the Bedrock profile lookup below, so the permission is checked before + // any of them rather than only by dbauthz on the write. + if !api.Authorize(r, policy.ActionCreate, rbac.ResourceAIProvider) { + httpapi.Forbidden(rw) + return + } + var req codersdk.CreateAIProviderRequest if !httpapi.Read(ctx, rw, r, &req) { return @@ -298,6 +308,13 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { ) defer commitAudit() + // Matches the create path: the Bedrock profile lookup below runs before + // dbauthz sees the write, so gate on the permission first. + if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceAIProvider) { + httpapi.Forbidden(rw) + return + } + var req codersdk.UpdateAIProviderRequest if !httpapi.Read(ctx, rw, r, &req) { return diff --git a/coderd/ai_providers_bedrock_test.go b/coderd/ai_providers_bedrock_test.go index 8af28f74bc9..7beb168a3ca 100644 --- a/coderd/ai_providers_bedrock_test.go +++ b/coderd/ai_providers_bedrock_test.go @@ -317,4 +317,44 @@ func TestAIProvidersBedrockProfileResolution(t *testing.T) { require.Equal(t, "anthropic.claude-opus-4-8", updated.Settings.Bedrock.ResolvedModel) require.Len(t, paths(), callsAfterCreate, "an unrelated update does not call AWS") }) + + t.Run("NonOwnerCannotDriveResolution", func(t *testing.T) { + url, _ := mockBedrock(t, func(http.ResponseWriter, *http.Request) { + t.Error("Bedrock called for an unauthorized request") + }) + t.Setenv("AWS_ENDPOINT_URL_BEDROCK", url) + + ownerClient := coderdtest.New(t, nil) + firstUser := coderdtest.CreateFirstUser(t, ownerClient) + ctx := testutil.Context(t, testutil.WaitLong) + + //nolint:gocritic // Owner role is the audience for this endpoint. + created, err := ownerClient.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-owner-only", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings("eu.anthropic.claude-opus-4-8", "anthropic.claude-haiku-4-5"), + }) + require.NoError(t, err) + + memberClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, firstUser.OrganizationID) + + _, err = memberClient.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{ + Name: "bedrock-member", + Type: codersdk.AIProviderTypeBedrock, + BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com", + Enabled: true, + Settings: *bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + + _, err = memberClient.UpdateAIProvider(ctx, created.ID.String(), codersdk.UpdateAIProviderRequest{ + Settings: bedrockSettings(testProfileARN, "anthropic.claude-haiku-4-5"), + }) + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusForbidden, sdkErr.StatusCode()) + }) } From b5d4195b474266ca2663562b05a92fa0ed7adc08 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 19:51:12 +0000 Subject: [PATCH 23/25] docs(coderd): clarify bedrock resolution endpoint --- coderd/ai_providers_bedrock.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/coderd/ai_providers_bedrock.go b/coderd/ai_providers_bedrock.go index f57770d683e..a4f3373418c 100644 --- a/coderd/ai_providers_bedrock.go +++ b/coderd/ai_providers_bedrock.go @@ -18,7 +18,8 @@ import ( // an ARN, which costs no AWS call. func resolveBedrockProfiles(ctx context.Context, settings codersdk.AIProviderSettings) (map[string]string, error) { resolved := map[string]string{} - // BaseURL is the runtime endpoint; resolution calls the control plane. + // BaseURL configures the runtime data-plane endpoint. GetInferenceProfile + // uses the Bedrock control-plane endpoint derived from Region instead. cfg := agplaibridge.BedrockConfig("", settings.Bedrock) if cfg == nil { return resolved, nil From a291c8084cdc7cc216fbf3bf8550b08c15ecc292 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Fri, 11 Sep 2026 15:21:05 +0000 Subject: [PATCH 24/25] refactor(cli): drop the bedrock config wrapper --- cli/aibridged.go | 8 +------- cli/aibridged_internal_test.go | 8 ++++---- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/cli/aibridged.go b/cli/aibridged.go index 030fbc00e65..9f64aee9850 100644 --- a/cli/aibridged.go +++ b/cli/aibridged.go @@ -284,7 +284,7 @@ func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBrid }), nil case database.AIProviderTypeAnthropic, database.AIProviderTypeBedrock: - bedrock := bedrockConfig(spec.BaseURL, spec.Bedrock) + bedrock := agplaibridge.BedrockConfig(spec.BaseURL, spec.Bedrock) // A spec typed 'bedrock' authenticates exclusively via settings; // without populated Bedrock credentials it cannot make upstream // calls, so refuse rather than falling back to an unsigned @@ -336,12 +336,6 @@ func buildAIProviderKeyPool(providerName string, keys []string, metrics *aibridg return keypool.New(providerName, keys, quartz.NewReal(), metrics) } -// bedrockConfig is [agplaibridge.BedrockConfig], shared with the provider -// write path so both map stored settings the same way. -func bedrockConfig(baseURL string, bedrock *codersdk.AIProviderBedrockSettings) *aibridge.AWSBedrockConfig { - return agplaibridge.BedrockConfig(baseURL, bedrock) -} - // circuitBreakerConfig returns nil when the breaker is disabled. func circuitBreakerConfig(cfg codersdk.AIBridgeConfig) *config.CircuitBreaker { if !cfg.CircuitBreakerEnabled.Value() { diff --git a/cli/aibridged_internal_test.go b/cli/aibridged_internal_test.go index 7cd4f64d742..13a4a172178 100644 --- a/cli/aibridged_internal_test.go +++ b/cli/aibridged_internal_test.go @@ -280,7 +280,7 @@ func TestBuildProviders(t *testing.T) { Name: aibridge.ProviderAnthropic, BaseUrl: "https://api.anthropic.com/", } - assert.Nil(t, bedrockConfig(row.BaseUrl, codersdk.AIProviderSettings{}.Bedrock)) + assert.Nil(t, agplaibridge.BedrockConfig(row.BaseUrl, codersdk.AIProviderSettings{}.Bedrock)) }) t.Run("NativeAnthropicCustomBaseURL", func(t *testing.T) { @@ -290,7 +290,7 @@ func TestBuildProviders(t *testing.T) { Name: "anthropic-proxy", BaseUrl: "https://internal-proxy.example.com/anthropic/", } - assert.Nil(t, bedrockConfig(row.BaseUrl, codersdk.AIProviderSettings{}.Bedrock)) + assert.Nil(t, agplaibridge.BedrockConfig(row.BaseUrl, codersdk.AIProviderSettings{}.Bedrock)) }) t.Run("BedrockSettingsPresent", func(t *testing.T) { @@ -315,7 +315,7 @@ func TestBuildProviders(t *testing.T) { RoleARN: roleARN, }, } - got := bedrockConfig(row.BaseUrl, settings.Bedrock) + got := agplaibridge.BedrockConfig(row.BaseUrl, settings.Bedrock) require.NotNil(t, got) assert.Equal(t, row.BaseUrl, got.BaseURL) assert.Equal(t, "us-west-2", got.Region) @@ -339,7 +339,7 @@ func TestBuildProviders(t *testing.T) { settings := codersdk.AIProviderSettings{ Bedrock: &codersdk.AIProviderBedrockSettings{}, } - assert.Nil(t, bedrockConfig(row.BaseUrl, settings.Bedrock)) + assert.Nil(t, agplaibridge.BedrockConfig(row.BaseUrl, settings.Bedrock)) }) } From 9d5979b0a1d4457c25634b5e7b91c584fd49b081 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Fri, 11 Sep 2026 15:38:59 +0000 Subject: [PATCH 25/25] refactor(coderd): rename the merged settings variable in provider updates --- coderd/ai_providers.go | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/coderd/ai_providers.go b/coderd/ai_providers.go index f83f7ef69bf..93b7908b270 100644 --- a/coderd/ai_providers.go +++ b/coderd/ai_providers.go @@ -359,36 +359,36 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) { keyChanges aiProviderKeyChanges ) err := api.Database.InTx(func(tx database.Store) error { - old, existing, err := lookupAndMergeSettings(ctx, tx, idOrName, req.Settings) + old, merged, err := lookupAndMergeSettings(ctx, tx, idOrName, req.Settings) if err != nil { return err } aReq.Old = old if req.Settings != nil { - if err := validateBedrockExternalIDUnchanged(existing, *req.Settings); err != nil { + if err := validateBedrockExternalIDUnchanged(merged, *req.Settings); err != nil { return err } - applyBedrockResolution(&existing, resolved) + applyBedrockResolution(&merged, resolved) } // Bedrock settings are only meaningful for anthropic- or // bedrock-typed providers; rejecting the mismatch keeps a // misconfiguration from sitting silently in the encrypted // blob. - if existing.Bedrock != nil && + if merged.Bedrock != nil && old.Type != database.AIProviderTypeAnthropic && old.Type != database.AIProviderTypeBedrock { return errAIProviderBedrockTypeMismatch } - ensureBedrockExternalID(&existing) - settings, err := encodeAIProviderSettings(existing) + ensureBedrockExternalID(&merged) + settings, err := encodeAIProviderSettings(merged) if err != nil { return xerrors.Errorf("encode settings: %w", err) } // Reject keys against Bedrock providers (whether the existing // row is Bedrock or the patch would make it so). - if req.APIKeys != nil && existing.Bedrock != nil && len(*req.APIKeys) > 0 { + if req.APIKeys != nil && merged.Bedrock != nil && len(*req.APIKeys) > 0 { return errBedrockRejectsAPIKeys }