diff --git a/pkg/cmd/codespace/common.go b/pkg/cmd/codespace/common.go index 003b5043f..202f1447b 100644 --- a/pkg/cmd/codespace/common.go +++ b/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) diff --git a/pkg/cmd/codespace/jupyter.go b/pkg/cmd/codespace/jupyter.go new file mode 100644 index 000000000..77bedb301 --- /dev/null +++ b/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 + } +} diff --git a/pkg/cmd/codespace/logs.go b/pkg/cmd/codespace/logs.go index 6feab3080..b3289de41 100644 --- a/pkg/cmd/codespace/logs.go +++ b/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) - go func() { - authkeys <- checkAuthorizedKeys(ctx, a.apiClient) - }() - - session, err := codespaces.ConnectToLiveshare(ctx, a, noopLogger(), a.apiClient, codespace) + session, err := startLiveShareSession(ctx, codespace, a, false, "") 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 diff --git a/pkg/cmd/codespace/ports.go b/pkg/cmd/codespace/ports.go index f8662ae37..9bbe28f79 100644 --- a/pkg/cmd/codespace/ports.go +++ b/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) diff --git a/pkg/cmd/codespace/root.go b/pkg/cmd/codespace/root.go index 662f4c2de..10dc870f3 100644 --- a/pkg/cmd/codespace/root.go +++ b/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)) diff --git a/pkg/cmd/codespace/ssh.go b/pkg/cmd/codespace/ssh.go index 7721f7977..10ea30473 100644 --- a/pkg/cmd/codespace/ssh.go +++ b/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() - 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) + session, err := startLiveShareSession(ctx, codespace, a, opts.debug, opts.debugFile) if err != nil { - if authErr := <-authkeys; authErr != nil { - return authErr - } - return fmt.Errorf("error connecting to codespace: %w", err) + return 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 { diff --git a/pkg/liveshare/session.go b/pkg/liveshare/session.go index 1bffa096f..e2648c1c8 100644 --- a/pkg/liveshare/session.go +++ b/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