Feature/websocket lifecycle callbacks

.github/actions/uv-python-install/action.yaml

@@ -32,6 +32,6 @@ runs:
       with:
         python-version: ${{ inputs.python-version }}
     - name: Install Python dependencies
-      run: uv sync ${{ inputs.extra-dependencies }}
+      run: uv sync --no-install-project --no-editable ${{ inputs.extra-dependencies }}
       shell: bash
       working-directory: ${{ inputs.working-directory }}

.github/workflows/lint-backend.yaml

@@ -28,5 +28,5 @@ jobs:
           changed-files: "true"
           args: "format --check"
       - name: Run Mypy
-        run: uv run mypy chainlit/
+        run: uv run --no-project mypy chainlit/
         working-directory: ${{ env.BACKEND_DIR }}

.github/workflows/publish.yaml

@@ -35,15 +35,11 @@ jobs:
         with:
           working-directory: ${{ env.BACKEND_DIR }}
 
-      - name: Build frontend and prepare assets
-        run: uv run python build.py
-        working-directory: ${{ env.BACKEND_DIR }}
-
       - name: Build Python distribution
         run: uv build
         working-directory: ${{ env.BACKEND_DIR }}
 
-      - name: List wheel contents
+      - name: Check frontend and copilot folder included
         run: |
           pip install wheel
           python -m wheel unpack dist/chainlit-*.whl -d unpacked

.github/workflows/pytest.yaml

@@ -26,5 +26,5 @@ jobs:
         run: pnpm run buildUi
         timeout-minutes: 5
       - name: Run Pytest
-        run: uv run pytest --cov=chainlit/
+        run: uv run --no-project pytest --cov=chainlit/
         working-directory: ${{ env.BACKEND_DIR }}

backend/build.py

@@ -6,6 +6,7 @@
 import subprocess
 import sys
 
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
 
 class BuildError(Exception):
     """Custom exception for build failures"""
@@ -98,5 +99,6 @@ def build():
         sys.exit(1)
 
 
-if __name__ == "__main__":
-    build()
+class CustomBuildHook(BuildHookInterface):
+    def initialize(self, version, build_data):
+        build()

backend/chainlit/__init__.py

@@ -120,6 +120,32 @@ def acall(self):
         "LlamaIndexCallbackHandler": "chainlit.llama_index.callbacks",
         "instrument_openai": "chainlit.openai",
         "instrument_mistralai": "chainlit.mistralai",
+        "HaystackAgentCallbackHandler": "chainlit.haystack.callbacks",
+        "AsyncHaystackAgentCallbackHandler": "chainlit.haystack.callbacks",
+        "HaystackCallbackHandler": "chainlit.haystack.callbacks",
+        "AsyncHaystackCallbackHandler": "chainlit.haystack.callbacks",
+        "on_chat_start": "chainlit.lifecycle",
+        "on_chat_resume": "chainlit.lifecycle",
+        "on_chat_end": "chainlit.lifecycle",
+        "on_stop": "chainlit.lifecycle",
+        "on_message": "chainlit.lifecycle",
+        "on_audio_chunk": "chainlit.lifecycle",
+        "on_audio_start": "chainlit.lifecycle",
+        "on_audio_end": "chainlit.lifecycle",
+        "on_settings_update": "chainlit.lifecycle",
+        "on_action": "chainlit.lifecycle",
+        "on_feedback": "chainlit.lifecycle",
+        "on_app_start": "chainlit.lifecycle",
+        "on_app_shutdown": "chainlit.lifecycle",
+        "on_connect": "chainlit.socket",
+        "on_disconnect": "chainlit.socket",
+        "on_socket_connect": "chainlit.callbacks",
+        "on_socket_disconnect": "chainlit.callbacks",
+        "action_callback": "chainlit.action",
+        "author_rename": "chainlit.message",
+        "set_chat_profiles": "chainlit.chat_profile",
+        "set_starters": "chainlit.starters",
+        "data_layer": "chainlit.data",
         "SemanticKernelFilter": "chainlit.semantic_kernel",
         "server": "chainlit.server",
     }
@@ -190,6 +216,8 @@ def acall(self):
     "on_mcp_disconnect",
     "on_message",
     "on_settings_update",
+    "on_socket_connect",
+    "on_socket_disconnect",
     "on_stop",
     "on_window_message",
     "password_auth_callback",

