From 18b44f41fee38faff84e4825732ebf1cc411a7b8 Mon Sep 17 00:00:00 2001
From: Andrew Phelps <andrew.phelps@canonical.com>
Date: Wed, 5 Feb 2025 20:08:10 -0500
Subject: [PATCH 1/4] o/snapstate, o/devicestate, tests: support components in
 snapstate.LinkNewBaseOrKernel and snapstate.AddLinkNewBaseOrKernel

---
 .../devicestate/devicestate_remodel_test.go   | 151 ++++++--------
 overlord/snapstate/snapstate.go               | 154 +++++++-------
 overlord/snapstate/snapstate_test.go          | 181 +++++++++++++----
 ...-snapd-remodel-to-installed-kernel-24.json |  38 ++++
 ...-remodel-to-installed-kernel-rev-1-24.json |  42 ++++
 ...-remodel-to-installed-kernel-rev-2-24.json |  42 ++++
 ...-remodel-to-installed-kernel-rev-3-24.json |  42 ++++
 tests/lib/tools/build_kernel_with_comps.sh    |   8 +-
 .../remodel-to-installed-kernel/task.yaml     | 189 ++++++++++++++++++
 9 files changed, 649 insertions(+), 198 deletions(-)
 create mode 100644 tests/lib/assertions/test-snapd-remodel-to-installed-kernel-24.json
 create mode 100644 tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-1-24.json
 create mode 100644 tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-2-24.json
 create mode 100644 tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-3-24.json
 create mode 100644 tests/nested/manual/remodel-to-installed-kernel/task.yaml

diff --git a/overlord/devicestate/devicestate_remodel_test.go b/overlord/devicestate/devicestate_remodel_test.go
index ca81ea8d571..e1e70132b02 100644
--- a/overlord/devicestate/devicestate_remodel_test.go
+++ b/overlord/devicestate/devicestate_remodel_test.go
@@ -3100,8 +3100,8 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnaps(c *C) {
 		c.Logf("%s: %s", t.Kind(), t.Summary())
 	}
 
-	// 3 snaps (2 tasks for each) + assets update and setup from kernel + gadget (3 tasks) + recovery system (2 tasks) + set-model
-	c.Assert(tl, HasLen, 3*2+2+3+2+1)
+	// 3 snaps (2 tasks for each) + assets update + gadget (3 tasks) + recovery system (2 tasks) + set-model
+	c.Assert(tl, HasLen, 3*2+1+3+2+1)
 
 	deviceCtx, err := devicestate.DeviceCtx(s.state, tl[0], nil)
 	c.Assert(err, IsNil)
@@ -3115,26 +3115,23 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnaps(c *C) {
 
 	// check the tasks
 	tPrepareKernel := tl[0]
-	tSetupKernelSnap := tl[1]
-	tUpdateAssetsKernel := tl[2]
-	tLinkKernel := tl[3]
-	tPrepareBase := tl[4]
-	tLinkBase := tl[5]
-	tPrepareGadget := tl[6]
-	tUpdateAssets := tl[7]
-	tUpdateCmdline := tl[8]
-	tValidateApp := tl[9]
-	tInstallApp := tl[10]
-	tCreateRecovery := tl[11]
-	tFinalizeRecovery := tl[12]
-	tSetModel := tl[13]
+	tUpdateAssetsKernel := tl[1]
+	tLinkKernel := tl[2]
+	tPrepareBase := tl[3]
+	tLinkBase := tl[4]
+	tPrepareGadget := tl[5]
+	tUpdateAssets := tl[6]
+	tUpdateCmdline := tl[7]
+	tValidateApp := tl[8]
+	tInstallApp := tl[9]
+	tCreateRecovery := tl[10]
+	tFinalizeRecovery := tl[11]
+	tSetModel := tl[12]
 
 	// check the tasks
 	c.Assert(tPrepareKernel.Kind(), Equals, "prepare-snap")
 	c.Assert(tPrepareKernel.Summary(), Equals, `Prepare snap "pc-kernel-new" (222) for remodel`)
 	c.Assert(tPrepareKernel.WaitTasks(), HasLen, 0)
-	c.Assert(tSetupKernelSnap.Kind(), Equals, "prepare-kernel-snap")
-	c.Assert(tSetupKernelSnap.Summary(), Equals, `Prepare kernel driver tree for "pc-kernel-new" (222) for remodel`)
 	c.Assert(tLinkKernel.Kind(), Equals, "link-snap")
 	c.Assert(tLinkKernel.Summary(), Equals, `Make snap "pc-kernel-new" (222) available to the system during remodel`)
 	c.Assert(tUpdateAssetsKernel.Kind(), Equals, "update-gadget-assets")
@@ -3169,15 +3166,12 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnaps(c *C) {
 	c.Assert(tLinkKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tUpdateAssetsKernel,
 	})
-	c.Assert(tSetupKernelSnap.WaitTasks(), DeepEquals, []*state.Task{
+	c.Assert(tUpdateAssetsKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tPrepareKernel,
 		tValidateApp,
 		tCreateRecovery,
 		tFinalizeRecovery,
 	})
-	c.Assert(tUpdateAssetsKernel.WaitTasks(), DeepEquals, []*state.Task{
-		tSetupKernelSnap,
-	})
 	c.Assert(tPrepareBase.WaitTasks(), DeepEquals, []*state.Task{
 		tPrepareKernel,
 	})
@@ -3208,7 +3202,7 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnaps(c *C) {
 	})
 	// setModel waits for everything in the change
 	c.Assert(tSetModel.WaitTasks(), DeepEquals, []*state.Task{
-		tPrepareKernel, tSetupKernelSnap, tUpdateAssetsKernel,
+		tPrepareKernel, tUpdateAssetsKernel,
 		tLinkKernel, tPrepareBase, tLinkBase,
 		tPrepareGadget, tUpdateAssets, tUpdateCmdline,
 		tValidateApp, tInstallApp,
@@ -3369,8 +3363,8 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnapsChannelSwitch
 		c.Logf("%s: %s", t.Kind(), t.Summary())
 	}
 
-	// 3 snaps (2 tasks for each) + assets update and setup from kernel + gadget (3 tasks) + recovery system (2 tasks) + set-model
-	c.Assert(tl, HasLen, 3*2+2+3+2+1)
+	// 3 snaps (2 tasks for each) + assets update from kernel + gadget (3 tasks) + recovery system (2 tasks) + set-model
+	c.Assert(tl, HasLen, 3*2+1+3+2+1)
 
 	deviceCtx, err := devicestate.DeviceCtx(s.state, tl[0], nil)
 	c.Assert(err, IsNil)
@@ -3384,26 +3378,23 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnapsChannelSwitch
 
 	// check the tasks
 	tSwitchKernel := tl[0]
-	tSetupKernelSnap := tl[1]
-	tUpdateAssetsKernel := tl[2]
-	tLinkKernel := tl[3]
-	tPrepareBase := tl[4]
-	tLinkBase := tl[5]
-	tSwitchGadget := tl[6]
-	tUpdateAssets := tl[7]
-	tUpdateCmdline := tl[8]
-	tValidateApp := tl[9]
-	tInstallApp := tl[10]
-	tCreateRecovery := tl[11]
-	tFinalizeRecovery := tl[12]
-	tSetModel := tl[13]
+	tUpdateAssetsKernel := tl[1]
+	tLinkKernel := tl[2]
+	tPrepareBase := tl[3]
+	tLinkBase := tl[4]
+	tSwitchGadget := tl[5]
+	tUpdateAssets := tl[6]
+	tUpdateCmdline := tl[7]
+	tValidateApp := tl[8]
+	tInstallApp := tl[9]
+	tCreateRecovery := tl[10]
+	tFinalizeRecovery := tl[11]
+	tSetModel := tl[12]
 
 	// check the tasks
 	c.Assert(tSwitchKernel.Kind(), Equals, "switch-snap")
 	c.Assert(tSwitchKernel.Summary(), Equals, `Switch snap "pc-kernel-new" from channel "20/stable" to "20/edge"`)
 	c.Assert(tSwitchKernel.WaitTasks(), HasLen, 0)
-	c.Assert(tSetupKernelSnap.Kind(), Equals, "prepare-kernel-snap")
-	c.Assert(tSetupKernelSnap.Summary(), Equals, `Prepare kernel driver tree for "pc-kernel-new" (222) for remodel`)
 	c.Assert(tLinkKernel.Kind(), Equals, "link-snap")
 	c.Assert(tLinkKernel.Summary(), Equals, `Make snap "pc-kernel-new" (222) available to the system during remodel`)
 	c.Assert(tUpdateAssetsKernel.Kind(), Equals, "update-gadget-assets")
@@ -3438,15 +3429,12 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnapsChannelSwitch
 	c.Assert(tLinkKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tUpdateAssetsKernel,
 	})
-	c.Assert(tSetupKernelSnap.WaitTasks(), DeepEquals, []*state.Task{
+	c.Assert(tUpdateAssetsKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tSwitchKernel,
 		tValidateApp,
 		tCreateRecovery,
 		tFinalizeRecovery,
 	})
-	c.Assert(tUpdateAssetsKernel.WaitTasks(), DeepEquals, []*state.Task{
-		tSetupKernelSnap,
-	})
 	c.Assert(tPrepareBase.WaitTasks(), DeepEquals, []*state.Task{
 		tSwitchKernel,
 	})
@@ -3477,7 +3465,7 @@ func (s *deviceMgrRemodelSuite) TestRemodelOfflineUseInstalledSnapsChannelSwitch
 	})
 	// setModel waits for everything in the change
 	c.Assert(tSetModel.WaitTasks(), DeepEquals, []*state.Task{
-		tSwitchKernel, tSetupKernelSnap, tUpdateAssetsKernel,
+		tSwitchKernel, tUpdateAssetsKernel,
 		tLinkKernel, tPrepareBase, tLinkBase,
 		tSwitchGadget, tUpdateAssets, tUpdateCmdline,
 		tValidateApp, tInstallApp,
@@ -3635,7 +3623,7 @@ func (s *deviceMgrRemodelSuite) TestRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 
 	tl := chg.Tasks()
 	// 2 snaps (2 tasks for each) + assets update and setup from kernel + gadget (3 tasks) + recovery system (2 tasks) + set-model
-	c.Assert(tl, HasLen, 2*2+2+3+2+1)
+	c.Assert(tl, HasLen, 2*2+1+3+2+1)
 
 	deviceCtx, err := devicestate.DeviceCtx(s.state, tl[0], nil)
 	c.Assert(err, IsNil)
@@ -3649,24 +3637,21 @@ func (s *deviceMgrRemodelSuite) TestRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 
 	// check the tasks
 	tPrepareKernel := tl[0]
-	tSetupKernelSnap := tl[1]
-	tUpdateAssetsKernel := tl[2]
-	tLinkKernel := tl[3]
-	tPrepareBase := tl[4]
-	tLinkBase := tl[5]
-	tPrepareGadget := tl[6]
-	tUpdateAssets := tl[7]
-	tUpdateCmdline := tl[8]
-	tCreateRecovery := tl[9]
-	tFinalizeRecovery := tl[10]
-	tSetModel := tl[11]
+	tUpdateAssetsKernel := tl[1]
+	tLinkKernel := tl[2]
+	tPrepareBase := tl[3]
+	tLinkBase := tl[4]
+	tPrepareGadget := tl[5]
+	tUpdateAssets := tl[6]
+	tUpdateCmdline := tl[7]
+	tCreateRecovery := tl[8]
+	tFinalizeRecovery := tl[9]
+	tSetModel := tl[10]
 
 	// check the tasks
 	c.Assert(tPrepareKernel.Kind(), Equals, "prepare-snap")
 	c.Assert(tPrepareKernel.Summary(), Equals, `Prepare snap "pc-kernel-new" (222) for remodel`)
 	c.Assert(tPrepareKernel.WaitTasks(), HasLen, 0)
-	c.Assert(tSetupKernelSnap.Kind(), Equals, "prepare-kernel-snap")
-	c.Assert(tSetupKernelSnap.Summary(), Equals, `Prepare kernel driver tree for "pc-kernel-new" (222) for remodel`)
 	c.Assert(tLinkKernel.Kind(), Equals, "link-snap")
 	c.Assert(tLinkKernel.Summary(), Equals, `Make snap "pc-kernel-new" (222) available to the system during remodel`)
 	c.Assert(tUpdateAssetsKernel.Kind(), Equals, "update-gadget-assets")
@@ -3697,15 +3682,12 @@ func (s *deviceMgrRemodelSuite) TestRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 	c.Assert(tLinkKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tUpdateAssetsKernel,
 	})
-	c.Assert(tSetupKernelSnap.WaitTasks(), DeepEquals, []*state.Task{
+	c.Assert(tUpdateAssetsKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tPrepareKernel,
 		tPrepareGadget,
 		tCreateRecovery,
 		tFinalizeRecovery,
 	})
-	c.Assert(tUpdateAssetsKernel.WaitTasks(), DeepEquals, []*state.Task{
-		tSetupKernelSnap,
-	})
 	c.Assert(tPrepareBase.WaitTasks(), DeepEquals, []*state.Task{
 		tPrepareKernel,
 	})
@@ -3736,7 +3718,7 @@ func (s *deviceMgrRemodelSuite) TestRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 	})
 	// setModel waits for everything in the change
 	c.Assert(tSetModel.WaitTasks(), DeepEquals, []*state.Task{
-		tPrepareKernel, tSetupKernelSnap, tUpdateAssetsKernel,
+		tPrepareKernel, tUpdateAssetsKernel,
 		tLinkKernel, tPrepareBase, tLinkBase,
 		tPrepareGadget, tUpdateAssets, tUpdateCmdline,
 		tCreateRecovery, tFinalizeRecovery,
@@ -3922,14 +3904,17 @@ func (s *deviceMgrRemodelSuite) testRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 	for _, alreadyInstalledName := range []string{"pc-kernel-new", "core24-new", "pc-new"} {
 		snapYaml := "name: pc-kernel-new\nversion: 1\ntype: kernel\n"
 		channel := "other/edge"
+		rev := snap.R(222)
 		if alreadyInstalledName == "core24-new" {
 			snapYaml = "name: core24-new\nversion: 1\ntype: base\n"
+			rev = snap.R(223)
 		} else if alreadyInstalledName == "pc-new" {
 			snapYaml = "name: pc-new\nversion: 1\ntype: gadget\nbase: core24-new\n"
+			rev = snap.R(224)
 		}
 		si := &snap.SideInfo{
 			RealName: alreadyInstalledName,
-			Revision: snap.R(222),
+			Revision: rev,
 			SnapID:   snaptest.AssertedSnapID(alreadyInstalledName),
 		}
 		info := snaptest.MakeSnapFileAndDir(c, snapYaml, nil, si)
@@ -3997,10 +3982,10 @@ func (s *deviceMgrRemodelSuite) testRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 	}
 
 	tl := chg.Tasks()
-	// 2 snaps with (snap switch channel + link snap) + assets update and setup
-	// for the kernel snap + gadget snap (switch channel, assets update, cmdline update) +
-	// recovery system (2 tasks) + set-model
-	c.Assert(tl, HasLen, 2*2+2+3+2+1)
+	// 2 snaps with (snap switch channel + link snap) + assets update for the
+	// kernel snap + gadget snap (switch channel, assets update, cmdline update)
+	// + recovery system (2 tasks) + set-model
+	c.Assert(tl, HasLen, 2*2+1+3+2+1)
 
 	deviceCtx, err := devicestate.DeviceCtx(s.state, tl[0], nil)
 	c.Assert(err, IsNil)
@@ -4014,24 +3999,21 @@ func (s *deviceMgrRemodelSuite) testRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 
 	// check the tasks
 	tSwitchChannelKernel := tl[0]
-	tSetupKernelSnap := tl[1]
-	tUpdateAssetsFromKernel := tl[2]
-	tLinkKernel := tl[3]
-	tSwitchChannelBase := tl[4]
-	tLinkBase := tl[5]
-	tSwitchChannelGadget := tl[6]
-	tUpdateAssetsFromGadget := tl[7]
-	tUpdateCmdlineFromGadget := tl[8]
-	tCreateRecovery := tl[9]
-	tFinalizeRecovery := tl[10]
-	tSetModel := tl[11]
+	tUpdateAssetsFromKernel := tl[1]
+	tLinkKernel := tl[2]
+	tSwitchChannelBase := tl[3]
+	tLinkBase := tl[4]
+	tSwitchChannelGadget := tl[5]
+	tUpdateAssetsFromGadget := tl[6]
+	tUpdateCmdlineFromGadget := tl[7]
+	tCreateRecovery := tl[8]
+	tFinalizeRecovery := tl[9]
+	tSetModel := tl[10]
 
 	// check the tasks
 	c.Assert(tSwitchChannelKernel.Kind(), Equals, "switch-snap-channel")
 	c.Assert(tSwitchChannelKernel.Summary(), Equals, `Switch pc-kernel-new channel to 20/stable`)
 	c.Assert(tSwitchChannelKernel.WaitTasks(), HasLen, 0)
-	c.Assert(tSetupKernelSnap.Kind(), Equals, "prepare-kernel-snap")
-	c.Assert(tSetupKernelSnap.Summary(), Equals, `Prepare kernel driver tree for "pc-kernel-new" (222) for remodel`)
 	c.Assert(tUpdateAssetsFromKernel.Kind(), Equals, "update-gadget-assets")
 	c.Assert(tUpdateAssetsFromKernel.Summary(), Equals, `Update assets from kernel "pc-kernel-new" (222) for remodel`)
 	c.Assert(tLinkKernel.Kind(), Equals, "link-snap")
@@ -4069,13 +4051,10 @@ func (s *deviceMgrRemodelSuite) testRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 		tCreateRecovery,
 		tSwitchChannelGadget,
 	})
