diff --git a/src/makecode/utils.ts b/src/makecode/utils.ts
index f208e8329..7fe55ecf8 100644
--- a/src/makecode/utils.ts
+++ b/src/makecode/utils.ts
@@ -5,6 +5,7 @@ import upperFirst from "lodash.upperfirst";
 import { DatasetEditorJsonFormat, GestureData } from "../model";
 import { getMainScript } from "./generate-main-scripts";
 import { getAutogeneratedTs, getDatasetJson } from "./generate-custom-scripts";
+import { Project } from "@microbit/makecode-embed/react";
 
 export const filenames = {
   mainTs: "main.ts",
@@ -15,6 +16,9 @@ export const filenames = {
   readme: "README.md",
 };
 
+const extensionName = "machine-learning";
+const extensionURL = "github:microbit-foundation/pxt-microbit-ml#v0.4.3";
+
 const pxt = {
   name: "Untitled",
   description: "",
@@ -22,7 +26,7 @@ const pxt = {
     core: "*",
     microphone: "*",
     radio: "*", // Needed to compile.
-    "machine-learning": "github:microbit-foundation/pxt-microbit-ml#v0.4.3",
+    [extensionName]: extensionURL,
   },
   files: Object.values(filenames),
   preferredEditor: "blocksprj",
@@ -56,17 +60,38 @@ export const generateProject = (
 
 export const generateCustomFiles = (
   gestureState: DatasetEditorJsonFormat,
-  model: LayersModel | undefined
+  model: LayersModel | undefined,
+  project?: Project
 ) => {
   const { data: gestures } = gestureState;
   const useableGestures = model ? gestures : [];
-  return {
+
+  const customFiles = {
     [filenames.autogenerated]: model
       ? getAutogeneratedTs(useableGestures, model)
       : "",
     // Save all gestures to dataset.json.
     [filenames.datasetJson]: getDatasetJson(gestureState),
   };
+
+  if (!project) {
+    return customFiles;
+  }
+
+  const currentPxtJSON = project.text?.[filenames.pxtJson];
+  if (currentPxtJSON) {
+    try {
+      const updatedPxt = JSON.parse(currentPxtJSON) as typeof pxt;
+      updatedPxt.dependencies[extensionName] = extensionURL;
+      return {
+        ...customFiles,
+        [filenames.pxtJson]: JSON.stringify(updatedPxt),
+      };
+    } catch (e) {
+      // If we reach this case, the project is in theory already very broken
+      // as it will be missing a pxt file and the extension.
+    }
+  }
 };
 
 export interface ActionName {
diff --git a/src/store.ts b/src/store.ts
index 98a2ea1a8..b4171abbd 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -48,7 +48,7 @@ const updateProject = (
     text: {
       ...project.text,
       ...(projectEdited
-        ? generateCustomFiles(gestureData, model)
+        ? generateCustomFiles(gestureData, model, project)
         : generateProject(gestureData, model).text),
     },
   };