feat: method extensions prototyping

a2aclient/client.go

@@ -16,10 +16,12 @@ package a2aclient
 
 import (
 	"context"
+	"fmt"
 	"iter"
 	"sync/atomic"
 
 	"github.com/a2aproject/a2a-go/a2a"
+	"github.com/a2aproject/a2a-go/a2aext"
 	"github.com/a2aproject/a2a-go/internal/utils"
 )
 
@@ -270,6 +272,26 @@ func (c *Client) GetAgentCard(ctx context.Context) (*a2a.AgentCard, error) {
 	return resp, err
 }
 
+func Invoke[Req, Resp any](ctx context.Context, client *Client, method *a2aext.UnaryClientMethod[Req, Resp], arg *Req) (*Resp, error) {
+	ctx, err := client.interceptBefore(ctx, method.Name(), arg)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := client.transport.Invoke(ctx, method, arg)
+	if errOverride := client.interceptAfter(ctx, method.Name(), resp, err); errOverride != nil {
+		return nil, errOverride
+	}
+
+	typedResp, ok := resp.(*Resp)
+	if !ok {
+		var want Resp
+		return nil, fmt.Errorf("unexpected extension method result type %T, want pointer to %T", resp, want)
+	}
+
+	return typedResp, err
+}
+
 func (c *Client) Destroy() error {
 	return c.transport.Destroy()
 }

a2aclient/jsonrpc.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/a2aproject/a2a-go/a2a"
+	"github.com/a2aproject/a2a-go/a2aext"
 	"github.com/a2aproject/a2a-go/internal/jsonrpc"
 	"github.com/a2aproject/a2a-go/internal/sse"
 	"github.com/a2aproject/a2a-go/log"
@@ -47,6 +48,28 @@ type jsonrpcResponse struct {
 	Error   *jsonrpc.Error  `json:"error,omitempty"`
 }
 
+type jsonrpcExtensionBinding struct {
+	Method            string
+	ResponseFromBytes func(resp []byte) (any, error)
+}
+
+func NewJSONRPCExtensionBinding[Resp any](name string) a2aext.Binding {
+	return &jsonrpcExtensionBinding{
+		Method: name,
+		ResponseFromBytes: func(resp []byte) (any, error) {
+			var zero Resp
+			if err := json.Unmarshal(resp, &zero); err != nil {
+				return nil, err
+			}
+			return &zero, nil
+		},
+	}
+}
+
+func (*jsonrpcExtensionBinding) Protocol() a2a.TransportProtocol {
+	return a2a.TransportProtocolJSONRPC
+}
+
 // JSONRPCOption configures optional parameters for the JSONRPC transport.
 // Options are applied during NewJSONRPCTransport initialization.
 type JSONRPCOption func(*jsonrpcTransport)
@@ -351,6 +374,29 @@ func (t *jsonrpcTransport) DeleteTaskPushConfig(ctx context.Context, params *a2a
 	return err
 }
 
