mirror of
https://github.com/coder/coder.git
synced 2026-09-01 14:53:15 +08:00
fix: backfill legacy Bedrock AI provider rows and stale model config strings (#26155)
Fixes CODAGT-548
Adds two idempotent startup backfills run after `newAPI():
- `BackfillBedrockProviderType`: promotes `ai_providers` rows from
`type=anthropic` with Bedrock settings to `type=bedrock`.
- `BackfillChatModelConfigProviderStrings`: fixes stale
`chat_model_configs.provider = "anthropic"` strings on rows whose linked
provider was just promoted.
- `UpdateAIProvider` query now also writes the `type` column, so the
fix persists on any subsequent PATCH.
> 🤖 Generated by Claude with oversight from a human.
This commit is contained in:
@@ -1048,6 +1048,11 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
|
||||
); err != nil {
|
||||
return xerrors.Errorf("seed ai providers from env: %w", err)
|
||||
}
|
||||
// Must run after newAPI so options.Database is dbcrypt-wrapped.
|
||||
coderd.BackfillBedrockProviderType(aibridgeInitCtx, options.Database, logger.Named("aibridge.backfill"))
|
||||
// Must run after BackfillBedrockProviderType; shares aibridgeInitCtx so
|
||||
// a timeout on the first backfill will skip this one until next startup.
|
||||
coderd.BackfillChatModelConfigProviderStrings(aibridgeInitCtx, options.Database, logger.Named("aibridge.backfill"))
|
||||
|
||||
// In-memory aibridge daemon. Registered on coderd so chatd can
|
||||
// dispatch LLM requests via the in-process transport without
|
||||
|
||||
@@ -351,6 +351,7 @@ func (api *API) aiProvidersUpdate(rw http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
params := database.UpdateAIProviderParams{
|
||||
ID: old.ID,
|
||||
Type: old.Type,
|
||||
DisplayName: displayName,
|
||||
Enabled: ptr.NilToDefault(req.Enabled, old.Enabled),
|
||||
BaseUrl: ptr.NilToDefault(req.BaseURL, old.BaseUrl),
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
package coderd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/db2sdk"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
// BackfillBedrockProviderType promotes legacy ai_providers rows stored as
|
||||
// type=anthropic with Bedrock settings to type=bedrock. Must run after newAPI
|
||||
// so options.Database is dbcrypt-wrapped. Idempotent; errors are logged and
|
||||
// startup continues.
|
||||
//
|
||||
// BackfillChatModelConfigProviderStrings must run after this function so
|
||||
// provider types are correct when its JOIN executes.
|
||||
func BackfillBedrockProviderType(ctx context.Context, db database.Store, logger slog.Logger) {
|
||||
//nolint:gocritic // Startup-only backfill; no user actor is present.
|
||||
sysCtx := dbauthz.AsSystemRestricted(ctx)
|
||||
providers, err := db.GetAIProviders(sysCtx, database.GetAIProvidersParams{
|
||||
IncludeDeleted: false,
|
||||
IncludeDisabled: true,
|
||||
})
|
||||
if err != nil {
|
||||
logger.Error(ctx, "backfill bedrock provider type: list providers", slog.Error(err))
|
||||
return
|
||||
}
|
||||
var promoted int
|
||||
for _, provider := range providers {
|
||||
if provider.Type != database.AiProviderTypeAnthropic {
|
||||
continue
|
||||
}
|
||||
settings, err := db2sdk.AIProviderSettings(provider.Settings)
|
||||
if err != nil {
|
||||
logger.Warn(ctx, "backfill bedrock provider type: skip provider with unparsable settings",
|
||||
slog.F("provider_id", provider.ID), slog.Error(err))
|
||||
continue
|
||||
}
|
||||
if settings.Bedrock == nil {
|
||||
continue
|
||||
}
|
||||
_, err = db.UpdateAIProvider(sysCtx, database.UpdateAIProviderParams{
|
||||
ID: provider.ID,
|
||||
Type: database.AiProviderTypeBedrock,
|
||||
DisplayName: provider.DisplayName,
|
||||
Enabled: provider.Enabled,
|
||||
BaseUrl: provider.BaseUrl,
|
||||
Settings: provider.Settings,
|
||||
// SettingsKeyID is re-set by the dbcrypt wrapper on write.
|
||||
SettingsKeyID: sql.NullString{},
|
||||
})
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
logger.Debug(ctx, "backfill bedrock provider type: provider deleted during backfill",
|
||||
slog.F("provider_id", provider.ID))
|
||||
continue
|
||||
}
|
||||
logger.Error(ctx, "backfill bedrock provider type: provider update failed and will re-attempt on next server startup",
|
||||
slog.F("provider_id", provider.ID), slog.Error(err))
|
||||
continue
|
||||
}
|
||||
promoted++
|
||||
}
|
||||
if promoted > 0 {
|
||||
logger.Info(ctx, "backfilled bedrock provider types", slog.F("count", promoted))
|
||||
}
|
||||
}
|
||||
|
||||
// BackfillChatModelConfigProviderStrings fixes stale chat_model_configs.provider
|
||||
// strings left as "anthropic" when the linked provider was promoted from
|
||||
// type=anthropic to type=bedrock by BackfillBedrockProviderType. Errors are
|
||||
// logged and startup continues.
|
||||
func BackfillChatModelConfigProviderStrings(ctx context.Context, db database.Store, logger slog.Logger) {
|
||||
//nolint:gocritic // Startup-only backfill; no user actor is present.
|
||||
sysCtx := dbauthz.AsSystemRestricted(ctx)
|
||||
result, err := db.BackfillChatModelConfigProvider(sysCtx, database.BackfillChatModelConfigProviderParams{
|
||||
OldProvider: string(codersdk.AIProviderTypeAnthropic),
|
||||
NewProvider: string(codersdk.AIProviderTypeBedrock),
|
||||
})
|
||||
if err != nil {
|
||||
logger.Error(ctx, "backfill chat model config provider strings", slog.Error(err))
|
||||
return
|
||||
}
|
||||
if result != nil {
|
||||
if n, _ := result.RowsAffected(); n > 0 {
|
||||
logger.Info(ctx, "backfilled chat model config provider strings", slog.F("count", n))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,368 @@
|
||||
package coderd_test
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"github.com/coder/coder/v2/coderd"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
// TestBackfillBedrockProviderType runs all DB-backed cases against a single
|
||||
// database instance. Subtests are intentionally sequential so that each one
|
||||
// builds on the state left by the previous, which proves idempotency without
|
||||
// extra setup: a second backfill call on an already-promoted DB must be a
|
||||
// no-op. Failure-path tests use a mock and stay parallel.
|
||||
func TestBackfillBedrockProviderType(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
bedrockSettings := sql.NullString{
|
||||
String: `{"_type":"bedrock","_version":1,"region":"us-east-1"}`,
|
||||
Valid: true,
|
||||
}
|
||||
|
||||
// All DB subtests share one database instance and run sequentially.
|
||||
t.Run("DB", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||
logger := testLogger(t)
|
||||
|
||||
t.Run("NoLegacyRows", func(t *testing.T) {
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
all, err := db.GetAIProviders(ctx, database.GetAIProvidersParams{
|
||||
IncludeDeleted: true,
|
||||
IncludeDisabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, all)
|
||||
})
|
||||
|
||||
t.Run("PromotesLegacyRow", func(t *testing.T) {
|
||||
legacy := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, legacy.Type, "pre-condition: row must start as anthropic")
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
row, err := db.GetAIProviderByName(ctx, legacy.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeBedrock, row.Type)
|
||||
})
|
||||
|
||||
t.Run("Idempotent", func(t *testing.T) {
|
||||
// DB already has one bedrock row from the previous subtest.
|
||||
// A second run must be a no-op: no type changes, no new rows.
|
||||
before, err := db.GetAIProviders(ctx, database.GetAIProvidersParams{
|
||||
IncludeDeleted: true,
|
||||
IncludeDisabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
for _, r := range before {
|
||||
require.Equal(t, database.AiProviderTypeBedrock, r.Type,
|
||||
"pre-condition: all rows must already be promoted before testing idempotency")
|
||||
}
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
after, err := db.GetAIProviders(ctx, database.GetAIProvidersParams{
|
||||
IncludeDeleted: true,
|
||||
IncludeDisabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, len(before), len(after), "second run must not create rows")
|
||||
for i := range after {
|
||||
require.Equal(t, before[i].Type, after[i].Type, "second run must not change types")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("PreservesNativeAnthropicRow", func(t *testing.T) {
|
||||
native := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, native.Type, "pre-condition")
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
row, err := db.GetAIProviderByName(ctx, native.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, row.Type)
|
||||
})
|
||||
|
||||
t.Run("PreservesNativeBedrockRow", func(t *testing.T) {
|
||||
native := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeBedrock,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeBedrock, native.Type, "pre-condition")
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
row, err := db.GetAIProviderByName(ctx, native.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeBedrock, row.Type)
|
||||
})
|
||||
|
||||
t.Run("SkipsDeletedRows", func(t *testing.T) {
|
||||
deleted := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, deleted.Type, "pre-condition")
|
||||
require.NoError(t, db.DeleteAIProviderByID(ctx, deleted.ID))
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
row, err := db.GetAIProviders(ctx, database.GetAIProvidersParams{
|
||||
IncludeDeleted: true,
|
||||
IncludeDisabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
var found bool
|
||||
for _, r := range row {
|
||||
if r.ID == deleted.ID {
|
||||
found = true
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, r.Type, "deleted row must not be promoted")
|
||||
}
|
||||
}
|
||||
require.True(t, found, "deleted row must appear in IncludeDeleted result set")
|
||||
})
|
||||
|
||||
t.Run("IncludesDisabledRows", func(t *testing.T) {
|
||||
disabled := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Enabled: false,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, disabled.Type, "pre-condition")
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
row, err := db.GetAIProviderByName(ctx, disabled.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeBedrock, row.Type, "disabled legacy row must be promoted")
|
||||
})
|
||||
|
||||
t.Run("PreservesAnthropicRowWithNonBedrockSettings", func(t *testing.T) {
|
||||
// {} has no _type discriminator, so UnmarshalJSON returns an error
|
||||
// and the row is skipped via the unparsable-settings path, not the
|
||||
// settings.Bedrock == nil guard. Either way the row must stay anthropic.
|
||||
nonBedrock := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: sql.NullString{String: "{}", Valid: true},
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, nonBedrock.Type, "pre-condition")
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
row, err := db.GetAIProviderByName(ctx, nonBedrock.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, row.Type, "anthropic row with non-bedrock settings must not be promoted")
|
||||
})
|
||||
|
||||
t.Run("SkipsUnparsableSettings", func(t *testing.T) {
|
||||
malformed := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: sql.NullString{String: "{", Valid: true},
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, malformed.Type, "pre-condition")
|
||||
good := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, good.Type, "pre-condition")
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, logger)
|
||||
|
||||
malformedRow, err := db.GetAIProviderByName(ctx, malformed.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, malformedRow.Type, "row with unparsable settings must not be touched")
|
||||
|
||||
goodRow, err := db.GetAIProviderByName(ctx, good.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeBedrock, goodRow.Type, "valid row alongside unparsable one must still be promoted")
|
||||
})
|
||||
|
||||
// --- chat_model_configs.provider backfill ---
|
||||
// These subtests rely on the DB already having type=bedrock providers
|
||||
// from the provider backfill subtests above.
|
||||
|
||||
t.Run("FixesStaleModelConfigProvider", func(t *testing.T) {
|
||||
// Simulate a model config created when the linked provider was still
|
||||
// type=anthropic. The stored provider string is "anthropic" but the
|
||||
// linked provider row now has type=bedrock.
|
||||
bedrockProvider := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeBedrock,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
staleConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Provider: "anthropic",
|
||||
AIProviderID: uuid.NullUUID{UUID: bedrockProvider.ID, Valid: true},
|
||||
})
|
||||
|
||||
coderd.BackfillChatModelConfigProviderStrings(ctx, db, logger)
|
||||
|
||||
updated, err := db.GetChatModelConfigByID(ctx, staleConfig.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "bedrock", updated.Provider, "stale anthropic provider string must be fixed to bedrock")
|
||||
|
||||
// Second run must be a no-op: the same config must still be "bedrock".
|
||||
coderd.BackfillChatModelConfigProviderStrings(ctx, db, logger)
|
||||
|
||||
updated, err = db.GetChatModelConfigByID(ctx, staleConfig.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "bedrock", updated.Provider, "provider must remain bedrock after second run")
|
||||
})
|
||||
|
||||
t.Run("ModelConfigIdempotent", func(t *testing.T) {
|
||||
before, err := db.GetChatModelConfigs(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
coderd.BackfillChatModelConfigProviderStrings(ctx, db, logger)
|
||||
|
||||
after, err := db.GetChatModelConfigs(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, len(before), len(after), "second run must not create or delete rows")
|
||||
})
|
||||
|
||||
t.Run("PreservesNonAnthropicModelConfig", func(t *testing.T) {
|
||||
// A model config with provider="openai" linked to a Bedrock provider
|
||||
// must not be touched. Only "anthropic" → "bedrock" is in scope.
|
||||
bedrockProvider := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeBedrock,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
openAIConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Provider: "openai",
|
||||
AIProviderID: uuid.NullUUID{UUID: bedrockProvider.ID, Valid: true},
|
||||
})
|
||||
|
||||
coderd.BackfillChatModelConfigProviderStrings(ctx, db, logger)
|
||||
|
||||
row, err := db.GetChatModelConfigByID(ctx, openAIConfig.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "openai", row.Provider, "non-anthropic provider string must not be changed")
|
||||
})
|
||||
|
||||
t.Run("SkipsModelConfigWithDeletedProvider", func(t *testing.T) {
|
||||
// Verifies the EXISTS subquery excludes soft-deleted providers.
|
||||
// The model config provider string must stay "anthropic" because
|
||||
// the linked provider is deleted and therefore excluded by the
|
||||
// AND deleted = FALSE condition in the query.
|
||||
deletedProvider := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeBedrock,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
staleConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Provider: "anthropic",
|
||||
AIProviderID: uuid.NullUUID{UUID: deletedProvider.ID, Valid: true},
|
||||
})
|
||||
require.NoError(t, db.DeleteAIProviderByID(ctx, deletedProvider.ID))
|
||||
|
||||
coderd.BackfillChatModelConfigProviderStrings(ctx, db, logger)
|
||||
|
||||
row, err := db.GetChatModelConfigByID(ctx, staleConfig.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "anthropic", row.Provider, "config linked to deleted provider must not be updated")
|
||||
})
|
||||
|
||||
t.Run("SkipsDeletedModelConfig", func(t *testing.T) {
|
||||
// The SQL query guards on deleted = FALSE. Capture the config ID
|
||||
// before deletion so we delete the right row regardless of ordering.
|
||||
bedrockProvider := dbgen.AIProvider(t, db, database.AIProvider{
|
||||
Type: database.AiProviderTypeBedrock,
|
||||
Settings: bedrockSettings,
|
||||
})
|
||||
cfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Provider: "anthropic",
|
||||
AIProviderID: uuid.NullUUID{UUID: bedrockProvider.ID, Valid: true},
|
||||
})
|
||||
|
||||
before, err := db.GetChatModelConfigs(ctx)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.DeleteChatModelConfigByID(ctx, cfg.ID))
|
||||
|
||||
coderd.BackfillChatModelConfigProviderStrings(ctx, db, logger)
|
||||
|
||||
after, err := db.GetChatModelConfigs(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, len(before)-1, len(after), "deleted config must not reappear after backfill")
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("ListFailure", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
|
||||
db.EXPECT().
|
||||
GetAIProviders(gomock.Any(), gomock.Any()).
|
||||
Return(nil, sql.ErrConnDone)
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, testLogger(t))
|
||||
})
|
||||
|
||||
t.Run("UpdateFailure", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
|
||||
db.EXPECT().
|
||||
GetAIProviders(gomock.Any(), gomock.Any()).
|
||||
Return([]database.AIProvider{{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: bedrockSettings,
|
||||
}}, nil)
|
||||
db.EXPECT().
|
||||
UpdateAIProvider(gomock.Any(), gomock.Any()).
|
||||
Return(database.AIProvider{}, sql.ErrConnDone)
|
||||
|
||||
coderd.BackfillBedrockProviderType(ctx, db, testLogger(t))
|
||||
})
|
||||
|
||||
t.Run("ProviderDeletedDuringBackfill", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
|
||||
db.EXPECT().
|
||||
GetAIProviders(gomock.Any(), gomock.Any()).
|
||||
Return([]database.AIProvider{{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: bedrockSettings,
|
||||
}}, nil)
|
||||
db.EXPECT().
|
||||
UpdateAIProvider(gomock.Any(), gomock.Any()).
|
||||
Return(database.AIProvider{}, sql.ErrNoRows)
|
||||
|
||||
// ErrNoRows is benign: provider was deleted between list and update.
|
||||
coderd.BackfillBedrockProviderType(ctx, db, testLogger(t))
|
||||
})
|
||||
|
||||
t.Run("ModelConfigQueryFailure", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
|
||||
db.EXPECT().
|
||||
BackfillChatModelConfigProvider(gomock.Any(), gomock.Any()).
|
||||
Return(nil, sql.ErrConnDone)
|
||||
|
||||
coderd.BackfillChatModelConfigProviderStrings(ctx, db, testLogger(t))
|
||||
})
|
||||
}
|
||||
@@ -126,8 +126,15 @@ func SeedAIProvidersFromEnv(
|
||||
for _, k := range existingKeyRows {
|
||||
existingKeys = append(existingKeys, k.APIKey)
|
||||
}
|
||||
// Use the canonical type so that a row promoted from
|
||||
// type=anthropic to type=bedrock by the startup backfill
|
||||
// is not mistaken for drift on the next startup.
|
||||
existingType := existing.Type
|
||||
if existingSettings.Bedrock != nil && existing.Type == database.AiProviderTypeAnthropic {
|
||||
existingType = database.AiProviderTypeBedrock
|
||||
}
|
||||
existingDP := desiredAIProvider{
|
||||
Type: existing.Type,
|
||||
Type: existingType,
|
||||
BaseURL: existing.BaseUrl,
|
||||
Bedrock: existingSettings.Bedrock,
|
||||
Keys: existingKeys,
|
||||
@@ -330,6 +337,7 @@ func providersFromEnv(ctx context.Context, cfg codersdk.AIBridgeConfig, logger s
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
}
|
||||
if hasLegacyBedrock {
|
||||
dp.Type = database.AiProviderTypeBedrock
|
||||
if hasAnthropicKey {
|
||||
logger.Warn(ctx, "ignoring legacy Anthropic API key because Bedrock credentials are configured; Bedrock authenticates via access keys or credential chain",
|
||||
slog.F("provider", aibridge.ProviderAnthropic),
|
||||
|
||||
@@ -2,6 +2,7 @@ package coderd_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"database/sql"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -145,8 +146,8 @@ func TestSeedAIProvidersFromEnv(t *testing.T) {
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
// Bedrock fields without an Anthropic key produce a Bedrock-
|
||||
// authenticated Anthropic provider with no bearer keys.
|
||||
// Bedrock fields without an Anthropic key produce a type=bedrock
|
||||
// provider named "anthropic" with no bearer keys.
|
||||
cfg := codersdk.AIBridgeConfig{
|
||||
LegacyBedrock: codersdk.AIBridgeBedrockConfig{
|
||||
Region: serpent.String("us-west-2"),
|
||||
@@ -160,7 +161,7 @@ func TestSeedAIProvidersFromEnv(t *testing.T) {
|
||||
|
||||
row, err := db.GetAIProviderByName(ctx, "anthropic")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, row.Type)
|
||||
require.Equal(t, database.AiProviderTypeBedrock, row.Type)
|
||||
require.Contains(t, row.Settings.String, "us-west-2")
|
||||
require.Contains(t, row.Settings.String, "anthropic.claude-3-5-sonnet")
|
||||
require.Contains(t, row.Settings.String, "anthropic.claude-3-5-haiku")
|
||||
@@ -600,6 +601,44 @@ func TestSeedAIProvidersFromEnv(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "conflicting fields")
|
||||
})
|
||||
|
||||
t.Run("SeedIsIdempotentAfterBedrockBackfill", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
// Regression: seed must not treat a type=anthropic row promoted to
|
||||
// type=bedrock by the backfill as drift.
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
cfg := codersdk.AIBridgeConfig{
|
||||
LegacyBedrock: codersdk.AIBridgeBedrockConfig{
|
||||
Region: serpent.String("us-east-1"),
|
||||
AccessKey: serpent.String("AKIA"),
|
||||
AccessKeySecret: serpent.String("secret"),
|
||||
Model: serpent.String("anthropic.claude-3-5-sonnet"),
|
||||
},
|
||||
}
|
||||
|
||||
// Seed to get a row with correct settings, then set type=anthropic to
|
||||
// simulate the pre-upgrade state where the old seed stored that type.
|
||||
require.NoError(t, coderd.SeedAIProvidersFromEnv(ctx, db, cfg, testLogger(t)))
|
||||
row, err := db.GetAIProviderByName(ctx, "anthropic")
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpdateAIProvider(ctx, database.UpdateAIProviderParams{
|
||||
ID: row.ID,
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
DisplayName: row.DisplayName,
|
||||
Enabled: row.Enabled,
|
||||
BaseUrl: row.BaseUrl,
|
||||
Settings: row.Settings,
|
||||
SettingsKeyID: sql.NullString{},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
row, err = db.GetAIProviderByName(ctx, "anthropic")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeAnthropic, row.Type, "pre-condition: row must be anthropic before seed runs")
|
||||
|
||||
require.NoError(t, coderd.SeedAIProvidersFromEnv(ctx, db, cfg, testLogger(t)))
|
||||
})
|
||||
}
|
||||
|
||||
func testLogger(t *testing.T) slog.Logger {
|
||||
|
||||
@@ -1701,6 +1701,13 @@ func (q *querier) AutoArchiveInactiveChats(ctx context.Context, arg database.Aut
|
||||
return q.db.AutoArchiveInactiveChats(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) BackfillChatModelConfigProvider(ctx context.Context, arg database.BackfillChatModelConfigProviderParams) (sql.Result, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.BackfillChatModelConfigProvider(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
|
||||
// This is a system-level operation used by the gitsync
|
||||
// background worker to reschedule failed refreshes. Same
|
||||
|
||||
@@ -6570,6 +6570,7 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
provider := testutil.Fake(s.T(), faker, database.AIProvider{})
|
||||
arg := database.UpdateAIProviderParams{
|
||||
ID: provider.ID,
|
||||
Type: provider.Type,
|
||||
Enabled: true,
|
||||
BaseUrl: "https://api.example.com/",
|
||||
}
|
||||
@@ -6581,6 +6582,14 @@ func (s *MethodTestSuite) TestAIBridge() {
|
||||
dbm.EXPECT().DeleteAIProviderByID(gomock.Any(), provider.ID).Return(nil).AnyTimes()
|
||||
check.Args(provider.ID).Asserts(rbac.ResourceAIProvider, policy.ActionDelete).Returns()
|
||||
}))
|
||||
s.Run("BackfillChatModelConfigProvider", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.BackfillChatModelConfigProviderParams{
|
||||
OldProvider: "anthropic",
|
||||
NewProvider: "bedrock",
|
||||
}
|
||||
dbm.EXPECT().BackfillChatModelConfigProvider(gomock.Any(), arg).Return(nil, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("UpdateEncryptedAIProviderSettings", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
provider := testutil.Fake(s.T(), faker, database.AIProvider{})
|
||||
arg := database.UpdateEncryptedAIProviderSettingsParams{
|
||||
|
||||
+9
@@ -5,6 +5,7 @@ package dbmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"slices"
|
||||
"time"
|
||||
@@ -185,6 +186,14 @@ func (m queryMetricsStore) AutoArchiveInactiveChats(ctx context.Context, arg dat
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BackfillChatModelConfigProvider(ctx context.Context, arg database.BackfillChatModelConfigProviderParams) (sql.Result, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.BackfillChatModelConfigProvider(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("BackfillChatModelConfigProvider").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BackfillChatModelConfigProvider").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.BackoffChatDiffStatus(ctx, arg)
|
||||
|
||||
Generated
+16
@@ -11,6 +11,7 @@ package dbmock
|
||||
|
||||
import (
|
||||
context "context"
|
||||
sql "database/sql"
|
||||
json "encoding/json"
|
||||
reflect "reflect"
|
||||
time "time"
|
||||
@@ -193,6 +194,21 @@ func (mr *MockStoreMockRecorder) AutoArchiveInactiveChats(ctx, arg any) *gomock.
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AutoArchiveInactiveChats", reflect.TypeOf((*MockStore)(nil).AutoArchiveInactiveChats), ctx, arg)
|
||||
}
|
||||
|
||||
// BackfillChatModelConfigProvider mocks base method.
|
||||
func (m *MockStore) BackfillChatModelConfigProvider(ctx context.Context, arg database.BackfillChatModelConfigProviderParams) (sql.Result, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "BackfillChatModelConfigProvider", ctx, arg)
|
||||
ret0, _ := ret[0].(sql.Result)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// BackfillChatModelConfigProvider indicates an expected call of BackfillChatModelConfigProvider.
|
||||
func (mr *MockStoreMockRecorder) BackfillChatModelConfigProvider(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BackfillChatModelConfigProvider", reflect.TypeOf((*MockStore)(nil).BackfillChatModelConfigProvider), ctx, arg)
|
||||
}
|
||||
|
||||
// BackoffChatDiffStatus mocks base method.
|
||||
func (m *MockStore) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+6
@@ -6,6 +6,7 @@ package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
@@ -70,6 +71,11 @@ type sqlcQuerier interface {
|
||||
// created_at ASC flows through to dbpurge's digest truncation; see
|
||||
// buildDigestData in dbpurge.go for the tradeoff rationale.
|
||||
AutoArchiveInactiveChats(ctx context.Context, arg AutoArchiveInactiveChatsParams) ([]AutoArchiveInactiveChatsRow, error)
|
||||
// old_provider is matched as text; new_provider is also cast to ai_provider_type
|
||||
// for the EXISTS check against ai_providers.type.
|
||||
// ai_provider_id IS NOT NULL is defensive; the check constraint already
|
||||
// enforces that non-deleted rows always have a provider ID.
|
||||
BackfillChatModelConfigProvider(ctx context.Context, arg BackfillChatModelConfigProviderParams) (sql.Result, error)
|
||||
BackoffChatDiffStatus(ctx context.Context, arg BackoffChatDiffStatusParams) error
|
||||
BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg BatchUpdateWorkspaceAgentMetadataParams) error
|
||||
BatchUpdateWorkspaceLastUsedAt(ctx context.Context, arg BatchUpdateWorkspaceLastUsedAtParams) error
|
||||
|
||||
Generated
+40
-6
@@ -738,19 +738,21 @@ const updateAIProvider = `-- name: UpdateAIProvider :one
|
||||
UPDATE
|
||||
ai_providers
|
||||
SET
|
||||
display_name = $1::text,
|
||||
enabled = $2::boolean,
|
||||
base_url = $3::text,
|
||||
settings = $4::text,
|
||||
settings_key_id = $5::text,
|
||||
type = $1::ai_provider_type,
|
||||
display_name = $2::text,
|
||||
enabled = $3::boolean,
|
||||
base_url = $4::text,
|
||||
settings = $5::text,
|
||||
settings_key_id = $6::text,
|
||||
updated_at = NOW()
|
||||
WHERE
|
||||
id = $6::uuid AND deleted = FALSE
|
||||
id = $7::uuid AND deleted = FALSE
|
||||
RETURNING
|
||||
id, type, name, display_name, enabled, deleted, base_url, settings, settings_key_id, created_at, updated_at
|
||||
`
|
||||
|
||||
type UpdateAIProviderParams struct {
|
||||
Type AIProviderType `db:"type" json:"type"`
|
||||
DisplayName sql.NullString `db:"display_name" json:"display_name"`
|
||||
Enabled bool `db:"enabled" json:"enabled"`
|
||||
BaseUrl string `db:"base_url" json:"base_url"`
|
||||
@@ -761,6 +763,7 @@ type UpdateAIProviderParams struct {
|
||||
|
||||
func (q *sqlQuerier) UpdateAIProvider(ctx context.Context, arg UpdateAIProviderParams) (AIProvider, error) {
|
||||
row := q.db.QueryRowContext(ctx, updateAIProvider,
|
||||
arg.Type,
|
||||
arg.DisplayName,
|
||||
arg.Enabled,
|
||||
arg.BaseUrl,
|
||||
@@ -5551,6 +5554,37 @@ func (q *sqlQuerier) GetPRInsightsTimeSeries(ctx context.Context, arg GetPRInsig
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const backfillChatModelConfigProvider = `-- name: BackfillChatModelConfigProvider :execresult
|
||||
UPDATE
|
||||
chat_model_configs
|
||||
SET
|
||||
provider = $1::text,
|
||||
updated_at = NOW()
|
||||
WHERE
|
||||
provider = $2::text
|
||||
AND deleted = FALSE
|
||||
AND ai_provider_id IS NOT NULL
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM ai_providers
|
||||
WHERE id = chat_model_configs.ai_provider_id
|
||||
AND type = $1::ai_provider_type
|
||||
AND deleted = FALSE
|
||||
)
|
||||
`
|
||||
|
||||
type BackfillChatModelConfigProviderParams struct {
|
||||
NewProvider string `db:"new_provider" json:"new_provider"`
|
||||
OldProvider string `db:"old_provider" json:"old_provider"`
|
||||
}
|
||||
|
||||
// old_provider is matched as text; new_provider is also cast to ai_provider_type
|
||||
// for the EXISTS check against ai_providers.type.
|
||||
// ai_provider_id IS NOT NULL is defensive; the check constraint already
|
||||
// enforces that non-deleted rows always have a provider ID.
|
||||
func (q *sqlQuerier) BackfillChatModelConfigProvider(ctx context.Context, arg BackfillChatModelConfigProviderParams) (sql.Result, error) {
|
||||
return q.db.ExecContext(ctx, backfillChatModelConfigProvider, arg.NewProvider, arg.OldProvider)
|
||||
}
|
||||
|
||||
const deleteChatModelConfigByID = `-- name: DeleteChatModelConfigByID :exec
|
||||
UPDATE
|
||||
chat_model_configs
|
||||
|
||||
@@ -66,6 +66,7 @@ RETURNING
|
||||
UPDATE
|
||||
ai_providers
|
||||
SET
|
||||
type = @type::ai_provider_type,
|
||||
display_name = sqlc.narg('display_name')::text,
|
||||
enabled = @enabled::boolean,
|
||||
base_url = @base_url::text,
|
||||
|
||||
@@ -144,6 +144,27 @@ WHERE
|
||||
provider = @provider::text
|
||||
AND deleted = FALSE;
|
||||
|
||||
-- name: BackfillChatModelConfigProvider :execresult
|
||||
-- old_provider is matched as text; new_provider is also cast to ai_provider_type
|
||||
-- for the EXISTS check against ai_providers.type.
|
||||
-- ai_provider_id IS NOT NULL is defensive; the check constraint already
|
||||
-- enforces that non-deleted rows always have a provider ID.
|
||||
UPDATE
|
||||
chat_model_configs
|
||||
SET
|
||||
provider = @new_provider::text,
|
||||
updated_at = NOW()
|
||||
WHERE
|
||||
provider = @old_provider::text
|
||||
AND deleted = FALSE
|
||||
AND ai_provider_id IS NOT NULL
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM ai_providers
|
||||
WHERE id = chat_model_configs.ai_provider_id
|
||||
AND type = @new_provider::ai_provider_type
|
||||
AND deleted = FALSE
|
||||
);
|
||||
|
||||
-- name: DeleteChatModelConfigsByAIProviderID :exec
|
||||
UPDATE
|
||||
chat_model_configs
|
||||
|
||||
@@ -6310,6 +6310,7 @@ func setOpenAIProviderBaseURL(
|
||||
}
|
||||
_, err = db.UpdateAIProvider(ctx, database.UpdateAIProviderParams{
|
||||
ID: provider.ID,
|
||||
Type: provider.Type,
|
||||
DisplayName: provider.DisplayName,
|
||||
Enabled: provider.Enabled,
|
||||
BaseUrl: baseURL,
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
package coderd_test
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
agplcoderd "github.com/coder/coder/v2/coderd"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/enterprise/dbcrypt"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
func TestBackfillBedrockProviderTypeEncryptedSettings(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
rawDB, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
|
||||
key := make([]byte, 32)
|
||||
_, _ = rand.Read(key)
|
||||
ciphers, err := dbcrypt.NewCiphers(key)
|
||||
require.NoError(t, err)
|
||||
cryptDB, err := dbcrypt.New(ctx, rawDB, ciphers...)
|
||||
require.NoError(t, err)
|
||||
|
||||
rawSettings, err := json.Marshal(codersdk.AIProviderSettings{
|
||||
Bedrock: &codersdk.AIProviderBedrockSettings{Region: "us-east-1"},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
provider := dbgen.AIProvider(t, cryptDB, database.AIProvider{
|
||||
Type: database.AiProviderTypeAnthropic,
|
||||
Settings: sql.NullString{String: string(rawSettings), Valid: true},
|
||||
})
|
||||
|
||||
agplcoderd.BackfillBedrockProviderType(ctx, cryptDB, logger)
|
||||
|
||||
// Verify via raw DB: type is not encrypted so it is directly readable.
|
||||
row, err := rawDB.GetAIProviderByName(ctx, provider.Name)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.AiProviderTypeBedrock, row.Type, "encrypted legacy row must be promoted")
|
||||
require.True(t, row.SettingsKeyID.Valid, "settings must remain encrypted after backfill")
|
||||
}
|
||||
@@ -200,6 +200,7 @@ func setOpenAIProviderBaseURL(
|
||||
}
|
||||
_, err = db.UpdateAIProvider(ctx, database.UpdateAIProviderParams{
|
||||
ID: provider.ID,
|
||||
Type: provider.Type,
|
||||
DisplayName: provider.DisplayName,
|
||||
Enabled: provider.Enabled,
|
||||
BaseUrl: baseURL,
|
||||
|
||||
@@ -1195,6 +1195,7 @@ func TestAIProviders(t *testing.T) {
|
||||
const newSettings = `{"_type":"bedrock","_version":1,"region":"us-east-1","model":"anthropic.claude-sonnet-4-5-20250929-v1:0","access_key":"AKIA-test","access_key_secret":"test-secret"}`
|
||||
updated, err := crypt.UpdateAIProvider(ctx, database.UpdateAIProviderParams{
|
||||
ID: provider.ID,
|
||||
Type: provider.Type,
|
||||
DisplayName: provider.DisplayName,
|
||||
Enabled: provider.Enabled,
|
||||
BaseUrl: provider.BaseUrl,
|
||||
@@ -1211,6 +1212,7 @@ func TestAIProviders(t *testing.T) {
|
||||
provider := insertProvider(t, crypt, ciphers)
|
||||
updated, err := crypt.UpdateAIProvider(ctx, database.UpdateAIProviderParams{
|
||||
ID: provider.ID,
|
||||
Type: provider.Type,
|
||||
DisplayName: provider.DisplayName,
|
||||
Enabled: provider.Enabled,
|
||||
BaseUrl: provider.BaseUrl,
|
||||
|
||||
Reference in New Issue
Block a user