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

Skip to content

feat: fetch providers over DRPC #26650

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
dannykopping merged 4 commits into from
Jun 29, 2026
Merged

feat: fetch providers over DRPC #26650

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d1ef25c
feat(coderd/aibridged): fetch AI providers over DRPC
dannykopping Jun 26, 2026
33ce3b8
test(coderd/aibridgedtest): build providers over DRPC in daemon helper
dannykopping Jun 29, 2026
0bd385b
refactor(coderd/aibridged): move websocket dialer to the gateway-star…
dannykopping Jun 29, 2026
3ddcfb0
test(coderd/aibridgedserver): cover skip of provider with undecodable…
dannykopping Jun 29, 2026
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
340 changes: 171 additions & 169 deletions cli/aibridged.go
View file
Open in desktop

Large diffs are not rendered by default.

72 changes: 50 additions & 22 deletions cli/aibridged_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,22 @@
package cli

import (
"context"
"database/sql"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/coderd"
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
"github.com/coder/coder/v2/coderd/aibridged"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/coderd/aibridgedserver"
agplaiseats "github.com/coder/coder/v2/coderd/aiseats"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
Expand All @@ -24,10 +29,12 @@ import (

// buildFromEnv exercises the same env-config-in/providers-out path that
// production uses on boot: SeedAIProvidersFromEnv writes the env-derived
// rows to the database, and BuildProviders reads them back as runtime
// [aibridge.Provider] instances. This keeps the existing TestBuildProviders
// table intact while reflecting the post-refactor flow where the database
// is the single source of truth.
// rows to the database, the server's GetAIProviders handler reads them back
// over the (post-refactor) DB-read path and maps them to proto, and
// BuildProvidersFromProto constructs the runtime [aibridge.Provider]
// instances. This keeps the existing TestBuildProviders table intact while
// reflecting the post-refactor flow where the database is the single source
// of truth and the gateway fetches providers over DRPC.
func buildFromEnv(t *testing.T, cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) {
t.Helper()
db, _ := dbtestutil.NewDB(t)
Expand All @@ -36,10 +43,28 @@ func buildFromEnv(t *testing.T, cfg codersdk.AIBridgeConfig) ([]aibridge.Provide
if err := coderd.SeedAIProvidersFromEnv(ctx, db, cfg, logger); err != nil {
return nil, err
}
providers, _, err := BuildProviders(ctx, db, cfg, logger, nil)
providers, _, err := buildFromDB(ctx, t, db, cfg, logger)
return providers, err
}

// buildFromDB runs the production fetch path against a database: it calls the
// server's GetAIProviders handler (DB read + proto mapping) and then
// BuildProvidersFromProto (proto -> runtime providers), returning the same
// (providers, outcomes) the embedded reloader would observe.
func buildFromDB(ctx context.Context, t *testing.T, db database.Store, cfg codersdk.AIBridgeConfig, logger slog.Logger) ([]aibridge.Provider, []aibridged.ProviderOutcome, error) {
t.Helper()
srv, err := aibridgedserver.NewServer(ctx, db, logger, "/", cfg, nil, nil, agplaiseats.Noop{})
if err != nil {
return nil, nil, err
}
resp, err := srv.GetAIProviders(ctx, &proto.GetAIProvidersRequest{})
if err != nil {
return nil, nil, err
}
providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), cfg, logger, nil)
return providers, outcomes, nil
}

func TestBuildProviders(t *testing.T) {
t.Parallel()

Expand Down Expand Up @@ -243,7 +268,7 @@ func TestBuildProviders(t *testing.T) {
Name: aibridge.ProviderAnthropic,
BaseUrl: "https://api.anthropic.com/",
}
assert.Nil(t, bedrockConfigFromRow(row, codersdk.AIProviderSettings{}))
assert.Nil(t, bedrockConfig(row.BaseUrl, codersdk.AIProviderSettings{}.Bedrock))
})

