diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go index 8fdb9ca2757..8692e8f63ef 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go +++ b/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,12 +712,12 @@ 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) } } @@ -728,9 +725,6 @@ func handleInvocationSSE(body io.Reader, agentName string) error { return fmt.Errorf("error reading response stream: %w", err) } - if printed { - fmt.Println() - } return nil } diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_test.go index 6d2168530f2..a15e9369342 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_test.go +++ b/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,42 +320,62 @@ 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", }, } @@ -362,8 +383,9 @@ func TestHandleInvocationSSE(t *testing.T) { 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) + } } }) }