mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
+66
-77
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user