Skip to content

Commit 483adc5

Browse files
kylecarbskylecarbs
authored
feat: replace InsertChatMessage with batch InsertChatMessages (coder#23220)
Replaces the singular `InsertChatMessage` query with `InsertChatMessages` that uses PostgreSQL's `unnest()` for batch inserts. This reduces the number of database round-trips when inserting multiple messages in a single transaction. ## Changes - **SQL**: New `InsertChatMessages :many` query using `unnest()` arrays following the existing codebase pattern (e.g., `InsertWorkspaceAgentStats`). Preserves the CTE that updates `chats.last_model_config_id` using the last non-null model config from the batch. Uses `NULLIF` for UUID columns to handle NULL foreign keys. - **Go layers**: Updated `querier.go`, `dbauthz.go`, `dbmetrics/querymetrics.go`, `dbmock/dbmock.go`, and `queries.sql.go` to use the new batch signature (`[]ChatMessage` return type, array params). - **chatd.go**: All call sites converted to batch inserts: - **CreateChat**: System prompt + user message batched into one call - **persistStep**: Assistant message + tool messages batched into one call - **persistSummary**: Hidden summary + assistant + tool messages batched into one call - Single-message sites use the same API with single-element arrays - **Helper**: New `appendChatMessage` function simplifies building batch params at each call site. - **Tests**: All test files updated to use the new API. Builds on top of coder#23213.
1 parent 1f5f6c9 commit 483adc5

11 files changed

Lines changed: 701 additions & 593 deletions

File tree

‎coderd/chatd/chatd.go‎

Lines changed: 208 additions & 255 deletions
Large diffs are not rendered by default.

‎coderd/chatd/chatd_test.go‎

Lines changed: 125 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -710,23 +710,24 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) {
710710
})
711711
require.NoError(t, err)
712712

713-
_, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{
713+
_, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
714714
ChatID: existingChat.ID,
715-
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
716-
Role: database.ChatMessageRoleAssistant,
717-
ContentVersion: chatprompt.CurrentContentVersion,
718-
Content: assistantContent,
719-
Visibility: database.ChatMessageVisibilityBoth,
720-
InputTokens: sql.NullInt64{},
721-
OutputTokens: sql.NullInt64{},
722-
TotalTokens: sql.NullInt64{},
723-
ReasoningTokens: sql.NullInt64{},
724-
CacheCreationTokens: sql.NullInt64{},
725-
CacheReadTokens: sql.NullInt64{},
726-
ContextLimit: sql.NullInt64{},
727-
Compressed: sql.NullBool{},
728-
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
729-
RuntimeMs: sql.NullInt64{},
715+
CreatedBy: []uuid.UUID{uuid.Nil},
716+
ModelConfigID: []uuid.UUID{model.ID},
717+
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
718+
ContentVersion: []int16{chatprompt.CurrentContentVersion},
719+
Content: []string{string(assistantContent.RawMessage)},
720+
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
721+
InputTokens: []int64{0},
722+
OutputTokens: []int64{0},
723+
TotalTokens: []int64{0},
724+
ReasoningTokens: []int64{0},
725+
CacheCreationTokens: []int64{0},
726+
CacheReadTokens: []int64{0},
727+
ContextLimit: []int64{0},
728+
Compressed: []bool{false},
729+
TotalCostMicros: []int64{100},
730+
RuntimeMs: []int64{0},
730731
})
731732
require.NoError(t, err)
732733

@@ -809,23 +810,24 @@ func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing
809810
})
810811
require.NoError(t, err)
811812

812-
_, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{
813+
_, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
813814
ChatID: chat.ID,
814-
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
815-
Role: database.ChatMessageRoleAssistant,
816-
ContentVersion: chatprompt.CurrentContentVersion,
817-
Content: assistantContent,
818-
Visibility: database.ChatMessageVisibilityBoth,
819-
InputTokens: sql.NullInt64{},
820-
OutputTokens: sql.NullInt64{},
821-
TotalTokens: sql.NullInt64{},
822-
ReasoningTokens: sql.NullInt64{},
823-
CacheCreationTokens: sql.NullInt64{},
824-
CacheReadTokens: sql.NullInt64{},
825-
ContextLimit: sql.NullInt64{},
826-
Compressed: sql.NullBool{},
827-
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
828-
RuntimeMs: sql.NullInt64{},
815+
CreatedBy: []uuid.UUID{uuid.Nil},
816+
ModelConfigID: []uuid.UUID{model.ID},
817+
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
818+
ContentVersion: []int16{chatprompt.CurrentContentVersion},
819+
Content: []string{string(assistantContent.RawMessage)},
820+
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
821+
InputTokens: []int64{0},
822+
OutputTokens: []int64{0},
823+
TotalTokens: []int64{0},
824+
ReasoningTokens: []int64{0},
825+
CacheCreationTokens: []int64{0},
826+
CacheReadTokens: []int64{0},
827+
ContextLimit: []int64{0},
828+
Compressed: []bool{false},
829+
TotalCostMicros: []int64{100},
830+
RuntimeMs: []int64{0},
829831
})
830832
require.NoError(t, err)
831833

