From 77127d34769ac2ca67f8ef0f263fd4686744f979 Mon Sep 17 00:00:00 2001
From: DarkCat09 <gh@dc09.ru>
Date: Fri, 3 May 2024 14:39:29 +0400
Subject: [PATCH] Fix: keep connection

---
 backend/main.py | 50 ++++++++++++++++++++++++-------------------------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/backend/main.py b/backend/main.py
index 2d722e4..01b81a1 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -1,6 +1,5 @@
 import asyncio
 import json
-from typing import Any
 
 import secrets
 import time
@@ -25,31 +24,32 @@ def generate_key() -> str:
 
 
 async def handler(socket: websockets.WebSocketServerProtocol) -> None:
-    data = json.loads(await socket.recv())
-    if 'action' not in data:
-        await socket.send(response.error_field('action'))
-        return
-    match data['action']:
-        case 'init':  # create session
-            if data.get('site') not in SITES:
-                await socket.send(response.error_field('site'))
-                return
-            key = generate_key()
-            sessions[key] = ""  # TODO
-            await socket.send(response.ok_init(key))
-        case 'list':  # list tracks in album
-            if 'url' not in data:
-                await socket.send(response.error_field('url'))
-                return
-            await socket.send(response.ok_playlist(['title 1', 'title 2']))  # TODO
-        case 'download':  # download by URL
-            if 'url' not in data:
-                await socket.send(response.error_field('url'))
-                return
-            # TODO: pass command to the thread started in `init`
-            await socket.send(response.OK)
-        case _:
+    async for message in socket:
+        data = json.loads(await socket.recv())
+        if 'action' not in data:
             await socket.send(response.error_field('action'))
+            return
+        match data['action']:
+            case 'init':  # create session
+                if data.get('site') not in SITES:
+                    await socket.send(response.error_field('site'))
+                    return
+                key = generate_key()
+                sessions[key] = ""  # TODO
+                await socket.send(response.ok_init(key))
+            case 'list':  # list tracks in album
+                if 'url' not in data:
+                    await socket.send(response.error_field('url'))
+                    return
+                await socket.send(response.ok_playlist(['title 1', 'title 2']))  # TODO
+            case 'download':  # download by URL
+                if 'url' not in data:
+                    await socket.send(response.error_field('url'))
+                    return
+                # TODO: start thread and send all yt-dlp's output to client
+                await socket.send(response.OK)
+            case _:
+                await socket.send(response.error_field('action'))
 
 
 async def main() -> None: