diff --git a/sandbox/v2/docs/API.md b/sandbox/v2/docs/API.md index 2b7b711d..468270ec 100644 --- a/sandbox/v2/docs/API.md +++ b/sandbox/v2/docs/API.md @@ -519,6 +519,38 @@ fmt.Printf("exit=%d stdout=%s duration=%dms\n", result.ExitCode, string(result.Stdout), result.DurationMs) ``` +### Stream + +```go +func (h *Host) Stream(ctx context.Context, cmd string, args []string, opts ...HostExecOption) (*HostExecStream, error) +``` + +Runs a command on the Tai host and streams stdout/stderr in real time via HostExec gRPC +ExecStream. Returns a `HostExecStream` with separate channels for stdout and stderr. + +```go +host, _ := sandbox.M().Host(ctx, "remote") +stream, err := host.Stream(ctx, "tail", []string{"-f", "/var/log/app.log"}, + sandbox.WithHostWorkDir("/data"), + sandbox.WithHostTimeout(60000), +) +go func() { + for chunk := range stream.Stderr { + fmt.Fprintf(os.Stderr, "%s", chunk) + } +}() +for chunk := range stream.Stdout { + fmt.Printf("%s", chunk) +} +exitCode, err := stream.Wait() +``` + +To cancel a long-running stream early: + +```go +stream.Cancel() +``` + ### Workspace ```go @@ -709,6 +741,17 @@ type HostExecResult struct { } ``` +### HostExecStream + +```go +type HostExecStream struct { + Stdout <-chan []byte + Stderr <-chan []byte + Wait func() (int, error) // blocks until exit; returns exit code + Cancel func() // cancels the stream context +} +``` + --- ## Errors diff --git a/sandbox/v2/host.go b/sandbox/v2/host.go index f6c970ed..89ccff4d 100644 --- a/sandbox/v2/host.go +++ b/sandbox/v2/host.go @@ -66,6 +66,89 @@ func (h *Host) Exec(ctx context.Context, cmd string, args []string, opts ...Host }, nil } +// Stream runs a command on the Tai host and streams stdout/stderr in real time +// via HostExec gRPC ExecStream. Returns a HostExecStream with separate channels +// for stdout and stderr, plus Wait (blocks until exit) and Cancel. +func (h *Host) Stream(ctx context.Context, cmd string, args []string, opts ...HostExecOption) (*HostExecStream, error) { + client, err := h.manager.getPool(h.pool) + if err != nil { + return nil, err + } + + he := client.HostExec() + if he == nil { + return nil, fmt.Errorf("sandbox: host_exec not available on pool %q", h.pool) + } + + cfg := &hostExecConfig{} + for _, o := range opts { + o(cfg) + } + + req := &hepb.ExecRequest{ + Command: cmd, + Args: args, + WorkingDir: cfg.WorkDir, + Stdin: cfg.Stdin, + TimeoutMs: cfg.TimeoutMs, + MaxOutputBytes: cfg.MaxOutputBytes, + } + if cfg.Env != nil { + req.Env = cfg.Env + } + + streamCtx, cancel := context.WithCancel(ctx) + rpcStream, err := he.ExecStream(streamCtx, req) + if err != nil { + cancel() + return nil, fmt.Errorf("hostexec stream rpc: %w", err) + } + + stdoutCh := make(chan []byte, 64) + stderrCh := make(chan []byte, 64) + doneCh := make(chan struct{}) + var exitCode int + var exitErr error + + go func() { + defer close(stdoutCh) + defer close(stderrCh) + defer close(doneCh) + for { + msg, err := rpcStream.Recv() + if err != nil { + exitErr = fmt.Errorf("hostexec stream recv: %w", err) + return + } + if len(msg.Data) > 0 { + switch msg.Stream { + case hepb.ExecOutput_STDOUT: + stdoutCh <- msg.Data + case hepb.ExecOutput_STDERR: + stderrCh <- msg.Data + } + } + if msg.Done { + exitCode = int(msg.ExitCode) + if msg.Error != "" { + exitErr = fmt.Errorf("hostexec: %s", msg.Error) + } + return + } + } + }() + + return &HostExecStream{ + Stdout: stdoutCh, + Stderr: stderrCh, + Wait: func() (int, error) { + <-doneCh + return exitCode, exitErr + }, + Cancel: cancel, + }, nil +} + // Workspace returns a filesystem interface for the given session on the host. // The sessionID typically corresponds to a workspace ID; files are stored // under dataDir/{sessionID}/ on the Tai host, accessed via Volume gRPC. diff --git a/sandbox/v2/host_test.go b/sandbox/v2/host_test.go index 4a56beb0..de5d090d 100644 --- a/sandbox/v2/host_test.go +++ b/sandbox/v2/host_test.go @@ -157,6 +157,210 @@ func TestHost_Workspace(t *testing.T) { } } +func TestHost_Stream_Incremental(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostExecTargets() { + if tgt.IsWinNative { + continue + } + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, tgt) + + host, err := m.Host(context.Background(), tgt.Name) + if err != nil { + t.Skipf("Host(%s): %v", tgt.Name, err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + stream, err := host.Stream(ctx, "sh", []string{"-c", + "for i in 1 2 3 4 5; do echo chunk$i; sleep 0.2; done"}) + if err != nil { + t.Fatalf("Stream: %v", err) + } + + var chunks []string + for chunk := range stream.Stdout { + chunks = append(chunks, string(chunk)) + } + + exitCode, err := stream.Wait() + if err != nil && !strings.Contains(err.Error(), "EOF") { + if strings.Contains(err.Error(), "not in the allowed list") { + t.Skipf("command not allowed on %s", tgt.Name) + } + t.Fatalf("Wait: %v", err) + } + if exitCode != 0 { + t.Errorf("exit_code = %d, want 0", exitCode) + } + + combined := strings.Join(chunks, "") + for _, expect := range []string{"chunk1", "chunk3", "chunk5"} { + if !strings.Contains(combined, expect) { + t.Errorf("output = %q, want contains %q", combined, expect) + } + } + + if len(chunks) < 2 { + t.Errorf("received %d chunks, want >= 2 (proves streaming, not buffered)", len(chunks)) + } + t.Logf("received %d chunks over stream", len(chunks)) + }) + } +} + +func TestHost_Stream_MultiLine(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostExecTargets() { + if tgt.IsWinNative { + continue + } + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, tgt) + + host, err := m.Host(context.Background(), tgt.Name) + if err != nil { + t.Skipf("Host(%s): %v", tgt.Name, err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + stream, err := host.Stream(ctx, "sh", []string{"-c", "for i in 1 2 3; do echo line$i; done"}) + if err != nil { + t.Fatalf("Stream: %v", err) + } + + var stdout []byte + for chunk := range stream.Stdout { + stdout = append(stdout, chunk...) + } + + exitCode, err := stream.Wait() + if err != nil && !strings.Contains(err.Error(), "EOF") { + if strings.Contains(err.Error(), "not in the allowed list") { + t.Skipf("command not allowed on %s", tgt.Name) + } + t.Fatalf("Wait: %v", err) + } + if exitCode != 0 { + t.Errorf("exit_code = %d, want 0", exitCode) + } + got := strings.TrimSpace(string(stdout)) + for _, expect := range []string{"line1", "line2", "line3"} { + if !strings.Contains(got, expect) { + t.Errorf("stdout = %q, want contains %q", got, expect) + } + } + }) + } +} + +func TestHost_Stream_Stderr(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostExecTargets() { + if tgt.IsWinNative { + continue + } + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, tgt) + + host, err := m.Host(context.Background(), tgt.Name) + if err != nil { + t.Skipf("Host(%s): %v", tgt.Name, err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + stream, err := host.Stream(ctx, "sh", []string{"-c", "echo err-msg >&2"}) + if err != nil { + t.Fatalf("Stream: %v", err) + } + + var stderr []byte + done := make(chan struct{}) + go func() { + for chunk := range stream.Stdout { + _ = chunk + } + close(done) + }() + for chunk := range stream.Stderr { + stderr = append(stderr, chunk...) + } + <-done + + exitCode, err := stream.Wait() + if err != nil && !strings.Contains(err.Error(), "EOF") { + if strings.Contains(err.Error(), "not in the allowed list") { + t.Skipf("command not allowed on %s", tgt.Name) + } + t.Fatalf("Wait: %v", err) + } + if exitCode != 0 { + t.Errorf("exit_code = %d, want 0", exitCode) + } + got := strings.TrimSpace(string(stderr)) + if !strings.Contains(got, "err-msg") { + t.Errorf("stderr = %q, want contains 'err-msg'", got) + } + }) + } +} + +func TestHost_Stream_Cancel(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostExecTargets() { + if tgt.IsWinNative { + continue + } + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, tgt) + + host, err := m.Host(context.Background(), tgt.Name) + if err != nil { + t.Skipf("Host(%s): %v", tgt.Name, err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + stream, err := host.Stream(ctx, "sh", []string{"-c", "while true; do echo tick; sleep 0.1; done"}) + if err != nil { + if strings.Contains(err.Error(), "not in the allowed list") { + t.Skipf("command not allowed on %s", tgt.Name) + } + t.Fatalf("Stream: %v", err) + } + + received := 0 + for chunk := range stream.Stdout { + _ = chunk + received++ + if received >= 3 { + stream.Cancel() + break + } + } + + _, waitErr := stream.Wait() + if waitErr != nil && strings.Contains(waitErr.Error(), "not in the allowed list") { + t.Skipf("command not allowed on %s", tgt.Name) + } + if received < 3 && waitErr == nil { + t.Errorf("received %d chunks before cancel, want >= 3", received) + } + }) + } +} + func TestHost_CreateRejectsNoContainerPool(t *testing.T) { // Use the Windows native HostExec target which has no Docker. tgt := findHostExecOnly(t) diff --git a/sandbox/v2/types.go b/sandbox/v2/types.go index 5ed8a4f7..d6263300 100644 --- a/sandbox/v2/types.go +++ b/sandbox/v2/types.go @@ -187,6 +187,15 @@ type HostExecResult struct { Truncated bool } +// HostExecStream provides real-time streaming output from a command running +// on the Tai host machine via HostExec gRPC ExecStream. +type HostExecStream struct { + Stdout <-chan []byte + Stderr <-chan []byte + Wait func() (int, error) // blocks until exit; returns exit code + Cancel func() // cancels the stream context +} + type hostExecConfig struct { WorkDir string Env map[string]string