From da91216c31a864c1ad13595d9e5635ec33848bb2 Mon Sep 17 00:00:00 2001
From: JP Ungaretti <hello@jungaretti.com>
Date: Sat, 8 Oct 2022 00:04:45 +0000
Subject: [PATCH] Add new Rebuild function

---
 pkg/cmd/codespace/common.go   |  1 +
 pkg/liveshare/session.go      | 14 ++++++++++++++
 pkg/liveshare/session_test.go | 18 ++++++++++++++++++
 3 files changed, 33 insertions(+)

diff --git a/pkg/cmd/codespace/common.go b/pkg/cmd/codespace/common.go
index 46e69d97c..a2d1a2316 100644
--- a/pkg/cmd/codespace/common.go
+++ b/pkg/cmd/codespace/common.go
@@ -66,6 +66,7 @@ type liveshareSession interface {
 	StartSharing(context.Context, string, int) (liveshare.ChannelID, error)
 	StartSSHServer(context.Context) (int, string, error)
 	StartSSHServerWithOptions(context.Context, liveshare.StartSSHServerOptions) (int, string, error)
+	Rebuild(context.Context) error
 }
 
 // Connects to a codespace using Live Share and returns that session
diff --git a/pkg/liveshare/session.go b/pkg/liveshare/session.go
index f912fa5ea..8518afd56 100644
--- a/pkg/liveshare/session.go
+++ b/pkg/liveshare/session.go
@@ -123,6 +123,20 @@ func (s *Session) StartJupyterServer(ctx context.Context) (int, string, error) {
 	return port, response.ServerUrl, nil
 }
 
+func (s *Session) Rebuild(ctx context.Context) error {
+	var success bool
+	err := s.rpc.do(ctx, "IEnvironmentConfigurationService.rebuildContainer", []string{}, &success)
+	if err != nil {
+		return fmt.Errorf("invoking rebuild RPC: %w", err)
+	}
+
+	if !success {
+		return fmt.Errorf("couldn't rebuild codespace")
+	} else {
+		return 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
diff --git a/pkg/liveshare/session_test.go b/pkg/liveshare/session_test.go
index cfe8ccd11..2160caf6e 100644
--- a/pkg/liveshare/session_test.go
+++ b/pkg/liveshare/session_test.go
@@ -399,6 +399,24 @@ func TestSessionHeartbeat(t *testing.T) {
 	}
 }
 
+func TestRebuild(t *testing.T) {
+	getSharedServers := func(conn *jsonrpc2.Conn, req *jsonrpc2.Request) (interface{}, error) {
+		return true, nil
+	}
+	testServer, session, err := makeMockSession(
+		livesharetest.WithService("IEnvironmentConfigurationService.rebuildContainer", getSharedServers),
+	)
+	if err != nil {
+		t.Fatalf("creating mock session: %v", err)
+	}
+	defer testServer.Close()
+
+	err = session.Rebuild(context.Background())
+	if err != nil {
+		t.Fatalf("rebuilding codespace over mock session: %v", err)
+	}
+}
+
 type mockLogger struct {
 	sync.Mutex
 	buf *bytes.Buffer