diff --git a/internal/README.md b/internal/README.md index b3e2cca..b577054 100644 --- a/internal/README.md +++ b/internal/README.md @@ -168,11 +168,19 @@ a2a send -a --context-id "Related question" a2a task get -a a2a task get -a --history 10 a2a task get -a -o json + +# Follow a task started with `send --async` until it finishes +a2a task get -a --wait +a2a task get -a --wait --poll-interval 2s --timeout 60s ``` | Flag | Description | |---|---| | `--history ` | 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 ` | 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 diff --git a/internal/cli/cli_test.go b/internal/cli/cli_test.go index 4a00667..85a12ac 100644 --- a/internal/cli/cli_test.go +++ b/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) } diff --git a/internal/cli/task_get.go b/internal/cli/task_get.go index ee2aa4e..b086a84 100644 --- a/internal/cli/task_get.go +++ b/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 ", @@ -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 +} diff --git a/internal/polling/polling.go b/internal/polling/polling.go index fccc153..0c571b5 100644 --- a/internal/polling/polling.go +++ b/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 } diff --git a/internal/polling/polling_test.go b/internal/polling/polling_test.go index 112d6b1..d3d163a 100644 --- a/internal/polling/polling_test.go +++ b/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) {