+func (t *jsonrpcTransport) Invoke(ctx context.Context, method a2aext.Method, req any) (any, error) {
+	binding, ok := method.Binding(a2a.TransportProtocolJSONRPC)
+	if !ok {
+		return nil, fmt.Errorf("method %s is not bound to JSON-RPC", method.Name())
+	}
+
+	typedBinding, ok := binding.(*jsonrpcExtensionBinding)
+	if !ok {
+		return nil, fmt.Errorf("method %s is not bound to JSON-RPC", method.Name())
+	}
+
+	result, err := t.sendRequest(ctx, typedBinding.Method, req)
+	if err != nil {
+		return nil, err
+	}
+
+	response, err := typedBinding.ResponseFromBytes(result)
+	if err != nil {
+		return nil, fmt.Errorf("result violates A2A spec - could not determine type: %w; data: %s", err, string(result))
+	}
+	return response, nil
+}
+
 // GetAgentCard retrieves the agent's card.
 func (t *jsonrpcTransport) GetAgentCard(ctx context.Context) (*a2a.AgentCard, error) {
 	result, err := t.sendRequest(ctx, jsonrpc.MethodGetExtendedAgentCard, nil)

a2aclient/transport.go

@@ -20,6 +20,7 @@ import (
 	"iter"
 
 	"github.com/a2aproject/a2a-go/a2a"
+	"github.com/a2aproject/a2a-go/a2aext"
 )
 
 // A2AClient defines a transport-agnostic interface for making A2A requests.
@@ -56,6 +57,9 @@ type Transport interface {
 	// If extended card is supported calls the 'agent/getAuthenticatedExtendedCard' protocol method.
 	GetAgentCard(ctx context.Context) (*a2a.AgentCard, error)
 
+	// Invoke calls the provided extension method.
+	Invoke(ctx context.Context, client a2aext.Method, arg any) (any, error)
+
 	// Clean up resources associated with the transport (eg. close a gRPC channel).
 	Destroy() error
 }
@@ -120,6 +124,10 @@ func (unimplementedTransport) GetAgentCard(ctx context.Context) (*a2a.AgentCard,
 	return nil, errNotImplemented
 }
 
+func (unimplementedTransport) Invoke(ctx context.Context, client a2aext.Method, arg any) (any, error) {
+	return nil, errNotImplemented
+}
+
 func (unimplementedTransport) Destroy() error {
 	return nil
 }

a2aext/method.go

@@ -0,0 +1,92 @@
+package a2aext
+
+import (
+	"context"
+	"errors"
+	"iter"
+
+	"github.com/a2aproject/a2a-go/a2a"
+)
+
+type Binding interface {
+	Protocol() a2a.TransportProtocol
+}
+
+type Method interface {
+	Name() string
+	Binding(a2a.TransportProtocol) (Binding, bool)
+	Streaming() bool
+}
+
+type ServerMethod interface {
+	Method
+
+	InvokeUnary(ctx context.Context, arg any) (any, error)
+	InvokeStreaming(ctx context.Context, arg any) iter.Seq2[any, error]
+}
+
+type methodDescriptor struct {
+	name      string
+	bindings  map[a2a.TransportProtocol]Binding
+	streaming bool
+}
+
+func (d *methodDescriptor) Name() string { return d.name }
+
+func (d *methodDescriptor) Binding(protocol a2a.TransportProtocol) (Binding, bool) {
+	b, ok := d.bindings[protocol]
+	return b, ok
+}
+
+func (d *methodDescriptor) Streaming() bool { return d.streaming }
+
+type UnaryClientMethod[Req, Resp any] struct {
+	methodDescriptor
+}
+
+var _ Method = (*UnaryClientMethod[any, any])(nil)
+
+// NewUnaryClientMethod creates a new unary extension method definition.
+func NewUnaryClientMethod[Arg, Res any](
+	name string,
+	bindings ...Binding,
+) *UnaryClientMethod[Arg, Res] {
+	bmap := map[a2a.TransportProtocol]Binding{}
+	for _, b := range bindings {
+		bmap[b.Protocol()] = b
+	}
+	return &UnaryClientMethod[Arg, Res]{
+		methodDescriptor: methodDescriptor{name: name, bindings: bmap, streaming: false},
+	}
+}
+
+type unaryServerMethod[Arg, Res any] struct {
+	methodDescriptor
+	call func(context.Context, *Arg) (*Res, error)
+}
+
+// NewUnaryMethod creates a new unary extension method definition.
+func NewUnaryServerMethod[Arg, Res any](
+	name string,
+	call func(context.Context, *Arg) (*Res, error),
+	bindings ...Binding,
+) ServerMethod {
+	bmap := map[a2a.TransportProtocol]Binding{}
+	for _, b := range bindings {
+		bmap[b.Protocol()] = b
+	}
+	return &unaryServerMethod[Arg, Res]{
+		methodDescriptor: methodDescriptor{name: name, bindings: bmap, streaming: false},
+		call:             call,
+	}
+}
+
+func (m *unaryServerMethod[Arg, Res]) InvokeUnary(ctx context.Context, arg any) (any, error) {
+	return m.call(ctx, arg.(*Arg))
+}
+
+func (m *unaryServerMethod[Arg, Res]) InvokeStreaming(ctx context.Context, arg any) iter.Seq2[any, error] {
+	return func(yield func(any, error) bool) {
+		yield(nil, errors.New("unary method invoked as streaming"))
+	}
+}

a2asrv/handler.go

@@ -21,6 +21,7 @@ import (
 	"log/slog"
 
 	"github.com/a2aproject/a2a-go/a2a"
+	"github.com/a2aproject/a2a-go/a2aext"
 	"github.com/a2aproject/a2a-go/a2asrv/eventqueue"
 	"github.com/a2aproject/a2a-go/a2asrv/limiter"
 	"github.com/a2aproject/a2a-go/a2asrv/push"
@@ -59,6 +60,15 @@ type RequestHandler interface {
 
 	// GetAgentCard returns an extended a2a.AgentCard if configured.
 	OnGetExtendedAgentCard(ctx context.Context) (*a2a.AgentCard, error)
+
+	// GetAgentCard returns an extended a2a.AgentCard if configured.
+	GetExtensionMethods() []a2aext.ServerMethod
+
+	// Invoke calls a registered extension method.
+	Invoke(ctx context.Context, method a2aext.ServerMethod, arg any) (any, error)
+
+	// InvokeStreaming calls a registered streaming extension method.
+	InvokeStreaming(ctx context.Context, method a2aext.ServerMethod, arg any) iter.Seq2[any, error]
 }
 
 // Implements a2asrv.RequestHandler.
@@ -75,6 +85,8 @@ type defaultRequestHandler struct {
 	reqContextInterceptors []RequestContextInterceptor
 
 	authenticatedCardProducer AgentCardProducer
+
+	extensionMethods []a2aext.ServerMethod
 }
 
 // RequestHandlerOption can be used to customize the default [RequestHandler] implementation behavior.
@@ -104,12 +116,20 @@ func WithConcurrencyConfig(config limiter.ConcurrencyConfig) RequestHandlerOptio
 	}
 }
 
