Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions go/api/database/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
8 changes: 8 additions & 0 deletions go/api/database/errors.go
Original file line number Diff line number Diff line change
@@ -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")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this makes sense overall since it's probably better to hide the pgx errors, however this would only be one instance of that with still many pgx errors around. Can you leave this change for a follow-up where we clean-up all pgx leakage?

26 changes: 23 additions & 3 deletions go/core/internal/a2a/a2a_handler_mux.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ type handlerMux struct {
agentPathPrefix string
sandboxPathPrefix string
authenticator auth.AuthProvider
taskStore TaskStore
}

var _ A2AHandlerMux = &handlerMux{}
Expand All @@ -45,24 +46,43 @@ 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,
card a2atype.AgentCard,
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, "/")

Expand Down
164 changes: 164 additions & 0 deletions go/core/internal/a2a/a2av0_tasks_list.go
Original file line number Diff line number Diff line change
@@ -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, &params); 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)
}
}
Loading
Loading