feat: interruptible python (#52587)

Co-authored-by: Shaun Hamilton <shauhami020@gmail.com>

client/src/templates/Challenges/classic/xterm.tsx

@@ -41,13 +41,13 @@ export const XtermTerminal = ({
       if (termContainerRef.current) term.open(termContainerRef.current);
       fitAddon.fit();
 
-      const print = (text: string) => term?.writeln(`>>> ${text}`);
+      const print = (text?: string) => term?.writeln(`>>> ${text ?? ''}`);
 
       // TODO: prevent user from moving cursor outside the current input line and
       // handle insertion and deletion properly. While backspace and delete don't
       // seem to work, we can use "\x1b[0K" to clear from the cursor to the end.
       // Also, we should not add special characters to the userinput string.
-      const input = (text: string) => {
+      const input = (text?: string) => {
         print(text);
         let userinput = '';
         // Eslint is correct that this only gets assigned once, but we can't use
@@ -58,7 +58,12 @@ export const XtermTerminal = ({
 
         const done = () => {
           disposable?.dispose();
-          navigator.serviceWorker.controller?.postMessage(userinput);
+          navigator.serviceWorker.controller?.postMessage(
+            JSON.stringify({
+              type: 'msg',
+              value: userinput
+            })
+          );
         };
 
         const keyListener = (key: string) => {
@@ -80,11 +85,15 @@ export const XtermTerminal = ({
         if (disposable) disposables.push(disposable);
       };
       const reset = () => {
-        term?.reset();
+        // Ironically, term.reset(), while synchronous, is not a reliable way to
+        // reset the terminal. It does not clear the input buffer, so old print
+        // statements can still appear. The \x1bc (ESC c) escape sequence triggers
+        // a full terminal reset, which is what we want.
+        term?.write('\x1bc');
         disposables.forEach(disposable => disposable.dispose());
         disposables.length = 0;
       };
-      registerTerminal({ print, input }, reset);
+      registerTerminal({ print, input, reset });
     }
 
     void createTerminal();

client/src/templates/Challenges/redux/execute-challenge-saga.js

@@ -34,8 +34,8 @@ import {
   updateProjectPreview
 } from '../utils/build';
 import {
-  getPythonWorker,
-  resetPythonWorker
+  interruptCodeExecution,
+  runPythonCode
 } from '../utils/python-worker-handler';
 import { executeGA } from '../../../redux/actions';
 import { fireConfetti } from '../../../utils/fire-confetti';
@@ -311,14 +311,14 @@ function* updatePython(challengeData) {
   // functions to handle transforming code, embedding it and building the
   // final html. Then we can just use the transformation function here.
   const buildData = yield buildChallengeData(challengeData);
-  resetPythonWorker();
-  const worker = getPythonWorker();
+  interruptCodeExecution();
   const code = {
     contents: buildData.sources.index,
     editableContents: buildData.sources.editableContents,
     original: buildData.sources.original
   };
-  worker.postMessage({ code });
+
+  runPythonCode(code);
   // TODO: proxy errors to the console
 }
 

client/src/templates/Challenges/utils/python-worker-handler.ts

@@ -1,3 +1,4 @@
+// TODO: This might be cleaner as a class.
 import pythonWorkerData from '../../../../config/browser-scripts/python-worker.json';
 
 const pythonWorkerSrc = `/js/${pythonWorkerData.filename}.js`;
@@ -5,9 +6,16 @@ const pythonWorkerSrc = `/js/${pythonWorkerData.filename}.js`;
 let worker: Worker | null = null;
 let testWorker: Worker | null = null;
 let listener: ((event: MessageEvent) => void) | null = null;
-let resetTerminal: (() => void) | null = null;
+type Code = {
+  contents: string;
+  editableContents: string;
+  original: string;
+};
+// We need to keep track of the last code message so we can re-run it if the
+// worker is reset.
+let lastCodeMessage: Code | null = null;
 
-export function getPythonWorker(): Worker {
+function getPythonWorker(): Worker {
   if (!worker) {
     worker = new Worker(pythonWorkerSrc);
   }
@@ -23,8 +31,14 @@ export function getPythonTestWorker(): Worker {
 
 type PythonWorkerEvent = {
   data: {
-    type: 'print' | 'input' | 'contentLoaded';
-    text: string;
+    type:
+      | 'print'
+      | 'input'
+      | 'contentLoaded'
+      | 'reset'
+      | 'stopped'
+      | 'is-alive';
+    text?: string;
   };
 };
 
@@ -35,31 +49,75 @@ type PythonWorkerEvent = {
  * @param handlers.input - A function that handles input messages from the python worker
  * @param reset - A function that resets the terminal
  */
-export function registerTerminal(
-  handlers: {
-    print: (text: string) => void;
-    input: (text: string) => void;
-  },
-  reset: () => void
-): void {
+export function registerTerminal(handlers: {
+  print: (text?: string) => void;
+  input: (text?: string) => void;
+  reset: () => void;
+}): void {
   const pythonWorker = getPythonWorker();
   if (listener) pythonWorker.removeEventListener('message', listener);
   listener = (event: PythonWorkerEvent) => {
+    // TODO: refactor text -> value or msg.
     const { type, text } = event.data;
-    // Ignore contentLoaded messages for now.
-    if (type === 'contentLoaded') return;
-    handlers[type](text);
+
+    // TODO: this is a bit messy with the 'handlers' as well as the implicit
+    // handlers reacting to stopped and contentLoaded messages.
+    if (type === 'contentLoaded') return; // Ignore contentLoaded messages for now.
+    if (type === 'is-alive') {
+      clearTimeout(Number(text));
+      return;
+    }
+    // 'stopped' means the worker is ignoring 'run' messages.
+    if (type === 'stopped') {
+      clearTimeout(Number(text));
+      sendListenMessage();
+      // Generally, we get here if the learner changes their code while the
+      // worker is busy. In that case, we want to re-run the code on receipt of
+      // the 'stopped' message.
+      if (lastCodeMessage) runPythonCode(lastCodeMessage);
+    } else {
+      handlers[type](text);
+    }
   };
   pythonWorker.addEventListener('message', listener);
-  resetTerminal = reset;
 }
 
 /**
- * Terminates the existing python worker and creates a new one.
+ * Tries to cancel the currently running code and, if it cannot, terminate the worker.
  */
-export function resetPythonWorker(): void {
-  if (resetTerminal) resetTerminal();
-  worker?.terminate();
-  worker = new Worker(pythonWorkerSrc);
-  if (listener) worker.addEventListener('message', listener);
+export function interruptCodeExecution(): void {
+  const resetId = setTimeout(() => {
+    getPythonWorker().terminate();
+    worker = new Worker(pythonWorkerSrc);
+    if (listener) getPythonWorker().addEventListener('message', listener);
+  }, 1000) as unknown as number; // This is running the browser, so setTimeout returns a number, but TS doesn't know that.
+  navigator.serviceWorker.controller?.postMessage(
+    JSON.stringify({
+      type: 'cancel',
+      value: '' + resetId // Converting to string for convenience. (TODO: check if necesary)
+    })
+  );
+
+  // TODO: Since loading pyodide is slow, there's a risk that this will
+  // terminate the worker before it's finished loading. As such we should check
+  // if the worker has loaded before attempting to reset it (or send run
+  // messages).
+
+  // TODO: sort out the terminology.
+  getPythonWorker().postMessage({ type: 'cancel', value: resetId });
+}
+
+export function runPythonCode(code: {
+  contents: string;
+  editableContents: string;
+  original: string;
+}): void {
+  lastCodeMessage = code;
+  getPythonWorker().postMessage({ type: 'run', code });
+}
+
+// If the python worker reports that it has stopped, we need to send a listen
+// message to get it to listen to run messages again.
+function sendListenMessage(): void {
+  getPythonWorker().postMessage({ type: 'listen' });
 }

tools/client-plugins/browser-scripts/python-worker.ts

@@ -2,6 +2,7 @@
 // and 'import' defaults to .mjs
 import { loadPyodide, type PyodideInterface } from 'pyodide/pyodide.js';
 import pkg from 'pyodide/package.json';
+import type { PyProxy, PythonError } from 'pyodide/ffi';
 
 const ctx: Worker & typeof globalThis = self as unknown as Worker &
   typeof globalThis;
@@ -10,6 +11,7 @@ let pyodide: PyodideInterface;
 
 interface PythonRunEvent extends MessageEvent {
   data: {
+    type: 'run';
     code: {
       contents: string;
       editableContents: string;
@@ -18,6 +20,26 @@ interface PythonRunEvent extends MessageEvent {
   };
 }
 
+interface ListenRequestEvent extends MessageEvent {
+  data: {
+    type: 'listen';
+  };
+}
+
+interface CancelEvent extends MessageEvent {
+  data: {
+    type: 'cancel';
+    value: number;
+  };
+}
+
+// Since messages are buffered, it needs to be possible to discard 'run'
+// messages. Otherwise messages could build up while the worker is busy (for
+// example, while loading pyodide) and the work would try to process them in
+// sequence. Instead, it will ignore messages until it receives a 'listen'
+// message and will inform the client every time it starts ignoring messages.
+let ignoreRunMessages = true;
+
 async function setupPyodide() {
   if (pyodide) return pyodide;
 
@@ -31,6 +53,13 @@ async function setupPyodide() {
   // pyodide modifies self while loading.
   Object.freeze(self);
 
+  ignoreRunMessages = true;
+  postMessage({ type: 'stopped' });
+}
+
+void setupPyodide();
+
+function initRunPython() {
   // eslint-disable-next-line @typescript-eslint/no-unsafe-call
   const str = pyodide.globals.get('str') as (x: unknown) => string;
 
@@ -47,7 +76,13 @@ async function setupPyodide() {
     request.open('POST', '/python/intercept-input/', false);
     request.send(null);
 
-    return request.responseText;
+    // We want to raise a KeyboardInterrupt if the user cancels. To do that,
+    // this function returns a JS object with the 'type' property set to
+    // 'cancel'. Then the python code can actually raise the exception.
+    return JSON.parse(request.responseText) as {
+      type: 'msg' | 'cancel';
+      value?: string;
+    };
   }
 
   // I tried setting jsglobals here, to provide 'input' and 'print' to python,
@@ -60,22 +95,94 @@ async function setupPyodide() {
     print,
     input
   });
-  // TODO: use a fresh global object for each runPython call if we stop terminating
-  // the worker when the user input changes. (See python-test-evaluator.ts)
-  pyodide.runPython(`
+  // Create fresh globals each time user code is run.
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-call
+  const globals = pyodide.globals.get('dict')() as PyProxy;
+  // Some tests rely on __name__ being set to __main__ and we new dicts do not
+  // have this set by default.
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-call
+  globals.set('__name__', '__main__');
+  // The runPython helper is a shortcut for running python code with our
+  // custom globals.
+  const runPython = (pyCode: string) =>
+    pyodide.runPython(pyCode, { globals }) as unknown;
+  runPython(`
   import jscustom
   from jscustom import print
   from jscustom import input
-`);
+  def __wrap(func):
+    def fn(*args):
+      data = func(*args)
+      if data.type == 'cancel':
+        raise KeyboardInterrupt(data.value)
+      return data.value
+    return fn
+  input = __wrap(input)
+  `);
 
-  return pyodide;
+  // Exposing sys.last_value can create memory leaks, so this just returns a
+  // string instead of the actual exception. args[0] is what was passed to the
+  // exception constructor. In our case, that's the id we want.
+  // TODO: I'm using 'join' to make sure we're not leaking a reference to the
+  // exception. This might be excessive, but I don't know enough about pyodide
+  // to be sure.
+  runPython(`
+  import sys
+  def __get_reset_id():
+    if sys.last_value and sys.last_value.args:
+      return "".join(str(sys.last_value.args[0]))
+    else:
+      return ""
+  `);
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-call
+  const getResetId = globals.get('__get_reset_id') as () => string;
+  return { runPython, getResetId };
 }
 
-void setupPyodide();
-
-ctx.onmessage = async (e: PythonRunEvent) => {
-  const code = (e.data.code.contents || '').slice();
-  const pyodide = await setupPyodide();
-  // use pyodide.runPythonAsync if we want top-level await
-  pyodide.runPython(code);
+ctx.onmessage = (e: PythonRunEvent | ListenRequestEvent | CancelEvent) => {
+  const { data } = e;
+  if (data.type === 'listen') {
+    handleListenRequest();
+  } else if (data.type === 'cancel') {
+    handleCancelRequest(data);
+  } else {
+    handleRunRequest(data);
+  }
 };
+
+// This lets the client know that there is nothing to cancel.
+function handleCancelRequest({ value }: { value: number }) {
+  postMessage({ type: 'is-alive', text: value });
+}
+
+function handleListenRequest() {
+  ignoreRunMessages = false;
+}
+
+function handleRunRequest(data: PythonRunEvent['data']) {
+  if (ignoreRunMessages) return;
+  const code = (data.code.contents || '').slice();
+  // TODO: use reset-terminal for clarity?
+  postMessage({ type: 'reset' });
+
+  const { runPython, getResetId } = initRunPython();
+  // use pyodide.runPythonAsync if we want top-level await
+  try {
+    runPython(code);
+  } catch (e) {
+    const err = e as PythonError;
+    console.error(e);
+    const resetId = getResetId();
+    // TODO: if a user raises a KeyboardInterrupt with a custom message this
+    // will be treated as a reset, the client will resend their code and this
+    // will loop. Can we fix that? Perhaps by using a custom exception?
+    if (err.type === 'KeyboardInterrupt' && resetId) {
+      // If the client sends a lot of run messages, it's easy for them to build
+      // up while the worker is busy. As such, we both ignore any queued run
+      // messages...
+      ignoreRunMessages = true;
+      // ...and tell the client that we're ignoring them.
+      postMessage({ type: 'stopped', text: getResetId() });
+    }
+  }
+}