feat: blocking flag handling (a2aproject#97)

* Add `Polling` option to `a2aclient.Config`. If set to true `Blocking:
false` message configuration is sent to the server.
* Add `Blocking: false` handling on `a2asrv`. When execution results in
a Task, the Task is returned to the caller immediately.
* Make `Blocking` optional to match the "default to true if not set"
behavior all other SDKs implement.

Considered returning immediately when a requests references an existing
Task, but decided that'd be confusing to always receive the task in
`input-required` immediately after providing a follow-up.

re a2aproject#96

Author: yarolegovich

a2a/core.go

@@ -691,8 +691,8 @@ type MessageSendConfig struct {
 	AcceptedOutputModes []string `json:"acceptedOutputModes,omitempty" yaml:"acceptedOutputModes,omitempty" mapstructure:"acceptedOutputModes,omitempty"`
 
 	// Blocking indicates if the client will wait for the task to complete. The server may reject
-	// this if the task is long-running.
-	Blocking bool `json:"blocking,omitempty" yaml:"blocking,omitempty" mapstructure:"blocking,omitempty"`
+	// this if the task is long-running. Server might choose to default to true.
+	Blocking *bool `json:"blocking,omitempty" yaml:"blocking,omitempty" mapstructure:"blocking,omitempty"`
 
 	// HistoryLength is the number of most recent messages from the task's history to retrieve in the response.
 	HistoryLength *int `json:"historyLength,omitempty" yaml:"historyLength,omitempty" mapstructure:"historyLength,omitempty"`

a2aclient/client.go

@@ -20,6 +20,7 @@ import (
 	"sync/atomic"
 
 	"github.com/a2aproject/a2a-go/a2a"
+	"github.com/a2aproject/a2a-go/internal/utils"
 )
 
 // Config exposes options for customizing Client behavior.
@@ -38,6 +39,10 @@ type Config struct {
 	// If there's no overlap in supported Transport Factory will return an error on Client
 	// creation attempt.
 	PreferredTransports []a2a.TransportProtocol
+	// Whether client prefers to poll for task updates instead of blocking until a terminal state is reached.
+	// If set to true, non-streaming send message result might be a Message or a Task in any (including non-terminal) state.
+	// Callers are responsible for running the polling loop. This configuration does not apply to streaming requests.
+	Polling bool
 }
 
 // Client represents a transport-agnostic implementation of A2A client.
@@ -94,7 +99,7 @@ func (c *Client) CancelTask(ctx context.Context, id *a2a.TaskIDParams) (*a2a.Tas
 func (c *Client) SendMessage(ctx context.Context, message *a2a.MessageSendParams) (a2a.SendMessageResult, error) {
 	method := "SendMessage"
 
-	message = c.withDefaultSendConfig(message)
+	message = c.withDefaultSendConfig(message, blocking(!c.config.Polling))
 
 	ctx, err := c.interceptBefore(ctx, method, message)
 	if err != nil {
@@ -113,7 +118,7 @@ func (c *Client) SendStreamingMessage(ctx context.Context, message *a2a.MessageS
 	return func(yield func(a2a.Event, error) bool) {
 		method := "SendStreamingMessage"
 
-		message = c.withDefaultSendConfig(message)
+		message = c.withDefaultSendConfig(message, blocking(true))
 
 		ctx, err := c.interceptBefore(ctx, method, message)
 		if err != nil {
@@ -269,8 +274,10 @@ func (c *Client) Destroy() error {
 	return c.transport.Destroy()
 }
 
-func (c *Client) withDefaultSendConfig(message *a2a.MessageSendParams) *a2a.MessageSendParams {
-	if c.config.PushConfig == nil && c.config.AcceptedOutputModes == nil {
+type blocking bool
+
+func (c *Client) withDefaultSendConfig(message *a2a.MessageSendParams, blocking blocking) *a2a.MessageSendParams {
+	if c.config.PushConfig == nil && c.config.AcceptedOutputModes == nil && blocking {
 		return message
 	}
 	result := *message
@@ -286,6 +293,7 @@ func (c *Client) withDefaultSendConfig(message *a2a.MessageSendParams) *a2a.Mess
 	if result.Config.AcceptedOutputModes == nil {
 		result.Config.AcceptedOutputModes = c.config.AcceptedOutputModes
 	}
+	result.Config.Blocking = utils.Ptr(bool(blocking))
 	return &result
 }
 

a2aclient/client_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 
 	"github.com/a2aproject/a2a-go/a2a"
+	"github.com/a2aproject/a2a-go/internal/utils"
 	"github.com/google/go-cmp/cmp"
 )
 
@@ -175,7 +176,7 @@ func TestClient_DefaultSendMessageConfig(t *testing.T) {
 	}
 	interceptor := &testInterceptor{}
 	client := &Client{
-		config:       Config{PushConfig: pushConfig, AcceptedOutputModes: acceptedModes},
+		config:       Config{PushConfig: pushConfig, AcceptedOutputModes: acceptedModes, Polling: true},
 		transport:    transport,
 		interceptors: []CallInterceptor{interceptor},
 	}
@@ -187,7 +188,7 @@ func TestClient_DefaultSendMessageConfig(t *testing.T) {
 			t.Fatalf("client.SendMessage() error = %v", err)
 		}
 		want := &a2a.MessageSendParams{
-			Config: &a2a.MessageSendConfig{AcceptedOutputModes: acceptedModes, PushConfig: pushConfig},
+			Config: &a2a.MessageSendConfig{AcceptedOutputModes: acceptedModes, PushConfig: pushConfig, Blocking: utils.Ptr(false)},
 		}
 		if diff := cmp.Diff(want, interceptor.lastReq.Payload); diff != "" {
 			t.Fatalf("client.SendMessage() wrong result (+got,-want) diff = %s", diff)
@@ -214,7 +215,7 @@ func TestClient_DefaultSendStreamingMessageConfig(t *testing.T) {
 	}
 	interceptor := &testInterceptor{}
 	client := &Client{
-		config:       Config{PushConfig: pushConfig, AcceptedOutputModes: acceptedModes},
+		config:       Config{PushConfig: pushConfig, AcceptedOutputModes: acceptedModes, Polling: true},
 		transport:    transport,
 		interceptors: []CallInterceptor{interceptor},
 	}
@@ -225,7 +226,7 @@ func TestClient_DefaultSendStreamingMessageConfig(t *testing.T) {
 		}
 	}
 	want := &a2a.MessageSendParams{
-		Config: &a2a.MessageSendConfig{AcceptedOutputModes: acceptedModes, PushConfig: pushConfig},
+		Config: &a2a.MessageSendConfig{AcceptedOutputModes: acceptedModes, PushConfig: pushConfig, Blocking: utils.Ptr(true)},
 	}
 	if diff := cmp.Diff(want, interceptor.lastReq.Payload); diff != "" {
 		t.Fatalf("client.SendStreamingMessage() wrong result (+got,-want) diff = %s", diff)

a2apb/pbconv/from_proto.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/a2aproject/a2a-go/a2a"
 	"github.com/a2aproject/a2a-go/a2apb"
+	"google.golang.org/protobuf/proto"
 	structpb "google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -162,7 +163,7 @@ func fromProtoSendMessageConfig(conf *a2apb.SendMessageConfiguration) (*a2a.Mess
 
 	result := &a2a.MessageSendConfig{
 		AcceptedOutputModes: conf.GetAcceptedOutputModes(),
-		Blocking:            conf.GetBlocking(),
+		Blocking:            proto.Bool(conf.GetBlocking()),
 		PushConfig:          pConf,
 	}
 

a2apb/pbconv/from_proto_test.go

@@ -20,6 +20,8 @@ import (
 
 	"github.com/a2aproject/a2a-go/a2a"
 	"github.com/a2aproject/a2a-go/a2apb"
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -187,7 +189,7 @@ func TestFromProto_fromProtoSendMessageConfig(t *testing.T) {
 			},
 			want: &a2a.MessageSendConfig{
 				AcceptedOutputModes: []string{"text/plain"},
-				Blocking:            true,
+				Blocking:            proto.Bool(true),
 				HistoryLength:       &a2aHistoryLen,
 				PushConfig: &a2a.PushConfig{
 					ID:    "test-push-config",
@@ -205,14 +207,14 @@ func TestFromProto_fromProtoSendMessageConfig(t *testing.T) {
 			in: &a2apb.SendMessageConfiguration{
 				HistoryLength: 0,
 			},
-			want: &a2a.MessageSendConfig{},
+			want: &a2a.MessageSendConfig{Blocking: proto.Bool(false)},
 		},
 		{
 			name: "config with no push notification",
 			in: &a2apb.SendMessageConfiguration{
 				PushNotification: nil,
 			},
-			want: &a2a.MessageSendConfig{},
+			want: &a2a.MessageSendConfig{Blocking: proto.Bool(false)},
 		},
 		{
 			name: "nil config",
@@ -221,13 +223,15 @@ func TestFromProto_fromProtoSendMessageConfig(t *testing.T) {
 		},
 	}
 	for _, tt := range tests {
-		got, err := fromProtoSendMessageConfig(tt.in)
-		if (err != nil) != tt.wantErr {
-			t.Fatalf("fromProtoSendMessageConfig() error = %v", err)
-		}
-		if !reflect.DeepEqual(got, tt.want) {
-			t.Errorf("fromProtoSendMessageConfig() got = %v, want %v", got, tt.want)
-		}
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := fromProtoSendMessageConfig(tt.in)
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("fromProtoSendMessageConfig() error = %v", err)
+			}
+			if diff := cmp.Diff(tt.want, got); diff != "" {
+				t.Errorf("fromProtoSendMessageConfig() wrong result (+got,-want)\ngot = %v\n want %v\ndiff = %s", got, tt.want, diff)
+			}
+		})
 	}
 }
 
@@ -259,7 +263,7 @@ func TestFromProto_fromProtoSendMessageRequest(t *testing.T) {
 		},
 	}
 	a2aConf := &a2a.MessageSendConfig{
-		Blocking:      true,
+		Blocking:      proto.Bool(true),
 		HistoryLength: &a2aHistoryLen,
 		PushConfig: &a2a.PushConfig{
 			ID:    "push-config",
@@ -324,7 +328,7 @@ func TestFromProto_fromProtoSendMessageRequest(t *testing.T) {
 			},
 			want: &a2a.MessageSendParams{
 				Message: &a2aMsg,
-				Config:  &a2a.MessageSendConfig{PushConfig: &a2a.PushConfig{}},
+				Config:  &a2a.MessageSendConfig{PushConfig: &a2a.PushConfig{}, Blocking: proto.Bool(false)},
 			},
 		},
 	}

a2apb/pbconv/to_proto.go

@@ -93,9 +93,11 @@ func toProtoSendMessageConfig(config *a2a.MessageSendConfig) (*a2apb.SendMessage
 
 	pConf := &a2apb.SendMessageConfiguration{
 		AcceptedOutputModes: config.AcceptedOutputModes,
-		Blocking:            config.Blocking,
 		PushNotification:    pushConf,
 	}
+	if config.Blocking != nil {
+		pConf.Blocking = *config.Blocking
+	}
 	if config.HistoryLength != nil {
 		pConf.HistoryLength = int32(*config.HistoryLength)
 	}

a2asrv/handler.go

@@ -214,10 +214,11 @@ func (h *defaultRequestHandler) OnSendMessage(ctx context.Context, params *a2a.M
 		if err != nil {
 			return nil, err
 		}
-		if taskID, required := isAuthRequired(event); required {
+
+		if taskID, interrupt := shouldInterruptNonStreaming(params, event); interrupt {
 			task, err := h.taskStore.Get(ctx, taskID)
 			if err != nil {
-				return nil, fmt.Errorf("failed to load task in auth-required state: %w", err)
+				return nil, fmt.Errorf("failed to load task on event processing interrupt: %w", err)
 			}
 			return task, nil
 		}
@@ -349,12 +350,23 @@ func (h *defaultRequestHandler) OnGetExtendedAgentCard(ctx context.Context) (*a2
 	return h.authenticatedCardProducer.Card(ctx)
 }
 
-func isAuthRequired(event a2a.Event) (a2a.TaskID, bool) {
+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) {
+		if _, ok := event.(*a2a.Message); ok {
+			return "", false
+		}
+		taskInfo := event.TaskInfo()
+		return taskInfo.TaskID, true
+	}
+
+	// Non-streaming clients need to be notified when auth is required
 	switch v := event.(type) {
 	case *a2a.Task:
 		return v.ID, v.Status.State == a2a.TaskStateAuthRequired
 	case *a2a.TaskStatusUpdateEvent:
 		return v.TaskID, v.Status.State == a2a.TaskStateAuthRequired
 	}
+
 	return "", false
 }