feat(sandbox/v2): add Host.Stream for real-time streaming exec on Tai host
Add HostExecStream type and Host.Stream method that wraps the HostExec gRPC ExecStream RPC, providing separate stdout/stderr channels with Wait and Cancel support. Includes 4 tests covering incremental output, multi-line, stderr separation, and cancel behavior across k8s and win-linux platforms. Update API.md documentation accordingly. Made-with: Cursor
This commit is contained in:
Max
parent
1ffdcc8817
commit
1e454a38ad
4 changed files with 339 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue