diff --git a/packages/core/test/tool-shell.test.ts b/packages/core/test/tool-shell.test.ts
index 6fd43ee2a2e..217414aa2ce 100644
--- a/packages/core/test/tool-shell.test.ts
+++ b/packages/core/test/tool-shell.test.ts
@@ -171,9 +171,11 @@ const withSession = <A, E, R>(directory: string, body: (registry: ToolRegistry.I
     })
     const locations = yield* LocationServiceMap.Service
     const locationLayer = locations.get(location)
-    const registry = yield* ToolRegistry.Service.pipe(Effect.provide(locationLayer))
-    yield* waitForTool(registry, ShellTool.name)
-    return yield* body(registry).pipe(Effect.provide(locationLayer))
+    return yield* Effect.gen(function* () {
+      const registry = yield* ToolRegistry.Service
+      yield* waitForTool(registry, ShellTool.name)
+      return yield* body(registry)
+    }).pipe(Effect.provide(locationLayer), Effect.ensuring(locations.invalidate(location)))
   })
 
 describe("ShellTool", () => {