t.Run("NativeAnthropicCustomBaseURL", func(t *testing.T) {
Expand All @@ -253,7 +278,7 @@ func TestBuildProviders(t *testing.T) {
Name: "anthropic-proxy",
BaseUrl: "https://internal-proxy.example.com/anthropic/",
}
assert.Nil(t, bedrockConfigFromRow(row, codersdk.AIProviderSettings{}))
assert.Nil(t, bedrockConfig(row.BaseUrl, codersdk.AIProviderSettings{}.Bedrock))
})

t.Run("BedrockSettingsPresent", func(t *testing.T) {
Expand All @@ -278,7 +303,7 @@ func TestBuildProviders(t *testing.T) {
RoleARN: roleARN,
},
}
got := bedrockConfigFromRow(row, settings)
got := bedrockConfig(row.BaseUrl, settings.Bedrock)
require.NotNil(t, got)
assert.Equal(t, row.BaseUrl, got.BaseURL)
assert.Equal(t, "us-west-2", got.Region)
Expand All @@ -302,7 +327,7 @@ func TestBuildProviders(t *testing.T) {
settings := codersdk.AIProviderSettings{
Bedrock: &codersdk.AIProviderBedrockSettings{},
}
assert.Nil(t, bedrockConfigFromRow(row, settings))
assert.Nil(t, bedrockConfig(row.BaseUrl, settings.Bedrock))
})
}

Expand All @@ -328,13 +353,14 @@ func TestBuildProvidersSkipsBadRows(t *testing.T) {
Settings: sql.NullString{String: "not-json", Valid: true},
})

providers, outcomes, err := BuildProviders(ctx, db, codersdk.AIBridgeConfig{}, logger, nil)
// A row whose settings blob cannot be decoded is dropped server-side
// in GetAIProviders, so it never reaches the client: no provider and
// no outcome. This keeps one corrupt row from breaking the fetch (and
// thus provider configuration) for every gateway.
providers, outcomes, err := buildFromDB(ctx, t, db, codersdk.AIBridgeConfig{}, logger)
require.NoError(t, err)
assert.Empty(t, providers)
require.Len(t, outcomes, 1)
assert.Equal(t, "anthropic-broken", outcomes[0].Name)
assert.Equal(t, aibridged.ProviderStatusError, outcomes[0].Status)
assert.Error(t, outcomes[0].Err)
assert.Empty(t, outcomes)
})

t.Run("EnabledButNoKeys", func(t *testing.T) {
Expand All @@ -352,7 +378,7 @@ func TestBuildProvidersSkipsBadRows(t *testing.T) {
BaseUrl: "https://example.openai.azure.com/",
})

providers, outcomes, err := BuildProviders(ctx, db, codersdk.AIBridgeConfig{}, logger, nil)
providers, outcomes, err := buildFromDB(ctx, t, db, codersdk.AIBridgeConfig{}, logger)
require.NoError(t, err)
assert.Empty(t, providers)
require.Len(t, outcomes, 1)
Expand All @@ -365,11 +391,13 @@ func TestBuildProvidersSkipsBadRows(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitShort)
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})

// An enabled provider with no keys (and BYOK disabled) fails to build
// on the client side, yielding a ProviderStatusError outcome. It must
// not prevent the good provider from being built.
dbgen.AIProvider(t, db, database.AIProvider{
Type: database.AIProviderTypeAnthropic,
Name: "anthropic-broken",
BaseUrl: "https://api.anthropic.com/",
Settings: sql.NullString{String: "{not valid json", Valid: true},
Type: database.AIProviderTypeAzure,
Name: "azure-broken",
BaseUrl: "https://example.openai.azure.com/",
})
good := dbgen.AIProvider(t, db, database.AIProvider{
Type: database.AIProviderTypeOpenai,
Expand All @@ -381,7 +409,7 @@ func TestBuildProvidersSkipsBadRows(t *testing.T) {
APIKey: "sk-good",
})

providers, outcomes, err := BuildProviders(ctx, db, codersdk.AIBridgeConfig{}, logger, nil)
providers, outcomes, err := buildFromDB(ctx, t, db, codersdk.AIBridgeConfig{}, logger)
require.NoError(t, err)
require.Len(t, providers, 1)
assert.Equal(t, "openai-good", providers[0].Name())
Expand All @@ -390,7 +418,7 @@ func TestBuildProvidersSkipsBadRows(t *testing.T) {
for _, o := range outcomes {
byName[o.Name] = o
}
assert.Equal(t, aibridged.ProviderStatusError, byName["anthropic-broken"].Status)
assert.Equal(t, aibridged.ProviderStatusError, byName["azure-broken"].Status)
assert.Equal(t, aibridged.ProviderStatusEnabled, byName["openai-good"].Status)
})

