include session/sess param in ws handlers

AnswerDotAI · Nov 4, 2024 · d9a4c1a · d9a4c1a
1 parent 3b6c700
commit d9a4c1a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 7 deletions.
diff --git a/nbs/api/00_core.ipynb b/nbs/api/00_core.ipynb
@@ -796,6 +796,7 @@
     "        if arg.lower()=='htmx': return _get_htmx(hdrs)\n",
     "        if arg.lower()=='app': return ws.scope['app']\n",
     "        if arg.lower()=='send': return partial(_send_ws, ws)\n",
+    "        if 'session'.startswith(arg.lower()): return ws.scope.get('session', {})\n",
     "        return None\n",
     "    res = data.get(arg, None)\n",
     "    if res is empty or res is None: res = hdrs.get(arg, None)\n",
@@ -1228,7 +1229,7 @@
     {
      "data": {
       "text/plain": [
-       "'a604e4a2-08e8-462d-aff9-15468891fe09'"
+       "'77486da1-c613-48be-80c4-9cae89eeec48'"
       ]
      },
      "execution_count": null,
@@ -1748,17 +1749,24 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Message text was: Hi!, from client: Address(host='testclient', port=50000)\n"
+      "Message text was: Hi! with session bar, from client: Address(host='testclient', port=50000)\n"
      ]
     }
    ],
    "source": [
+    "@app.get(\"/setsess\")\n",
+    "def set_sess(session):\n",
+    "   session['foo'] = 'bar'\n",
+    "   return 'ok'\n",
+    "\n",
     "@app.ws(\"/ws\")\n",
-    "def ws(self, msg:str, ws:WebSocket): return f\"Message text was: {msg}, from client: {ws.client}\"\n",
+    "def ws(self, msg:str, ws:WebSocket, session): return f\"Message text was: {msg} with session {session.get('foo')}, from client: {ws.client}\"\n",
+    "\n",
+    "cli.get('/setsess')\n",
     "with cli.websocket_connect('/ws') as ws:\n",
     "    ws.send_text('{\"msg\":\"Hi!\"}')\n",
     "    data = ws.receive_text()\n",
-    "assert 'Message text was: Hi!' in data\n",
+    "assert 'Message text was: Hi! with session bar' in data\n",
     "print(data)"
    ]
   },
@@ -2415,13 +2423,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Set to 2024-10-28 20:22:34.772989\n"
+      "Set to 2024-11-04 15:30:23.038930\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "'Session time: 2024-10-28 20:22:34.772989'"
+       "'Session time: 2024-11-04 15:30:23.038930'"
       ]
      },
      "execution_count": null,
@@ -2786,7 +2794,7 @@
     {
      "data": {
       "text/plain": [
-       "'Cookie was set at time 20:22:35.467691'"
+       "'Cookie was set at time 15:30:23.174646'"
       ]
      },
      "execution_count": null,

diff --git a/nbs/explains/websockets.ipynb b/nbs/explains/websockets.ipynb
@@ -93,6 +93,33 @@
     "This is a fairly simple example and could be done just as easily with standard HTTP requests, but it illustrates the basic idea of how websockets work. Let's look at a more complex example next."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Session data in Websockets\n",
+    "\n",
+    "Session data is shared between standard HTTP routes and Websockets. This means you can access, for example, logged in user ID inside websocket handler:\n",
+    "\n",
+    "```python\n",
+    "from fasthtml.common import *\n",
+    "\n",
+    "app = FastHTML(exts='ws')\n",
+    "rt = app.route\n",
+    "\n",
+    "@rt('/login')\n",
+    "def get(session):\n",
+    "    session[\"person\"] = \"Bob\"\n",
+    "    return \"ok\"\n",
+    "\n",
+    "@app.ws('/ws')\n",
+    "async def ws(msg:str, send, session):\n",
+    "    await send(Div(f'Hello {session.get(\"person\")}' + msg, id='notifications'))\n",
+    "\n",
+    "serve()\n",
+    "```"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},