diff --git a/src/lib/components/start/RobotSelector.svelte b/src/lib/components/start/RobotSelector.svelte
index bd4d366..4c13c06 100644
--- a/src/lib/components/start/RobotSelector.svelte
+++ b/src/lib/components/start/RobotSelector.svelte
@@ -1,7 +1,7 @@
 <script lang="ts">
 import { type Robot, robots as allRobots } from "$domain/robots";
 import { Screen, screen, selected } from "$state/app.svelte";
-import { Mode, code, mode, robot } from "$state/workspace.svelte";
+import { Mode, code, mode, robot, saveState } from "$state/workspace.svelte";
 
 interface Props {
 	robots: Robot[][];
@@ -18,6 +18,7 @@ function select(type: Robot) {
 		robot.set(allRobots[type.defaultRobot]);
 		mode.set(type.mode);
 		screen.set(Screen.WORKSPACE);
+		saveState.set(true);
 		return;
 	}
 
diff --git a/src/lib/components/workspace/python/Python.svelte b/src/lib/components/workspace/python/Python.svelte
index ea613da..5d37cd8 100644
--- a/src/lib/components/workspace/python/Python.svelte
+++ b/src/lib/components/workspace/python/Python.svelte
@@ -3,8 +3,13 @@ import defaultProgram from "$assets/default-program.py?raw";
 import CodeEditor from "$components/ui/CodeEditor.svelte";
 import Tree from "$components/ui/Tree.svelte";
 import type { Tree as TreeType } from "$components/ui/Tree.types";
-import { PythonHandle } from "$domain/handles";
-import { code, handle, microPythonIO } from "$state/workspace.svelte";
+import { FileHandle, PythonHandle } from "$domain/handles";
+import {
+	code,
+	handle,
+	microPythonIO,
+	saveState,
+} from "$state/workspace.svelte";
 import { get } from "svelte/store";
 import type MicroPythonIO from "../../../micropython";
 import Terminal from "./Terminal.svelte";
@@ -41,6 +46,7 @@ async function getTree(io: MicroPythonIO, folder: string): Promise<TreeType> {
 
 microPythonIO.subscribe(async (io) => {
 	if (!io) return;
+	if (!get(saveState) || get(handle)) return;
 
 	if (!(await io.fs.exists("main.py"))) {
 		await io.fs.write("main.py", defaultProgram);
@@ -53,6 +59,10 @@ microPythonIO.subscribe(async (io) => {
 	handle.set(new PythonHandle("main.py"));
 });
 
+handle.subscribe((handle) => {
+	if (handle instanceof FileHandle) tree = undefined;
+});
+
 async function select(tree: string[]) {
 	selected = tree;