From 2f60b146494c6d33b0815bda3efd1cc0e8807e1d Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Tue, 2 Jun 2026 09:52:58 -0400 Subject: [PATCH] fix(coderd/x/chatd): preserve chat API key after compaction (#25930) (#25974) --- coderd/x/chatd/chatd.go | 34 ++--- coderd/x/chatd/chatd_internal_test.go | 83 +++++++----- coderd/x/chatd/model_routing_internal_test.go | 118 ++++++++++++++++-- 3 files changed, 182 insertions(+), 53 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 013567e3e0..5a5ba7fb60 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -6468,22 +6468,14 @@ type runChatResult struct { HistoryTipMessageID int64 } -func contextWithActiveTurnAPIKeyID(ctx context.Context, messages []database.ChatMessage) context.Context { - apiKeyID, ok := activeTurnAPIKeyIDFromMessages(messages) - if !ok { - return ctx - } - return aibridge.WithDelegatedAPIKeyID(ctx, apiKeyID) -} - func activeTurnAPIKeyIDFromMessages(messages []database.ChatMessage) (string, bool) { for i := len(messages) - 1; i >= 0; i-- { message := messages[i] if message.Role != database.ChatMessageRoleUser { continue } - if message.Visibility != database.ChatMessageVisibilityBoth && - message.Visibility != database.ChatMessageVisibilityUser { + if !isUserVisibleChatMessage(message) && + !(message.Visibility == database.ChatMessageVisibilityModel && message.Compressed) { continue } if !message.APIKeyID.Valid || message.APIKeyID.String == "" { @@ -6494,6 +6486,11 @@ func activeTurnAPIKeyIDFromMessages(messages []database.ChatMessage) (string, bo return "", false } +func isUserVisibleChatMessage(message database.ChatMessage) bool { + return message.Visibility == database.ChatMessageVisibilityBoth || + message.Visibility == database.ChatMessageVisibilityUser +} + func allToolNames(allTools []fantasy.AgentTool) []string { toolNames := make([]string, 0, len(allTools)) for _, tool := range allTools { @@ -7124,7 +7121,9 @@ func (p *Server) runChat( return result, xerrors.Errorf("get chat messages: %w", err) } modelOpts := modelBuildOptionsFromMessages(messages) - ctx = contextWithActiveTurnAPIKeyID(ctx, messages) + if modelOpts.ActiveAPIKeyID != "" { + ctx = aibridge.WithDelegatedAPIKeyID(ctx, modelOpts.ActiveAPIKeyID) + } // Load MCP server configs and user tokens in parallel with model // resolution. These queries have no dependencies on each other and all @@ -7831,6 +7830,7 @@ func (p *Server) runChat( persistCtx, chat.ID, modelConfig.ID, + modelOpts.ActiveAPIKeyID, compactionToolCallID, result, ); err != nil { @@ -8460,12 +8460,14 @@ func buildProviderTools(options *codersdk.ChatModelProviderOptions) []chatloop.P return tools } -// persistChatContextSummary persists a chat context summary to the database. -// This is invoked via the chat loop's compaction callback. +// persistChatContextSummary is called from the chat loop's compaction +// callback. activeAPIKeyID is stamped onto the summary user message. When +// empty, it falls back to the delegated key in ctx. func (p *Server) persistChatContextSummary( ctx context.Context, chatID uuid.UUID, modelConfigID uuid.UUID, + activeAPIKeyID string, toolCallID string, result chatloop.CompactionResult, ) error { @@ -8514,6 +8516,11 @@ func (p *Server) persistChatContextSummary( return xerrors.Errorf("encode summary tool result: %w", err) } + summaryAPIKeyID := activeAPIKeyID + if summaryAPIKeyID == "" { + summaryAPIKeyID, _ = aibridge.DelegatedAPIKeyIDFromContext(ctx) + } + var insertedMessages []database.ChatMessage txErr := p.db.InTx(func(tx database.Store) error { @@ -8522,7 +8529,6 @@ func (p *Server) persistChatContextSummary( } // Hidden summary user message (not published to subscribers). - summaryAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) summaryUserMsg := newUserChatMessage( summaryAPIKeyID, systemContent, diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 69aebbfe30..965b6b474e 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -6651,42 +6651,63 @@ func TestPersistChatContextSummarySetsAPIKeyID(t *testing.T) { UserID: user.ID, }) - ctx = aibridge.WithDelegatedAPIKeyID(ctx, apiKey.ID) - server := &Server{db: db} + persistAndAssertSummaryKey := func( + summaryCtx context.Context, + chatID uuid.UUID, + activeAPIKeyID string, + wantAPIKeyID string, + toolCallID string, + ) { + t.Helper() - 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) + err := server.persistChatContextSummary( + summaryCtx, + chatID, + modelConfig.ID, + activeAPIKeyID, + toolCallID, + 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) + msgs, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + 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) + // 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") + 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, wantAPIKeyID, msg.APIKeyID.String, "summary user message APIKeyID must match") + } } + require.True(t, foundUserSummary, "expected to find compressed user summary message") } - require.True(t, foundUserSummary, "expected to find compressed user summary message") + + persistAndAssertSummaryKey(ctx, chat.ID, apiKey.ID, apiKey.ID, "tool-call-id-1") + + fallbackChat := dbgen.Chat(t, db, database.Chat{ + OwnerID: user.ID, + OrganizationID: org.ID, + LastModelConfigID: modelConfig.ID, + }) + fallbackKey, _ := dbgen.APIKey(t, db, database.APIKey{ + UserID: user.ID, + }) + fallbackCtx := aibridge.WithDelegatedAPIKeyID(ctx, fallbackKey.ID) + persistAndAssertSummaryKey(fallbackCtx, fallbackChat.ID, "", fallbackKey.ID, "tool-call-id-2") } diff --git a/coderd/x/chatd/model_routing_internal_test.go b/coderd/x/chatd/model_routing_internal_test.go index 0d2f317204..70e0797883 100644 --- a/coderd/x/chatd/model_routing_internal_test.go +++ b/coderd/x/chatd/model_routing_internal_test.go @@ -405,7 +405,7 @@ func TestActiveTurnAPIKeyIDFromMessages(t *testing.T) { }, }, { - name: "SkipsModelOnlyUserMessages", + name: "SkipsUncompressedModelOnlyUserMessages", messages: []database.ChatMessage{ {ID: 1, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, APIKeyID: sqlNullString(oldKeyID)}, {ID: 2, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, APIKeyID: sqlNullString(currentKeyID)}, @@ -413,6 +413,54 @@ func TestActiveTurnAPIKeyIDFromMessages(t *testing.T) { wantKey: oldKeyID, wantOK: true, }, + { + name: "CompressedSummaryFallback", + messages: []database.ChatMessage{ + {ID: 1, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, Compressed: true, APIKeyID: sqlNullString(currentKeyID)}, + {ID: 2, Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth}, + }, + wantKey: currentKeyID, + wantOK: true, + }, + { + name: "LatestCompressedSummaryWins", + messages: []database.ChatMessage{ + {ID: 1, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, Compressed: true, APIKeyID: sqlNullString(oldKeyID)}, + {ID: 2, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, Compressed: true, APIKeyID: sqlNullString(currentKeyID)}, + {ID: 3, Role: database.ChatMessageRoleAssistant, Visibility: database.ChatMessageVisibilityBoth}, + }, + wantKey: currentKeyID, + wantOK: true, + }, + { + name: "VisibleUserWinsOverCompressedSummary", + messages: []database.ChatMessage{ + {ID: 1, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, Compressed: true, APIKeyID: sqlNullString(oldKeyID)}, + {ID: 2, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, APIKeyID: sqlNullString(currentKeyID)}, + }, + wantKey: currentKeyID, + wantOK: true, + }, + { + name: "MissingVisibleUserKeyDoesNotFallBackToCompressedSummary", + messages: []database.ChatMessage{ + {ID: 1, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, Compressed: true, APIKeyID: sqlNullString(oldKeyID)}, + {ID: 2, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth}, + }, + }, + { + name: "UncompressedModelOnlyUserIgnored", + messages: []database.ChatMessage{ + {ID: 1, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, APIKeyID: sqlNullString(currentKeyID)}, + }, + }, + { + name: "CompressedSummaryMissingKeyDoesNotFallBack", + messages: []database.ChatMessage{ + {ID: 1, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, APIKeyID: sqlNullString(oldKeyID)}, + {ID: 2, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityModel, Compressed: true}, + }, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -421,15 +469,11 @@ func TestActiveTurnAPIKeyIDFromMessages(t *testing.T) { gotKey, gotOK := activeTurnAPIKeyIDFromMessages(tt.messages) require.Equal(t, tt.wantOK, gotOK) require.Equal(t, tt.wantKey, gotKey) - ctx := contextWithActiveTurnAPIKeyID(t.Context(), tt.messages) - ctxKey, ctxOK := aibridge.DelegatedAPIKeyIDFromContext(ctx) - require.Equal(t, tt.wantOK, ctxOK) - require.Equal(t, tt.wantKey, ctxKey) }) } } -func TestActiveTurnContextUsesPromptMessages(t *testing.T) { +func TestPromptMessagesForVisibleUserPreserveActiveAPIKeyID(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) @@ -477,12 +521,70 @@ func TestActiveTurnContextUsesPromptMessages(t *testing.T) { messages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) require.NoError(t, err) - ctx = contextWithActiveTurnAPIKeyID(ctx, messages) - gotKey, ok := aibridge.DelegatedAPIKeyIDFromContext(ctx) + gotKey, ok := activeTurnAPIKeyIDFromMessages(messages) require.True(t, ok) require.Equal(t, currentKey.ID, gotKey) } +func TestPromptMessagesForCompactedChatPreserveActiveAPIKeyID(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := t.Context() + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: model.ID}) + key, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + + visibleUser := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Visibility: database.ChatMessageVisibilityBoth, + APIKeyID: sqlNullString(key.ID), + }) + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Visibility: database.ChatMessageVisibilityBoth, + }) + compressedSummary := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Visibility: database.ChatMessageVisibilityModel, + Compressed: true, + APIKeyID: sqlNullString(key.ID), + }) + afterSummary := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Visibility: database.ChatMessageVisibilityBoth, + }) + + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + + ids := make(map[int64]struct{}, len(messages)) + for _, message := range messages { + ids[message.ID] = struct{}{} + } + _, hasVisibleUser := ids[visibleUser.ID] + require.False(t, hasVisibleUser) + _, hasSummary := ids[compressedSummary.ID] + require.True(t, hasSummary) + _, hasAfterSummary := ids[afterSummary.ID] + require.True(t, hasAfterSummary) + + gotKey, ok := activeTurnAPIKeyIDFromMessages(messages) + require.True(t, ok) + require.Equal(t, key.ID, gotKey) +} + func sqlNullString(value string) sql.NullString { return sql.NullString{String: value, Valid: value != ""} }