Pass code edits through CRDT.

examples/word2vec.py

@@ -22,3 +22,8 @@ def word2vec_1000():
 @op(ENV, "Take first N")
 def first_n(df: pd.DataFrame, *, n=10):
     return df.head(n)
+
+
+@op(ENV, "Sample N")
+def sample_n(df: pd.DataFrame, *, n=10):
+    return df.sample(n)

lynxkite-app/src/lynxkite_app/crdt.py

@@ -26,11 +26,11 @@ def ws_exception_handler(exception, log):
     return True
 
 
-class WebsocketServer(pycrdt_websocket.WebsocketServer):
+class WorkspaceWebsocketServer(pycrdt_websocket.WebsocketServer):
     async def init_room(self, name: str) -> pycrdt_websocket.YRoom:
         """Initialize a room for the workspace with the given name.
 
-        The workspace is loaded from "crdt_data" if it exists there, or from "data", or a new workspace is created.
+        The workspace is loaded from ".crdt" if it exists there, or from a JSON file, or a new workspace is created.
         """
         crdt_path = pathlib.Path(".crdt")
         path = crdt_path / f"{name}.crdt"
@@ -82,6 +82,37 @@ class WebsocketServer(pycrdt_websocket.WebsocketServer):
         return room
 
 
+class CodeWebsocketServer(WorkspaceWebsocketServer):
+    async def init_room(self, name: str) -> pycrdt_websocket.YRoom:
+        """Initialize a room for a text document with the given name."""
+        crdt_path = pathlib.Path(".crdt")
+        path = crdt_path / f"{name}.crdt"
+        assert path.is_relative_to(crdt_path)
+        ystore = pycrdt_websocket.ystore.FileYStore(path)
+        ydoc = pycrdt.Doc()
+        ydoc["text"] = text = pycrdt.Text()
+        # Replay updates from the store.
+        try:
+            for update, timestamp in [(item[0], item[-1]) async for item in ystore.read()]:
+                ydoc.apply_update(update)
+        except pycrdt_websocket.ystore.YDocNotFound:
+            pass
+        if len(text) == 0:
+            if os.path.exists(name):
+                with open(name) as f:
+                    text += f.read()
+        room = pycrdt_websocket.YRoom(
+            ystore=ystore, ydoc=ydoc, exception_handler=ws_exception_handler
+        )
+        room.text = text
+
+        def on_change(changes):
+            asyncio.create_task(code_changed(name, changes, text))
+
+        text.observe(on_change)
+        return room
+
+
 last_ws_input = None
 
 
@@ -247,14 +278,21 @@ async def execute(name: str, ws_crdt: pycrdt.Map, ws_pyd: workspace.Workspace, d
     print(f"Finished running {name} in {ws_pyd.env}.")
 
 
+async def code_changed(name: str, changes: pycrdt.TextEvent, text: pycrdt.Text):
+    # TODO: Make this more fancy?
+    with open(name, "w") as f:
+        f.write(str(text))
+
+
 @contextlib.asynccontextmanager
 async def lifespan(app):
-    global websocket_server
-    websocket_server = WebsocketServer(
-        auto_clean_rooms=False,
-    )
-    async with websocket_server:
-        yield
+    global ws_websocket_server
+    global code_websocket_server
+    ws_websocket_server = WorkspaceWebsocketServer(auto_clean_rooms=False)
+    code_websocket_server = CodeWebsocketServer(auto_clean_rooms=False)
+    async with ws_websocket_server:
+        async with code_websocket_server:
+            yield
     print("closing websocket server")
 
 
@@ -265,5 +303,12 @@ def sanitize_path(path):
 @router.websocket("/ws/crdt/{room_name}")
 async def crdt_websocket(websocket: fastapi.WebSocket, room_name: str):
     room_name = sanitize_path(room_name)
-    server = pycrdt_websocket.ASGIServer(websocket_server)
+    server = pycrdt_websocket.ASGIServer(ws_websocket_server)
+    await server({"path": room_name}, websocket._receive, websocket._send)
+
+
+@router.websocket("/ws/code/crdt/{room_name}")
+async def code_crdt_websocket(websocket: fastapi.WebSocket, room_name: str):
+    room_name = sanitize_path(room_name)
+    server = pycrdt_websocket.ASGIServer(code_websocket_server)
     await server({"path": room_name}, websocket._receive, websocket._send)

lynxkite-app/src/lynxkite_app/main.py

@@ -39,15 +39,6 @@ def get_catalog(workspace: str):
     return {k: {op.name: op.model_dump() for op in v.values()} for k, v in ops.CATALOGS.items()}
 
 
-@app.get("/api/getCode")
-def get_code(path: str):
-    path = data_path / path
-    assert path.is_relative_to(data_path)
-    with open(path) as f:
-        code = f.read()
-    return {"code": code}
-
-
 class SaveRequest(workspace.BaseConfig):
     path: str
     ws: workspace.Workspace

lynxkite-app/web/package-lock.json

@@ -29,6 +29,7 @@
         "react-router-dom": "^7.0.2",
         "swr": "^2.2.5",
         "unplugin-icons": "^0.21.0",
+        "y-monaco": "^0.1.6",
         "y-websocket": "^2.0.4",
         "yjs": "^13.6.20"
       },