backend/chainlit/callbacks.py

@@ -17,6 +17,36 @@
 from chainlit.user import User
 from chainlit.utils import wrap_user_function
 
+__all__ = [
+    "action_callback",
+    "author_rename",
+    "data_layer",
+    "header_auth_callback",
+    "oauth_callback",
+    "on_app_shutdown",
+    "on_app_startup",
+    "on_audio_chunk",
+    "on_audio_end",
+    "on_audio_start",
+    "on_chat_end",
+    "on_chat_resume",
+    "on_chat_start",
+    "on_feedback",
+    "on_logout",
+    "on_mcp_connect",
+    "on_mcp_disconnect",
+    "on_message",
+    "on_settings_update",
+    "on_socket_connect",
+    "on_socket_disconnect",
+    "on_stop",
+    "on_window_message",
+    "password_auth_callback",
+    "send_window_message",
+    "set_chat_profiles",
+    "set_starters",
+]
+
 
 def on_app_startup(func: Callable[[], Union[None, Awaitable[None]]]) -> Callable:
     """
@@ -191,6 +221,37 @@ def on_window_message(func: Callable[[str], Any]) -> Callable:
     return func
 
 
+def on_socket_connect(func: Callable) -> Callable:
+    """
+    Hook to react to the socket connection event.
+    This is called when a new WebSocket connection is established.
+
+    Args:
+        func (Callable[[], Any]): The connection hook to execute.
+
+    Returns:
+        Callable[[], Any]: The decorated hook.
+    """
+    config.code.on_socket_connect = wrap_user_function(func)
+    return func
+
+
+def on_socket_disconnect(func: Callable) -> Callable:
+    """
+    Hook to react to the socket disconnection event.
+    This is called when a WebSocket connection is closed.
+
+    Args:
+        func (Callable[[], Any]): The disconnection hook to execute.
+
+    Returns:
+        Callable[[], Any]: The decorated hook.
+    """
+    config.code.on_socket_disconnect = wrap_user_function(func)
+
+    return func
+
+
 def on_chat_start(func: Callable) -> Callable:
     """
     Hook to react to the user websocket connection event.
