Merge pull request #25 from davidbrochart/httpx_ws

Fix authentication, use httpx-ws instead of websockets

README.md

@@ -34,17 +34,24 @@ pip install --pre jupyterlab
 Then launch it with:
 
 ```console
-jupyter lab --port=8000 --ServerApp.token='' --ServerApp.password='' --ServerApp.disable_check_xsrf=True --no-browser
+jupyter lab --port=8000 --no-browser
+
+# it will print a URL like: http://127.0.0.1:8000/lab?token=972cbd440db4b35581b25f90c0a88e3a1095534e18251ca8
+# you will need the token when launching jpterm, but if you don't want to be bothered with authentication:
+# jupyter lab --port=8000 --no-browser --ServerApp.token='' --ServerApp.password='' --ServerApp.disable_check_xsrf=True
 ```
 
-Then launch jpterm in another terminal and pass it the URL to the Jupyter server:
+Then launch jpterm in another terminal:
 
 ```console
-jpterm --server http://127.0.0.1:8000
+jpterm --server http://127.0.0.1:8000/?token=972cbd440db4b35581b25f90c0a88e3a1095534e18251ca8
+
+# if you launched JupyterLab without authentication:
+# jpterm --server http://127.0.0.1:8000
 ```
 
 If JupyterLab and jpterm are launched with `--collaborative`, you can open a document in