Expand Down Expand Up @@ -439,7 +467,7 @@ func TestBuildProvidersSkipsBadRows(t *testing.T) {
p.Enabled = false
})

providers, outcomes, err := BuildProviders(ctx, db, codersdk.AIBridgeConfig{}, logger, nil)
providers, outcomes, err := buildFromDB(ctx, t, db, codersdk.AIBridgeConfig{}, logger)
require.NoError(t, err)
require.Len(t, providers, 1, "disabled providers stay in the snapshot so the bridge can serve a 503 sentinel")
assert.Equal(t, tc.row.Name, providers[0].Name())
Expand Down
6 changes: 1 addition & 5 deletions cli/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1134,12 +1134,8 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
// https://linear.app/codercom/issue/AIGOV-447/remove-legacy-ai-gateway-metric-aliases
aibridgeReg := prometheusmetrics.NewMetricAliasRegisterer(coderAPI.PrometheusRegistry, "coder_ai_gateway_", "coder_aibridged_")
aibridgeMetrics := aibridge.NewMetrics(aibridgeReg)
aibridgeProviders, _, err := BuildProviders(aibridgeInitCtx, options.Database, vals.AI.BridgeConfig, logger.Named("aibridge.providers"), aibridgeMetrics)
if err != nil {
return xerrors.Errorf("build AI providers: %w", err)
}
var unsubscribeProviderReload func()
aibridgeDaemon, unsubscribeProviderReload, err = newAIBridgeDaemon(coderAPI, aibridgeProviders, vals.AI.BridgeConfig, aibridgeReg, aibridgeMetrics)
aibridgeDaemon, unsubscribeProviderReload, err = newAIBridgeDaemon(coderAPI, vals.AI.BridgeConfig, aibridgeReg, aibridgeMetrics)
if err != nil {
return xerrors.Errorf("create aibridged: %w", err)
}
Expand Down
74 changes: 31 additions & 43 deletions cli/server_aibridge_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,6 @@ package cli

import (
"context"
"database/sql"
"encoding/json"
"fmt"
"testing"

Expand All @@ -13,8 +11,8 @@ import (
"cdr.dev/slog/v3"
"cdr.dev/slog/v3/sloggers/slogtest"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/util/ptr"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/testutil"
"github.com/coder/serpent"
Expand Down Expand Up @@ -628,110 +626,108 @@ func TestWarnIfAIProvidersConfiguredFromEnv(t *testing.T) {
})
}

