coderd/aibridgedserver/aibridgedserver.go
Rama main · réplica de lectura del repositorio original.
package aibridgedserver
import (
"context"
"database/sql"
"encoding/json"
"errors"
"net/url"
"slices"
"strings"
"sync"
"time"
"github.com/google/uuid"
"github.com/hashicorp/go-multierror"
"golang.org/x/xerrors"
"google.golang.org/protobuf/types/known/anypb"
"google.golang.org/protobuf/types/known/structpb"
"google.golang.org/protobuf/types/known/timestamppb"
"storj.io/drpc/drpcerr"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge/recorder"
"github.com/coder/coder/v2/coderd/aibridge/budget"
"github.com/coder/coder/v2/coderd/aibridged"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/coderd/aiseats"
"github.com/coder/coder/v2/coderd/apikey"
"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/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/externalauth"
"github.com/coder/coder/v2/coderd/httpmw"
codermcp "github.com/coder/coder/v2/coderd/mcp"
"github.com/coder/coder/v2/coderd/notifications"
"github.com/coder/coder/v2/coderd/provisionerdserver"
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
"github.com/coder/coder/v2/coderd/util/ptr"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/quartz"
)
var (
ErrExpiredOrInvalidOAuthToken = xerrors.New("expired or invalid OAuth2 token")
ErrNoMCPConfigFound = xerrors.New("no MCP config found")
// These errors are returned by IsAuthorized. Since they're just returned as
// a generic dRPC error, it's difficult to tell them apart without string
// matching.
// TODO: return these errors to the client in a more structured/comparable
// way.
ErrInvalidKey = rejectionError("invalid key")
ErrUnknownKey = rejectionError("unknown key")
ErrExpired = rejectionError("expired")
ErrUnknownUser = rejectionError("unknown user")
ErrDeletedUser = rejectionError("deleted user")
ErrInactiveUser = rejectionError("inactive user")
ErrSystemUser = rejectionError("system user")
ErrAmbiguousAuth = rejectionError("direct and delegated authentication are mutually exclusive")
ErrInvalidDelegatedOrganization = rejectionError("invalid delegated_organization_id")
ErrWorkspaceAttribution = rejectionError("invalid workspace attribution")
ErrUnknownWorkspace = rejectionError("unknown workspace")
ErrUnknownOrganizationMember = rejectionError("user is not a member of the request organization")
ErrNoOrganization = rejectionError("no organization could be resolved for the user")
ErrNoExternalAuthLinkFound = xerrors.New("no external auth link found")
)
// rejectionError is an IsAuthorized error that refuses the request. The gateway
// answers it with a 403. Every other IsAuthorized error is a failure of coderd
// itself and carries [proto.ErrCodeInternal], so the gateway answers a 500.
type rejectionError string
func (r rejectionError) Error() string { return string(r) }
var _ aibridged.DRPCServer = &Server{}
type store interface {
// Recorder-related queries.
InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error)
InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) (database.AIBridgeTokenUsage, error)
InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) (database.AIBridgeUserPrompt, error)
InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) (database.AIBridgeToolUsage, error)
InsertAIBridgeModelThought(ctx context.Context, arg database.InsertAIBridgeModelThoughtParams) (database.AIBridgeModelThought, error)
UpdateAIBridgeInterceptionEnded(ctx context.Context, intcID database.UpdateAIBridgeInterceptionEndedParams) (database.AIBridgeInterception, error)
GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (database.GetAIBridgeInterceptionLineageByToolCallIDRow, error)
// Cost-attribution queries, used to snapshot price and effective group on
// each token usage record.
GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error)
GetAIProviderByName(ctx context.Context, name string) (database.AIProvider, error)
GetAIModelPriceForProvider(ctx context.Context, arg database.GetAIModelPriceForProviderParams) (database.GetAIModelPriceForProviderRow, error)
GetUserAIBudgetOverrideWithOrganization(ctx context.Context, arg database.GetUserAIBudgetOverrideWithOrganizationParams) (database.GetUserAIBudgetOverrideWithOrganizationRow, error)
GetHighestGroupAIBudgetByUser(ctx context.Context, arg database.GetHighestGroupAIBudgetByUserParams) (database.GetHighestGroupAIBudgetByUserRow, error)
GetUserEveryoneFallbackGroup(ctx context.Context, arg database.GetUserEveryoneFallbackGroupParams) (uuid.UUID, error)
GetUserAISpendSince(ctx context.Context, arg database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error)
GetGroupByID(ctx context.Context, id uuid.UUID) (database.Group, error)
GetOrganizationByID(ctx context.Context, id uuid.UUID) (database.Organization, error)
GetUsers(ctx context.Context, arg database.GetUsersParams) ([]database.GetUsersRow, error)
// MCPConfigurator-related queries.
GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error)
// Authorizer-related queries.
GetAPIKeyByID(ctx context.Context, id string) (database.APIKey, error)
GetUserByID(ctx context.Context, id uuid.UUID) (database.User, error)
GetWorkspaceByID(ctx context.Context, id uuid.UUID) (database.Workspace, error)
GetOrganizationsByUserID(ctx context.Context, arg database.GetOrganizationsByUserIDParams) ([]database.Organization, error)
// ProviderConfigurator-related queries. InTx wraps the provider and key
// reads in a single read-only transaction.
GetAIProviders(ctx context.Context, arg database.GetAIProvidersParams) ([]database.AIProvider, error)
GetAIProviderKeysByProviderIDs(ctx context.Context, providerIDs []uuid.UUID) ([]database.AIProviderKey, error)
InTx(func(database.Store) error, *database.TxOptions) error
}
type Server struct {
// lifecycleCtx must be tied to the API server's lifecycle
// as when the API server shuts down, we want to cancel any
// long-running operations.
lifecycleCtx context.Context
store store
pubsub pubsub.Pubsub
logger slog.Logger
externalAuthConfigs map[string]*externalauth.Config
coderMCPConfig *proto.MCPServerConfig // may be nil if not available
structuredLogging bool
aiSeatTracker aiseats.SeatTracker
experiments codersdk.Experiments
// budgetPolicy selects the effective group when a user belongs to multiple
// budgeted groups, used for cost attribution on token usage records.
budgetPolicy codersdk.AIBudgetPolicy
// budgetPeriod is the deployment-configured budgeting period used to
// derive the window over which user AI spend is aggregated.
budgetPeriod codersdk.AIBudgetPeriod
clock quartz.Clock
notifEnqueuer notifications.Enqueuer
// metrics records cost-control metrics. May be nil.
metrics *Metrics
}
// Options carries the dependencies required to construct an aibridged Server.
type Options struct {
Store store
Pubsub pubsub.Pubsub
AISeatTracker aiseats.SeatTracker
// Enqueuer enqueues notifications. When nil, NewServer substitutes a no-op
// enqueuer.
Enqueuer notifications.Enqueuer
AccessURL string
GatewayCfg codersdk.AIBridgeConfig
ExternalAuthConfigs []*externalauth.Config
Experiments codersdk.Experiments
// OAuth2ProviderEnabled gates the internal MCP server, which is
// unavailable when it is off.
OAuth2ProviderEnabled bool
Logger slog.Logger
Clock quartz.Clock
Metrics *Metrics
}
func NewServer(lifecycleCtx context.Context, opts Options) (*Server, error) {
enqueuer := opts.Enqueuer
if enqueuer == nil {
enqueuer = notifications.NewNoopEnqueuer()
}
eac := make(map[string]*externalauth.Config, len(opts.ExternalAuthConfigs))
for _, cfg := range opts.ExternalAuthConfigs {
// Only External Auth configs which are configured with an MCP URL are relevant to aibridged.
if cfg.MCPURL == "" {
continue
}
eac[cfg.ID] = cfg
}
srv := &Server{
lifecycleCtx: lifecycleCtx,
store: opts.Store,
pubsub: opts.Pubsub,
logger: opts.Logger,
externalAuthConfigs: eac,
structuredLogging: opts.GatewayCfg.EmitsStructuredLogs(codersdk.AIStructuredLoggingSourceCoderd),
aiSeatTracker: opts.AISeatTracker,
experiments: opts.Experiments,
budgetPolicy: codersdk.NewAIBudgetPolicyFromString(opts.GatewayCfg.BudgetPolicy),
budgetPeriod: codersdk.NewAIBudgetPeriodFromString(opts.GatewayCfg.BudgetPeriod),
clock: opts.Clock,
notifEnqueuer: enqueuer,
metrics: opts.Metrics,
}
if opts.GatewayCfg.InjectCoderMCPTools {
opts.Logger.Warn(lifecycleCtx, "inject MCP tools option is deprecated and will be removed in a future release")
coderMCPConfig, err := getCoderMCPServerConfig(opts.Experiments, opts.OAuth2ProviderEnabled, opts.AccessURL)
if err != nil {
opts.Logger.Warn(lifecycleCtx, "failed to retrieve coder MCP server config, Coder MCP will not be available", slog.Error(err))
}
srv.coderMCPConfig = coderMCPConfig
}
return srv, nil
}
// recordTime resolves the time a gateway reported for a record. A gateway is a
// separate process, possibly of a different version, so an unset time is
// replaced with the server's own rather than persisted as the zero time, which
// would place the record in year one and skew budget periods and retention.
func (s *Server) recordTime(ctx context.Context, field string, reported *timestamppb.Timestamp, interceptionID string) time.Time {
if reported.IsValid() && !reported.AsTime().IsZero() {
return reported.AsTime()
}
now := dbtime.Now()
s.logger.Warn(ctx, "record reported no time, using server time",
slog.F("field", field),
slog.F("interception_id", interceptionID),
slog.F("server_time", now),
)
return now
}
func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
intcID, err := uuid.Parse(in.GetId())
if err != nil {
return nil, xerrors.Errorf("invalid interception ID %q: %w", in.GetId(), err)
}
initID, err := uuid.Parse(in.GetInitiatorId())
if err != nil {
return nil, xerrors.Errorf("invalid initiator ID %q: %w", in.GetInitiatorId(), err)
}
if in.ApiKeyId == "" {
return nil, xerrors.Errorf("empty API key ID")
}
metadata := metadataToMap(in.GetMetadata())
if in.UserAgent != "" {
if _, ok := metadata[recorder.MetadataUserAgentKey]; ok {
s.logger.Warn(ctx, "interception metadata contains user agent key, will be overwritten")
}
metadata[recorder.MetadataUserAgentKey] = in.UserAgent
}
// Look up the interception lineage using the correlating tool call ID.
parentID, rootID := s.findInterceptionLineage(ctx, in.GetCorrelatingToolCallId())
out, err := json.Marshal(metadata)
if err != nil {
s.logger.Warn(ctx, "failed to marshal ai gateway metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
}
providerName := strings.TrimSpace(in.ProviderName)
if providerName == "" {
providerName = in.Provider
}
agentFirewallSessionID, err := parseOptionalUUID(in.AgentFirewallSessionId)
if err != nil {
s.logger.Warn(ctx, "invalid agent firewall session ID in interception request",
slog.F("agent_firewall_session_id", in.GetAgentFirewallSessionId()), slog.Error(err))
}
workspaceID, err := interceptionAttribution(in)
if err != nil {
s.logger.Warn(ctx, "failed to parse interception attribution", slog.F("metadata", in), slog.Error(err))
}
organizationID, organizationSource, err := interceptionOrganization(in)
if err != nil {
s.logger.Warn(ctx, "invalid organization attribution in interception request, storing none",
slog.F("organization_id", in.GetOrganizationId()),
slog.F("organization_source", in.GetOrganizationSource()),
slog.Error(err))
}
if s.structuredLogging {
fields := []slog.Field{
slog.F("record_type", recorder.RecordTypeInterceptionStart),
slog.F("interception_id", intcID.String()),
slog.F("initiator_id", initID.String()),
slog.F("api_key_id", in.ApiKeyId),
slog.F("provider", in.Provider),
slog.F("model", in.Model),
slog.F("client", in.Client),
slog.F("client_session_id", in.GetClientSessionId()),
slog.F("started_at", in.StartedAt.AsTime()),
slog.F("metadata", metadata),
slog.F("correlating_tool_call_id", in.GetCorrelatingToolCallId()),
slog.F("thread_parent_id", parentID),
slog.F("thread_root_id", rootID),
}
if workspaceID.Valid {
fields = append(fields, slog.F("workspace_id", workspaceID.UUID))
}
if organizationID.Valid {
fields = append(fields,
slog.F("organization_id", organizationID.UUID),
slog.F("organization_source", organizationSource.AIBridgeInterceptionOrganizationSource),
)
}
s.logger.Info(ctx, recorder.InterceptionLogMarker, fields...)
}
_, err = s.store.InsertAIBridgeInterception(ctx, database.InsertAIBridgeInterceptionParams{
ID: intcID,
APIKeyID: sql.NullString{String: in.ApiKeyId, Valid: true},
Client: sql.NullString{String: in.Client, Valid: in.Client != ""},
ClientSessionID: sql.NullString{String: in.GetClientSessionId(), Valid: in.GetClientSessionId() != ""},
InitiatorID: initID,
Provider: in.Provider,
ProviderName: providerName,
Model: in.Model,
Metadata: out,
StartedAt: s.recordTime(ctx, "started_at", in.GetStartedAt(), intcID.String()),
ThreadParentInterceptionID: uuid.NullUUID{UUID: parentID, Valid: parentID != uuid.Nil},
ThreadRootInterceptionID: uuid.NullUUID{UUID: rootID, Valid: rootID != uuid.Nil},
CredentialKind: credentialKindOrDefault(in.CredentialKind),
CredentialHint: in.CredentialHint,
AgentFirewallSessionID: agentFirewallSessionID,
AgentFirewallSequenceNumber: parseOptionalInt32(in.AgentFirewallSequenceNumber),
WorkspaceID: workspaceID,
OrganizationID: organizationID,
OrganizationSource: organizationSource,
})
if err != nil {
return nil, xerrors.Errorf("start interception: %w", err)
}
if !s.experiments.Enabled(codersdk.ExperimentAIGatewaySeatExclusion) {
reason := aiseats.ReasonAIBridge("provider=" + in.Provider + ", model=" + in.Model)
s.aiSeatTracker.RecordUsage(ctx, initID, reason)
}
return &proto.RecordInterceptionResponse{}, nil
}
func (s *Server) RecordInterceptionEnded(ctx context.Context, in *proto.RecordInterceptionEndedRequest) (*proto.RecordInterceptionEndedResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
intcID, err := uuid.Parse(in.GetId())
if err != nil {
return nil, xerrors.Errorf("invalid interception ID %q: %w", in.GetId(), err)
}
if s.structuredLogging {
s.logger.Info(ctx, recorder.InterceptionLogMarker,
slog.F("record_type", recorder.RecordTypeInterceptionEnd),
slog.F("interception_id", intcID.String()),
slog.F("ended_at", in.EndedAt.AsTime()),
)
}
// The error type and message form one logical unit: the terminal error.
// Gate the message on the type so the row never carries a message without a
// type (the migration treats both-NULL as a successful interception).
errType := interceptionErrorType(in.GetErrorType())
var errMsg sql.NullString
if errType.Valid && in.GetErrorMessage() != "" {
errMsg = sql.NullString{String: truncateErrorMessage(in.GetErrorMessage()), Valid: true}
}
_, err = s.store.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
ID: intcID,
EndedAt: s.recordTime(ctx, "ended_at", in.GetEndedAt(), intcID.String()),
CredentialHint: in.CredentialHint,
ErrorType: errType,
ErrorMessage: errMsg,
})
if err != nil {
return nil, xerrors.Errorf("end interception: %w", err)
}
return &proto.RecordInterceptionEndedResponse{}, nil
}
func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
intcID, err := uuid.Parse(in.GetInterceptionId())
if err != nil {
return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err)
}
metadata := metadataToMap(in.GetMetadata())
if s.structuredLogging {
s.logger.Info(ctx, recorder.InterceptionLogMarker,
slog.F("record_type", recorder.RecordTypeTokenUsage),
slog.F("interception_id", intcID.String()),
slog.F("msg_id", in.GetMsgId()),
slog.F("provider_model", in.GetProviderModel()),
slog.F("input_tokens", in.GetInputTokens()),
slog.F("output_tokens", in.GetOutputTokens()),
slog.F("cache_read_input_tokens", in.GetCacheReadInputTokens()),
slog.F("cache_write_input_tokens", in.GetCacheWriteInputTokens()),
slog.F("created_at", in.GetCreatedAt().AsTime()),
slog.F("metadata", metadata),
)
}
if err := validateTokenUsage(in); err != nil {
s.logger.Error(ctx, "implausible token usage, discarding record",
slog.F("interception_id", intcID),
slog.F("input_tokens", in.GetInputTokens()),
slog.F("output_tokens", in.GetOutputTokens()),
slog.F("cache_read_input_tokens", in.GetCacheReadInputTokens()),
slog.F("cache_write_input_tokens", in.GetCacheWriteInputTokens()),
slog.Error(err))
return nil, xerrors.Errorf("validate token usage for interception %q: %w", intcID, err)
}
out, err := json.Marshal(metadata)
if err != nil {
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
}
// The interception is always recorded before any of its token usages,
// so it must exist. It carries the provider, model, and initiator needed
// for cost attribution.
intc, err := s.store.GetAIBridgeInterceptionByID(ctx, intcID)
if err != nil {
return nil, xerrors.Errorf("get interception %q: %w", intcID, err)
}
// Snapshot the effective group, per-token prices and compute cost. A
// missing price row or no effective group yields NULL columns.
cost, err := s.resolveTokenUsageCost(ctx, intc, in)
if err != nil {
return nil, xerrors.Errorf("resolve token usage cost: %w", err)
}
if err := s.recordTokenUsageAndSpend(ctx, intc, cost, in, out); err != nil {
return nil, xerrors.Errorf("record token usage and spend: %w", err)
}
return &proto.RecordTokenUsageResponse{}, nil
}
// recordTokenUsageAndSpend atomically records the token usage (including the
// interception's cost) and, when the user is budgeted and the computed cost is
// positive, accumulates that cost into the user's daily spend.
func (s *Server) recordTokenUsageAndSpend(ctx context.Context, intc database.AIBridgeInterception, cost tokenUsageCost, in *proto.RecordTokenUsageRequest, metadataJSON []byte) error {
createdAt := s.recordTime(ctx, "created_at", in.GetCreatedAt(), intc.ID.String())
// Populated inside the transaction with any budget thresholds this
// interception crossed.
var crossings []budgetThresholdCrossing
err := s.store.InTx(func(tx database.Store) error {
if _, err := tx.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{
ID: uuid.New(),
InterceptionID: intc.ID,
ProviderResponseID: in.GetMsgId(),
InputTokens: in.GetInputTokens(),
OutputTokens: in.GetOutputTokens(),
CacheReadInputTokens: in.GetCacheReadInputTokens(),
CacheWriteInputTokens: in.GetCacheWriteInputTokens(),
Metadata: metadataJSON,
CreatedAt: createdAt,
EffectiveGroupID: cost.effectiveGroupID,
InputPriceMicros: cost.inputPriceMicros,
OutputPriceMicros: cost.outputPriceMicros,
CacheReadPriceMicros: cost.cacheReadPriceMicros,
CacheWritePriceMicros: cost.cacheWritePriceMicros,
CostMicros: cost.costMicros,
ProviderModel: sql.NullString{String: in.GetProviderModel(), Valid: in.GetProviderModel() != ""},
PricedModel: cost.pricedModel,
}); err != nil {
return xerrors.Errorf("insert token usage: %w", err)
}
// Skip the spend update when there is no effective group or the interception has no cost.
if !cost.effectiveGroupID.Valid || !cost.costMicros.Valid || cost.costMicros.Int64 <= 0 {
s.logger.Debug(ctx, "skipping spend update",
slog.F("interception_id", intc.ID),
slog.F("initiator_id", intc.InitiatorID),
slog.F("has_effective_group", cost.effectiveGroupID.Valid),
slog.F("has_cost", cost.costMicros.Valid),
slog.F("cost_micros", cost.costMicros.Int64),
)
return nil
}
if _, err := tx.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: intc.InitiatorID,
EffectiveGroupID: cost.effectiveGroupID.UUID,
// Day is derived from the record usage request CreatedAt
// so it matches the token usage row's created_at column.
Day: dbtime.StartOfDay(createdAt.UTC()),
CostMicros: cost.costMicros.Int64,
}); err != nil {
return xerrors.Errorf("increment user daily spend: %w", err)
}
// Threshold detection is best-effort: a failed read must not roll back
// the committed spend, so the error is logged rather than propagated.
var detectErr error
crossings, detectErr = s.detectBudgetThresholdCrossings(ctx, tx, intc, cost, createdAt)
if detectErr != nil {
s.logger.Error(ctx, "failed to detect AI budget threshold crossing",
slog.F("interception_id", intc.ID),
slog.F("initiator_id", intc.InitiatorID),
slog.Error(detectErr))
}
return nil
}, nil)
if err != nil {
return err
}
if err := s.notifyBudgetThresholdCrossings(ctx, crossings); err != nil {
s.logger.Error(ctx, "failed to send AI budget notifications",
slog.F("initiator_id", intc.InitiatorID),
slog.Error(err))
}
return nil
}
func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
intcID, err := uuid.Parse(in.GetInterceptionId())
if err != nil {
return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err)
}
metadata := metadataToMap(in.GetMetadata())
if s.structuredLogging {
s.logger.Info(ctx, recorder.InterceptionLogMarker,
slog.F("record_type", recorder.RecordTypePromptUsage),
slog.F("interception_id", intcID.String()),
slog.F("msg_id", in.GetMsgId()),
slog.F("prompt", in.GetPrompt()),
slog.F("created_at", in.GetCreatedAt().AsTime()),
slog.F("metadata", metadata),
)
}
out, err := json.Marshal(metadata)
if err != nil {
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
}
_, err = s.store.InsertAIBridgeUserPrompt(ctx, database.InsertAIBridgeUserPromptParams{
ID: uuid.New(),
InterceptionID: intcID,
ProviderResponseID: in.GetMsgId(),
Prompt: in.GetPrompt(),
Metadata: out,
CreatedAt: s.recordTime(ctx, "created_at", in.GetCreatedAt(), intcID.String()),
})
if err != nil {
return nil, xerrors.Errorf("insert user prompt: %w", err)
}
return &proto.RecordPromptUsageResponse{}, nil
}
func (s *Server) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageRequest) (*proto.RecordToolUsageResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
intcID, err := uuid.Parse(in.GetInterceptionId())
if err != nil {
return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err)
}
metadata := metadataToMap(in.GetMetadata())
if s.structuredLogging {
s.logger.Info(ctx, recorder.InterceptionLogMarker,
slog.F("record_type", recorder.RecordTypeToolUsage),
slog.F("interception_id", intcID.String()),
slog.F("msg_id", in.GetMsgId()),
slog.F("tool_call_id", in.GetToolCallId()),
slog.F("item_id", in.GetItemId()),
slog.F("tool", in.GetTool()),
slog.F("input", in.GetInput()),
slog.F("server_url", in.GetServerUrl()),
slog.F("injected", in.GetInjected()),
slog.F("invocation_error", in.GetInvocationError()),
slog.F("created_at", in.GetCreatedAt().AsTime()),
slog.F("metadata", metadata),
)
}
out, err := json.Marshal(metadata)
if err != nil {
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
}
_, err = s.store.InsertAIBridgeToolUsage(ctx, database.InsertAIBridgeToolUsageParams{
ID: uuid.New(),
InterceptionID: intcID,
ProviderResponseID: in.GetMsgId(),
ProviderToolCallID: sql.NullString{String: in.GetToolCallId(), Valid: in.GetToolCallId() != ""},
ProviderItemID: sql.NullString{String: in.GetItemId(), Valid: in.GetItemId() != ""},
ServerUrl: sql.NullString{String: in.GetServerUrl(), Valid: in.ServerUrl != nil},
Tool: in.GetTool(),
Input: in.GetInput(),
Injected: in.GetInjected(),
InvocationError: sql.NullString{String: in.GetInvocationError(), Valid: in.InvocationError != nil},
Metadata: out,
CreatedAt: s.recordTime(ctx, "created_at", in.GetCreatedAt(), intcID.String()),
})
if err != nil {
return nil, xerrors.Errorf("insert tool usage: %w", err)
}
return &proto.RecordToolUsageResponse{}, nil
}
func (s *Server) RecordModelThought(ctx context.Context, in *proto.RecordModelThoughtRequest) (*proto.RecordModelThoughtResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
intcID, err := uuid.Parse(in.GetInterceptionId())
if err != nil {
return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err)
}
metadata := metadataToMap(in.GetMetadata())
if s.structuredLogging {
s.logger.Info(ctx, recorder.InterceptionLogMarker,
slog.F("record_type", recorder.RecordTypeModelThought),
slog.F("interception_id", intcID.String()),
slog.F("content", in.GetContent()),
slog.F("created_at", in.GetCreatedAt().AsTime()),
slog.F("metadata", metadata),
)
}
out, err := json.Marshal(metadata)
if err != nil {
s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
}
_, err = s.store.InsertAIBridgeModelThought(ctx, database.InsertAIBridgeModelThoughtParams{
InterceptionID: intcID,
Content: in.GetContent(),
Metadata: out,
CreatedAt: s.recordTime(ctx, "created_at", in.GetCreatedAt(), intcID.String()),
})
if err != nil {
return nil, xerrors.Errorf("insert model thought: %w", err)
}
return &proto.RecordModelThoughtResponse{}, nil
}
// findInterceptionLineage looks up the parent interception and the root
// of the thread by finding which interception recorded a tool usage with
// the given tool call ID. Returns (parentID, rootID); both will be
// uuid.Nil if no match is found or the tool call ID is empty.
func (s *Server) findInterceptionLineage(ctx context.Context, toolCallID string) (parent uuid.UUID, root uuid.UUID) {
if toolCallID == "" {
return uuid.Nil, uuid.Nil
}
lineage, err := s.store.GetAIBridgeInterceptionLineageByToolCallID(ctx, toolCallID)
if err != nil {
s.logger.Warn(ctx, "failed to retrieve interception lineage",
slog.Error(err), slog.F("tool_call_id", toolCallID))
return uuid.Nil, uuid.Nil
}
return lineage.ThreadParentID, lineage.ThreadRootID
}
func (s *Server) GetMCPServerConfigs(_ context.Context, _ *proto.GetMCPServerConfigsRequest) (*proto.GetMCPServerConfigsResponse, error) {
cfgs := make([]*proto.MCPServerConfig, 0, len(s.externalAuthConfigs))
for _, eac := range s.externalAuthConfigs {
var allowlist, denylist string
if eac.MCPToolAllowRegex != nil {
allowlist = eac.MCPToolAllowRegex.String()
}
if eac.MCPToolDenyRegex != nil {
denylist = eac.MCPToolDenyRegex.String()
}
cfgs = append(cfgs, &proto.MCPServerConfig{
Id: eac.ID,
Url: eac.MCPURL,
ToolAllowRegex: allowlist,
ToolDenyRegex: denylist,
})
}
return &proto.GetMCPServerConfigsResponse{
CoderMcpConfig: s.coderMCPConfig, // it's fine if this is nil
ExternalAuthMcpConfigs: cfgs,
}, nil
}
func (s *Server) GetMCPServerAccessTokensBatch(ctx context.Context, in *proto.GetMCPServerAccessTokensBatchRequest) (*proto.GetMCPServerAccessTokensBatchResponse, error) {
if len(in.GetMcpServerConfigIds()) == 0 {
return &proto.GetMCPServerAccessTokensBatchResponse{}, nil
}
userID, err := uuid.Parse(in.GetUserId())
if err != nil {
return nil, xerrors.Errorf("parse user_id: %w", err)
}
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
links, err := s.store.GetExternalAuthLinksByUserID(ctx, userID)
if err != nil {
return nil, xerrors.Errorf("fetch external auth links: %w", err)
}
if len(links) == 0 {
return &proto.GetMCPServerAccessTokensBatchResponse{}, nil
}
// Ensure unique to prevent unnecessary effort.
ids := in.GetMcpServerConfigIds()
slices.Sort(ids)
ids = slices.Compact(ids)
var (
wg sync.WaitGroup
errs error
mu sync.Mutex
tokens = make(map[string]string, len(ids))
tokenErrs = make(map[string]string)
)
externalAuthLoop:
for _, id := range ids {
eac, ok := s.externalAuthConfigs[id]
if !ok {
mu.Lock()
s.logger.Warn(ctx, "no MCP server config found by given ID", slog.F("id", id))
tokenErrs[id] = ErrNoMCPConfigFound.Error()
mu.Unlock()
continue
}
for _, link := range links {
if link.ProviderID != eac.ID {
continue
}
// Validate all configured External Auth links concurrently.
wg.Add(1)
go func() {
defer wg.Done()
// TODO: timeout.
valid, _, validateErr := eac.ValidateToken(ctx, link.OAuthToken())
mu.Lock()
defer mu.Unlock()
if !valid {
// TODO: attempt refresh.
s.logger.Warn(ctx, "invalid/expired access token, cannot auto-configure MCP", slog.F("provider", link.ProviderID), slog.Error(validateErr))
tokenErrs[id] = ErrExpiredOrInvalidOAuthToken.Error()
return
}
if validateErr != nil {
errs = multierror.Append(errs, validateErr)
tokenErrs[id] = validateErr.Error()
} else {
tokens[id] = link.OAuthAccessToken
}
}()
continue externalAuthLoop
}
// No link found for this external auth config, so include a generic
// error.
mu.Lock()
tokenErrs[id] = ErrNoExternalAuthLinkFound.Error()
mu.Unlock()
}
wg.Wait()
return &proto.GetMCPServerAccessTokensBatchResponse{
AccessTokens: tokens,
Errors: tokenErrs,
}, errs
}
// IsAuthorized validates a given Coder API key and returns the user ID to which it belongs (if valid).
// It also returns the organization the request resolves to, as determined by [Server.resolveOrganization], and
// rejects the request when that resolution fails. An error caused by coderd itself, such as a failed database
// query, carries [proto.ErrCodeInternal] so the gateway can tell it from a rejected request.
//
// SECURITY: when in.KeyId is set (the "delegated" path), this method trusts the
// caller's claim of identity and skips the key-secret check. This DRPCServer is
// reachable both in-process via [aibridged.MemTransportPipe] and over the network
// via the /api/v2/ai-gateway/serve endpoint. That endpoint admits only holders of
// AI Gateway key, which are fully trusted. Standalone AI Gateway authenticates its
// own users and acts on their behalf, much like a provisioner daemon. A Gateway key
// holder can therefore act as any user without that user's secret. Per-user
// authorization on this surface is a known gap.
//
// NOTE: this should really be using the code from [httpmw.ExtractAPIKey]. That function not only validates the key
// but handles many other cases like updating last used, expiry, etc. This code does not currently use it for
// a few reasons:
//
// 1. [httpmw.ExtractAPIKey] relies on keys being given in specific headers [httpmw.APITokenFromRequest] which AI
// bridge requests will not conform to.
// 2. The code mixes many different concerns, and handles HTTP responses too, which is undesirable here.
// 3. The core logic would need to be extracted, but that will surely be a complex & time-consuming distraction right now.
// 4. Once we have an Early Access release of AI Bridge, we need to return to this.
//
// TODO: replace with logic from [httpmw.ExtractAPIKey].
func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest) (*proto.IsAuthorizedResponse, error) {
resp, err := s.authorize(ctx, in)
var r rejectionError
if err != nil && !errors.As(err, &r) {
return nil, drpcerr.WithCode(err, proto.ErrCodeInternal)
}
return resp, err
}
func (s *Server) authorize(ctx context.Context, in *proto.IsAuthorizedRequest) (*proto.IsAuthorizedResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
var (
keyID string
keySecret string
// delegated requests skip the secret check: the caller never
// has the secret. Trust is established at the in-process
// transport boundary, not in this RPC.
delegated bool
delegatedOrgID uuid.NullUUID
)
switch {
case in.GetKey() != "" && in.GetKeyId() != "":
return nil, ErrAmbiguousAuth
case in.GetKey() != "" && in.GetDelegatedOrganizationId() != "":
return nil, ErrAmbiguousAuth
case in.GetKeyId() != "":
keyID = in.GetKeyId()
delegated = true
if raw := in.GetDelegatedOrganizationId(); raw != "" {
id, err := uuid.Parse(raw)
if err != nil || id == uuid.Nil {
return nil, ErrInvalidDelegatedOrganization
}
delegatedOrgID = uuid.NullUUID{UUID: id, Valid: true}
}
default:
var err error
keyID, keySecret, err = httpmw.SplitAPIToken(in.GetKey())
if err != nil {
return nil, ErrInvalidKey
}
}
// Key exists.
key, err := s.store.GetAPIKeyByID(ctx, keyID)
if errors.Is(err, sql.ErrNoRows) {
return nil, ErrUnknownKey
}
if err != nil {
s.logger.Warn(ctx, "failed to retrieve API key by id", slog.F("key_id", keyID), slog.Error(err))
return nil, xerrors.Errorf("get API key %s: %w", keyID, err)
}
// Key has not expired.
now := dbtime.Now()
if key.ExpiresAt.Before(now) {
return nil, ErrExpired
}
// Key secret matches (skipped for delegated callers).
if !delegated && !apikey.ValidateHash(key.HashedSecret, keySecret) {
return nil, ErrInvalidKey
}
// User exists.
user, err := s.store.GetUserByID(ctx, key.UserID)
if errors.Is(err, sql.ErrNoRows) {
return nil, ErrUnknownUser
}
if err != nil {
s.logger.Warn(ctx, "failed to retrieve API key user", slog.F("key_id", keyID), slog.F("user_id", key.UserID), slog.Error(err))
return nil, xerrors.Errorf("get user %s: %w", key.UserID, err)
}
// User is active, not deleted, and not a system user.
if user.Deleted {
return nil, ErrDeletedUser
}
if user.Status != database.UserStatusActive {
return nil, ErrInactiveUser
}
if user.IsSystem {
return nil, ErrSystemUser
}
resp := &proto.IsAuthorizedResponse{
OwnerId: key.UserID.String(),
ApiKeyId: key.ID,
Username: user.Username,
Email: user.Email,
}
var workspaceID uuid.NullUUID
if !delegated {
id, ok, err := workspaceAttribution(key)
if err != nil {
return nil, err
}
if ok {
workspaceID = uuid.NullUUID{UUID: id, Valid: true}
resp.WorkspaceId = id.String()
}
}
org, err := s.resolveOrganization(ctx, user.ID, orgResolverInput{WorkspaceID: workspaceID, DelegatedOrganizationID: delegatedOrgID})
if err != nil {
s.logger.Warn(ctx, "failed to resolve request organization", slog.F("key_id", keyID), slog.F("user_id", user.ID), slog.Error(err))
return nil, err
}
resp.OrganizationId = org.ID.String()
resp.OrganizationSource = string(org.Source)
return resp, nil
}
// workspaceAttribution parses the workspace UUID from a strict server-minted
// token name without performing any database lookup. It returns the workspace
// UUID and true when the token name matches the expected pattern and the
// embedded owner UUID equals key.UserID. It fails closed (ErrWorkspaceAttribution)
// when the embedded owner does not match.
func workspaceAttribution(key database.APIKey) (uuid.UUID, bool, error) {
if key.LoginType == database.LoginTypeToken {
return uuid.Nil, false, nil
}
ownerID, workspaceID, ok := provisionerdserver.ParseWorkspaceSessionTokenName(key.TokenName)
if !ok {
return uuid.Nil, false, nil
}
if ownerID != key.UserID {
return uuid.Nil, false, ErrWorkspaceAttribution
}
return workspaceID, true, nil
}
// interceptionAttribution parses the optional workspace_id from the
// interception request. The returned NullUUID is valid only when a non-empty
// workspace_id was provided.
func interceptionAttribution(in *proto.RecordInterceptionRequest) (workspaceID uuid.NullUUID, err error) {
if in.GetWorkspaceId() == "" {
return uuid.NullUUID{}, nil
}
id, err := uuid.Parse(in.GetWorkspaceId())
if err != nil {
return uuid.NullUUID{}, xerrors.Errorf("invalid workspace ID: %w", err)
}
return uuid.NullUUID{UUID: id, Valid: true}, nil
}
// interceptionOrganization parses the optional organization attribution. It
// returns both values valid or neither, so an invalid pair never fails the record.
// It checks only the format. The caller must send an existing organization: a
// well-formed ID for an unknown organization passes here and fails the foreign
// key on insert.
func interceptionOrganization(in *proto.RecordInterceptionRequest) (uuid.NullUUID, database.NullAIBridgeInterceptionOrganizationSource, error) {
rawID, rawSource := in.GetOrganizationId(), in.GetOrganizationSource()
if rawID == "" && rawSource == "" {
return uuid.NullUUID{}, database.NullAIBridgeInterceptionOrganizationSource{}, nil
}
id, err := uuid.Parse(rawID)
if err != nil {
return uuid.NullUUID{}, database.NullAIBridgeInterceptionOrganizationSource{}, xerrors.Errorf("invalid organization ID: %w", err)
}
source := database.AIBridgeInterceptionOrganizationSource(rawSource)
if !source.Valid() {
return uuid.NullUUID{}, database.NullAIBridgeInterceptionOrganizationSource{}, xerrors.Errorf("unknown organization source %q", rawSource)
}
return uuid.NullUUID{UUID: id, Valid: true},
database.NullAIBridgeInterceptionOrganizationSource{AIBridgeInterceptionOrganizationSource: source, Valid: true},
nil
}
// IsBudgetExceeded reports whether the user's AI spend has reached their
// effective limit over [periodStart, now], where periodStart is the start of
// the current deployment-configured budget period.
func (s *Server) IsBudgetExceeded(ctx context.Context, in *proto.IsBudgetExceededRequest) (resp *proto.IsBudgetExceededResponse, retErr error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
start := s.clock.Now()
defer func() {
if s.metrics == nil {
return
}
outcome := "allowed"
switch {
case retErr != nil:
outcome = "error"
case resp != nil && resp.Exceeded:
outcome = "blocked"
}
s.metrics.EnforcementDuration.WithLabelValues(outcome).Observe(s.clock.Since(start).Seconds())
}()
userID, err := uuid.Parse(in.GetUserId())
if err != nil {
return nil, xerrors.Errorf("invalid user_id %q: %w", in.GetUserId(), err)
}
// The organization is the one IsAuthorized resolved for this request, and
// the gateway sends it back unchanged. It is trusted to be valid and to
// contain the user, so it is not checked again here. From v1.6 a gateway
// always sends it. Only gateways older than v1.6 send it empty, which
// means site-wide resolution.
var organizationID uuid.NullUUID
if raw := in.GetOrganizationId(); raw != "" {
id, err := uuid.Parse(raw)
if err != nil {
return nil, xerrors.Errorf("invalid organization_id %q: %w", raw, err)
}
organizationID = uuid.NullUUID{UUID: id, Valid: true}
}
periodWindow, err := budget.CurrentPeriod(s.clock.Now(), s.budgetPeriod)
if err != nil {
return nil, xerrors.Errorf("compute AI budget period: %w", err)
}
userBudget, err := s.checkUserAIBudget(ctx, userID, organizationID, periodWindow.Start)
if err != nil {
return nil, err
}
if userBudget.Exceeded && s.metrics != nil {
s.metrics.BlockedRequests.WithLabelValues(userBudget.GroupID.String()).Inc()
}
return &proto.IsBudgetExceededResponse{
Exceeded: userBudget.Exceeded,
SpendLimitMicros: userBudget.SpendLimitMicros,
}, nil
}
// userAIBudget is a snapshot of a user's AI budget status. SpendLimitMicros
// is nil when no budget is configured for the user (unlimited). GroupID is the
// effective group the limit resolved to, set only when a limit applies.
type userAIBudget struct {
Exceeded bool
SpendLimitMicros *int64
GroupID uuid.UUID
}
// checkUserAIBudget evaluates the user's AI budget status aggregated over
// [periodStart, now]. When organizationID is set, the budget is resolved
// inside that organization; otherwise it is resolved site-wide.
//
// Note: there is a potential race condition where two concurrent requests
// from the same user can both pass the check if processed in parallel,
// allowing brief overage. This is acceptable because:
// - Cost is only known after the LLM API returns.
// - Overage is bounded by request cost × concurrency; once the accumulated
// spend crosses the limit, subsequent requests are blocked.
// - Cost accounting is advisory, not strict. The goal is to prevent
// overages, not build an accounting system.
// - Fail-open is acceptable for this case.
func (s *Server) checkUserAIBudget(ctx context.Context, userID uuid.UUID, organizationID uuid.NullUUID, periodStart time.Time) (userAIBudget, error) {
effectiveGroup, ok, err := budget.ResolveUserAIBudget(ctx, s.store, userID, organizationID, s.budgetPolicy)
if err != nil {
return userAIBudget{}, xerrors.Errorf("resolve effective AI budget for user %q with budget policy %q: %w", userID, s.budgetPolicy, err)
}
// ok is false when no budget is configured. The nil Limit check keeps
// enforcement failing open if a caller resolves via the unlimited
// Everyone fallback.
if !ok || effectiveGroup.Limit == nil {
// No enforceable spend limit for the user; return zero-valued status.
return userAIBudget{}, nil
}
spend, err := s.store.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{
UserID: userID,
EffectiveGroupID: effectiveGroup.GroupID,
PeriodStart: periodStart,
})
if err != nil {
return userAIBudget{}, xerrors.Errorf("get user AI spend for user %q in group %q: %w", userID, effectiveGroup.GroupID, err)
}
exceeded := spend.SpendMicros >= effectiveGroup.Limit.SpendLimitMicros
logger := s.logger.With(
slog.F("user_id", userID),
slog.F("organization_id", organizationID),
slog.F("effective_group_id", effectiveGroup.GroupID),
slog.F("period_start", periodStart),
slog.F("current_spend_micros", spend.SpendMicros),
slog.F("spend_limit_micros", effectiveGroup.Limit.SpendLimitMicros),
slog.F("exceeded", exceeded),
)
logger.Debug(ctx, "user AI spend status")
if exceeded {
logger.Warn(ctx, "user AI budget exceeded")
}
return userAIBudget{
Exceeded: exceeded,
SpendLimitMicros: new(effectiveGroup.Limit.SpendLimitMicros),
GroupID: effectiveGroup.GroupID,
}, nil
}
// GetAIProviders returns the full AI provider set (enabled and disabled) from
// the database. Embedded and standalone AI Gateway daemons call this over DRPC
// to build their provider pool instead of reading the database directly.
//
// Keys are populated only for enabled providers; disabled providers never call
// upstream, so their secrets are withheld.
//
// SECURITY: the response carries plaintext API keys and Bedrock credentials.
// Do not log the response struct.
func (s *Server) GetAIProviders(ctx context.Context, _ *proto.GetAIProvidersRequest) (*proto.GetAIProvidersResponse, error) {
//nolint:gocritic // AIBridged has a minimal permission set scoped to AI Bridge queries.
ctx = dbauthz.AsAIBridged(ctx)
var (
rows []database.AIProvider
keysByProvider map[uuid.UUID][]database.AIProviderKey
)
// Wrap both reads in a read-only transaction so the provider list and the
// key list are consistent with each other.
err := s.store.InTx(func(tx database.Store) error {
var err error
rows, err = tx.GetAIProviders(ctx, database.GetAIProvidersParams{IncludeDisabled: true})
if err != nil {
return xerrors.Errorf("get ai providers: %w", err)
}
// Load keys only for enabled providers to avoid materializing secrets
// for disabled rows.
ids := make([]uuid.UUID, 0, len(rows))
for _, row := range rows {
if !row.Enabled {
continue
}
ids = append(ids, row.ID)
}
keysByProvider = make(map[uuid.UUID][]database.AIProviderKey, len(ids))
if len(ids) == 0 {
return nil
}
keyRows, err := tx.GetAIProviderKeysByProviderIDs(ctx, ids)
if err != nil {
return xerrors.Errorf("get ai provider keys: %w", err)
}
for _, k := range keyRows {
keysByProvider[k.ProviderID] = append(keysByProvider[k.ProviderID], k)
}
return nil
}, &database.TxOptions{ReadOnly: true, TxIdentifier: "get_ai_providers"})
if err != nil {
return nil, err
}
providers := make([]*proto.AIProvider, 0, len(rows))
for _, row := range rows {
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
// configuration for every gateway, which would otherwise loop
// forever on the empty pool.
s.logger.Error(ctx, "skipping ai provider with invalid settings; it will be absent from the gateway pool",
slog.F("provider_id", row.ID),
slog.F("provider_name", row.Name),
slog.F("provider_type", string(row.Type)),
slog.Error(err),
)
continue
}
providers = append(providers, p)
}
return &proto.GetAIProvidersResponse{Providers: providers}, nil
}
// WatchAIProviders streams a payload-free change signal on each
// AIProvidersChangedChannel event, plus one immediately on subscribe. Pubsub
// drop errors produce a signal rather than failing the stream. Blocks until the
// stream context or the server lifecycle is canceled.
func (s *Server) WatchAIProviders(_ *proto.WatchAIProvidersRequest, stream proto.DRPCProviderConfigurator_WatchAIProvidersStream) error {
if s.pubsub == nil {
return xerrors.New("pubsub not configured")
}
ctx, cancel := context.WithCancel(stream.Context())
defer cancel()
// Cancel when the server lifecycle ends, not just when the stream closes.
stop := context.AfterFunc(s.lifecycleCtx, cancel)
defer stop()
// Buffered to one so a burst of events collapses into a single pending
// signal.
signals := make(chan struct{}, 1)
notify := func() {
select {
case signals <- struct{}{}:
default:
}
}
// Every event signals, including dropped-message errors.
unsubscribe, err := s.pubsub.SubscribeWithErr(coderdpubsub.AIProvidersChangedChannel, func(cbCtx context.Context, _ []byte, err error) {
if err != nil {
s.logger.Warn(cbCtx, "ai providers changed event delivered with error", slog.Error(err))
}
notify()
})
if err != nil {
return xerrors.Errorf("subscribe to %s: %w", coderdpubsub.AIProvidersChangedChannel, err)
}
defer unsubscribe()
// Initial signal on subscribe.
notify()
for {
select {
case <-ctx.Done():
return nil
case <-signals:
if err := stream.Send(&proto.WatchAIProvidersResponse{}); err != nil {
return xerrors.Errorf("send ai providers change signal: %w", err)
}
}
}
}
// Deprecated: Injected MCP in AI Bridge is deprecated and will be removed in a future release.
//
//nolint:revive // The flag is fixed for the life of the process.
func getCoderMCPServerConfig(experiments codersdk.Experiments, oauth2ProviderEnabled bool, accessURL string) (*proto.MCPServerConfig, error) {
// The internal MCP server needs the MCP experiment and the OAuth2 provider.
if !experiments.Enabled(codersdk.ExperimentMCPServerHTTP) {
return nil, xerrors.Errorf("%q experiment not enabled", codersdk.ExperimentMCPServerHTTP)
}
if !oauth2ProviderEnabled {
return nil, xerrors.New("OAuth2 provider is disabled; set CODER_OAUTH2_PROVIDER_ENABLE=true")
}
u, err := url.JoinPath(accessURL, codermcp.MCPEndpoint)
if err != nil {
return nil, xerrors.Errorf("build MCP URL with %q: %w", accessURL, err)
}
return &proto.MCPServerConfig{
Id: aibridged.InternalMCPServerID,
Url: u,
}, nil
}
// credentialKindOrDefault converts the proto credential kind string to
// the database enum, defaulting to "centralized" when the value is
// empty or not a valid enum member.
func credentialKindOrDefault(kind string) database.CredentialKind {
ck := database.CredentialKind(kind)
if !ck.Valid() {
return database.CredentialKindCentralized
}
return ck
}
// maxErrorMessageBytes caps the interception error message stored in the
// database, enforced at this trust boundary regardless of caller behavior.
const maxErrorMessageBytes = 1024
// truncateErrorMessage caps msg to maxErrorMessageBytes, dropping any partial
// trailing rune so the stored value stays valid UTF-8.
func truncateErrorMessage(msg string) string {
if len(msg) <= maxErrorMessageBytes {
return msg
}
return strings.ToValidUTF8(msg[:maxErrorMessageBytes], "")
}
// interceptionErrorType maps the wire error type onto the nullable DB enum. An
// empty value yields NULL (a successful interception). A non-empty but
// unrecognized value (e.g. version skew where the client knows an enum the DB
// migration does not yet) is stored as 'unknown' rather than NULL, so the row's
// error columns stay consistent with a recorded error_message.
func interceptionErrorType(errType string) database.NullAIBridgeInterceptionErrorType {
if errType == "" {
return database.NullAIBridgeInterceptionErrorType{}
}
et := database.AIBridgeInterceptionErrorType(errType)
if !et.Valid() {
et = database.AibridgeInterceptionErrorTypeUnknown
}
return database.NullAIBridgeInterceptionErrorType{
AIBridgeInterceptionErrorType: et,
Valid: true,
}
}
func metadataToMap(in map[string]*anypb.Any) map[string]any {
meta := make(map[string]any, len(in))
for k, v := range in {
if v == nil {
continue
}
var sv structpb.Value
if err := v.UnmarshalTo(&sv); err == nil {
meta[k] = sv.AsInterface()
}
}
return meta
}
// parseOptionalUUID converts an optional proto string to uuid.NullUUID.
// Returns a zero NullUUID if s is nil. If s is non-nil but not a valid UUID, it
// returns a zero NullUUID along with the parse error so the caller can decide
// how to surface it.
func parseOptionalUUID(s *string) (uuid.NullUUID, error) {
if s == nil {
return uuid.NullUUID{}, nil
}
id, err := uuid.Parse(*s)
if err != nil {
return uuid.NullUUID{}, err
}
return uuid.NullUUID{UUID: id, Valid: true}, nil
}
// parseOptionalInt32 converts an optional proto int32 to sql.NullInt32.
func parseOptionalInt32(n *int32) sql.NullInt32 {
if n == nil {
return sql.NullInt32{}
}
return sql.NullInt32{Int32: *n, Valid: true}
}
// 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.
func aiProviderToProto(row database.AIProvider, keys []database.AIProviderKey) (*proto.AIProvider, error) {
p := &proto.AIProvider{
Name: row.Name,
Type: string(row.Type),
Enabled: row.Enabled,
BaseUrl: row.BaseUrl,
}
// Disabled providers are rendered as stubs by the client and never call
// upstream, so only the identity fields are returned; keys and settings
// (including Bedrock credentials) are withheld.
if !row.Enabled {
return p, nil
}
p.Keys = make([]string, 0, len(keys))
for _, k := range keys {
p.Keys = append(p.Keys, k.APIKey)
}
settings, err := db2sdk.AIProviderSettings(row.Settings)
if err != nil {
return nil, xerrors.Errorf("decode settings: %w", err)
}
if settings.Bedrock != nil {
p.Bedrock = &proto.AIProviderKindBedrock{
Region: settings.Bedrock.Region,
AccessKey: ptr.NilToEmpty(settings.Bedrock.AccessKey),
AccessKeySecret: ptr.NilToEmpty(settings.Bedrock.AccessKeySecret),
Model: settings.Bedrock.Model,
SmallFastModel: settings.Bedrock.SmallFastModel,
RoleArn: settings.Bedrock.RoleARN,
ExternalId: settings.Bedrock.ExternalID,
Protocol: string(settings.Bedrock.Protocol),
ResolvedModel: settings.Bedrock.ResolvedModel,
ResolvedSmallFastModel: settings.Bedrock.ResolvedSmallFastModel,
}
}
if settings.ClaudePlatformAWS != nil {
p.ClaudePlatformAws = &proto.AIProviderKindClaudePlatformAWS{
Region: settings.ClaudePlatformAWS.Region,
WorkspaceId: settings.ClaudePlatformAWS.WorkspaceID,
}
}
return p, nil
}