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)
|
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
|
### Workspace
|
||||||
|
|
||||||
```go
|
```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
|
## Errors
|
||||||
|
|
|
||||||
|
|
@ -66,6 +66,89 @@ func (h *Host) Exec(ctx context.Context, cmd string, args []string, opts ...Host
|
||||||
}, nil
|
}, 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.
|
// Workspace returns a filesystem interface for the given session on the host.
|
||||||
// The sessionID typically corresponds to a workspace ID; files are stored
|
// The sessionID typically corresponds to a workspace ID; files are stored
|
||||||
// under dataDir/{sessionID}/ on the Tai host, accessed via Volume gRPC.
|
// 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) {
|
func TestHost_CreateRejectsNoContainerPool(t *testing.T) {
|
||||||
// Use the Windows native HostExec target which has no Docker.
|
// Use the Windows native HostExec target which has no Docker.
|
||||||
tgt := findHostExecOnly(t)
|
tgt := findHostExecOnly(t)
|
||||||
|
|
|
||||||
|
|
@ -187,6 +187,15 @@ type HostExecResult struct {
|
||||||
Truncated bool
|
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 {
|
type hostExecConfig struct {
|
||||||
WorkDir string
|
WorkDir string
|
||||||
Env map[string]string
|
Env map[string]string
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue