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>

Author: FrostGod

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 the Rust and Go 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

@@ -104,17 +104,8 @@ func (e *execClient) Interactive(ctx context.Context, workspace, sandboxName str
 	startReq.WorkspaceScope = namedWorkspaceScope(workspace)
 
 	streamCtx, cancel := context.WithCancel(ctx)
-	stream, err := openRouted(streamCtx, owner, func(ctx context.Context) (grpc.BidiStreamingClient[pb.ExecSandboxInput, pb.ExecSandboxEvent], error) {
-		s, err := e.client.ExecSandboxInteractive(ctx)
-		if err != nil {
-			return nil, err
-		}
-		// io.EOF means the stream already ended; its status comes from Recv.
-		if err := s.Send(&pb.ExecSandboxInput{Payload: &pb.ExecSandboxInput_Start{Start: startReq}}); err != nil && err != io.EOF {
-			return s, err
-		}
-		return s, nil
-	})
+	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)

sdk/go/openshell/v1/routing.go

@@ -8,15 +8,15 @@ import (
 	"io"
 	"strings"
 
-	pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1"
 	"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.
+// 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"
@@ -42,34 +42,41 @@ func sandboxOwner(ctx context.Context, sandboxes SandboxInterface, workspace, na
 // 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 := open(metadata.AppendToOutgoingContext(ctx, routeReplicaHeader, owner))
+	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 {
-		// Header blocks until the proxy answers; nil metadata means the stream
-		// ended without headers, and RecvMsg then returns its status.
 		var md metadata.MD
 		if md, err = stream.Header(); err == nil && md == nil {
 			err = stream.RecvMsg(nil)
 		}
 	}
-	if retryUnrouted(err) {
-		return open(ctx)
-	}
 	return stream, err
 }
 
-// openForward opens a ForwardTcp stream routed to owner and sends its init frame.
-func openForward(ctx context.Context, client pb.OpenShellClient, owner string, init *pb.TcpForwardFrame) (grpc.BidiStreamingClient[pb.TcpForwardFrame, pb.TcpForwardFrame], error) {
-	return openRouted(ctx, owner, func(ctx context.Context) (grpc.BidiStreamingClient[pb.TcpForwardFrame, pb.TcpForwardFrame], error) {
-		s, err := client.ForwardTcp(ctx)
+// 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; its status comes from Recv.
-		if err := s.Send(init); err != nil && err != io.EOF {
+		// 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

sdk/go/openshell/v1/sandbox_client.go

@@ -117,32 +117,32 @@ func validateTemplateCreateSpec(spec *SandboxSpec) error {
 }
 
 func (s *sandboxClient) Get(ctx context.Context, workspace, name string) (*Sandbox, error) {
-	resp, err := s.client.GetSandbox(ctx, &pb.GetSandboxRequest{
-		Name:           name,
-		WorkspaceScope: namedWorkspaceScope(workspace),
-	})
-	if err != nil {
-		return nil, converter.FromGRPCError(err)
-	}
-	return converter.SandboxFromProto(resp.GetSandbox()), nil
+	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
-	resp, err := s.client.GetSandbox(
-		metadata.AppendToOutgoingContext(ctx, ownerRequestHeader, "1"),
-		&pb.GetSandboxRequest{Name: name, WorkspaceScope: namedWorkspaceScope(workspace)},
-		grpc.Header(&md),
-	)
+	sb, err := s.get(metadata.AppendToOutgoingContext(ctx, ownerRequestHeader, "1"), workspace, name, grpc.Header(&md))
 	if err != nil {
-		return nil, "", converter.FromGRPCError(err)
+		return nil, "", err
 	}
 	var owner string
 	if v := md.Get(ownerReplicaHeader); len(v) > 0 {
 		owner = v[0]
 	}
-	return converter.SandboxFromProto(resp.GetSandbox()), owner, nil
+	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)
+	}
+	return converter.SandboxFromProto(resp.GetSandbox()), nil
 }
 
 func (s *sandboxClient) List(workspace string, opts ...ListOptions) (*Pager[*Sandbox], error) {

sdk/go/openshell/v1/ssh_client.go

@@ -100,7 +100,7 @@ func (s *sshClient) Tunnel(ctx context.Context, workspace, sandboxName string, p
 	}
 
 	streamCtx, cancel := context.WithCancel(ctx)
-	stream, err := openForward(streamCtx, s.client, owner, initFrame)
+	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

@@ -62,7 +62,7 @@ func (t *tcpClient) Forward(ctx context.Context, workspace, sandboxName string,
 	}
 
 	streamCtx, cancel := context.WithCancel(ctx)
-	stream, err := openForward(streamCtx, t.client, owner, initFrame)
+	stream, err := openBidiRouted(streamCtx, owner, initFrame, t.client.ForwardTcp) //nolint:revive // proto-generated method name
 	if err != nil {
 		cancel()
 		return nil, converter.FromGRPCError(err)