func TestBuildAIProviderFromRowSetsAPIDumpDir(t *testing.T) {
func TestBuildProviderFromProtoSetsAPIDumpDir(t *testing.T) {
t.Parallel()

const dumpDir = "/tmp/coder-aibridge-dumps"

tests := []struct {
name string
row database.AIProvider
provider *proto.AIProvider
expectedType string
}{
{
name: "OpenAI",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeOpenai,
Type: string(database.AIProviderTypeOpenai),
Name: "openai",
BaseUrl: "https://api.openai.com/",
},
expectedType: aibridge.ProviderOpenAI,
},
{
name: "Anthropic",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeAnthropic,
Type: string(database.AIProviderTypeAnthropic),
Name: "anthropic",
BaseUrl: "https://api.anthropic.com/",
},
expectedType: aibridge.ProviderAnthropic,
},
{
name: "Copilot",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeCopilot,
Type: string(database.AIProviderTypeCopilot),
Name: "copilot",
BaseUrl: "https://api.githubcopilot.com/",
},
expectedType: aibridge.ProviderCopilot,
},
{
name: "Azure",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeAzure,
Type: string(database.AIProviderTypeAzure),
Name: "azure",
BaseUrl: "https://example.openai.azure.com/",
},
expectedType: aibridge.ProviderOpenAI,
},
{
name: "Google",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeGoogle,
Type: string(database.AIProviderTypeGoogle),
Name: "google",
BaseUrl: "https://generativelanguage.googleapis.com/v1beta/openai/",
},
expectedType: aibridge.ProviderOpenAI,
},
{
name: "OpenAICompat",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeOpenaiCompat,
Type: string(database.AIProviderTypeOpenaiCompat),
Name: "openai-compat",
BaseUrl: "https://compat.example.com/v1/",
},
expectedType: aibridge.ProviderOpenAI,
},
{
name: "OpenRouter",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeOpenrouter,
Type: string(database.AIProviderTypeOpenrouter),
Name: "openrouter",
BaseUrl: "https://openrouter.ai/api/v1/",
},
expectedType: aibridge.ProviderOpenAI,
},
{
name: "Vercel",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeVercel,
Type: string(database.AIProviderTypeVercel),
Name: "vercel",
BaseUrl: "https://api.v0.dev/v1/",
},
expectedType: aibridge.ProviderOpenAI,
},
{
name: "Bedrock",
row: database.AIProvider{
provider: &proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeBedrock,
Type: string(database.AIProviderTypeBedrock),
Name: "bedrock",
BaseUrl: "https://bedrock-runtime.us-east-1.amazonaws.com/",
Settings: mustMarshalSettings(codersdk.AIProviderSettings{
Bedrock: &codersdk.AIProviderBedrockSettings{
Region: "us-east-1",
AccessKey: ptr.Ref("AKID"),
AccessKeySecret: ptr.Ref("secret"),
},
}),
Bedrock: &proto.AIProviderKindBedrock{
Region: "us-east-1",
AccessKey: "AKID",
AccessKeySecret: "secret",
},
},
expectedType: aibridge.ProviderAnthropic,
},
Expand All @@ -741,7 +737,7 @@ func TestBuildAIProviderFromRowSetsAPIDumpDir(t *testing.T) {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()

provider, err := buildAIProviderFromRow(t.Context(), tt.row, nil, codersdk.AIBridgeConfig{
provider, err := buildProvider(t.Context(), protoToProviderSpec(tt.provider), codersdk.AIBridgeConfig{
AllowBYOK: serpent.Bool(true),
APIDumpDir: serpent.String(dumpDir),
}, nil)
Expand All @@ -752,29 +748,21 @@ func TestBuildAIProviderFromRowSetsAPIDumpDir(t *testing.T) {
}
}

func TestBuildAIProviderFromRowBedrockWithoutSettings(t *testing.T) {
func TestBuildProviderFromProtoBedrockWithoutSettings(t *testing.T) {
t.Parallel()

_, err := buildAIProviderFromRow(t.Context(), database.AIProvider{
_, err := buildProvider(t.Context(), protoToProviderSpec(&proto.AIProvider{
Enabled: true,
Type: database.AIProviderTypeBedrock,
Type: string(database.AIProviderTypeBedrock),
Name: "bedrock-no-settings",
BaseUrl: "https://bedrock-runtime.us-east-1.amazonaws.com/",
}, nil, codersdk.AIBridgeConfig{
}), codersdk.AIBridgeConfig{
AllowBYOK: serpent.Bool(true),
}, nil)
require.Error(t, err)
assert.Contains(t, err.Error(), "bedrock provider has no bedrock credentials configured")
}

func mustMarshalSettings(s codersdk.AIProviderSettings) sql.NullString {
data, err := json.Marshal(s)
if err != nil {
panic(err)
}
return sql.NullString{String: string(data), Valid: true}
}

func assertFieldValue(t *testing.T, fields slog.Map, name string, expected interface{}) {
t.Helper()
for _, f := range fields {
Expand Down
Loading
Loading