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) {