feat: route chatd provider traffic through aibridge (#25629)

## Summary

Routes chatd model calls backed by concrete AI Provider rows through the
in-process aibridge transport by default, with deployment options to use
direct provider routing when AI Gateway is disabled or chat AI Gateway
routing is disabled.

- Splits model routing into common, direct provider, and AI Gateway
paths behind a single deployment-mode entry point.
- Builds chatd models through explicit request, route, and options data.
Active API key attribution is passed explicitly instead of being hidden
inside generic model construction.
- For AI Gateway BYOK routes, resolves the user's provider key in chatd,
forwards it through provider-specific auth headers, and sets
`X-Coder-AI-Governance-Token` to the `delegated` marker so aibridge
preserves those headers while still stripping Coder-specific metadata.
- Keeps central provider credentials and deployment fallback credentials
out of forwarded provider auth headers, so AI Gateway central policy
remains authoritative.
- Redacts delegated provider auth from default string formatting to
avoid accidental plaintext logging of user BYOK credentials.
- Covers selected chat models, advisor overrides, title and quickgen
paths, subagent overrides, computer use model selection, and an
integration-style chat turn through the aibridge transport path.
- Persists initiating API key IDs on chat and queued user messages,
including subagent child messages, and fails closed for AI
Gateway-routed model builds without an active key.
- Removes unused `api_key_id` indexes while keeping the persistence
columns and foreign keys.
- Keeps the deployment option available through config and env parsing,
but hides it from CLI help and generated docs.
- Stabilizes the subagent poll fallback test so background CreateChat
processing cannot win the state transition under slower CI environments.

## Tests

- `go test ./coderd/x/chatd -run
'TestAIGatewayProviderAuthForUser|TestAIGatewayProviderAuthRedactsFormatting|TestResolveModelRouteForConfigAIGatewayProviderAuth|TestAIGatewayModelForwardsProviderAuth|TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey|TestAwaitSubagentCompletion'
-count=1`
- `go test ./coderd/aibridged -run
'TestServeHTTP_DelegatedAPIKey|TestServeHTTP_StripCoderToken' -count=1`
- `git diff --check HEAD~1..HEAD`
- `make lint`

> Mux working on behalf of Mike.
This commit is contained in:
Michael Suchacz
2026-05-26 19:31:52 +00:00
committed by GitHub
parent a56c88a0cc
commit 8b1705eb65
31 changed files with 2463 additions and 377 deletions
+66 -77
View File
@@ -4,7 +4,6 @@ import (
"context"
"errors"
"fmt"
"net/http"
"slices"
"strings"
"time"
@@ -69,9 +68,39 @@ var preferredTitleModels = []struct {
type shortTextCandidate struct {
provider string
model string
route resolvedModelRoute
lm fantasy.LanguageModel
}
func (p *Server) preferredShortTextCandidates(
chat database.Chat,
keys chatprovider.ProviderAPIKeys,
) []shortTextCandidate {
if p.shouldUseAIGatewayRouting() {
return nil
}
candidates := make([]shortTextCandidate, 0, len(preferredTitleModels)+1)
userAgent := chatprovider.UserAgent()
extraHeaders := chatprovider.CoderHeaders(chat)
for _, candidate := range preferredTitleModels {
model, err := chatprovider.ModelFromConfig(
candidate.provider, candidate.model, keys, userAgent,
extraHeaders,
nil,
)
if err == nil {
candidates = append(candidates, shortTextCandidate{
provider: candidate.provider,
model: candidate.model,
route: newDirectModelRoute(candidate.provider, keys),
lm: model,
})
}
}
return candidates
}
func selectPreferredConfiguredShortTextModelConfig(
configs []database.ChatModelConfig,
) (database.ChatModelConfig, bool) {
@@ -121,7 +150,9 @@ func (p *Server) maybeGenerateChatTitle(
fallbackProvider string,
fallbackModelName string,
fallbackModel fantasy.LanguageModel,
fallbackRoute resolvedModelRoute,
keys chatprovider.ProviderAPIKeys,
modelOpts modelBuildOptions,
generatedTitle *generatedChatTitle,
logger slog.Logger,
debugSvc *chatdebug.Service,
@@ -135,10 +166,11 @@ func (p *Server) maybeGenerateChatTitle(
titleCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
overrideConfig, overrideModel, overrideKeys, overrideSet, overrideErr := p.resolveTitleGenerationModelOverride(
overrideConfig, overrideModel, _, overrideRoute, overrideSet, overrideErr := p.resolveTitleGenerationModelOverride(
titleCtx,
chat,
keys,
modelOpts,
)
if overrideErr != nil {
if overrideSet {
@@ -161,29 +193,15 @@ func (p *Server) maybeGenerateChatTitle(
candidates = []shortTextCandidate{{
provider: overrideConfig.Provider,
model: overrideConfig.Model,
route: overrideRoute,
lm: overrideModel,
}}
} else {
// Build candidate list: preferred lightweight models first,
// then the user's chat model as last resort.
candidates = make([]shortTextCandidate, 0, len(preferredTitleModels)+1)
for _, c := range preferredTitleModels {
m, err := chatprovider.ModelFromConfig(
c.provider, c.model, keys, chatprovider.UserAgent(),
chatprovider.CoderHeaders(chat),
nil,
)
if err == nil {
candidates = append(candidates, shortTextCandidate{
provider: c.provider,
model: c.model,
lm: m,
})
}
}
candidates = p.preferredShortTextCandidates(chat, keys)
candidates = append(candidates, shortTextCandidate{
provider: fallbackProvider,
model: fallbackModelName,
route: fallbackRoute,
lm: fallbackModel,
})
}
@@ -213,17 +231,13 @@ func (p *Server) maybeGenerateChatTitle(
candidateCtx := titleCtx
candidateModel := candidate.lm
finishDebugRun := func(error) {}
candidateKeys := keys
if overrideSet {
candidateKeys = overrideKeys
}
if debugEnabled {
candidateCtx, candidateModel, finishDebugRun = prepareQuickgenDebugCandidate(
candidateCtx, candidateModel, finishDebugRun = p.prepareQuickgenDebugCandidate(
titleCtx,
chat,
candidateKeys,
debugSvc,
candidate,
modelOpts,
chatdebug.KindTitleGeneration,
triggerMessageID,
historyTipMessageID,
@@ -293,32 +307,26 @@ func (p *Server) maybeGenerateChatTitle(
}
}
func newQuickgenDebugModel(
func (p *Server) newQuickgenDebugModel(
ctx context.Context,
chat database.Chat,
keys chatprovider.ProviderAPIKeys,
debugSvc *chatdebug.Service,
provider string,
model string,
route resolvedModelRoute,
modelOpts modelBuildOptions,
) (fantasy.LanguageModel, error) {
httpClient := &http.Client{Transport: &chatdebug.RecordingTransport{}}
debugModel, err := chatprovider.ModelFromConfig(
provider,
model,
keys,
chatprovider.UserAgent(),
chatprovider.CoderHeaders(chat),
httpClient,
)
debugOpts := modelOpts
debugOpts.RecordHTTP = true
debugModel, err := p.newModel(ctx, modelClientRequest{
Chat: chat,
ModelName: model,
UserAgent: chatprovider.UserAgent(),
ExtraHeaders: chatprovider.CoderHeaders(chat),
}, route, debugOpts)
if err != nil {
return nil, err
}
if debugModel == nil {
return nil, xerrors.Errorf(
"create model for %s/%s returned nil",
provider,
model,
)
}
return chatdebug.WrapModel(debugModel, debugSvc, chatdebug.RecorderOptions{
ChatID: chat.ID,
@@ -328,12 +336,12 @@ func newQuickgenDebugModel(
}), nil
}
func prepareQuickgenDebugCandidate(
func (p *Server) prepareQuickgenDebugCandidate(
ctx context.Context,
chat database.Chat,
keys chatprovider.ProviderAPIKeys,
debugSvc *chatdebug.Service,
candidate shortTextCandidate,
modelOpts modelBuildOptions,
kind chatdebug.RunKind,
triggerMessageID int64,
historyTipMessageID int64,
@@ -345,12 +353,14 @@ func prepareQuickgenDebugCandidate(
return ctx, candidate.lm, finishDebugRun
}
debugModel, err := newQuickgenDebugModel(
debugModel, err := p.newQuickgenDebugModel(
ctx,
chat,
keys,
debugSvc,
candidate.provider,
candidate.model,
candidate.route,
modelOpts,
)
if err != nil {
logger.Warn(ctx, "failed to build short-text debug model",
@@ -393,18 +403,8 @@ func prepareQuickgenDebugCandidate(
return ctx, candidate.lm, finishDebugRun
}
runCtx := chatdebug.ContextWithRun(
ctx,
&chatdebug.RunContext{
RunID: run.ID,
ChatID: chat.ID,
TriggerMessageID: triggerMessageID,
HistoryTipMessageID: historyTipMessageID,
Kind: kind,
Provider: candidate.provider,
Model: candidate.model,
},
)
runContext := chatdebugRunContext(run)
runCtx := chatdebug.ContextWithRun(ctx, &runContext)
finishDebugRun = func(runErr error) {
if finalizeErr := debugSvc.FinalizeRun(ctx, chatdebug.FinalizeRunParams{
RunID: run.ID,
@@ -824,7 +824,7 @@ const turnStatusLabelPrompt = "You write compact chat status labels for a sideba
// message text. It follows the same candidate-selection strategy
// as title generation: try preferred lightweight models first, then
// fall back to the provided model. Returns "" on any failure.
func generateTurnStatusLabel(
func (p *Server) generateTurnStatusLabel(
ctx context.Context,
chat database.Chat,
status database.ChatStatus,
@@ -832,7 +832,9 @@ func generateTurnStatusLabel(
fallbackProvider string,
fallbackModelName string,
fallbackModel fantasy.LanguageModel,
fallbackRoute resolvedModelRoute,
keys chatprovider.ProviderAPIKeys,
modelOpts modelBuildOptions,
logger slog.Logger,
debugSvc *chatdebug.Service,
triggerMessageID int64,
@@ -848,24 +850,11 @@ func generateTurnStatusLabel(
"\nChat title: " + chat.Title +
"\n\nAgent's latest message:\n" + assistantText
candidates := make([]shortTextCandidate, 0, len(preferredTitleModels)+1)
for _, c := range preferredTitleModels {
m, err := chatprovider.ModelFromConfig(
c.provider, c.model, keys, chatprovider.UserAgent(),
chatprovider.CoderHeaders(chat),
nil,
)
if err == nil {
candidates = append(candidates, shortTextCandidate{
provider: c.provider,
model: c.model,
lm: m,
})
}
}
candidates := p.preferredShortTextCandidates(chat, keys)
candidates = append(candidates, shortTextCandidate{
provider: fallbackProvider,
model: fallbackModelName,
route: fallbackRoute,
lm: fallbackModel,
})
@@ -876,12 +865,12 @@ func generateTurnStatusLabel(
candidateModel := candidate.lm
finishDebugRun := func(error) {}
if debugEnabled {
candidateCtx, candidateModel, finishDebugRun = prepareQuickgenDebugCandidate(
candidateCtx, candidateModel, finishDebugRun = p.prepareQuickgenDebugCandidate(
labelCtx,
chat,
keys,
debugSvc,
candidate,
modelOpts,
chatdebug.KindQuickgen,
triggerMessageID,
historyTipMessageID,