diff --git a/lib/utils/device/deploy.ts b/lib/utils/device/deploy.ts
index 6717735b..cae5c955 100644
--- a/lib/utils/device/deploy.ts
+++ b/lib/utils/device/deploy.ts
@@ -627,17 +627,10 @@ export function generateTargetState(
 		};
 
 		opts.environment = _.merge(opts.environment, env[name]);
-		// This function can be called with a subset of the
-		// build tasks, when a single dockerfile has changed
-		// when livepushing, so check the build task exists for
-		// this composition entry (everything else in this
-		// function comes from the composition which doesn't
-		// change)
-		let contract;
-		if (name in keyedBuildTasks) {
-			contract = keyedBuildTasks[name].contract;
-		}
-
+		// This function should always be called with all the build tasks
+		// so we can construct the correct target state so we don't really need
+		// to check that the key exists on the `keyedBuildTasks` object
+		const contract = keyedBuildTasks[name].contract;
 		const task = keyedBuildTasks[name];
 
 		services[idx] = {
diff --git a/lib/utils/device/live.ts b/lib/utils/device/live.ts
index 6bc83cc9..48bafea5 100644
--- a/lib/utils/device/live.ts
+++ b/lib/utils/device/live.ts
@@ -423,7 +423,12 @@ export class LivepushManager {
 			// If we re-apply the target state, the supervisor
 			// should recreate the container
 			await this.api.setTargetState(
-				generateTargetState(currentState, this.composition, [buildTask], {}),
+				generateTargetState(
+					currentState,
+					this.composition,
+					this.buildTasks,
+					{},
+				),
 			);
 
 			await this.awaitDeviceStateSettle();