diff --git a/go/core/internal/database/client_langgraph_test.go b/go/core/internal/database/client_langgraph_test.go new file mode 100644 index 0000000000..d0aaf034dd --- /dev/null +++ b/go/core/internal/database/client_langgraph_test.go @@ -0,0 +1,137 @@ +package database + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + dbpkg "github.com/kagent-dev/kagent/go/api/database" + "github.com/stretchr/testify/require" +) + +// countingTracer is a pgx.QueryTracer that counts how many executed statements +// contain a given SQL substring. It is used to guard against the N+1 query +// pattern in ListCheckpoints. +type countingTracer struct { + substr string + mu sync.Mutex + count int +} + +func (t *countingTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context { + if strings.Contains(data.SQL, t.substr) { + t.mu.Lock() + t.count++ + t.mu.Unlock() + } + return ctx +} + +func (t *countingTracer) TraceQueryEnd(context.Context, *pgx.Conn, pgx.TraceQueryEndData) {} + +func (t *countingTracer) reset() { + t.mu.Lock() + t.count = 0 + t.mu.Unlock() +} + +func (t *countingTracer) value() int { + t.mu.Lock() + defer t.mu.Unlock() + return t.count +} + +// newCountingClient opens a single-connection pool against the shared test +// database with a query-counting tracer attached, so a test can assert how many +// times a particular query runs. +func newCountingClient(t *testing.T, tracer *countingTracer) dbpkg.Client { + t.Helper() + + config, err := pgxpool.ParseConfig(sharedConnStr) + require.NoError(t, err) + config.MaxConns = 1 + config.ConnConfig.Tracer = tracer + + pool, err := pgxpool.NewWithConfig(context.Background(), config) + require.NoError(t, err) + t.Cleanup(pool.Close) + + return NewClient(pool) +} + +// TestListCheckpointsBatchesWrites verifies that reading a thread's checkpoint +// history fetches all checkpoint writes in a single query instead of one query +// per checkpoint (the N+1 pattern). It also checks the writes are grouped onto +// the correct checkpoints and keep their per-checkpoint ordering. +func TestListCheckpointsBatchesWrites(t *testing.T) { + setupTestDB(t) + + tracer := &countingTracer{substr: "FROM lg_checkpoint_write"} + client := newCountingClient(t, tracer) + ctx := context.Background() + + const ( + userID = "test-user" + threadID = "test-thread" + checkpointNS = "" + numCP = 5 + writesPerCP = 3 + ) + + // Seed several checkpoints, each with multiple writes. + for i := range numCP { + cpID := fmt.Sprintf("cp-%02d", i) + err := client.StoreCheckpoint(ctx, &dbpkg.LangGraphCheckpoint{ + UserID: userID, + ThreadID: threadID, + CheckpointNS: checkpointNS, + CheckpointID: cpID, + Metadata: "{}", + Checkpoint: "{}", + Version: 1, + }) + require.NoError(t, err) + + writes := make([]*dbpkg.LangGraphCheckpointWrite, writesPerCP) + for j := range writesPerCP { + writes[j] = &dbpkg.LangGraphCheckpointWrite{ + UserID: userID, + ThreadID: threadID, + CheckpointNS: checkpointNS, + CheckpointID: cpID, + WriteIdx: int64(j), + Value: fmt.Sprintf("value-%d", j), + ValueType: "json", + Channel: fmt.Sprintf("channel-%d", j), + TaskID: "task-0", + } + } + require.NoError(t, client.StoreCheckpointWrites(ctx, writes)) + } + + // Read the whole thread history (checkpointID nil, limit 0 -> unbounded). + tracer.reset() + tuples, err := client.ListCheckpoints(ctx, userID, threadID, checkpointNS, nil, 0) + require.NoError(t, err) + + // Exactly one query against lg_checkpoint_write, regardless of how many + // checkpoints the thread has. Before batching this was one query per + // checkpoint (numCP). + writeQueries := tracer.value() + require.Equal(t, 1, writeQueries, + "expected a single checkpoint-writes query, got %d (N+1 regression)", writeQueries) + + // Every checkpoint's writes are attached, in write_idx order. + require.Len(t, tuples, numCP) + for _, tuple := range tuples { + require.Len(t, tuple.Writes, writesPerCP, "checkpoint %s", tuple.Checkpoint.CheckpointID) + for j, w := range tuple.Writes { + require.Equal(t, int64(j), w.WriteIdx) + require.Equal(t, tuple.Checkpoint.CheckpointID, w.CheckpointID) + } + } +} diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 4e0bdc66d9..4b56d1f96f 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -550,17 +550,33 @@ func (c *postgresClient) ListCheckpoints(ctx context.Context, userID, threadID, } } - tuples = make([]*dbpkg.LangGraphCheckpointTuple, 0, len(checkpoints)) - for _, cp := range checkpoints { - writes, err := q.ListCheckpointWrites(ctx, dbgen.ListCheckpointWritesParams{ - UserID: userID, ThreadID: threadID, CheckpointNs: checkpointNS, CheckpointID: cp.CheckpointID, + // Fetch all writes for the returned checkpoints in a single query and + // bucket them by checkpoint ID, instead of issuing one query per + // checkpoint (which turned reading a thread's history into 1+N round + // trips that grew with the conversation length). + checkpointIDs := make([]string, len(checkpoints)) + for i, cp := range checkpoints { + checkpointIDs[i] = cp.CheckpointID + } + + writesByCheckpoint := make(map[string][]*dbpkg.LangGraphCheckpointWrite, len(checkpoints)) + if len(checkpointIDs) > 0 { + writes, err := q.ListCheckpointWritesForCheckpoints(ctx, dbgen.ListCheckpointWritesForCheckpointsParams{ + UserID: userID, ThreadID: threadID, CheckpointNs: checkpointNS, Column4: checkpointIDs, }) if err != nil { return fmt.Errorf("failed to get checkpoint writes: %w", err) } - dbWrites := make([]*dbpkg.LangGraphCheckpointWrite, len(writes)) - for i, w := range writes { - dbWrites[i] = toCheckpointWrite(w) + for _, w := range writes { + writesByCheckpoint[w.CheckpointID] = append(writesByCheckpoint[w.CheckpointID], toCheckpointWrite(w)) + } + } + + tuples = make([]*dbpkg.LangGraphCheckpointTuple, 0, len(checkpoints)) + for _, cp := range checkpoints { + dbWrites := writesByCheckpoint[cp.CheckpointID] + if dbWrites == nil { + dbWrites = []*dbpkg.LangGraphCheckpointWrite{} } tuples = append(tuples, &dbpkg.LangGraphCheckpointTuple{ Checkpoint: toCheckpoint(cp), diff --git a/go/core/internal/database/gen/langgraph.sql.go b/go/core/internal/database/gen/langgraph.sql.go index eb483f1387..c720acc2db 100644 --- a/go/core/internal/database/gen/langgraph.sql.go +++ b/go/core/internal/database/gen/langgraph.sql.go @@ -100,6 +100,58 @@ func (q *Queries) ListCheckpointWrites(ctx context.Context, arg ListCheckpointWr return items, nil } +const listCheckpointWritesForCheckpoints = `-- name: ListCheckpointWritesForCheckpoints :many +SELECT user_id, thread_id, checkpoint_ns, checkpoint_id, write_idx, value, value_type, channel, task_id, created_at, updated_at, deleted_at 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 +` + +type ListCheckpointWritesForCheckpointsParams struct { + UserID string + ThreadID string + CheckpointNs string + Column4 []string +} + +func (q *Queries) ListCheckpointWritesForCheckpoints(ctx context.Context, arg ListCheckpointWritesForCheckpointsParams) ([]LgCheckpointWrite, error) { + rows, err := q.db.Query(ctx, listCheckpointWritesForCheckpoints, + arg.UserID, + arg.ThreadID, + arg.CheckpointNs, + arg.Column4, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []LgCheckpointWrite + for rows.Next() { + var i LgCheckpointWrite + if err := rows.Scan( + &i.UserID, + &i.ThreadID, + &i.CheckpointNs, + &i.CheckpointID, + &i.WriteIdx, + &i.Value, + &i.ValueType, + &i.Channel, + &i.TaskID, + &i.CreatedAt, + &i.UpdatedAt, + &i.DeletedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const listCheckpoints = `-- name: ListCheckpoints :many SELECT user_id, thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, created_at, updated_at, deleted_at, metadata, checkpoint, checkpoint_type, version FROM lg_checkpoint WHERE user_id = $1 AND thread_id = $2 AND checkpoint_ns = $3 diff --git a/go/core/internal/database/gen/querier.go b/go/core/internal/database/gen/querier.go index 078a065021..cbf459618e 100644 --- a/go/core/internal/database/gen/querier.go +++ b/go/core/internal/database/gen/querier.go @@ -33,6 +33,7 @@ type Querier interface { ListAgentMemories(ctx context.Context, arg ListAgentMemoriesParams) ([]Memory, error) ListAgents(ctx context.Context) ([]Agent, error) ListCheckpointWrites(ctx context.Context, arg ListCheckpointWritesParams) ([]LgCheckpointWrite, error) + ListCheckpointWritesForCheckpoints(ctx context.Context, arg ListCheckpointWritesForCheckpointsParams) ([]LgCheckpointWrite, error) ListCheckpoints(ctx context.Context, arg ListCheckpointsParams) ([]LgCheckpoint, error) ListCheckpointsLimit(ctx context.Context, arg ListCheckpointsLimitParams) ([]LgCheckpoint, error) ListEventsByContextID(ctx context.Context, sessionID *string) ([]Event, error) diff --git a/go/core/internal/database/queries/langgraph.sql b/go/core/internal/database/queries/langgraph.sql index bd20481cf4..71423c3b3c 100644 --- a/go/core/internal/database/queries/langgraph.sql +++ b/go/core/internal/database/queries/langgraph.sql @@ -37,6 +37,12 @@ WHERE user_id = $1 AND thread_id = $2 AND checkpoint_ns = $3 AND checkpoint_id = $4 AND deleted_at IS NULL ORDER BY task_id, write_idx; +-- 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; + -- name: UpsertCheckpointWrite :exec INSERT INTO lg_checkpoint_write ( user_id, thread_id, checkpoint_ns, checkpoint_id, write_idx,