server middleware API

Author: yarolegovich

a2a/core.go

@@ -34,6 +34,8 @@ func (*Message) isSendMessageResult() {}
 // Event interface is used to represent types that can be sent over a streaming connection.
 type Event interface {
 	isEvent()
+
+	Meta() map[string]any
 }
 
 func (*Message) isEvent()                 {}
@@ -54,6 +56,8 @@ func NewMessageID() string {
 	return uuid.NewString()
 }
 
+var _ Event = (*Message)(nil)
+
 // Message represents a single message in the conversation between a user and an agent.
 type Message struct {
 	// ID is a unique identifier for the message, typically a UUID, generated by the sender.
@@ -105,6 +109,10 @@ func NewMessageForTask(role MessageRole, task *Task, parts ...Part) *Message {
 	}
 }
 
+func (m *Message) Meta() map[string]any {
+	return m.Metadata
+}
+
 // TaskID is a unique identifier for the task, generated by the server for a new task.
 type TaskID string
 
@@ -142,6 +150,8 @@ func (ts TaskState) Terminal() bool {
 		ts == TaskStateRejected
 }
 
+var _ Event = (*Task)(nil)
+
 // Task represents a single, stateful operation or conversation between a client and an agent.
 type Task struct {
 	// ID is a unique identifier for the task, generated by the server for a new task.
@@ -176,6 +186,10 @@ type TaskStatus struct {
 	Timestamp *time.Time `json:"timestamp,omitempty" yaml:"timestamp,omitempty" mapstructure:"timestamp,omitempty"`
 }
 
+func (m *Task) Meta() map[string]any {
+	return m.Metadata
+}
+
 // ArtifactID is a unique identifier for the artifact within the scope of the task.
 type ArtifactID string
 
@@ -205,6 +219,8 @@ type Artifact struct {
 	Parts ContentParts `json:"parts" yaml:"parts" mapstructure:"parts"`
 }
 
+var _ Event = (*TaskArtifactUpdateEvent)(nil)
+
 // TaskArtifactUpdateEvent is an event sent by the agent to notify the client that an artifact has been
 // generated or updated. This is typically used in streaming models.
 type TaskArtifactUpdateEvent struct {
@@ -228,6 +244,10 @@ type TaskArtifactUpdateEvent struct {
 	Metadata map[string]any `json:"metadata,omitempty" yaml:"metadata,omitempty" mapstructure:"metadata,omitempty"`
 }
 
+func (a *TaskArtifactUpdateEvent) Meta() map[string]any {
+	return a.Metadata
+}
+
 // NewArtifactEvent create a TaskArtifactUpdateEvent for an Artifact with a random ID.
 func NewArtifactEvent(task *Task, parts ...Part) *TaskArtifactUpdateEvent {
 	return &TaskArtifactUpdateEvent{
@@ -253,6 +273,8 @@ func NewArtifactUpdateEvent(task *Task, id ArtifactID, parts ...Part) *TaskArtif
 	}
 }
 
+var _ Event = (*TaskStatusUpdateEvent)(nil)
+
 // TaskStatusUpdateEvent is an event sent by the agent to notify the client of a change in a task's status.
 // This is typically used in streaming or subscription models.
 type TaskStatusUpdateEvent struct {
@@ -286,6 +308,10 @@ func NewStatusUpdateEvent(task *Task, state TaskState, msg *Message) *TaskStatus
 	}
 }
 
+func (a *TaskStatusUpdateEvent) Meta() map[string]any {
+	return a.Metadata
+}
+
 // ContentParts is an array of content parts that form the message body or an artifact.
 type ContentParts []Part
 

a2asrv/auth.go

@@ -0,0 +1,34 @@
+package a2asrv
+
+// User can be attached to call context by authentication middleware.
+type User interface {
+	// Name returns a username.
+	Name() string
+	// Authenticated returns true if requested was authenticated.
+	Authenticated() bool
+}
+
+// AuthenticatedUser is a simple implementation of User interface which can be configured with a username.
+type AuthenticatedUser struct {
+	UserName string
+}
+
+func (u *AuthenticatedUser) Name() string {
+	return u.UserName
+}
+
+func (u *AuthenticatedUser) Authenticated() bool {
+	return true
+}
+
+type unauthenticatedUser struct {
+	UserName string
+}
+
+func (unauthenticatedUser) Name() string {
+	return ""
+}
+
+func (unauthenticatedUser) Authenticated() bool {
+	return false
+}

a2asrv/extensions.go

@@ -0,0 +1,58 @@
+package a2asrv
+
+import (
+	"context"
+	"slices"
+
+	"github.com/a2aproject/a2a-go/a2a"
+)
+
+const ExtensionsMetaKey = "X-A2A-Extensions"
+
+// Extensions provides utility methods for accessing extensions requested by the client and keeping track of extensions
+// activated during request processing.
+type Extensions struct {
+	callCtx *CallContext
+}
+
+// ExtensionsFrom is a helper function for quick access to Extensions in the current CallContext.
+func ExtensionsFrom(ctx context.Context) (*Extensions, bool) {
+	serverCallCtx, ok := CallContextFrom(ctx)
+	if !ok {
+		return nil, false
+	}
+	return serverCallCtx.Extensions(), true
+}
+
+// Active returns true if an extension has already been activated in the current CallContext using ExtensionContext.Activate.
+func (e *Extensions) Active(extension *a2a.AgentExtension) bool {
+	return slices.Contains(e.callCtx.activatedExtensions, extension.URI)
+}
+
+// Activate marks extension as activated in the current CallContext. A list of activated extensions might be attached as
+// response metadata by a transport implementation.
+func (e *Extensions) Activate(extension *a2a.AgentExtension) {
+	if e.Active(extension) {
+		return
+	}
+	e.callCtx.activatedExtensions = append(e.callCtx.activatedExtensions, extension.URI)
+}
+
+// ActivatedURIs returns all URIs activated during call execution.
+func (e *Extensions) ActivatedURIs() []string {
+	return slices.Clone(e.callCtx.activatedExtensions)
+}
+
+// Requested returns true if the provided extension was requested by the client.
+func (e *Extensions) Requested(extension *a2a.AgentExtension) bool {
+	return slices.Contains(e.RequestedURIs(), extension.URI)
+}
+
+// RequestedURIs returns all URIs of extensions requested by the client.
+func (e *Extensions) RequestedURIs() []string {
+	requested, ok := e.callCtx.RequestMeta().Get(ExtensionsMetaKey)
+	if !ok {
+		return []string{}
+	}
+	return slices.Clone(requested)
+}

a2asrv/intercepted_handler.go

@@ -0,0 +1,185 @@
+package a2asrv
+
+import (
+	"context"
+	"iter"
+
+	"github.com/a2aproject/a2a-go/a2a"
+)
+
+// InterceptedHandler implements RequestHandler. It can be used to attach call interceptors and initialize
+// call context for every method of the wrapped handler.
+type InterceptedHandler struct {
+	// Handler is responsible for the actual processing of every call.
+	Handler RequestHandler
+	// Interceptors is a list of call interceptors which will be applied before and after each call.
+	Interceptors []CallInterceptor
+}
+
+func (h *InterceptedHandler) OnGetTask(ctx context.Context, query *a2a.TaskQueryParams) (*a2a.Task, error) {
+	ctx, callCtx := withMethodCallContext(ctx, "OnGetTask")
+	ctx, err := h.interceptBefore(ctx, callCtx, query)
+	if err != nil {
+		return nil, err
+	}
+	response, err := h.Handler.OnGetTask(ctx, query)
+	if errOverride := h.interceptAfter(ctx, callCtx, response, err); errOverride != nil {
+		return nil, errOverride
+	}
+	return response, err
+}
+
+func (h *InterceptedHandler) OnCancelTask(ctx context.Context, params *a2a.TaskIDParams) (*a2a.Task, error) {
+	ctx, callCtx := withMethodCallContext(ctx, "OnCancelTask")
+	ctx, err := h.interceptBefore(ctx, callCtx, params)
+	if err != nil {
+		return nil, err
+	}
+	response, err := h.Handler.OnCancelTask(ctx, params)
+	if errOverride := h.interceptAfter(ctx, callCtx, response, err); errOverride != nil {
+		return nil, errOverride
+	}
+	return response, err
+}
+
+func (h *InterceptedHandler) OnSendMessage(ctx context.Context, params *a2a.MessageSendParams) (a2a.SendMessageResult, error) {
+	ctx, callCtx := withMethodCallContext(ctx, "OnSendMessage")
+	ctx, err := h.interceptBefore(ctx, callCtx, params)
+	if err != nil {
+		return nil, err
+	}
+	response, err := h.Handler.OnSendMessage(ctx, params)
+	if errOverride := h.interceptAfter(ctx, callCtx, response, err); errOverride != nil {
+		return nil, errOverride
+	}
+	return response, err
+}
+
+func (h *InterceptedHandler) OnSendMessageStream(ctx context.Context, params *a2a.MessageSendParams) iter.Seq2[a2a.Event, error] {
+	return func(yield func(a2a.Event, error) bool) {
+		ctx, callCtx := withMethodCallContext(ctx, "OnSendMessageStream")
+		ctx, err := h.interceptBefore(ctx, callCtx, params)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+		for event, err := range h.Handler.OnSendMessageStream(ctx, params) {
+			if errOverride := h.interceptAfter(ctx, callCtx, event, err); errOverride != nil {
+				yield(nil, errOverride)
+				return
+			}
+			if !yield(event, err) {
+				return
+			}
+		}
+	}
+}
+
+func (h *InterceptedHandler) OnResubscribeToTask(ctx context.Context, params *a2a.TaskIDParams) iter.Seq2[a2a.Event, error] {
+	return func(yield func(a2a.Event, error) bool) {
+		ctx, callCtx := withMethodCallContext(ctx, "OnResubscribeToTask")
+		ctx, err := h.interceptBefore(ctx, callCtx, params)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+		for event, err := range h.Handler.OnResubscribeToTask(ctx, params) {
+			if errOverride := h.interceptAfter(ctx, callCtx, event, err); errOverride != nil {
+				yield(nil, errOverride)
+				return
+			}
+			if !yield(event, err) {
+				return
+			}
+		}
+	}
+}
+
+func (h *InterceptedHandler) OnGetTaskPushConfig(ctx context.Context, params *a2a.GetTaskPushConfigParams) (*a2a.TaskPushConfig, error) {
+	ctx, callCtx := withMethodCallContext(ctx, "OnGetTaskPushConfig")
+	ctx, err := h.interceptBefore(ctx, callCtx, params)
+	if err != nil {
+		return nil, err
+	}
+	response, err := h.Handler.OnGetTaskPushConfig(ctx, params)
+	if errOverride := h.interceptAfter(ctx, callCtx, response, err); errOverride != nil {
+		return nil, errOverride
+	}
+	return response, err
+}
+
+func (h *InterceptedHandler) OnListTaskPushConfig(ctx context.Context, params *a2a.ListTaskPushConfigParams) ([]*a2a.TaskPushConfig, error) {
+	ctx, callCtx := withMethodCallContext(ctx, "OnListTaskPushConfig")
+	ctx, err := h.interceptBefore(ctx, callCtx, params)
+	if err != nil {
+		return nil, err
+	}
+	response, err := h.Handler.OnListTaskPushConfig(ctx, params)
+	if errOverride := h.interceptAfter(ctx, callCtx, response, err); errOverride != nil {
+		return nil, errOverride
+	}
+	return response, err
+}
+
+func (h *InterceptedHandler) OnSetTaskPushConfig(ctx context.Context, params *a2a.TaskPushConfig) (*a2a.TaskPushConfig, error) {
+	ctx, callCtx := withMethodCallContext(ctx, "OnSetTaskPushConfig")
+	ctx, err := h.interceptBefore(ctx, callCtx, params)
+	if err != nil {
+		return nil, err
+	}
+	response, err := h.Handler.OnSetTaskPushConfig(ctx, params)
+	if errOverride := h.interceptAfter(ctx, callCtx, response, err); errOverride != nil {
+		return nil, errOverride
+	}
+	return response, err
+}
+
+func (h *InterceptedHandler) OnDeleteTaskPushConfig(ctx context.Context, params *a2a.DeleteTaskPushConfigParams) error {
+	ctx, callCtx := withMethodCallContext(ctx, "OnDeleteTaskPushConfig")
+	ctx, err := h.interceptBefore(ctx, callCtx, params)
+	if err != nil {
+		return err
+	}
+	err = h.Handler.OnDeleteTaskPushConfig(ctx, params)
+	if errOverride := h.interceptAfter(ctx, callCtx, nil, err); errOverride != nil {
+		return errOverride
+	}
+	return err
+}
+
+func (h *InterceptedHandler) interceptBefore(ctx context.Context, callCtx *CallContext, payload any) (context.Context, error) {
+	request := &Request{Payload: payload}
+
+	for _, interceptor := range h.Interceptors {
+		localCtx, err := interceptor.Before(ctx, callCtx, request)
+		if err != nil {
+			return ctx, err
+		}
+		ctx = localCtx
+	}
+
+	return ctx, nil
+}
+
+func (h *InterceptedHandler) interceptAfter(ctx context.Context, callCtx *CallContext, payload any, responseErr error) error {
+	response := &Response{Payload: payload, Err: responseErr}
+
+	for _, interceptor := range h.Interceptors {
+		if err := interceptor.After(ctx, callCtx, response); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// withMethodCallContext is a private utility function which modifies CallContext.method if a CallContext
+// was passed by a transport implementation or initializes a new CallContext with the provided method.
+func withMethodCallContext(ctx context.Context, method string) (context.Context, *CallContext) {
+	callCtx, ok := CallContextFrom(ctx)
+	if !ok {
+		ctx, callCtx = WithCallContext(ctx, nil)
+	}
+	callCtx.method = method
+	return ctx, callCtx
+}