opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(coderd): enforce api_key_id on user messages at type level (#25729)

Browse Source 
- 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
This commit is contained in:
Cian Johnston
2026-05-27 17:00:23 +01:00
committed by GitHub
parent 2d40a40c57
commit b278be7361
5 changed files with 239 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/workspaceagents.go
+2 -2
View File
@@ -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,
coderd/x/chatd/chatd.go
+113 -48
View File
@@ -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(&params, msg)
appendUserChatMessage(&params, 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,
coderd/x/chatd/chatd_internal_test.go
+75 -1
View File
@@ -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")
}
coderd/x/chatd/subagent.go
+25 -12
View File
@@ -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)
}
coderd/x/chatd/subagent_context_internal_test.go
+24
View File
@@ -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) {