perf(a2a): push ListTasks all-sessions query down to SQL

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)

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).

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) {

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 = &params.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 _, 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 %s: %w", r.ID, 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)