fix(cli): completion with direct connections (--connect, --context, --uncloud-config) (fixes #377)

cmd/uc/main.go

@@ -24,6 +24,7 @@ import (
 	"github.com/psviderski/uncloud/internal/machine"
 	"github.com/psviderski/uncloud/internal/version"
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 )
 
 type globalOptions struct {
@@ -42,6 +43,13 @@ func main() {
 		SilenceUsage:  true,
 		SilenceErrors: true,
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+			// Shell completion runs through the hidden __complete command which has flag parsing disabled,
+			// so the global flags from the completed command line are never parsed. Apply them manually to make
+			// completion work with --connect, --context, and --uncloud-config.
+			if cmd.Name() == cobra.ShellCompRequestCmd {
+				applyGlobalFlagsFromCompletionArgs(cmd.Root().PersistentFlags(), os.Args[1:])
+			}
+
 			cli.BindEnvToFlag(cmd, "connect", "UNCLOUD_CONNECT")
 			cli.BindEnvToFlag(cmd, "context", "UNCLOUD_CONTEXT")
 			cli.BindEnvToFlag(cmd, "uncloud-config", "UNCLOUD_CONFIG")
@@ -159,3 +167,27 @@ func main() {
 		cobra.CheckErr(err)
 	}
 }
+
+// applyGlobalFlagsFromCompletionArgs parses the global flags from the raw arguments of a __complete command and applies
+// the ones found to flags. The trailing word being completed, unknown flags, and positional arguments are ignored.
+func applyGlobalFlagsFromCompletionArgs(flags *pflag.FlagSet, args []string) {
+	// The shell always passes the word being completed as the last argument, even if it's empty.
+	// Exclude it from parsing as its value may not be complete yet.
+	if len(args) == 0 {
+		return
+	}
+	args = args[:len(args)-1]
+
+	fset := pflag.NewFlagSet("global", pflag.ContinueOnError)
+	fset.ParseErrorsAllowlist.UnknownFlags = true
+	fset.String("connect", "", "")
+	fset.StringP("context", "c", "", "")
+	fset.String("uncloud-config", "", "")
+	// Parsing an incomplete command line may fail, apply the flags parsed so far anyway.
+	_ = fset.Parse(args)
+
+	fset.Visit(func(f *pflag.Flag) {
+		// Setting the flag marks it as changed so it takes precedence over environment variables.
+		_ = flags.Set(f.Name, f.Value.String())
+	})
+}

cmd/uc/main_test.go

@@ -0,0 +1,116 @@
+package main
+
+import (
+	"slices"
+	"testing"
+
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestApplyGlobalFlagsFromCompletionArgs(t *testing.T) {
+	defaultConfigPath := "~/.config/uncloud/config.yaml"
+
+	tests := []struct {
+		name           string
+		args           []string
+		wantConnect    string
+		wantContext    string
+		wantConfigPath string
+		// Flag names expected to be marked as changed on the target flag set.
+		wantChanged []string
+	}{
+		{
+			name: "no flags",
+			args: []string{"__complete", "inspect", ""},
+		},
+		{
+			name:        "connect with space",
+			args:        []string{"__complete", "--connect", "ssh://user@host", "inspect", ""},
+			wantConnect: "ssh://user@host",
+			wantChanged: []string{"connect"},
+		},
+		{
+			name:        "connect with equals",
+			args:        []string{"__complete", "--connect=tcp://127.0.0.1:51000", "inspect", ""},
+			wantConnect: "tcp://127.0.0.1:51000",
+			wantChanged: []string{"connect"},
+		},
+		{
+			name:        "context shorthand",
+			args:        []string{"__complete", "-c", "prod", "inspect", ""},
+			wantContext: "prod",
+			wantChanged: []string{"context"},
+		},
+		{
+			name:           "all flags",
+			args:           []string{"__complete", "--connect", "user@host", "-c", "prod", "--uncloud-config", "/tmp/uncloud.yaml", "inspect", ""},
+			wantConnect:    "user@host",
+			wantContext:    "prod",
+			wantConfigPath: "/tmp/uncloud.yaml",
+			wantChanged:    []string{"connect", "context", "uncloud-config"},
+		},
+		{
+			name:        "unknown flags are ignored",
+			args:        []string{"__complete", "--quiet", "-n", "5", "--connect", "user@host", "logs", ""},
+			wantConnect: "user@host",
+			wantChanged: []string{"connect"},
+		},
+		{
+			name: "flags after double dash are ignored",
+			args: []string{"__complete", "exec", "svc", "--", "sh", "--connect", "user@host"},
+		},
+		{
+			name:        "flags before double dash are applied",
+			args:        []string{"__complete", "--connect", "user@host", "exec", "svc", "--", "sh", "-c", "env"},
+			wantConnect: "user@host",
+			wantChanged: []string{"connect"},
+		},
+		{
+			name:        "partial flag name being completed is excluded",
+			args:        []string{"__complete", "--connect", "user@host", "inspect", "--context"},
+			wantConnect: "user@host",
+			wantChanged: []string{"connect"},
+		},
+		{
+			name: "partial flag value being completed is excluded",
+			args: []string{"__complete", "--uncloud-config", "/tmp/"},
+		},
+		{
+			name: "partial connect value being completed is excluded",
+			args: []string{"__complete", "--connect", "tcp://127.0.0.1:5"},
+		},
+		{
+			name:           "completed flag value with partial command word",
+			args:           []string{"__complete", "--uncloud-config", "/tmp/uncloud.yaml", "insp"},
+			wantConfigPath: "/tmp/uncloud.yaml",
+			wantChanged:    []string{"uncloud-config"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Mirror the global persistent flags defined on the root command.
+			var opts globalOptions
+			flags := pflag.NewFlagSet("test", pflag.ContinueOnError)
+			flags.StringVar(&opts.connect, "connect", "", "")
+			flags.StringVarP(&opts.context, "context", "c", "", "")
+			flags.StringVar(&opts.configPath, "uncloud-config", defaultConfigPath, "")
+
+			applyGlobalFlagsFromCompletionArgs(flags, tt.args)
+
+			assert.Equal(t, tt.wantConnect, opts.connect)
+			assert.Equal(t, tt.wantContext, opts.context)
+			wantConfigPath := tt.wantConfigPath
+			if wantConfigPath == "" {
+				wantConfigPath = defaultConfigPath
+			}
+			assert.Equal(t, wantConfigPath, opts.configPath)
+
+			for _, name := range []string{"connect", "context", "uncloud-config"} {
+				assert.Equal(t, slices.Contains(tt.wantChanged, name), flags.Changed(name),
+					"changed status of flag '%s'", name)
+			}
+		})
+	}
+}

cmd/uc/proxy.go

@@ -132,35 +132,34 @@ func runProxy(ctx context.Context, uncli *cli.CLI, opts proxyOptions) error {
 	// endpoint and shuffles the data, *it* will actually experience errors.
 	remoteAddr := net.JoinHostPort(ip.String(), strconv.Itoa(opts.remotePort))
 
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
 	p := &proxy.Proxy{
 		Listener:    listener,
 		RemoteAddr:  remoteAddr,
 		DialContext: dialer.DialContext,
 		OnError: func(err error) {
-			fmt.Printf("Failed to proxy to '%s': %v\n", remoteAddr, err)
-			cancel()
+			if proxy.IsConnectionClosedError(err) {
+				return
+			}
+			// A more actionable error instead of the cryptic [ssh -W] command error.
+			if strings.Contains(err.Error(), "Session open refused by peer") {
+				fmt.Printf("Could not connect to '%s': connection refused. "+
+					"Check that the service is running and listening on port %d inside the container.\n",
+					remoteAddr, opts.remotePort)
+				return
+			}
+
+			fmt.Printf("Failed to proxy a connection to '%s': %v\n", remoteAddr, err)
 		},
 	}
 
-	// Run the proxy in the background and signal when it has fully shut down.
-	done := make(chan struct{})
-	go func() {
-		p.Run(ctx)
-		close(done)
-	}()
-
 	// Prefix the local address with the scheme for common HTTP ports so it becomes control-clickable in most
 	// terminals. We assume plain HTTP since TLS is typically terminated by Caddy in front of the service.
 	fmt.Printf("%s%s → %s (%s%s%s)\n", schemeForPort(opts.remotePort), p.Listener.Addr().String(),
 		remoteAddr, opts.service, tui.Faint.Render("/"), containerID)
 
-	<-ctx.Done()
-	// Wait for the proxy to drain in-flight connections and shut down gracefully.
-	<-done
-
+	if err = p.Run(ctx); err != nil {
+		return fmt.Errorf("run proxy to '%s': %w", remoteAddr, err)
+	}
 	return nil
 }
 

internal/cli/completion/context.go

@@ -11,6 +11,11 @@ import (
 )
 
 func Contexts(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
+	// There are no contexts to complete when the CLI uses a direct machine connection (--connect) without a config.
+	if uncli.Config == nil {
+		return nil, cobra.ShellCompDirectiveNoFileComp
+	}
+
 	contexts := slices.Sorted(maps.Keys(uncli.Config.Contexts))
 
 	names := []cobra.Completion{}
@@ -21,7 +26,6 @@ func Contexts(ctx context.Context, uncli *cli.CLI, args []string, toComplete str
 		if strings.HasPrefix(context, toComplete) {
 			names = append(names, context)
 		}
-		names = append(names, context)
 	}
 
 	return names, cobra.ShellCompDirectiveNoFileComp

internal/cli/completion/machine.go

@@ -10,7 +10,8 @@ import (
 )
 
 func Machines(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
-	client, err := uncli.ConnectCluster(ctx)
+	// Disable the connection progress output to not interfere with the shell completion output.
+	client, err := uncli.ConnectClusterWithOptions(ctx, cli.ConnectOptions{})
 	if err != nil {
 		return nil, cobra.ShellCompDirectiveError
 	}

internal/cli/completion/service.go

@@ -11,7 +11,8 @@ import (
 )
 
 func Services(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
-	client, err := uncli.ConnectCluster(ctx)
+	// Disable the connection progress output to not interfere with the shell completion output.
+	client, err := uncli.ConnectClusterWithOptions(ctx, cli.ConnectOptions{})
 	if err != nil {
 		return nil, cobra.ShellCompDirectiveError
 	}

internal/cli/completion/volume.go

@@ -11,7 +11,8 @@ import (
 )
 
 func Volumes(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
-	client, err := uncli.ConnectCluster(ctx)
+	// Disable the connection progress output to not interfere with the shell completion output.
+	client, err := uncli.ConnectClusterWithOptions(ctx, cli.ConnectOptions{})
 	if err != nil {
 		return nil, cobra.ShellCompDirectiveError
 	}

internal/proxy/proxy.go

@@ -2,11 +2,12 @@ package proxy
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net"
-	"os"
 	"sync"
+	"syscall"
 	"time"
 )
 
@@ -15,58 +16,53 @@ type Proxy struct {
 	Listener    net.Listener
 	RemoteAddr  string
 	DialContext func(ctx context.Context, network, address string) (net.Conn, error)
+	// OnError is called for errors that occur during proxying individual connections. It may be called concurrently
+	// for different connections.
 	OnError     func(error)
 	activeConns sync.WaitGroup
 }
 
-// deadliner is an interface for listeners that support setting deadlines.
-type deadliner interface {
-	SetDeadline(t time.Time) error
-}
-
 // halfCloser is an interface for connections that support half-close.
 type halfCloser interface {
 	CloseWrite() error
 }
 
-// Run starts the proxy and runs until the context is canceled.
-func (p *Proxy) Run(ctx context.Context) {
+// IsConnectionClosedError reports whether err indicates that a connection was closed or aborted by either peer.
+// Callers can use it to ignore routine connection shutdown or broken pipe errors reported to Proxy.OnError.
+func IsConnectionClosedError(err error) bool {
+	return errors.Is(err, net.ErrClosed) || errors.Is(err, io.ErrClosedPipe) ||
+		errors.Is(err, syscall.EPIPE) || errors.Is(err, syscall.ECONNRESET)
+}
+
+// Run starts the proxy and runs until the context is canceled or the listener fails. It returns nil when the context
+// is canceled. Errors handling individual connections are reported to OnError and do not stop the proxy.
+func (p *Proxy) Run(ctx context.Context) error {
 	if p.DialContext == nil {
 		p.DialContext = (&net.Dialer{}).DialContext
 	}
 
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 	defer p.Listener.Close()
 
-	// Handle incoming connections until context is canceled.
-Loop:
+	// Closing the listener unblocks Accept when the context is canceled. This works for both TCP and Unix listeners
+	// and avoids polling with listener deadlines.
+	stopClose := context.AfterFunc(ctx, func() {
+		p.Listener.Close()
+	})
+	defer stopClose()
+
+	var runErr error
 	for {
-		select {
-		case <-ctx.Done():
-			break Loop
-		default:
-		}
-
-		// Set a deadline on the listener if supported to check context periodically.
-		if dl, ok := p.Listener.(deadliner); ok {
-			dl.SetDeadline(time.Now().Add(1 * time.Second))
-		}
-
 		conn, err := p.Listener.Accept()
 		if err != nil {
-			if os.IsTimeout(err) {
-				// Just a timeout, continue to check context and accept again.
-				continue
+			if ctx.Err() != nil {
+				break
 			}
 
-			select {
-			case <-ctx.Done():
-				break Loop
-			default:
-				if p.OnError != nil {
-					p.OnError(fmt.Errorf("accept local connection: %w", err))
-				}
-				continue
-			}
+			runErr = fmt.Errorf("accept local connection: %w", err)
+			cancel()
+			break
 		}
 
 		p.activeConns.Add(1)
@@ -75,6 +71,7 @@ Loop:
 
 	// Wait for all connections to finish.
 	p.activeConns.Wait()
+	return runErr
 }
 
 func (p *Proxy) handleConnection(ctx context.Context, localConn net.Conn) {
@@ -87,46 +84,62 @@ func (p *Proxy) handleConnection(ctx context.Context, localConn net.Conn) {
 
 	remoteConn, err := p.DialContext(dialCtx, "tcp", p.RemoteAddr)
 	if err != nil {
-		if p.OnError != nil {
+		if ctx.Err() == nil && p.OnError != nil {
 			p.OnError(fmt.Errorf("connect remote address '%s': %w", p.RemoteAddr, err))
 		}
 		return
 	}
 	defer remoteConn.Close()
 
-	// Bidirectional copy with proper half-close handling.
+	// Closing both connections aborts both copies after cancellation or a copy error. A clean EOF still uses
+	// half-close so the other direction can finish sending any remaining data.
+	closeConnections := func() {
+		localConn.Close()
+		remoteConn.Close()
+	}
+	stopClose := context.AfterFunc(ctx, closeConnections)
+	defer stopClose()
+
 	done := make(chan error, 2)
 
 	go func() {
 		_, err := io.Copy(remoteConn, localConn)
+		if err != nil {
+			done <- err
+			closeConnections()
+			return
+		}
 		// Close write half of remote connection if supported.
 		if hc, ok := remoteConn.(halfCloser); ok {
 			hc.CloseWrite()
 		}
-		done <- err
+		done <- nil
 	}()
 
 	go func() {
 		_, err := io.Copy(localConn, remoteConn)
+		if err != nil {
+			done <- err
+			closeConnections()
+			return
+		}
 		// Close write half of local connection if supported.
 		if hc, ok := localConn.(halfCloser); ok {
 			hc.CloseWrite()
 		}
-		done <- err
+		done <- nil
 	}()
 
-	// Wait for both copies to complete or context cancel.
+	// Wait for both copies to complete. The first error is the original failure because a copy reports it before
+	// closing the connections to unblock the other copy.
+	var copyErr error
 	for range 2 {
-		select {
-		case <-ctx.Done():
-			// Close connections to abort ongoing copies.
-			localConn.Close()
-			remoteConn.Close()
-			return
-		case err = <-done:
-			if err != nil && p.OnError != nil {
-				p.OnError(fmt.Errorf("data copy: %w", err))
-			}
+		if err = <-done; err != nil && copyErr == nil {
+			copyErr = err
 		}
 	}
+
+	if copyErr != nil && ctx.Err() == nil && p.OnError != nil {
+		p.OnError(fmt.Errorf("data copy: %w", copyErr))
+	}
 }

internal/proxy/proxy_test.go

@@ -0,0 +1,188 @@
+package proxy
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"sync/atomic"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsConnectionClosedError(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{name: "closed network connection", err: net.ErrClosed, want: true},
+		{name: "closed pipe", err: io.ErrClosedPipe, want: true},
+		{name: "broken pipe", err: syscall.EPIPE, want: true},
+		{name: "connection reset", err: syscall.ECONNRESET, want: true},
+		{name: "wrapped connection error", err: fmt.Errorf("copy data: %w", syscall.EPIPE), want: true},
+		{name: "other error", err: errors.New("copy failed"), want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.want, IsConnectionClosedError(tt.err))
+		})
+	}
+}
+
+func TestRunContinuesAfterClosedConnectionError(t *testing.T) {
+	t.Parallel()
+
+	listener := newTestListener()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	closedErrCh := make(chan error, 1)
+	unexpectedErrCh := make(chan error, 1)
+	var dialCount atomic.Int32
+
+	p := &Proxy{
+		Listener:   listener,
+		RemoteAddr: "remote:80",
+		DialContext: func(context.Context, string, string) (net.Conn, error) {
+			if dialCount.Add(1) == 1 {
+				return readErrorConn{err: syscall.EPIPE}, nil
+			}
+
+			proxyConn, upstreamConn := net.Pipe()
+			go func() {
+				defer upstreamConn.Close()
+				_, _ = upstreamConn.Write([]byte("ok"))
+			}()
+			return proxyConn, nil
+		},
+		OnError: func(err error) {
+			if IsConnectionClosedError(err) {
+				closedErrCh <- err
+				return
+			}
+			unexpectedErrCh <- err
+		},
+	}
+
+	runErrCh := make(chan error, 1)
+	go func() {
+		runErrCh <- p.Run(ctx)
+	}()
+
+	firstConn := listener.connect()
+	defer firstConn.Close()
+
+	select {
+	case connErr := <-closedErrCh:
+		require.ErrorIs(t, connErr, syscall.EPIPE)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for closed connection error")
+	}
+
+	secondConn := listener.connect()
+	defer secondConn.Close()
+	require.NoError(t, secondConn.SetReadDeadline(time.Now().Add(time.Second)))
+
+	got := make([]byte, 2)
+	_, err := io.ReadFull(secondConn, got)
+	require.NoError(t, err)
+	require.Equal(t, "ok", string(got))
+
+	select {
+	case unexpectedErr := <-unexpectedErrCh:
+		t.Fatalf("unexpected connection error: %v", unexpectedErr)
+	default:
+	}
+
+	cancel()
+	select {
+	case runErr := <-runErrCh:
+		require.NoError(t, runErr)
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for proxy to stop")
+	}
+}
+
+func TestRunReturnsListenerError(t *testing.T) {
+	t.Parallel()
+
+	listenerErr := errors.New("listener failed")
+	listener := errorListener{err: listenerErr}
+
+	p := &Proxy{Listener: listener}
+	err := p.Run(context.Background())
+	require.Error(t, err)
+	require.ErrorContains(t, err, "accept local connection")
+	require.ErrorIs(t, err, listenerErr)
+}
+
+type testListener struct {
+	conns     chan net.Conn
+	closed    chan struct{}
+	closeOnce sync.Once
+}
+
+func newTestListener() *testListener {
+	return &testListener{
+		conns:  make(chan net.Conn),
+		closed: make(chan struct{}),
+	}
+}
+
+func (l *testListener) connect() net.Conn {
+	clientConn, proxyConn := net.Pipe()
+	l.conns <- proxyConn
+	return clientConn
+}
+
+func (l *testListener) Accept() (net.Conn, error) {
+	select {
+	case conn := <-l.conns:
+		return conn, nil
+	case <-l.closed:
+		return nil, net.ErrClosed
+	}
+}
+
+func (l *testListener) Close() error {
+	l.closeOnce.Do(func() {
+		close(l.closed)
+	})
+	return nil
+}
+
+func (l *testListener) Addr() net.Addr {
+	return &net.TCPAddr{}
+}
+
+type errorListener struct {
+	err error
+}
+
+func (l errorListener) Accept() (net.Conn, error) { return nil, l.err }
+func (errorListener) Close() error                { return nil }
+func (errorListener) Addr() net.Addr              { return &net.TCPAddr{} }
+
+// readErrorConn fails reads immediately so tests can deterministically exercise a proxy copy failure.
+type readErrorConn struct {
+	err error
+}
+
+func (c readErrorConn) Read([]byte) (int, error)       { return 0, c.err }
+func (readErrorConn) Write(p []byte) (int, error)      { return len(p), nil }
+func (readErrorConn) Close() error                     { return nil }
+func (readErrorConn) LocalAddr() net.Addr              { return &net.TCPAddr{} }
+func (readErrorConn) RemoteAddr() net.Addr             { return &net.TCPAddr{} }
+func (readErrorConn) SetDeadline(time.Time) error      { return nil }
+func (readErrorConn) SetReadDeadline(time.Time) error  { return nil }
+func (readErrorConn) SetWriteDeadline(time.Time) error { return nil }

pkg/client/image.go

@@ -252,13 +252,19 @@ func (cli *Client) pushImageToMachine(
 	// The proxy runs in a goroutine. Capture the first error in a channel
 	// so we can surface it alongside the push error if push fails.
 	proxyErrCh := make(chan error, 1)
-	onProxyError := func(err error) {
+	recordProxyError := func(err error) {
 		select {
 		case proxyErrCh <- fmt.Errorf("proxy to unregistry: %w", err):
 		default:
 		}
 		pw.Event(progress.NewEvent(proxyEventID, progress.Error, err.Error()))
 	}
+	onProxyError := func(err error) {
+		if proxy.IsConnectionClosedError(err) {
+			return
+		}
+		recordProxyError(err)
+	}
 
 	// socketPath is set for plain rootless Docker (not running inside a VM): the Go proxy listens on a unix
 	// socket that is bind-mounted into the socat container, bypassing slirp4netns network routing entirely.
@@ -315,7 +321,11 @@ func (cli *Client) pushImageToMachine(
 	}
 	defer cleanup()
 
-	go unregProxy.Run(proxyCtx)
+	go func() {
+		if err := unregProxy.Run(proxyCtx); err != nil {
+			recordProxyError(err)
+		}
+	}()
 
 	if dockerEnv.Virtualised {
 		// VM-based Docker (Docker Desktop, Rancher Desktop, etc.): run a socat container inside the VM