Thanks to visit codestin.com
Credit goes to github.com

Skip to content

refactor(coderd): store bedrock model resolution in provider settings #29175

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
evgeniy-scherbina merged 1 commit into from
Sep 10, 2026
Merged

refactor(coderd): store bedrock model resolution in provider settings #29175

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions coderd/ai_providers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
56 changes: 41 additions & 15 deletions coderd/ai_providers_bedrock.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package coderd

import (
"context"
"database/sql"
"net/http"

"golang.org/x/xerrors"
Expand All @@ -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
Expand Down
121 changes: 50 additions & 71 deletions coderd/ai_providers_bedrock_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package coderd_test

import (
"context"
"net/http"
"net/http/httptest"
"slices"
Expand All @@ -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"
)
Expand Down Expand Up @@ -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.
Expand All @@ -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)

Expand All @@ -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) {
Expand All @@ -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)

Expand All @@ -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")
Expand All @@ -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)

Expand All @@ -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)

Expand All @@ -197,23 +203,23 @@ 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{
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", 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)

Expand All @@ -226,23 +232,23 @@ 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")
})

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)

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)

Expand All @@ -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))
})
}
Loading
Loading