shishantbiswas/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-06-02 20:48:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(coderd/x/chatd): preserve chat API key after compaction (#25930)

Browse Source 
> Mux updated this PR on behalf of Mike.

AI Gateway chat retries after context compaction could lose active turn
API key routing metadata because the prompt query keeps the compressed
model-only summary but omits the original visible user turn.

Persist the active API key ID onto compaction summaries explicitly.
Model construction now uses one active-turn lookup helper for visible
user turns and compressed summary boundaries, so prompt model
construction can recover the key when no later visible user turn exists.
Added unit and DB-backed coverage for the compacted prompt path.
This commit is contained in:
Michael Suchacz
2026-06-02 12:19:06 +02:00
committed by GitHub
parent 2269cec830
commit dd22086734
3 changed files with 182 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/x/chatd/chatd.go
+20 -14
View File
@@ -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,
coderd/x/chatd/chatd_internal_test.go
+52 -31
View File
@@ -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")
}
coderd/x/chatd/model_routing_internal_test.go
+110 -8
View File
@@ -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 != ""}
}