Address comments

internal/codespaces/grpc/client.go

@@ -6,58 +6,71 @@ package grpc
 import (
 	"context"
 	"fmt"
+	"net"
 	"strconv"
 	"time"
 
-	"github.com/cli/cli/v2/pkg/grpc/jupyter"
+	"github.com/cli/cli/v2/internal/codespaces/grpc/jupyter"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/metadata"
 )
 
 const (
-	requestTimeout = 30 * time.Second
+	connectionTimeout = 5 * time.Second
+	requestTimeout    = 30 * time.Second
 )
 
-type GrpcClient struct {
+type Client struct {
 	conn          *grpc.ClientConn
 	token         string
+	listener      net.Listener
 	jupyterClient jupyter.JupyterServerHostClient
 }
 
-func New() *GrpcClient {
-	return &GrpcClient{}
+func NewClient() *Client {
+	return &Client{}
 }
 
 // Connects to the gRPC server on the given port
-func (g *GrpcClient) Connect(ctx context.Context, port int, token string) error {
+func (g *Client) Connect(ctx context.Context, listener net.Listener, port int, token string) error {
 	// Attempt to connect to the given port
-	conn, err := grpc.Dial(fmt.Sprintf("localhost:%d", port), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock())
-
+	conn, err := grpc.Dial(fmt.Sprintf("127.0.0.1:%d", port), grpc.WithTimeout(connectionTimeout), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock())
 	if err != nil {
-		return fmt.Errorf("Failed to connect to the internal server on port %d", port)
+		return err
 	}
 
 	g.conn = conn
 	g.token = token
+	g.listener = listener
 	g.jupyterClient = jupyter.NewJupyterServerHostClient(conn)
 
 	return nil
 }
 
+// Closes the gRPC connection
+func (g *Client) Close() error {
+	// Closing the local listener effectively closes the gRPC connection
+	if err := g.listener.Close(); err != nil {
+		g.conn.Close() // If we fail to close the listener, explicitly close the gRPC connection and ignore any error
+		return fmt.Errorf("failed to close local tcp port listener: %w", err)
+	}
+
+	return nil
+}
+
 // Appends the authentication token to the gRPC context
-func (g *GrpcClient) appendMetadata(ctx context.Context) context.Context {
+func (g *Client) appendMetadata(ctx context.Context) context.Context {
 	return metadata.AppendToOutgoingContext(ctx, "Authorization", "Bearer "+g.token)
 }
 
 // Starts a remote JupyterLab server to allow the user to connect to the codespace via JupyterLab in their browser
-func (g *GrpcClient) GetRunningServer() (int, string, error) {
+func (g *Client) StartJupyterServer() (port int, serverUrl string, err error) {
 	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
 	ctx = g.appendMetadata(ctx)
 	defer cancel()
 
 	response, err := g.jupyterClient.GetRunningServer(ctx, &jupyter.GetRunningServerRequest{})
-
 	if err != nil {
 		return 0, "", fmt.Errorf("failed to invoke JupyterLab RPC: %w", err)
 	}
@@ -66,8 +79,7 @@ func (g *GrpcClient) GetRunningServer() (int, string, error) {
 		return 0, "", fmt.Errorf("failed to start JupyterLab: %s", response.Message)
 	}
 
-	port, err := strconv.Atoi(response.Port)
-
+	port, err = strconv.Atoi(response.Port)
 	if err != nil {
 		return 0, "", fmt.Errorf("failed to parse JupyterLab port: %w", err)
 	}

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