-JupyterLab, by opening your browser at http://127.0.0.1:8000, modify it, and see the changes live
+JupyterLab (go to http://127.0.0.1:8000 in your browser), modify it, and see the changes live
 in jpterm.
 
 ## Development install

jpterm/cli.py

@@ -11,6 +11,7 @@ def jpterm_main(kwargs):
         set_.append("component.enable=[txl_remote_contents,txl_remote_terminals]")
         set_.append(f"component.components.contents.url={server}")
         set_.append(f"component.components.contents.collaborative={collaborative}")
+        set_.append(f"component.components.terminals.url={server}")
     else:
         set_.append("component.disable=[txl_remote_contents,txl_remote_terminals]")
         set_.append("component.enable=[txl_local_contents,txl_local_terminals]")

plugins/remote_contents/pyproject.toml

@@ -25,9 +25,9 @@ classifiers = [
 ]
 dependencies = [
   "txl",
-  "httpx",
+  "httpx>=0.23.1",
+  "httpx-ws>=0.2.4",
   "jupyter_ydoc >=0.3.0a1,<1",
-  "websockets >=10.4,<11",
   "ypy-websocket >=0.3.2,<1",
 ]
 dynamic = ["version"]

plugins/remote_contents/txl_remote_contents/components.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Union
@@ -6,13 +7,39 @@
 import httpx
 import y_py as Y
 from asphalt.core import Component, Context
+from httpx_ws import aconnect_ws
 from jupyter_ydoc import ydocs
 from txl.base import Contents
 from txl.hooks import register_component
-from websockets import connect
 from ypy_websocket import WebsocketProvider
 
 
+class Websocket:
+    def __init__(self, websocket, roomid: str):
+        self.websocket = websocket
+        self.roomid = roomid
+
+    @property
+    def path(self) -> str:
+        return self.roomid
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> bytes:
+        try:
+            message = await self.recv()
+        except:
+            raise StopAsyncIteration()
+        return message
+
+    async def send(self, message: bytes):
+        await self.websocket.send_bytes(message)
+
+    async def recv(self) -> bytes:
+        return await self.websocket.receive_bytes()
+
+
 class Entry:
     """Provide a scandir-like API"""
 
@@ -48,6 +75,8 @@ def __init__(
         self.query_params = query_params
         self.cookies = cookies
         self.collaborative = collaborative
+        i = base_url.find(":")
+        self.ws_url = ("wss" if base_url[i - 1] == "s" else "ws") + base_url[i:]
 
     async def get(
         self, path: str, is_dir: bool = False, on_change: Optional[Callable] = None
@@ -78,19 +107,16 @@ async def get(
                 r = await client.put(
                     f"{self.base_url}/api/yjs/roomid/{path}",
                     json={"format": doc_format, "type": doc_type},
+                    params={**self.query_params},
+                    cookies=self.cookies,
                 )
+            self.cookies.update(r.cookies)
             roomid = r.text
-            ws_url = f"ws{self.base_url[self.base_url.find(':'):]}/api/yjs/{roomid}"
-            ws_cookies = "; ".join([f"{k}={v}" for k, v in self.cookies.items()])
             ydoc = Y.YDoc()
             jupyter_ydoc = ydocs[doc_type](ydoc)
             if on_change:
                 jupyter_ydoc.observe(partial(self.on_change, jupyter_ydoc, on_change))
-            self.websocket = await connect(
-                ws_url, extra_headers=[("Cookie", ws_cookies)]
-            )
-            # AT_EXIT.append(self.websocket.close)
-            WebsocketProvider(ydoc, self.websocket)
+            asyncio.create_task(self._websocket_provider(roomid, ydoc))
             if doc_type == "notebook":
                 return {}
             else:
@@ -102,6 +128,12 @@ async def get(
         else:
             return document
 
+    async def _websocket_provider(self, roomid, ydoc):
+        ws_url = f"{self.ws_url}/api/yjs/{roomid}"
+        async with aconnect_ws(ws_url, cookies=self.cookies) as websocket:
+            WebsocketProvider(ydoc, Websocket(websocket, roomid))
+            await asyncio.Future()
+
     def on_change(self, jupyter_ydoc, on_change: Callable, events) -> None:
         content = jupyter_ydoc.get()
         on_change(content)

plugins/remote_terminals/pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 dependencies = [
   "txl",
   "httpx>=0.23.1",
-  "websockets>=10.4",
+  "httpx-ws>=0.2.4",
 ]
 dynamic = ["version"]
 

plugins/remote_terminals/txl_remote_terminals/components.py

@@ -1,10 +1,9 @@
 import asyncio
-import json
 from typing import Dict, List
 
 import httpx
-import websockets
 from asphalt.core import Component, Context
+from httpx_ws import aconnect_ws
 from textual.widget import Widget
 from textual.widgets._header import HeaderTitle
 from txl.base import TerminalFactory, Terminals, Header, Launcher
@@ -34,45 +33,49 @@ def __init__(
         self.ws_url = ("wss" if base_url[i - 1] == "s" else "ws") + base_url[i:]
         self._recv_queue = asyncio.Queue()
         self._send_queue = asyncio.Queue()
+        self._done = asyncio.Event()
         super().__init__()
 
     async def open(self):
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/api/terminals",
                 json={"cwd": ""},
-                # cookies=self.cookies,
+                params={**self.query_params},
+                cookies=self.cookies,
             )
-            # self.cookies.update(response.cookies)
+            self.cookies.update(response.cookies)
             name = response.json()["name"]
             response = await client.get(
                 f"{self.base_url}/api/terminals",
                 cookies=self.cookies,
             )
         if name in [terminal["name"] for terminal in response.json()]:
+            self.header.query_one(HeaderTitle).text = "Terminal"
             terminal = self.terminal(self._send_queue, self._recv_queue)
             terminal.focus()
             await self.mount(terminal)
             terminal.set_size(self.size)
-            self.websocket = await websockets.connect(
-                f"{self.ws_url}/terminals/websocket/{name}"
-            )
-            asyncio.create_task(self._recv())
-            asyncio.create_task(self._send())
-        self.header.query_one(HeaderTitle).text = "Terminal"
+            async with aconnect_ws(
+                f"{self.ws_url}/terminals/websocket/{name}", cookies=self.cookies
+            ) as self.websocket:
+                asyncio.create_task(self._recv())
+                asyncio.create_task(self._send())
+                await self._done.wait()
 
     async def _send(self):
         while True:
             message = await self._send_queue.get()
-            await self.websocket.send(json.dumps(message))
-
-    async def _recv(self):
-        while True:
             try:
-                message = await self.websocket.recv()
+                await self.websocket.send_json(message)
             except BaseException:
+                self._done.set()
                 break
-            await self._recv_queue.put(json.loads(message))
+
+    async def _recv(self):
+        while True:
+            message = await self.websocket.receive_json()
+            await self._recv_queue.put(message)
 
 
 class RemoteTerminalsComponent(Component):