feat(sdk): route Go SDK exec and forward to the sandbox owner replica (#4852)

* feat(sdk): route Go SDK exec and forward to the sandbox owner replica

The Go SDK now does what the CLI already does for long-lived sandbox
calls. The GetSandbox check that Exec.Run, Exec.Stream, Exec.Interactive,
TCP.Forward and SSH.Tunnel already make also asks the gateway for the
owner replica. The call then sends it in the x-openshell-replica header
and retries once without the header when the edge proxy reports that the
owner's route has no ready endpoints.

No extra round trip is added. Without a reported owner the calls are
unchanged.

Part of #4223

Signed-off-by: divesh <dgude@nvidia.com>

* refactor(sdk): address Go replica routing review

Share the header check between the routed attempt and its retry, keep the
first-frame EOF rule in one bidi helper, reuse Get's request for the owner
lookup, and document the header contract next to the replica routing docs.

Part of #4223

Signed-off-by: divesh <dgude@nvidia.com>

* docs(kubernetes): refer to the SDKs collectively in replica routing

Part of #4223

Signed-off-by: divesh <dgude@nvidia.com>

---------

Signed-off-by: divesh <dgude@nvidia.com>

docs/kubernetes/ingress.mdx

@@ -257,7 +257,7 @@ helm upgrade <release-name> oci://ghcr.io/nvidia/openshell/helm-chart \
   --set grpcRoute.replicaRouting.enabled=true
 ```
 
-The chart renders one Service per replica and a GRPCRoute rule that matches the `x-openshell-replica` request header. The CLI and the Rust SDK set that header from the owner the gateway reports and retry without it if the named replica is unavailable. Replica routing requires `workload.kind=statefulset`, `workload.allowMultiReplicaStatefulSet=true`, and at most 15 replicas.
+The chart renders one Service per replica and a GRPCRoute rule that matches the `x-openshell-replica` request header. The CLI and SDKs send `x-openshell-want-owner: 1` on `GetSandbox`, read the owner from the `x-openshell-owner` response header, and set it as `x-openshell-replica`. If the proxy answers `UNAVAILABLE` with an empty message, `no healthy upstream`, or `connection failure`, they retry once without it. Replica routing requires `workload.kind=statefulset`, `workload.allowMultiReplicaStatefulSet=true`, and at most 15 replicas.
 
 If the release uses `workload.kind=deployment`, this upgrade replaces the Deployment with a StatefulSet, so every gateway pod restarts at once. The StatefulSet also requests a 1Gi PersistentVolumeClaim per replica, which stays `Pending` on clusters without a default StorageClass.
 

sdk/go/openshell/v1/exec_client.go

@@ -27,7 +27,8 @@ func (e *execClient) Run(ctx context.Context, workspace, sandboxName string, com
 	if sandboxName == "" {
 		return nil, &StatusError{Code: ErrorInvalidArgument, Message: "sandbox name must not be empty"}
 	}
-	if _, err := e.sandboxes.Get(ctx, workspace, sandboxName); err != nil {
+	owner, err := sandboxOwner(ctx, e.sandboxes, workspace, sandboxName)
+	if err != nil {
 		return nil, err
 	}
 	var opt *ExecOptions
@@ -37,7 +38,9 @@ func (e *execClient) Run(ctx context.Context, workspace, sandboxName string, com
 	req := converter.ExecRequestToProto(sandboxName, command, opt)
 	req.WorkspaceScope = namedWorkspaceScope(workspace)
 
-	stream, err := e.client.ExecSandbox(ctx, req)
+	stream, err := openRouted(ctx, owner, func(ctx context.Context) (grpc.ServerStreamingClient[pb.ExecSandboxEvent], error) {
+		return e.client.ExecSandbox(ctx, req)
+	})
 	if err != nil {
 		return nil, converter.FromGRPCError(err)
 	}
@@ -61,7 +64,8 @@ func (e *execClient) Stream(ctx context.Context, workspace, sandboxName string,
 	if sandboxName == "" {
 		return nil, &StatusError{Code: ErrorInvalidArgument, Message: "sandbox name must not be empty"}
 	}
-	if _, err := e.sandboxes.Get(ctx, workspace, sandboxName); err != nil {
+	owner, err := sandboxOwner(ctx, e.sandboxes, workspace, sandboxName)
+	if err != nil {
 		return nil, err
 	}
 	var opt *ExecOptions
@@ -72,7 +76,9 @@ func (e *execClient) Stream(ctx context.Context, workspace, sandboxName string,
 	req.WorkspaceScope = namedWorkspaceScope(workspace)
 
 	streamCtx, cancel := context.WithCancel(ctx)
-	stream, err := e.client.ExecSandbox(streamCtx, req)
+	stream, err := openRouted(streamCtx, owner, func(ctx context.Context) (grpc.ServerStreamingClient[pb.ExecSandboxEvent], error) {
+		return e.client.ExecSandbox(ctx, req)
+	})
 	if err != nil {
 		cancel()
 		return nil, converter.FromGRPCError(err)
@@ -85,7 +91,8 @@ func (e *execClient) Interactive(ctx context.Context, workspace, sandboxName str
 	if sandboxName == "" {
 		return nil, &StatusError{Code: ErrorInvalidArgument, Message: "sandbox name must not be empty"}
 	}
-	if _, err := e.sandboxes.Get(ctx, workspace, sandboxName); err != nil {
+	owner, err := sandboxOwner(ctx, e.sandboxes, workspace, sandboxName)
+	if err != nil {
 		return nil, err
 	}
 	var opt *ExecOptions
@@ -93,22 +100,17 @@ func (e *execClient) Interactive(ctx context.Context, workspace, sandboxName str
 		opt = &opts[0]
 	}
 
+	startReq := converter.ExecInteractiveRequestToProto(sandboxName, command, cols, rows, opt)
+	startReq.WorkspaceScope = namedWorkspaceScope(workspace)
+
 	streamCtx, cancel := context.WithCancel(ctx)
-	stream, err := e.client.ExecSandboxInteractive(streamCtx)
+	start := &pb.ExecSandboxInput{Payload: &pb.ExecSandboxInput_Start{Start: startReq}}
+	stream, err := openBidiRouted(streamCtx, owner, start, e.client.ExecSandboxInteractive)
 	if err != nil {
 		cancel()
 		return nil, converter.FromGRPCError(err)
 	}
 
-	startReq := converter.ExecInteractiveRequestToProto(sandboxName, command, cols, rows, opt)
-	startReq.WorkspaceScope = namedWorkspaceScope(workspace)
-	if sendErr := stream.Send(&pb.ExecSandboxInput{
-		Payload: &pb.ExecSandboxInput_Start{Start: startReq},
-	}); sendErr != nil {
-		cancel()
-		return nil, converter.FromGRPCError(sendErr)
-	}
-
 	return newInteractiveSession(streamCtx, cancel, stream), nil
 }
 

sdk/go/openshell/v1/routing.go

@@ -0,0 +1,95 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package v1
+
+import (
+	"context"
+	"io"
+	"strings"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+)
+
+// Metadata that lets an edge proxy send a sandbox's long-lived calls straight
+// to the gateway replica holding its supervisor session. The headers and the
+// retry rule are documented in docs/kubernetes/ingress.mdx.
+const (
+	ownerReplicaHeader = "x-openshell-owner"
+	ownerRequestHeader = "x-openshell-want-owner"
+	routeReplicaHeader = "x-openshell-replica"
+)
+
+// ownerGetter is implemented by the gateway-backed sandbox client.
+type ownerGetter interface {
+	getWithOwner(ctx context.Context, workspace, name string) (*Sandbox, string, error)
+}
+
+// sandboxOwner checks that the sandbox exists and returns the replica that
+// owns its supervisor session, or "" when the gateway reports none.
+func sandboxOwner(ctx context.Context, sandboxes SandboxInterface, workspace, name string) (string, error) {
+	if g, ok := sandboxes.(ownerGetter); ok {
+		_, owner, err := g.getWithOwner(ctx, workspace, name)
+		return owner, err
+	}
+	_, err := sandboxes.Get(ctx, workspace, name)
+	return "", err
+}
+
+// openRouted opens a stream tagged with the owner replica. It opens once more
+// untagged when the proxy reports that the owner's route has no ready endpoints.
+// A stale owner is safe: any replica relays the call to the real owner.
+// With an owner, it returns only after the gateway answers, as the CLI does.
+func openRouted[S grpc.ClientStream](ctx context.Context, owner string, open func(context.Context) (S, error)) (S, error) {
+	if owner == "" {
+		return open(ctx)
+	}
+	stream, err := answered(open(metadata.AppendToOutgoingContext(ctx, routeReplicaHeader, owner)))
+	if retryUnrouted(err) {
+		return answered(open(ctx))
+	}
+	return stream, err
+}
+
+// answered waits for the response headers. Nil metadata means the stream ended
+// without headers, and RecvMsg then returns its status.
+func answered[S grpc.ClientStream](stream S, err error) (S, error) {
+	if err == nil {
+		var md metadata.MD
+		if md, err = stream.Header(); err == nil && md == nil {
+			err = stream.RecvMsg(nil)
+		}
+	}
+	return stream, err
+}
+
+// openBidiRouted is openRouted for a bidi stream whose first frame is sent on open.
+func openBidiRouted[I, O any](ctx context.Context, owner string, first *I, open func(context.Context, ...grpc.CallOption) (grpc.BidiStreamingClient[I, O], error)) (grpc.BidiStreamingClient[I, O], error) {
+	return openRouted(ctx, owner, func(ctx context.Context) (grpc.BidiStreamingClient[I, O], error) {
+		s, err := open(ctx)
+		if err != nil {
+			return nil, err
+		}
+		// io.EOF means the stream already ended; answered reports its status.
+		if err := s.Send(first); err == io.EOF {
+			return answered(s, nil)
+		} else if err != nil {
+			return s, err
+		}
+		return s, nil
+	})
+}
+
+// retryUnrouted matches what Envoy Gateway answers for a replica route with no
+// ready endpoints: UNAVAILABLE with no message. Gateway errors always carry one.
+func retryUnrouted(err error) bool {
+	s, ok := status.FromError(err)
+	if !ok || s.Code() != codes.Unavailable {
+		return false
+	}
+	msg := s.Message()
+	return msg == "" || strings.Contains(msg, "no healthy upstream") || strings.Contains(msg, "connection failure")
+}

sdk/go/openshell/v1/routing_test.go

@@ -0,0 +1,192 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package v1
+
+import (
+	"context"
+	"io"
+	"net"
+	"sync"
+	"testing"
+
+	pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+// routingServer reports owner from GetSandbox and records the replica header
+// of every streaming call ("" when unrouted). Routed calls fail with routedErr.
+type routingServer struct {
+	pb.UnimplementedOpenShellServer
+	owner     string
+	routedErr error
+
+	mu       sync.Mutex
+	wantedBy []string
+	routes   []string
+}
+
+func (s *routingServer) GetSandbox(ctx context.Context, _ *pb.GetSandboxRequest) (*pb.SandboxResponse, error) {
+	md, _ := metadata.FromIncomingContext(ctx)
+	s.mu.Lock()
+	s.wantedBy = append(s.wantedBy, firstValue(md, ownerRequestHeader))
+	s.mu.Unlock()
+	if s.owner != "" {
+		_ = grpc.SetHeader(ctx, metadata.Pairs(ownerReplicaHeader, s.owner))
+	}
+	return &pb.SandboxResponse{Sandbox: &pb.Sandbox{}}, nil
+}
+
+// route records the call and returns the error a routed call should fail with.
+func (s *routingServer) route(ctx context.Context) error {
+	md, _ := metadata.FromIncomingContext(ctx)
+	replica := firstValue(md, routeReplicaHeader)
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.routes = append(s.routes, replica)
+	if replica != "" {
+		return s.routedErr
+	}
+	return nil
+}
+
+func (s *routingServer) ExecSandbox(_ *pb.ExecSandboxRequest, stream grpc.ServerStreamingServer[pb.ExecSandboxEvent]) error {
+	if err := s.route(stream.Context()); err != nil {
+		return err
+	}
+	return stream.Send(exitEvent())
+}
+
+func (s *routingServer) ExecSandboxInteractive(stream grpc.BidiStreamingServer[pb.ExecSandboxInput, pb.ExecSandboxEvent]) error {
+	if _, err := stream.Recv(); err != nil {
+		return err
+	}
+	if err := s.route(stream.Context()); err != nil {
+		return err
+	}
+	return stream.Send(exitEvent())
+}
+
+func (s *routingServer) ForwardTcp(stream grpc.BidiStreamingServer[pb.TcpForwardFrame, pb.TcpForwardFrame]) error { //nolint:revive // proto-generated method name
+	if _, err := stream.Recv(); err != nil {
+		return err
+	}
+	if err := s.route(stream.Context()); err != nil {
+		return err
+	}
+	return stream.Send(&pb.TcpForwardFrame{Payload: &pb.TcpForwardFrame_Data{Data: []byte("pong")}})
+}
+
+func exitEvent() *pb.ExecSandboxEvent {
+	return &pb.ExecSandboxEvent{Payload: &pb.ExecSandboxEvent_Exit{Exit: &pb.ExecSandboxExit{}}}
+}
+
+func firstValue(md metadata.MD, key string) string {
+	if v := md.Get(key); len(v) > 0 {
+		return v[0]
+	}
+	return ""
+}
+
+func routingClients(t *testing.T, srv *routingServer) (*execClient, *tcpClient) {
+	t.Helper()
+	lis := bufconn.Listen(bufSize)
+	gs := grpc.NewServer()
+	pb.RegisterOpenShellServer(gs, srv)
+	go func() { _ = gs.Serve(lis) }()
+	conn, err := grpc.NewClient("passthrough:///bufconn",
+		grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return lis.Dial() }),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = conn.Close()
+		gs.Stop()
+	})
+	sandboxes := newSandboxClient(conn)
+	return newExecClient(conn, sandboxes), newTCPClient(conn, sandboxes, nil)
+}
+
+func TestExecRoutesToOwner(t *testing.T) {
+	srv := &routingServer{owner: "openshell-1"}
+	exec, _ := routingClients(t, srv)
+
+	_, err := exec.Run(context.Background(), "default", "box", []string{"true"})
+	require.NoError(t, err)
+
+	assert.Equal(t, []string{"1"}, srv.wantedBy)
+	assert.Equal(t, []string{"openshell-1"}, srv.routes)
+}
+
+func TestExecWithoutOwnerIsUnrouted(t *testing.T) {
+	srv := &routingServer{}
+	exec, _ := routingClients(t, srv)
+
+	_, err := exec.Run(context.Background(), "default", "box", []string{"true"})
+	require.NoError(t, err)
+
+	assert.Equal(t, []string{""}, srv.routes)
+}
+
+func TestRoutedCallsRetryOnceWithoutHintWhenOwnerHasNoEndpoints(t *testing.T) {
+	srv := &routingServer{owner: "openshell-1", routedErr: status.Error(codes.Unavailable, "")}
+	exec, tcp := routingClients(t, srv)
+	ctx := context.Background()
+
+	_, err := exec.Run(ctx, "default", "box", []string{"true"})
+	require.NoError(t, err)
+
+	stream, err := exec.Stream(ctx, "default", "box", []string{"true"})
+	require.NoError(t, err)
+	code, err := stream.ExitCode()
+	require.NoError(t, err)
+	assert.Zero(t, code)
+
+	session, err := exec.Interactive(ctx, "default", "box", []string{"sh"}, 80, 24)
+	require.NoError(t, err)
+	_, err = session.ExitCode()
+	require.NoError(t, err)
+
+	conn, err := tcp.Forward(ctx, "default", "box", 8080)
+	require.NoError(t, err)
+	buf := make([]byte, 4)
+	_, err = io.ReadFull(conn, buf)
+	require.NoError(t, err)
+	assert.Equal(t, "pong", string(buf))
+	_ = conn.Close()
+
+	// Each call is attempted routed, then once unrouted.
+	assert.Equal(t, []string{
+		"openshell-1", "", "openshell-1", "", "openshell-1", "", "openshell-1", "",
+	}, srv.routes)
+}
+
+func TestRoutedCallDoesNotRetryOnGatewayError(t *testing.T) {
+	srv := &routingServer{owner: "openshell-1", routedErr: status.Error(codes.Unavailable, "supervisor not connected")}
+	exec, tcp := routingClients(t, srv)
+	ctx := context.Background()
+
+	_, err := exec.Run(ctx, "default", "box", []string{"true"})
+	require.Error(t, err)
+	_, err = tcp.Forward(ctx, "default", "box", 8080)
+	require.Error(t, err)
+
+	assert.Equal(t, []string{"openshell-1", "openshell-1"}, srv.routes)
+}
+
+func TestRetryUnrouted(t *testing.T) {
+	assert.True(t, retryUnrouted(status.Error(codes.Unavailable, "")))
+	assert.True(t, retryUnrouted(status.Error(codes.Unavailable, "no healthy upstream")))
+	assert.True(t, retryUnrouted(status.Error(codes.Unavailable, "upstream connect error: connection failure")))
+	assert.False(t, retryUnrouted(status.Error(codes.Unavailable, "supervisor not connected")))
+	assert.False(t, retryUnrouted(status.Error(codes.NotFound, "")))
+	assert.False(t, retryUnrouted(io.EOF))
+	assert.False(t, retryUnrouted(nil))
+}

sdk/go/openshell/v1/sandbox_client.go

@@ -13,6 +13,7 @@ import (
 	"github.com/NVIDIA/OpenShell/sdk/go/openshell/v1/types"
 	pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 const defaultPollInterval = 500 * time.Millisecond
@@ -116,10 +117,28 @@ func validateTemplateCreateSpec(spec *SandboxSpec) error {
 }
 
 func (s *sandboxClient) Get(ctx context.Context, workspace, name string) (*Sandbox, error) {
+	return s.get(ctx, workspace, name)
+}
+
+// getWithOwner is Get that also asks the gateway for the owner replica.
+func (s *sandboxClient) getWithOwner(ctx context.Context, workspace, name string) (*Sandbox, string, error) {
+	var md metadata.MD
+	sb, err := s.get(metadata.AppendToOutgoingContext(ctx, ownerRequestHeader, "1"), workspace, name, grpc.Header(&md))
+	if err != nil {
+		return nil, "", err
+	}
+	var owner string
+	if v := md.Get(ownerReplicaHeader); len(v) > 0 {
+		owner = v[0]
+	}
+	return sb, owner, nil
+}
+
+func (s *sandboxClient) get(ctx context.Context, workspace, name string, opts ...grpc.CallOption) (*Sandbox, error) {
 	resp, err := s.client.GetSandbox(ctx, &pb.GetSandboxRequest{
 		Name:           name,
 		WorkspaceScope: namedWorkspaceScope(workspace),
-	})
+	}, opts...)
 	if err != nil {
 		return nil, converter.FromGRPCError(err)
 	}

sdk/go/openshell/v1/ssh_client.go

@@ -65,7 +65,8 @@ func (s *sshClient) Tunnel(ctx context.Context, workspace, sandboxName string, p
 			Message: fmt.Sprintf("port must be in range 1-65535, got %d", port),
 		}
 	}
-	if _, err := s.sandboxes.Get(ctx, workspace, sandboxName); err != nil {
+	owner, err := sandboxOwner(ctx, s.sandboxes, workspace, sandboxName)
+	if err != nil {
 		return nil, err
 	}
 
@@ -84,13 +85,6 @@ func (s *sshClient) Tunnel(ctx context.Context, workspace, sandboxName string, p
 		}
 	}()
 
-	streamCtx, cancel := context.WithCancel(ctx)
-	stream, err := s.client.ForwardTcp(streamCtx)
-	if err != nil {
-		cancel()
-		return nil, converter.FromGRPCError(err)
-	}
-
 	initFrame := &pb.TcpForwardFrame{
 		Payload: &pb.TcpForwardFrame_Init{
 			Init: &pb.TcpForwardInit{
@@ -105,7 +99,9 @@ func (s *sshClient) Tunnel(ctx context.Context, workspace, sandboxName string, p
 		},
 	}
 
-	if err := stream.Send(initFrame); err != nil {
+	streamCtx, cancel := context.WithCancel(ctx)
+	stream, err := openBidiRouted(streamCtx, owner, initFrame, s.client.ForwardTcp) //nolint:revive // proto-generated method name
+	if err != nil {
 		cancel()
 		return nil, converter.FromGRPCError(err)
 	}

sdk/go/openshell/v1/tcp_client.go

@@ -37,20 +37,14 @@ func (t *tcpClient) Forward(ctx context.Context, workspace, sandboxName string,
 			Message: fmt.Sprintf("port must be in range 1-65535, got %d", port),
 		}
 	}
-	if _, err := t.sandboxes.Get(ctx, workspace, sandboxName); err != nil {
+	owner, err := sandboxOwner(ctx, t.sandboxes, workspace, sandboxName)
+	if err != nil {
 		return nil, err
 	}
 
 	var cfg forwardConfig
 	options.Apply(&cfg, opts)
 
-	streamCtx, cancel := context.WithCancel(ctx)
-	stream, err := t.client.ForwardTcp(streamCtx)
-	if err != nil {
-		cancel()
-		return nil, converter.FromGRPCError(err)
-	}
-
 	initFrame := &pb.TcpForwardFrame{
 		Payload: &pb.TcpForwardFrame_Init{
 			Init: &pb.TcpForwardInit{
@@ -67,7 +61,9 @@ func (t *tcpClient) Forward(ctx context.Context, workspace, sandboxName string,
 		},
 	}
 
-	if err := stream.Send(initFrame); err != nil {
+	streamCtx, cancel := context.WithCancel(ctx)
+	stream, err := openBidiRouted(streamCtx, owner, initFrame, t.client.ForwardTcp) //nolint:revive // proto-generated method name
+	if err != nil {
 		cancel()
 		return nil, converter.FromGRPCError(err)
 	}