Refactor worker message handling to use Comlink

app/terminal/worker/jsEval.worker.ts

@@ -1,7 +1,8 @@
 /// <reference lib="webworker" />
 
+import { expose } from "comlink";
 import type { ReplOutput } from "../repl";
-import type { MessageType, WorkerRequest, WorkerResponse } from "./runtime";
+import type { WorkerCapabilities } from "./runtime";
 import inspect from "object-inspect";
 
 function format(...args: unknown[]): string {
@@ -29,12 +30,12 @@ self.console = {
   },
 };
 
-async function init({ id }: WorkerRequest["init"]) {
+async function init(/*_interruptBuffer?: Uint8Array*/): Promise<{
+  capabilities: WorkerCapabilities;
+}> {
   // Initialize the worker and report capabilities
-  self.postMessage({
-    id,
-    payload: { capabilities: { interrupt: "restart" } },
-  } satisfies WorkerResponse["init"]);
+  // interruptBuffer is not used for JavaScript (restart-based interruption)
+  return { capabilities: { interrupt: "restart" } };
 }
 
 async function replLikeEval(code: string): Promise<unknown> {
@@ -72,8 +73,10 @@ async function replLikeEval(code: string): Promise<unknown> {
   }
 }
 
-async function runCode({ id, payload }: WorkerRequest["runCode"]) {
-  const { code } = payload;
+async function runCode(code: string): Promise<{
+  output: ReplOutput[];
+  updatedFiles: Record<string, string>;
+}> {
   try {
     const result = await replLikeEval(code);
     jsOutput.push({
@@ -99,14 +102,13 @@ async function runCode({ id, payload }: WorkerRequest["runCode"]) {
   const output = [...jsOutput];
   jsOutput = []; // Clear output
 
-  self.postMessage({
-    id,
-    payload: { output, updatedFiles: [] },
-  } satisfies WorkerResponse["runCode"]);
+  return { output, updatedFiles: {} as Record<string, string> };
 }
 
-function runFile({ id, payload }: WorkerRequest["runFile"]) {
-  const { name, files } = payload;
+function runFile(
+  name: string,
+  files: Record<string, string>
+): { output: ReplOutput[]; updatedFiles: Record<string, string> } {
   // pyodide worker などと異なり、複数ファイルを読み込んでimportのようなことをするのには対応していません。
   try {
     self.eval(files[name]);
@@ -129,23 +131,17 @@ function runFile({ id, payload }: WorkerRequest["runFile"]) {
   const output = [...jsOutput];
   jsOutput = []; // Clear output
 
-  self.postMessage({
-    id,
-    payload: { output, updatedFiles: [] },
-  } satisfies WorkerResponse["runFile"]);
+  return { output, updatedFiles: {} as Record<string, string> };
 }
 
-async function checkSyntax({ id, payload }: WorkerRequest["checkSyntax"]) {
-  const { code } = payload;
-
+async function checkSyntax(
+  code: string
+): Promise<{ status: "complete" | "incomplete" | "invalid" }> {
   try {
     // Try to create a Function to check syntax
     // new Function(code); // <- not working
     self.eval(`() => {${code}}`);
-    self.postMessage({
-      id,
-      payload: { status: "complete" },
-    } satisfies WorkerResponse["checkSyntax"]);
+    return { status: "complete" };
   } catch (e) {
     // Check if it's a syntax error or if more input is expected
     if (e instanceof SyntaxError) {
@@ -154,28 +150,18 @@ async function checkSyntax({ id, payload }: WorkerRequest["checkSyntax"]) {
         e.message.includes("Unexpected token '}'") ||
         e.message.includes("Unexpected end of input")
       ) {
-        self.postMessage({
-          id,
-          payload: { status: "incomplete" },
-        } satisfies WorkerResponse["checkSyntax"]);
+        return { status: "incomplete" };
       } else {
-        self.postMessage({
-          id,
-          payload: { status: "invalid" },
-        } satisfies WorkerResponse["checkSyntax"]);
+        return { status: "invalid" };
       }
     } else {
-      self.postMessage({
-        id,
-        payload: { status: "invalid" },
-      } satisfies WorkerResponse["checkSyntax"]);
+      return { status: "invalid" };
     }
   }
 }
 
-async function restoreState({ id, payload }: WorkerRequest["restoreState"]) {
+async function restoreState(commands: string[]): Promise<object> {
   // Re-execute all previously successful commands to restore state
-  const { commands } = payload;
   jsOutput = []; // Clear output for restoration
 
   for (const command of commands) {
@@ -188,32 +174,15 @@ async function restoreState({ id, payload }: WorkerRequest["restoreState"]) {
   }
 
   jsOutput = []; // Clear any output from restoration
-  self.postMessage({
-    id,
-    payload: {},
-  } satisfies WorkerResponse["restoreState"]);
+  return {};
 }
 
-self.onmessage = async (event: MessageEvent<WorkerRequest[MessageType]>) => {
-  switch (event.data.type) {
-    case "init":
-      await init(event.data);
-      return;
-    case "runCode":
-      await runCode(event.data);
-      return;
-    case "runFile":
-      runFile(event.data);
-      return;
-    case "checkSyntax":
-      await checkSyntax(event.data);
-      return;
-    case "restoreState":
-      await restoreState(event.data);
-      return;
-    default:
-      event.data satisfies never;
-      originalConsole.error(`Unknown message: ${event.data}`);
-      return;
-  }
+const api = {
+  init,
+  runCode,
+  runFile,
+  checkSyntax,
+  restoreState,
 };
+
+expose(api);

app/terminal/worker/pyodide.worker.ts

@@ -1,11 +1,12 @@
 /// <reference lib="webworker" />
 /// <reference lib="ES2023" />
 
+import { expose } from "comlink";
 import type { PyodideInterface } from "pyodide";
 // import { loadPyodide } from "pyodide"; -> Reading from "node:child_process" is not handled by plugins
 import { version as pyodideVersion } from "pyodide/package.json";
 import type { PyCallable } from "pyodide/ffi";
-import type { MessageType, WorkerRequest, WorkerResponse } from "./runtime";
+import type { WorkerCapabilities } from "./runtime";
 import type { ReplOutput } from "../repl";
 
 import execfile_py from "./pyodide/execfile.py?raw";
@@ -36,8 +37,9 @@ function readAllFiles(): Record<string, string> {
   return updatedFiles;
 }
 
-async function init({ id, payload }: WorkerRequest["init"]) {
-  const { interruptBuffer } = payload;
+async function init(
+  interruptBuffer: Uint8Array
+): Promise<{ capabilities: WorkerCapabilities }> {
   if (!pyodide) {
     self.importScripts(`${PYODIDE_CDN}pyodide.js`);
 
@@ -57,20 +59,15 @@ async function init({ id, payload }: WorkerRequest["init"]) {
 
     pyodide.setInterruptBuffer(interruptBuffer);
   }
-  self.postMessage({
-    id,
-    payload: { capabilities: { interrupt: "buffer" } },
-  } satisfies WorkerResponse["init"]);
+  return { capabilities: { interrupt: "buffer" } };
 }
 
-async function runCode({ id, payload }: WorkerRequest["runCode"]) {
-  const { code } = payload;
+async function runCode(code: string): Promise<{
+  output: ReplOutput[];
+  updatedFiles: Record<string, string>;
+}> {
   if (!pyodide) {
-    self.postMessage({
-      id,
-      error: "Pyodide not initialized",
-    } satisfies WorkerResponse["runCode"]);
-    return;
+    throw new Error("Pyodide not initialized");
   }
   try {
     const result = await pyodide.runPythonAsync(code);
@@ -115,20 +112,15 @@ async function runCode({ id, payload }: WorkerRequest["runCode"]) {
   const output = [...pyodideOutput];
   pyodideOutput = []; // 出力をクリア
 
-  self.postMessage({
-    id,
-    payload: { output, updatedFiles },
-  } satisfies WorkerResponse["runCode"]);
+  return { output, updatedFiles };
 }
 
-async function runFile({ id, payload }: WorkerRequest["runFile"]) {
-  const { name, files } = payload;
+async function runFile(
+  name: string,
+  files: Record<string, string>
+): Promise<{ output: ReplOutput[]; updatedFiles: Record<string, string> }> {
   if (!pyodide) {
-    self.postMessage({
-      id,
-      error: "Pyodide not initialized",
-    } satisfies WorkerResponse["runFile"]);
-    return;
+    throw new Error("Pyodide not initialized");
   }
   try {
     // Use Pyodide FS API to write files to the file system
@@ -175,70 +167,41 @@ async function runFile({ id, payload }: WorkerRequest["runFile"]) {
   const updatedFiles = readAllFiles();
   const output = [...pyodideOutput];
   pyodideOutput = []; // 出力をクリア
-  self.postMessage({
-    id,
-    payload: { output, updatedFiles },
-  } satisfies WorkerResponse["runFile"]);
+  return { output, updatedFiles };
 }
 
-async function checkSyntax({ id, payload }: WorkerRequest["checkSyntax"]) {
-  const { code } = payload;
+async function checkSyntax(
+  code: string
+): Promise<{ status: "complete" | "incomplete" | "invalid" }> {
   if (!pyodide) {
-    self.postMessage({
-      id,
-      payload: { status: "invalid" },
-    } satisfies WorkerResponse["checkSyntax"]);
-    return;
+    return { status: "invalid" };
   }
 
   // 複数行コマンドは最後に空行を入れないと完了しないものとする
   if (code.includes("\n") && code.split("\n").at(-1) !== "") {
-    self.postMessage({
-      id,
-      payload: { status: "incomplete" },
-    } satisfies WorkerResponse["checkSyntax"]);
-    return;
+    return { status: "incomplete" };
   }
 
   try {
     // Pythonのコードを実行して結果を受け取る
     const status = (pyodide.runPython(check_syntax_py) as PyCallable)(code);
-    self.postMessage({
-      id,
-      payload: { status },
-    } satisfies WorkerResponse["checkSyntax"]);
+    return { status };
   } catch (e) {
     console.error("Syntax check error:", e);
-    self.postMessage({
-      id,
-      payload: { status: "invalid" },
-    } satisfies WorkerResponse["checkSyntax"]);
+    return { status: "invalid" };
   }
 }
 
-self.onmessage = async (event: MessageEvent<WorkerRequest[MessageType]>) => {
-  switch (event.data.type) {
-    case "init":
-      await init(event.data);
-      return;
-    case "runCode":
-      await runCode(event.data);
-      return;
-    case "runFile":
-      await runFile(event.data);
-      return;
-    case "checkSyntax":
-      await checkSyntax(event.data);
-      return;
-    case "restoreState":
-      self.postMessage({
-        id: event.data.id,
-        error: "not implemented",
-      } satisfies WorkerResponse["restoreState"]);
-      return;
-    default:
-      event.data satisfies never;
-      console.error(`Unknown message: ${event.data}`);
-      return;
-  }
+async function restoreState(): Promise<object> {
+  throw new Error("not implemented");
+}
+
+const api = {
+  init,
+  runCode,
+  runFile,
+  checkSyntax,
+  restoreState,
 };
+
+expose(api);