Add command to open codespaces in JupyterLab

Add command to open codespaces in JupyterLab

pkg/cmd/codespace/common.go

@@ -14,8 +14,10 @@ import (
 
 	"github.com/AlecAivazis/survey/v2"
 	"github.com/AlecAivazis/survey/v2/terminal"
+	"github.com/cli/cli/v2/internal/codespaces"
 	"github.com/cli/cli/v2/internal/codespaces/api"
 	"github.com/cli/cli/v2/pkg/iostreams"
+	"github.com/cli/cli/v2/pkg/liveshare"
 	"github.com/spf13/cobra"
 	"golang.org/x/term"
 )
@@ -58,6 +60,37 @@ func (a *App) StopProgressIndicator() {
 	a.io.StopProgressIndicator()
 }
 
+// Connects to a codespace using Live Share and returns that session
+func startLiveShareSession(ctx context.Context, codespace *api.Codespace, a *App, debug bool, debugFile string) (session *liveshare.Session, err error) {
+	// While connecting, ensure in the background that the user has keys installed.
+	// That lets us report a more useful error message if they don't.
+	authkeys := make(chan error, 1)
+	go func() {
+		authkeys <- checkAuthorizedKeys(ctx, a.apiClient)
+	}()
+
+	liveshareLogger := noopLogger()
+	if debug {
+		debugLogger, err := newFileLogger(debugFile)
+		if err != nil {
+			return nil, fmt.Errorf("couldn't create file logger: %w", err)
+		}
+		defer safeClose(debugLogger, &err)
+
+		liveshareLogger = debugLogger.Logger
+		a.errLogger.Printf("Debug file located at: %s", debugLogger.Name())
+	}
+
+	session, err = codespaces.ConnectToLiveshare(ctx, a, liveshareLogger, a.apiClient, codespace)
+	if err != nil {
+		if authErr := <-authkeys; authErr != nil {
+			return nil, fmt.Errorf("failed to fetch authorization keys: %w", authErr)
+		}
+		return nil, fmt.Errorf("failed to connect to Live Share: %w", err)
+	}
+	return
+}
+
 //go:generate moq -fmt goimports -rm -skip-ensure -out mock_api.go . apiClient
 type apiClient interface {
 	GetUser(ctx context.Context) (*api.User, error)

pkg/cmd/codespace/jupyter.go

@@ -0,0 +1,82 @@
+package codespace
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"strings"
+
+	"github.com/cli/cli/v2/pkg/liveshare"
+	"github.com/spf13/cobra"
+)
+
+func newJupyterCmd(app *App) *cobra.Command {
+	var codespace string
+
+	jupyterCmd := &cobra.Command{
+		Use:   "jupyter",
+		Short: "Open a codespace in JupyterLab",
+		Args:  noArgsConstraint,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return app.Jupyter(cmd.Context(), codespace)
+		},
+	}
+
+	jupyterCmd.Flags().StringVarP(&codespace, "codespace", "c", "", "Name of the codespace")
+
+	return jupyterCmd
+}
+
+func (a *App) Jupyter(ctx context.Context, codespaceName string) (err error) {
+	// Ensure all child tasks (e.g. port forwarding) terminate before return.
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	codespace, err := getOrChooseCodespace(ctx, a.apiClient, codespaceName)
+	if err != nil {
+		return err
+	}
+
+	session, err := startLiveShareSession(ctx, codespace, a, false, "")
+	if err != nil {
+		return err
+	}
+	defer safeClose(session, &err)
+
+	a.StartProgressIndicatorWithLabel("Starting JupyterLab on codespace")
+	serverPort, serverUrl, err := session.StartJupyterServer(ctx)
+	a.StopProgressIndicator()
+	if err != nil {
+		return fmt.Errorf("failed to start JupyterLab server: %w", err)
+	}
+
+	// Pass 0 to pick a random port
+	listen, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", 0))
+	if err != nil {
+		return err
+	}
+	defer listen.Close()
+	destPort := listen.Addr().(*net.TCPAddr).Port
+
+	tunnelClosed := make(chan error, 1)
+	go func() {
+		fwd := liveshare.NewPortForwarder(session, "jupyter", serverPort, true)
+		tunnelClosed <- fwd.ForwardToListener(ctx, listen) // always non-nil
+	}()
+
+	// Server URL contains an authentication token that must be preserved
+	targetUrl := strings.Replace(serverUrl, fmt.Sprintf("%d", serverPort), fmt.Sprintf("%d", destPort), 1)
+	err = a.browser.Browse(targetUrl)
+	if err != nil {
+		return fmt.Errorf("failed to open JupyterLab in browser: %w", err)
+	}
+
+	fmt.Fprintln(a.io.Out, targetUrl)
+
+	select {
+	case err := <-tunnelClosed:
+		return fmt.Errorf("tunnel closed: %w", err)
+	case <-ctx.Done():
+		return nil // success
+	}
+}

pkg/cmd/codespace/logs.go

@@ -41,20 +41,11 @@ func (a *App) Logs(ctx context.Context, codespaceName string, follow bool) (err
 		return err
 	}
 
