From f14f58a58edfa71493022954600b6c9291877840 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Wed, 25 Mar 2026 13:34:29 +0000 Subject: [PATCH] feat(coderd/x/chatd): send Coder identity headers to upstream LLM providers (#23578) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add `X-Coder-Owner-Id`, `X-Coder-Chat-Id`, `X-Coder-Subchat-Id`, `X-Coder-Workspace-Id` headers to all outgoing LLM API requests from chatd - Extend `ModelFromConfig` with `extraHeaders` param, forwarded via Fantasy `WithHeaders` on all 8 providers - Add `CoderHeaders(database.Chat)` helper to build the header map from chat state - Update all 4 `ModelFromConfig` call sites (resolveChatModel, computer-use override, title gen, push summary) - Thread `database.Chat` into `generatePushSummary` (was `chatTitle string`) - Tests: `TestCoderHeaders` (4 subtests), `TestModelFromConfig_ExtraHeaders` (OpenAI + Anthropic), `TestModelFromConfig_NilExtraHeaders` - Refactor existing `TestModelFromConfig_UserAgent` to use channel-based signaling > 🤖 This PR was generated by Coder Agents and self-reviewed by a human. --- coderd/x/chatd/chatd.go | 4 +- coderd/x/chatd/chatprovider/chatprovider.go | 99 ++++++++- .../x/chatd/chatprovider/chatprovider_test.go | 208 ++++++++++++++++++ coderd/x/chatd/chatprovider/useragent_test.go | 30 +-- coderd/x/chatd/quickgen.go | 6 +- 5 files changed, 317 insertions(+), 30 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index e0018b5ed8..70fc336d60 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -3340,6 +3340,7 @@ func (p *Server) runChat( chattool.ComputerUseModelName, providerKeys, chatprovider.UserAgent(), + chatprovider.CoderHeaders(chat), ) if cuErr != nil { return result, xerrors.Errorf("resolve computer use model: %w", cuErr) @@ -3795,6 +3796,7 @@ func (p *Server) resolveChatModel( model, err := chatprovider.ModelFromConfig( dbConfig.Provider, dbConfig.Model, keys, chatprovider.UserAgent(), + chatprovider.CoderHeaders(chat), ) if err != nil { return nil, database.ChatModelConfig{}, chatprovider.ProviderAPIKeys{}, xerrors.Errorf( @@ -4091,7 +4093,7 @@ func (p *Server) maybeSendPushNotification( if assistantText != "" && runResult.PushSummaryModel != nil { if summary := generatePushSummary( pushCtx, - chat.Title, + chat, assistantText, runResult.PushSummaryModel, runResult.ProviderKeys, diff --git a/coderd/x/chatd/chatprovider/chatprovider.go b/coderd/x/chatd/chatprovider/chatprovider.go index fa92bfd977..f88774858e 100644 --- a/coderd/x/chatd/chatprovider/chatprovider.go +++ b/coderd/x/chatd/chatprovider/chatprovider.go @@ -14,8 +14,10 @@ import ( fantasyopenaicompat "charm.land/fantasy/providers/openaicompat" fantasyopenrouter "charm.land/fantasy/providers/openrouter" fantasyvercel "charm.land/fantasy/providers/vercel" + "github.com/google/uuid" "golang.org/x/xerrors" + "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/codersdk" ) @@ -888,15 +890,71 @@ func MergeMissingProviderOptions( } } +// Header constants sent on upstream LLM API requests so that +// intermediaries (e.g. aibridged) can correlate traffic back to +// Coder entities. +const ( + // HeaderCoderOwnerID identifies the Coder user who owns the chat. + HeaderCoderOwnerID = "X-Coder-Owner-Id" + // HeaderCoderChatID identifies the top-level (parent) chat. + // For root chats this is the chat's own ID; for subchats it + // is the parent chat's ID. + HeaderCoderChatID = "X-Coder-Chat-Id" + // HeaderCoderSubchatID identifies the current subchat. Only + // present when the request originates from a child chat. + HeaderCoderSubchatID = "X-Coder-Subchat-Id" + // HeaderCoderWorkspaceID identifies the workspace associated + // with the chat, if any. + HeaderCoderWorkspaceID = "X-Coder-Workspace-Id" +) + +// CoderHeaders builds the set of Coder identity headers to attach +// to outgoing LLM API requests for the given chat. +func CoderHeaders(chat database.Chat) map[string]string { + chatID := chat.ID + if chat.ParentChatID.Valid { + chatID = chat.ParentChatID.UUID + } + h := map[string]string{ + HeaderCoderOwnerID: chat.OwnerID.String(), + HeaderCoderChatID: chatID.String(), + } + if chat.ParentChatID.Valid { + h[HeaderCoderSubchatID] = chat.ID.String() + } + if chat.WorkspaceID.Valid { + h[HeaderCoderWorkspaceID] = chat.WorkspaceID.UUID.String() + } + return h +} + +// CoderHeadersFromIDs is a convenience form of CoderHeaders for call +// sites that do not have a full database.Chat in scope. +func CoderHeadersFromIDs( + ownerID uuid.UUID, + chatID uuid.UUID, + parentChatID uuid.NullUUID, + workspaceID uuid.NullUUID, +) map[string]string { + return CoderHeaders(database.Chat{ + ID: chatID, + OwnerID: ownerID, + ParentChatID: parentChatID, + WorkspaceID: workspaceID, + }) +} + // ModelFromConfig resolves a provider/model pair and constructs a fantasy // language model client using the provided provider credentials. The // userAgent is sent as the User-Agent header on every outgoing LLM -// API request. +// API request. extraHeaders, when non-nil, are sent as additional +// HTTP headers on every request. func ModelFromConfig( providerHint string, modelName string, providerKeys ProviderAPIKeys, userAgent string, + extraHeaders map[string]string, ) (fantasy.LanguageModel, error) { provider, modelID, err := ResolveModelWithProviderHint(modelName, providerHint) if err != nil { @@ -916,6 +974,9 @@ func ModelFromConfig( fantasyanthropic.WithAPIKey(apiKey), fantasyanthropic.WithUserAgent(userAgent), } + if len(extraHeaders) > 0 { + options = append(options, fantasyanthropic.WithHeaders(extraHeaders)) + } if baseURL != "" { options = append(options, fantasyanthropic.WithBaseURL(baseURL)) } @@ -924,22 +985,33 @@ func ModelFromConfig( if baseURL == "" { return nil, xerrors.New("AZURE_OPENAI_BASE_URL is not set") } - providerClient, err = fantasyazure.New( + azureOpts := []fantasyazure.Option{ fantasyazure.WithAPIKey(apiKey), fantasyazure.WithBaseURL(baseURL), fantasyazure.WithUseResponsesAPI(), fantasyazure.WithUserAgent(userAgent), - ) + } + if len(extraHeaders) > 0 { + azureOpts = append(azureOpts, fantasyazure.WithHeaders(extraHeaders)) + } + providerClient, err = fantasyazure.New(azureOpts...) case fantasybedrock.Name: - providerClient, err = fantasybedrock.New( + bedrockOpts := []fantasybedrock.Option{ fantasybedrock.WithAPIKey(apiKey), fantasybedrock.WithUserAgent(userAgent), - ) + } + if len(extraHeaders) > 0 { + bedrockOpts = append(bedrockOpts, fantasybedrock.WithHeaders(extraHeaders)) + } + providerClient, err = fantasybedrock.New(bedrockOpts...) case fantasygoogle.Name: options := []fantasygoogle.Option{ fantasygoogle.WithGeminiAPIKey(apiKey), fantasygoogle.WithUserAgent(userAgent), } + if len(extraHeaders) > 0 { + options = append(options, fantasygoogle.WithHeaders(extraHeaders)) + } if baseURL != "" { options = append(options, fantasygoogle.WithBaseURL(baseURL)) } @@ -950,6 +1022,9 @@ func ModelFromConfig( fantasyopenai.WithUseResponsesAPI(), fantasyopenai.WithUserAgent(userAgent), } + if len(extraHeaders) > 0 { + options = append(options, fantasyopenai.WithHeaders(extraHeaders)) + } if baseURL != "" { options = append(options, fantasyopenai.WithBaseURL(baseURL)) } @@ -959,20 +1034,30 @@ func ModelFromConfig( fantasyopenaicompat.WithAPIKey(apiKey), fantasyopenaicompat.WithUserAgent(userAgent), } + if len(extraHeaders) > 0 { + options = append(options, fantasyopenaicompat.WithHeaders(extraHeaders)) + } if baseURL != "" { options = append(options, fantasyopenaicompat.WithBaseURL(baseURL)) } providerClient, err = fantasyopenaicompat.New(options...) case fantasyopenrouter.Name: - providerClient, err = fantasyopenrouter.New( + routerOpts := []fantasyopenrouter.Option{ fantasyopenrouter.WithAPIKey(apiKey), fantasyopenrouter.WithUserAgent(userAgent), - ) + } + if len(extraHeaders) > 0 { + routerOpts = append(routerOpts, fantasyopenrouter.WithHeaders(extraHeaders)) + } + providerClient, err = fantasyopenrouter.New(routerOpts...) case fantasyvercel.Name: options := []fantasyvercel.Option{ fantasyvercel.WithAPIKey(apiKey), fantasyvercel.WithUserAgent(userAgent), } + if len(extraHeaders) > 0 { + options = append(options, fantasyvercel.WithHeaders(extraHeaders)) + } if baseURL != "" { options = append(options, fantasyvercel.WithBaseURL(baseURL)) } diff --git a/coderd/x/chatd/chatprovider/chatprovider_test.go b/coderd/x/chatd/chatprovider/chatprovider_test.go index 8ebce7efbc..f312ddc9fd 100644 --- a/coderd/x/chatd/chatprovider/chatprovider_test.go +++ b/coderd/x/chatd/chatprovider/chatprovider_test.go @@ -1,17 +1,24 @@ package chatprovider_test import ( + "net/http" "testing" + "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" fantasyopenai "charm.land/fantasy/providers/openai" fantasyopenrouter "charm.land/fantasy/providers/openrouter" fantasyvercel "charm.land/fantasy/providers/vercel" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" ) func TestReasoningEffortFromChat(t *testing.T) { @@ -84,6 +91,207 @@ func TestReasoningEffortFromChat(t *testing.T) { } } +func TestCoderHeaders(t *testing.T) { + t.Parallel() + + t.Run("RootChatNoWorkspace", func(t *testing.T) { + t.Parallel() + chatID := uuid.New() + ownerID := uuid.New() + chat := database.Chat{ + ID: chatID, + OwnerID: ownerID, + } + h := chatprovider.CoderHeaders(chat) + require.Equal(t, ownerID.String(), h[chatprovider.HeaderCoderOwnerID]) + require.Equal(t, chatID.String(), h[chatprovider.HeaderCoderChatID]) + require.NotContains(t, h, chatprovider.HeaderCoderSubchatID) + require.NotContains(t, h, chatprovider.HeaderCoderWorkspaceID) + }) + + t.Run("RootChatWithWorkspace", func(t *testing.T) { + t.Parallel() + chatID := uuid.New() + ownerID := uuid.New() + workspaceID := uuid.New() + chat := database.Chat{ + ID: chatID, + OwnerID: ownerID, + WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, + } + h := chatprovider.CoderHeaders(chat) + require.Equal(t, ownerID.String(), h[chatprovider.HeaderCoderOwnerID]) + require.Equal(t, chatID.String(), h[chatprovider.HeaderCoderChatID]) + require.NotContains(t, h, chatprovider.HeaderCoderSubchatID) + require.Equal(t, workspaceID.String(), h[chatprovider.HeaderCoderWorkspaceID]) + }) + + t.Run("SubchatWithWorkspace", func(t *testing.T) { + t.Parallel() + parentID := uuid.New() + subchatID := uuid.New() + ownerID := uuid.New() + workspaceID := uuid.New() + chat := database.Chat{ + ID: subchatID, + OwnerID: ownerID, + ParentChatID: uuid.NullUUID{UUID: parentID, Valid: true}, + WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, + } + h := chatprovider.CoderHeaders(chat) + require.Equal(t, ownerID.String(), h[chatprovider.HeaderCoderOwnerID]) + require.Equal(t, parentID.String(), h[chatprovider.HeaderCoderChatID]) + require.Equal(t, subchatID.String(), h[chatprovider.HeaderCoderSubchatID]) + require.Equal(t, workspaceID.String(), h[chatprovider.HeaderCoderWorkspaceID]) + }) + + t.Run("SubchatNoWorkspace", func(t *testing.T) { + t.Parallel() + parentID := uuid.New() + subchatID := uuid.New() + ownerID := uuid.New() + chat := database.Chat{ + ID: subchatID, + OwnerID: ownerID, + ParentChatID: uuid.NullUUID{UUID: parentID, Valid: true}, + } + h := chatprovider.CoderHeaders(chat) + require.Equal(t, ownerID.String(), h[chatprovider.HeaderCoderOwnerID]) + require.Equal(t, parentID.String(), h[chatprovider.HeaderCoderChatID]) + require.Equal(t, subchatID.String(), h[chatprovider.HeaderCoderSubchatID]) + require.NotContains(t, h, chatprovider.HeaderCoderWorkspaceID) + }) +} + +// TestModelFromConfig_ExtraHeaders verifies that extra headers passed +// to ModelFromConfig are sent on outgoing LLM API requests. Only the +// OpenAI and Anthropic providers are tested end-to-end because the +// WithHeaders injection is the same mechanical pattern across all +// eight provider cases, and these are the only two providers with +// chattest test servers. CoderHeaders construction is tested +// separately in TestCoderHeaders. +func TestModelFromConfig_ExtraHeaders(t *testing.T) { + t.Parallel() + + parentID := uuid.New() + subchatID := uuid.New() + ownerID := uuid.New() + workspaceID := uuid.New() + + chat := database.Chat{ + ID: subchatID, + OwnerID: ownerID, + ParentChatID: uuid.NullUUID{UUID: parentID, Valid: true}, + WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, + } + headers := chatprovider.CoderHeaders(chat) + + assertCoderHeaders := func(t *testing.T, got http.Header) { + t.Helper() + assert.Equal(t, ownerID.String(), got.Get(chatprovider.HeaderCoderOwnerID)) + assert.Equal(t, parentID.String(), got.Get(chatprovider.HeaderCoderChatID)) + assert.Equal(t, subchatID.String(), got.Get(chatprovider.HeaderCoderSubchatID)) + assert.Equal(t, workspaceID.String(), got.Get(chatprovider.HeaderCoderWorkspaceID)) + } + + t.Run("OpenAI", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + + called := make(chan struct{}) + serverURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + assertCoderHeaders(t, req.Header) + close(called) + return chattest.OpenAINonStreamingResponse("hello") + }) + + keys := chatprovider.ProviderAPIKeys{ + ByProvider: map[string]string{"openai": "test-key"}, + BaseURLByProvider: map[string]string{"openai": serverURL}, + } + + model, err := chatprovider.ModelFromConfig("openai", "gpt-4", keys, chatprovider.UserAgent(), headers) + require.NoError(t, err) + + _, err = model.Generate(ctx, fantasy.Call{ + Prompt: []fantasy.Message{ + { + Role: fantasy.MessageRoleUser, + Content: []fantasy.MessagePart{fantasy.TextPart{Text: "hello"}}, + }, + }, + }) + require.NoError(t, err) + _ = testutil.TryReceive(ctx, t, called) + }) + + t.Run("Anthropic", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + + called := make(chan struct{}) + serverURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + assertCoderHeaders(t, req.Header) + close(called) + return chattest.AnthropicNonStreamingResponse("hello") + }) + + keys := chatprovider.ProviderAPIKeys{ + ByProvider: map[string]string{"anthropic": "test-key"}, + BaseURLByProvider: map[string]string{"anthropic": serverURL}, + } + + model, err := chatprovider.ModelFromConfig("anthropic", "claude-sonnet-4-20250514", keys, chatprovider.UserAgent(), headers) + require.NoError(t, err) + + _, err = model.Generate(ctx, fantasy.Call{ + Prompt: []fantasy.Message{ + { + Role: fantasy.MessageRoleUser, + Content: []fantasy.MessagePart{fantasy.TextPart{Text: "hello"}}, + }, + }, + }) + require.NoError(t, err) + _ = testutil.TryReceive(ctx, t, called) + }) +} + +func TestModelFromConfig_NilExtraHeaders(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + + called := make(chan struct{}) + serverURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + // Coder headers must be absent when nil is passed. + assert.Empty(t, req.Header.Get(chatprovider.HeaderCoderOwnerID)) + assert.Empty(t, req.Header.Get(chatprovider.HeaderCoderChatID)) + assert.Empty(t, req.Header.Get(chatprovider.HeaderCoderSubchatID)) + assert.Empty(t, req.Header.Get(chatprovider.HeaderCoderWorkspaceID)) + close(called) + return chattest.OpenAINonStreamingResponse("hello") + }) + + keys := chatprovider.ProviderAPIKeys{ + ByProvider: map[string]string{"openai": "test-key"}, + BaseURLByProvider: map[string]string{"openai": serverURL}, + } + + model, err := chatprovider.ModelFromConfig("openai", "gpt-4", keys, chatprovider.UserAgent(), nil) + require.NoError(t, err) + + _, err = model.Generate(ctx, fantasy.Call{ + Prompt: []fantasy.Message{ + { + Role: fantasy.MessageRoleUser, + Content: []fantasy.MessagePart{fantasy.TextPart{Text: "hello"}}, + }, + }, + }) + require.NoError(t, err) + _ = testutil.TryReceive(ctx, t, called) +} + func TestMergeMissingProviderOptions_OpenRouterNested(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatprovider/useragent_test.go b/coderd/x/chatd/chatprovider/useragent_test.go index 7df3f3a39a..58ee18fffe 100644 --- a/coderd/x/chatd/chatprovider/useragent_test.go +++ b/coderd/x/chatd/chatprovider/useragent_test.go @@ -1,10 +1,8 @@ package chatprovider_test import ( - "context" "runtime" "strings" - "sync" "testing" "charm.land/fantasy" @@ -14,6 +12,7 @@ import ( "github.com/coder/coder/v2/buildinfo" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" "github.com/coder/coder/v2/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/testutil" ) func TestUserAgent(t *testing.T) { @@ -34,29 +33,27 @@ func TestUserAgent(t *testing.T) { func TestModelFromConfig_UserAgent(t *testing.T) { t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) - var mu sync.Mutex - var capturedUA string - + expectedUA := chatprovider.UserAgent() + called := make(chan struct{}) serverURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - mu.Lock() - capturedUA = req.Header.Get("User-Agent") - mu.Unlock() + assert.Equal(t, expectedUA, req.Header.Get("User-Agent")) + close(called) return chattest.OpenAINonStreamingResponse("hello") }) - expectedUA := chatprovider.UserAgent() keys := chatprovider.ProviderAPIKeys{ ByProvider: map[string]string{"openai": "test-key"}, BaseURLByProvider: map[string]string{"openai": serverURL}, } - model, err := chatprovider.ModelFromConfig("openai", "gpt-4", keys, expectedUA) + model, err := chatprovider.ModelFromConfig("openai", "gpt-4", keys, expectedUA, nil) require.NoError(t, err) // Make a real call so Fantasy sends an HTTP request to the - // fake server, which captures the User-Agent header. - _, err = model.Generate(context.Background(), fantasy.Call{ + // fake server, which asserts the User-Agent header. + _, err = model.Generate(ctx, fantasy.Call{ Prompt: []fantasy.Message{ { Role: fantasy.MessageRoleUser, @@ -67,12 +64,5 @@ func TestModelFromConfig_UserAgent(t *testing.T) { }, }) require.NoError(t, err) - - mu.Lock() - got := capturedUA - mu.Unlock() - - require.NotEmpty(t, got, "User-Agent header was not sent") - require.Equal(t, expectedUA, got, - "User-Agent header should match chatprovider.UserAgent()") + _ = testutil.TryReceive(ctx, t, called) } diff --git a/coderd/x/chatd/quickgen.go b/coderd/x/chatd/quickgen.go index 77c9ff5bd0..b23a394b55 100644 --- a/coderd/x/chatd/quickgen.go +++ b/coderd/x/chatd/quickgen.go @@ -83,6 +83,7 @@ func (p *Server) maybeGenerateChatTitle( for _, c := range preferredTitleModels { m, err := chatprovider.ModelFromConfig( c.provider, c.model, keys, chatprovider.UserAgent(), + chatprovider.CoderHeaders(chat), ) if err == nil { candidates = append(candidates, m) @@ -271,7 +272,7 @@ const pushSummaryPrompt = "You are a notification assistant. Given a chat title // fall back to the provided model. Returns "" on any failure. func generatePushSummary( ctx context.Context, - chatTitle string, + chat database.Chat, assistantText string, fallbackModel fantasy.LanguageModel, keys chatprovider.ProviderAPIKeys, @@ -280,12 +281,13 @@ func generatePushSummary( summaryCtx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - input := "Chat title: " + chatTitle + "\n\nAgent's last message:\n" + assistantText + input := "Chat title: " + chat.Title + "\n\nAgent's last message:\n" + assistantText candidates := make([]fantasy.LanguageModel, 0, len(preferredTitleModels)+1) for _, c := range preferredTitleModels { m, err := chatprovider.ModelFromConfig( c.provider, c.model, keys, chatprovider.UserAgent(), + chatprovider.CoderHeaders(chat), ) if err == nil { candidates = append(candidates, m)