Print each SSE data object on its own line in invoke streaming output

cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go

@@ -633,7 +633,7 @@ func handleInvocationResponse(
 
 	contentType := resp.Header.Get("Content-Type")
 	if strings.HasPrefix(contentType, "text/event-stream") {
-		return handleInvocationSSE(resp.Body, agentName)
+		return handleInvocationSSE(os.Stdout, resp.Body, agentName)
 	}
 
 	return handleInvocationSync(resp.Body, agentName)
@@ -679,7 +679,7 @@ func handleInvocationSync(body io.Reader, agentName string) error {
 
 // handleInvocationSSE handles a streaming (200 OK, text/event-stream) invocations response.
 // The invocations protocol has a developer-defined SSE format, so we print data lines as they arrive.
-func handleInvocationSSE(body io.Reader, agentName string) error {
+func handleInvocationSSE(w io.Writer, body io.Reader, agentName string) error {
 	scanner := bufio.NewScanner(body)
 	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
 
@@ -702,9 +702,6 @@ func handleInvocationSSE(body io.Reader, agentName string) error {
 				} `json:"error"`
 			}
 			if json.Unmarshal([]byte(data), &errEnvelope) == nil && errEnvelope.Error.Message != "" {
-				if printed {
-					fmt.Println()
-				}
 				label := errEnvelope.Error.Code
 				if label == "" {
 					label = errEnvelope.Error.Type
@@ -715,22 +712,19 @@ func handleInvocationSSE(body io.Reader, agentName string) error {
 				return fmt.Errorf("agent error: %s", errEnvelope.Error.Message)
 			}
 
-			// Print data as-is
+			// Print data as-is, one line per SSE data object
 			if !printed {
-				fmt.Printf("[%s] ", agentName)
+				fmt.Fprintf(w, "[%s] ", agentName)
 				printed = true
 			}
-			fmt.Print(data)
+			fmt.Fprintln(w, data)
 		}
 	}
 
 	if err := scanner.Err(); err != nil {
 		return fmt.Errorf("error reading response stream: %w", err)
 	}
 
-	if printed {
-		fmt.Println()
-	}
 	return nil
 }
 

cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_test.go

@@ -4,6 +4,7 @@
 package cmd
 
 import (
+	"bytes"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -319,51 +320,72 @@ func TestHandleInvocationSSE(t *testing.T) {
 	t.Parallel()
 
 	tests := []struct {
-		name    string
-		input   string
-		wantErr bool
-		errMsg  string
+		name       string
+		input      string
+		agentName  string
+		wantErr    bool
+		errMsg     string
+		wantOutput string
 	}{
 		{
-			name:    "simple data lines",
-			input:   "data: Hello \ndata: world!\n\n",
-			wantErr: false,
+			name:       "simple data lines produce separate output lines with prefix",
+			input:      "data: Hello \ndata: world!\n\n",
+			agentName:  "test-agent",
+			wantOutput: "[test-agent] Hello \nworld!\n",
 		},
 		{
-			name:    "DONE signal ends stream",
-			input:   "data: Hello\ndata: [DONE]\ndata: ignored\n\n",
-			wantErr: false,
+			name:       "single data line gets prefix and newline",
+			input:      "data: only-one\n\n",
+			agentName:  "my-bot",
+			wantOutput: "[my-bot] only-one\n",
 		},
 		{
-			name:    "error envelope in data",
-			input:   `data: {"error": {"code": "rate_limit", "message": "too many requests"}}` + "\n\n",
-			wantErr: true,
-			errMsg:  "agent error (rate_limit): too many requests",
+			name:       "DONE signal ends stream, only preceding data printed",
+			input:      "data: Hello\ndata: [DONE]\ndata: ignored\n\n",
+			agentName:  "test-agent",
+			wantOutput: "[test-agent] Hello\n",
 		},
 		{
-			name:    "error envelope with type only",
-			input:   `data: {"error": {"type": "server_error", "message": "crash"}}` + "\n\n",
-			wantErr: true,
-			errMsg:  "agent error (server_error): crash",
+			name:      "error envelope in data",
+			input:     `data: {"error": {"code": "rate_limit", "message": "too many requests"}}` + "\n\n",
+			agentName: "test-agent",
+			wantErr:   true,
+			errMsg:    "agent error (rate_limit): too many requests",
 		},
 		{
-			name:    "empty stream",
-			input:   "",
-			wantErr: false,
+			name:      "error envelope with type only",
+			input:     `data: {"error": {"type": "server_error", "message": "crash"}}` + "\n\n",
+			agentName: "test-agent",
+			wantErr:   true,
+			errMsg:    "agent error (server_error): crash",
 		},
 		{
-			name:    "non-data lines ignored",
-			input:   "event: custom\nid: 123\ndata: content\n\n",
-			wantErr: false,
+			name:       "empty stream produces no output",
+			input:      "",
+			agentName:  "test-agent",
+			wantOutput: "",
+		},
+		{
+			name:       "non-data lines ignored",
+			input:      "event: custom\nid: 123\ndata: content\n\n",
+			agentName:  "test-agent",
+			wantOutput: "[test-agent] content\n",
+		},
+		{
+			name:       "three data lines produce three output lines",
+			input:      "data: line1\ndata: line2\ndata: line3\n\n",
+			agentName:  "agent",
+			wantOutput: "[agent] line1\nline2\nline3\n",
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
+			var buf bytes.Buffer
 			reader := strings.NewReader(tt.input)
-			err := handleInvocationSSE(reader, "test-agent")
+			err := handleInvocationSSE(&buf, reader, tt.agentName)
 
 			if tt.wantErr {
 				if err == nil {
@@ -372,8 +394,13 @@ func TestHandleInvocationSSE(t *testing.T) {
 				if tt.errMsg != "" && err.Error() != tt.errMsg {
 					t.Errorf("error = %q, want %q", err.Error(), tt.errMsg)
 				}
-			} else if err != nil {
-				t.Fatalf("unexpected error: %v", err)
+			} else {
+				if err != nil {
+					t.Fatalf("unexpected error: %v", err)
+				}
+				if got := buf.String(); got != tt.wantOutput {
+					t.Errorf("output mismatch\ngot:  %q\nwant: %q", got, tt.wantOutput)
+				}
 			}
 		})
 	}