diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 9771f464..b3283615 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -22,6 +22,7 @@ import ( "io" "net" "os" + "os/signal" "sort" "strconv" "strings" @@ -1106,10 +1107,23 @@ func runSSH(serverURL, atespace, kubeContext string, args []string) error { } defer guestClient.Close() + // A bare shell on a terminal runs with -i, so it prompts and Ctrl-C stops the + // command it is running rather than the shell itself. + if stdinIsTerminal() && len(cmdToRun) == 1 && cmdToRun[0] == "/bin/sh" { + cmdToRun = []string{"/bin/sh", "-i"} + } + + // Ctrl-C is passed on to the remote command instead of ending ax ssh. + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, os.Interrupt) + defer signal.Stop(sigs) + exitCode, err := guestClient.Exec(context.Background(), guest.ExecOptions{ Command: cmdToRun, + Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr, + Signals: sigs, }) if err != nil { return err @@ -1121,3 +1135,10 @@ func runSSH(serverURL, atespace, kubeContext string, args []string) error { return nil } + +// stdinIsTerminal reports whether standard input is a terminal rather than a +// pipe or file. +func stdinIsTerminal() bool { + fi, err := os.Stdin.Stat() + return err == nil && fi.Mode()&os.ModeCharDevice != 0 +} diff --git a/internal/guest/client.go b/internal/guest/client.go index 959e1fe6..ffb74802 100644 --- a/internal/guest/client.go +++ b/internal/guest/client.go @@ -19,7 +19,9 @@ import ( "errors" "fmt" "io" + "os" "strings" + "syscall" "time" ateenvv1alpha "github.com/agent-substrate/env/proto/ateenv/v1alpha" @@ -85,12 +87,18 @@ type ExecOptions struct { Command []string Cwd string Env map[string]string - Stdout io.Writer - Stderr io.Writer + // Stdin, if set, is copied to the process's standard input, which is closed + // once Stdin reaches EOF. If nil, the process reads an empty stdin. + Stdin io.Reader + Stdout io.Writer + Stderr io.Writer + // Signals, if set, delivers each signal received on it to the process group. + // Only SIGHUP, SIGINT, SIGQUIT and SIGTERM are forwarded. + Signals <-chan os.Signal } -// Exec runs a command inside the task container and streams stdout/stderr until completion. -// It returns the process exit code. +// Exec runs a command inside the task container, feeding it opts.Stdin and +// streaming stdout/stderr until completion. It returns the process exit code. func (c *Client) Exec(ctx context.Context, opts ExecOptions) (int, error) { if len(opts.Command) == 0 { return 1, errors.New("exec: command cannot be empty") @@ -100,12 +108,24 @@ func (c *Client) Exec(ctx context.Context, opts ExecOptions) (int, error) { Command: opts.Command, Cwd: opts.Cwd, Env: opts.Env, + Stdin: opts.Stdin != nil, }) if err != nil { return 1, fmt.Errorf("starting process: %w", err) } pid := proc.GetProcessId() + + // Input and signal forwarding stop when the command finishes. + ctx, cancel := context.WithCancel(ctx) + defer cancel() + if opts.Stdin != nil { + go c.forwardStdin(ctx, pid, opts.Stdin) + } + if opts.Signals != nil { + go c.forwardSignals(ctx, pid, opts.Signals) + } + stream, err := c.process.StreamProcessOutput(ctx, &ateenvv1alpha.StreamProcessOutputRequest{ ProcessId: pid, Follow: true, @@ -150,3 +170,51 @@ func (c *Client) Exec(ctx context.Context, opts ExecOptions) (int, error) { } } } + +// forwardStdin copies r to the stdin of process pid and closes it when r ends. +// Errors are dropped: they mean the process has exited or the connection is +// gone, and Exec reports either from the output stream. +func (c *Client) forwardStdin(ctx context.Context, pid string, r io.Reader) { + stream, err := c.process.WriteProcessInput(ctx) + if err != nil { + return + } + buf := make([]byte, 32*1024) + for { + n, readErr := r.Read(buf) + if n > 0 { + if err := stream.Send(&ateenvv1alpha.WriteProcessInputRequest{ProcessId: pid, Data: buf[:n]}); err != nil { + return + } + } + if readErr != nil { + if err := stream.Send(&ateenvv1alpha.WriteProcessInputRequest{ProcessId: pid, Close: true}); err != nil { + return + } + _, _ = stream.CloseAndRecv() + return + } + } +} + +// forwardedSignals maps the local signals Exec forwards to their guest equivalents. +var forwardedSignals = map[os.Signal]ateenvv1alpha.Signal{ + syscall.SIGHUP: ateenvv1alpha.Signal_SIGNAL_HUP, + syscall.SIGINT: ateenvv1alpha.Signal_SIGNAL_INT, + syscall.SIGQUIT: ateenvv1alpha.Signal_SIGNAL_QUIT, + syscall.SIGTERM: ateenvv1alpha.Signal_SIGNAL_TERM, +} + +// forwardSignals delivers each signal from sigs to process pid until ctx ends. +func (c *Client) forwardSignals(ctx context.Context, pid string, sigs <-chan os.Signal) { + for { + select { + case <-ctx.Done(): + return + case sig := <-sigs: + if guestSig, ok := forwardedSignals[sig]; ok { + _, _ = c.process.SignalProcess(ctx, &ateenvv1alpha.SignalProcessRequest{ProcessId: pid, Signal: guestSig}) + } + } + } +} diff --git a/internal/guest/client_test.go b/internal/guest/client_test.go new file mode 100644 index 00000000..7d8109dd --- /dev/null +++ b/internal/guest/client_test.go @@ -0,0 +1,240 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package guest_test + +import ( + "bytes" + "context" + "io" + "net" + "os" + "strings" + "sync" + "testing" + "time" + + envguest "github.com/agent-substrate/env/guest" + "github.com/google/ax/internal/guest" +) + +// startGuest serves the real guest services on a local port and returns a +// client connected to them. +func startGuest(t *testing.T) *guest.Client { + t.Helper() + cfg := envguest.DefaultConfig() + cfg.Workspace = t.TempDir() + cfg.LogDir = t.TempDir() + srv, cleanup, err := envguest.NewServer(cfg) + if err != nil { + t.Fatalf("creating guest server: %v", err) + } + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listening: %v", err) + } + go func() { _ = srv.Serve(lis) }() + t.Cleanup(func() { + srv.Stop() + cleanup() + }) + + client, err := guest.Dial(lis.Addr().String()) + if err != nil { + t.Fatalf("dialing guest: %v", err) + } + t.Cleanup(func() { _ = client.Close() }) + return client +} + +func execContext(t *testing.T) context.Context { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + t.Cleanup(cancel) + return ctx +} + +// newPipe returns an OS pipe standing in for a terminal's stdin. +func newPipe(t *testing.T) (*os.File, *os.File) { + t.Helper() + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("creating pipe: %v", err) + } + t.Cleanup(func() { + _ = r.Close() + _ = w.Close() + }) + return r, w +} + +// syncBuffer is a bytes.Buffer that can be written and read from different goroutines. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *syncBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +func TestExecForwardsStdin(t *testing.T) { + client := startGuest(t) + + var stdout bytes.Buffer + code, err := client.Exec(execContext(t), guest.ExecOptions{ + Command: []string{"cat"}, + Stdin: strings.NewReader("hello from stdin\n"), + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if code != 0 { + t.Errorf("exit code = %d, want 0", code) + } + if got := stdout.String(); got != "hello from stdin\n" { + t.Errorf("stdout = %q, want %q", got, "hello from stdin\n") + } +} + +func TestExecStreamsInteractiveInput(t *testing.T) { + client := startGuest(t) + + // Each line is only written after the reply to the previous one arrives, + // so this passes only if input and output stream while the process runs. + stdinR, stdinW := newPipe(t) + stdout := &syncBuffer{} + done := make(chan struct{}) + var ( + code int + err error + ) + go func() { + defer close(done) + code, err = client.Exec(execContext(t), guest.ExecOptions{ + Command: []string{"sh", "-c", "while read line; do echo \"got $line\"; done; echo bye"}, + Stdin: stdinR, + Stdout: stdout, + }) + }() + + waitFor := func(want string) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for !strings.Contains(stdout.String(), want) { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %q, stdout so far %q", want, stdout.String()) + } + time.Sleep(10 * time.Millisecond) + } + } + _, _ = io.WriteString(stdinW, "one\n") + waitFor("got one\n") + _, _ = io.WriteString(stdinW, "two\n") + waitFor("got two\n") + + // Closing stdin ends the read loop, so the shell prints bye and exits. + _ = stdinW.Close() + <-done + if err != nil { + t.Fatalf("Exec: %v", err) + } + if code != 0 { + t.Errorf("exit code = %d, want 0", code) + } + if got, want := stdout.String(), "got one\ngot two\nbye\n"; got != want { + t.Errorf("stdout = %q, want %q", got, want) + } +} + +func TestExecPropagatesExitCode(t *testing.T) { + client := startGuest(t) + + code, err := client.Exec(execContext(t), guest.ExecOptions{ + Command: []string{"sh", "-c", "read code; exit $code"}, + Stdin: strings.NewReader("7\n"), + }) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if code != 7 { + t.Errorf("exit code = %d, want 7", code) + } +} + +func TestExecWithoutStdinReadsEmpty(t *testing.T) { + client := startGuest(t) + + var stdout bytes.Buffer + code, err := client.Exec(execContext(t), guest.ExecOptions{ + Command: []string{"sh", "-c", "cat; echo done"}, + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if code != 0 || stdout.String() != "done\n" { + t.Errorf("got exit code %d, stdout %q; want 0, %q", code, stdout.String(), "done\n") + } +} + +func TestExecForwardsSignals(t *testing.T) { + client := startGuest(t) + + // The process prints once it is running, then waits on stdin, which the + // test keeps open. Only the forwarded SIGINT can end it. + stdinR, _ := newPipe(t) + stdout := &syncBuffer{} + sigs := make(chan os.Signal, 1) + done := make(chan struct{}) + var ( + code int + err error + ) + go func() { + defer close(done) + code, err = client.Exec(execContext(t), guest.ExecOptions{ + Command: []string{"sh", "-c", "echo ready; cat"}, + Stdin: stdinR, + Stdout: stdout, + Signals: sigs, + }) + }() + + deadline := time.Now().Add(5 * time.Second) + for !strings.Contains(stdout.String(), "ready") { + if time.Now().After(deadline) { + t.Fatalf("process never became ready") + } + time.Sleep(10 * time.Millisecond) + } + sigs <- os.Interrupt + <-done + + if err != nil { + t.Fatalf("Exec: %v", err) + } + if code != 130 { + t.Errorf("exit code = %d, want 130 (killed by SIGINT)", code) + } +}