Refactor Hooks

packages/engine/src/CellEvaluator.ts

@@ -1,5 +1,6 @@
 import { TypeCellContext } from "./context";
 import { ModuleExecution, runModule } from "./executor";
+import { HookExecution } from "./HookExecution";
 import { createExecutionScope, getModulesFromTypeCellCode } from "./modules";
 import { isReactView } from "./reactView";
 
@@ -49,7 +50,11 @@ export function createCellEvaluator(
     onOutputChanged(error);
   }
 
-  const executionScope = createExecutionScope(typecellContext);
+  const hookExecution = new HookExecution();
+  const executionScope = createExecutionScope(
+    typecellContext,
+    hookExecution.context
+  );
   let moduleExecution: ModuleExecution | undefined;
 
   async function evaluate(compiledCode: string) {
@@ -69,6 +74,7 @@ export function createCellEvaluator(
         modules[0],
         typecellContext,
         resolveImport,
+        hookExecution,
         beforeExecuting,
         onExecuted,
         onError,

packages/engine/src/HookExecution.ts

@@ -0,0 +1,115 @@
+const glob = typeof window === "undefined" ? global : window;
+
+const overrideFunctions = [
+  "setTimeout",
+  "setInterval",
+  "console",
+  "EventTarget",
+  "window",
+  "global",
+] as const;
+
+const originalReferences: HookContext = {
+  setTimeout: glob.setTimeout,
+  setInterval: glob.setInterval,
+  console: glob.console,
+  EventTarget: glob.EventTarget,
+  window: window,
+  global: global,
+};
+
+export type HookContext = { [K in typeof overrideFunctions[number]]: any };
+
+export class HookExecution {
+  public disposers: Array<() => void> = [];
+  public context: HookContext = {
+    setTimeout: this.applyDisposer(setTimeout, this.disposers, (ret) => {
+      clearTimeout(ret);
+    }),
+    setInterval: this.applyDisposer(setInterval, this.disposers, (ret) => {
+      clearInterval(ret);
+    }),
+    console: {
+      ...originalReferences.console,
+      log: (...args: any) => {
+        // TODO: broadcast output to console view
+        originalReferences.console.log(...args);
+      },
+    },
+    EventTarget: undefined,
+    window: undefined,
+    global: undefined,
+  };
+
+  constructor() {
+    if (typeof EventTarget !== "undefined") {
+      this.context.EventTarget = {
+        EventTarget,
+        prototype: {
+          ...EventTarget.prototype,
+          addEventListener: this.applyDisposer(
+            EventTarget.prototype.addEventListener as any,
+            this.disposers,
+            function (this: any, _ret, args) {
+              this.removeEventListener(args[0], args[1]);
+            }
+          ),
+        },
+      };
+    }
+
+    if (typeof window !== "undefined") {
+      this.context.window = {
+        ...window,
+        setTimeout: this.context.setTimeout,
+        setInterval: this.context.setInterval,
+        console: this.context.console,
+        EventTarget: this.context.EventTarget,
+      };
+    }
+
+    if (typeof global !== "undefined") {
+      this.context.global = {
+        ...global,
+        setTimeout: this.context.setTimeout,
+        setInterval: this.context.setInterval,
+        console: this.context.console,
+        EventTarget: this.context.EventTarget,
+      };
+    }
+  }
+
+  public attachToWindow() {
+    overrideFunctions
+      .filter((name) => name !== "window" && name !== "global")
+      .forEach((name) => {
+        (glob as any)[name] = this.context[name];
+      });
+  }
+
+  public detachFromWindow() {
+    overrideFunctions
+      .filter((name) => name !== "window" && name !== "global")
+      .forEach((name) => {
+        (glob as any)[name] = originalReferences[name];
+      });
+  }
+
+  public dispose() {
+    this.disposers.forEach((d) => d());
+  }
+
+  private applyDisposer<T, Y>(
+    original: (...args: T[]) => Y,
+    disposes: Array<() => void>,
+    disposer: (ret: Y, args: T[]) => void
+  ) {
+    return function newFunction(this: any): Y {
+      const callerArguments = arguments;
+      const ret = original.apply(this, callerArguments as any); // TODO: fix any?
+      const ctx = this;
+      disposes.push(() => disposer.call(ctx, ret, callerArguments as any));
+      return ret;
+    };
+  }
+}

packages/engine/src/executor.ts

@@ -1,6 +1,6 @@
 import { autorun, runInAction } from "mobx";
 import { TypeCellContext } from "./context";
-import { installHooks } from "./hookDisposables";
+import { HookExecution } from "./HookExecution";
 import { Module } from "./modules";
 import { isStored } from "./storage/stored";
 import { isView } from "./view";
@@ -61,6 +61,7 @@ export async function runModule(
   mod: Module,
   context: TypeCellContext<any>,
   resolveImport: (module: string) => any,
+  hookExecution: HookExecution,
   beforeExecuting: () => void,
   onExecuted: (exports: any) => void,
   onError: (error: any) => void,
@@ -96,7 +97,7 @@ export async function runModule(
     );
   }
 
-  const execute = async () => {
+  async function execute(this: any) {
     try {
       if (wouldLoopOnAutorun) {
         detectedLoop = true;
@@ -114,25 +115,20 @@ export async function runModule(
       disposeEveryRun.length = 0; // clear existing array in this way, because we've passed the reference to resolveDependencyArray and want to keep it intact
 
       beforeExecuting();
-      const hooks = installHooks();
-      disposeEveryRun.push(hooks.disposeAll);
-      let executionPromise: Promise<any>;
+
+      disposeEveryRun.push(hookExecution.dispose);
+      hookExecution.attachToWindow();
+
       try {
-        executionPromise = mod.factoryFunction.apply(
-          undefined,
-          argsToCallFunctionWith
-        ); // TODO: what happens with disposers if a rerun of this function is slow / delayed?
+        await mod.factoryFunction.apply(undefined, argsToCallFunctionWith);
       } finally {
-        // Hooks are only installed for sync code. Ideally, we'd want to run it for all code, but then we have the chance hooks will affect other parts of the TypeCell (non-user) code
-        // (we ran into this that notebooks wouldn't be saved (_.debounce), and also that setTimeout of Monaco blink cursor would be hooked)
-        hooks.unHookAll();
+        hookExecution.detachFromWindow();
+
         if (previousVariableDisposer) {
           previousVariableDisposer(exports);
         }
       }
 
-      await executionPromise;
-
       // Running the assignments to `context` in action should be a performance improvement to prevent triggering observers one-by-one
       wouldLoopOnAutorun = true;
       runInAction(() => {
@@ -211,7 +207,7 @@ export async function runModule(
       //reject(e);
       resolve();
     }
-  };
+  }
 
   const autorunDisposer = autorun(() => execute());
 

packages/engine/src/modules.ts

@@ -1,5 +1,6 @@
 import { TypeCellContext } from "./context";
 import { observable, untracked, computed, autorun } from "mobx";
+import { HookContext } from "./HookExecution";
 // import { stored } from "./storage/stored";
 // import { view } from "./view";
 
@@ -71,7 +72,10 @@ function createDefine(modules: Module[]) {
   };
 }
 
-export function createExecutionScope(context: TypeCellContext<any>) {
+export function createExecutionScope(
+  context: TypeCellContext<any>,
+  hookContext: HookContext
+) {
   const scope = {
     autorun,
     $: context.context,
@@ -82,6 +86,7 @@ export function createExecutionScope(context: TypeCellContext<any>) {
     // stored,
     // view,
     observable,
+    ...hookContext,
   };
   return scope;
 }