feat: replace InsertChatMessage with batch InsertChatMessages (#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 #23213.
This commit is contained in:
Kyle Carberry
2026-03-18 16:27:07 +00:00
committed by GitHub
parent 1f5f6c9ccb
commit 483adc59fe
11 changed files with 692 additions and 584 deletions
+40 -22
View File
@@ -178,16 +178,34 @@ INSERT INTO chats (
RETURNING
*;
-- name: InsertChatMessage :one
-- name: InsertChatMessages :many
WITH updated_chat AS (
UPDATE
chats
SET
last_model_config_id = sqlc.narg('model_config_id')::uuid
last_model_config_id = (
SELECT val
FROM unnest(@model_config_id::uuid[])
WITH ORDINALITY AS t(val, ord)
WHERE val != '00000000-0000-0000-0000-000000000000'::uuid
ORDER BY ord DESC
LIMIT 1
)
WHERE
id = @chat_id::uuid
AND sqlc.narg('model_config_id')::uuid IS NOT NULL
AND chats.last_model_config_id IS DISTINCT FROM sqlc.narg('model_config_id')::uuid
AND EXISTS (
SELECT 1
FROM unnest(@model_config_id::uuid[])
WHERE unnest != '00000000-0000-0000-0000-000000000000'::uuid
)
AND chats.last_model_config_id IS DISTINCT FROM (
SELECT val
FROM unnest(@model_config_id::uuid[])
WITH ORDINALITY AS t(val, ord)
WHERE val != '00000000-0000-0000-0000-000000000000'::uuid
ORDER BY ord DESC
LIMIT 1
)
)
INSERT INTO chat_messages (
chat_id,
@@ -207,25 +225,25 @@ INSERT INTO chat_messages (
compressed,
total_cost_micros,
runtime_ms
) VALUES (
@chat_id::uuid,
sqlc.narg('created_by')::uuid,
sqlc.narg('model_config_id')::uuid,
@role::chat_message_role,
sqlc.narg('content')::jsonb,
@content_version::smallint,
@visibility::chat_message_visibility,
sqlc.narg('input_tokens')::bigint,
sqlc.narg('output_tokens')::bigint,
sqlc.narg('total_tokens')::bigint,
sqlc.narg('reasoning_tokens')::bigint,
sqlc.narg('cache_creation_tokens')::bigint,
sqlc.narg('cache_read_tokens')::bigint,
sqlc.narg('context_limit')::bigint,
COALESCE(sqlc.narg('compressed')::boolean, FALSE),
sqlc.narg('total_cost_micros')::bigint,
sqlc.narg('runtime_ms')::bigint
)
SELECT
@chat_id::uuid,
NULLIF(unnest(@created_by::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
NULLIF(unnest(@model_config_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid),
unnest(@role::chat_message_role[]),
unnest(@content::text[])::jsonb,
unnest(@content_version::smallint[]),
unnest(@visibility::chat_message_visibility[]),
NULLIF(unnest(@input_tokens::bigint[]), 0),
NULLIF(unnest(@output_tokens::bigint[]), 0),
NULLIF(unnest(@total_tokens::bigint[]), 0),
NULLIF(unnest(@reasoning_tokens::bigint[]), 0),
NULLIF(unnest(@cache_creation_tokens::bigint[]), 0),
NULLIF(unnest(@cache_read_tokens::bigint[]), 0),
NULLIF(unnest(@context_limit::bigint[]), 0),
unnest(@compressed::boolean[]),
NULLIF(unnest(@total_cost_micros::bigint[]), 0),
NULLIF(unnest(@runtime_ms::bigint[]), 0)
RETURNING
*;