@@ -7397,6 +7398,23 @@
         "yjs": "^13.0.0"
       }
     },
+    "node_modules/y-monaco": {
+      "version": "0.1.6",
+      "resolved": "https://registry.npmjs.org/y-monaco/-/y-monaco-0.1.6.tgz",
+      "integrity": "sha512-sYRywMmcylt+Nupl+11AvizD2am06ST8lkVbUXuaEmrtV6Tf+TD4rsEm6u9YGGowYue+Vfg1IJ97SUP2J+PVXg==",
+      "license": "MIT",
+      "dependencies": {
+        "lib0": "^0.2.43"
+      },
+      "engines": {
+        "node": ">=12.0.0",
+        "npm": ">=6.0.0"
+      },
+      "peerDependencies": {
+        "monaco-editor": ">=0.20.0",
+        "yjs": "^13.3.1"
+      }
+    },
     "node_modules/y-protocols": {
       "version": "1.0.6",
       "resolved": "https://registry.npmjs.org/y-protocols/-/y-protocols-1.0.6.tgz",

lynxkite-app/web/package.json

@@ -32,6 +32,7 @@
     "react-router-dom": "^7.0.2",
     "swr": "^2.2.5",
     "unplugin-icons": "^0.21.0",
+    "y-monaco": "^0.1.6",
     "y-websocket": "^2.0.4",
     "yjs": "^13.6.20"
   },

lynxkite-app/web/src/Code.tsx

@@ -1,10 +1,11 @@
 // Full-page editor for code files.
 
-import Editor from "@monaco-editor/react";
-import { loader } from "@monaco-editor/react";
-import { useEffect } from "react";
+import Editor, { type Monaco } from "@monaco-editor/react";
+import type { editor } from "monaco-editor";
+import { useEffect, useRef } from "react";
 import { useParams } from "react-router";
-import useSWR, { type Fetcher } from "swr";
+import { WebsocketProvider } from "y-websocket";
+import * as Y from "yjs";
 // @ts-ignore
 import Atom from "~icons/tabler/atom.jsx";
 // @ts-ignore
@@ -13,20 +14,41 @@ import Backspace from "~icons/tabler/backspace.jsx";
 import Close from "~icons/tabler/x.jsx";
 import favicon from "./assets/favicon.ico";
 import theme from "./code-theme.ts";
+// For some reason y-monaco is gigantic. The other Monaco packages are small.
+const MonacoBinding = await import("y-monaco").then((m) => m.MonacoBinding);
 
 export default function Code() {
-  useEffect(() => {
-    const initMonaco = async () => {
-      const monaco = await loader.init();
-      monaco.editor.defineTheme("lynxkite", theme);
-    };
-    initMonaco();
-  }, []);
   const { path } = useParams();
   const parentDir = path!.split("/").slice(0, -1).join("/");
-  const fetcher: Fetcher<{ code: string }> = (resource: string, init?: RequestInit) =>
-    fetch(resource, init).then((res) => res.json());
-  const code = useSWR(`/api/getCode?path=${path}`, fetcher);
+  const ydoc = useRef<any>();
+  const wsProvider = useRef<any>();
+  const monacoBinding = useRef<any>();
+  function beforeMount(monaco: Monaco) {
+    monaco.editor.defineTheme("lynxkite", theme);
+  }
+  function onMount(_editor: editor.IStandaloneCodeEditor) {
+    ydoc.current = new Y.Doc();
+    const text = ydoc.current.getText("text");
+    const proto = location.protocol === "https:" ? "wss:" : "ws:";
+    wsProvider.current = new WebsocketProvider(
+      `${proto}//${location.host}/ws/code/crdt`,
+      path!,
+      ydoc.current,
+    );
+    monacoBinding.current = new MonacoBinding(
+      text,
+      _editor.getModel()!,
+      new Set([_editor]),
+      wsProvider.current.awareness,
+    );
+  }
+  useEffect(() => {
+    return () => {
+      ydoc.current?.destroy();
+      wsProvider.current?.destroy();
+      monacoBinding.current?.destroy();
+    };
+  });
   return (
     <div className="workspace">
       <div className="top-bar bg-neutral">
@@ -46,29 +68,19 @@ export default function Code() {
           </a>
         </div>
       </div>
-      {code.isLoading && (
-        <div className="loading">
-          <div className="loading-text">Loading...</div>
-        </div>
-      )}
-      {code.error && (
-        <div className="error">
-          <div className="error-text">Error: {code.error}</div>
-        </div>
-      )}
-      {code.data && (
-        <Editor
-          defaultLanguage="python"
-          defaultValue={code.data.code}
-          theme="lynxkite"
-          options={{
-            cursorStyle: "block",
-            cursorBlinking: "solid",
-            minimap: { enabled: false },
-            renderLineHighlight: "none",
-          }}
-        />
-      )}
+      <Editor
+        defaultLanguage="python"
+        theme="lynxkite"
+        path={path}
+        beforeMount={beforeMount}
+        onMount={onMount}
+        options={{
+          cursorStyle: "block",
+          cursorBlinking: "solid",
+          minimap: { enabled: false },
+          renderLineHighlight: "none",
+        }}
+      />
     </div>
   );
 }