feat: restore files after build failure

internal/librarian/generate.go

@@ -294,14 +294,22 @@ func (r *generateRunner) runBuildCommand(ctx context.Context, libraryID string)
 		State:     r.state,
 	}
 	slog.Info("Performing build for library", "id", libraryID)
-	if err := r.containerClient.Build(ctx, buildRequest); err != nil {
-		return err
+	if containerErr := r.containerClient.Build(ctx, buildRequest); containerErr != nil {
+		if restoreErr := r.restoreLibrary(libraryID); restoreErr != nil {
+			return errors.Join(containerErr, restoreErr)
+		}
+
+		return errors.Join(containerErr)
 	}
 
 	// Read the library state from the response.
-	if _, err := readLibraryState(
-		filepath.Join(buildRequest.RepoDir, config.LibrarianDir, config.BuildResponse)); err != nil {
-		return err
+	if _, responseErr := readLibraryState(
+		filepath.Join(buildRequest.RepoDir, config.LibrarianDir, config.BuildResponse)); responseErr != nil {
+		if restoreErr := r.restoreLibrary(libraryID); restoreErr != nil {
+			return errors.Join(responseErr, restoreErr)
+		}
+
+		return responseErr
 	}
 
 	slog.Info("Build succeeds", "id", libraryID)
@@ -388,6 +396,16 @@ func (r *generateRunner) runConfigureCommand(ctx context.Context) (string, error
 	return libraryState.ID, nil
 }
 
+func (r *generateRunner) restoreLibrary(libraryID string) error {
+	// At this point, we should have a library in the state.
+	library := findLibraryByID(r.state, libraryID)
+	if err := r.repo.Restore(library.SourceRoots); err != nil {
+		return err
+	}
+
+	return r.repo.CleanUntracked(library.SourceRoots)
+}
+
 func setAllAPIStatus(state *config.LibrarianState, status string) {
 	for _, library := range state.Libraries {
 		for _, api := range library.APIs {

internal/librarian/generate_test.go

@@ -110,67 +110,51 @@ func TestRunBuildCommand(t *testing.T) {
 		name           string
 		build          bool
 		libraryID      string
-		repo           gitrepo.Repository
-		state          *config.LibrarianState
 		container      *mockContainerClient
 		wantBuildCalls int
 		wantErr        bool
 	}{
 		{
-			name:           "build flag not specified",
+			name:           "build_flag_not_specified",
 			build:          false,
 			container:      &mockContainerClient{},
 			wantBuildCalls: 0,
 		},
 		{
-			name:      "build with library id",
-			build:     true,
-			libraryID: "some-library",
-			repo:      newTestGitRepo(t),
-			state: &config.LibrarianState{
-				Libraries: []*config.LibraryState{
-					{
-						ID: "some-library",
-					},
-				},
-			},
+			name:           "build_with_library_id",
+			build:          true,
+			libraryID:      "some-library",
 			container:      &mockContainerClient{},
 			wantBuildCalls: 1,
 		},
 		{
-			name:      "build with no library id",
-			build:     true,
-			container: &mockContainerClient{},
+			name:           "build_with_no_library_id",
+			build:          true,
+			container:      &mockContainerClient{},
+			wantBuildCalls: 0,
 		},
 		{
-			name:      "build with no response",
+			name:      "build_with_no_response",
 			build:     true,
 			libraryID: "some-library",
-			repo:      newTestGitRepo(t),
-			state: &config.LibrarianState{
-				Libraries: []*config.LibraryState{
-					{
-						ID: "some-library",
-					},
-				},
-			},
 			container: &mockContainerClient{
 				noBuildResponse: true,
 			},
 			wantBuildCalls: 1,
 		},
 		{
-			name:      "build with error response in response",
+			name:      "build_with_docker_command_error_files_restored",
 			build:     true,
 			libraryID: "some-library",
-			repo:      newTestGitRepo(t),
-			state: &config.LibrarianState{
-				Libraries: []*config.LibraryState{
-					{
-						ID: "some-library",
-					},
-				},
+			container: &mockContainerClient{
+				buildErr: errors.New("simulate build error"),
 			},
+			wantErr: true,
+		},
+		{
+			name:      "build_with_error_response_in_response",
+			build:     true,
+			libraryID: "some-library",
 			container: &mockContainerClient{
 				wantErrorMsg: true,
 			},
@@ -179,20 +163,87 @@ func TestRunBuildCommand(t *testing.T) {
 	} {
 		t.Run(test.name, func(t *testing.T) {
 			t.Parallel()
+			repo := newTestGitRepo(t)
+			state := &config.LibrarianState{
+				Libraries: []*config.LibraryState{
+					{
+						ID: "some-library",
+						SourceRoots: []string{
+							"a/path",
+							"another/path",
+						},
+					},
+				},
+			}
 			r := &generateRunner{
 				build:           test.build,
-				repo:            test.repo,
-				state:           test.state,
+				repo:            repo,
+				state:           state,
 				containerClient: test.container,
 			}
 
+			// Create library files and commit the change.
+			repoDir := r.repo.GetDir()
+			for _, library := range r.state.Libraries {
+				for _, srcPath := range library.SourceRoots {
+					relPath := filepath.Join(repoDir, srcPath)
+					if err := os.MkdirAll(relPath, 0755); err != nil {
+						t.Fatal(err)
+					}
+					file := filepath.Join(relPath, "example.txt")
+					if err := os.WriteFile(file, []byte("old content"), 0755); err != nil {
+						t.Fatal(err)
+					}
+				}
+			}
+			if _, err := r.repo.AddAll(); err != nil {
+				t.Fatal(err)
+			}
+			if err := r.repo.Commit("test commit"); err != nil {
+				t.Fatal(err)
+			}
+			// Modify library files and add untacked files.
+			for _, library := range r.state.Libraries {
+				for _, srcPath := range library.SourceRoots {
+					file := filepath.Join(repoDir, srcPath, "example.txt")
+					if err := os.WriteFile(file, []byte("new content"), 0755); err != nil {
+						t.Fatal(err)
+					}
+
+					newFile := filepath.Join(repoDir, srcPath, "another_example.txt")
+					if err := os.WriteFile(newFile, []byte("new content"), 0755); err != nil {
+						t.Fatal(err)
+					}
+				}
+			}
+
 			err := r.runBuildCommand(context.Background(), test.libraryID)
 			if test.wantErr {
 				if err == nil {
-					t.Fatalf("%s should return error", test.name)
+					t.Fatal(err)
 				}
+				// Verify the library files are restore.
+				for _, library := range r.state.Libraries {
+					for _, srcPath := range library.SourceRoots {
+						file := filepath.Join(repoDir, srcPath, "example.txt")
+						readFile, err := os.ReadFile(file)
+						if err != nil {
+							t.Fatal(err)
+						}
+						if diff := cmp.Diff("old content", string(readFile)); diff != "" {
+							t.Errorf("file content mismatch (-want +got):%s", diff)
+						}
+
+						newFile := filepath.Join(repoDir, srcPath, "another_example.txt")
+						if _, err := os.Stat(newFile); !os.IsNotExist(err) {
+							t.Fatal(err)
+						}
+					}
+				}
+
 				return
 			}
+
 			if err != nil {
 				t.Fatal(err)
 			}

internal/librarian/state.go

@@ -184,7 +184,7 @@ func saveLibrarianState(repoDir string, state *config.LibrarianState) error {
 // readLibraryState reads the library state from a container response, if it exists.
 // If the response file does not exist, readLibraryState succeeds but returns a nil pointer.
 //
-// The response file is removed afterwards.
+// The response file is removed afterward.
 func readLibraryState(jsonFilePath string) (*config.LibraryState, error) {
 	data, err := os.ReadFile(jsonFilePath)
 	defer func() {