+// WithConcurrencyConfig allows to set limits on the number of concurrent executions.
+func WithExtensionMethod(method a2aext.ServerMethod) RequestHandlerOption {
+	return func(ih *InterceptedHandler, h *defaultRequestHandler) {
+		h.extensionMethods = append(h.extensionMethods, method)
+	}
+}
+
 // NewHandler creates a new request handler.
 func NewHandler(executor AgentExecutor, options ...RequestHandlerOption) RequestHandler {
 	h := &defaultRequestHandler{
-		agentExecutor: executor,
-		queueManager:  eventqueue.NewInMemoryManager(),
-		taskStore:     taskstore.NewMem(),
+		agentExecutor:    executor,
+		extensionMethods: []a2aext.ServerMethod{},
+		queueManager:     eventqueue.NewInMemoryManager(),
+		taskStore:        taskstore.NewMem(),
 		// push notifications are not supported by default
 	}
 	ih := &InterceptedHandler{Handler: h, Logger: slog.Default()}
@@ -307,6 +327,18 @@ func (h *defaultRequestHandler) OnGetExtendedAgentCard(ctx context.Context) (*a2
 	return h.authenticatedCardProducer.Card(ctx)
 }
 
+func (h *defaultRequestHandler) GetExtensionMethods() []a2aext.ServerMethod {
+	return h.extensionMethods
+}
+
+func (h *defaultRequestHandler) Invoke(ctx context.Context, method a2aext.ServerMethod, arg any) (any, error) {
+	return method.InvokeUnary(ctx, arg)
+}
+
+func (h *defaultRequestHandler) InvokeStreaming(ctx context.Context, method a2aext.ServerMethod, arg any) iter.Seq2[any, error] {
+	return method.InvokeStreaming(ctx, arg)
+}
+
 func shouldInterruptNonStreaming(params *a2a.MessageSendParams, event a2a.Event) (a2a.TaskID, bool) {
 	// Non-blocking clients receive a result on the first task event, default Blocking to TRUE
 	if params.Config != nil && params.Config.Blocking != nil && !(*params.Config.Blocking) {

a2asrv/intercepted_handler.go

@@ -22,6 +22,7 @@ import (
 	"github.com/google/uuid"
 
 	"github.com/a2aproject/a2a-go/a2a"
+	"github.com/a2aproject/a2a-go/a2aext"
 	"github.com/a2aproject/a2a-go/log"
 )
 
@@ -225,6 +226,45 @@ func (h *InterceptedHandler) OnGetExtendedAgentCard(ctx context.Context) (*a2a.A
 	return response, err
 }
 
+func (h *InterceptedHandler) GetExtensionMethods() []a2aext.ServerMethod {
+	return h.Handler.GetExtensionMethods()
+}
+
+func (h *InterceptedHandler) Invoke(ctx context.Context, method a2aext.ServerMethod, arg any) (any, error) {
+	ctx, callCtx := withMethodCallContext(ctx, method.Name())
+	ctx = h.withLoggerContext(ctx)
+	ctx, err := h.interceptBefore(ctx, callCtx, arg)
+	if err != nil {
+		return nil, err
+	}
+	response, err := h.Handler.Invoke(ctx, method, arg)
+	if errOverride := h.interceptAfter(ctx, callCtx, response, err); errOverride != nil {
+		return nil, errOverride
+	}
+	return response, err
+}
+
+func (h *InterceptedHandler) InvokeStreaming(ctx context.Context, method a2aext.ServerMethod, arg any) iter.Seq2[any, error] {
+	return func(yield func(any, error) bool) {
+		ctx, callCtx := withMethodCallContext(ctx, method.Name())
+		ctx = h.withLoggerContext(ctx)
+		ctx, err := h.interceptBefore(ctx, callCtx, arg)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+		for event, err := range h.Handler.InvokeStreaming(ctx, method, arg) {
+			if errOverride := h.interceptAfter(ctx, callCtx, event, err); errOverride != nil {
+				yield(nil, errOverride)
+				return
+			}
+			if !yield(event, err) {
+				return
+			}
+		}
+	}
+}
+
 func (h *InterceptedHandler) interceptBefore(ctx context.Context, callCtx *CallContext, payload any) (context.Context, error) {
 	request := &Request{Payload: payload}