From e44d1ef13d95c58f2f2269cbcdf05809191954aa Mon Sep 17 00:00:00 2001
From: Yuta Hayashibe <yuta@hayashibe.jp>
Date: Sat, 1 Oct 2022 23:55:12 +0900
Subject: [PATCH] Use different context for each connection

---
 whispering/serve.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/whispering/serve.py b/whispering/serve.py
index 198a4c8..20d50fb 100644
--- a/whispering/serve.py
+++ b/whispering/serve.py
@@ -5,6 +5,7 @@ from logging import getLogger
 
 import numpy as np
 import websockets
+from websockets.exceptions import ConnectionClosedOK
 
 from whispering.transcriber import Context, WhisperStreamingTranscriber
 
@@ -15,10 +16,16 @@ async def serve_with_websocket_main(websocket):
     global g_wsp
     global g_ctx
     idx: int = 0
+    ctx: Context = g_ctx.copy(
+        deep=True,
+    )
 
     while True:
         logger.debug(f"Segment #: {idx}")
-        message = await websocket.recv()
+        try:
+            message = await websocket.recv()
+        except ConnectionClosedOK:
+            break
 
         if isinstance(message, str):
             logger.debug(f"Got str: {message}")
@@ -28,7 +35,7 @@ async def serve_with_websocket_main(websocket):
         segment = np.frombuffer(message, dtype=np.float32)
         for chunk in g_wsp.transcribe(
             segment=segment,  # type: ignore
-            ctx=g_ctx,
+            ctx=ctx,
         ):
             await websocket.send(chunk.json())
         idx += 1