perf(database): batch checkpoint writes to remove N+1 query

[anxkhn]

What

Reading a LangGraph thread's checkpoint history through
postgresClient.ListCheckpoints
(go/core/internal/database/client_postgres.go) issued 1+N database queries.
After loading the thread's checkpoints, it looped over each one and ran a
separate ListCheckpointWrites query per checkpoint:

for _, cp := range checkpoints {
    writes, err := q.ListCheckpointWrites(ctx, dbgen.ListCheckpointWritesParams{
        UserID: userID, ThreadID: threadID, CheckpointNs: checkpointNS, CheckpointID: cp.CheckpointID,
    })
    ...
}

All of these run inside the single transaction opened at the top of
ListCheckpoints, so they are serialized on one pooled connection.

This is the normal history-read path, not an edge case:

• GET /api/langgraph/checkpoints (route in
  go/core/internal/httpserver/server.go) calls HandleListCheckpoints
  (go/core/internal/httpserver/handlers/checkpoints.go), which reads the
  limit query param defaulting to 0 and passes it straight through.
• When no single checkpointID is requested and limit is not > 0,
  ListCheckpoints takes the unbounded list branch (q.ListCheckpoints, whose
  SQL has no LIMIT).
• The Python LangGraph checkpointer sends limit=-1 whenever no limit is set
  (python/packages/kagent-langgraph/src/kagent/langgraph/_checkpointer.py:
  limit = limit if limit else -1), which is not > 0, so it hits that branch.

A LangGraph thread accumulates one checkpoint per superstep, so N grows with
conversation length. Reading a long thread's history therefore issues 1+N
queries and holds the connection for the whole fan-out.

Fix

Add a single :many query that loads the writes for all returned checkpoints
at once, using checkpoint_id = ANY(...)
(go/core/internal/database/queries/langgraph.sql):

-- name: ListCheckpointWritesForCheckpoints :many
SELECT * FROM lg_checkpoint_write
WHERE user_id = $1 AND thread_id = $2 AND checkpoint_ns = $3
  AND checkpoint_id = ANY($4::text[]) AND deleted_at IS NULL
ORDER BY checkpoint_id, task_id, write_idx;

ListCheckpoints now runs that one query, buckets the rows into a
map[checkpointID][]write in a single pass, and attaches each slice to its
tuple. This turns the write fan-out into one query regardless of history length
(so the whole call is 2 queries instead of 1+N).

Behavior is preserved:

• The ANY(...) pattern already exists in this codebase (see
  IncrementMemoryAccessCount in queries/memory.sql, which uses
  ANY($1::text[])).
• Per-checkpoint write ordering is unchanged: the original singular query
  ordered by task_id, write_idx; the batched query orders by
  checkpoint_id, task_id, write_idx, so within each checkpoint the relative
  order is identical.
• The non-nil empty-slice contract is kept: a checkpoint with no writes still
  gets []*LangGraphCheckpointWrite{}, matching the previous
  make(..., len(writes)) result.
• The single-checkpoint path (GetCheckpoint) flows through the same code with
  a one-element batch and yields the same result.

The generated code (gen/langgraph.sql.go, gen/querier.go) is the committed
output of sqlc v1.30.0 (make -C go sqlc-generate produces no diff).

Tests

Added TestListCheckpointsBatchesWrites in
go/core/internal/database/client_langgraph_test.go. It seeds a thread with 5
checkpoints (3 writes each), reads the full history, and asserts:

• exactly one lg_checkpoint_write query runs (a pgx.QueryTracer counts
  statements containing FROM lg_checkpoint_write) - this is the N+1 guard, and
  it was numCP before the change, and
• the writes are grouped onto the correct checkpoints in write_idx order.

The test uses the same testcontainers-backed Postgres as the existing
database tests, so it skips under go test -short and runs in CI.

Locally (this change):

• make -C go sqlc-generate -> no diff
• go build ./..., go vet ./core/internal/database/... -> pass
• make -C go lint (golangci-lint on the database package) -> 0 issues
• go test -short ./core/internal/database/... ./core/internal/httpserver/...
  -> pass

Note

The pre-existing singular ListCheckpointWrites query (and its generated
method) is left in place: it is no longer called from Go after this change, but
removing it would enlarge the generated interface diff. Happy to drop it in a
follow-up (or in this PR) if you prefer.

[Copilot]

Pull request overview

This PR removes an N+1 query pattern in the Go Postgres database client when listing LangGraph checkpoint history by batching checkpoint-write reads into a single SQL query, improving performance and reducing connection hold time in the normal history-read path.

Changes:

• Add a new batched SQL query (ListCheckpointWritesForCheckpoints) that loads writes for multiple checkpoints using checkpoint_id = ANY(...).
• Update postgresClient.ListCheckpoints to fetch all writes in one query and bucket them by checkpoint ID before assembling tuples.
• Add a database test that guards against regressions back to the N+1 query pattern and validates write grouping/order.

Reviewed changes

Copilot reviewed 3 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
go/core/internal/database/queries/langgraph.sql	Adds a new :many query to fetch writes for multiple checkpoints in one call.
go/core/internal/database/gen/querier.go	Extends the generated Querier interface with the new batched query method.
go/core/internal/database/gen/langgraph.sql.go	Adds sqlc-generated method/params for ListCheckpointWritesForCheckpoints.
go/core/internal/database/client_postgres.go	Replaces per-checkpoint write fan-out with a single batched query and bucketing logic.
go/core/internal/database/client_langgraph_test.go	Adds a regression test to ensure only one checkpoint-write SELECT runs and writes are correctly grouped/ordered.

Files not reviewed (2)

• go/core/internal/database/gen/langgraph.sql.go: Generated file
• go/core/internal/database/gen/querier.go: Generated file

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[EItanya]

The pre-existing singular ListCheckpointWrites query (and its generated
method) is left in place: it is no longer called from Go after this change, but
removing it would enlarge the generated interface diff. Happy to drop it in a
follow-up (or in this PR) if you prefer.

Let's just do it here since it's unused

[mesutoezdil]

the singular listcheckpointwrites query above is now dead code since this pr removed its only caller, drop it and regenerate so the sqlc surface stays in sync.