cli/aibridged.go
Rama main · réplica de lectura del repositorio original.
//go:build !slim
package cli
import (
"context"
"slices"
"github.com/prometheus/client_golang/prometheus"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/config"
aibheaders "github.com/coder/coder/v2/aibridge/headers"
"github.com/coder/coder/v2/aibridge/keypool"
"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/database"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/tracing"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/quartz"
)
// AIBridgeDaemonOptions configures [NewAIBridgeDaemon].
type AIBridgeDaemonOptions struct {
// API is the coderd instance the daemon serves and dials over the
// in-memory RPC.
API *coderd.API
// Config supplies both the provider build settings and the pool's record
// policy.
Config codersdk.AIBridgeConfig
// Registerer receives the daemon's collectors, so it must not already
// hold them. Nil skips registration.
Registerer prometheus.Registerer
// BridgeMetrics is the metrics instance the bridge and key pools report
// to. Nil disables those metrics.
BridgeMetrics *aibridge.Metrics
// ProviderMetrics receives provider reload events. Nil builds one from
// Registerer; supply one to assert on reload metrics without owning the
// registerer.
ProviderMetrics *aibridged.Metrics
// Pubsub triggers provider hot reloads. Nil uses API.Pubsub.
Pubsub pubsub.Pubsub
}
// NewAIBridgeDaemon constructs the in-memory aibridge daemon and wires
// up a subscription that hot-reloads providers over the in-memory
// RPC on every ai_providers change event. The returned unsubscribe
// function tears down the subscription; callers must invoke it
// alongside Server.Close on shutdown.
//
// Tests reach this through [github.com/coder/coder/v2/coderd/aibridgedtest],
// so a deployment option that does not reach the serving pool fails there too.
//
// Reloads fetch the provider set from coderd over the in-memory DRPC
// (GetAIProviders) rather than reading the database directly, so embedded and
// standalone gateways construct providers identically. Pubsub remains the
// hot-reload trigger.
//
// SubscribeProviderReload performs a best-effort initial reload synchronously,
// so providers are populated before this returns whenever the fetch succeeds.
// That reload blocks while acquiring a client, and it passes a background
// context, so only the daemon lifecycle bounds the wait. That is acceptable
// here: the embedded daemon's connection is an in-memory pipe that comes up
// immediately.
func NewAIBridgeDaemon(ctx context.Context, opts AIBridgeDaemonOptions) (*aibridged.Server, func(), error) {
coderAPI, cfg, reg, metrics := opts.API, opts.Config, opts.Registerer, opts.BridgeMetrics
coderAPI.Logger.Debug(ctx, "starting in-memory aibridge daemon")
logger := coderAPI.Logger.Named("ai-gateway")
providerMetrics := opts.ProviderMetrics
if providerMetrics == nil {
providerMetrics = aibridged.NewMetrics(reg)
}
ps := opts.Pubsub
if ps == nil {
ps = coderAPI.Pubsub
}
tracer := coderAPI.TracerProvider.Tracer(tracing.TracerName)
// Create daemon. Construct it before subscribing so the reloader can use
// srv.Client to fetch providers over the in-memory RPC.
srv, err := aibridged.New(ctx, func(dialCtx context.Context) (aibridged.DRPCClient, error) {
return coderAPI.CreateInMemoryAIBridgeServer(dialCtx)
}, logger, tracer, coderAPI.Experiments, metrics, aibridged.WithPoolOptions(aibridged.PoolOptionsFromConfig(ctx, logger, cfg)))
if err != nil {
return nil, nil, xerrors.Errorf("start in-memory aibridge daemon: %w", err)
}
if reg != nil {
reg.MustRegister(srv.KeyPoolStateCollector())
}
// Subscribe to ai_providers change events so the backend tracks the database
// without a restart, and perform the initial reload. The reload data path
// is the in-memory RPC.
reloader := NewProviderRPCReloader(srv.ReplaceProviders, srv.Client, cfg, logger.Named("provider-loader"), metrics, providerMetrics)
unsubscribe, err := aibridged.SubscribeProviderReload(ctx, ps, reloader, logger.Named("provider-reload"))
if err != nil {
// Without the subscription the backend cannot track provider changes,
// so fail startup rather than serve a permanently stale snapshot.
_ = srv.Close()
return nil, nil, xerrors.Errorf("subscribe to ai providers change channel: %w", err)
}
return srv, unsubscribe, nil
}
// providerRPCReloader implements [aibridged.ProviderReloader] by fetching the
// live provider set from coderd over a DRPC client and publishing it to the
// selected backend. Embedded (in-memory RPC, pubsub-triggered) and standalone
// (WebSocket RPC, retried at startup) gateways share the fetch,
// build, replace, and reload-metric accounting live in one place.
type providerRPCReloader struct {
replaceProviders func(context.Context, []aibridge.Provider) error
client aibridged.ClientFunc
cfg codersdk.AIBridgeConfig
logger slog.Logger
aibridgeMetrics *aibridge.Metrics
providerMetrics *aibridged.Metrics
}
// NewProviderRPCReloader fetches and builds provider snapshots over DRPC and
// publishes them through replaceProviders. Failed publication is reported as a
// failed reload, preserving the selected backend's previous snapshot.
func NewProviderRPCReloader(
replaceProviders func(context.Context, []aibridge.Provider) error,
client aibridged.ClientFunc,
cfg codersdk.AIBridgeConfig,
logger slog.Logger,
aibridgeMetrics *aibridge.Metrics,
providerMetrics *aibridged.Metrics,
) aibridged.ProviderReloader {
return &providerRPCReloader{
replaceProviders: replaceProviders,
client: client,
cfg: cfg,
logger: logger,
aibridgeMetrics: aibridgeMetrics,
providerMetrics: providerMetrics,
}
}
func (r *providerRPCReloader) Reload(ctx context.Context) error {
r.providerMetrics.RecordReloadAttempt()
// r.client blocks until the daemon connects to coderd or ctx is canceled.
client, err := r.client(ctx)
if err != nil {
return xerrors.Errorf("get ai-gateway client: %w", err)
}
resp, err := client.GetAIProviders(ctx, &proto.GetAIProvidersRequest{})
if err != nil {
// Keep the previous snapshot in place: dropping all providers
// because the fetch failed would compound the visible failure mode
// beyond the operator's actual misconfiguration.
return xerrors.Errorf("fetch ai providers: %w", err)
}
providers, outcomes := BuildProvidersFromProto(ctx, resp.GetProviders(), r.cfg, r.logger, r.aibridgeMetrics)
if err := r.replaceProviders(ctx, providers); err != nil {
return xerrors.Errorf("replace ai providers: %w", err)
}
r.providerMetrics.RecordReloadSuccess(outcomes)
return nil
}
// BuildProvidersFromProto constructs the runtime [aibridge.Provider] set from
// proto provider configuration.
//
// Disabled entries produce a Provider stub with Enabled() == false so the
// bridge can answer requests targeting them with a 503 sentinel.
//
// Per-provider construction errors are logged and the offending entry is
// excluded from the returned snapshot; this keeps a single misconfigured
// provider from taking the whole daemon down. The returned outcomes mirror the
// per-provider status for metrics reporting.
func BuildProvidersFromProto(ctx context.Context, protoProviders []*proto.AIProvider, cfg codersdk.AIBridgeConfig, logger slog.Logger, metrics *aibridge.Metrics) ([]aibridge.Provider, []aibridged.ProviderOutcome) {
providers := make([]aibridge.Provider, 0, len(protoProviders))
outcomes := make([]aibridged.ProviderOutcome, 0, len(protoProviders))
enabledCount := 0
for _, pp := range protoProviders {
spec := protoToProviderSpec(pp)
outcome := aibridged.ProviderOutcome{
Name: spec.Name,
Type: string(spec.Type),
}
if spec.Enabled {
enabledCount++
}
prov, err := buildProvider(ctx, spec, cfg, logger, metrics)
if err != nil {
outcome.Status = aibridged.ProviderStatusError
outcome.Err = err
outcomes = append(outcomes, outcome)
logger.Error(ctx, "skipping misconfigured ai provider",
slog.F("provider_name", spec.Name),
slog.F("provider_type", string(spec.Type)),
slog.Error(err),
)
continue
}
if spec.Enabled {
outcome.Status = aibridged.ProviderStatusEnabled
} else {
outcome.Status = aibridged.ProviderStatusDisabled
}
outcomes = append(outcomes, outcome)
providers = append(providers, prov)
}
if enabledCount > 0 && !slices.ContainsFunc(providers, func(p aibridge.Provider) bool { return p.Enabled() }) {
logger.Warn(ctx, "all enabled ai providers failed to build; only disabled providers remain")
}
return providers, outcomes
}
// protoToProviderSpec maps a proto [proto.AIProvider] into the database-neutral
// [aiProviderSpec] consumed by [buildProvider]. Keys and Bedrock settings are
// only meaningful for enabled providers; disabled providers carry neither over
// the wire.
func protoToProviderSpec(pp *proto.AIProvider) aiProviderSpec {
spec := aiProviderSpec{
Type: database.AIProviderType(pp.GetType()),
Name: pp.GetName(),
Enabled: pp.GetEnabled(),
BaseURL: pp.GetBaseUrl(),
Keys: pp.GetKeys(),
}
if b := pp.GetBedrock(); b != nil {
bedrock := codersdk.NewAIProviderBedrockSettings(
b.GetRegion(),
b.GetAccessKey(),
b.GetAccessKeySecret(),
b.GetModel(),
b.GetSmallFastModel(),
)
bedrock.RoleARN = b.GetRoleArn()
bedrock.ExternalID = b.GetExternalId()
bedrock.Protocol = codersdk.AIProviderBedrockProtocol(b.GetProtocol())
bedrock.ResolvedModel = b.GetResolvedModel()
bedrock.ResolvedSmallFastModel = b.GetResolvedSmallFastModel()
spec.Bedrock = new(bedrock)
}
if cp := pp.GetClaudePlatformAws(); cp != nil {
claudePlatform := codersdk.AIProviderClaudePlatformAWSSettings{
Region: cp.GetRegion(),
WorkspaceID: cp.GetWorkspaceId(),
}
spec.ClaudePlatformAWS = &claudePlatform
}
return spec
}
// aiProviderSpec is a database-neutral description of a single provider,
// carrying exactly the inputs [buildProvider] needs. The RPC path
// ([protoToProviderSpec]) maps the proto provider into this shape so the
// per-type construction logic stays in one place.
type aiProviderSpec struct {
Type database.AIProviderType
Name string
Enabled bool
BaseURL string
// Keys holds bearer API keys for non-Bedrock providers.
Keys []string
// Bedrock holds Bedrock-specific settings when the provider targets
// AWS Bedrock; nil otherwise.
Bedrock *codersdk.AIProviderBedrockSettings
// ClaudePlatformAWS holds Claude Platform for AWS settings when the
// provider targets Anthropic's AWS-hosted Messages API; nil otherwise.
// Mutually exclusive with Bedrock.
ClaudePlatformAWS *codersdk.AIProviderClaudePlatformAWSSettings
}
// buildProvider constructs the appropriate [aibridge.Provider] for a
// single provider spec, independent of where the spec was sourced from.
func buildProvider(ctx context.Context, spec aiProviderSpec, cfg codersdk.AIBridgeConfig, logger slog.Logger, metrics *aibridge.Metrics) (aibridge.Provider, error) {
if !spec.Enabled {
return aibridge.NewDisabledProviderStub(spec.Name, string(spec.Type)), nil
}
cbCfg := circuitBreakerConfig(cfg)
dumpDir := cfg.APIDumpDir.Value()
var actorHeaderNames map[string]string
if cfg.SendActorHeaders.Value() {
actorHeaderNames = map[string]string{}
if name := cfg.ActorHeaderID.Value(); name != "" {
actorHeaderNames[aibheaders.ActorAttributeID] = name
}
if name := cfg.ActorHeaderUsername.Value(); name != "" {
actorHeaderNames[aibheaders.ActorAttributeUsername] = name
}
if name := cfg.ActorHeaderEmail.Value(); name != "" {
actorHeaderNames[aibheaders.ActorAttributeEmail] = name
}
}
// aibridge currently has native support for OpenAI and Anthropic
// only. The other ai_provider_type values (azure, google,
// openai-compat, openrouter, vercel) route through the OpenAI
// provider because chatd configures them against their
// OpenAI-compatible endpoints. Bedrock routes through the Anthropic
// provider with a Bedrock discriminator in Settings.
switch spec.Type {
case database.AIProviderTypeOpenai,
database.AIProviderTypeAzure,
database.AIProviderTypeGoogle,
database.AIProviderTypeOpenaiCompat,
database.AIProviderTypeOpenrouter,
database.AIProviderTypeVercel:
if len(spec.Keys) == 0 && !cfg.AllowBYOK.Value() {
return nil, xerrors.Errorf("%s provider has no api keys configured and BYOK is not enabled", spec.Type)
}
var pool *keypool.Pool
if len(spec.Keys) > 0 {
var err error
pool, err = buildAIProviderKeyPool(spec.Name, spec.Keys, metrics)
if err != nil {
return nil, xerrors.Errorf("%s key pool: %w", spec.Type, err)
}
}
return aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
Name: spec.Name,
BaseURL: spec.BaseURL,
KeyPool: pool,
APIDumpDir: dumpDir,
CircuitBreaker: cbCfg,
ActorHeaderNames: actorHeaderNames,
}), nil
case database.AIProviderTypeAnthropic:
claudePlatform := claudePlatformConfig(spec.BaseURL, spec.ClaudePlatformAWS)
// A bearer-token Anthropic without any key cannot make upstream calls.
// Claude Platform in IAM mode authenticates by signing, so it counts
// as configured here too.
if len(spec.Keys) == 0 && !cfg.AllowBYOK.Value() && claudePlatform == nil {
return nil, xerrors.New("anthropic provider has no api keys configured and BYOK is not enabled")
}
var pool *keypool.Pool
if len(spec.Keys) > 0 {
var err error
pool, err = buildAIProviderKeyPool(spec.Name, spec.Keys, metrics)
if err != nil {
return nil, xerrors.Errorf("anthropic key pool: %w", err)
}
}
return aibridge.NewAnthropicProvider(ctx, aibridge.AnthropicConfig{
Logger: logger.With(slog.F("provider", spec.Name)),
Name: spec.Name,
BaseURL: spec.BaseURL,
KeyPool: pool,
APIDumpDir: dumpDir,
CircuitBreaker: cbCfg,
ActorHeaderNames: actorHeaderNames,
}, nil, claudePlatform)
case database.AIProviderTypeBedrock:
// A spec typed 'bedrock' authenticates exclusively via settings;
// without populated Bedrock credentials it cannot make upstream
// calls, so refuse rather than falling back to an unsigned
// Anthropic client.
bedrock := agplaibridge.BedrockConfig(spec.BaseURL, spec.Bedrock)
if bedrock == nil {
return nil, xerrors.New("bedrock provider has no bedrock credentials configured")
}
return aibridge.NewBedrockProvider(ctx, aibridge.AnthropicConfig{
Name: spec.Name,
BaseURL: spec.BaseURL,
APIDumpDir: dumpDir,
CircuitBreaker: cbCfg,
ActorHeaderNames: actorHeaderNames,
}, *bedrock)
case database.AIProviderTypeCopilot:
// Copilot is always BYOK; the per-user token is supplied on each
// request via the Authorization header, so no keypool is built.
return aibridge.NewCopilotProvider(aibridge.CopilotConfig{
Name: spec.Name,
BaseURL: spec.BaseURL,
APIDumpDir: dumpDir,
CircuitBreaker: cbCfg,
}), nil
default:
return nil, xerrors.Errorf("unsupported provider type: %q", spec.Type)
}
}
// buildAIProviderKeyPool builds a [keypool.Pool]. Callers must check
// len(keys) > 0 first; keypool.New rejects empty input.
func buildAIProviderKeyPool(providerName string, keys []string, metrics *aibridge.Metrics) (*keypool.Pool, error) {
return keypool.New(providerName, keys, quartz.NewReal(), metrics)
}
// claudePlatformConfig maps Claude Platform for AWS settings into the gateway
// config. baseURL, when set, overrides the default regional endpoint; the
// region still determines the SigV4 signing scope. Returns nil when the
// settings are absent or incomplete, so the provider falls back to a plain
// bearer-token Anthropic client.
func claudePlatformConfig(baseURL string, cp *codersdk.AIProviderClaudePlatformAWSSettings) *aibridge.AWSClaudePlatformConfig {
if cp == nil || !cp.IsConfigured() {
return nil
}
return &aibridge.AWSClaudePlatformConfig{
Region: cp.Region,
WorkspaceID: cp.WorkspaceID,
BaseURL: baseURL,
}
}
// circuitBreakerConfig returns nil when the breaker is disabled.
func circuitBreakerConfig(cfg codersdk.AIBridgeConfig) *config.CircuitBreaker {
if !cfg.CircuitBreakerEnabled.Value() {
return nil
}
return &config.CircuitBreaker{
FailureThreshold: uint32(cfg.CircuitBreakerFailureThreshold.Value()), //nolint:gosec // Validated by serpent.Validate in deployment options.
Interval: cfg.CircuitBreakerInterval.Value(),
Timeout: cfg.CircuitBreakerTimeout.Value(),
MaxRequests: uint32(cfg.CircuitBreakerMaxRequests.Value()), //nolint:gosec // Validated by serpent.Validate in deployment options.
}
}