From 05b727779843f4f30615ef55bd9697b0bceda140 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 10 Sep 2026 12:41:26 +0000 Subject: [PATCH] refactor(coderd): store bedrock model resolution in provider settings --- 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