-	c.Assert(tSetupKernelSnap.WaitTasks(), DeepEquals, []*state.Task{
+	c.Assert(tUpdateAssetsFromKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tSwitchChannelKernel, tSwitchChannelGadget,
 		tCreateRecovery, tFinalizeRecovery,
 	})
-	c.Assert(tUpdateAssetsFromKernel.WaitTasks(), DeepEquals, []*state.Task{
-		tSetupKernelSnap,
-	})
 	c.Check(tLinkKernel.WaitTasks(), DeepEquals, []*state.Task{
 		tUpdateAssetsFromKernel,
 	})
@@ -4084,7 +4063,7 @@ func (s *deviceMgrRemodelSuite) testRemodelUC20SwitchKernelBaseGadgetSnapsInstal
 	})
 	// setModel waits for everything in the change
 	c.Assert(tSetModel.WaitTasks(), DeepEquals, []*state.Task{
-		tSwitchChannelKernel, tSetupKernelSnap, tUpdateAssetsFromKernel,
+		tSwitchChannelKernel, tUpdateAssetsFromKernel,
 		tLinkKernel, tSwitchChannelBase, tLinkBase,
 		tSwitchChannelGadget, tUpdateAssetsFromGadget, tUpdateCmdlineFromGadget,
 		tCreateRecovery, tFinalizeRecovery,
diff --git a/overlord/snapstate/snapstate.go b/overlord/snapstate/snapstate.go
index 149da37b177..295b1bd703c 100644
--- a/overlord/snapstate/snapstate.go
+++ b/overlord/snapstate/snapstate.go
@@ -3404,43 +3404,79 @@ func LinkNewBaseOrKernel(st *state.State, name string, fromChange string) (*stat
 		InstanceKey: snapst.InstanceKey,
 	}
 
+	// note that prepare-snap doesn't actually do anything here, and is mostly
+	// used as a task to carry the snap-setup information.
 	prepareSnap := st.NewTask("prepare-snap", fmt.Sprintf(i18n.G("Prepare snap %q (%s) for remodel"), snapsup.InstanceName(), snapst.Current))
 	prepareSnap.Set("snap-setup", &snapsup)
-	prev := prepareSnap
+
 	ts := state.NewTaskSet(prepareSnap)
+	ts.MarkEdge(prepareSnap, LastBeforeLocalModificationsEdge)
+	ts.MarkEdge(prepareSnap, SnapSetupEdge)
+
+	if err := addLinkNewBaseOrKernelTasks(st, snapst, ts, prepareSnap); err != nil {
+		return nil, err
+	}
+
+	return ts, nil
+}
+
+func addLinkNewBaseOrKernelTasks(st *state.State, snapst SnapState, ts *state.TaskSet, snapsupTask *state.Task) error {
+	tasks := ts.Tasks()
+	if len(tasks) == 0 {
+		return errors.New("internal error: task set must be seeded with at least one task")
+	}
+
+	prev := tasks[len(tasks)-1]
+	add := func(t *state.Task) {
+		t.Set("snap-setup-task", snapsupTask.ID())
+		t.WaitFor(prev)
+		ts.AddTask(t)
+		prev = t
+	}
+
+	info, err := snapst.CurrentInfo()
+	if err != nil {
+		return err
+	}
+
 	// preserve the same order as during the update
 	if info.Type() == snap.TypeKernel {
-		// TODO in a remodel this would use the old model, we need to fix this
-		// as needsKernelSetup needs to know the new model for UC2{0,2} -> UC24
-		// remodel case.
-		deviceCtx, err := DeviceCtx(st, nil, nil)
-		if err != nil {
-			return nil, err
-		}
-		if kernel.NeedsKernelDriversTree(deviceCtx.Model()) {
-			setupKernel := st.NewTask("prepare-kernel-snap", fmt.Sprintf(i18n.G("Prepare kernel driver tree for %q (%s) for remodel"), snapsup.InstanceName(), snapst.Current))
-			ts.AddTask(setupKernel)
-			setupKernel.Set("snap-setup-task", prepareSnap.ID())
-			setupKernel.WaitFor(prev)
-			prev = setupKernel
-		}
+		// this previously created a prepare-kernel-snap task, but this was not
+		// needed since this function is only used to swap to kernel snaps that
+		// are still installed on the system (but not in use. this can happen
+		// because remodeling leaves around old snaps). if the snap is
+		// installed, then it is expected that the drivers tree is present.
 
 		// kernel snaps can carry boot assets
-		gadgetUpdate := st.NewTask("update-gadget-assets", fmt.Sprintf(i18n.G("Update assets from %s %q (%s) for remodel"), snapsup.Type, snapsup.InstanceName(), snapst.Current))
-		gadgetUpdate.Set("snap-setup-task", prepareSnap.ID())
-		gadgetUpdate.WaitFor(prev)
-		ts.AddTask(gadgetUpdate)
-		prev = gadgetUpdate
+		gadgetUpdate := st.NewTask("update-gadget-assets", fmt.Sprintf(i18n.G("Update assets from %s %q (%s) for remodel"), info.Type(), info.InstanceName(), snapst.Current))
+		add(gadgetUpdate)
 	}
-	linkSnap := st.NewTask("link-snap", fmt.Sprintf(i18n.G("Make snap %q (%s) available to the system during remodel"), snapsup.InstanceName(), snapst.Current))
-	linkSnap.Set("snap-setup-task", prepareSnap.ID())
-	linkSnap.WaitFor(prev)
-	ts.AddTask(linkSnap)
+
+	linkSnap := st.NewTask("link-snap", fmt.Sprintf(i18n.G("Make snap %q (%s) available to the system during remodel"), info.InstanceName(), snapst.Current))
+	add(linkSnap)
 	ts.MarkEdge(linkSnap, MaybeRebootEdge)
-	// prepare-snap is the last task that carries no system modifications
-	ts.MarkEdge(prepareSnap, LastBeforeLocalModificationsEdge)
-	ts.MarkEdge(prepareSnap, SnapSetupEdge)
-	return ts, nil
+
+	components := snapst.Sequence.ComponentsForRevision(snapst.Current)
+	compsupTasks := make([]string, 0, len(components))
+	for _, cs := range components {
+		compsup := ComponentSetup{
+			CompSideInfo: cs.SideInfo,
+			CompType:     cs.CompType,
+		}
+
+		cref := compsup.CompSideInfo.Component
+		compRev := compsup.CompSideInfo.Revision
+
+		link := st.NewTask("link-component", fmt.Sprintf(i18n.G("Make component %q (%s) available to the system during remodel"), cref, compRev))
+		link.Set("component-setup", compsup)
+		add(link)
+
+		compsupTasks = append(compsupTasks, link.ID())
+	}
+
+	snapsupTask.Set("component-setup-tasks", compsupTasks)
+
+	return nil
 }
 
 func findSnapSetupTask(tasks []*state.Task) (*state.Task, *SnapSetup, error) {
@@ -3459,55 +3495,37 @@ func findSnapSetupTask(tasks []*state.Task) (*state.Task, *SnapSetup, error) {
 // AddLinkNewBaseOrKernel creates the same tasks as LinkNewBaseOrKernel but adds
 // them to the provided task set.
 func AddLinkNewBaseOrKernel(st *state.State, ts *state.TaskSet) (*state.TaskSet, error) {
+	if ts.MaybeEdge(LastBeforeLocalModificationsEdge) != nil {
+		return nil, errors.New("internal error: cannot add tasks to link new base or kernel to task set that introduces local modifications")
+	}
+
 	allTasks := ts.Tasks()
 	snapSetupTask, snapsup, err := findSnapSetupTask(allTasks)
 	if err != nil {
 		return nil, err
 	}
 	if snapSetupTask == nil {
-		return nil, fmt.Errorf("internal error: cannot identify task with snap-setup")
+		return nil, errors.New("internal error: cannot identify task with snap-setup")
 	}
-	// the first task added here waits for the last task in the existing set
-	prev := allTasks[len(allTasks)-1]
-	// preserve the same order as during the update
-	if snapsup.Type == snap.TypeKernel {
-		// TODO in a remodel this would use the old model, we need to fix this
-		// as needsKernelSetup needs to know the new model for UC2{0,2} -> UC24
-		// remodel case.
-		deviceCtx, err := DeviceCtx(st, nil, nil)
-		if err != nil {
-			return nil, err
-		}
-		if kernel.NeedsKernelDriversTree(deviceCtx.Model()) {
-			setupKernel := st.NewTask("prepare-kernel-snap", fmt.Sprintf(i18n.G("Prepare kernel driver tree for %q (%s) for remodel"), snapsup.InstanceName(), snapsup.Revision()))
-			setupKernel.Set("snap-setup-task", snapSetupTask.ID())
-			setupKernel.WaitFor(prev)
-			ts.AddTask(setupKernel)
-			prev = setupKernel
-		}
 
-		// kernel snaps can carry boot assets
-		gadgetUpdate := st.NewTask("update-gadget-assets", fmt.Sprintf(i18n.G("Update assets from %s %q (%s) for remodel"), snapsup.Type, snapsup.InstanceName(), snapsup.Revision()))
-		gadgetUpdate.Set("snap-setup-task", snapSetupTask.ID())
-		// wait for the last task in existing set
-		gadgetUpdate.WaitFor(prev)
-		ts.AddTask(gadgetUpdate)
-		prev = gadgetUpdate
-	}
-	linkSnap := st.NewTask("link-snap",
-		fmt.Sprintf(i18n.G("Make snap %q (%s) available to the system during remodel"), snapsup.InstanceName(), snapsup.SideInfo.Revision))
-	linkSnap.Set("snap-setup-task", snapSetupTask.ID())
-	linkSnap.WaitFor(prev)
-	ts.AddTask(linkSnap)
-	ts.MarkEdge(linkSnap, MaybeRebootEdge)
-	// make sure that remodel can identify which tasks introduce actual
-	// changes to the system and order them correctly
-	if edgeTask := ts.MaybeEdge(LastBeforeLocalModificationsEdge); edgeTask == nil {
-		// no task in the task set is marked as last before system
-		// modifications are introduced, so we need to mark the last
-		// task in the set, as tasks introduced here modify system state
-		ts.MarkEdge(allTasks[len(allTasks)-1], LastBeforeLocalModificationsEdge)
+	var snapst SnapState
+	if err := Get(st, snapsup.InstanceName(), &snapst); err != nil {
+		return nil, err
+	}
+
+	if snapst.Current != snapsup.Revision() {
+		return nil, errors.New("internal error: cannot add tasks to link new base or kernel to task set that changes the snap revision")
+	}
+
+	// no task in the task set is marked as last before system modifications are
+	// introduced, so we need to mark the last task in the original set, as
+	// tasks introduced here modify system state
+	ts.MarkEdge(allTasks[len(allTasks)-1], LastBeforeLocalModificationsEdge)
+
+	if err := addLinkNewBaseOrKernelTasks(st, snapst, ts, snapSetupTask); err != nil {
+		return nil, err
 	}
+
 	return ts, nil
 }
 
diff --git a/overlord/snapstate/snapstate_test.go b/overlord/snapstate/snapstate_test.go
index cdef3a33d16..a138b9a6176 100644
--- a/overlord/snapstate/snapstate_test.go
+++ b/overlord/snapstate/snapstate_test.go
@@ -8061,7 +8061,7 @@ version: 1`), &snap.SideInfo{Revision: snap.R("5")})
 	c.Check(ignoreValidation, DeepEquals, map[string]bool{"bar": true})
 }
 
-func (s *snapmgrTestSuite) addSnapsForRemodel(c *C) {
+func (s *snapmgrTestSuite) addSnapsForRemodel(c *C, withComponents bool) {
 	si := &snap.SideInfo{
 		RealName: "some-base", Revision: snap.R(1),
 	}
@@ -8077,12 +8077,25 @@ func (s *snapmgrTestSuite) addSnapsForRemodel(c *C) {
 		RealName: "some-kernel", Revision: snap.R(2),
 	}
 	snaptest.MockSnapCurrent(c, "name: some-kernel\nversion: 1.0\ntype: kernel\n", si)
+	seq := snapstatetest.NewSequenceFromSnapSideInfos([]*snap.SideInfo{si})
+	if withComponents {
+		seq.AddComponentForRevision(si.Revision, sequence.NewComponentState(&snap.ComponentSideInfo{
+			Component: naming.NewComponentRef("some-kernel", "comp-1"),
+			Revision:  snap.R(22),
+		}, snap.KernelModulesComponent))
+
+		seq.AddComponentForRevision(si.Revision, sequence.NewComponentState(&snap.ComponentSideInfo{
+			Component: naming.NewComponentRef("some-kernel", "comp-2"),
+			Revision:  snap.R(33),
+		}, snap.KernelModulesComponent))
+	}
 	snapstate.Set(s.state, "some-kernel", &snapstate.SnapState{
 		Active:   true,
-		Sequence: snapstatetest.NewSequenceFromSnapSideInfos([]*snap.SideInfo{si}),
+		Sequence: seq,
 		Current:  si.Revision,
 		SnapType: "kernel",
 	})
+
 	si = &snap.SideInfo{
 		RealName: "some-gadget", Revision: snap.R(3),
 	}
@@ -8096,6 +8109,7 @@ func (s *snapmgrTestSuite) addSnapsForRemodel(c *C) {
 }
 
 var nonReLinkKinds = []string{
+	"prepare-kernel-snap",
 	"copy-snap-data",
 	"setup-profiles",
 	"auto-connect",
@@ -8107,6 +8121,7 @@ var nonReLinkKinds = []string{
 	"run-hook[configure]",
 	"run-hook[check-health]",
 	"discard-old-kernel-snap-setup",
+	"mount-component",
 }
 
 func kindsToSet(kinds []string) map[string]bool {
@@ -8125,6 +8140,10 @@ func (s *snapmgrTestSuite) TestRemodelLinkNewBaseOrUC20KernelHappy(c *C) {
 	s.testRemodelLinkNewBaseOrKernelHappy(c, MakeModel20("brand-gadget", nil), 0)
 }
 
+func (s *snapmgrTestSuite) TestRemodelLinkNewBaseOrUC20KernelHappyWithKmodComponent(c *C) {
+	s.testRemodelLinkNewBaseOrKernelHappy(c, MakeModel20("brand-gadget", nil), 0)
+}
+
 func (s *snapmgrTestSuite) TestRemodelLinkNewBaseOrUC24KernelHappy(c *C) {
 	// UC24 model has additional tasks for the kernel
 	s.testRemodelLinkNewBaseOrKernelHappy(c,
@@ -8142,27 +8161,20 @@ func (s *snapmgrTestSuite) testRemodelLinkNewBaseOrKernelHappy(c *C, model *asse
 
 	defer snapstatetest.MockDeviceModel(model)()
 
-	s.addSnapsForRemodel(c)
+	const withComponents = false
+	s.addSnapsForRemodel(c, withComponents)
 
 	ts, err := snapstate.LinkNewBaseOrKernel(s.state, "some-kernel", "")
 	c.Assert(err, IsNil)
+
 	tasks := ts.Tasks()
 	c.Check(taskKinds(tasks), DeepEquals, expectedDoInstallTasks(snap.TypeKernel, opts, 0, 0, []string{"prepare-snap"}, nil, kindsToSet(nonReLinkKinds)))
+	c.Assert(tasks, HasLen, 3)
+
 	tPrepare := tasks[0]
-	var tLink, tUpdateGadgetAssets *state.Task
-	if opts&needsKernelSetup != 0 {
-		c.Assert(tasks, HasLen, 4)
-		tSetupKernelSnap := tasks[1]
-		c.Assert(tSetupKernelSnap.Kind(), Equals, "prepare-kernel-snap")
-		c.Assert(tSetupKernelSnap.Summary(), Equals, `Prepare kernel driver tree for "some-kernel" (2) for remodel`)
-		c.Assert(tSetupKernelSnap.WaitTasks(), DeepEquals, []*state.Task{tPrepare})
-		tUpdateGadgetAssets = tasks[2]
-		tLink = tasks[3]
-	} else {
-		c.Assert(tasks, HasLen, 3)
-		tUpdateGadgetAssets = tasks[1]
-		tLink = tasks[2]
-	}
+	tUpdateGadgetAssets := tasks[1]
+	tLink := tasks[2]
+
 	c.Assert(tPrepare.Kind(), Equals, "prepare-snap")
 	c.Assert(tPrepare.Summary(), Equals, `Prepare snap "some-kernel" (2) for remodel`)
 	c.Assert(tPrepare.Has("snap-setup"), Equals, true)
@@ -8188,6 +8200,92 @@ func (s *snapmgrTestSuite) testRemodelLinkNewBaseOrKernelHappy(c *C, model *asse
 	c.Assert(ts.MaybeEdge(snapstate.MaybeRebootEdge), Equals, tLink)
 }
 
+func (s *snapmgrTestSuite) TestRemodelLinkNewBaseOrKernelWithComponent(c *C) {
+	model := MakeModel20("brand-gadget", nil)
+
+	restore := release.MockOnClassic(false)
+	defer restore()
+
+	s.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	defer snapstatetest.MockDeviceModel(model)()
+
+	const withComponents = true
+	s.addSnapsForRemodel(c, withComponents)
+
+	ts, err := snapstate.LinkNewBaseOrKernel(s.state, "some-kernel", "")
+	c.Assert(err, IsNil)
+
+	comps := []string{"comp-1", "comp-2"}
+	startTasks := []string{"prepare-snap"}
+
+	tasks := ts.Tasks()
+	c.Check(taskKinds(tasks), DeepEquals, expectedDoInstallTasks(snap.TypeKernel, 0, 0, 0, startTasks, comps, kindsToSet(nonReLinkKinds)))
+	c.Assert(tasks, HasLen, 5)
+
+	prepare := ts.MaybeEdge(snapstate.SnapSetupEdge)
+	c.Assert(prepare, NotNil)
+	c.Assert(prepare.Kind(), Equals, "prepare-snap")
+	c.Assert(prepare.Has("snap-setup"), Equals, true)
+
+	link := ts.MaybeEdge(snapstate.MaybeRebootEdge)
+	c.Assert(link, NotNil)
+	c.Assert(link.Kind(), Equals, "link-snap")
+
+	for _, t := range ts.Tasks() {
+		if t.Kind() == "link-component" {
+			c.Assert(t.Has("component-setup"), Equals, true)
+		}
+	}
+}
+
+func (s *snapmgrTestSuite) TestRemodelAddLinkNewBaseOrKernelWithComponent(c *C) {
+	model := MakeModel20("brand-gadget", nil)
+
+	restore := release.MockOnClassic(false)
+	defer restore()
+
+	s.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	defer snapstatetest.MockDeviceModel(model)()
+
+	const withComponents = true
+	s.addSnapsForRemodel(c, withComponents)
+
+	switchSnap := s.state.NewTask("switch-snap", "switch snap")
+	switchSnap.Set("snap-setup", &snapstate.SnapSetup{
+		SideInfo: &snap.SideInfo{RealName: "some-kernel", Revision: snap.R(2)},
+		Type:     "kernel",
+		Channel:  "new-channel",
+	})
+
+	ts := state.NewTaskSet(switchSnap)
+
+	ts, err := snapstate.AddLinkNewBaseOrKernel(s.state, ts)
+	c.Assert(err, IsNil)
+
+	comps := []string{"comp-1", "comp-2"}
+	startTasks := []string{"switch-snap"}
+
+	tasks := ts.Tasks()
+	c.Check(taskKinds(tasks), DeepEquals, expectedDoInstallTasks(snap.TypeKernel, 0, 0, 0, startTasks, comps, kindsToSet(nonReLinkKinds)))
+	c.Assert(tasks, HasLen, 5)
+
+	link := ts.MaybeEdge(snapstate.MaybeRebootEdge)
+	c.Assert(link, NotNil)
+	c.Assert(link.Kind(), Equals, "link-snap")
+
+	for _, t := range ts.Tasks() {
+		if t.Kind() == "link-component" {
+			c.Assert(t.Has("component-setup"), Equals, true)
+		}
+	}
+}
+
 func (s *snapmgrTestSuite) TestRemodelLinkNewBaseOrKernelBadType(c *C) {
 	restore := release.MockOnClassic(false)
 	defer restore()
@@ -8195,7 +8293,9 @@ func (s *snapmgrTestSuite) TestRemodelLinkNewBaseOrKernelBadType(c *C) {
 	s.BaseTest.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
 	s.state.Lock()
 	defer s.state.Unlock()
-	s.addSnapsForRemodel(c)
+
+	const withComponents = false
+	s.addSnapsForRemodel(c, withComponents)
 
 	si := &snap.SideInfo{RealName: "some-snap", Revision: snap.R(3)}
 	snaptest.MockSnapCurrent(c, "name: snap-gadget\nversion: 1.0\n", si)
@@ -8221,7 +8321,9 @@ func (s *snapmgrTestSuite) TestRemodelLinkNewBaseOrKernelNoRemodelConflict(c *C)
 	s.BaseTest.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
 	s.state.Lock()
 	defer s.state.Unlock()
-	s.addSnapsForRemodel(c)
+
+	const withComponents = false
+	s.addSnapsForRemodel(c, withComponents)
 
 	tugc := s.state.NewTask("update-managed-boot-config", "update managed boot config")
 	chg := s.state.NewChange("remodel", "remodel")
@@ -8256,6 +8358,8 @@ func (s *snapmgrTestSuite) testRemodelAddLinkNewBaseOrKernel(c *C, model *assert
 
 	defer snapstatetest.MockDeviceModel(model)()
 
+	s.addSnapsForRemodel(c, false)
+
 	// try a kernel snap first
 	si := &snap.SideInfo{RealName: "some-kernel", Revision: snap.R(2)}
 	tPrepare := s.state.NewTask("prepare-snap", "test task")
@@ -8273,22 +8377,9 @@ func (s *snapmgrTestSuite) testRemodelAddLinkNewBaseOrKernel(c *C, model *assert
 	tasks := tsNew.Tasks()
 	c.Check(taskKinds(tasks), DeepEquals, expectedDoInstallTasks(snap.TypeKernel, opts, 0, 0, []string{"prepare-snap", "test-task"}, nil, kindsToSet(nonReLinkKinds)))
 	// since this is the kernel, we have our task + test task + update-gadget-assets + link-snap
-	var tLink, tUpdateGadgetAssets *state.Task
-	if opts&needsKernelSetup != 0 {
-		c.Assert(tasks, HasLen, 5)
-		tSetupKernelSnap := tasks[2]
-		c.Assert(tSetupKernelSnap.Kind(), Equals, "prepare-kernel-snap")
-		c.Assert(tSetupKernelSnap.Summary(), Equals, `Prepare kernel driver tree for "some-kernel" (2) for remodel`)
-		c.Assert(tSetupKernelSnap.WaitTasks(), DeepEquals, []*state.Task{
-			testTask,
-		})
-		tUpdateGadgetAssets = tasks[3]
-		tLink = tasks[4]
-	} else {
-		c.Assert(tasks, HasLen, 4)
-		tUpdateGadgetAssets = tasks[2]
-		tLink = tasks[3]
-	}
+	c.Assert(tasks, HasLen, 4)
+	tUpdateGadgetAssets := tasks[2]
+	tLink := tasks[3]
 	c.Assert(tUpdateGadgetAssets.Kind(), Equals, "update-gadget-assets")
 	c.Assert(tUpdateGadgetAssets.Summary(), Equals, `Update assets from kernel "some-kernel" (2) for remodel`)
 	c.Assert(tLink.Kind(), Equals, "link-snap")
@@ -8341,7 +8432,9 @@ func (s *snapmgrTestSuite) TestRemodelSwitchNewGadget(c *C) {
 	s.BaseTest.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
 	s.state.Lock()
 	defer s.state.Unlock()
-	s.addSnapsForRemodel(c)
+
+	const withComponents = false
+	s.addSnapsForRemodel(c, withComponents)
 
 	ts, err := snapstate.SwitchToNewGadget(s.state, "some-gadget", "")
 	c.Assert(err, IsNil)
@@ -8371,7 +8464,9 @@ func (s *snapmgrTestSuite) TestRemodelSwitchNewGadgetNoRemodelConflict(c *C) {
 	s.BaseTest.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
 	s.state.Lock()
 	defer s.state.Unlock()
-	s.addSnapsForRemodel(c)
+
+	const withComponents = false
+	s.addSnapsForRemodel(c, withComponents)
 
 	tugc := s.state.NewTask("update-managed-boot-config", "update managed boot config")
 	chg := s.state.NewChange("remodel", "remodel")
@@ -8388,7 +8483,9 @@ func (s *snapmgrTestSuite) TestRemodelSwitchNewGadgetBadType(c *C) {
 	s.BaseTest.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
 	s.state.Lock()
 	defer s.state.Unlock()
-	s.addSnapsForRemodel(c)
+
+	const withComponent = false
+	s.addSnapsForRemodel(c, withComponent)
 
 	si := &snap.SideInfo{RealName: "some-snap", Revision: snap.R(3)}
 	snaptest.MockSnapCurrent(c, "name: snap-gadget\nversion: 1.0\n", si)
@@ -8416,7 +8513,9 @@ func (s *snapmgrTestSuite) TestRemodelSwitchNewGadgetConflict(c *C) {
 	s.BaseTest.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
 	s.state.Lock()
 	defer s.state.Unlock()
-	s.addSnapsForRemodel(c)
+
+	const withComponents = false
+	s.addSnapsForRemodel(c, withComponents)
 
 	tugc := s.state.NewTask("update-gadget-cmdline", "update gadget cmdline")
 	chg := s.state.NewChange("optional-kernel-cmdline", "optional kernel cmdline")
@@ -8442,7 +8541,9 @@ func (s *snapmgrTestSuite) TestRemodelSwitchNewGadgetConflictExclusiveKind(c *C)
 	s.BaseTest.AddCleanup(snapstate.MockSnapReadInfo(snap.ReadInfo))
 	s.state.Lock()
 	defer s.state.Unlock()
-	s.addSnapsForRemodel(c)
+
+	const withComponent = false
+	s.addSnapsForRemodel(c, withComponent)
 
 	tugc := s.state.NewTask("some-random-task", "...")
 	chg := s.state.NewChange("transition-ubuntu-core", "...")
diff --git a/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-24.json b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-24.json
new file mode 100644
index 00000000000..93975e0f39f
--- /dev/null
+++ b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-24.json
@@ -0,0 +1,38 @@
+{
+    "type": "model",
+    "authority-id": "developer1",
+    "series": "16",
+    "brand-id": "developer1",
+    "model": "model",
+    "architecture": "amd64",
+    "timestamp": "2024-04-24T00:00:00+00:00",
+    "grade": "dangerous",
+    "base": "core24",
+    "serial-authority": ["generic"],
+    "snaps": [
+        {
+            "default-channel": "24/edge",
+            "id": "UqFziVZDHLSyO3TqSWgNBoAdHbLI4dAH",
+            "name": "pc",
+            "type": "gadget"
+        },
+        {
+            "default-channel": "24/edge",
+            "id": "pYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza",
+            "name": "pc-kernel",
+            "type": "kernel"
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "dwTAh7MZZ01zyriOZErqd1JynQLiOGvM",
+            "name": "core24",
+            "type": "base"
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "PMrrV4ml8uWuEUDBT8dSGnKUYbevVhc4",
+            "name": "snapd",
+            "type": "snapd"
+        }
+    ]
+}
diff --git a/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-1-24.json b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-1-24.json
new file mode 100644
index 00000000000..f3dbef88fdb
--- /dev/null
+++ b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-1-24.json
@@ -0,0 +1,42 @@
+{
+    "type": "model",
+    "authority-id": "developer1",
+    "series": "16",
+    "brand-id": "developer1",
+    "model": "model",
+    "revision": "1",
+    "architecture": "amd64",
+    "timestamp": "2024-04-24T00:00:00+00:00",
+    "grade": "dangerous",
+    "base": "core24",
+    "serial-authority": ["generic"],
+    "snaps": [
+        {
+            "default-channel": "24/edge",
+            "id": "UqFziVZDHLSyO3TqSWgNBoAdHbLI4dAH",
+            "name": "pc",
+            "type": "gadget"
+        },
+        {
+            "default-channel": "24/edge",
+            "id": "qYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza",
+            "name": "pc-kernel-mac80211-hwsim",
+            "type": "kernel",
+            "components": {
+              "wifi-comp": "required"
+            }
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "dwTAh7MZZ01zyriOZErqd1JynQLiOGvM",
+            "name": "core24",
+            "type": "base"
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "PMrrV4ml8uWuEUDBT8dSGnKUYbevVhc4",
+            "name": "snapd",
+            "type": "snapd"
+        }
+    ]
+}
diff --git a/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-2-24.json b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-2-24.json
new file mode 100644
index 00000000000..5f7ffa04c2b
--- /dev/null
+++ b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-2-24.json
@@ -0,0 +1,42 @@
+{
+    "type": "model",
+    "authority-id": "developer1",
+    "series": "16",
+    "brand-id": "developer1",
+    "model": "model",
+    "revision": "2",
+    "architecture": "amd64",
+    "timestamp": "2024-04-24T00:00:00+00:00",
+    "grade": "dangerous",
+    "base": "core24",
+    "serial-authority": ["generic"],
+    "snaps": [
+        {
+            "default-channel": "24/edge",
+            "id": "UqFziVZDHLSyO3TqSWgNBoAdHbLI4dAH",
+            "name": "pc",
+            "type": "gadget"
+        },
+        {
+            "default-channel": "24/edge",
+            "id": "rYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza",
+            "name": "pc-kernel-efi-pstore",
+            "type": "kernel",
+            "components": {
+              "efi-comp": "required"
+            }
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "dwTAh7MZZ01zyriOZErqd1JynQLiOGvM",
+            "name": "core24",
+            "type": "base"
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "PMrrV4ml8uWuEUDBT8dSGnKUYbevVhc4",
+            "name": "snapd",
+            "type": "snapd"
+        }
+    ]
+}
diff --git a/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-3-24.json b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-3-24.json
new file mode 100644
index 00000000000..44a02e52ba0
--- /dev/null
+++ b/tests/lib/assertions/test-snapd-remodel-to-installed-kernel-rev-3-24.json
@@ -0,0 +1,42 @@
+{
+    "type": "model",
+    "authority-id": "developer1",
+    "series": "16",
+    "brand-id": "developer1",
+    "model": "model",
+    "revision": "3",
+    "architecture": "amd64",
+    "timestamp": "2024-04-24T00:00:00+00:00",
+    "grade": "dangerous",
+    "base": "core24",
+    "serial-authority": ["generic"],
+    "snaps": [
+        {
+            "default-channel": "24/edge",
+            "id": "UqFziVZDHLSyO3TqSWgNBoAdHbLI4dAH",
+            "name": "pc",
+            "type": "gadget"
+        },
+        {
+            "default-channel": "24/edge",
+            "id": "qYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza",
+            "name": "pc-kernel-mac80211-hwsim",
+            "type": "kernel",
+            "components": {
+              "wifi-comp": "required"
+            }
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "dwTAh7MZZ01zyriOZErqd1JynQLiOGvM",
+            "name": "core24",
+            "type": "base"
+        },
+        {
+            "default-channel": "latest/edge",
+            "id": "PMrrV4ml8uWuEUDBT8dSGnKUYbevVhc4",
+            "name": "snapd",
+            "type": "snapd"
+        }
+    ]
+}
diff --git a/tests/lib/tools/build_kernel_with_comps.sh b/tests/lib/tools/build_kernel_with_comps.sh
index 654da60b83d..94b63eca9e7 100755
--- a/tests/lib/tools/build_kernel_with_comps.sh
+++ b/tests/lib/tools/build_kernel_with_comps.sh
@@ -19,12 +19,13 @@ build_kernel_with_comp() {
         kernel_snap_file="pc-kernel.snap"
     fi
     unsquashfs -d kernel "${kernel_snap_file}"
+    kernel_name="$(grep 'name:' kernel/meta/snap.yaml | awk '{ print $2 }')"
     kern_ver=$(find kernel/modules/* -maxdepth 0 -printf "%f\n")
     comp_ko_dir=$comp_name/modules/"$kern_ver"/kmod/
     mkdir -p "$comp_ko_dir"
     mkdir -p "$comp_name"/meta/
     cat << EOF > "$comp_name"/meta/component.yaml
-component: pc-kernel+$comp_name
+component: ${kernel_name}+${comp_name}
 type: kernel-modules
 version: 1.0
 summary: kernel component
@@ -34,7 +35,7 @@ EOF
     glob_mod_name=$(printf '%s' "$mod_name" | sed -r 's/[-_]/[-_]/g')
     module_path=$(find kernel -name "${glob_mod_name}.ko*")
     cp "$module_path" "$comp_ko_dir"
-    snap pack --filename=pc-kernel+"$comp_name".comp "$comp_name"
+    snap pack --filename="${kernel_name}+${comp_name}".comp "$comp_name"
 
     # Create kernel without the kernel module
     rm "$module_path"
@@ -47,8 +48,7 @@ EOF
     # append component meta-information
     printf 'components:\n  %s:\n    type: kernel-modules\n' "$comp_name" >> kernel/meta/snap.yaml
     snap pack --filename="${kernel_snap_file}" kernel
-    # Just so that nested_prepare_kernel does not recopy the old one
-    cp "${kernel_snap_file}" "${NESTED_ASSETS_DIR}/pc-kernel.snap"
+    rm -r kernel
 }
 
 build_kernel_with_comp "$@"
diff --git a/tests/nested/manual/remodel-to-installed-kernel/task.yaml b/tests/nested/manual/remodel-to-installed-kernel/task.yaml
new file mode 100644
index 00000000000..a7689b90105
--- /dev/null
+++ b/tests/nested/manual/remodel-to-installed-kernel/task.yaml
@@ -0,0 +1,189 @@
+summary: Remodel to a model with a different kernel and then back to the original kernel again
+
+details: |
+  This test covers a few remodeling edge cases:
+  * Switching the kernel to a new snap (changing kernel snap names, not just
+    revision) works properly.
+  * The drivers tree mounted at /usr/lib/modules is correctly updated and
+    maintained across remodels.
+  * When returning to a kernel that was left behind on the system because of a
+    previous remodel, the drivers tree is reused.
+
+  Note that the implementation of this last feature relies on the fact that
+  switching the kernel snap during a remodel does not remove the old kernel's
+  snap or drivers tree. We reuse the drivers tree that was originally created
+  when the kernel was installed.
+
+systems: [-ubuntu-1*, -ubuntu-20*, -ubuntu-22*]
+
+environment:
+  INITIAL_MODEL_JSON: $TESTSLIB/assertions/test-snapd-remodel-to-installed-kernel-24.json
+  SECOND_MODEL_JSON: $TESTSLIB/assertions/test-snapd-remodel-to-installed-kernel-rev-1-24.json
+  THIRD_MODEL_JSON: $TESTSLIB/assertions/test-snapd-remodel-to-installed-kernel-rev-2-24.json
+  FOURTH_MODEL_JSON: $TESTSLIB/assertions/test-snapd-remodel-to-installed-kernel-rev-3-24.json
+  NESTED_ENABLE_TPM: true
+  NESTED_ENABLE_SECURE_BOOT: true
+  NESTED_BUILD_SNAPD_FROM_CURRENT: true
+  NESTED_REPACK_GADGET_SNAP: true
+  NESTED_REPACK_KERNEL_SNAP: true
+  NESTED_REPACK_BASE_SNAP: true
+  NESTED_REPACK_FOR_FAKESTORE: true
+  NESTED_FAKESTORE_BLOB_DIR: $(pwd)/fake-store-blobdir
+  NESTED_SIGN_SNAPS_FAKESTORE: true
+  NESTED_UBUNTU_IMAGE_SNAPPY_FORCE_SAS_URL: http://localhost:11028
+
+prepare: |
+    if [ "${TRUST_TEST_KEYS}" = "false" ]; then
+        echo "This test needs test keys to be trusted"
+        exit
+    fi
+
+    # although nested_start_core_vm_unit usually installs this, the fake store
+    # will already have been set up, so we need to install it here
+    snap install test-snapd-swtpm --edge
+
+    "${TESTSTOOLS}/store-state" setup-fake-store "${NESTED_FAKESTORE_BLOB_DIR}"
+
+    gendeveloper1 sign-model < "${INITIAL_MODEL_JSON}" > initial-model.assert
+
+    cp "${TESTSLIB}/assertions/testrootorg-store.account-key" "${NESTED_FAKESTORE_BLOB_DIR}/asserts"
+    cp "${TESTSLIB}/assertions/developer1.account" "${NESTED_FAKESTORE_BLOB_DIR}/asserts"
+    cp "${TESTSLIB}/assertions/developer1.account-key" "${NESTED_FAKESTORE_BLOB_DIR}/asserts"
+    cp initial-model.assert "${NESTED_FAKESTORE_BLOB_DIR}/asserts"
+
+    tests.nested prepare-essential-snaps
+
+    export SNAPPY_FORCE_API_URL="${NESTED_UBUNTU_IMAGE_SNAPPY_FORCE_SAS_URL}"
+    ubuntu-image snap --image-size 10G ./initial-model.assert
+
+    image_dir=$(tests.nested get images-path)
+    image_name=$(tests.nested get image-name core)
+    cp ./pc.img "${image_dir}/${image_name}"
+    tests.nested configure-default-user
+
+    # run the fake device service too, so that the device can be initialised
+    systemd-run --collect --unit fakedevicesvc fakedevicesvc localhost:11029
+
+    tests.nested build-image core
+    tests.nested create-vm core
+
+    #shellcheck source=tests/lib/core-config.sh
+    . "$TESTSLIB"/core-config.sh
+    wait_for_first_boot_change
+
+    remote.exec 'sudo systemctl stop snapd snapd.socket'
+
+    remote.exec 'sudo cat /var/lib/snapd/state.json' | gojq '.data.auth.device."session-macaroon"="fake-session"' > state.json
+    remote.push state.json
+    remote.exec 'sudo mv state.json /var/lib/snapd/state.json'
+    remote.exec 'sudo systemctl start snapd snapd.socket'
+
+    # here unpack the kernel snap that was used to build the original image. we
+    # take that kernel and create a component from one of its kernel modules.
+    # upload the new kernel and component to the fake store.
+    unsquashfs -d mac80211-hwsim "${NESTED_FAKESTORE_BLOB_DIR}/pc-kernel.snap"
+
+    # we will use this copy of the original kernel soon
+    cp -ar mac80211-hwsim efi-pstore
+
+    sed -i -e '/^name:/ s/$/-mac80211-hwsim/' mac80211-hwsim/meta/snap.yaml
+    snap pack --filename=pc-kernel-mac80211-hwsim.snap ./mac80211-hwsim
+    "${TESTSTOOLS}"/build_kernel_with_comps.sh mac80211_hwsim wifi-comp pc-kernel-mac80211-hwsim.snap
+
+    "${TESTSTOOLS}"/store-state make-snap-installable --noack \
+      --revision 1 \
+      "${NESTED_FAKESTORE_BLOB_DIR}" \
+      ./pc-kernel-mac80211-hwsim.snap \
+      'qYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza'
+
+    "${TESTSTOOLS}"/store-state make-component-installable --noack \
+      --snap-revision 1 \
+      --component-revision 1 \
+      --snap-id 'qYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza' \
+      "${NESTED_FAKESTORE_BLOB_DIR}" \
+      ./pc-kernel-mac80211-hwsim+wifi-comp.comp
+
+    # do the same thing for a different kernel module.
+    sed -i -e '/^name:/ s/$/-efi-pstore/' efi-pstore/meta/snap.yaml
+    snap pack --filename=pc-kernel-efi-pstore.snap ./efi-pstore
+    "${TESTSTOOLS}"/build_kernel_with_comps.sh efi_pstore efi-comp pc-kernel-efi-pstore.snap
+
+    "${TESTSTOOLS}"/store-state make-snap-installable --noack \
+      --revision 1 \
+      "${NESTED_FAKESTORE_BLOB_DIR}" \
+      ./pc-kernel-efi-pstore.snap \
+      'rYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza'
+
+    "${TESTSTOOLS}"/store-state make-component-installable --noack \
+      --snap-revision 1 \
+      --component-revision 1 \
+      --snap-id 'rYVQrBcKmBa0mZ4CCN7ExT6jH8rY1hza' \
+      "${NESTED_FAKESTORE_BLOB_DIR}" \
+      ./pc-kernel-efi-pstore+efi-comp.comp
+
+restore: |
+    systemctl stop fakedevicesvc
+    "${TESTSTOOLS}/store-state" teardown-fake-store "${NESTED_FAKESTORE_BLOB_DIR}"
+
+execute: |
+  echo 'Swap to the first new kernel and component pc-kernel-mac80211-hwsim+wifi-comp'
+  gendeveloper1 sign-model < "${SECOND_MODEL_JSON}" > second-model.assert
+  remote.push second-model.assert
+
+  boot_id="$(tests.nested boot-id)"
+  change_id="$(remote.exec 'sudo snap remodel --no-wait second-model.assert')"
+  remote.wait-for reboot "${boot_id}"
+
+  # this remodel expects two reboots, once for testing the recovery system
+  # and once for rebooting into the new kernel
+  boot_id="$(tests.nested boot-id)"
+  remote.wait-for reboot "${boot_id}"
+
+  remote.exec "snap watch ${change_id}"
+  remote.exec 'snap list pc-kernel-mac80211-hwsim' | awk '$NR != 1 { print $3 }' | MATCH '1'
+  remote.exec 'snap components pc-kernel-mac80211-hwsim' | sed 1d | MATCH 'pc-kernel-mac80211-hwsim\+wifi-comp\s+installed'
+
+  # make sure that the kernel module got installed and is loaded from our
+  # component
+  remote.exec sudo modprobe mac80211_hwsim
+  remote.exec ip link show wlan0
+  remote.exec modinfo --filename mac80211_hwsim | MATCH '/lib/modules/.*/updates/wifi-comp'
+
+  echo 'Swap to the second new kernel and component pc-kernel-efi-pstore+efi-comp'
+  gendeveloper1 sign-model < "${THIRD_MODEL_JSON}" > third-model.assert
+  remote.push third-model.assert
+
+  boot_id="$(tests.nested boot-id)"
+  change_id="$(remote.exec 'sudo snap remodel --no-wait third-model.assert')"
+  remote.wait-for reboot "${boot_id}"
+
+  # same deal here, we expect two reboots
+  boot_id="$(tests.nested boot-id)"
+  remote.wait-for reboot "${boot_id}"
+
+  remote.exec "snap watch ${change_id}"
+  remote.exec 'snap list pc-kernel-efi-pstore' | awk '$NR != 1 { print $3 }' | MATCH '1'
+  remote.exec 'snap components pc-kernel-efi-pstore' | sed 1d | MATCH 'pc-kernel-efi-pstore\+efi-comp\s+installed'
+
+  remote.exec sudo modprobe efi_pstore
+  remote.exec modinfo --filename efi_pstore | MATCH '/lib/modules/.*/updates/efi-comp'
+
+  echo 'Swap back to the first new kernel and component pc-kernel-mac80211-hwsim+wifi-comp again'
+  gendeveloper1 sign-model < "${FOURTH_MODEL_JSON}" > fourth-model.assert
+  remote.push fourth-model.assert
+
+  boot_id="$(tests.nested boot-id)"
+  change_id="$(remote.exec 'sudo snap remodel --no-wait fourth-model.assert')"
+  remote.wait-for reboot "${boot_id}"
+
+  # same deal here, we expect two reboots
+  boot_id="$(tests.nested boot-id)"
+  remote.wait-for reboot "${boot_id}"
+
+  remote.exec "snap watch ${change_id}"
+  remote.exec 'snap list pc-kernel-mac80211-hwsim' | awk '$NR != 1 { print $3 }' | MATCH '1'
+  remote.exec 'snap components pc-kernel-mac80211-hwsim' | sed 1d | MATCH 'pc-kernel-mac80211-hwsim\+wifi-comp\s+installed'
+
+  remote.exec sudo modprobe mac80211_hwsim
+  remote.exec ip link show wlan0
+  remote.exec modinfo --filename mac80211_hwsim | MATCH '/lib/modules/.*/updates/wifi-comp'

From 4512f34cf207401a0ce5e0e8395728382aa45b5c Mon Sep 17 00:00:00 2001
From: Andrew Phelps <andrew.phelps@canonical.com>
Date: Wed, 5 Feb 2025 21:11:47 -0500
Subject: [PATCH 2/4] o/snapstate: fix unused slice in test setup that is not
 being used

---
 overlord/snapstate/snapstate_install_test.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/overlord/snapstate/snapstate_install_test.go b/overlord/snapstate/snapstate_install_test.go
index b503313cef5..c223aae28ff 100644
--- a/overlord/snapstate/snapstate_install_test.go
+++ b/overlord/snapstate/snapstate_install_test.go
@@ -178,6 +178,9 @@ func expectedDoInstallTasks(typ snap.Type, opts, compOpts, discards int, startTa
 		expected = append(expected, "run-hook[default-configure]")
 	}
 
+	// TODO: it seems that removing this line doesn't break any tests
+	expected = append(expected, tasksBeforeDiscard...)
+
 	expected = append(expected, "start-snap-services")
 	for i := 0; i < discards; i++ {
 		expected = append(expected,

From a2ebb88efb2834cde3c6a4521808d1b08fa56c0d Mon Sep 17 00:00:00 2001
From: Andrew Phelps <andrew.phelps@canonical.com>
Date: Fri, 14 Feb 2025 13:43:15 -0500
Subject: [PATCH 3/4] o/devicestate: add test for case where we swap back to
 already installed kernel during remodel

---
 .../devicestate/devicestate_remodel_test.go   | 147 ++++++++++++++++++
 1 file changed, 147 insertions(+)

diff --git a/overlord/devicestate/devicestate_remodel_test.go b/overlord/devicestate/devicestate_remodel_test.go
index e1e70132b02..fc7d311e393 100644
--- a/overlord/devicestate/devicestate_remodel_test.go
+++ b/overlord/devicestate/devicestate_remodel_test.go
@@ -7682,6 +7682,153 @@ func (s *deviceMgrRemodelSuite) TestRemodelWithComponentsNewComponentSwitchSnap(
 	})
 }
 
+func (s *deviceMgrRemodelSuite) TestRemodelWithComponentsToAlreadyInstalledKernelWithComponents(c *C) {
+	// since remodeling can leave around old kernels, we need to handle the case
+	// where we're remodeling to a kernel that happens to already be installed.
+	// this case covers that.
+
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	s.state.Set("seeded", true)
+	s.state.Set("refresh-privacy-key", "some-privacy-key")
+
+	now := time.Now()
+	restore := devicestate.MockTimeNow(func() time.Time { return now })
+	defer restore()
+
+	currentModel := s.brands.Model("canonical", "pc-model", map[string]interface{}{
+		"architecture": "amd64",
+		"base":         "core24",
+		"snaps": []interface{}{
+			map[string]interface{}{
+				"name":            "pc-kernel",
+				"id":              fakeSnapID("pc-kernel"),
+				"type":            "kernel",
+				"default-channel": "latest/stable",
+			},
+			map[string]interface{}{
+				"name":            "pc",
+				"id":              fakeSnapID("pc"),
+				"type":            "gadget",
+				"default-channel": "latest/stable",
+			},
+			map[string]interface{}{
+				"name":            "core24",
+				"id":              fakeSnapID("core24"),
+				"type":            "base",
+				"default-channel": "latest/stable",
+			},
+			map[string]interface{}{
+				"name":            "snapd",
+				"id":              fakeSnapID("snapd"),
+				"type":            "snapd",
+				"default-channel": "latest/stable",
+			},
+		},
+	})
+	err := assertstate.Add(s.state, currentModel)
+	c.Assert(err, IsNil)
+
+	err = devicestatetest.SetDevice(s.state, &auth.DeviceState{
+		Brand: "canonical",
+		Model: "pc-model",
+	})
+	c.Assert(err, IsNil)
+
+	snapstatetest.InstallEssentialSnaps(c, s.state, "core24", nil, nil)
+
+	kmodComps := map[string]snap.ComponentType{
+		"kmod": snap.KernelModulesComponent,
+	}
+	snapstatetest.InstallSnap(c, s.state, withComponents("name: iot-kernel\nversion: 1\ntype: kernel\n", kmodComps), nil, &snap.SideInfo{
+		SnapID:   fakeSnapID("iot-kernel"),
+		Revision: snap.R(2),
+		RealName: "iot-kernel",
+		Channel:  "latest/stable",
+	}, snapstatetest.InstallSnapOptions{Required: true})
+
+	var snapst snapstate.SnapState
+	err = snapstate.Get(s.state, "iot-kernel", &snapst)
+	c.Assert(err, IsNil)
+
+	err = snapst.Sequence.AddComponentForRevision(snapst.Current, sequence.NewComponentState(&snap.ComponentSideInfo{
+		Component: naming.NewComponentRef("iot-kernel", "kmod"),
+		Revision:  snap.R(22),
+	}, snap.KernelModulesComponent))
+	c.Assert(err, IsNil)
+
+	snapstate.Set(s.state, "iot-kernel", &snapst)
+
+	newModel := s.brands.Model("canonical", "pc-model", map[string]interface{}{
+		"architecture": "amd64",
+		"base":         "core24",
+		"snaps": []interface{}{
+			map[string]interface{}{
+				"name":            "iot-kernel",
+				"id":              fakeSnapID("iot-kernel"),
+				"type":            "kernel",
+				"default-channel": "latest/stable",
+				"components": map[string]interface{}{
+					"kmod": map[string]interface{}{
+						"presence": "required",
+					},
+				},
+			},
+			map[string]interface{}{
+				"name":            "pc",
+				"id":              fakeSnapID("pc"),
+				"type":            "gadget",
+				"default-channel": "latest/stable",
+			},
+			map[string]interface{}{
+				"name":            "core24",
+				"id":              fakeSnapID("core24"),
+				"type":            "base",
+				"default-channel": "latest/stable",
+			},
+			map[string]interface{}{
+				"name":            "snapd",
+				"id":              fakeSnapID("snapd"),
+				"type":            "snapd",
+				"default-channel": "latest/stable",
+			},
+		},
+	})
+
+	testDeviceCtx := snapstatetest.TrivialDeviceContext{
+		Remodeling:     true,
+		DeviceModel:    newModel,
+		OldDeviceModel: currentModel,
+	}
+
+	tss, err := devicestate.RemodelTasks(
+		context.Background(),
+		s.state,
+		currentModel,
+		newModel,
+		&testDeviceCtx,
+		"99",
+		devicestate.RemodelOptions{},
+	)
+	c.Assert(err, IsNil)
+
+	// re-linking already installed kernel and components, create recovery
+	// system, set model
+	c.Assert(tss, HasLen, 3)
+
+	updateTS := tss[0]
+	checkTaskSetKinds(c, updateTS, []string{
+		"prepare-snap",
+		"update-gadget-assets",
+		"link-snap",
+		"link-component",
+	})
+
+	createRecoverySystem := tss[1].Tasks()[0]
+	checkRecoverySystemSetup(createRecoverySystem, c, now, []interface{}{"1"}, []interface{}{"4"})
+}
+
 func checkRecoverySystemSetup(t *state.Task, c *C, now time.Time, snapsups, compsups []interface{}) {
 	var data map[string]interface{}
 	err := t.Get("recovery-system-setup", &data)

From 82bdecbfe1b3069ad7d53de94cb1f6b43be3dd9d Mon Sep 17 00:00:00 2001
From: Andrew Phelps <andrew.phelps@canonical.com>
Date: Fri, 14 Feb 2025 13:55:30 -0500
Subject: [PATCH 4/4] o/snapstate: make doc comments on
 (Add)LinkNewBaseOrKernel a bit better

---
 overlord/snapstate/snapstate.go | 45 ++++++++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 9 deletions(-)

diff --git a/overlord/snapstate/snapstate.go b/overlord/snapstate/snapstate.go
index 295b1bd703c..88222e2446d 100644
--- a/overlord/snapstate/snapstate.go
+++ b/overlord/snapstate/snapstate.go
@@ -3366,8 +3366,21 @@ func MigrateHome(st *state.State, snaps []string) ([]*state.TaskSet, error) {
 	return tss, nil
 }
 
-// LinkNewBaseOrKernel creates a new task set with prepare/link-snap, and
-// additionally update-gadget-assets for the kernel snap, tasks for a remodel.
+// LinkNewBaseOrKernel creates a new task set that enables swapping to a base or
+// kernel snap that is already installed on the system. The primary use case for
+// this function is remodeling.
+//
+// For bases, we create prepare-snap and link-snap tasks. Technically this would
+// create link-component tasks for any installed components, but bases do not
+// currently use components.
+//
+// For kernels, we create prepare-snap, an update-gadget-assets task (if
+// needed), link-snap, and link-component tasks for any installed components.
+//
+// Note that this function previously created a prepare-kernel-snap task, but
+// this was not needed. Since this function is only used if the snap is
+// installed already installed, then it is expected that the drivers tree is
+// present. Thus, the prepare-kernel-snap task would be redundant.
 func LinkNewBaseOrKernel(st *state.State, name string, fromChange string) (*state.TaskSet, error) {
 	var snapst SnapState
 	err := Get(st, name, &snapst)
@@ -3441,11 +3454,9 @@ func addLinkNewBaseOrKernelTasks(st *state.State, snapst SnapState, ts *state.Ta
 
 	// preserve the same order as during the update
 	if info.Type() == snap.TypeKernel {
-		// this previously created a prepare-kernel-snap task, but this was not
-		// needed since this function is only used to swap to kernel snaps that
-		// are still installed on the system (but not in use. this can happen
-		// because remodeling leaves around old snaps). if the snap is
-		// installed, then it is expected that the drivers tree is present.
+		// this previously created a prepare-kernel-snap task. however, this
+		// isn't needed since we're only using this function to swap to already
+		// installed kernels. thus, the drivers tree should be present.
 
 		// kernel snaps can carry boot assets
 		gadgetUpdate := st.NewTask("update-gadget-assets", fmt.Sprintf(i18n.G("Update assets from %s %q (%s) for remodel"), info.Type(), info.InstanceName(), snapst.Current))
@@ -3492,8 +3503,24 @@ func findSnapSetupTask(tasks []*state.Task) (*state.Task, *SnapSetup, error) {
 	return nil, nil, nil
 }
 
-// AddLinkNewBaseOrKernel creates the same tasks as LinkNewBaseOrKernel but adds
-// them to the provided task set.
+// AddLinkNewBaseOrKernel appends tasks to a given task set. This enables
+// swapping to a base or kernel snap that is already installed on the system.
+// The primary use case for this function is remodeling.
+//
+// It is expected that the given task set contains a snap setup task.
+// Additionally, it should not perform any modifications to the local system.
+//
+// For bases, we create a link-snap task. Technically this would create
+// link-component tasks for any installed components, but bases do not currently
+// use components.
+//
+// For kernels, we create an update-gadget-assets task (if needed), link-snap,
+// and link-component tasks for any installed components.
+//
+// Note that this function previously created a prepare-kernel-snap task, but
+// this was not needed. Since this function is only used if the snap is
+// installed already installed, then it is expected that the drivers tree is
+// present. Thus, the prepare-kernel-snap task would be redundant.
 func AddLinkNewBaseOrKernel(st *state.State, ts *state.TaskSet) (*state.TaskSet, error) {
 	if ts.MaybeEdge(LastBeforeLocalModificationsEdge) != nil {
 		return nil, errors.New("internal error: cannot add tasks to link new base or kernel to task set that introduces local modifications")