feat: implement get --wait support for polling-simulated blocking

internal/README.md

@@ -168,11 +168,19 @@ a2a send -a <url> --context-id <context-id> "Related question"
 a2a task get -a <url> <id>
 a2a task get -a <url> <id> --history 10
 a2a task get -a <url> <id> -o json
+
+# Follow a task started with `send --async` until it finishes
+a2a task get -a <url> <id> --wait
+a2a task get -a <url> <id> --wait --poll-interval 2s --timeout 60s
 ```
 
 | Flag | Description |
 |---|---|
 | `--history <n>` | Include up to `n` history messages. |
+| `--wait` | Poll until the task reaches a terminal (`completed`/`failed`/`canceled`/`rejected`) state, returning early on an interrupted (`input-required`/`auth-required`) state so the caller can act. |
+| `--poll-interval <duration>` | Delay between polls while waiting (default `2s`); only used with `--wait`. |
+
+The overall wait budget is the global `--timeout` (default `30s`); when it expires before the task settles, the command reports a timeout error and exits non-zero.
 
 ### `task list` - List Tasks
 

internal/cli/cli_test.go

@@ -418,7 +418,7 @@ func partsFromArgs(t *testing.T, args ...string) *sendFlags {
 func TestSendStreaming(t *testing.T) {
 	t.Parallel()
 	url := startTestServer(t)
-	nonStreamingServerURL := startTestServerWith(t, a2a.AgentCapabilities{Streaming: false})
+	nonStreamingServerURL := startTestServerWith(t, a2a.AgentCapabilities{Streaming: false}, localsrv.NewEchoExecutor())
 
 	testCases := []struct {
 		name               string
@@ -547,7 +547,7 @@ func TestSendOutputInvalidFormat(t *testing.T) {
 
 func TestSendStreamingFallbackUsesDefaultPoller(t *testing.T) {
 	t.Parallel()
-	nonStreamingURL := startTestServerWith(t, a2a.AgentCapabilities{Streaming: false})
+	nonStreamingURL := startTestServerWith(t, a2a.AgentCapabilities{Streaming: false}, localsrv.NewEchoExecutor())
 
 	out, err := runCMDWithConfig(t, deps{cfgLoader: clicfg.LoadEmpty},
 		"send", "-a", nonStreamingURL, "-o", "json", "--stream", "stream me", "--poll-interval", "5ms")
@@ -707,6 +707,21 @@ func TestGetTask(t *testing.T) {
 		}
 	})
 
+	t.Run("get task with --wait polls to terminal state", func(t *testing.T) {
+		t.Parallel()
+		out := mustRunCMD(t, "task", "get", "-a", url, string(taskID), "--wait", "--poll-interval", "5ms", "-o", "json")
+		var task a2a.Task
+		if err := json.Unmarshal([]byte(out), &task); err != nil {
+			t.Fatalf("json.Unmarshal(task get --wait output) error = %v", err)
+		}
+		if task.ID != taskID {
+			t.Fatalf("a2a task get --wait ID = %q, want %q", task.ID, taskID)
+		}
+		if task.Status.State != a2a.TaskStateCompleted {
+			t.Fatalf("a2a task get --wait Status.State = %q, want %q", task.Status.State, a2a.TaskStateCompleted)
+		}
+	})
+
 	t.Run("missing args fails", func(t *testing.T) {
 		t.Parallel()
 		if _, err := runCMD(t, "task", "get", "-a", url); err == nil {
@@ -715,6 +730,30 @@ func TestGetTask(t *testing.T) {
 	})
 }
 
+func TestGetTaskWait_Timeout(t *testing.T) {
+	t.Parallel()
+	url := startTestServerWith(t, a2a.AgentCapabilities{Streaming: false},
+		a2asrv.AgentExecutorFunc(func(ctx context.Context, ec *a2asrv.ExecutorContext) iter.Seq2[a2a.Event, error] {
+			return func(yield func(a2a.Event, error) bool) {
+				if ec.StoredTask == nil {
+					if !yield(a2a.NewSubmittedTask(ec, ec.Message), nil) {
+						return
+					}
+				}
+				<-ctx.Done()
+			}
+		}),
+	)
+	taskID := sendTestMessageWithConfig(t, url, &a2a.SendMessageConfig{ReturnImmediately: true}, "hello")
+	_, err := runCMD(t, "task", "get", "-a", url, string(taskID), "--wait", "--poll-interval", "1ms", "--timeout", "5ms")
+	if err == nil {
+		t.Fatal("a2a task get --wait against a never-terminal task should time out")
+	}
+	if !strings.Contains(err.Error(), "timed out") {
+		t.Fatalf("a2a task get --wait error = %v, want a timeout error", err)
+	}
+}
+
 func TestServe_ModeValidation(t *testing.T) {
 	t.Parallel()
 	for _, tt := range []struct {
@@ -801,13 +840,13 @@ func TestConfigApplied(t *testing.T) {
 
 func startTestServer(t *testing.T) string {
 	t.Helper()
-	return startTestServerWith(t, a2a.AgentCapabilities{Streaming: true})
+	return startTestServerWith(t, a2a.AgentCapabilities{Streaming: true}, localsrv.NewEchoExecutor())
 }
 
-func startTestServerWith(t *testing.T, capabilities a2a.AgentCapabilities) string {
+func startTestServerWith(t *testing.T, capabilities a2a.AgentCapabilities, executor a2asrv.AgentExecutor) string {
 	t.Helper()
 
-	handler := a2asrv.NewHandler(localsrv.NewEchoExecutor(), a2asrv.WithCapabilityChecks(&capabilities))
+	handler := a2asrv.NewHandler(executor, a2asrv.WithCapabilityChecks(&capabilities))
 
 	mux := http.NewServeMux()
 	mux.Handle("/", a2asrv.NewRESTHandler(handler))
@@ -866,6 +905,11 @@ func mustDecodeTask(t *testing.T, out string) *a2a.Task {
 }
 
 func sendTestMessage(t *testing.T, url, text string) a2a.TaskID {
+	t.Helper()
+	return sendTestMessageWithConfig(t, url, nil, text)
+}
+
+func sendTestMessageWithConfig(t *testing.T, url string, config *a2a.SendMessageConfig, text string) a2a.TaskID {
 	t.Helper()
 	ctx := t.Context()
 
@@ -878,7 +922,7 @@ func sendTestMessage(t *testing.T, url, text string) a2a.TaskID {
 	defer func() { _ = client.Destroy() }()
 
 	msg := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart(text))
-	result, err := client.SendMessage(ctx, &a2a.SendMessageRequest{Message: msg})
+	result, err := client.SendMessage(ctx, &a2a.SendMessageRequest{Message: msg, Config: config})
 	if err != nil {
 		t.Fatalf("client.SendMessage() error = %v", err)
 	}

internal/cli/task_get.go

@@ -16,15 +16,21 @@ package cli
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"time"
 
 	"github.com/spf13/cobra"
 
+	"github.com/a2aproject/a2a-cli/internal/polling"
 	"github.com/a2aproject/a2a-go/v2/a2a"
+	"github.com/a2aproject/a2a-go/v2/a2aclient"
 )
 
 func newTaskGetCmd(cfg *globalConfig) *cobra.Command {
 	var history int
+	var wait bool
+	var pollInterval time.Duration
 
 	cmd := &cobra.Command{
 		Use:   "get <task-id>",
@@ -49,9 +55,9 @@ func newTaskGetCmd(cfg *globalConfig) *cobra.Command {
 				req.HistoryLength = &history
 			}
 
-			task, err := client.GetTask(ctx, req)
+			task, err := getTask(ctx, cfg, client, req, wait, pollInterval)
 			if err != nil {
-				return fmt.Errorf("failed to get task %s: %w", args[0], err)
+				return err
 			}
 
 			if err := cfg.PrintTask(task); err != nil {
@@ -61,6 +67,30 @@ func newTaskGetCmd(cfg *globalConfig) *cobra.Command {
 		},
 	}
 
-	cmd.Flags().IntVar(&history, "history", 0, "Include up to n history messages")
+	f := cmd.Flags()
+	f.IntVar(&history, "history", 0, "Include up to n history messages")
+	f.BoolVar(&wait, "wait", false, "Poll until the task reaches a terminal or interrupted (input/auth-required) state")
+	f.DurationVar(&pollInterval, "poll-interval", 2*time.Second, "Duration between polls while waiting; only used with --wait. Overall wait budget is --timeout.")
 	return cmd
 }
+
+func getTask(ctx context.Context, cfg *globalConfig, client *a2aclient.Client, req *a2a.GetTaskRequest, wait bool, interval time.Duration) (*a2a.Task, error) {
+	if !wait {
+		task, err := client.GetTask(ctx, req)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get task %s: %w", req.ID, err)
+		}
+		return task, nil
+	}
+
+	cfg.logf("waiting for task %s (poll interval %v, timeout %v)", req.ID, interval, cfg.timeout)
+
+	task, err := polling.WaitForTask(ctx, client, req, interval)
+	if err != nil {
+		if errors.Is(err, context.DeadlineExceeded) {
+			return nil, fmt.Errorf("timed out after %s waiting for task %s: %w", cfg.timeout, req.ID, err)
+		}
+		return nil, fmt.Errorf("failed to wait for task %s: %w", req.ID, err)
+	}
+	return task, nil
+}

internal/polling/polling.go

@@ -27,6 +27,42 @@ import (
 	"github.com/a2aproject/a2a-go/v2/a2aevent"
 )
 
+// maxSuccessiveFailures is the number of consecutive GetTask failures tolerated
+// before polling gives up.
+const maxSuccessiveFailures = 3
+
+// WaitForTask polls the task identified by req at the given interval until it
+// reaches a terminal state (completed/failed/canceled/rejected) or an
+// interrupted state that needs the caller to act (input-required/auth-required),
+// returning the final task. The first poll happens immediately.
+func WaitForTask(ctx context.Context, client *a2aclient.Client, req *a2a.GetTaskRequest, interval time.Duration) (*a2a.Task, error) {
+	successiveFailures := 0
+	for {
+		task, err := client.GetTask(ctx, req)
+		if err != nil {
+			if ctx.Err() != nil {
+				return nil, ctx.Err()
+			}
+			successiveFailures++
+			if successiveFailures >= maxSuccessiveFailures {
+				return nil, fmt.Errorf("successive polling failure threshold exceeded for task %q: %w", req.ID, err)
+			}
+		} else {
+			successiveFailures = 0
+			if task.Status.State.Terminal() ||
+				task.Status.State == a2a.TaskStateInputRequired ||
+				task.Status.State == a2a.TaskStateAuthRequired {
+				return task, nil
+			}
+		}
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(interval):
+		}
+	}
+}
+
 // Stream sends the original message and then polls the resulting task at the
 // given interval, yielding task events until it reaches a terminal state.
 func Stream(ctx context.Context, client *a2aclient.Client, original *a2a.SendMessageRequest, interval time.Duration) iter.Seq2[a2a.Event, error] {
@@ -70,7 +106,7 @@ func Stream(ctx context.Context, client *a2aclient.Client, original *a2a.SendMes
 			task, err := client.GetTask(ctx, &a2a.GetTaskRequest{ID: tid, Tenant: req.Tenant})
 			if err != nil {
 				successiveFailures++
-				if successiveFailures == 3 {
+				if successiveFailures >= maxSuccessiveFailures {
 					yield(nil, fmt.Errorf("successive polling failure threshold exceeded for task %q", tid))
 					return
 				}

internal/polling/polling_test.go

@@ -187,6 +187,114 @@ func TestHandlePolling(t *testing.T) {
 	}
 }
 
+func TestWaitForTask(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name         string
+		getResponses []getTaskResponse
+		wantState    a2a.TaskState
+		wantErr      string
+	}{
+		{
+			name:         "returns immediately when already terminal",
+			getResponses: []getTaskResponse{{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}}}},
+			wantState:    a2a.TaskStateCompleted,
+		},
+		{
+			name: "polls until completion",
+			getResponses: []getTaskResponse{
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateSubmitted}}},
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateWorking}}},
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}}},
+			},
+			wantState: a2a.TaskStateCompleted,
+		},
+		{
+			name: "stops on input-required",
+			getResponses: []getTaskResponse{
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateWorking}}},
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateInputRequired}}},
+			},
+			wantState: a2a.TaskStateInputRequired,
+		},
+		{
+			name: "stops on auth-required",
+			getResponses: []getTaskResponse{
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateWorking}}},
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateAuthRequired}}},
+			},
+			wantState: a2a.TaskStateAuthRequired,
+		},
+		{
+			name: "tolerates transient failures below threshold",
+			getResponses: []getTaskResponse{
+				{err: errors.New("temporary 1")},
+				{err: errors.New("temporary 2")},
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateWorking}}},
+				{err: errors.New("temporary 3")},
+				{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}}},
+			},
+			wantState: a2a.TaskStateCompleted,
+		},
+		{
+			name: "successive failures exceed threshold",
+			getResponses: []getTaskResponse{
+				{err: errors.New("temporary 1")},
+				{err: errors.New("temporary 2")},
+				{err: errors.New("temporary 3")},
+			},
+			wantErr: "successive polling failure threshold exceeded",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			transport := &fakePollingTransport{getResponses: tt.getResponses}
+			client := newPollingClient(t, transport)
+			req := &a2a.GetTaskRequest{ID: "task-1"}
+
+			task, err := WaitForTask(t.Context(), client, req, 0)
+
+			if tt.wantErr != "" {
+				if err == nil {
+					t.Fatalf("WaitForTask() error = nil, want error containing %q", tt.wantErr)
+				}
+				if !strings.Contains(err.Error(), tt.wantErr) {
+					t.Fatalf("WaitForTask() error = %v, want error containing %q", err, tt.wantErr)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("WaitForTask() error = %v, want nil", err)
+			}
+			if task.Status.State != tt.wantState {
+				t.Fatalf("WaitForTask() state = %v, want %v", task.Status.State, tt.wantState)
+			}
+		})
+	}
+}
+
+func TestWaitForTask_CancelledWithContext(t *testing.T) {
+	t.Parallel()
+
+	transport := &fakePollingTransport{
+		getResponses: []getTaskResponse{{task: &a2a.Task{Status: a2a.TaskStatus{State: a2a.TaskStateWorking}}}},
+	}
+	client := newPollingClient(t, transport)
+	req := &a2a.GetTaskRequest{ID: "task-1"}
+
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	_, err := WaitForTask(ctx, client, req, time.Hour)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("WaitForTask() error = %v, want context.Canceled", err)
+	}
+}
+
 func TestStream_SleepCancelledWithContext(t *testing.T) {
 	t.Parallel()
 
@@ -253,6 +361,7 @@ type fakePollingTransport struct {
 	sendErr    error
 
 	getResponses []getTaskResponse
+	getAlways    *a2a.Task
 	getCalls     int
 
 	sendRequest *a2a.SendMessageRequest
@@ -267,6 +376,9 @@ func (f *fakePollingTransport) SendMessage(ctx context.Context, c a2aclient.Serv
 }
 
 func (f *fakePollingTransport) GetTask(context.Context, a2aclient.ServiceParams, *a2a.GetTaskRequest) (*a2a.Task, error) {
+	if f.getAlways != nil {
+		return f.getAlways, nil
+	}
 	i := f.getCalls
 	f.getCalls++
 	if i >= len(f.getResponses) {