From 99dd08b29057e9042242fe90ad14d2da43a91831 Mon Sep 17 00:00:00 2001 From: Yaroslav Shevchuk Date: Wed, 14 Oct 2026 15:23:36 +0000 Subject: [PATCH] fix printing --- internal/cli/send.go | 1 + internal/cli/task_subscribe.go | 1 + internal/output/output.go | 118 +++++++++++++++++++++++++-------- internal/output/output_test.go | 99 +++++++++++++++++++++++++++ 4 files changed, 193 insertions(+), 26 deletions(-) diff --git a/internal/cli/send.go b/internal/cli/send.go index ac96e86..5af9e64 100644 --- a/internal/cli/send.go +++ b/internal/cli/send.go @@ -72,6 +72,7 @@ func newSendCmd(cfg *globalConfig, poller pollerFunc) *cobra.Command { defer destroyClient(cfg, client) if flags.stream { + defer func() { _ = cfg.Printer.Flush() }() ctx, debounceTimeout := utils.WithInactivityTimeout(baseCtx, cfg.timeout) err = handleStreaming(ctx, cfg, client, req, debounceTimeout) diff --git a/internal/cli/task_subscribe.go b/internal/cli/task_subscribe.go index b955b05..20cd541 100644 --- a/internal/cli/task_subscribe.go +++ b/internal/cli/task_subscribe.go @@ -40,6 +40,7 @@ func newTaskSubscribeCmd(cfg *globalConfig) *cobra.Command { defer destroyClient(cfg, client) cfg.logf("subscribing to task %s", args[0]) + defer func() { _ = cfg.Printer.Flush() }() for event, err := range client.SubscribeToTask(ctx, &a2a.SubscribeToTaskRequest{ ID: a2a.TaskID(args[0]), diff --git a/internal/output/output.go b/internal/output/output.go index 4b94a3b..4f1b3da 100644 --- a/internal/output/output.go +++ b/internal/output/output.go @@ -54,6 +54,12 @@ var taskStateNames = map[a2a.TaskState]string{ type Printer struct { Out io.Writer Mode Mode + + // openArtifact is the artifact whose streamed text is being written inline + // in text mode; its line is not terminated yet. + openArtifact a2a.ArtifactID + // atLineStart reports whether the inline text written so far ends with a newline. + atLineStart bool } // NewPrinter returns a Printer that writes to out using the given Mode. @@ -115,6 +121,13 @@ func (p *Printer) PrintEvent(event a2a.Event) error { return p.jsonEncoder().Encode(a2a.StreamResponse{Event: event}) } + if e, ok := event.(*a2a.TaskArtifactUpdateEvent); ok { + return p.printArtifactChunk(e) + } + if err := p.Flush(); err != nil { + return err + } + var s string switch e := event.(type) { case *a2a.TaskStatusUpdateEvent: @@ -124,13 +137,6 @@ func (p *Printer) PrintEvent(event a2a.Event) error { } else { s = fmt.Sprintf("[status] %s\n", state) } - case *a2a.TaskArtifactUpdateEvent: - text := partsText(e.Artifact.Parts) - if e.Append { - s = fmt.Sprintf("[artifact+] %s\n", text) - } else { - s = fmt.Sprintf("[artifact] %s\n", text) - } case *a2a.Task: s = formatTask(e) + formatResumeHint(e) case *a2a.Message: @@ -142,6 +148,49 @@ func (p *Printer) PrintEvent(event a2a.Event) error { return err } +// printArtifactChunk writes streamed artifact text inline: a new artifact +// starts a "[artifact] " line, and chunks appended to it continue that line +// until the last chunk or another event arrives. +func (p *Printer) printArtifactChunk(e *a2a.TaskArtifactUpdateEvent) error { + text := artifactText(e.Artifact.Parts) + continues := e.Append && e.Artifact.ID == p.openArtifact + if !continues { + if err := p.Flush(); err != nil { + return err + } + prefix := "[artifact] " + if e.Append { + prefix = "[artifact+] " + } + text = prefix + text + } + if _, err := io.WriteString(p.Out, text); err != nil { + return err + } + p.openArtifact = e.Artifact.ID + if text != "" { + p.atLineStart = strings.HasSuffix(text, "\n") + } + if e.LastChunk { + return p.Flush() + } + return nil +} + +// Flush terminates the line of an artifact streamed inline in text mode, if +// any. Call it when a stream of events ends. +func (p *Printer) Flush() error { + if p.openArtifact == "" { + return nil + } + p.openArtifact = "" + if p.atLineStart { + return nil + } + _, err := io.WriteString(p.Out, "\n") + return err +} + // PrintTaskList writes a list of tasks in the configured Mode. func (p *Printer) PrintTaskList(resp *a2a.ListTasksResponse) error { if p.IsJSON() { @@ -257,7 +306,7 @@ func formatTask(task *a2a.Task) string { if art.Name != "" { label = art.Name } - fmt.Fprintf(&sb, " [%s] %s\n", label, partsText(art.Parts)) + fmt.Fprintf(&sb, " [%s] %s\n", label, artifactText(art.Parts)) } } @@ -313,33 +362,50 @@ func MessageText(msg *a2a.Message) string { return partsText(msg.Parts) } +// partsText renders parts separated by spaces. func partsText(parts a2a.ContentParts) string { + return joinParts(parts, false) +} + +// artifactText renders artifact parts. Adjacent text parts are concatenated +// as is, because streamed chunks appended to an artifact are pieces of one text. +func artifactText(parts a2a.ContentParts) string { + return joinParts(parts, true) +} + +func joinParts(parts a2a.ContentParts, concatText bool) string { var sb strings.Builder for i, p := range parts { - if i > 0 { + if i > 0 && !(concatText && isText(parts[i-1]) && isText(p)) { sb.WriteString(" ") } - if t := p.Text(); t != "" { - sb.WriteString(t) - continue - } - if p.URL() != "" || p.Raw() != nil { - sb.WriteString(filePartText(p)) - continue - } - if p.Data() != nil { - b, err := json.Marshal(p.Data()) - if err != nil { - fmt.Fprintf(&sb, "[data: %v]", err) - continue - } - sb.WriteString(string(b)) - continue - } + sb.WriteString(partText(p)) } return sb.String() } +func isText(p *a2a.Part) bool { + _, ok := p.Content.(a2a.Text) + return ok +} + +func partText(p *a2a.Part) string { + if t := p.Text(); t != "" { + return t + } + if p.URL() != "" || p.Raw() != nil { + return filePartText(p) + } + if p.Data() != nil { + b, err := json.Marshal(p.Data()) + if err != nil { + return fmt.Sprintf("[data: %v]", err) + } + return string(b) + } + return "" +} + // filePartText renders a file part by name, media type and size or URL. func filePartText(p *a2a.Part) string { var sb strings.Builder diff --git a/internal/output/output_test.go b/internal/output/output_test.go index d41a1f8..bb13aa1 100644 --- a/internal/output/output_test.go +++ b/internal/output/output_test.go @@ -111,3 +111,102 @@ func TestMessageTextFileParts(t *testing.T) { }) } } + +func TestPrintEventStreamsArtifactTextInline(t *testing.T) { + t.Parallel() + + info := a2a.TaskInfo{TaskID: "t1", ContextID: "c1"} + chunk := func(id a2a.ArtifactID, text string, appendChunk, last bool) *a2a.TaskArtifactUpdateEvent { + e := a2a.NewArtifactUpdateEvent(info, id, a2a.NewTextPart(text)) + e.Append, e.LastChunk = appendChunk, last + return e + } + + tests := []struct { + name string + events []a2a.Event + want string + }{ + { + name: "appended chunks continue the line until a status update", + events: []a2a.Event{ + chunk("a1", "Hel", false, false), + chunk("a1", "lo wor", true, false), + chunk("a1", "ld", true, false), + a2a.NewStatusUpdateEvent(info, a2a.TaskStateCompleted, nil), + }, + want: "[artifact] Hello world\n[status] completed\n", + }, + { + name: "last chunk ends the line without doubling a trailing newline", + events: []a2a.Event{ + chunk("a1", "line one\n", false, false), + chunk("a1", "line two\n", true, true), + a2a.NewStatusUpdateEvent(info, a2a.TaskStateCompleted, nil), + }, + want: "[artifact] line one\nline two\n[status] completed\n", + }, + { + name: "a new artifact starts a new line", + events: []a2a.Event{ + chunk("a1", "first", false, false), + chunk("a2", "second", false, false), + }, + want: "[artifact] first\n[artifact] second\n", + }, + { + name: "an append to another artifact is labeled", + events: []a2a.Event{ + chunk("a1", "first", false, false), + chunk("a2", " more", true, false), + }, + want: "[artifact] first\n[artifact+] more\n", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + var buf bytes.Buffer + p := NewPrinter(&buf, ModeText) + for _, e := range tt.events { + if err := p.PrintEvent(e); err != nil { + t.Fatalf("PrintEvent() error = %v, want nil", err) + } + } + if err := p.Flush(); err != nil { + t.Fatalf("Flush() error = %v, want nil", err) + } + if diff := cmp.Diff(tt.want, buf.String()); diff != "" { + t.Fatalf("PrintEvent() wrong result (-want +got) diff = %s", diff) + } + }) + } +} + +func TestTaskArtifactTextConcatenatesChunks(t *testing.T) { + t.Parallel() + + task := &a2a.Task{ + ID: "t1", + Status: a2a.TaskStatus{State: a2a.TaskStateCompleted}, + Artifacts: []*a2a.Artifact{{ + ID: "a1", + Name: "response", + Parts: a2a.ContentParts{ + a2a.NewTextPart("Hel"), + a2a.NewTextPart("lo"), + a2a.NewDataPart(map[string]any{"k": "v"}), + }, + }}, + } + var buf bytes.Buffer + if err := NewPrinter(&buf, ModeText).PrintEvent(task); err != nil { + t.Fatalf("PrintEvent() error = %v, want nil", err) + } + want := "Task: t1\nStatus: completed\nArtifacts:\n [response] Hello {\"k\":\"v\"}\n" + if diff := cmp.Diff(want, buf.String()); diff != "" { + t.Fatalf("PrintEvent() wrong result (-want +got) diff = %s", diff) + } +}