Force close_session on ConnectionClosedOK

Author: blast-hardcheese

src/replit_river/v2/session.py

@@ -20,13 +20,13 @@
 )
 
 import nanoid
-from websockets import ConnectionClosedOK
 import websockets.asyncio.client
 from aiochannel import Channel, ChannelEmpty, ChannelFull
 from aiochannel.errors import ChannelClosed
 from opentelemetry.trace import Span, use_span
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from pydantic import ValidationError
+from websockets import ConnectionClosedOK
 from websockets.asyncio.client import ClientConnection
 from websockets.exceptions import ConnectionClosed
 
@@ -1331,6 +1331,9 @@ async def _recv_from_ws(
                 # is no @overrides in `websockets` to hint this.
                 try:
                     message = await ws.recv(decode=False)
+                except ConnectionClosedOK as e:
+                    close_session(e)
+                    continue
                 except ConnectionClosed:
                     # This triggers a break in the inner loop so we can get back to
                     # the outer loop.