@@ -998,23 +1000,24 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
9981000
})
9991001
require.NoError(t, err)
10001002

1001-
_, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{
1003+
_, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
10021004
ChatID: spendChat.ID,
1003-
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
1004-
Role: database.ChatMessageRoleAssistant,
1005-
ContentVersion: chatprompt.CurrentContentVersion,
1006-
Content: assistantContent,
1007-
Visibility: database.ChatMessageVisibilityBoth,
1008-
InputTokens: sql.NullInt64{},
1009-
OutputTokens: sql.NullInt64{},
1010-
TotalTokens: sql.NullInt64{},
1011-
ReasoningTokens: sql.NullInt64{},
1012-
CacheCreationTokens: sql.NullInt64{},
1013-
CacheReadTokens: sql.NullInt64{},
1014-
ContextLimit: sql.NullInt64{},
1015-
Compressed: sql.NullBool{},
1016-
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
1017-
RuntimeMs: sql.NullInt64{},
1005+
CreatedBy: []uuid.UUID{uuid.Nil},
1006+
ModelConfigID: []uuid.UUID{model.ID},
1007+
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
1008+
ContentVersion: []int16{chatprompt.CurrentContentVersion},
1009+
Content: []string{string(assistantContent.RawMessage)},
1010+
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
1011+
InputTokens: []int64{0},
1012+
OutputTokens: []int64{0},
1013+
TotalTokens: []int64{0},
1014+
ReasoningTokens: []int64{0},
1015+
CacheCreationTokens: []int64{0},
1016+
CacheReadTokens: []int64{0},
1017+
ContextLimit: []int64{0},
1018+
Compressed: []bool{false},
1019+
TotalCostMicros: []int64{100},
1020+
RuntimeMs: []int64{0},
10181021
})
10191022
require.NoError(t, err)
10201023

@@ -1094,23 +1097,24 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) {
10941097
})
10951098
require.NoError(t, err)
10961099

1097-
_, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{
1100+
_, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
10981101
ChatID: chat.ID,
1099-
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
1100-
Role: database.ChatMessageRoleAssistant,
1101-
ContentVersion: chatprompt.CurrentContentVersion,
1102-
Content: assistantContent,
1103-
Visibility: database.ChatMessageVisibilityBoth,
1104-
InputTokens: sql.NullInt64{},
1105-
OutputTokens: sql.NullInt64{},
1106-
TotalTokens: sql.NullInt64{},
1107-
ReasoningTokens: sql.NullInt64{},
1108-
CacheCreationTokens: sql.NullInt64{},
1109-
CacheReadTokens: sql.NullInt64{},
1110-
ContextLimit: sql.NullInt64{},
1111-
Compressed: sql.NullBool{},
1112-
TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true},
1113-
RuntimeMs: sql.NullInt64{},
1102+
CreatedBy: []uuid.UUID{uuid.Nil},
1103+
ModelConfigID: []uuid.UUID{model.ID},
1104+
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
1105+
ContentVersion: []int16{chatprompt.CurrentContentVersion},
1106+
Content: []string{string(assistantContent.RawMessage)},
1107+
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
1108+
InputTokens: []int64{0},
1109+
OutputTokens: []int64{0},
1110+
TotalTokens: []int64{0},
1111+
ReasoningTokens: []int64{0},
1112+
CacheCreationTokens: []int64{0},
1113+
CacheReadTokens: []int64{0},
1114+
ContextLimit: []int64{0},
1115+
Compressed: []bool{false},
1116+
TotalCostMicros: []int64{100},
1117+
RuntimeMs: []int64{0},
11141118
})
11151119
require.NoError(t, err)
11161120

@@ -1185,24 +1189,27 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) {
11851189
})
11861190
require.NoError(t, err)
11871191

1188-
assistantMessage, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{
1192+
assistantMessages, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
11891193
ChatID: chat.ID,
1190-
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
1191-
Role: database.ChatMessageRoleAssistant,
1192-
ContentVersion: chatprompt.CurrentContentVersion,
1193-
Content: assistantContent,
1194-
Visibility: database.ChatMessageVisibilityBoth,
1195-
InputTokens: sql.NullInt64{},
1196-
OutputTokens: sql.NullInt64{},
1197-
TotalTokens: sql.NullInt64{},
1198-
ReasoningTokens: sql.NullInt64{},
1199-
CacheCreationTokens: sql.NullInt64{},
1200-
CacheReadTokens: sql.NullInt64{},
1201-
ContextLimit: sql.NullInt64{},
1202-
Compressed: sql.NullBool{},
1203-
RuntimeMs: sql.NullInt64{},
1204-
})
1205-
require.NoError(t, err)
1194+
CreatedBy: []uuid.UUID{uuid.Nil},
1195+
ModelConfigID: []uuid.UUID{model.ID},
1196+
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
1197+
ContentVersion: []int16{chatprompt.CurrentContentVersion},
1198+
Content: []string{string(assistantContent.RawMessage)},
1199+
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
1200+
InputTokens: []int64{0},
1201+
OutputTokens: []int64{0},
1202+
TotalTokens: []int64{0},
1203+
ReasoningTokens: []int64{0},
1204+
CacheCreationTokens: []int64{0},
1205+
CacheReadTokens: []int64{0},
1206+
ContextLimit: []int64{0},
1207+
Compressed: []bool{false},
1208+
TotalCostMicros: []int64{0},
1209+
RuntimeMs: []int64{0},
1210+
})
1211+
require.NoError(t, err)
1212+
assistantMessage := assistantMessages[0]
12061213

