fix(socket): set environ when restoring websocket session (#2689)

Hi,

I encountered this behavior when using `cl.context.session.environ` to
get the HTTP headers.
It seemed like the headers were outdated when I used them - and after
going through the `connect` function it seemed like `environ` is not
passed to the `restore_existing_session` function, which means that the
updated socket won't necessarily have the updated environment variables.

If there's an intentional reason for this behavior that I'm missing, I'd
appreciate some context :)
And one more question - I didn't dive deeply into the authentication
flow, but do you think I should also update the user when restoring the
session? for now I left it unchanged.

Thanks for the help!







<!-- This is an auto-generated description by cubic. -->
---
## Summary by cubic
Restored websocket sessions now receive the current request environ so
headers and context stay up to date after reconnects. This prevents
stale header usage in session logic.

- **Bug Fixes**
- Added an environ parameter to restore_existing_session and set
session.environ during restore.
- Passed environ when restoring a session and updated tests to assert
environ propagation and not-found behavior.

<sup>Written for commit 0b840fe.
Summary will update automatically on new commits.</sup>

<!-- End of auto-generated description by cubic. -->

---------

Co-authored-by: Josh Hayes <[email protected]>

Author: daniel-oronsi

backend/chainlit/socket.py

@@ -38,12 +38,13 @@ class WebSocketSessionAuth(TypedDict):
     threadId: str | None
 
 
-def restore_existing_session(sid, session_id, emit_fn, emit_call_fn):
+def restore_existing_session(sid, session_id, emit_fn, emit_call_fn, environ):
     """Restore a session from the sessionId provided by the client."""
     if session := WebsocketSession.get_by_id(session_id):
         session.restore(new_socket_id=sid)
         session.emit = emit_fn
         session.emit_call = emit_call_fn
+        session.environ = environ
         return True
     return False
 
@@ -149,7 +150,7 @@ def emit_call_fn(event: Literal["ask", "call_fn"], data, timeout):
         return sio.call(event, data, timeout=timeout, to=sid)
 
     session_id = auth["sessionId"]
-    if restore_existing_session(sid, session_id, emit_fn, emit_call_fn):
+    if restore_existing_session(sid, session_id, emit_fn, emit_call_fn, environ):
         return True
 
     user_env_string = auth.get("userEnv", None)

backend/tests/test_socket.py

@@ -118,25 +118,29 @@ def test_restore_existing_session_success(self):
         mock_session = Mock(spec=WebsocketSession)
         emit_fn = Mock()
         emit_call_fn = Mock()
+        environ = {"HTTP_COOKIE": "token=token"}
 
         with patch.object(WebsocketSession, "get_by_id") as mock_get:
             mock_get.return_value = mock_session
 
             result = restore_existing_session(
-                "new_sid", "session_123", emit_fn, emit_call_fn
+                "new_sid", "session_123", emit_fn, emit_call_fn, environ
             )
 
             assert result is True
             mock_session.restore.assert_called_once_with(new_socket_id="new_sid")
             assert mock_session.emit == emit_fn
             assert mock_session.emit_call == emit_call_fn
+            assert mock_session.environ == environ
 
     def test_restore_existing_session_not_found(self):
         """Test when session is not found."""
         with patch.object(WebsocketSession, "get_by_id") as mock_get:
             mock_get.return_value = None
 
-            result = restore_existing_session("new_sid", "session_123", Mock(), Mock())
+            result = restore_existing_session(
+                "new_sid", "session_123", Mock(), Mock(), {"HTTP_COOKIE": "token=token"}
+            )
 
             assert result is False
 
@@ -407,7 +411,7 @@ def test_restore_existing_session_with_none_session_id(self):
         with patch.object(WebsocketSession, "get_by_id") as mock_get:
             mock_get.return_value = None
 
-            result = restore_existing_session(None, None, Mock(), Mock())
+            result = restore_existing_session(None, None, Mock(), Mock(), None)
 
             assert result is False
 

cypress/e2e/header_auth/spec.cy.ts

@@ -15,11 +15,6 @@ describe('Header auth', () => {
         req.headers['test-header'] = 'test header value';
         req.reply();
       }).as('auth');
-
-      // Only intercept /user _after_ we're logged in.
-      cy.wait('@auth').then(() => {
-        cy.intercept('GET', '/user').as('user');
-      });
     };
 
     beforeEach(() => {
@@ -55,6 +50,10 @@ describe('Header auth', () => {
     shouldBeLoggedIn();
 
     it('should request and have access to /user', () => {
+      // Only intercept /user _after_ we're logged in.
+      cy.wait('@auth').then(() => {
+        cy.intercept('GET', '/user').as('user');
+      });
       cy.wait('@user').then((interception) => {
         expect(interception.response, 'Intercepted response').to.satisfy(
           () => true