From 02ddebf18dbb08f3cd0a5ccf9fdc2eb24fbc7241 Mon Sep 17 00:00:00 2001 From: Anzal Husain Abidi Date: Fri, 25 Sep 2026 22:49:51 +0530 Subject: [PATCH] feat(cli): forward stdin and Ctrl-C in ax ssh ax ssh with no command started /bin/sh remotely, but Exec never opened the process's stdin, so the shell read EOF and exited straight away. guest.ExecOptions now takes Stdin, which is streamed to the process via WriteProcessInput and closed on EOF, and Signals, which are delivered with SignalProcess. ax ssh passes os.Stdin and forwards SIGINT, so Ctrl-C interrupts the remote command. When stdin is a terminal the default shell runs as /bin/sh -i so it shows a prompt and survives Ctrl-C. The guest protocol has no PTY support yet, so full-screen programs and job control still don't work. Part of #374 --- cmd/ax/main.go | 21 +++ internal/guest/client.go | 76 ++++++++++- internal/guest/client_test.go | 240 ++++++++++++++++++++++++++++++++++ 3 files changed, 333 insertions(+), 4 deletions(-) create mode 100644 internal/guest/client_test.go 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) + } +}