@@ -201,7 +262,6 @@ def on_chat_start(func: Callable) -> Callable:
     Returns:
         Callable[], Any]: The decorated hook.
     """
-
     config.code.on_chat_start = wrap_user_function(
         step(func, name="on_chat_start", type="run"), with_task=True
     )

backend/chainlit/config.py

@@ -358,6 +358,8 @@ class CodeSettings(BaseModel):
     on_logout: Optional[Callable[["Request", "Response"], Any]] = None
     on_stop: Optional[Callable[[], Any]] = None
     on_chat_start: Optional[Callable[[], Any]] = None
+    on_socket_connect: Optional[Callable[[], Any]] = None
+    on_socket_disconnect: Optional[Callable[[], Any]] = None
     on_chat_end: Optional[Callable[[], Any]] = None
     on_chat_resume: Optional[Callable[["ThreadDict"], Any]] = None
     on_message: Optional[Callable[["Message"], Any]] = None

backend/chainlit/socket.py

@@ -129,31 +129,51 @@ def emit_call_fn(event: Literal["ask", "call_fn"], data, timeout):
         return sio.call(event, data, timeout=timeout, to=sid)
 
     session_id = auth.get("sessionId")
+
+    # Try to restore existing session first
     if restore_existing_session(sid, session_id, emit_fn, emit_call_fn):
-        return True
+        session = WebsocketSession.get(sid)
+        if not session:
+            logger.error("Failed to restore existing session")
+            return False
 
-    user_env_string = auth.get("userEnv")
-    user_env = load_user_env(user_env_string)
-
-    client_type = auth.get("clientType")
-    url_encoded_chat_profile = auth.get("chatProfile")
-    chat_profile = (
-        unquote(url_encoded_chat_profile) if url_encoded_chat_profile else None
-    )
-
-    WebsocketSession(
-        id=session_id,
-        socket_id=sid,
-        emit=emit_fn,
-        emit_call=emit_call_fn,
-        client_type=client_type,
-        user_env=user_env,
-        user=user,
-        token=token,
-        chat_profile=chat_profile,
-        thread_id=auth.get("threadId"),
-        environ=environ,
-    )
+        # Initialize WebSocket context for restored session
+        init_ws_context(session)
+    else:
+        # Create new session
+        user_env_string = auth.get("userEnv")
+        user_env = load_user_env(user_env_string)
+
+        client_type = auth.get("clientType")
+        url_encoded_chat_profile = auth.get("chatProfile")
+        chat_profile = (
+            unquote(url_encoded_chat_profile) if url_encoded_chat_profile else None
+        )
+
+        session = WebsocketSession(
+            id=session_id,
+            socket_id=sid,
+            emit=emit_fn,
+            emit_call=emit_call_fn,
+            client_type=client_type,
+            user_env=user_env,
+            user=user,
+            token=token,
+            chat_profile=chat_profile,
+            thread_id=auth.get("threadId"),
+            environ=environ,
+        )
+
+        # Initialize WebSocket context with the session object
+        init_ws_context(session)
+
+    # Call on_socket_connect if defined
+    if config.code.on_socket_connect:
+        task = asyncio.create_task(
+            config.code.on_socket_connect(),
+            name="on_socket_connect",
+        )
+        session.current_task = task
 
     return True
 
@@ -207,10 +227,16 @@ async def disconnect(sid):
     if not session:
         return
 
+    # Re-initialize context after error
     init_ws_context(session)
 
-    if config.code.on_chat_end:
-        await config.code.on_chat_end()
+    # Call on_socket_disconnect if defined
+    if config.code.on_socket_disconnect:
+        try:
+            # Call the disconnect callback
+            await config.code.on_socket_disconnect()
+        except Exception as e:
+            logger.error("Error in on_socket_disconnect: %s", str(e))
 
     if session.thread_id and session.has_first_interaction:
         await persist_user_session(session.thread_id, session.to_persistable())
@@ -233,6 +259,13 @@ async def clear_on_timeout(_sid):
 
         asyncio.ensure_future(clear_on_timeout(sid))
 
+        # Call on_chat_end if defined
+        if config.code.on_chat_end:
+            try:
+                await config.code.on_chat_end()
+            except Exception as e:
+                logger.error("Error in on_chat_end: %s", str(e))
+
 
 @sio.on("stop")  # pyright: ignore [reportOptionalCall]
 async def stop(sid):

backend/pyproject.toml

@@ -116,6 +116,9 @@ exclude = [
   "chainlit/copilot/**/**/"
 ]
 
+[tool.hatch.build.hooks.custom]
+path = "build.py"
+
 [tool.hatch.build.targets.wheel]
 packages = ["chainlit"]
 artifacts = [

backend/tests/test_callbacks.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from unittest.mock import AsyncMock, Mock
+from unittest.mock import AsyncMock, Mock, patch
 
 from chainlit import config
 from chainlit.callbacks import password_auth_callback
@@ -61,8 +61,6 @@ async def auth_func(headers: Headers) -> User | None:
 
 
 async def test_oauth_callback(test_config: config.ChainlitConfig):
-    from unittest.mock import patch
-
     from chainlit.callbacks import oauth_callback
     from chainlit.user import User
 
@@ -541,6 +539,56 @@ async def handle_chat_end():
         context.session.emit.assert_called()
 
 
+async def test_on_socket_connect(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import on_socket_connect
+
+    async with mock_chainlit_context:
+        # Setup test data
+        socket_connected = False
+
+        @on_socket_connect
+        async def handle_socket_connect():
+            nonlocal socket_connected
+            socket_connected = True
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_socket_connect is not None
+
+        # Call the registered callback
+        await test_config.code.on_socket_connect()
+
+        # Check that the callback was executed
+        assert socket_connected
+        # Socket connect callbacks don't emit steps, so no emit call expected
+
+
+async def test_on_socket_disconnect(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import on_socket_disconnect
+
+    async with mock_chainlit_context:
+        # Setup test data
+        socket_disconnected = False
+
+        @on_socket_disconnect
+        async def handle_socket_disconnect():
+            nonlocal socket_disconnected
+            socket_disconnected = True
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_socket_disconnect is not None
+
+        # Call the registered callback
+        await test_config.code.on_socket_disconnect()
+
+        # Check that the callback was executed
+        assert socket_disconnected
+        # Socket disconnect callbacks don't emit steps, so no emit call expected
+
+
 async def test_data_layer_config(
     mock_data_layer: AsyncMock,
     test_config: config.ChainlitConfig,