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.scopeHooks
+  );
   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,102 @@
+const glob = typeof window === "undefined" ? global : window;
+
+// These functions will be added to the scope of the cell
+const onScopeFunctions = ["setTimeout", "setInterval", "console"] as const;
+
+// These functions will be attached to the window during cell execution
+const onWindowFunctions = [
+  ...onScopeFunctions,
+  "EventTarget.prototype.addEventListener",
+] as const;
+
+export const originalReferences: ScopeHooks & WindowHooks = {
+  setTimeout: glob.setTimeout,
+  setInterval: glob.setInterval,
+  console: glob.console,
+  "EventTarget.prototype.addEventListener":
+    glob.EventTarget.prototype.addEventListener,
+};
+
+export type ScopeHooks = { [K in typeof onScopeFunctions[number]]: any };
+
+export type WindowHooks = { [K in typeof onWindowFunctions[number]]: any };
+
+function setProperty(base: Object, path: string, value: any) {
+  const layers = path.split(".");
+  if (layers.length > 1) {
+    const toOverride = layers.pop()!;
+    // Returns second last path member
+    const layer = layers.reduce((o, i) => o[i], base as any);
+    layer[toOverride] = value;
+  } else {
+    (base as any)[path] = value;
+  }
+}
+
+export class HookExecution {
+  public disposers: Array<() => void> = [];
+  public scopeHooks: ScopeHooks = {
+    setTimeout: this.createHookedFunction(setTimeout, (ret) => {
+      clearTimeout(ret);
+    }),
+    setInterval: this.createHookedFunction(setInterval, (ret) => {
+      clearInterval(ret);
+    }),
+    console: {
+      ...originalReferences.console,
+      log: (...args: any) => {
+        // TODO: broadcast output to console view
+        originalReferences.console.log(...args);
+      },
+    },
+  };
+
+  public windowHooks: WindowHooks = {
+    ...this.scopeHooks,
+    ["EventTarget.prototype.addEventListener"]: undefined,
+  };
+
+  constructor() {
+    if (typeof EventTarget !== "undefined") {
+      this.windowHooks["EventTarget.prototype.addEventListener"] =
+        this.createHookedFunction(
+          EventTarget.prototype.addEventListener as any,
+          function (this: any, _ret, args) {
+            this.removeEventListener(args[0], args[1]);
+          }
+        );
+    }
+  }
+
+  public attachToWindow() {
+    onWindowFunctions.forEach((path) =>
+      setProperty(glob, path, this.windowHooks[path])
+    );
+  }
+
+  public detachFromWindow() {
+    onWindowFunctions.forEach((path) =>
+      setProperty(glob, path, originalReferences[path])
+    );
+  }
+
+  public dispose() {
+    this.disposers.forEach((d) => d());
+  }
+
+  private createHookedFunction<T, Y>(
+    original: (...args: T[]) => Y,
+    disposer: (ret: Y, args: T[]) => void
+  ) {
+    const self = this;
+    return function newFunction(this: any): Y {
+      const callerArguments = arguments;
+      const ret = original.apply(this, callerArguments as any); // TODO: fix any?
+      const ctx = this;
+      self.disposers.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,18 +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);
+
+      disposeEveryRun.push(() => hookExecution.dispose());
+      hookExecution.attachToWindow();
+
       let executionPromise: Promise<any>;
+
       try {
         executionPromise = mod.factoryFunction.apply(
           undefined,
           argsToCallFunctionWith
-        ); // TODO: what happens with disposers if a rerun of this function is slow / delayed?
+        );
       } 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);
         }
@@ -211,7 +214,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 { ScopeHooks } 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: ScopeHooks
+) {
   const scope = {
     autorun,
     $: context.context,
@@ -82,7 +86,9 @@ export function createExecutionScope(context: TypeCellContext<any>) {
     // stored,
     // view,
     observable,
+    ...hookContext,
   };
+
   return scope;
 }