From d079879e0af40c2fc946518fe547c4ce62e86a41 Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Mon, 20 Jul 2026 15:26:46 +0200 Subject: [PATCH 1/2] perf(a2a): push ListTasks all-sessions query down to SQL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit collectUserTasks loaded every session for a user, then every task per session, then filtered, sorted, and paginated in memory (an N+1 query whose result set and app-side memory scaled with the user's total task count even for a small page). Replace it with a single ListUserTasks sqlc query that joins session to task, filters, orders by task id, and paginates server-side, returning COUNT(*) OVER() as the total. Status and statusTimestampAfter are filtered in SQL via a jsonb cast on task.data (no schema migration). Rows exist in two shapes — v1 (protocol_version 'v1', state "TASK_STATE_WORKING") and legacy (protocol_version NULL, state "working") — so both spellings are passed and matched; the vocabularies never overlap. The timestamp cast is guarded by a CASE so a malformed value cannot error the query. COUNT(*) OVER() rides on the returned rows, so a page requested past the end of the set (e.g. rows deleted between paginated calls) comes back empty and carries no total; CountUserTasks recovers the true filtered count in that case. The single-session path keeps its fail-closed GetSession ownership check; the join scopes every path to the caller's own sessions. Adds a Postgres integration test covering both row shapes, cross-user isolation, the session predicate, status/timestamp filters, pagination, and the empty over-range page total. Signed-off-by: QuentinBisson --- go/api/database/client.go | 14 ++ go/core/internal/a2a/task_query_store.go | 93 +++------- go/core/internal/a2a/task_query_store_test.go | 55 ++++-- go/core/internal/database/client_postgres.go | 52 ++++++ go/core/internal/database/client_test.go | 132 ++++++++++++++ go/core/internal/database/gen/querier.go | 26 +++ go/core/internal/database/gen/tasks.sql.go | 162 ++++++++++++++++++ go/core/internal/database/queries/tasks.sql | 78 +++++++++ go/core/pkg/a2acompat/trpcv0/convert.go | 7 + 9 files changed, 534 insertions(+), 85 deletions(-) diff --git a/go/api/database/client.go b/go/api/database/client.go index b58942479..272165865 100644 --- a/go/api/database/client.go +++ b/go/api/database/client.go @@ -19,6 +19,19 @@ type QueryOptions struct { After time.Time OrderAsc bool // When true, order results by created_at ASC (chronological). Default is DESC (newest first). } + +// ListUserTasksParams scopes and paginates a user's tasks. SessionID empty +// lists across every session the user owns. Status empty (TaskStateUnspecified) +// disables the status filter; StatusTimestampAfter nil disables the timestamp +// filter. Results are ordered by task id. +type ListUserTasksParams struct { + UserID string + SessionID string + Status a2a.TaskState + StatusTimestampAfter *time.Time + Limit int + Offset int +} type LangGraphCheckpointTuple struct { Checkpoint *LangGraphCheckpoint Writes []*LangGraphCheckpointWrite @@ -55,6 +68,7 @@ type Client interface { ListTools(ctx context.Context) ([]Tool, error) ListFeedback(ctx context.Context, userID string) ([]Feedback, error) ListTasksForSession(ctx context.Context, sessionID string, userID string) ([]*a2a.Task, error) + ListUserTasks(ctx context.Context, params ListUserTasksParams) (tasks []*a2a.Task, total int, err error) ListSessions(ctx context.Context, userID string) ([]Session, error) ListSessionsForAgent(ctx context.Context, agentID string, userID string) ([]SessionWithShareToken, error) ListSessionsForAgentAllUsers(ctx context.Context, agentID string) ([]Session, error) diff --git a/go/core/internal/a2a/task_query_store.go b/go/core/internal/a2a/task_query_store.go index a121af135..432dda833 100644 --- a/go/core/internal/a2a/task_query_store.go +++ b/go/core/internal/a2a/task_query_store.go @@ -1,11 +1,9 @@ package a2a import ( - "cmp" "context" "encoding/base64" "fmt" - "slices" "strconv" a2atype "github.com/a2aproject/a2a-go/v2/a2a" @@ -21,11 +19,11 @@ const ( // TaskStore is the subset of the persistent store ListTasks reads from. // *database.Client satisfies it. GetSession errors (including a missing or -// other-user session) surface to the caller. +// other-user session) surface to the caller. ListUserTasks filters, orders, +// and paginates server-side and returns the full filtered count. type TaskStore interface { GetSession(ctx context.Context, sessionID, userID string) (*dbpkg.Session, error) - ListSessions(ctx context.Context, userID string) ([]dbpkg.Session, error) - ListTasksForSession(ctx context.Context, sessionID, userID string) ([]*a2atype.Task, error) + ListUserTasks(ctx context.Context, params dbpkg.ListUserTasksParams) (tasks []*a2atype.Task, total int, err error) } // storeTaskQueryHandler answers ListTasks from kagent's task store, which is @@ -75,34 +73,39 @@ func (h *storeTaskQueryHandler) ListTasks(ctx context.Context, req *a2atype.List return &a2atype.ListTasksResponse{Tasks: []*a2atype.Task{}, PageSize: pageSize}, nil } - tasks, err := h.collectUserTasks(ctx, userID, req.ContextID) + offset, err := decodePageToken(req.PageToken) if err != nil { - return nil, err + return nil, a2atype.NewError(a2atype.ErrInvalidParams, "invalid pageToken") } - filtered := filterTasks(tasks, req) - // Order by task id so the page-token offset is stable across calls: task - // ids are immutable, unlike session updated_at (which reorders on writes). - slices.SortFunc(filtered, func(a, b *a2atype.Task) int { return cmp.Compare(a.ID, b.ID) }) + // A single-session query fails closed: a missing or other-user session is an + // error, not an empty page. The join in ListUserTasks also scopes by user, so + // this only distinguishes the error case and keeps share-context validation. + if req.ContextID != "" { + if _, err := h.store.GetSession(ctx, req.ContextID, userID); err != nil { + return nil, fmt.Errorf("get session %s: %w", req.ContextID, err) + } + } - offset, err := decodePageToken(req.PageToken) + tasks, total, err := h.store.ListUserTasks(ctx, dbpkg.ListUserTasksParams{ + UserID: userID, + SessionID: req.ContextID, + Status: req.Status, + StatusTimestampAfter: req.StatusTimestampAfter, + Limit: pageSize, + Offset: offset, + }) if err != nil { - return nil, a2atype.NewError(a2atype.ErrInvalidParams, "invalid pageToken") - } - total := len(filtered) - if offset > total { - offset = total + return nil, err } - end := min(offset+pageSize, total) - page := filtered[offset:end] - shaped := make([]*a2atype.Task, 0, len(page)) - for _, t := range page { + shaped := make([]*a2atype.Task, 0, len(tasks)) + for _, t := range tasks { shaped = append(shaped, shapeTask(t, req.HistoryLength, req.IncludeArtifacts)) } nextToken := "" - if end < total { + if end := offset + len(tasks); end < total { nextToken = encodePageToken(end) } @@ -114,52 +117,6 @@ func (h *storeTaskQueryHandler) ListTasks(ctx context.Context, req *a2atype.List }, nil } -// collectUserTasks returns the caller's tasks, either for a single session -// (contextId) or across every session the user owns. Both paths are strictly -// scoped to userID. -func (h *storeTaskQueryHandler) collectUserTasks(ctx context.Context, userID, contextID string) ([]*a2atype.Task, error) { - if contextID != "" { - if _, err := h.store.GetSession(ctx, contextID, userID); err != nil { - return nil, fmt.Errorf("get session %s: %w", contextID, err) - } - return h.store.ListTasksForSession(ctx, contextID, userID) - } - - sessions, err := h.store.ListSessions(ctx, userID) - if err != nil { - return nil, fmt.Errorf("list sessions: %w", err) - } - var all []*a2atype.Task - for _, s := range sessions { - tasks, err := h.store.ListTasksForSession(ctx, s.ID, userID) - if err != nil { - return nil, fmt.Errorf("list tasks for session %s: %w", s.ID, err) - } - all = append(all, tasks...) - } - return all, nil -} - -func filterTasks(tasks []*a2atype.Task, req *a2atype.ListTasksRequest) []*a2atype.Task { - filtered := make([]*a2atype.Task, 0, len(tasks)) - for _, t := range tasks { - if t == nil { - continue - } - if req.Status != a2atype.TaskStateUnspecified && t.Status.State != req.Status { - continue - } - if req.StatusTimestampAfter != nil { - ts := t.Status.Timestamp - if ts == nil || !ts.After(*req.StatusTimestampAfter) { - continue - } - } - filtered = append(filtered, t) - } - return filtered -} - // shapeTask returns a copy of task with history capped and artifacts included // only when requested. includeArtifacts defaults to false, in which case // artifacts are omitted entirely (nil slice + omitempty). diff --git a/go/core/internal/a2a/task_query_store_test.go b/go/core/internal/a2a/task_query_store_test.go index d2c650590..b2b98ccda 100644 --- a/go/core/internal/a2a/task_query_store_test.go +++ b/go/core/internal/a2a/task_query_store_test.go @@ -1,11 +1,13 @@ package a2a import ( + "cmp" "context" "encoding/json" "fmt" "net/http" "net/http/httptest" + "slices" "strings" "testing" "time" @@ -48,21 +50,44 @@ func (f *fakeTaskStore) GetSession(_ context.Context, sessionID, userID string) return &s, nil } -func (f *fakeTaskStore) ListSessions(_ context.Context, userID string) ([]dbpkg.Session, error) { - var out []dbpkg.Session - for _, s := range f.sessions { - if s.UserID == userID { - out = append(out, s) +// ListUserTasks mirrors the SQL query's semantics in memory: scope to the +// user's own sessions (optionally one), filter status/timestamp, order by task +// id, and paginate, returning the full filtered count as total. +func (f *fakeTaskStore) ListUserTasks(_ context.Context, params dbpkg.ListUserTasksParams) ([]*a2atype.Task, int, error) { + var filtered []*a2atype.Task + for sessionID, tasks := range f.tasks { + s, ok := f.sessions[sessionID] + if !ok || s.UserID != params.UserID { + continue + } + if params.SessionID != "" && sessionID != params.SessionID { + continue + } + for _, t := range tasks { + if t == nil { + continue + } + if params.Status != a2atype.TaskStateUnspecified && t.Status.State != params.Status { + continue + } + if params.StatusTimestampAfter != nil { + ts := t.Status.Timestamp + if ts == nil || !ts.After(*params.StatusTimestampAfter) { + continue + } + } + filtered = append(filtered, t) } } - return out, nil -} + slices.SortFunc(filtered, func(a, b *a2atype.Task) int { return cmp.Compare(a.ID, b.ID) }) -func (f *fakeTaskStore) ListTasksForSession(_ context.Context, sessionID, userID string) ([]*a2atype.Task, error) { - if s, ok := f.sessions[sessionID]; !ok || s.UserID != userID { - return nil, nil + total := len(filtered) + offset := min(params.Offset, total) + end := total + if params.Limit > 0 && offset+params.Limit < end { + end = offset + params.Limit } - return f.tasks[sessionID], nil + return filtered[offset:end], total, nil } // fakeSession injects a user principal into the request context. @@ -270,12 +295,8 @@ func (f failingTaskStore) GetSession(context.Context, string, string) (*dbpkg.Se return nil, f.err } -func (f failingTaskStore) ListSessions(context.Context, string) ([]dbpkg.Session, error) { - return nil, f.err -} - -func (f failingTaskStore) ListTasksForSession(context.Context, string, string) ([]*a2atype.Task, error) { - return nil, f.err +func (f failingTaskStore) ListUserTasks(context.Context, dbpkg.ListUserTasksParams) ([]*a2atype.Task, int, error) { + return nil, 0, f.err } func TestListTasks_BackendFailurePropagates(t *testing.T) { diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 064825d18..9f3462f3e 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -342,6 +342,58 @@ func (c *postgresClient) ListTasksForSession(ctx context.Context, sessionID, use return tasks, nil } +func (c *postgresClient) ListUserTasks(ctx context.Context, params dbpkg.ListUserTasksParams) ([]*a2a.Task, int, error) { + arg := dbgen.ListUserTasksParams{ + UserID: params.UserID, + StatusAfter: params.StatusTimestampAfter, + PageOffset: int32(params.Offset), + PageLimit: int32(params.Limit), + } + if params.SessionID != "" { + arg.SessionID = ¶ms.SessionID + } + if params.Status != a2a.TaskStateUnspecified { + v1 := string(params.Status) + legacy := trpcv0.LegacyTaskStateString(params.Status) + arg.StatusV1 = &v1 + arg.StatusLegacy = &legacy + } + + rows, err := c.q.ListUserTasks(ctx, arg) + if err != nil { + return nil, 0, fmt.Errorf("failed to list user tasks: %w", err) + } + + total := 0 + tasks := make([]*a2a.Task, 0, len(rows)) + for i, r := range rows { + total = int(r.Total) + task, err := parseVersionedTask(r.Data, r.ProtocolVersion) + if err != nil { + return nil, 0, fmt.Errorf("failed to parse task row %d: %w", i, err) + } + tasks = append(tasks, task) + } + + // COUNT(*) OVER() rides on the returned rows, so a page requested past the + // end of the set comes back empty and carries no total. Recover it directly + // so the caller still sees the true filtered count. + if len(rows) == 0 && params.Offset > 0 { + count, err := c.q.CountUserTasks(ctx, dbgen.CountUserTasksParams{ + UserID: arg.UserID, + SessionID: arg.SessionID, + StatusV1: arg.StatusV1, + StatusLegacy: arg.StatusLegacy, + StatusAfter: arg.StatusAfter, + }) + if err != nil { + return nil, 0, fmt.Errorf("failed to count user tasks: %w", err) + } + total = int(count) + } + return tasks, total, nil +} + func (c *postgresClient) DeleteTask(ctx context.Context, taskID, userID string) error { if err := c.q.SoftDeleteTask(ctx, dbgen.SoftDeleteTaskParams{ID: taskID, UserID: &userID}); err != nil { return fmt.Errorf("failed to delete task %s: %w", taskID, err) diff --git a/go/core/internal/database/client_test.go b/go/core/internal/database/client_test.go index 16f10bf90..e4726e8d8 100644 --- a/go/core/internal/database/client_test.go +++ b/go/core/internal/database/client_test.go @@ -2,6 +2,7 @@ package database import ( "context" + "encoding/json" "fmt" "sync" "testing" @@ -11,6 +12,7 @@ import ( "github.com/jackc/pgx/v5/pgxpool" dbpkg "github.com/kagent-dev/kagent/go/api/database" "github.com/kagent-dev/kagent/go/api/v1alpha2" + "github.com/kagent-dev/kagent/go/core/pkg/a2acompat/trpcv0" "github.com/pgvector/pgvector-go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -974,3 +976,133 @@ func TestSearchAgentMemoryConcurrentAccessCount(t *testing.T) { require.NoError(t, err, "concurrent memory search must not fail") } } + +// TestListUserTasks exercises the ListUserTasks SQL query against a real +// database: cross-user scoping via the session join, the optional single-session +// predicate, status filtering across both persisted row shapes (legacy rows +// written by StoreTask, and a v1 row inserted directly), the timestamp filter, +// and LIMIT/OFFSET pagination with a stable COUNT(*) OVER() total. +func TestListUserTasks(t *testing.T) { + db := setupTestDB(t) + client := NewClient(db) + ctx := context.Background() + + for _, s := range []struct{ id, user string }{ + {"s1", "alice"}, {"s2", "alice"}, {"s3", "bob"}, + } { + require.NoError(t, client.StoreSession(ctx, &dbpkg.Session{ID: s.id, UserID: s.user})) + } + + early := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC) + late := time.Date(2026, 7, 8, 0, 0, 0, 0, time.UTC) + mkTask := func(id, contextID string, state a2a.TaskState, ts time.Time) *a2a.Task { + return &a2a.Task{ + ID: a2a.TaskID(id), + ContextID: contextID, + Status: a2a.TaskStatus{State: state, Timestamp: &ts}, + } + } + + // StoreTask persists legacy-format rows (protocol_version NULL, lowercase state). + require.NoError(t, client.StoreTask(ctx, mkTask("t1", "s1", a2a.TaskStateWorking, early), "alice")) + require.NoError(t, client.StoreTask(ctx, mkTask("t2", "s1", a2a.TaskStateCompleted, late), "alice")) + require.NoError(t, client.StoreTask(ctx, mkTask("t3", "s2", a2a.TaskStateWorking, late), "alice")) + require.NoError(t, client.StoreTask(ctx, mkTask("t4", "s3", a2a.TaskStateWorking, late), "bob")) + + // A v1-format row (protocol_version 'v1', uppercase TASK_STATE_* state) to + // prove the JSON cast matches both row shapes. It carries no user_id, so it + // also covers the NULL-owner fallback: s1 resolves to alice alone. + v1Data, err := json.Marshal(mkTask("t5", "s1", a2a.TaskStateInputRequired, late)) + require.NoError(t, err) + _, err = db.Exec(ctx, + `INSERT INTO task (id, data, session_id, protocol_version, created_at, updated_at) + VALUES ($1, $2, $3, $4, NOW(), NOW())`, + "t5", string(v1Data), "s1", trpcv0.ProtocolVersionV1) + require.NoError(t, err) + + // A task owned by bob but parked in alice's session. The session join alone + // would hand it to alice; the task-owner predicate keeps it out. + foreignData, err := json.Marshal(mkTask("t6", "s1", a2a.TaskStateWorking, late)) + require.NoError(t, err) + _, err = db.Exec(ctx, + `INSERT INTO task (id, data, session_id, user_id, created_at, updated_at) + VALUES ($1, $2, $3, $4, NOW(), NOW())`, + "t6", string(foreignData), "s1", "bob") + require.NoError(t, err) + + ids := func(tasks []*a2a.Task) []string { + out := make([]string, len(tasks)) + for i, tk := range tasks { + out[i] = string(tk.ID) + } + return out + } + + t.Run("all sessions ordered by id", func(t *testing.T) { + tasks, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", Limit: 50}) + require.NoError(t, err) + require.Equal(t, 4, total) + require.Equal(t, []string{"t1", "t2", "t3", "t5"}, ids(tasks)) + }) + + t.Run("cross-user isolation", func(t *testing.T) { + tasks, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "bob", Limit: 50}) + require.NoError(t, err) + require.Equal(t, 1, total) + require.Equal(t, []string{"t4"}, ids(tasks), "t6 sits in a session bob does not own") + }) + + t.Run("foreign task in own session stays hidden", func(t *testing.T) { + tasks, _, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", SessionID: "s1", Limit: 50}) + require.NoError(t, err) + require.NotContains(t, ids(tasks), "t6") + }) + + t.Run("single session predicate", func(t *testing.T) { + tasks, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", SessionID: "s1", Limit: 50}) + require.NoError(t, err) + require.Equal(t, 3, total) + require.Equal(t, []string{"t1", "t2", "t5"}, ids(tasks)) + }) + + t.Run("status filter matches legacy rows", func(t *testing.T) { + tasks, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", Status: a2a.TaskStateWorking, Limit: 50}) + require.NoError(t, err) + require.Equal(t, 2, total) + require.Equal(t, []string{"t1", "t3"}, ids(tasks)) + }) + + t.Run("status filter matches v1 rows", func(t *testing.T) { + tasks, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", Status: a2a.TaskStateInputRequired, Limit: 50}) + require.NoError(t, err) + require.Equal(t, 1, total) + require.Equal(t, []string{"t5"}, ids(tasks)) + }) + + t.Run("status timestamp after", func(t *testing.T) { + cutoff := time.Date(2026, 7, 5, 0, 0, 0, 0, time.UTC) + tasks, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", StatusTimestampAfter: &cutoff, Limit: 50}) + require.NoError(t, err) + require.Equal(t, 3, total) // t1 is early and excluded + require.Equal(t, []string{"t2", "t3", "t5"}, ids(tasks)) + }) + + t.Run("pagination with stable total", func(t *testing.T) { + p1, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", Limit: 2, Offset: 0}) + require.NoError(t, err) + require.Equal(t, 4, total) + require.Equal(t, []string{"t1", "t2"}, ids(p1)) + + p2, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", Limit: 2, Offset: 2}) + require.NoError(t, err) + require.Equal(t, 4, total) + require.Equal(t, []string{"t3", "t5"}, ids(p2)) + }) + + t.Run("offset past end keeps the true total", func(t *testing.T) { + tasks, total, err := client.ListUserTasks(ctx, dbpkg.ListUserTasksParams{UserID: "alice", Limit: 2, Offset: 10}) + require.NoError(t, err) + require.Empty(t, tasks) + require.Equal(t, 4, total, "an empty over-range page must still report the full count") + }) +} diff --git a/go/core/internal/database/gen/querier.go b/go/core/internal/database/gen/querier.go index f897decfe..febfbaed8 100644 --- a/go/core/internal/database/gen/querier.go +++ b/go/core/internal/database/gen/querier.go @@ -9,6 +9,11 @@ import ( ) type Querier interface { + // The full filtered count for ListUserTasks, independent of LIMIT/OFFSET. Used + // to recover total when a requested page lands past the end of the set (an empty + // page carries no COUNT(*) OVER()). The WHERE clause is identical to + // ListUserTasks and must stay in sync with it. + CountUserTasks(ctx context.Context, arg CountUserTasksParams) (int64, error) CreateSessionShare(ctx context.Context, arg CreateSessionShareParams) (SessionShare, error) DeleteAgentMemory(ctx context.Context, arg DeleteAgentMemoryParams) error DeleteExpiredMemories(ctx context.Context) error @@ -65,6 +70,27 @@ type Querier interface { ListToolServers(ctx context.Context) ([]Toolserver, error) ListTools(ctx context.Context) ([]Tool, error) ListToolsForServer(ctx context.Context, arg ListToolsForServerParams) ([]Tool, error) + // Lists a user's tasks across every session they own (or a single session when + // session_id is set), filtering, ordering, and paginating server-side. total is + // the COUNT(*) OVER() of the full filtered set, before LIMIT/OFFSET. + // + // A row must clear both gates: the session belongs to the caller, and the task + // itself does under the ownership rule at the top of this file. The task gate is + // what keeps a foreign task parked in the caller's session out of the result. + // + // status_state matches task.data's persisted state string. Rows exist in two + // shapes: v1 (protocol_version = 'v1', state e.g. 'TASK_STATE_WORKING') and + // legacy (protocol_version NULL, state e.g. 'working'); the caller passes both + // spellings in status_v1/status_legacy so either row shape matches. The two + // vocabularies never overlap, so matching against both is unambiguous. + // data is always a JSON object (json.Marshal output), so ::jsonb never errors; + // the timestamp cast is guarded by a CASE (ordered evaluation) against a + // present-but-malformed value. + // + // COUNT(*) OVER() rides on the returned rows, so a page past the end of the set + // carries no total; callers recover it with CountUserTasks, whose WHERE clause + // must stay identical to this one. + ListUserTasks(ctx context.Context, arg ListUserTasksParams) ([]ListUserTasksRow, error) // Memory uses hard DELETE (not soft deletes), so no deleted_at filter is needed. // COALESCE guards against NULL embeddings (score=0 rather than NULL); rows are still ordered last by the ORDER BY clause. SearchAgentMemory(ctx context.Context, arg SearchAgentMemoryParams) ([]SearchAgentMemoryRow, error) diff --git a/go/core/internal/database/gen/tasks.sql.go b/go/core/internal/database/gen/tasks.sql.go index bde040957..f16e0c467 100644 --- a/go/core/internal/database/gen/tasks.sql.go +++ b/go/core/internal/database/gen/tasks.sql.go @@ -7,8 +7,61 @@ package dbgen import ( "context" + "time" ) +const countUserTasks = `-- name: CountUserTasks :one +SELECT COUNT(*) +FROM task +JOIN session ON session.id = task.session_id +WHERE session.user_id = $1 + AND task.deleted_at IS NULL + AND session.deleted_at IS NULL + AND (task.user_id = $1 OR (task.user_id IS NULL AND $1 = ( + SELECT MIN(s.user_id) FROM session s + WHERE s.id = task.session_id AND s.created_at <= task.created_at + HAVING COUNT(DISTINCT s.user_id) = 1))) + AND ($2::text IS NULL OR task.session_id = $2) + AND ( + $3::text IS NULL + OR (task.data::jsonb -> 'status' ->> 'state') IN ($3, $4) + ) + AND ( + $5::timestamptz IS NULL + OR ( + CASE + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz + END + ) > $5 + ) +` + +type CountUserTasksParams struct { + UserID string + SessionID *string + StatusV1 *string + StatusLegacy *string + StatusAfter *time.Time +} + +// The full filtered count for ListUserTasks, independent of LIMIT/OFFSET. Used +// to recover total when a requested page lands past the end of the set (an empty +// page carries no COUNT(*) OVER()). The WHERE clause is identical to +// ListUserTasks and must stay in sync with it. +func (q *Queries) CountUserTasks(ctx context.Context, arg CountUserTasksParams) (int64, error) { + row := q.db.QueryRow(ctx, countUserTasks, + arg.UserID, + arg.SessionID, + arg.StatusV1, + arg.StatusLegacy, + arg.StatusAfter, + ) + var count int64 + err := row.Scan(&count) + return count, err +} + const getTask = `-- name: GetTask :one SELECT id, created_at, updated_at, deleted_at, data, session_id, protocol_version, user_id FROM task @@ -109,6 +162,115 @@ func (q *Queries) ListTasksForSession(ctx context.Context, arg ListTasksForSessi return items, nil } +const listUserTasks = `-- name: ListUserTasks :many +SELECT task.id, task.created_at, task.updated_at, task.deleted_at, task.data, task.session_id, task.protocol_version, task.user_id, COUNT(*) OVER() AS total +FROM task +JOIN session ON session.id = task.session_id +WHERE session.user_id = $1 + AND task.deleted_at IS NULL + AND session.deleted_at IS NULL + AND (task.user_id = $1 OR (task.user_id IS NULL AND $1 = ( + SELECT MIN(s.user_id) FROM session s + WHERE s.id = task.session_id AND s.created_at <= task.created_at + HAVING COUNT(DISTINCT s.user_id) = 1))) + AND ($2::text IS NULL OR task.session_id = $2) + AND ( + $3::text IS NULL + OR (task.data::jsonb -> 'status' ->> 'state') IN ($3, $4) + ) + AND ( + $5::timestamptz IS NULL + OR ( + CASE + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz + END + ) > $5 + ) +ORDER BY task.id +LIMIT $7::int OFFSET $6::int +` + +type ListUserTasksParams struct { + UserID string + SessionID *string + StatusV1 *string + StatusLegacy *string + StatusAfter *time.Time + PageOffset int32 + PageLimit int32 +} + +type ListUserTasksRow struct { + ID string + CreatedAt *time.Time + UpdatedAt *time.Time + DeletedAt *time.Time + Data string + SessionID *string + ProtocolVersion *string + UserID *string + Total int64 +} + +// Lists a user's tasks across every session they own (or a single session when +// session_id is set), filtering, ordering, and paginating server-side. total is +// the COUNT(*) OVER() of the full filtered set, before LIMIT/OFFSET. +// +// A row must clear both gates: the session belongs to the caller, and the task +// itself does under the ownership rule at the top of this file. The task gate is +// what keeps a foreign task parked in the caller's session out of the result. +// +// status_state matches task.data's persisted state string. Rows exist in two +// shapes: v1 (protocol_version = 'v1', state e.g. 'TASK_STATE_WORKING') and +// legacy (protocol_version NULL, state e.g. 'working'); the caller passes both +// spellings in status_v1/status_legacy so either row shape matches. The two +// vocabularies never overlap, so matching against both is unambiguous. +// data is always a JSON object (json.Marshal output), so ::jsonb never errors; +// the timestamp cast is guarded by a CASE (ordered evaluation) against a +// present-but-malformed value. +// +// COUNT(*) OVER() rides on the returned rows, so a page past the end of the set +// carries no total; callers recover it with CountUserTasks, whose WHERE clause +// must stay identical to this one. +func (q *Queries) ListUserTasks(ctx context.Context, arg ListUserTasksParams) ([]ListUserTasksRow, error) { + rows, err := q.db.Query(ctx, listUserTasks, + arg.UserID, + arg.SessionID, + arg.StatusV1, + arg.StatusLegacy, + arg.StatusAfter, + arg.PageOffset, + arg.PageLimit, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ListUserTasksRow + for rows.Next() { + var i ListUserTasksRow + if err := rows.Scan( + &i.ID, + &i.CreatedAt, + &i.UpdatedAt, + &i.DeletedAt, + &i.Data, + &i.SessionID, + &i.ProtocolVersion, + &i.UserID, + &i.Total, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const softDeleteTask = `-- name: SoftDeleteTask :exec UPDATE task SET deleted_at = NOW() WHERE task.id = $1 AND task.deleted_at IS NULL diff --git a/go/core/internal/database/queries/tasks.sql b/go/core/internal/database/queries/tasks.sql index 14e20793f..469d92ef1 100644 --- a/go/core/internal/database/queries/tasks.sql +++ b/go/core/internal/database/queries/tasks.sql @@ -38,6 +38,84 @@ WHERE task.session_id = $1 AND task.deleted_at IS NULL HAVING COUNT(DISTINCT s.user_id) = 1))) ORDER BY created_at ASC; +-- name: ListUserTasks :many +-- Lists a user's tasks across every session they own (or a single session when +-- session_id is set), filtering, ordering, and paginating server-side. total is +-- the COUNT(*) OVER() of the full filtered set, before LIMIT/OFFSET. +-- +-- A row must clear both gates: the session belongs to the caller, and the task +-- itself does under the ownership rule at the top of this file. The task gate is +-- what keeps a foreign task parked in the caller's session out of the result. +-- +-- status_state matches task.data's persisted state string. Rows exist in two +-- shapes: v1 (protocol_version = 'v1', state e.g. 'TASK_STATE_WORKING') and +-- legacy (protocol_version NULL, state e.g. 'working'); the caller passes both +-- spellings in status_v1/status_legacy so either row shape matches. The two +-- vocabularies never overlap, so matching against both is unambiguous. +-- data is always a JSON object (json.Marshal output), so ::jsonb never errors; +-- the timestamp cast is guarded by a CASE (ordered evaluation) against a +-- present-but-malformed value. +-- +-- COUNT(*) OVER() rides on the returned rows, so a page past the end of the set +-- carries no total; callers recover it with CountUserTasks, whose WHERE clause +-- must stay identical to this one. +SELECT task.*, COUNT(*) OVER() AS total +FROM task +JOIN session ON session.id = task.session_id +WHERE session.user_id = @user_id + AND task.deleted_at IS NULL + AND session.deleted_at IS NULL + AND (task.user_id = @user_id OR (task.user_id IS NULL AND @user_id = ( + SELECT MIN(s.user_id) FROM session s + WHERE s.id = task.session_id AND s.created_at <= task.created_at + HAVING COUNT(DISTINCT s.user_id) = 1))) + AND (sqlc.narg('session_id')::text IS NULL OR task.session_id = sqlc.narg('session_id')) + AND ( + sqlc.narg('status_v1')::text IS NULL + OR (task.data::jsonb -> 'status' ->> 'state') IN (sqlc.narg('status_v1'), sqlc.narg('status_legacy')) + ) + AND ( + sqlc.narg('status_after')::timestamptz IS NULL + OR ( + CASE + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz + END + ) > sqlc.narg('status_after') + ) +ORDER BY task.id +LIMIT @page_limit::int OFFSET @page_offset::int; + +-- name: CountUserTasks :one +-- The full filtered count for ListUserTasks, independent of LIMIT/OFFSET. Used +-- to recover total when a requested page lands past the end of the set (an empty +-- page carries no COUNT(*) OVER()). The WHERE clause is identical to +-- ListUserTasks and must stay in sync with it. +SELECT COUNT(*) +FROM task +JOIN session ON session.id = task.session_id +WHERE session.user_id = @user_id + AND task.deleted_at IS NULL + AND session.deleted_at IS NULL + AND (task.user_id = @user_id OR (task.user_id IS NULL AND @user_id = ( + SELECT MIN(s.user_id) FROM session s + WHERE s.id = task.session_id AND s.created_at <= task.created_at + HAVING COUNT(DISTINCT s.user_id) = 1))) + AND (sqlc.narg('session_id')::text IS NULL OR task.session_id = sqlc.narg('session_id')) + AND ( + sqlc.narg('status_v1')::text IS NULL + OR (task.data::jsonb -> 'status' ->> 'state') IN (sqlc.narg('status_v1'), sqlc.narg('status_legacy')) + ) + AND ( + sqlc.narg('status_after')::timestamptz IS NULL + OR ( + CASE + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz + END + ) > sqlc.narg('status_after') + ); + -- UpsertTask returns the upserted id, or no rows when the write was rejected: -- the id belongs to another user, or it belongs to a soft-deleted task (a -- deleted id is never updated or resurrected, it stays burned). Callers map diff --git a/go/core/pkg/a2acompat/trpcv0/convert.go b/go/core/pkg/a2acompat/trpcv0/convert.go index e6a2384c7..a30c7f7b8 100644 --- a/go/core/pkg/a2acompat/trpcv0/convert.go +++ b/go/core/pkg/a2acompat/trpcv0/convert.go @@ -277,6 +277,13 @@ func toLegacyPart(part *a2av1.Part) (trpc.Part, error) { }, nil } +// LegacyTaskStateString returns the persisted state string a legacy (pre-v1) +// task row carries for the given v1 state, e.g. "working" for TaskStateWorking. +// A v1 row persists string(state) instead (e.g. "TASK_STATE_WORKING"). +func LegacyTaskStateString(state a2av1.TaskState) string { + return string(toLegacyTaskState(state)) +} + func toLegacyTaskState(state a2av1.TaskState) trpc.TaskState { switch state { case a2av1.TaskStateSubmitted: From 05e087d4908b5bf56b85328b184522ac2fa4811b Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Wed, 22 Jul 2026 22:30:58 +0200 Subject: [PATCH 2/2] fix(a2a): tighten task timestamp guard regex and name bad row in error The timestamp CASE guard only matched a YYYY-MM-DD prefix, so a value like "2026-07-01not-a-time" passed the guard and errored the ::timestamptz cast. Match a full RFC3339 timestamp (time + Z/offset) in both ListUserTasks and CountUserTasks so the guard actually protects the cast. Report the task id in the ListUserTasks parse error instead of the row index so operators can locate the bad row. Signed-off-by: QuentinBisson --- go/core/internal/database/client_postgres.go | 4 ++-- go/core/internal/database/gen/tasks.sql.go | 4 ++-- go/core/internal/database/queries/tasks.sql | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 9f3462f3e..554b4fbc8 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -366,11 +366,11 @@ func (c *postgresClient) ListUserTasks(ctx context.Context, params dbpkg.ListUse total := 0 tasks := make([]*a2a.Task, 0, len(rows)) - for i, r := range rows { + for _, r := range rows { total = int(r.Total) task, err := parseVersionedTask(r.Data, r.ProtocolVersion) if err != nil { - return nil, 0, fmt.Errorf("failed to parse task row %d: %w", i, err) + return nil, 0, fmt.Errorf("failed to parse task %s: %w", r.ID, err) } tasks = append(tasks, task) } diff --git a/go/core/internal/database/gen/tasks.sql.go b/go/core/internal/database/gen/tasks.sql.go index f16e0c467..67cd95394 100644 --- a/go/core/internal/database/gen/tasks.sql.go +++ b/go/core/internal/database/gen/tasks.sql.go @@ -30,7 +30,7 @@ WHERE session.user_id = $1 $5::timestamptz IS NULL OR ( CASE - WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:?[0-9]{2})$' THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz END ) > $5 @@ -182,7 +182,7 @@ WHERE session.user_id = $1 $5::timestamptz IS NULL OR ( CASE - WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:?[0-9]{2})$' THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz END ) > $5 diff --git a/go/core/internal/database/queries/tasks.sql b/go/core/internal/database/queries/tasks.sql index 469d92ef1..e6ebfd1a4 100644 --- a/go/core/internal/database/queries/tasks.sql +++ b/go/core/internal/database/queries/tasks.sql @@ -78,7 +78,7 @@ WHERE session.user_id = @user_id sqlc.narg('status_after')::timestamptz IS NULL OR ( CASE - WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:?[0-9]{2})$' THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz END ) > sqlc.narg('status_after') @@ -110,7 +110,7 @@ WHERE session.user_id = @user_id sqlc.narg('status_after')::timestamptz IS NULL OR ( CASE - WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}' + WHEN (task.data::jsonb -> 'status' ->> 'timestamp') ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:?[0-9]{2})$' THEN (task.data::jsonb -> 'status' ->> 'timestamp')::timestamptz END ) > sqlc.narg('status_after')