Add files via upload

internal/security/shell_execute_stream_test.go

@@ -0,0 +1,117 @@
+package security
+
+import (
+	"context"
+	"errors"
+	"io"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/cloudwego/eino/adk/filesystem"
+)
+
+func TestEinoStreamingShell_StreamsStderrBeforeStdoutEOF(t *testing.T) {
+	shell := NewEinoStreamingShell()
+	cmd := PrepareNonInteractiveShellCommand("echo err-only >&2; exit 1")
+	sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd})
+	if err != nil {
+		t.Fatalf("ExecuteStreaming: %v", err)
+	}
+	defer sr.Close()
+
+	start := time.Now()
+	var got strings.Builder
+	for {
+		resp, rerr := sr.Recv()
+		if errors.Is(rerr, io.EOF) {
+			break
+		}
+		if rerr != nil {
+			t.Fatalf("recv: %v", rerr)
+		}
+		if resp != nil && resp.Output != "" {
+			got.WriteString(resp.Output)
+		}
+	}
+	if time.Since(start) > 3*time.Second {
+		t.Fatalf("expected fast completion, took %v", time.Since(start))
+	}
+	if !strings.Contains(got.String(), "err-only") {
+		t.Fatalf("expected stderr in output, got: %q", got.String())
+	}
+}
+
+func TestEinoStreamingShell_SudoFailsFast(t *testing.T) {
+	shell := NewEinoStreamingShell()
+	cmd := PrepareNonInteractiveShellCommand("sudo whoami && sudo cat /etc/os-release")
+	sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd})
+	if err != nil {
+		t.Fatalf("ExecuteStreaming: %v", err)
+	}
+	defer sr.Close()
+
+	start := time.Now()
+	var got strings.Builder
+	for {
+		resp, rerr := sr.Recv()
+		if errors.Is(rerr, io.EOF) {
+			break
+		}
+		if rerr != nil {
+			t.Fatalf("recv: %v", rerr)
+		}
+		if resp == nil {
+			continue
+		}
+		got.WriteString(resp.Output)
+	}
+	if time.Since(start) > 5*time.Second {
+		t.Fatalf("sudo should fail quickly, took %v output=%q", time.Since(start), got.String())
+	}
+	out := got.String()
+	if strings.Contains(out, "command exited with non-zero code") {
+		t.Fatalf("legacy exit line present: %q", out)
+	}
+	if !strings.Contains(out, "sudo") && !strings.Contains(out, "password") && !strings.Contains(out, "terminal") {
+		t.Fatalf("expected sudo error text, got: %q", out)
+	}
+}
+
+func TestEinoStreamingShell_StderrWhileStdoutBlocks(t *testing.T) {
+	shell := NewEinoStreamingShell()
+	// 模拟 sudo：stderr 先有输出，stdout 侧进程仍挂起；旧 eino local 在首包 stderr 前不会向流写任何内容。
+	cmd := PrepareNonInteractiveShellCommand(`echo "password prompt" >&2; sleep 30`)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+
+	sr, err := shell.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: cmd})
+	if err != nil {
+		t.Fatalf("ExecuteStreaming: %v", err)
+	}
+	defer sr.Close()
+
+	start := time.Now()
+	var got strings.Builder
+	for {
+		resp, rerr := sr.Recv()
+		if errors.Is(rerr, io.EOF) {
+			break
+		}
+		if rerr != nil {
+			break
+		}
+		if resp != nil && resp.Output != "" {
+			got.WriteString(resp.Output)
+			if strings.Contains(got.String(), "password prompt") {
+				break
+			}
+		}
+	}
+	if time.Since(start) > 1500*time.Millisecond {
+		t.Fatalf("expected stderr promptly, took %v output=%q", time.Since(start), got.String())
+	}
+	if !strings.Contains(got.String(), "password prompt") {
+		t.Fatalf("expected early stderr, got: %q", got.String())
+	}
+}