From b278be7361616516a2189ce86717556e5e696a86 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Wed, 27 May 2026 17:00:23 +0100 Subject: [PATCH] fix(coderd): enforce api_key_id on user messages at type level (#25729) - Empty string is valid for `apiKeyID` in paths that genuinely lack a caller key (e.g. agent-initiated context injection in `workspaceAgentAddChatContext`). AI Gateway fail-closed check remains the runtime safety net. - Context injection paths (`persistInstructionFiles`, compaction) read the key from `aibridge.DelegatedAPIKeyIDFromContext(ctx)`, set upstream by `contextWithActiveTurnAPIKeyID`. - Subagent context copy branches on `copiedRole == database.ChatMessageRoleUser` to choose the right append function. > Generated by Coder Agents --- coderd/workspaceagents.go | 4 +- coderd/x/chatd/chatd.go | 161 ++++++++++++------ coderd/x/chatd/chatd_internal_test.go | 76 ++++++++- coderd/x/chatd/subagent.go | 37 ++-- .../x/chatd/subagent_context_internal_test.go | 24 +++ 5 files changed, 239 insertions(+), 63 deletions(-) diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index b3074740b1..097d4ff41b 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -2562,9 +2562,9 @@ func (api *API) workspaceAgentAddChatContext(rw http.ResponseWriter, r *http.Req if locked.OwnerID != workspace.OwnerID { return errChatDoesNotBelongToWorkspaceOwner } - if _, err := tx.InsertChatMessages(sysCtx, chatd.BuildSingleChatMessageInsertParams( + if _, err := tx.InsertChatMessages(sysCtx, chatd.BuildSingleUserChatMessageInsertParams( chat.ID, - database.ChatMessageRoleUser, + "", // Agent-initiated context injection has no caller API key. content, database.ChatMessageVisibilityBoth, locked.LastModelConfigID, diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index debee591e8..013567e3e0 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -1629,7 +1629,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C return xerrors.Errorf("marshal initial user content: %w", err) } - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by append[User]ChatMessage. ChatID: insertedChat.ID, } @@ -1673,13 +1673,15 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C chatprompt.CurrentContentVersion, )) - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleUser, + userMsg := newUserChatMessage( + opts.APIKeyID, userContent, database.ChatMessageVisibilityBoth, opts.ModelConfigID, chatprompt.CurrentContentVersion, - ).withCreatedBy(opts.OwnerID).withAPIKeyID(opts.APIKeyID)) + ) + userMsg = userMsg.withCreatedBy(opts.OwnerID) + appendUserChatMessage(&msgParams, userMsg) _, err = tx.InsertChatMessages(ctx, msgParams) if err != nil { @@ -2111,16 +2113,18 @@ func (p *Server) EditMessage( // InsertChatMessages CTE updates chats.last_model_config_id // when the new message's model differs, so the assistant turn // that follows picks up the new selection. - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: opts.ChatID, } - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleUser, + editUserMsg := newUserChatMessage( + opts.APIKeyID, content, editedMsg.Visibility, messageModelConfigID, chatprompt.CurrentContentVersion, - ).withCreatedBy(opts.CreatedBy).withAPIKeyID(opts.APIKeyID)) + ) + editUserMsg = editUserMsg.withCreatedBy(opts.CreatedBy) + appendUserChatMessage(&msgParams, editUserMsg) newMessages, err := insertChatMessageWithStore(ctx, tx, msgParams) if err != nil { return xerrors.Errorf("insert replacement message: %w", err) @@ -3899,19 +3903,16 @@ func insertChatMessageWithStore( return messages, nil } -// chatMessage describes a single message to insert as part of a batch. -// Use newChatMessage to create one, then chain builder methods for -// optional fields. For nullable UUID fields (ModelConfigID, CreatedBy), -// use uuid.Nil to represent NULL — the SQL uses NULLIF to convert zero -// UUIDs to NULL. For nullable int64 fields, use 0 to represent NULL — -// the SQL uses NULLIF to convert zeros to NULL. +// chatMessage is the base message type for batch inserts. Use directly +// only for non-user messages; for user messages, use userChatMessage. +// For nullable UUID fields (ModelConfigID, CreatedBy), use uuid.Nil to +// represent NULL. For nullable int64 fields, use 0 to represent NULL. type chatMessage struct { role database.ChatMessageRole content pqtype.NullRawMessage visibility database.ChatMessageVisibility modelConfigID uuid.UUID createdBy uuid.UUID - apiKeyID string contentVersion int16 compressed bool inputTokens int64 @@ -3926,6 +3927,23 @@ type chatMessage struct { providerResponseID string } +// userChatMessage wraps chatMessage with a required apiKeyID so that +// omitting it for user messages is a compile error, not a silent data bug. +type userChatMessage struct { + chatMessage + apiKeyID string +} + +func (m userChatMessage) withCreatedBy(id uuid.UUID) userChatMessage { + m.chatMessage = m.chatMessage.withCreatedBy(id) + return m +} + +func (m userChatMessage) withCompressed() userChatMessage { + m.chatMessage = m.chatMessage.withCompressed() + return m +} + func newChatMessage( role database.ChatMessageRole, content pqtype.NullRawMessage, @@ -3942,13 +3960,29 @@ func newChatMessage( } } -func (m chatMessage) withCreatedBy(id uuid.UUID) chatMessage { - m.createdBy = id - return m +// newUserChatMessage creates a user message. apiKeyID is required so +// that forgetting it is a compile error rather than a silent data bug. +func newUserChatMessage( + apiKeyID string, + content pqtype.NullRawMessage, + visibility database.ChatMessageVisibility, + modelConfigID uuid.UUID, + contentVersion int16, +) userChatMessage { + return userChatMessage{ + chatMessage: newChatMessage( + database.ChatMessageRoleUser, + content, + visibility, + modelConfigID, + contentVersion, + ), + apiKeyID: apiKeyID, + } } -func (m chatMessage) withAPIKeyID(id string) chatMessage { - m.apiKeyID = id +func (m chatMessage) withCreatedBy(id uuid.UUID) chatMessage { + m.createdBy = id return m } @@ -3990,13 +4024,16 @@ func (m chatMessage) withProviderResponseID(id string) chatMessage { return m } -// appendChatMessage appends a single message to the batch insert params. -func appendChatMessage( +// appendMessageFields writes all chatMessage fields into the batch insert +// params. apiKeyID is explicit so non-user messages always get "" while +// user messages carry the caller's key for AI Gateway routing. +func appendMessageFields( params *database.InsertChatMessagesParams, msg chatMessage, + apiKeyID string, ) { params.CreatedBy = append(params.CreatedBy, msg.createdBy) - params.APIKeyID = append(params.APIKeyID, msg.apiKeyID) + params.APIKeyID = append(params.APIKeyID, apiKeyID) params.ModelConfigID = append(params.ModelConfigID, msg.modelConfigID) params.Role = append(params.Role, msg.role) params.Content = append(params.Content, string(msg.content.RawMessage)) @@ -4015,25 +4052,44 @@ func appendChatMessage( params.ProviderResponseID = append(params.ProviderResponseID, msg.providerResponseID) } -// BuildSingleChatMessageInsertParams creates batch insert params for one -// message using the shared chat message builder. -func BuildSingleChatMessageInsertParams( +// appendChatMessage appends a non-user message to the batch insert params. +func appendChatMessage( + params *database.InsertChatMessagesParams, + msg chatMessage, +) { + if msg.role == database.ChatMessageRoleUser { + panic("developer error: use appendUserChatMessage for user-role messages") + } + appendMessageFields(params, msg, "") +} + +// appendUserChatMessage inserts a user message with its apiKeyID preserved. +func appendUserChatMessage( + params *database.InsertChatMessagesParams, + msg userChatMessage, +) { + appendMessageFields(params, msg.chatMessage, msg.apiKeyID) +} + +// BuildSingleUserChatMessageInsertParams creates batch insert params for +// one user message, requiring an apiKeyID for AI Gateway attribution. +func BuildSingleUserChatMessageInsertParams( chatID uuid.UUID, - role database.ChatMessageRole, + apiKeyID string, content pqtype.NullRawMessage, visibility database.ChatMessageVisibility, modelConfigID uuid.UUID, contentVersion int16, createdBy uuid.UUID, ) database.InsertChatMessagesParams { - params := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + params := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: chatID, } - msg := newChatMessage(role, content, visibility, modelConfigID, contentVersion) + msg := newUserChatMessage(apiKeyID, content, visibility, modelConfigID, contentVersion) if createdBy != uuid.Nil { msg = msg.withCreatedBy(createdBy) } - appendChatMessage(¶ms, msg) + appendUserChatMessage(¶ms, msg) return params } @@ -4048,16 +4104,18 @@ func insertUserMessageAndSetPending( createdBy uuid.UUID, apiKeyID string, ) (database.ChatMessage, database.Chat, error) { - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: lockedChat.ID, } - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleUser, + insertUserMsg := newUserChatMessage( + apiKeyID, content, database.ChatMessageVisibilityBoth, modelConfigID, chatprompt.CurrentContentVersion, - ).withCreatedBy(createdBy).withAPIKeyID(apiKeyID)) + ) + insertUserMsg = insertUserMsg.withCreatedBy(createdBy) + appendUserChatMessage(&msgParams, insertUserMsg) messages, err := insertChatMessageWithStore(ctx, store, msgParams) if err != nil { return database.ChatMessage{}, database.Chat{}, err @@ -5870,11 +5928,11 @@ func (p *Server) tryAutoPromoteQueuedMessage( return nil, nil, false, xerrors.New("popped queued message out of order") } - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: chat.ID, } - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleUser, + queuedUserMsg := newUserChatMessage( + nextQueued.APIKeyID.String, pqtype.NullRawMessage{ RawMessage: nextQueued.Content, Valid: len(nextQueued.Content) > 0, @@ -5882,7 +5940,9 @@ func (p *Server) tryAutoPromoteQueuedMessage( database.ChatMessageVisibilityBoth, effectiveModelConfigID, chatprompt.CurrentContentVersion, - ).withCreatedBy(chat.OwnerID).withAPIKeyID(nextQueued.APIKeyID.String)) + ) + queuedUserMsg = queuedUserMsg.withCreatedBy(chat.OwnerID) + appendUserChatMessage(&msgParams, queuedUserMsg) msgs, err := insertChatMessageWithStore(ctx, tx, msgParams) if err != nil { return nil, nil, false, xerrors.Errorf("insert promoted message: %w", err) @@ -8457,18 +8517,21 @@ func (p *Server) persistChatContextSummary( var insertedMessages []database.ChatMessage txErr := p.db.InTx(func(tx database.Store) error { - summaryParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + summaryParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by append[User]ChatMessage. ChatID: chatID, } // Hidden summary user message (not published to subscribers). - appendChatMessage(&summaryParams, newChatMessage( - database.ChatMessageRoleUser, + summaryAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) + summaryUserMsg := newUserChatMessage( + summaryAPIKeyID, systemContent, database.ChatMessageVisibilityModel, modelConfigID, chatprompt.CurrentContentVersion, - ).withCompressed()) + ) + summaryUserMsg = summaryUserMsg.withCompressed() + appendUserChatMessage(&summaryParams, summaryUserMsg) // Assistant tool-call message. appendChatMessage(&summaryParams, newChatMessage( @@ -8999,11 +9062,12 @@ func (p *Server) persistInstructionFiles( if err != nil { return "", nil, nil } - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + contextAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: chat.ID, } - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleUser, + appendUserChatMessage(&msgParams, newUserChatMessage( + contextAPIKeyID, content, database.ChatMessageVisibilityBoth, modelConfigID, @@ -9022,11 +9086,12 @@ func (p *Server) persistInstructionFiles( return "", nil, xerrors.Errorf("marshal context-file parts: %w", err) } - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + contextAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: chat.ID, } - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleUser, + appendUserChatMessage(&msgParams, newUserChatMessage( + contextAPIKeyID, content, database.ChatMessageVisibilityBoth, modelConfigID, diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index b28c0cdb0c..69aebbfe30 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -19,9 +19,12 @@ import ( "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/coderd/database/dbtestutil" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/rbac" @@ -1339,6 +1342,8 @@ func TestPersistInstructionFilesIncludesAgentMetadata(t *testing.T) { t.Parallel() ctx := context.Background() + testAPIKeyID := uuid.NewString() + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID) ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) @@ -1366,7 +1371,18 @@ func TestPersistInstructionFilesIncludesAgentMetadata(t *testing.T) { gomock.Any(), agentID, ).Return(workspaceAgent, nil).Times(1) - db.EXPECT().InsertChatMessages(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + db.EXPECT().InsertChatMessages(gomock.Any(), gomock.Cond(func(x any) bool { + params, ok := x.(database.InsertChatMessagesParams) + if !ok { + return false + } + for i, role := range params.Role { + if role == database.ChatMessageRoleUser && params.APIKeyID[i] != testAPIKeyID { + return false + } + } + return true + })).Return(nil, nil).AnyTimes() db.EXPECT().UpdateChatLastInjectedContext(gomock.Any(), gomock.Cond(func(x any) bool { arg, ok := x.(database.UpdateChatLastInjectedContextParams) @@ -6616,3 +6632,61 @@ func TestPrimeWorkspaceMCPCache_ExitsOnContextCancel(t *testing.T) { _, ok := server.workspaceMCPToolsCache.Load(chat.ID) require.False(t, ok, "primer must not cache anything when canceled") } + +func TestPersistChatContextSummarySetsAPIKeyID(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := context.Background() + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OwnerID: user.ID, + OrganizationID: org.ID, + LastModelConfigID: modelConfig.ID, + }) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{ + UserID: user.ID, + }) + + ctx = aibridge.WithDelegatedAPIKeyID(ctx, apiKey.ID) + + server := &Server{db: db} + + err := server.persistChatContextSummary( + ctx, + chat.ID, + modelConfig.ID, + "tool-call-id-1", + chatloop.CompactionResult{ + SystemSummary: "summarized context", + SummaryReport: "context was summarized", + ThresholdPercent: 70, + UsagePercent: 85.0, + ContextTokens: 8500, + ContextLimit: 10000, + }, + ) + require.NoError(t, err) + + msgs, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + + // GetChatMessagesForPromptByChatID uses a compaction boundary CTE + // that selects compressed=true, visibility='model'. Only the user + // summary qualifies; the assistant (visibility=user) and tool + // result (visibility=both) are excluded by the CTE filter. + require.NotEmpty(t, msgs) + + var foundUserSummary bool + for _, msg := range msgs { + if msg.Role == database.ChatMessageRoleUser { + foundUserSummary = true + require.True(t, msg.APIKeyID.Valid, "summary user message must have APIKeyID set") + require.Equal(t, apiKey.ID, msg.APIKeyID.String, "summary user message APIKeyID must match") + } + } + require.True(t, foundUserSummary, "expected to find compressed user summary message") +} diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index 8984ac86cd..450397416b 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -1090,16 +1090,18 @@ func (p *Server) createChildSubagentChatWithOptions( return xerrors.Errorf("update child injected context: %w", err) } - userParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + userParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: insertedChat.ID, } - appendChatMessage(&userParams, newChatMessage( - database.ChatMessageRoleUser, + childUserMsg := newUserChatMessage( + childAPIKeyID, userContent, database.ChatMessageVisibilityBoth, modelConfigID, chatprompt.CurrentContentVersion, - ).withCreatedBy(parent.OwnerID).withAPIKeyID(childAPIKeyID)) + ) + childUserMsg = childUserMsg.withCreatedBy(parent.OwnerID) + appendUserChatMessage(&userParams, childUserMsg) if _, err := tx.InsertChatMessages(ctx, userParams); err != nil { return xerrors.Errorf("insert initial child user message: %w", err) } @@ -1176,16 +1178,27 @@ func copyParentContextMessages( return nil, xerrors.Errorf("marshal filtered context parts: %w", err) } - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by append[User]ChatMessage. ChatID: child.ID, } - appendChatMessage(&msgParams, newChatMessage( - copiedRole, - filteredContent, - copiedVisibility, - child.LastModelConfigID, - copiedVersion, - )) + if copiedRole == database.ChatMessageRoleUser { + copiedAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) + appendUserChatMessage(&msgParams, newUserChatMessage( + copiedAPIKeyID, + filteredContent, + copiedVisibility, + child.LastModelConfigID, + copiedVersion, + )) + } else { + appendChatMessage(&msgParams, newChatMessage( + copiedRole, + filteredContent, + copiedVisibility, + child.LastModelConfigID, + copiedVersion, + )) + } if _, err := store.InsertChatMessages(ctx, msgParams); err != nil { return nil, xerrors.Errorf("insert context message: %w", err) } diff --git a/coderd/x/chatd/subagent_context_internal_test.go b/coderd/x/chatd/subagent_context_internal_test.go index dc60e3330f..5ccab312d6 100644 --- a/coderd/x/chatd/subagent_context_internal_test.go +++ b/coderd/x/chatd/subagent_context_internal_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/require" "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" @@ -446,10 +447,33 @@ func TestCreateChildSubagentChatUpdatesInheritedLastInjectedContext(t *testing.T ctx := chatdTestContext(t) parentChat := createParentChatWithInheritedContext(ctx, t, db, server) + // Set a delegated API key so that copied user-role context messages + // are stamped with api_key_id, preserving AI Gateway routing. + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: parentChat.OwnerID}) + ctx = aibridge.WithDelegatedAPIKeyID(ctx, apiKey.ID) + child, err := server.createChildSubagentChat(ctx, parentChat, "inspect bindings", "") require.NoError(t, err) assertChildInheritedContext(ctx, t, db, child.ID, "inspect bindings") + + // Verify that all user-role messages in the child chat carry + // api_key_id so activeTurnAPIKeyIDFromMessages resolves correctly. + childMessages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: child.ID, + AfterID: 0, + }) + require.NoError(t, err) + var userMsgCount int + for _, msg := range childMessages { + if msg.Role != database.ChatMessageRoleUser { + continue + } + userMsgCount++ + require.True(t, msg.APIKeyID.Valid, "child user message (id=%d) should have api_key_id set", msg.ID) + require.Equal(t, apiKey.ID, msg.APIKeyID.String, "child user message (id=%d) api_key_id mismatch", msg.ID) + } + require.Greater(t, userMsgCount, 0, "expected at least one user-role message in child chat") } func TestSpawnComputerUseAgentInheritsContext(t *testing.T) {