feat(cli): forward stdin and Ctrl-C in ax ssh

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
+}

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})
+			}
+		}
+	}
+}

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)
+	}
+}