12071214
_, err = replica.EditMessage(ctx, chatd.EditMessageOptions{
12081215
ChatID: chat.ID,
@@ -1534,46 +1541,51 @@ func TestSubscribeAfterMessageID(t *testing.T) {
15341541
})
15351542
require.NoError(t, err)
15361543

1537-
msg2, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{
1544+
msg2Results, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
15381545
ChatID: chat.ID,
1539-
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
1540-
Role: database.ChatMessageRoleAssistant,
1541-
ContentVersion: chatprompt.CurrentContentVersion,
1542-
Content: secondContent,
1543-
Visibility: database.ChatMessageVisibilityBoth,
1544-
InputTokens: sql.NullInt64{},
1545-
OutputTokens: sql.NullInt64{},
1546-
TotalTokens: sql.NullInt64{},
1547-
ReasoningTokens: sql.NullInt64{},
1548-
CacheCreationTokens: sql.NullInt64{},
1549-
CacheReadTokens: sql.NullInt64{},
1550-
ContextLimit: sql.NullInt64{},
1551-
Compressed: sql.NullBool{},
1552-
RuntimeMs: sql.NullInt64{},
1553-
})
1554-
require.NoError(t, err)
1546+
CreatedBy: []uuid.UUID{uuid.Nil},
1547+
ModelConfigID: []uuid.UUID{model.ID},
1548+
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
1549+
ContentVersion: []int16{chatprompt.CurrentContentVersion},
1550+
Content: []string{string(secondContent.RawMessage)},
1551+
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
1552+
InputTokens: []int64{0},
1553+
OutputTokens: []int64{0},
1554+
TotalTokens: []int64{0},
1555+
ReasoningTokens: []int64{0},
1556+
CacheCreationTokens: []int64{0},
1557+
CacheReadTokens: []int64{0},
1558+
ContextLimit: []int64{0},
1559+
Compressed: []bool{false},
1560+
TotalCostMicros: []int64{0},
1561+
RuntimeMs: []int64{0},
1562+
})
1563+
require.NoError(t, err)
1564+
msg2 := msg2Results[0]
15551565

15561566
thirdContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{
15571567
codersdk.ChatMessageText("third"),
15581568
})
15591569
require.NoError(t, err)
15601570

1561-
_, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{
1571+
_, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
15621572
ChatID: chat.ID,
1563-
ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true},
1564-
Role: database.ChatMessageRoleUser,
1565-
ContentVersion: chatprompt.CurrentContentVersion,
1566-
Content: thirdContent,
1567-
Visibility: database.ChatMessageVisibilityBoth,
1568-
InputTokens: sql.NullInt64{},
1569-
OutputTokens: sql.NullInt64{},
1570-
TotalTokens: sql.NullInt64{},
1571-
ReasoningTokens: sql.NullInt64{},
1572-
CacheCreationTokens: sql.NullInt64{},
1573-
CacheReadTokens: sql.NullInt64{},
1574-
ContextLimit: sql.NullInt64{},
1575-
Compressed: sql.NullBool{},
1576-
RuntimeMs: sql.NullInt64{},
1573+
CreatedBy: []uuid.UUID{uuid.Nil},
1574+
ModelConfigID: []uuid.UUID{model.ID},
1575+
Role: []database.ChatMessageRole{database.ChatMessageRoleUser},
1576+
ContentVersion: []int16{chatprompt.CurrentContentVersion},
1577+
Content: []string{string(thirdContent.RawMessage)},
1578+
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
1579+
InputTokens: []int64{0},
1580+
OutputTokens: []int64{0},
1581+
TotalTokens: []int64{0},
1582+
ReasoningTokens: []int64{0},
1583+
CacheCreationTokens: []int64{0},
1584+
CacheReadTokens: []int64{0},
1585+
ContextLimit: []int64{0},
1586+
Compressed: []bool{false},
1587+
TotalCostMicros: []int64{0},
1588+
RuntimeMs: []int64{0},
15771589
})
15781590
require.NoError(t, err)
15791591

0 commit comments

Comments
 (0)