diff --git a/go/api/database/client.go b/go/api/database/client.go index e80dc8647..67614ebda 100644 --- a/go/api/database/client.go +++ b/go/api/database/client.go @@ -38,6 +38,9 @@ type Client interface { DeleteToolsForServer(ctx context.Context, serverName string, groupKind string) error // Get methods + + // GetSession returns an error wrapping ErrNotFound when no session with + // that id exists for the user. GetSession(ctx context.Context, sessionID string, userID string) (*Session, error) GetAgent(ctx context.Context, name string) (*Agent, error) GetTask(ctx context.Context, id string) (*a2a.Task, error) diff --git a/go/api/database/errors.go b/go/api/database/errors.go new file mode 100644 index 000000000..aee9734c8 --- /dev/null +++ b/go/api/database/errors.go @@ -0,0 +1,8 @@ +package database + +import "errors" + +// ErrNotFound reports that the requested record does not exist (or is not +// visible to the given user). Match with errors.Is; implementations wrap it +// with call-site context. +var ErrNotFound = errors.New("record not found") diff --git a/go/core/internal/a2a/a2a_handler_mux.go b/go/core/internal/a2a/a2a_handler_mux.go index 1bd767e41..26dffe949 100644 --- a/go/core/internal/a2a/a2a_handler_mux.go +++ b/go/core/internal/a2a/a2a_handler_mux.go @@ -37,6 +37,7 @@ type handlerMux struct { agentPathPrefix string sandboxPathPrefix string authenticator auth.AuthProvider + taskStore TaskStore } var _ A2AHandlerMux = &handlerMux{} @@ -45,15 +46,33 @@ type middleware interface { Wrap(next http.Handler) http.Handler } -func NewA2AHttpMux(agentPathPrefix, sandboxPathPrefix string, authenticator auth.AuthProvider) *handlerMux { +func NewA2AHttpMux(agentPathPrefix, sandboxPathPrefix string, authenticator auth.AuthProvider, taskStore TaskStore) *handlerMux { return &handlerMux{ handlers: make(map[string]http.Handler), agentPathPrefix: agentPathPrefix, sandboxPathPrefix: sandboxPathPrefix, authenticator: authenticator, + taskStore: taskStore, } } +// newTaskQueryHandlers builds the request handler and the legacy (v0) JSON-RPC +// handler for one agent. kagent persists tasks and is their source of truth, +// so with a store ListTasks is served from it instead of proxying to the agent +// runtime (whose legacy 0.3 transport returns ErrUnsupportedOperation for it); +// every other method, GetTask included, still delegates to the passthrough +// proxy. v0 has no native tasks/list, so the legacy handler is wrapped to +// serve that method from the store too (lowercase TaskState). Without a store +// both wires keep their native behavior, including v0's method-not-found for +// tasks/list. +func newTaskQueryHandlers(requestHandler a2asrv.RequestHandler, store TaskStore) (a2asrv.RequestHandler, http.Handler) { + if store == nil { + return requestHandler, a2av0.NewJSONRPCHandler(requestHandler) + } + taskHandler := newStoreTaskQueryHandler(requestHandler, store) + return taskHandler, newV0TasksListInterceptor(a2av0.NewJSONRPCHandler(taskHandler), taskHandler) +} + func (a *handlerMux) SetAgentHandler( agentRef string, client *a2aclient.Client, @@ -61,8 +80,9 @@ func (a *handlerMux) SetAgentHandler( tracing middleware, ) error { requestHandler := NewPassthroughRequestHandler(client, &card) - legacyJSONRPCHandler := a2av0.NewJSONRPCHandler(requestHandler) - v1JSONRPCHandler := a2asrv.NewJSONRPCHandler(requestHandler) + + taskHandler, legacyJSONRPCHandler := newTaskQueryHandlers(requestHandler, a.taskStore) + v1JSONRPCHandler := a2asrv.NewJSONRPCHandler(taskHandler) cardHandler := a2asrv.NewAgentCardHandler(a2av0.NewStaticAgentCardProducer(&card)) wellKnownPath := "/" + strings.TrimPrefix(a2asrv.WellKnownAgentCardPath, "/") diff --git a/go/core/internal/a2a/a2av0_tasks_list.go b/go/core/internal/a2a/a2av0_tasks_list.go new file mode 100644 index 000000000..624b9ca2c --- /dev/null +++ b/go/core/internal/a2a/a2av0_tasks_list.go @@ -0,0 +1,164 @@ +package a2a + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "net/http" + "time" + + a2alegacy "github.com/a2aproject/a2a-go/a2a" + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2acompat/a2av0" + "github.com/a2aproject/a2a-go/v2/a2asrv" +) + +// v0MethodTasksList is a kagent compatibility extension: the A2A 0.3 wire +// predates the task-query methods, so a2av0.NewJSONRPCHandler returns +// method-not-found for it. We intercept it here and serve it from the task +// store with the legacy lowercase TaskState encoding. tasks/get is left to the +// underlying v0 handler, which already routes it to the store-backed GetTask. +const v0MethodTasksList = "tasks/list" + +// JSON-RPC 2.0 error codes mirrored from a2a-go's internal mapping. +const ( + codeTaskNotFound = -32001 + codeInvalidParams = -32602 + codeInternalError = -32603 +) + +type v0RPCRequest struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` + ID json.RawMessage `json:"id,omitempty"` +} + +type v0RPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id,omitempty"` + Result any `json:"result,omitempty"` + Error *v0RPCError `json:"error,omitempty"` +} + +type v0RPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +// v0ListTasksParams mirrors the v1 ListTasks params but carries the legacy +// lowercase TaskState string for the status filter. +type v0ListTasksParams struct { + ContextID string `json:"contextId,omitempty"` + Status string `json:"status,omitempty"` + PageSize int `json:"pageSize,omitempty"` + PageToken string `json:"pageToken,omitempty"` + HistoryLength *int `json:"historyLength,omitempty"` + StatusTimestampAfter *time.Time `json:"statusTimestampAfter,omitempty"` + IncludeArtifacts bool `json:"includeArtifacts,omitempty"` +} + +type v0ListTasksResult struct { + Tasks []*a2alegacy.Task `json:"tasks"` + PageSize int `json:"pageSize"` + TotalSize int `json:"totalSize"` + NextPageToken string `json:"nextPageToken"` +} + +// v0TasksListInterceptor serves tasks/list from handler and delegates every +// other request to next unchanged. +type v0TasksListInterceptor struct { + next http.Handler + handler a2asrv.RequestHandler +} + +func newV0TasksListInterceptor(next http.Handler, handler a2asrv.RequestHandler) http.Handler { + return &v0TasksListInterceptor{next: next, handler: handler} +} + +func (h *v0TasksListInterceptor) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + h.next.ServeHTTP(w, r) + return + } + + body, err := io.ReadAll(r.Body) + if err != nil { + h.next.ServeHTTP(w, r) + return + } + // Restore the body so the delegate can read it when this is not tasks/list. + r.Body = io.NopCloser(bytes.NewReader(body)) + + var rpcReq v0RPCRequest + if err := json.Unmarshal(body, &rpcReq); err != nil || rpcReq.Method != v0MethodTasksList { + h.next.ServeHTTP(w, r) + return + } + + h.serveTasksList(w, r, &rpcReq) +} + +func (h *v0TasksListInterceptor) serveTasksList(w http.ResponseWriter, r *http.Request, rpcReq *v0RPCRequest) { + w.Header().Set("Content-Type", "application/json") + + var params v0ListTasksParams + if len(rpcReq.Params) > 0 { + if err := json.Unmarshal(rpcReq.Params, ¶ms); err != nil { + writeV0Error(w, rpcReq.ID, codeInvalidParams, "invalid params") + return + } + } + + listReq := &a2atype.ListTasksRequest{ + ContextID: params.ContextID, + PageSize: params.PageSize, + PageToken: params.PageToken, + HistoryLength: params.HistoryLength, + StatusTimestampAfter: params.StatusTimestampAfter, + IncludeArtifacts: params.IncludeArtifacts, + } + if params.Status != "" { + listReq.Status = a2av0.ToV1TaskState(a2alegacy.TaskState(params.Status)) + } + + resp, err := h.handler.ListTasks(r.Context(), listReq) + if err != nil { + writeV0Error(w, rpcReq.ID, v0ErrorCode(err), err.Error()) + return + } + + result := v0ListTasksResult{ + Tasks: make([]*a2alegacy.Task, 0, len(resp.Tasks)), + PageSize: resp.PageSize, + TotalSize: resp.TotalSize, + NextPageToken: resp.NextPageToken, + } + for _, t := range resp.Tasks { + result.Tasks = append(result.Tasks, a2av0.FromV1Task(t)) + } + + writeV0Response(w, v0RPCResponse{JSONRPC: "2.0", ID: rpcReq.ID, Result: result}) +} + +func v0ErrorCode(err error) int { + switch { + case errors.Is(err, a2atype.ErrInvalidParams): + return codeInvalidParams + case errors.Is(err, a2atype.ErrTaskNotFound): + return codeTaskNotFound + default: + return codeInternalError + } +} + +func writeV0Error(w http.ResponseWriter, id json.RawMessage, code int, message string) { + writeV0Response(w, v0RPCResponse{JSONRPC: "2.0", ID: id, Error: &v0RPCError{Code: code, Message: message}}) +} + +func writeV0Response(w http.ResponseWriter, resp v0RPCResponse) { + if err := json.NewEncoder(w).Encode(resp); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } +} diff --git a/go/core/internal/a2a/task_query_store.go b/go/core/internal/a2a/task_query_store.go new file mode 100644 index 000000000..6285932d2 --- /dev/null +++ b/go/core/internal/a2a/task_query_store.go @@ -0,0 +1,227 @@ +package a2a + +import ( + "cmp" + "context" + "encoding/base64" + "errors" + "fmt" + "slices" + "strconv" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2asrv" + dbpkg "github.com/kagent-dev/kagent/go/api/database" + "github.com/kagent-dev/kagent/go/core/pkg/auth" +) + +const ( + defaultTaskPageSize = 50 + maxTaskPageSize = 100 +) + +// TaskStore is the subset of the persistent store ListTasks reads from. +// *database.Client satisfies it. GetSession reports a missing (or +// other-user) session by wrapping database.ErrNotFound; any other error is +// a backend failure and is surfaced to the caller. +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 string) ([]*a2atype.Task, error) +} + +// storeTaskQueryHandler answers ListTasks from kagent's task store, which is +// the source of truth for persisted tasks. Every other method (including +// GetTask, which already resolves to the same store via the passthrough) is +// delegated to the embedded handler unchanged. +type storeTaskQueryHandler struct { + a2asrv.RequestHandler + store TaskStore +} + +func newStoreTaskQueryHandler(delegate a2asrv.RequestHandler, store TaskStore) *storeTaskQueryHandler { + return &storeTaskQueryHandler{RequestHandler: delegate, store: store} +} + +// callerUserID returns the authenticated principal's user id, or "" when the +// request carries no user identity. Task queries are scoped to this id so a +// caller can never read another user's tasks. +func callerUserID(ctx context.Context) string { + session, ok := auth.AuthSessionFrom(ctx) + if !ok || session == nil { + return "" + } + return session.Principal().User.ID +} + +// effectiveUserIDForContext returns the user id tasks are read as. A valid +// share token grants access to a single session, so when the request carries a +// ShareContext for exactly the requested contextId the owner's id is used +// (mirroring getEffectiveUserIDForSession in the session handlers); otherwise +// the caller's own id is used. A share token never widens an all-sessions +// query, so contextId must be set for it to apply. +func effectiveUserIDForContext(ctx context.Context, contextID string) string { + if contextID != "" { + if sc, ok := auth.ShareContextFrom(ctx); ok && sc.SessionID == contextID { + return sc.UserID + } + } + return callerUserID(ctx) +} + +func (h *storeTaskQueryHandler) ListTasks(ctx context.Context, req *a2atype.ListTasksRequest) (*a2atype.ListTasksResponse, error) { + pageSize := clampPageSize(req.PageSize) + + userID := effectiveUserIDForContext(ctx, req.ContextID) + if userID == "" { + return &a2atype.ListTasksResponse{Tasks: []*a2atype.Task{}, PageSize: pageSize}, nil + } + + tasks, err := h.collectUserTasks(ctx, userID, req.ContextID) + if err != nil { + return nil, err + } + + 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) }) + + offset, err := decodePageToken(req.PageToken) + if err != nil { + return nil, a2atype.NewError(a2atype.ErrInvalidParams, "invalid pageToken") + } + total := len(filtered) + if offset > total { + offset = total + } + end := min(offset+pageSize, total) + + page := filtered[offset:end] + shaped := make([]*a2atype.Task, 0, len(page)) + for _, t := range page { + shaped = append(shaped, shapeTask(t, req.HistoryLength, req.IncludeArtifacts)) + } + + nextToken := "" + if end < total { + nextToken = encodePageToken(end) + } + + return &a2atype.ListTasksResponse{ + Tasks: shaped, + PageSize: pageSize, + TotalSize: total, + NextPageToken: nextToken, + }, 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 { + // Session does not exist or is not the caller's: no tasks to return. + if errors.Is(err, dbpkg.ErrNotFound) { + return nil, nil + } + return nil, fmt.Errorf("get session %s: %w", contextID, err) + } + return h.store.ListTasksForSession(ctx, contextID) + } + + 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) + 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). +func shapeTask(task *a2atype.Task, historyLength *int, includeArtifacts bool) *a2atype.Task { + shaped := *task + shaped.History = truncateHistory(task.History, historyLength) + if includeArtifacts { + shaped.Artifacts = task.Artifacts + } else { + shaped.Artifacts = nil + } + return &shaped +} + +// truncateHistory keeps the most recent n messages. A nil historyLength keeps +// the full history; n <= 0 drops it entirely. +func truncateHistory(history []*a2atype.Message, historyLength *int) []*a2atype.Message { + if historyLength == nil { + return history + } + n := *historyLength + if n <= 0 { + return nil + } + if n >= len(history) { + return history + } + return history[len(history)-n:] +} + +func clampPageSize(size int) int { + switch { + case size <= 0: + return defaultTaskPageSize + case size > maxTaskPageSize: + return maxTaskPageSize + default: + return size + } +} + +func encodePageToken(offset int) string { + return base64.RawURLEncoding.EncodeToString([]byte(strconv.Itoa(offset))) +} + +func decodePageToken(token string) (int, error) { + if token == "" { + return 0, nil + } + raw, err := base64.RawURLEncoding.DecodeString(token) + if err != nil { + return 0, err + } + offset, err := strconv.Atoi(string(raw)) + if err != nil || offset < 0 { + return 0, fmt.Errorf("invalid page token") + } + return offset, nil +} diff --git a/go/core/internal/a2a/task_query_store_test.go b/go/core/internal/a2a/task_query_store_test.go new file mode 100644 index 000000000..7128d9df6 --- /dev/null +++ b/go/core/internal/a2a/task_query_store_test.go @@ -0,0 +1,370 @@ +package a2a + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + a2atype "github.com/a2aproject/a2a-go/v2/a2a" + "github.com/a2aproject/a2a-go/v2/a2acompat/a2av0" + "github.com/a2aproject/a2a-go/v2/a2asrv" + dbpkg "github.com/kagent-dev/kagent/go/api/database" + "github.com/kagent-dev/kagent/go/core/pkg/auth" + "github.com/stretchr/testify/require" +) + +// fakeTaskStore is an in-memory TaskStore. Sessions are keyed by (id, userID) +// so cross-user isolation is exercised the same way the real store enforces it. +type fakeTaskStore struct { + sessions map[string]dbpkg.Session // key: sessionID -> session (carries UserID) + tasks map[string][]*a2atype.Task +} + +func newFakeStore() *fakeTaskStore { + return &fakeTaskStore{ + sessions: map[string]dbpkg.Session{}, + tasks: map[string][]*a2atype.Task{}, + } +} + +func (f *fakeTaskStore) addSession(id, userID string) { + f.sessions[id] = dbpkg.Session{ID: id, UserID: userID} +} + +func (f *fakeTaskStore) addTask(sessionID string, task *a2atype.Task) { + f.tasks[sessionID] = append(f.tasks[sessionID], task) +} + +func (f *fakeTaskStore) GetSession(_ context.Context, sessionID, userID string) (*dbpkg.Session, error) { + s, ok := f.sessions[sessionID] + if !ok || s.UserID != userID { + return nil, fmt.Errorf("session %s for user %s: %w", sessionID, userID, dbpkg.ErrNotFound) + } + 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) + } + } + return out, nil +} + +func (f *fakeTaskStore) ListTasksForSession(_ context.Context, sessionID string) ([]*a2atype.Task, error) { + return f.tasks[sessionID], nil +} + +// fakeSession injects a user principal into the request context. +type fakeSession struct{ user string } + +func (f fakeSession) Principal() auth.Principal { + return auth.Principal{User: auth.User{ID: f.user}} +} + +func userCtx(user string) context.Context { + return auth.AuthSessionTo(context.Background(), fakeSession{user: user}) +} + +func newTask(id, contextID string, state a2atype.TaskState, history, artifacts int) *a2atype.Task { + ts := time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC) + t := &a2atype.Task{ + ID: a2atype.TaskID(id), + ContextID: contextID, + Status: a2atype.TaskStatus{State: state, Timestamp: &ts}, + } + for i := range history { + t.History = append(t.History, &a2atype.Message{ID: fmt.Sprintf("%s-msg-%d", id, i), Role: a2atype.MessageRoleUser}) + } + for i := range artifacts { + t.Artifacts = append(t.Artifacts, &a2atype.Artifact{ID: a2atype.ArtifactID(fmt.Sprintf("%s-art-%d", id, i))}) + } + return t +} + +func storeWith(t *testing.T, user, session string, tasks ...*a2atype.Task) *storeTaskQueryHandler { + t.Helper() + store := newFakeStore() + store.addSession(session, user) + for _, tk := range tasks { + store.addTask(session, tk) + } + return newStoreTaskQueryHandler(&PassthroughRequestHandler{}, store) +} + +func TestListTasks_Pagination(t *testing.T) { + tasks := []*a2atype.Task{ + newTask("t1", "s1", a2atype.TaskStateWorking, 0, 0), + newTask("t2", "s1", a2atype.TaskStateWorking, 0, 0), + newTask("t3", "s1", a2atype.TaskStateWorking, 0, 0), + newTask("t4", "s1", a2atype.TaskStateWorking, 0, 0), + newTask("t5", "s1", a2atype.TaskStateWorking, 0, 0), + } + h := storeWith(t, "alice", "s1", tasks...) + ctx := userCtx("alice") + + var seen []string + token := "" + pages := 0 + for { + resp, err := h.ListTasks(ctx, &a2atype.ListTasksRequest{ContextID: "s1", PageSize: 2, PageToken: token}) + require.NoError(t, err) + require.Equal(t, 5, resp.TotalSize) + require.Equal(t, 2, resp.PageSize) + for _, tk := range resp.Tasks { + seen = append(seen, string(tk.ID)) + } + pages++ + if resp.NextPageToken == "" { + break + } + token = resp.NextPageToken + require.LessOrEqual(t, pages, 5, "pagination did not terminate") + } + require.Equal(t, 3, pages) + require.Equal(t, []string{"t1", "t2", "t3", "t4", "t5"}, seen) +} + +func TestListTasks_NextPageTokenAlwaysPresentEmptyOnLastPage(t *testing.T) { + h := storeWith(t, "alice", "s1", + newTask("t1", "s1", a2atype.TaskStateWorking, 0, 0), + newTask("t2", "s1", a2atype.TaskStateWorking, 0, 0), + ) + resp, err := h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{ContextID: "s1", PageSize: 10}) + require.NoError(t, err) + require.Len(t, resp.Tasks, 2) + require.Equal(t, "", resp.NextPageToken, "nextPageToken must be empty string on the final page") + + // Empty result set still reports an (empty) nextPageToken. + empty, err := h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{ContextID: "does-not-exist"}) + require.NoError(t, err) + require.Empty(t, empty.Tasks) + require.Equal(t, "", empty.NextPageToken) +} + +func TestListTasks_IncludeArtifacts(t *testing.T) { + h := storeWith(t, "alice", "s1", newTask("t1", "s1", a2atype.TaskStateCompleted, 0, 3)) + ctx := userCtx("alice") + + off, err := h.ListTasks(ctx, &a2atype.ListTasksRequest{ContextID: "s1"}) + require.NoError(t, err) + require.Nil(t, off.Tasks[0].Artifacts, "artifacts must be omitted when includeArtifacts is false (default)") + + on, err := h.ListTasks(ctx, &a2atype.ListTasksRequest{ContextID: "s1", IncludeArtifacts: true}) + require.NoError(t, err) + require.Len(t, on.Tasks[0].Artifacts, 3) +} + +func TestListTasks_StatusFilter(t *testing.T) { + h := storeWith(t, "alice", "s1", + newTask("t1", "s1", a2atype.TaskStateWorking, 0, 0), + newTask("t2", "s1", a2atype.TaskStateInputRequired, 0, 0), + newTask("t3", "s1", a2atype.TaskStateInputRequired, 0, 0), + ) + resp, err := h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{ContextID: "s1", Status: a2atype.TaskStateInputRequired}) + require.NoError(t, err) + require.Equal(t, 2, resp.TotalSize) + for _, tk := range resp.Tasks { + require.Equal(t, a2atype.TaskStateInputRequired, tk.Status.State) + } +} + +func TestListTasks_StatusTimestampAfter(t *testing.T) { + early := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC) + late := time.Date(2026, 7, 8, 0, 0, 0, 0, time.UTC) + old := newTask("t1", "s1", a2atype.TaskStateWorking, 0, 0) + old.Status.Timestamp = &early + recent := newTask("t2", "s1", a2atype.TaskStateWorking, 0, 0) + recent.Status.Timestamp = &late + + h := storeWith(t, "alice", "s1", old, recent) + cutoff := time.Date(2026, 7, 5, 0, 0, 0, 0, time.UTC) + resp, err := h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{ContextID: "s1", StatusTimestampAfter: &cutoff}) + require.NoError(t, err) + require.Equal(t, 1, resp.TotalSize) + require.Equal(t, "t2", string(resp.Tasks[0].ID)) +} + +func TestListTasks_HistoryLength(t *testing.T) { + h := storeWith(t, "alice", "s1", newTask("t1", "s1", a2atype.TaskStateWorking, 5, 0)) + n := 2 + resp, err := h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{ContextID: "s1", HistoryLength: &n}) + require.NoError(t, err) + require.Len(t, resp.Tasks[0].History, 2) + // The most recent messages are kept. + require.Equal(t, "t1-msg-3", resp.Tasks[0].History[0].ID) + require.Equal(t, "t1-msg-4", resp.Tasks[0].History[1].ID) +} + +func TestCrossUserIsolation(t *testing.T) { + store := newFakeStore() + store.addSession("s1", "alice") + store.addTask("s1", newTask("t1", "s1", a2atype.TaskStateWorking, 0, 0)) + h := newStoreTaskQueryHandler(&PassthroughRequestHandler{}, store) + + // mallory asks for alice's context id: gets nothing. + resp, err := h.ListTasks(userCtx("mallory"), &a2atype.ListTasksRequest{ContextID: "s1"}) + require.NoError(t, err) + require.Empty(t, resp.Tasks) + require.Equal(t, 0, resp.TotalSize) + + // Unauthenticated context: empty, never a leak. + respAnon, err := h.ListTasks(context.Background(), &a2atype.ListTasksRequest{ContextID: "s1"}) + require.NoError(t, err) + require.Empty(t, respAnon.Tasks) +} + +func TestListTasks_ShareContextGrantsOwnerSession(t *testing.T) { + store := newFakeStore() + store.addSession("s1", "alice") + store.addTask("s1", newTask("t1", "s1", a2atype.TaskStateWorking, 0, 0)) + h := newStoreTaskQueryHandler(&PassthroughRequestHandler{}, store) + + // A visitor (bob) holding a share token for alice's session s1 lists it. + ctx := auth.ShareContextTo(userCtx("bob"), &auth.ShareContext{SessionID: "s1", UserID: "alice"}) + resp, err := h.ListTasks(ctx, &a2atype.ListTasksRequest{ContextID: "s1"}) + require.NoError(t, err) + require.Equal(t, 1, resp.TotalSize) + require.Equal(t, "t1", string(resp.Tasks[0].ID)) + + // A share token for a different session must not grant s1. + ctxOther := auth.ShareContextTo(userCtx("bob"), &auth.ShareContext{SessionID: "s2", UserID: "alice"}) + respOther, err := h.ListTasks(ctxOther, &a2atype.ListTasksRequest{ContextID: "s1"}) + require.NoError(t, err) + require.Empty(t, respOther.Tasks) + + // The share token must not widen an all-sessions query to the owner's tasks. + respAll, err := h.ListTasks(ctx, &a2atype.ListTasksRequest{}) + require.NoError(t, err) + require.Empty(t, respAll.Tasks, "share token grants one session, not the owner's whole account") +} + +func TestListTasks_AcrossAllUserSessions(t *testing.T) { + store := newFakeStore() + store.addSession("s1", "alice") + store.addSession("s2", "alice") + store.addSession("s3", "bob") + store.addTask("s1", newTask("t1", "s1", a2atype.TaskStateWorking, 0, 0)) + store.addTask("s2", newTask("t2", "s2", a2atype.TaskStateWorking, 0, 0)) + store.addTask("s3", newTask("t3", "s3", a2atype.TaskStateWorking, 0, 0)) + h := newStoreTaskQueryHandler(&PassthroughRequestHandler{}, store) + + resp, err := h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{}) + require.NoError(t, err) + require.Equal(t, 2, resp.TotalSize) + got := []string{string(resp.Tasks[0].ID), string(resp.Tasks[1].ID)} + require.ElementsMatch(t, []string{"t1", "t2"}, got) +} + +// failingTaskStore fails every read with a backend error. +type failingTaskStore struct{ err error } + +func (f failingTaskStore) GetSession(context.Context, string, string) (*dbpkg.Session, error) { + 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) ([]*a2atype.Task, error) { + return nil, f.err +} + +func TestListTasks_BackendFailurePropagates(t *testing.T) { + backendErr := fmt.Errorf("failed to get session s1: connection refused") + h := newStoreTaskQueryHandler(&PassthroughRequestHandler{}, failingTaskStore{err: backendErr}) + + // A store failure must surface as an error, not an empty task list. + _, err := h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{ContextID: "s1"}) + require.ErrorContains(t, err, "connection refused") + + _, err = h.ListTasks(userCtx("alice"), &a2atype.ListTasksRequest{}) + require.ErrorContains(t, err, "connection refused") +} + +// ── Wire tests: identical filtering, different enum casing ────────────────── + +func withUser(next http.Handler, user string) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + next.ServeHTTP(w, r.WithContext(auth.AuthSessionTo(r.Context(), fakeSession{user: user}))) + }) +} + +func rpcCall(t *testing.T, h http.Handler, body string) map[string]any { + t.Helper() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + require.Equal(t, http.StatusOK, rec.Code, "body: %s", rec.Body.String()) + var out map[string]any + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &out), "raw: %s", rec.Body.String()) + return out +} + +func wireHandlers(user string, tasks ...*a2atype.Task) (v1 http.Handler, v0 http.Handler) { + store := newFakeStore() + store.addSession("s1", user) + for _, tk := range tasks { + store.addTask("s1", tk) + } + h := newStoreTaskQueryHandler(&PassthroughRequestHandler{}, store) + v1 = withUser(a2asrv.NewJSONRPCHandler(h), user) + v0 = withUser(newV0TasksListInterceptor(a2av0.NewJSONRPCHandler(h), h), user) + return v1, v0 +} + +func TestWire_ListTasksStateCasing(t *testing.T) { + tasks := []*a2atype.Task{ + newTask("t1", "s1", a2atype.TaskStateInputRequired, 0, 0), + newTask("t2", "s1", a2atype.TaskStateWorking, 0, 0), + } + v1, v0 := wireHandlers("alice", tasks...) + + // v1 wire: uppercase TaskState, method "ListTasks". + v1resp := rpcCall(t, v1, `{"jsonrpc":"2.0","id":1,"method":"ListTasks","params":{"contextId":"s1","status":"TASK_STATE_INPUT_REQUIRED"}}`) + v1result := v1resp["result"].(map[string]any) + v1list := v1result["tasks"].([]any) + require.Len(t, v1list, 1) + require.Equal(t, float64(1), v1result["totalSize"]) + require.Contains(t, v1result, "nextPageToken") + v1state := v1list[0].(map[string]any)["status"].(map[string]any)["state"].(string) + require.Equal(t, "TASK_STATE_INPUT_REQUIRED", v1state) + + // v0 wire: lowercase TaskState, method "tasks/list". + v0resp := rpcCall(t, v0, `{"jsonrpc":"2.0","id":1,"method":"tasks/list","params":{"contextId":"s1","status":"input-required"}}`) + v0result := v0resp["result"].(map[string]any) + v0list := v0result["tasks"].([]any) + require.Len(t, v0list, 1, "v0 must filter identically to v1") + require.Equal(t, float64(1), v0result["totalSize"]) + require.Contains(t, v0result, "nextPageToken") + v0state := v0list[0].(map[string]any)["status"].(map[string]any)["state"].(string) + require.Equal(t, "input-required", v0state) +} + +func TestWire_V0UnknownMethodDelegates(t *testing.T) { + _, v0 := wireHandlers("alice") + // A method the interceptor does not own must fall through to the v0 handler, + // which reports method-not-found rather than being swallowed here. + resp := rpcCall(t, v0, `{"jsonrpc":"2.0","id":1,"method":"tasks/frobnicate","params":{}}`) + require.Contains(t, resp, "error") +} + +func TestWire_V0TasksListWithoutStoreIsMethodNotFound(t *testing.T) { + // Without a task store the v0 interceptor must not be installed: tasks/list + // keeps the legacy wire's native method-not-found instead of hitting the + // passthrough (which would surface ErrUnsupportedOperation). + _, legacy := newTaskQueryHandlers(&PassthroughRequestHandler{}, nil) + resp := rpcCall(t, withUser(legacy, "alice"), `{"jsonrpc":"2.0","id":1,"method":"tasks/list","params":{}}`) + require.Contains(t, resp, "error") + errObj := resp["error"].(map[string]any) + require.Equal(t, float64(-32601), errObj["code"], "expected JSON-RPC method-not-found") +} diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 4e0bdc66d..c4bcfa3b6 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -101,6 +101,9 @@ func (c *postgresClient) StoreSession(ctx context.Context, session *dbpkg.Sessio func (c *postgresClient) GetSession(ctx context.Context, sessionID, userID string) (*dbpkg.Session, error) { row, err := c.q.GetSession(ctx, dbgen.GetSessionParams{ID: sessionID, UserID: userID}) if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, fmt.Errorf("session %s: %w", sessionID, dbpkg.ErrNotFound) + } return nil, fmt.Errorf("failed to get session %s: %w", sessionID, err) } return toSession(row), nil diff --git a/go/core/internal/database/client_test.go b/go/core/internal/database/client_test.go index a8c263d7d..ca752b972 100644 --- a/go/core/internal/database/client_test.go +++ b/go/core/internal/database/client_test.go @@ -232,6 +232,12 @@ func TestStoreSessionIdempotence(t *testing.T) { retrieved, err := client.GetSession(ctx, session.ID, userID) require.NoError(t, err) assert.Equal(t, "Updated", *retrieved.Name, "Session should have updated name") + + _, err = client.GetSession(ctx, "no-such-session", userID) + require.ErrorIs(t, err, dbpkg.ErrNotFound) + + _, err = client.GetSession(ctx, session.ID, "other-user") + require.ErrorIs(t, err, dbpkg.ErrNotFound, "another user's session must read as not found") } func TestListSessionsOrdersByRecentActivity(t *testing.T) { diff --git a/go/core/pkg/app/app.go b/go/core/pkg/app/app.go index 6cadbdd6d..2d75b0b03 100644 --- a/go/core/pkg/app/app.go +++ b/go/core/pkg/app/app.go @@ -652,7 +652,7 @@ func Start(getExtensionConfig GetExtensionConfig, extraSources []migrations.Sour } // Register A2A handlers on all replicas - a2aHandler := a2a.NewA2AHttpMux(httpserver.APIPathA2A, httpserver.APIPathA2ASandboxes, extensionCfg.Authenticator) + a2aHandler := a2a.NewA2AHttpMux(httpserver.APIPathA2A, httpserver.APIPathA2ASandboxes, extensionCfg.Authenticator, dbClient) ateneRouterURL := cfg.Substrate.AtenetRouterURL if ateneRouterURL == "" { ateneRouterURL = substrate.DefaultAtenetRouterURL