Address comments

pkg/liveshare/client.go

@@ -20,7 +20,7 @@ import (
 	"strings"
 	"time"
 
-	"github.com/cli/cli/v2/pkg/grpc"
+	"github.com/cli/cli/v2/internal/codespaces/grpc"
 	"github.com/opentracing/opentracing-go"
 )
 
@@ -119,7 +119,7 @@ func Connect(ctx context.Context, opts Options) (*Session, error) {
 	s := &Session{
 		ssh:             ssh,
 		rpc:             rpc,
-		grpc:            grpc.New(),
+		grpc:            grpc.NewClient(),
 		clientName:      opts.ClientName,
 		keepAliveReason: make(chan string, 1),
 		logger:          opts.Logger,
@@ -127,7 +127,10 @@ func Connect(ctx context.Context, opts Options) (*Session, error) {
 	go s.heartbeat(ctx, 1*time.Minute)
 
 	// Connect to the gRPC server so we can make requests anywhere we have access to the session
-	s.connectToGrpcServer(ctx, opts.SessionToken)
+	err = s.connectToGrpcServer(ctx, opts.SessionToken)
+	if err != nil {
+		return nil, fmt.Errorf("error connecting to internal server: %w", err)
+	}
 
 	return s, nil
 }
@@ -136,7 +139,7 @@ func Connect(ctx context.Context, opts Options) (*Session, error) {
 func (s *Session) connectToGrpcServer(ctx context.Context, token string) error {
 	listen, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", 0))
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to listen to local port over tcp: %w", err)
 	}
 
 	// Tunnel the remote gRPC server port to the local port
@@ -148,10 +151,10 @@ func (s *Session) connectToGrpcServer(ctx context.Context, token string) error {
 	}()
 
 	// Make a connection to the gRPC server
-	err = s.grpc.Connect(ctx, localGrpcServerPort, token)
+	err = s.grpc.Connect(ctx, listen, localGrpcServerPort, token)
 
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to establish connection on port %d: %w", localGrpcServerPort, err)
 	}
 
 	select {

pkg/liveshare/session.go

@@ -8,7 +8,7 @@ import (
 	"strings"
 	"time"
 
-	"github.com/cli/cli/v2/pkg/grpc"
+	"github.com/cli/cli/v2/internal/codespaces/grpc"
 	"github.com/opentracing/opentracing-go"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/sync/errgroup"
@@ -24,7 +24,7 @@ type ChannelID struct {
 type Session struct {
 	ssh  *sshSession
 	rpc  *rpcClient
-	grpc *grpc.GrpcClient
+	grpc *grpc.Client
 
 	clientName      string
 	keepAliveReason chan string
@@ -45,6 +45,11 @@ func (s *Session) Close() error {
 		return fmt.Errorf("error while closing Live Share session: %w", err)
 	}
 
+	// Close the connection to the gRPC server
+	if err := s.grpc.Close(); err != nil {
+		return fmt.Errorf("error while closing internal server connection: %w", err)
+	}
+
 	return nil
 }
 
@@ -102,7 +107,7 @@ func (s *Session) StartSSHServerWithOptions(ctx context.Context, options StartSS
 // StartJupyterServer starts a Juypyter server in the container and returns
 // the port on which it listens and the server URL.
 func (s *Session) StartJupyterServer(ctx context.Context) (int, string, error) {
-	return s.grpc.GetRunningServer()
+	return s.grpc.StartJupyterServer()
 }
 
 // heartbeat runs until context cancellation, periodically checking whether there is a