diff --git a/src/index.ts b/src/index.ts
index 0b46f19..2eb9c26 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -64,6 +64,17 @@ export default function skillReinject(pi: ExtensionAPI): void {
 		});
 	}
 
+	function restoreSessionState(ctx: ExtensionContext): void {
+		detectAndCachePiAutoCompact(pi, runtime);
+		const settings = readSettings(ctx);
+		const branch = ctx.sessionManager.getBranch();
+		const loaded = loadStateFromBranch(branch);
+		applyExtensionState(state, loaded ?? createInitialState());
+		if (!loaded) {
+			rescanSkillsFromBranch(state, branch, ctx.cwd, registeredSkills, settings.trackReadPaths);
+		}
+	}
+
 	function handleSessionCompact(event: SessionCompactEvent, ctx: ExtensionContext): void {
 		const settings = readSettings(ctx);
 		const shouldReinject = consumeCompactionOnSessionCompact(
@@ -99,14 +110,11 @@ export default function skillReinject(pi: ExtensionAPI): void {
 	}
 
 	pi.on("session_start", async (_event, ctx) => {
-		detectAndCachePiAutoCompact(pi, runtime);
-		const settings = readSettings(ctx);
-		const branch = ctx.sessionManager.getBranch();
-		const loaded = loadStateFromBranch(branch);
-		applyExtensionState(state, loaded ?? createInitialState());
-		if (!loaded) {
-			rescanSkillsFromBranch(state, branch, ctx.cwd, registeredSkills, settings.trackReadPaths);
-		}
+		restoreSessionState(ctx);
+	});
+
+	pi.on("session_tree", async (_event, ctx) => {
+		restoreSessionState(ctx);
 	});
 
 	pi.on("session_before_compact", async () => {