-	authkeys := make(chan error, 1)
+	session, err := startLiveShareSession(ctx, codespace, a, false, "")
-	go func() {
-		authkeys <- checkAuthorizedKeys(ctx, a.apiClient)
-	}()
-
-	session, err := codespaces.ConnectToLiveshare(ctx, a, noopLogger(), a.apiClient, codespace)
 	if err != nil {
-		return fmt.Errorf("connecting to codespace: %w", err)
-	}
-	defer safeClose(session, &err)
-
-	if err := <-authkeys; err != nil {
 		return err
 	}
+	defer safeClose(session, &err)
 
 	// Ensure local port is listening before client (getPostCreateOutput) connects.
 	listen, err := net.Listen("tcp", "127.0.0.1:0") // arbitrary port

pkg/cmd/codespace/ports.go

@@ -55,9 +55,9 @@ func (a *App) ListPorts(ctx context.Context, codespaceName string, exporter cmdu
 
 	devContainerCh := getDevContainer(ctx, a.apiClient, codespace)
 
-	session, err := codespaces.ConnectToLiveshare(ctx, a, noopLogger(), a.apiClient, codespace)
+	session, err := startLiveShareSession(ctx, codespace, a, false, "")
 	if err != nil {
-		return fmt.Errorf("error connecting to codespace: %w", err)
+		return err
 	}
 	defer safeClose(session, &err)
 

pkg/cmd/codespace/root.go

@@ -14,6 +14,7 @@ func NewRootCmd(app *App) *cobra.Command {
 	root.AddCommand(newCreateCmd(app))
 	root.AddCommand(newEditCmd(app))
 	root.AddCommand(newDeleteCmd(app))
+	root.AddCommand(newJupyterCmd(app))
 	root.AddCommand(newListCmd(app))
 	root.AddCommand(newLogsCmd(app))
 	root.AddCommand(newPortsCmd(app))

pkg/cmd/codespace/ssh.go

@@ -115,36 +115,14 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	// While connecting, ensure in the background that the user has keys installed.
-	// That lets us report a more useful error message if they don't.
-	authkeys := make(chan error, 1)
-	go func() {
-		authkeys <- checkAuthorizedKeys(ctx, a.apiClient)
-	}()
-
 	codespace, err := getOrChooseCodespace(ctx, a.apiClient, opts.codespace)
 	if err != nil {
 		return err
 	}
 
-	liveshareLogger := noopLogger()
+	session, err := startLiveShareSession(ctx, codespace, a, opts.debug, opts.debugFile)
-	if opts.debug {
-		debugLogger, err := newFileLogger(opts.debugFile)
-		if err != nil {
-			return fmt.Errorf("error creating debug logger: %w", err)
-		}
-		defer safeClose(debugLogger, &err)
-
-		liveshareLogger = debugLogger.Logger
-		a.errLogger.Printf("Debug file located at: %s", debugLogger.Name())
-	}
-
-	session, err := codespaces.ConnectToLiveshare(ctx, a, liveshareLogger, a.apiClient, codespace)
 	if err != nil {
-		if authErr := <-authkeys; authErr != nil {
+		return err
-			return authErr
-		}
-		return fmt.Errorf("error connecting to codespace: %w", err)
 	}
 	defer safeClose(session, &err)
 
@@ -208,11 +186,10 @@ func (a *App) SSH(ctx context.Context, sshArgs []string, opts sshOptions) (err e
 	}
 }
 
-func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) error {
+func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) (err error) {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	var err error
 	var csList []*api.Codespace
 	if opts.codespace == "" {
 		a.StartProgressIndicatorWithLabel("Fetching codespaces")
@@ -253,7 +230,7 @@ func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) error {
 			if err != nil {
 				result.err = fmt.Errorf("error connecting to codespace: %w", err)
 			} else {
-				defer session.Close()
+				defer safeClose(session, &err)
 
 				_, result.user, err = session.StartSSHServer(ctx)
 				if err != nil {

pkg/liveshare/session.go

@@ -62,6 +62,32 @@ func (s *Session) StartSSHServer(ctx context.Context) (int, string, error) {
 	return port, response.User, nil
 }
 
+// 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) {
+	var response struct {
+		Result    bool   `json:"result"`
+		Message   string `json:"message"`
+		Port      string `json:"port"`
+		ServerUrl string `json:"serverUrl"`
+	}
+
+	if err := s.rpc.do(ctx, "IJupyterServerHostService.getRunningServer", []string{}, &response); err != nil {
+		return 0, "", fmt.Errorf("failed to invoke JupyterLab RPC: %w", err)
+	}
+
+	if !response.Result {
+		return 0, "", fmt.Errorf("failed to start JupyterLab: %s", response.Message)
+	}
+
+	port, err := strconv.Atoi(response.Port)
+	if err != nil {
+		return 0, "", fmt.Errorf("failed to parse JupyterLab port: %w", err)
+	}
+
+	return port, response.ServerUrl, nil
+}
+
 // heartbeat runs until context cancellation, periodically checking whether there is a
 // reason to keep the connection alive, and if so, notifying the Live Share host to do so.
 // Heartbeat ensures it does not send more than one request every "interval" to ratelimit