chore: import upstream snapshot with attribution

backend/tests/__init__.py

@@ -0,0 +1 @@
+

backend/tests/auth/test_cookie.py

@@ -0,0 +1,165 @@
+import importlib
+
+import pytest
+from fastapi import FastAPI, Form
+from fastapi.testclient import TestClient
+from starlette.requests import Request
+from starlette.responses import Response
+
+import chainlit.auth.cookie as cookie_module
+from chainlit.auth import (
+    clear_auth_cookie,
+    get_token_from_cookies,
+    set_auth_cookie,
+)
+
+
+@pytest.fixture
+def test_app():
+    app = FastAPI()
+
+    @app.post("/set-cookie")
+    async def set_cookie_endpoint(request: Request, token: str = Form()):
+        response = Response()
+        set_auth_cookie(request, response, token)
+        return response
+
+    @app.get("/get-token")
+    async def get_token_endpoint(request: Request):
+        token = get_token_from_cookies(request.cookies)
+        return {"token": token}
+
+    @app.delete("/clear-cookie")
+    async def clear_cookie_endpoint(request: Request):
+        response = Response()
+        clear_auth_cookie(request, response)
+        return response
+
+    return app
+
+
+@pytest.fixture
+def client(test_app):
+    return TestClient(test_app)
+
+
+def test_short_token(client):
+    """Test with a <3000 shorter token."""
+
+    # Set a short token
+    short_token = "x" * 1000
+    set_response = client.post("/set-cookie", data={"token": short_token})
+    assert set_response.status_code == 200
+
+    # Verify cookies were set
+    cookies = set_response.cookies
+    assert cookies, "No cookies set"
+    assert "access_token" in cookies, f"No chunking for short cookies: {cookies}"
+
+    # Read back the token using client's cookie jar
+    get_response = client.get("/get-token")
+    assert get_response.status_code == 200
+    assert get_response.json()["token"] == short_token
+
+
+def test_set_and_read_4kb_token(client):
+    """Test full cookie lifecycle using actual client cookie handling."""
+
+    # Set a 4KB token
+    token_4kb = "x" * 4000
+    set_response = client.post("/set-cookie", data={"token": token_4kb})
+    assert set_response.status_code == 200
+
+    # Verify cookies were set
+    cookies = set_response.cookies
+    assert f"{cookies.keys()} should contain chunked cookies", any(
+        key.startswith("access_token_") for key in cookies.keys()
+    )
+
+    # Read back the token using client's cookie jar
+    get_response = client.get("/get-token")
+    assert get_response.status_code == 200
+
+    response_token = get_response.json()["token"]
+    assert len(response_token) == len(token_4kb)
+    assert response_token == token_4kb
+
+
+def test_overwrite_shorter_token_chunked(client):
+    """Test cookie chunk cleanup when replacing a large token with a smaller one."""
+    # Set initial long token
+    long_token = "LONG" * 2000  # 8000 characters
+    client.post("/set-cookie", data={"token": long_token})
+
+    # Verify initial chunks exist
+    first_cookies = client.cookies
+    assert len([k for k in first_cookies if k.startswith("access_token_")]) > 1
+
+    # Set shorter token (should clear previous chunks)
+    short_token = "SHORT" * 1000  # 4000 characters
+    client.post("/set-cookie", data={"token": short_token})
+
+    # Verify new cookie state
+    final_response = client.get("/get-token")
+    assert final_response.json()["token"] == short_token
+
+    # Verify only two chunks remain
+    final_cookies = client.cookies
+    chunk_cookies = [k for k in final_cookies if k.startswith("access_token_")]
+    assert len(chunk_cookies) == 2, f"Found {len(chunk_cookies)} residual cookies"
+
+
+def test_overwrite_shorter_token_unchunked(client):
+    """Test cookie chunk cleanup when replacing a large token with a smaller one."""
+    # Set initial long token
+    long_token = "LONG" * 1000  # 4000 characters
+    client.post("/set-cookie", data={"token": long_token})
+
+    # Verify initial chunks exist
+    first_cookies = client.cookies
+    assert len([k for k in first_cookies if k.startswith("access_token_")]) > 1
+
+    # Set shorter token (should clear previous chunks)
+    short_token = "SHORT"
+    client.post("/set-cookie", data={"token": short_token})
+
+    # Verify new cookie state
+    final_response = client.get("/get-token")
+    assert final_response.json()["token"] == short_token
+
+    # Verify no chunks remain
+    final_cookies = client.cookies
+    chunk_cookies = [k for k in final_cookies if k.startswith("access_token_")]
+    assert len(chunk_cookies) == 0, f"Found {len(chunk_cookies)} residual cookies"
+
+
+def test_state_cookie_lifetime_default(monkeypatch):
+    """Test that _state_cookie_lifetime defaults to 180 seconds (3 minutes)."""
+    monkeypatch.delenv("CHAINLIT_STATE_COOKIE_LIFETIME", raising=False)
+    importlib.reload(cookie_module)
+    assert cookie_module._state_cookie_lifetime == 180
+
+
+def test_state_cookie_lifetime_custom(monkeypatch):
+    """Test that _state_cookie_lifetime can be set via environment variable."""
+    monkeypatch.setenv("CHAINLIT_STATE_COOKIE_LIFETIME", "600")
+    importlib.reload(cookie_module)
+    assert cookie_module._state_cookie_lifetime == 600
+
+
+def test_clear_auth_cookie(client):
+    """Test cookie clearing removes all chunks."""
+    # Set initial token
+    client.post("/set-cookie", data={"token": "x" * 4000})
+
+    # Verify cookies exist
+    assert len(client.cookies) > 0
+
+    # Clear cookies
+    clear_response = client.delete("/clear-cookie")
+    assert clear_response.status_code == 200
+
+    # Verify cookies were cleared
+    assert len(clear_response.cookies) == 0
+    final_response = client.get("/get-token")
+    assert final_response.json()["token"] is None

backend/tests/conftest.py

@@ -0,0 +1,129 @@
+import datetime
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Callable
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+import pytest_asyncio
+
+from chainlit import config
+from chainlit.callbacks import data_layer
+from chainlit.context import ChainlitContext, context_var
+from chainlit.data.base import BaseDataLayer
+from chainlit.session import HTTPSession, WebsocketSession
+from chainlit.user import PersistedUser
+from chainlit.user_session import UserSession
+
+
+@pytest.fixture
+def persisted_test_user():
+    return PersistedUser(
+        id="test_user_id",
+        createdAt=datetime.datetime.now().isoformat(),
+        identifier="test_user_identifier",
+    )
+
+
+@pytest.fixture
+def mock_session_factory(persisted_test_user: PersistedUser) -> Callable[..., Mock]:
+    def create_mock_session(**kwargs) -> Mock:
+        mock = Mock(spec=WebsocketSession)
+        mock.user = kwargs.get("user", persisted_test_user)
+        mock.id = kwargs.get("id", "test_session_id")
+        mock.user_env = kwargs.get("user_env", {"test_env": "value"})
+        mock.chat_settings = kwargs.get("chat_settings", {})
+        mock.chat_profile = kwargs.get("chat_profile", None)
+        mock.environ = kwargs.get("environ", None)
+        mock.client_type = kwargs.get("client_type", "webapp")
+        mock.thread_id = kwargs.get("thread_id", "test_thread_id")
+        mock.emit = AsyncMock()
+        mock.has_first_interaction = kwargs.get("has_first_interaction", True)
+        mock.files = kwargs.get("files", {})
+        mock.files_spec = kwargs.get("files_spec", {})
+
+        return mock
+
+    return create_mock_session
+
+
+@pytest.fixture
+def mock_session(mock_session_factory) -> Mock:
+    return mock_session_factory()
+
+
+@asynccontextmanager
+async def create_chainlit_context(mock_session):
+    from chainlit.emitter import BaseChainlitEmitter
+
+    # Create a mock emitter with all necessary methods
+    mock_emitter = Mock(spec=BaseChainlitEmitter)
+    mock_emitter.send_step = AsyncMock()
+    mock_emitter.update_step = AsyncMock()
+    mock_emitter.delete_step = AsyncMock()
+    mock_emitter.stream_start = AsyncMock()
+    mock_emitter.send_element = AsyncMock()
+    mock_emitter.send_action = AsyncMock()
+    mock_emitter.remove_action = AsyncMock()
+    mock_emitter.emit = AsyncMock()
+    mock_emitter.set_chat_settings = Mock()  # Sync method, not async
+
+    context = ChainlitContext(mock_session, emitter=mock_emitter)
+    token = context_var.set(context)
+    try:
+        yield context
+    finally:
+        context_var.reset(token)
+
+
+@pytest_asyncio.fixture
+async def mock_chainlit_context(persisted_test_user, mock_session):
+    mock_session.user = persisted_test_user
+    return create_chainlit_context(mock_session)
+
+
+@pytest.fixture
+def user_session():
+    return UserSession()
+
+
+@pytest.fixture
+def mock_websocket_session():
+    session = Mock(spec=WebsocketSession)
+    session.emit = AsyncMock()
+
+    return session
+
+
+@pytest.fixture
+def mock_http_session():
+    return Mock(spec=HTTPSession)
+
+
+@pytest.fixture
+def mock_data_layer(monkeypatch: pytest.MonkeyPatch) -> AsyncMock:
+    mock_data_layer = AsyncMock(spec=BaseDataLayer)
+
+    return mock_data_layer
+
+
+@pytest.fixture
+def mock_get_data_layer(mock_data_layer: AsyncMock, test_config: config.ChainlitConfig):
+    # Instantiate mock data layer
+    mock_get_data_layer = Mock(return_value=mock_data_layer)
+
+    # Configure it using @data_layer decorator
+    return data_layer(mock_get_data_layer)
+
+
+@pytest.fixture
+def test_config(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
+    monkeypatch.setenv("CHAINLIT_ROOT_PATH", str(tmp_path))
+
+    test_config = config.load_config()
+
+    monkeypatch.setattr("chainlit.callbacks.config", test_config)
+    monkeypatch.setattr("chainlit.server.config", test_config)
+    monkeypatch.setattr("chainlit.config.config", test_config)
+
+    return test_config

backend/tests/data/conftest.py

@@ -0,0 +1,21 @@
+from unittest.mock import AsyncMock
+
+import pytest
+
+from chainlit.data.storage_clients.base import BaseStorageClient
+from chainlit.user import User
+
+
+@pytest.fixture
+def mock_storage_client():
+    mock_client = AsyncMock(spec=BaseStorageClient)
+    mock_client.upload_file.return_value = {
+        "url": "https://example.com/test.txt",
+        "object_key": "test_user/test_element/test.txt",
+    }
+    return mock_client
+
+
+@pytest.fixture
+def test_user() -> User:
+    return User(identifier="test_user_identifier", metadata={})

backend/tests/data/storage_clients/test_gcs.py

@@ -0,0 +1,337 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from chainlit.data.storage_clients.base import storage_expiry_time
+from chainlit.data.storage_clients.gcs import GCSStorageClient
+
+
+@pytest.fixture
+def mock_gcs_client():
+    """Create a mock Google Cloud Storage client."""
+    # First mock the service_account
+    with patch(
+        "chainlit.data.storage_clients.gcs.service_account"
+    ) as mock_service_account:
+        mock_credentials = MagicMock()
+        mock_service_account.Credentials.from_service_account_info.return_value = (
+            mock_credentials
+        )
+
+        # Then mock the storage client
+        with patch("chainlit.data.storage_clients.gcs.storage") as mock_storage:
+            mock_client = MagicMock()
+            mock_storage.Client.return_value = mock_client
+            mock_bucket = MagicMock()
+            mock_client.bucket.return_value = mock_bucket
+            mock_blob = MagicMock()
+            mock_bucket.blob.return_value = mock_blob
+
+            yield {
+                "storage": mock_storage,
+                "client": mock_client,
+                "bucket": mock_bucket,
+                "blob": mock_blob,
+                "service_account": mock_service_account,
+                "credentials": mock_credentials,
+            }
+
+
+class TestGCSStorageClient:
+    def test_init(self, mock_gcs_client):
+        """Test initialization of GCS client."""
+        # Remove client assignment, directly call the constructor
+        GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Verify service account credentials were created correctly
+        mock_gcs_client[
+            "service_account"
+        ].Credentials.from_service_account_info.assert_called_once_with(
+            {
+                "type": "service_account",
+                "project_id": "test-project",
+                "private_key": "test-key",
+                "client_email": "test@example.com",
+                "token_uri": "https://oauth2.googleapis.com/token",
+            }
+        )
+
+        # Verify storage client was initialized with credentials
+        mock_gcs_client["storage"].Client.assert_called_once_with(
+            project="test-project", credentials=mock_gcs_client["credentials"]
+        )
+
+        # Verify bucket was retrieved
+        mock_gcs_client["client"].bucket.assert_called_once_with("test-bucket")
+
+    def test_sync_get_read_url(self, mock_gcs_client):
+        """Test generating a signed URL."""
+        # Set up the mock to return a URL
+        mock_gcs_client[
+            "blob"
+        ].generate_signed_url.return_value = "https://signed-url.example.com"
+
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset mocks to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test the method
+        url = client.sync_get_read_url("test/path/file.txt")
+
+        # Verify correct methods were called
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].generate_signed_url.assert_called_once_with(
+            version="v4", expiration=storage_expiry_time, method="GET"
+        )
+
+        assert url == "https://signed-url.example.com"
+
+    @pytest.mark.asyncio
+    async def test_get_read_url(self, mock_gcs_client):
+        """Test the async wrapper for getting a read URL."""
+        # Set up the mock to return a URL
+        mock_gcs_client[
+            "blob"
+        ].generate_signed_url.return_value = "https://signed-url.example.com"
+
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset mocks to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test the async method
+        url = await client.get_read_url("test/path/file.txt")
+
+        # Verify correct methods were called
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].generate_signed_url.assert_called_once_with(
+            version="v4", expiration=storage_expiry_time, method="GET"
+        )
+
+        assert url == "https://signed-url.example.com"
+
+    def test_sync_upload_file(self, mock_gcs_client):
+        """Test uploading a file to GCS."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Mock the bucket name property
+        mock_gcs_client["bucket"].name = "test-bucket"
+
+        # Test with binary data
+        binary_data = b"test content"
+        object_key = "test/path/file.txt"
+
+        # Remove the unused result assignment
+        client.sync_upload_file(
+            object_key=object_key, data=binary_data, mime="text/plain", overwrite=True
+        )
+
+        # Check that the blob was called with the correct object key (using assert_any_call instead of assert_called_once_with)
+        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
+        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
+            binary_data, content_type="text/plain"
+        )
+
+    def test_sync_upload_file_string_data(self, mock_gcs_client):
+        """Test uploading string data to GCS."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+        mock_gcs_client["bucket"].name = "test-bucket"
+
+        # Test with string data that should be encoded
+        string_data = "test content"
+        object_key = "test/path/file.txt"
+
+        # Remove the unused result assignment
+        client.sync_upload_file(
+            object_key=object_key, data=string_data, mime="text/plain", overwrite=True
+        )
+
+        # Check that the correct methods were called
+        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
+        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
+            b"test content", content_type="text/plain"
+        )
+
+    def test_sync_upload_file_no_overwrite(self, mock_gcs_client):
+        """Test upload with overwrite=False and existing file."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Configure blob to indicate file exists
+        mock_gcs_client["blob"].exists.return_value = True
+
+        with pytest.raises(
+            Exception,
+            match=r"Failed to upload file to GCS: File test/path/existing\.txt already exists and overwrite is False",
+        ):
+            client.sync_upload_file(
+                object_key="test/path/existing.txt",
+                data=b"test content",
+                overwrite=False,
+            )
+
+        mock_gcs_client["blob"].exists.assert_called_once()
+        mock_gcs_client["blob"].upload_from_string.assert_not_called()
+
+    def test_sync_upload_file_error(self, mock_gcs_client):
+        """Test error handling during upload."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Configure upload to throw an exception
+        mock_gcs_client["blob"].upload_from_string.side_effect = ValueError(
+            "Upload failed"
+        )
+
+        with pytest.raises(
+            Exception, match="Failed to upload file to GCS: Upload failed"
+        ):
+            client.sync_upload_file(
+                object_key="test/path/file.txt", data=b"test content"
+            )
+
+    @pytest.mark.asyncio
+    async def test_upload_file(self, mock_gcs_client):
+        """Test the async wrapper for uploading a file."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+        mock_gcs_client["bucket"].name = "test-bucket"
+
+        # Test with binary data
+        binary_data = b"test content"
+        object_key = "test/path/file.txt"
+
+        # Remove the unused result assignment
+        await client.upload_file(
+            object_key=object_key, data=binary_data, mime="text/plain", overwrite=True
+        )
+
+        # Check that the correct methods were called
+        mock_gcs_client["bucket"].blob.assert_any_call(object_key)
+        mock_gcs_client["blob"].upload_from_string.assert_called_once_with(
+            binary_data, content_type="text/plain"
+        )
+
+    def test_sync_delete_file(self, mock_gcs_client):
+        """Test deleting a file from GCS."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test successful delete
+        result = client.sync_delete_file("test/path/file.txt")
+
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].delete.assert_called_once()
+        assert result is True
+
+    def test_sync_delete_file_error(self, mock_gcs_client):
+        """Test error handling during file deletion."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Configure delete to throw an exception
+        mock_gcs_client["blob"].delete.side_effect = ValueError("Delete failed")
+
+        # The method should catch the exception and return False
+        result = client.sync_delete_file("test/path/file.txt")
+
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].delete.assert_called_once()
+        assert result is False
+
+    @pytest.mark.asyncio
+    async def test_delete_file(self, mock_gcs_client):
+        """Test the async wrapper for deleting a file."""
+        client = GCSStorageClient(
+            bucket_name="test-bucket",
+            project_id="test-project",
+            client_email="test@example.com",
+            private_key="test-key",
+        )
+
+        # Reset the mock to ensure clean state
+        mock_gcs_client["bucket"].reset_mock()
+        mock_gcs_client["blob"].reset_mock()
+
+        # Test successful delete
+        result = await client.delete_file("test/path/file.txt")
+
+        mock_gcs_client["bucket"].blob.assert_called_once_with("test/path/file.txt")
+        mock_gcs_client["blob"].delete.assert_called_once()
+        assert result is True

backend/tests/data/storage_clients/test_s3.py

@@ -0,0 +1,47 @@
+import os
+
+import boto3  # type: ignore
+import pytest
+from moto import mock_aws
+
+from chainlit.data.storage_clients.s3 import S3StorageClient
+
+
+# Fixtures for setting up the DynamoDB table
+@pytest.fixture
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+
+
+@pytest.fixture
+def s3_mock(aws_credentials):
+    """Moto mock S3 setup."""
+    with mock_aws():
+        s3 = boto3.client("s3", region_name="us-east-1")
+        # Create a mock bucket
+        s3.create_bucket(Bucket="my-test-bucket")
+        yield s3
+
+
+@pytest.mark.asyncio
+async def test_upload_file(s3_mock):
+    # Initialize the S3StorageClient with the mock bucket
+    client = S3StorageClient(bucket="my-test-bucket")
+
+    # Call the upload_file method and await the result
+    result = await client.upload_file(
+        object_key="test.txt", data="This is a test file", mime="text/plain"
+    )
+
+    # Assert that the file upload returned the correct URL
+    assert result["object_key"] == "test.txt"
+    assert result["url"] == "https://my-test-bucket.s3.amazonaws.com/test.txt"
+
+    # Verify that the file exists in the mock S3
+    response = s3_mock.get_object(Bucket="my-test-bucket", Key="test.txt")
+    assert response["Body"].read().decode() == "This is a test file"

backend/tests/data/test_chainlit_data_layer.py

@@ -0,0 +1,224 @@
+import json
+from unittest.mock import AsyncMock
+
+import pytest
+
+from chainlit.data.chainlit_data_layer import ChainlitDataLayer
+
+
+@pytest.mark.asyncio
+async def test_update_thread_preserves_metadata_when_none_existing_thread():
+    """Test that update_thread preserves existing metadata when metadata=None on an existing thread."""
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    existing_metadata = {"important": "data", "keep": "me"}
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            return [{"metadata": json.dumps(existing_metadata)}]
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    await data_layer.update_thread(thread_id="test-thread-123", name="Updated Name")
+
+    assert data_layer.execute_query.call_count == 2
+
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_params = update_call[0][1]
+
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    assert metadata_json == existing_metadata
+
+
+@pytest.mark.asyncio
+async def test_update_thread_noop_skips_upsert_on_existing_thread():
+    """Test that update_thread with no arguments early-returns for an existing thread."""
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            return [{"metadata": json.dumps({"existing": "data"})}]
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    await data_layer.update_thread(thread_id="test-thread-123")
+
+    assert data_layer.execute_query.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_update_thread_new_thread_includes_empty_metadata():
+    """Test that update_thread includes metadata='{}' for a new thread (NOT NULL safe)."""
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            return []
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    await data_layer.update_thread(thread_id="new-thread", name="New Chat")
+
+    assert data_layer.execute_query.call_count == 2
+
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_query = update_call[0][0]
+    update_params = update_call[0][1]
+
+    assert "metadata" in update_query.lower()
+
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    assert metadata_json == {}
+
+
+@pytest.mark.asyncio
+async def test_update_thread_merges_metadata_when_provided():
+    """Test that update_thread merges metadata correctly when provided."""
+    # Create a mock data layer
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    # Mock the execute_query method to return existing metadata
+    existing_metadata = {"is_shared": True, "custom_field": "original"}
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            # Return existing thread metadata
+            return [{"metadata": json.dumps(existing_metadata)}]
+        # For the UPDATE/INSERT, return None
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    # Call update_thread with partial metadata update
+    new_metadata = {"custom_field": "updated", "new_field": "added"}
+    await data_layer.update_thread(
+        thread_id="test-thread-123", name="Updated Name", metadata=new_metadata
+    )
+
+    # Verify execute_query was called twice (once for SELECT, once for UPDATE)
+    assert data_layer.execute_query.call_count == 2
+
+    # Get the UPDATE call
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_params = update_call[0][1]
+
+    # The metadata should be merged
+    # Expected: {"is_shared": True, "custom_field": "updated", "new_field": "added"}
+    # Find the JSON metadata in the params
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    assert metadata_json is not None
+    assert metadata_json.get("is_shared") is True
+    assert metadata_json.get("custom_field") == "updated"
+    assert metadata_json.get("new_field") == "added"
+
+
+@pytest.mark.asyncio
+async def test_update_thread_deletes_keys_with_none_values():
+    """Test that update_thread deletes keys when value is None."""
+    # Create a mock data layer
+    data_layer = ChainlitDataLayer(
+        database_url="postgresql://test", storage_client=None, show_logger=False
+    )
+
+    # Mock the execute_query method to return existing metadata
+    existing_metadata = {
+        "is_shared": True,
+        "to_delete": "will be removed",
+        "keep": "stays",
+    }
+
+    async def mock_execute_query(query, params):
+        if "SELECT" in query and "metadata" in query:
+            # Return existing thread metadata
+            return [{"metadata": json.dumps(existing_metadata)}]
+        # For the UPDATE/INSERT, return None
+        return None
+
+    data_layer.execute_query = AsyncMock(side_effect=mock_execute_query)
+
+    # Call update_thread with None value to delete a key
+    new_metadata = {"to_delete": None, "new_field": "added"}
+    await data_layer.update_thread(thread_id="test-thread-123", metadata=new_metadata)
+
+    # Verify execute_query was called twice
+    assert data_layer.execute_query.call_count == 2
+
+    # Get the UPDATE call
+    update_call = data_layer.execute_query.call_args_list[1]
+    update_params = update_call[0][1]
+
+    # The metadata should have deleted "to_delete" key and added "new_field"
+    # Expected: {"is_shared": True, "keep": "stays", "new_field": "added"}
+    metadata_json = None
+    for value in update_params.values():
+        if isinstance(value, str) and value.startswith("{"):
+            try:
+                metadata_json = json.loads(value)
+                break
+            except json.JSONDecodeError:
+                pass
+
+    if metadata_json:
+        # Verify "to_delete" is not in the merged metadata
+        assert "to_delete" not in metadata_json
+        # Verify "new_field" was added
+        assert metadata_json.get("new_field") == "added"
+        # Verify "is_shared" and "keep" are preserved
+        assert metadata_json.get("is_shared") is True
+        assert metadata_json.get("keep") == "stays"
+
+
+@pytest.mark.asyncio
+async def test_create_step_uses_nullif_for_output_and_input():
+    """Empty-string output/input should not overwrite existing content.
+
+    Regression test for https://github.com/Chainlit/chainlit/issues/2789
+    The SQL uses NULLIF(EXCLUDED.output, '') so that an empty string from the
+    initial Step.send() is treated as NULL by COALESCE, preventing it from
+    overwriting non-empty content saved by a subsequent Step.update().
+    """
+    import inspect
+
+    source = inspect.getsource(ChainlitDataLayer.create_step)
+
+    assert "NULLIF(EXCLUDED.output, '')" in source, (
+        "output should use NULLIF to treat empty string as NULL"
+    )
+    assert "NULLIF(EXCLUDED.input, '')" in source, (
+        "input should use NULLIF to treat empty string as NULL"
+    )

backend/tests/data/test_get_data_layer.py

@@ -0,0 +1,18 @@
+from unittest.mock import AsyncMock, Mock
+
+from chainlit.data import get_data_layer
+
+
+async def test_get_data_layer(
+    mock_data_layer: AsyncMock,
+    mock_get_data_layer: Mock,
+):
+    # Check whether the data layer is properly set
+    assert mock_data_layer == get_data_layer()
+
+    mock_get_data_layer.assert_called_once()
+
+    # Getting the data layer again, should not result in additional call
+    assert mock_data_layer == get_data_layer()
+
+    mock_get_data_layer.assert_called_once()

backend/tests/data/test_sql_alchemy.py

@@ -0,0 +1,327 @@
+import json
+import uuid
+from pathlib import Path
+
+import pytest
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine
+
+from chainlit import User
+from chainlit.data.sql_alchemy import SQLAlchemyDataLayer
+from chainlit.data.storage_clients.base import BaseStorageClient
+from chainlit.element import Text
+
+
+@pytest.fixture
+async def data_layer(mock_storage_client: BaseStorageClient, tmp_path: Path):
+    db_file = tmp_path / "test_db.sqlite"
+    conninfo = f"sqlite+aiosqlite:///{db_file}"
+
+    # Create async engine
+    engine = create_async_engine(conninfo)
+
+    # Execute initialization statements
+    # Ref: https://docs.chainlit.io/data-persistence/custom#sql-alchemy-data-layer
+    async with engine.begin() as conn:
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE users (
+                    "id" UUID PRIMARY KEY,
+                    "identifier" TEXT NOT NULL UNIQUE,
+                    "metadata" JSONB NOT NULL,
+                    "createdAt" TEXT
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS threads (
+                    "id" UUID PRIMARY KEY,
+                    "createdAt" TEXT,
+                    "name" TEXT,
+                    "userId" UUID,
+                    "userIdentifier" TEXT,
+                    "tags" TEXT[],
+                    "metadata" JSONB NOT NULL DEFAULT '{}',
+                    FOREIGN KEY ("userId") REFERENCES users("id") ON DELETE CASCADE
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS steps (
+                    "id" UUID PRIMARY KEY,
+                    "name" TEXT NOT NULL,
+                    "type" TEXT NOT NULL,
+                    "threadId" UUID NOT NULL,
+                    "parentId" UUID,
+                    "disableFeedback" BOOLEAN NOT NULL,
+                    "streaming" BOOLEAN NOT NULL,
+                    "waitForAnswer" BOOLEAN,
+                    "isError" BOOLEAN,
+                    "metadata" JSONB,
+                    "tags" TEXT[],
+                    "input" TEXT,
+                    "output" TEXT,
+                    "createdAt" TEXT,
+                    "start" TEXT,
+                    "end" TEXT,
+                    "generation" JSONB,
+                    "showInput" TEXT,
+                    "language" TEXT,
+                    "indent" INT
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS elements (
+                    "id" UUID PRIMARY KEY,
+                    "threadId" UUID,
+                    "type" TEXT,
+                    "url" TEXT,
+                    "chainlitKey" TEXT,
+                    "name" TEXT NOT NULL,
+                    "display" TEXT,
+                    "objectKey" TEXT,
+                    "size" TEXT,
+                    "page" INT,
+                    "language" TEXT,
+                    "forId" UUID,
+                    "mime" TEXT
+                );
+        """
+            )
+        )
+
+        await conn.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS feedbacks (
+                    "id" UUID PRIMARY KEY,
+                    "forId" UUID NOT NULL,
+                    "threadId" UUID NOT NULL,
+                    "value" INT NOT NULL,
+                    "comment" TEXT
+                );
+        """
+            )
+        )
+
+    # Create SQLAlchemyDataLayer instance
+    data_layer = SQLAlchemyDataLayer(conninfo, storage_provider=mock_storage_client)
+
+    return data_layer
+
+
+async def test_create_and_get_element(
+    mock_chainlit_context, data_layer: SQLAlchemyDataLayer
+):
+    async with mock_chainlit_context:
+        text_element = Text(
+            id=str(uuid.uuid4()),
+            name="test.txt",
+            mime="text/plain",
+            content="test content",
+            for_id="test_step_id",
+        )
+
+        # Needs context because of wrapper in utils.py
+        await data_layer.create_element(text_element)
+
+    retrieved_element = await data_layer.get_element(
+        text_element.thread_id, text_element.id
+    )
+    assert retrieved_element is not None
+    assert retrieved_element["id"] == text_element.id
+    assert retrieved_element["name"] == text_element.name
+    assert retrieved_element["mime"] == text_element.mime
+    # The 'content' field is not part of the ElementDict, so we remove this assertion
+
+
+async def test_get_current_timestamp(data_layer: SQLAlchemyDataLayer):
+    timestamp = await data_layer.get_current_timestamp()
+    assert isinstance(timestamp, str)
+
+
+async def test_get_user(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    fetched_user = await data_layer.get_user(persisted_user.identifier)
+
+    assert fetched_user
+    assert fetched_user.createdAt == persisted_user.createdAt
+    assert fetched_user.id == persisted_user.id
+
+    nonexistent_user = await data_layer.get_user("nonexistent")
+    assert nonexistent_user is None
+
+
+async def test_create_user(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+
+    assert persisted_user
+    assert persisted_user.identifier == test_user.identifier
+    assert persisted_user.createdAt
+    assert persisted_user.id
+
+    # Assert id is valid uuid
+    assert uuid.UUID(persisted_user.id)
+
+
+async def test_update_thread(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread")
+
+
+async def test_get_thread_author(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread", user_id=persisted_user.id)
+    author = await data_layer.get_thread_author("test_thread")
+
+    assert author == persisted_user.identifier
+
+
+async def test_get_thread(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread")
+    result = await data_layer.get_thread("test_thread")
+    assert result is not None
+
+    result = await data_layer.get_thread("nonexisting_thread")
+    assert result is None
+
+
+async def test_delete_thread(test_user: User, data_layer: SQLAlchemyDataLayer):
+    persisted_user = await data_layer.create_user(test_user)
+    assert persisted_user
+
+    await data_layer.update_thread("test_thread", "test_user")
+    await data_layer.delete_thread("test_thread")
+    thread = await data_layer.get_thread("test_thread")
+    assert thread is None
+
+
+async def _get_thread_metadata_raw(
+    data_layer: SQLAlchemyDataLayer, thread_id: str
+) -> str | None:
+    result = await data_layer.execute_sql(
+        query='SELECT "metadata" FROM threads WHERE "id" = :id',
+        parameters={"id": thread_id},
+    )
+    if isinstance(result, list) and result:
+        return result[0].get("metadata")
+    return None
+
+
+async def test_update_thread_without_metadata_stores_empty_dict(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_no_meta")
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_no_meta")
+    assert raw is not None
+    assert json.loads(raw) == {}
+
+
+async def test_update_thread_with_explicit_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread(
+        "thread_meta", metadata={"key": "value", "count": 42}
+    )
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_meta")
+    assert raw is not None
+    assert json.loads(raw) == {"key": "value", "count": 42}
+
+
+async def test_update_thread_with_empty_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_empty_meta", metadata={})
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_empty_meta")
+    assert raw is not None
+    assert json.loads(raw) == {}
+
+
+async def test_update_thread_preserves_metadata_on_noop_call(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_preserve", metadata={"important": "data"})
+
+    await data_layer.update_thread("thread_preserve")
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_preserve")
+    assert raw is not None
+    assert json.loads(raw) == {"important": "data"}
+
+
+async def test_update_thread_merges_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread(
+        "thread_merge", metadata={"existing": "old", "keep": "me"}
+    )
+
+    await data_layer.update_thread(
+        "thread_merge", metadata={"existing": "new", "added": "field"}
+    )
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_merge")
+    assert raw is not None
+    merged = json.loads(raw)
+    assert merged == {"existing": "new", "keep": "me", "added": "field"}
+
+
+async def test_update_thread_deletes_metadata_keys_via_none(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread(
+        "thread_delete_key", metadata={"a": 1, "b": 2, "c": 3}
+    )
+
+    await data_layer.update_thread("thread_delete_key", metadata={"b": None})
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_delete_key")
+    assert raw is not None
+    result = json.loads(raw)
+    assert result == {"a": 1, "c": 3}
+
+
+async def test_update_thread_name_update_preserves_metadata(
+    data_layer: SQLAlchemyDataLayer,
+):
+    await data_layer.update_thread("thread_name_meta", metadata={"saved": True})
+
+    await data_layer.update_thread("thread_name_meta", name="New Name")
+
+    raw = await _get_thread_metadata_raw(data_layer, "thread_name_meta")
+    assert raw is not None
+    assert json.loads(raw) == {"saved": True}
+
+    result = await data_layer.execute_sql(
+        query='SELECT "name" FROM threads WHERE "id" = :id',
+        parameters={"id": "thread_name_meta"},
+    )
+    assert isinstance(result, list)
+    assert result
+    assert result[0]["name"] == "New Name"

backend/tests/langchain/test_async_callback.py

@@ -0,0 +1,168 @@
+"""Tests for async LangChain callback handlers."""
+
+from datetime import datetime
+from unittest.mock import AsyncMock, Mock, patch
+from uuid import uuid4
+
+import pytest
+from langchain_core.outputs import GenerationChunk
+
+from chainlit.langchain.callbacks import LangchainTracer
+from chainlit.step import Step
+
+
+def create_mock_run(**kwargs):
+    """Helper to create a properly mocked Run object with all required attributes."""
+    run = Mock()
+    run.id = kwargs.get("id", uuid4())
+    run.parent_run_id = kwargs.get("parent_run_id", None)
+    run.name = kwargs.get("name", "test_run")
+    run.run_type = kwargs.get("run_type", "llm")
+    run.tags = kwargs.get("tags", [])
+    run.inputs = kwargs.get("inputs", {})
+    run.outputs = kwargs.get("outputs", None)
+    run.serialized = kwargs.get("serialized", {})
+    run.extra = kwargs.get("extra", {})
+    run.start_time = kwargs.get("start_time", datetime.now())
+    run.end_time = kwargs.get("end_time", None)
+    run.error = kwargs.get("error", None)
+    return run
+
+
+@pytest.fixture
+def mock_run():
+    """Create a mock LangChain Run object."""
+    return create_mock_run(
+        name="test_run",
+        run_type="llm",
+        inputs={"input": "test input"},
+        serialized={"name": "test_llm"},
+    )
+
+
+async def test_tracer_initialization(mock_chainlit_context):
+    """Test LangchainTracer initialization."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+
+        assert tracer.steps == {}
+        assert tracer.parent_id_map == {}
+        assert tracer.ignored_runs == set()
+        assert tracer.stream_final_answer is False
+        assert tracer.answer_reached is False
+
+
+async def test_on_llm_start(mock_chainlit_context):
+    """Test on_llm_start callback."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+        run_id = uuid4()
+        prompts = ["Test prompt"]
+
+        await tracer.on_llm_start(
+            serialized={"name": "test_llm"},
+            prompts=prompts,
+            run_id=run_id,
+        )
+
+        assert str(run_id) in tracer.completion_generations
+        completion_gen = tracer.completion_generations[str(run_id)]
+        assert completion_gen["prompt"] == "Test prompt"
+        assert completion_gen["token_count"] == 0
+        assert completion_gen["tt_first_token"] is None
+        assert "start" in completion_gen
+
+
+async def test_on_llm_new_token(mock_chainlit_context):
+    """Test on_llm_new_token with token streaming."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+        run_id = uuid4()
+
+        await tracer.on_llm_start(
+            serialized={"name": "test_llm"},
+            prompts=["Test prompt"],
+            run_id=run_id,
+        )
+
+        chunk = GenerationChunk(text="Hello")
+        await tracer.on_llm_new_token(
+            token="Hello",
+            chunk=chunk,
+            run_id=run_id,
+        )
+
+        completion_gen = tracer.completion_generations[str(run_id)]
+        assert completion_gen["token_count"] == 1
+        assert completion_gen["tt_first_token"] is not None
+
+
+async def test_start_trace(mock_chainlit_context):
+    """Test _start_trace creates steps correctly."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+
+        # Test LLM run
+        llm_run = create_mock_run(
+            id=uuid4(),
+            name="test_llm",
+            run_type="llm",
+            inputs={"input": "test"},
+        )
+
+        with patch.object(Step, "send", new_callable=AsyncMock):
+            await tracer._start_trace(llm_run)
+
+        assert str(llm_run.id) in tracer.steps
+        assert tracer.steps[str(llm_run.id)].type == "llm"
+
+        # Test ignored run
+        ignored_run = create_mock_run(
+            id=uuid4(),
+            name="RunnableSequence",
+            run_type="chain",
+        )
+        tracer.to_ignore = ["RunnableSequence"]
+        await tracer._start_trace(ignored_run)
+
+        assert str(ignored_run.id) in tracer.ignored_runs
+
+
+async def test_on_run_update(mock_chainlit_context):
+    """Test _on_run_update updates steps."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+        run_id = uuid4()
+
+        step = Step(id=str(run_id), name="test_tool", type="tool")
+        tracer.steps[str(run_id)] = step
+
+        run = create_mock_run(
+            id=run_id,
+            name="test_tool",
+            run_type="tool",
+            outputs={"output": "result"},
+        )
+
+        with patch.object(step, "update", new_callable=AsyncMock):
+            await tracer._on_run_update(run)
+
+        assert step.output is not None
+
+
+async def test_error_handling(mock_chainlit_context):
+    """Test error handling in callbacks."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+        run_id = uuid4()
+
+        step = Step(id=str(run_id), name="test_step", type="llm")
+        tracer.steps[str(run_id)] = step
+
+        error = ValueError("Test error")
+
+        with patch.object(step, "update", new_callable=AsyncMock):
+            await tracer._on_error(error, run_id=run_id)
+
+        assert step.is_error is True
+        assert step.output == "Test error"

backend/tests/langchain/test_chain_types.py

@@ -0,0 +1,134 @@
+"""Tests for different LangChain chain types and their integration with Chainlit."""
+
+from datetime import datetime
+from unittest.mock import AsyncMock, Mock, patch
+from uuid import uuid4
+
+from chainlit.langchain.callbacks import LangchainTracer
+from chainlit.step import Step
+
+
+def create_mock_run(**kwargs):
+    """Helper to create a properly mocked Run object with all required attributes."""
+    run = Mock()
+    run.id = kwargs.get("id", uuid4())
+    run.parent_run_id = kwargs.get("parent_run_id", None)
+    run.name = kwargs.get("name", "test_run")
+    run.run_type = kwargs.get("run_type", "llm")
+    run.tags = kwargs.get("tags", [])
+    run.inputs = kwargs.get("inputs", {})
+    run.outputs = kwargs.get("outputs", None)
+    run.serialized = kwargs.get("serialized", {})
+    run.extra = kwargs.get("extra", {})
+    run.start_time = kwargs.get("start_time", datetime.now())
+    run.end_time = kwargs.get("end_time", None)
+    run.error = kwargs.get("error", None)
+    return run
+
+
+async def test_different_run_types(mock_chainlit_context):
+    """Test different LangChain run types create correct steps."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+
+        # Test agent run
+        agent_run = create_mock_run(
+            id=uuid4(),
+            name="test_agent",
+            run_type="agent",
+            inputs={"input": "test"},
+        )
+
+        with patch.object(Step, "send", new_callable=AsyncMock):
+            await tracer._start_trace(agent_run)
+
+        assert str(agent_run.id) in tracer.steps
+        assert tracer.steps[str(agent_run.id)].type == "run"
+
+        # Test tool run
+        tool_run = create_mock_run(
+            id=uuid4(),
+            name="test_tool",
+            run_type="tool",
+            inputs={"query": "test"},
+        )
+
+        with patch.object(Step, "send", new_callable=AsyncMock):
+            await tracer._start_trace(tool_run)
+
+        assert str(tool_run.id) in tracer.steps
+        assert tracer.steps[str(tool_run.id)].type == "tool"
+
+
+async def test_nested_chain_hierarchy(mock_chainlit_context):
+    """Test nested chain with LLM calls."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+
+        # Create parent chain
+        chain_run = create_mock_run(
+            id=uuid4(),
+            name="parent_chain",
+            run_type="chain",
+            inputs={"input": "test"},
+        )
+
+        with patch.object(Step, "send", new_callable=AsyncMock):
+            await tracer._start_trace(chain_run)
+
+        # Create nested LLM
+        llm_run = create_mock_run(
+            id=uuid4(),
+            parent_run_id=chain_run.id,
+            name="nested_llm",
+            run_type="llm",
+            inputs={},
+        )
+
+        with patch.object(Step, "send", new_callable=AsyncMock):
+            await tracer._start_trace(llm_run)
+
+        # Both should exist
+        assert str(chain_run.id) in tracer.steps
+        assert str(llm_run.id) in tracer.steps
+
+        # LLM should have chain as parent
+        llm_step = tracer.steps[str(llm_run.id)]
+        assert llm_step.parent_id == str(chain_run.id)
+
+
+async def test_ignored_runs(mock_chainlit_context):
+    """Test that default ignored runs are properly filtered."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+
+        # Test RunnableSequence is ignored
+        sequence_run = create_mock_run(
+            id=uuid4(),
+            name="RunnableSequence",
+            run_type="chain",
+            inputs={},
+        )
+
+        await tracer._start_trace(sequence_run)
+
+        assert str(sequence_run.id) not in tracer.steps
+        assert str(sequence_run.id) in tracer.ignored_runs
+
+
+async def test_custom_filtering(mock_chainlit_context):
+    """Test custom to_ignore and to_keep lists."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer(to_ignore=["CustomIgnore"], to_keep=["llm", "tool"])
+
+        # Test custom ignore
+        custom_run = create_mock_run(
+            id=uuid4(),
+            name="CustomIgnore",
+            run_type="chain",
+        )
+
+        await tracer._start_trace(custom_run)
+
+        assert str(custom_run.id) not in tracer.steps
+        assert str(custom_run.id) in tracer.ignored_runs

backend/tests/langchain/test_sync_callback.py

@@ -0,0 +1,147 @@
+"""Tests for synchronous LangChain callback operations and helper classes."""
+
+from datetime import datetime
+from unittest.mock import Mock
+from uuid import uuid4
+
+from langchain_core.messages import AIMessage, HumanMessage
+
+from chainlit.langchain.callbacks import (
+    FinalStreamHelper,
+    GenerationHelper,
+    LangchainTracer,
+)
+
+
+def create_mock_run(**kwargs):
+    """Helper to create a properly mocked Run object with all required attributes."""
+    run = Mock()
+    run.id = kwargs.get("id", uuid4())
+    run.parent_run_id = kwargs.get("parent_run_id", None)
+    run.name = kwargs.get("name", "test_run")
+    run.run_type = kwargs.get("run_type", "llm")
+    run.tags = kwargs.get("tags", [])
+    run.inputs = kwargs.get("inputs", {})
+    run.outputs = kwargs.get("outputs", None)
+    run.serialized = kwargs.get("serialized", {})
+    run.extra = kwargs.get("extra", {})
+    run.start_time = kwargs.get("start_time", datetime.now())
+    run.end_time = kwargs.get("end_time", None)
+    run.error = kwargs.get("error", None)
+    return run
+
+
+class TestFinalStreamHelper:
+    """Test FinalStreamHelper class."""
+
+    def test_initialization(self):
+        """Test FinalStreamHelper initialization."""
+        helper = FinalStreamHelper()
+
+        assert helper.answer_prefix_tokens == ["Final", "Answer", ":"]
+        assert helper.stream_final_answer is False
+        assert helper.answer_reached is False
+
+    def test_check_if_answer_reached(self):
+        """Test _check_if_answer_reached."""
+        helper = FinalStreamHelper()
+
+        helper._append_to_last_tokens("Final")
+        helper._append_to_last_tokens("Answer")
+        helper._append_to_last_tokens(":")
+
+        assert helper._check_if_answer_reached() is True
+
+
+class TestGenerationHelper:
+    """Test GenerationHelper class."""
+
+    def test_initialization(self):
+        """Test GenerationHelper initialization."""
+        helper = GenerationHelper()
+
+        assert helper.chat_generations == {}
+        assert helper.completion_generations == {}
+        assert helper.generation_inputs == {}
+
+    def test_ensure_values_serializable(self):
+        """Test ensure_values_serializable method."""
+        helper = GenerationHelper()
+
+        # Test dict
+        result = helper.ensure_values_serializable({"key": "value", "num": 42})
+        assert result == {"key": "value", "num": 42}
+
+        # Test list
+        result = helper.ensure_values_serializable([1, 2, "three"])
+        assert result == [1, 2, "three"]
+
+    def test_convert_message(self):
+        """Test _convert_message method."""
+        helper = GenerationHelper()
+
+        # Test HumanMessage
+        msg = HumanMessage(content="Hello")
+        result = helper._convert_message(msg)
+        assert result["role"] == "user"
+        assert result["content"] == "Hello"
+
+        # Test AIMessage
+        msg = AIMessage(content="Hi there")
+        result = helper._convert_message(msg)
+        assert result["role"] == "assistant"
+        assert result["content"] == "Hi there"
+
+    def test_build_llm_settings(self):
+        """Test _build_llm_settings method."""
+        helper = GenerationHelper()
+
+        serialized = {"name": "ChatOpenAI"}
+        invocation_params = {
+            "_type": "openai",
+            "model": "gpt-4",
+            "temperature": 0.7,
+        }
+
+        provider, model, _, settings = helper._build_llm_settings(
+            serialized, invocation_params
+        )
+
+        assert provider == "openai"
+        assert model == "gpt-4"
+        assert settings["temperature"] == 0.7
+
+
+async def test_should_ignore_run(mock_chainlit_context):
+    """Test _should_ignore_run method."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer(to_ignore=["RunnableSequence"])
+
+        # Test ignored by name
+        run = create_mock_run(name="RunnableSequence", run_type="chain")
+        ignore, _ = tracer._should_ignore_run(run)
+        assert ignore is True
+
+        # Test not ignored
+        run = create_mock_run(name="MyChain", run_type="chain")
+        ignore, _ = tracer._should_ignore_run(run)
+        assert ignore is False
+
+
+async def test_get_non_ignored_parent_id(mock_chainlit_context):
+    """Test _get_non_ignored_parent_id."""
+    async with mock_chainlit_context:
+        tracer = LangchainTracer()
+
+        # Setup parent chain
+        grandparent_id = str(uuid4())
+        parent_id = str(uuid4())
+
+        # Both parent and grandparent need to be in parent_id_map for the traversal to work
+        tracer.parent_id_map[parent_id] = grandparent_id
+        tracer.parent_id_map[grandparent_id] = None  # grandparent has no parent
+        tracer.ignored_runs.add(parent_id)
+
+        result = tracer._get_non_ignored_parent_id(parent_id)
+
+        assert result == grandparent_id

backend/tests/llama_index/test_callbacks.py

@@ -0,0 +1,97 @@
+from unittest.mock import patch
+
+from llama_index.core.callbacks.schema import CBEventType, EventPayload
+from llama_index.core.tools.types import ToolMetadata
+
+from chainlit.llama_index.callbacks import LlamaIndexCallbackHandler
+from chainlit.step import Step
+
+
+async def test_on_event_start_for_function_calls(mock_chainlit_context):
+    TEST_EVENT_ID = "test_event_id"
+    async with mock_chainlit_context:
+        handler = LlamaIndexCallbackHandler()
+
+        with patch.object(Step, "send") as mock_send:
+            result = handler.on_event_start(
+                CBEventType.FUNCTION_CALL,
+                {
+                    EventPayload.TOOL: ToolMetadata(
+                        name="test_tool", description="test_description"
+                    ),
+                    EventPayload.FUNCTION_CALL: {"arg1": "value1"},
+                },
+                TEST_EVENT_ID,
+            )
+
+        assert result == TEST_EVENT_ID
+        assert TEST_EVENT_ID in handler.steps
+        step = handler.steps[TEST_EVENT_ID]
+        assert isinstance(step, Step)
+        assert step.name == "test_tool"
+        assert step.type == "tool"
+        assert step.id == TEST_EVENT_ID
+        assert step.input == '{\n    "arg1": "value1"\n}'
+        mock_send.assert_called_once()
+
+
+async def test_on_event_start_for_function_calls_missing_payload(mock_chainlit_context):
+    TEST_EVENT_ID = "test_event_id"
+    async with mock_chainlit_context:
+        handler = LlamaIndexCallbackHandler()
+
+        with patch.object(Step, "send") as mock_send:
+            result = handler.on_event_start(
+                CBEventType.FUNCTION_CALL,
+                None,
+                TEST_EVENT_ID,
+            )
+
+        assert result == TEST_EVENT_ID
+        assert TEST_EVENT_ID in handler.steps
+        step = handler.steps[TEST_EVENT_ID]
+        assert isinstance(step, Step)
+        assert step.name == "function_call"
+        assert step.type == "tool"
+        assert step.id == TEST_EVENT_ID
+        assert step.input == "{}"
+        mock_send.assert_called_once()
+
+
+async def test_on_event_end_for_function_calls(mock_chainlit_context):
+    TEST_EVENT_ID = "test_event_id"
+    async with mock_chainlit_context:
+        handler = LlamaIndexCallbackHandler()
+        # Pretend that we have started a step before.
+        step = Step(name="test_tool", type="tool", id=TEST_EVENT_ID)
+        handler.steps[TEST_EVENT_ID] = step
+
+        with patch.object(step, "update") as mock_send:
+            handler.on_event_end(
+                CBEventType.FUNCTION_CALL,
+                payload={EventPayload.FUNCTION_OUTPUT: "test_output"},
+                event_id=TEST_EVENT_ID,
+            )
+
+        assert step.output == "test_output"
+        assert TEST_EVENT_ID not in handler.steps
+        mock_send.assert_called_once()
+
+
+async def test_on_event_end_for_function_calls_missing_payload(mock_chainlit_context):
+    TEST_EVENT_ID = "test_event_id"
+    async with mock_chainlit_context:
+        handler = LlamaIndexCallbackHandler()
+        # Pretend that we have started a step before.
+        step = Step(name="test_tool", type="tool", id=TEST_EVENT_ID)
+        handler.steps[TEST_EVENT_ID] = step
+
+        with patch.object(step, "update") as mock_send:
+            handler.on_event_end(
+                CBEventType.FUNCTION_CALL,
+                payload=None,
+                event_id=TEST_EVENT_ID,
+            )
+        # TODO: Is this the desired behavior? Shouldn't we still remove the step as long as we've been told it has ended, even if the payload is missing?
+        assert TEST_EVENT_ID in handler.steps
+        mock_send.assert_not_called()

backend/tests/test_action.py

@@ -0,0 +1,281 @@
+import uuid
+
+import pytest
+
+from chainlit.action import Action
+
+
+@pytest.mark.asyncio
+class TestAction:
+    """Test suite for the Action class."""
+
+    async def test_action_initialization_with_required_fields(self):
+        """Test Action initialization with only required fields."""
+        action = Action(
+            name="test_action",
+            payload={"key": "value"},
+        )
+
+        assert action.name == "test_action"
+        assert action.payload == {"key": "value"}
+        assert action.label == ""
+        assert action.tooltip == ""
+        assert action.icon is None
+        assert action.forId is None
+        assert isinstance(action.id, str)
+        # Verify ID is a valid UUID
+        uuid.UUID(action.id)
+
+    async def test_action_initialization_with_all_fields(self):
+        """Test Action initialization with all fields provided."""
+        test_id = str(uuid.uuid4())
+        action = Action(
+            name="custom_action",
+            payload={"param1": "value1", "param2": 42},
+            label="Click Me",
+            tooltip="This is a tooltip",
+            icon="check-circle",
+            id=test_id,
+        )
+
+        assert action.name == "custom_action"
+        assert action.payload == {"param1": "value1", "param2": 42}
+        assert action.label == "Click Me"
+        assert action.tooltip == "This is a tooltip"
+        assert action.icon == "check-circle"
+        assert action.id == test_id
+        assert action.forId is None
+
+    async def test_action_id_auto_generation(self):
+        """Test that Action generates unique IDs automatically."""
+        action1 = Action(name="action1", payload={})
+        action2 = Action(name="action2", payload={})
+
+        assert action1.id != action2.id
+        # Verify both are valid UUIDs
+        uuid.UUID(action1.id)
+        uuid.UUID(action2.id)
+
+    async def test_action_to_dict(self):
+        """Test Action serialization to dictionary."""
+        test_id = str(uuid.uuid4())
+        action = Action(
+            name="test_action",
+            payload={"data": "test"},
+            label="Test Label",
+            tooltip="Test Tooltip",
+            icon="star",
+            id=test_id,
+        )
+
+        action_dict = action.to_dict()
+
+        assert action_dict["name"] == "test_action"
+        assert action_dict["payload"] == {"data": "test"}
+        assert action_dict["label"] == "Test Label"
+        assert action_dict["tooltip"] == "Test Tooltip"
+        assert action_dict["icon"] == "star"
+        assert action_dict["id"] == test_id
+        assert action_dict.get("forId") is None
+
+    async def test_action_to_dict_with_for_id(self):
+        """Test Action serialization includes forId when set."""
+        action = Action(name="test", payload={})
+        action.forId = "message_123"
+
+        action_dict = action.to_dict()
+
+        assert action_dict["forId"] == "message_123"
+
+    async def test_action_send(self, mock_chainlit_context):
+        """Test Action.send() method emits correct event."""
+        async with mock_chainlit_context as ctx:
+            action = Action(
+                name="send_action",
+                payload={"test": "data"},
+                label="Send Test",
+            )
+
+            for_id = "target_message_id"
+            await action.send(for_id=for_id)
+
+            # Verify forId was set
+            assert action.forId == for_id
+
+            # Verify emit was called with correct parameters
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "action"
+
+            emitted_dict = call_args[0][1]
+            assert emitted_dict["name"] == "send_action"
+            assert emitted_dict["payload"] == {"test": "data"}
+            assert emitted_dict["label"] == "Send Test"
+            assert emitted_dict["forId"] == for_id
+
+    async def test_action_send_updates_for_id(self, mock_chainlit_context):
+        """Test that send() updates the forId field."""
+        async with mock_chainlit_context:
+            action = Action(name="test", payload={})
+
+            # Initially forId should be None
+            assert action.forId is None
+
+            # Send with first for_id
+            await action.send(for_id="first_id")
+            assert action.forId == "first_id"
+
+            # Send with different for_id should update
+            await action.send(for_id="second_id")
+            assert action.forId == "second_id"
+
+    async def test_action_remove(self, mock_chainlit_context):
+        """Test Action.remove() method emits correct event."""
+        async with mock_chainlit_context as ctx:
+            action = Action(
+                name="remove_action",
+                payload={"key": "value"},
+                label="Remove Test",
+            )
+            action.forId = "message_123"
+
+            await action.remove()
+
+            # Verify emit was called with correct parameters
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "remove_action"
+
+            emitted_dict = call_args[0][1]
+            assert emitted_dict["name"] == "remove_action"
+            assert emitted_dict["payload"] == {"key": "value"}
+            assert emitted_dict["forId"] == "message_123"
+
+    async def test_action_remove_without_for_id(self, mock_chainlit_context):
+        """Test Action.remove() works even without forId set."""
+        async with mock_chainlit_context as ctx:
+            action = Action(name="test", payload={})
+
+            await action.remove()
+
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "remove_action"
+
+    async def test_action_with_complex_payload(self):
+        """Test Action with complex nested payload."""
+        complex_payload = {
+            "nested": {
+                "data": [1, 2, 3],
+                "info": {"key": "value"},
+            },
+            "list": ["a", "b", "c"],
+            "number": 42,
+            "boolean": True,
+            "null": None,
+        }
+
+        action = Action(name="complex", payload=complex_payload)
+
+        assert action.payload == complex_payload
+        action_dict = action.to_dict()
+        assert action_dict["payload"] == complex_payload
+
+    async def test_action_with_empty_payload(self):
+        """Test Action with empty payload."""
+        action = Action(name="empty", payload={})
+
+        assert action.payload == {}
+        assert action.to_dict()["payload"] == {}
+
+    async def test_action_with_empty_strings(self):
+        """Test Action handles empty strings correctly."""
+        action = Action(
+            name="test",
+            payload={},
+            label="",
+            tooltip="",
+        )
+
+        assert action.label == ""
+        assert action.tooltip == ""
+
+        action_dict = action.to_dict()
+        assert action_dict["label"] == ""
+        assert action_dict["tooltip"] == ""
+
+    async def test_action_serialization_deserialization(self):
+        """Test Action can be serialized and deserialized."""
+        original = Action(
+            name="serialize_test",
+            payload={"data": "test"},
+            label="Test",
+            tooltip="Tooltip",
+            icon="icon-name",
+        )
+
+        # Serialize to dict
+        serialized = original.to_dict()
+
+        # Deserialize from dict
+        deserialized = Action.from_dict(serialized)
+
+        assert deserialized.name == original.name
+        assert deserialized.payload == original.payload
+        assert deserialized.label == original.label
+        assert deserialized.tooltip == original.tooltip
+        assert deserialized.icon == original.icon
+        assert deserialized.id == original.id
+
+    async def test_multiple_actions_with_same_name(self):
+        """Test that multiple actions can have the same name but different IDs."""
+        action1 = Action(name="duplicate", payload={"num": 1})
+        action2 = Action(name="duplicate", payload={"num": 2})
+
+        assert action1.name == action2.name
+        assert action1.id != action2.id
+        assert action1.payload != action2.payload
+
+    async def test_action_send_multiple_times(self, mock_chainlit_context):
+        """Test that an action can be sent multiple times."""
+        async with mock_chainlit_context as ctx:
+            action = Action(name="multi_send", payload={})
+
+            await action.send(for_id="id1")
+            await action.send(for_id="id2")
+            await action.send(for_id="id3")
+
+            # Should have been called 3 times
+            assert ctx.emitter.emit.call_count == 3
+
+            # Last forId should be id3
+            assert action.forId == "id3"
+
+    async def test_action_with_special_characters_in_payload(self):
+        """Test Action handles special characters in payload."""
+        special_payload = {
+            "unicode": "Hello 世界 🌍",
+            "quotes": 'He said "Hello"',
+            "newlines": "Line1\nLine2\nLine3",
+            "tabs": "Col1\tCol2\tCol3",
+        }
+
+        action = Action(name="special", payload=special_payload)
+
+        assert action.payload == special_payload
+        action_dict = action.to_dict()
+        assert action_dict["payload"] == special_payload
+
+    async def test_action_icon_variations(self):
+        """Test Action with different icon values."""
+        # With icon
+        action_with_icon = Action(name="test", payload={}, icon="check")
+        assert action_with_icon.icon == "check"
+
+        # Without icon (None)
+        action_no_icon = Action(name="test", payload={}, icon=None)
+        assert action_no_icon.icon is None
+
+        # Default (should be None)
+        action_default = Action(name="test", payload={})
+        assert action_default.icon is None

backend/tests/test_cache.py

@@ -0,0 +1,558 @@
+import sys
+import threading
+from unittest.mock import Mock, patch
+
+import pytest
+
+from chainlit.cache import cache, init_lc_cache
+
+# Import the actual cache module to access _cache dict
+cache_module = sys.modules["chainlit.cache"]
+
+
+class TestCacheDecorator:
+    """Test suite for the @cache decorator."""
+
+    def setup_method(self):
+        """Clear the cache before each test."""
+        cache_module._cache.clear()
+
+    def teardown_method(self):
+        """Clear the cache after each test."""
+        cache_module._cache.clear()
+
+    def test_cache_basic_function(self):
+        """Test that cache decorator caches function results."""
+        call_count = 0
+
+        @cache
+        def add(a, b):
+            nonlocal call_count
+            call_count += 1
+            return a + b
+
+        # First call
+        result1 = add(2, 3)
+        assert result1 == 5
+        assert call_count == 1
+
+        # Second call with same args should use cache
+        result2 = add(2, 3)
+        assert result2 == 5
+        assert call_count == 1  # Function not called again
+
+    def test_cache_different_arguments(self):
+        """Test that different arguments create different cache entries."""
+        call_count = 0
+
+        @cache
+        def multiply(x, y):
+            nonlocal call_count
+            call_count += 1
+            return x * y
+
+        result1 = multiply(2, 3)
+        assert result1 == 6
+        assert call_count == 1
+
+        result2 = multiply(4, 5)
+        assert result2 == 20
+        assert call_count == 2  # Different args, function called again
+
+        # Same args as first call, should use cache
+        result3 = multiply(2, 3)
+        assert result3 == 6
+        assert call_count == 2  # Cache hit, no new call
+
+    def test_cache_with_kwargs(self):
+        """Test that cache works with keyword arguments."""
+        call_count = 0
+
+        @cache
+        def greet(name, greeting="Hello"):
+            nonlocal call_count
+            call_count += 1
+            return f"{greeting}, {name}!"
+
+        result1 = greet("Alice", greeting="Hi")
+        assert result1 == "Hi, Alice!"
+        assert call_count == 1
+
+        # Same call should use cache
+        result2 = greet("Alice", greeting="Hi")
+        assert result2 == "Hi, Alice!"
+        assert call_count == 1
+
+        # Different kwargs should call function
+        result3 = greet("Alice", greeting="Hello")
+        assert result3 == "Hello, Alice!"
+        assert call_count == 2
+
+    def test_cache_kwargs_order_independence(self):
+        """Test that kwargs order doesn't affect cache key."""
+        call_count = 0
+
+        @cache
+        def func(a=1, b=2, c=3):
+            nonlocal call_count
+            call_count += 1
+            return a + b + c
+
+        result1 = func(a=1, b=2, c=3)
+        assert result1 == 6
+        assert call_count == 1
+
+        # Same kwargs, different order - should use cache
+        result2 = func(c=3, a=1, b=2)
+        assert result2 == 6
+        assert call_count == 1  # Cache hit
+
+    def test_cache_mixed_args_and_kwargs(self):
+        """Test cache with both positional and keyword arguments."""
+        call_count = 0
+
+        @cache
+        def compute(x, y, z=10):
+            nonlocal call_count
+            call_count += 1
+            return x + y + z
+
+        result1 = compute(1, 2, z=3)
+        assert result1 == 6
+        assert call_count == 1
+
+        result2 = compute(1, 2, z=3)
+        assert result2 == 6
+        assert call_count == 1  # Cache hit
+
+        result3 = compute(1, 2, z=5)
+        assert result3 == 8
+        assert call_count == 2  # Different z value
+
+    def test_cache_with_no_arguments(self):
+        """Test cache with functions that take no arguments."""
+        call_count = 0
+
+        @cache
+        def get_constant():
+            nonlocal call_count
+            call_count += 1
+            return 42
+
+        result1 = get_constant()
+        assert result1 == 42
+        assert call_count == 1
+
+        result2 = get_constant()
+        assert result2 == 42
+        assert call_count == 1  # Cache hit
+
+    def test_cache_with_mutable_return_value(self):
+        """Test that cache returns the same object reference."""
+        call_count = 0
+
+        @cache
+        def get_list():
+            nonlocal call_count
+            call_count += 1
+            return [1, 2, 3]
+
+        result1 = get_list()
+        assert result1 == [1, 2, 3]
+        assert call_count == 1
+
+        result2 = get_list()
+        assert result2 == [1, 2, 3]
+        assert call_count == 1
+
+        # Both results should be the same object
+        assert result1 is result2
+
+    def test_cache_thread_safety(self):
+        """Test that cache is thread-safe."""
+        call_count = 0
+        call_lock = threading.Lock()
+
+        @cache
+        def slow_function(x):
+            nonlocal call_count
+            with call_lock:
+                call_count += 1
+            return x * 2
+
+        results = []
+
+        def worker():
+            result = slow_function(5)
+            results.append(result)
+
+        threads = [threading.Thread(target=worker) for _ in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # All results should be the same
+        assert all(r == 10 for r in results)
+        # Function should only be called once despite multiple threads
+        assert call_count == 1
+
+    def test_cache_different_functions_same_args(self):
+        """Test that different functions with same args have separate cache."""
+        call_count_1 = 0
+        call_count_2 = 0
+
+        @cache
+        def func1(x):
+            nonlocal call_count_1
+            call_count_1 += 1
+            return x + 1
+
+        @cache
+        def func2(x):
+            nonlocal call_count_2
+            call_count_2 += 1
+            return x + 2
+
+        result1 = func1(5)
+        assert result1 == 6
+        assert call_count_1 == 1
+
+        result2 = func2(5)
+        assert result2 == 7
+        assert call_count_2 == 1
+
+        # Both should use their own cache
+        func1(5)
+        func2(5)
+        assert call_count_1 == 1
+        assert call_count_2 == 1
+
+    def test_cache_with_none_arguments(self):
+        """Test cache with None as argument."""
+        call_count = 0
+
+        @cache
+        def process(value):
+            nonlocal call_count
+            call_count += 1
+            return value
+
+        result1 = process(None)
+        assert result1 is None
+        assert call_count == 1
+
+        result2 = process(None)
+        assert result2 is None
+        assert call_count == 1  # Cache hit
+
+    def test_cache_preserves_function_behavior(self):
+        """Test that cache decorator preserves original function behavior."""
+
+        @cache
+        def divide(a, b):
+            return a / b
+
+        assert divide(10, 2) == 5.0
+        assert divide(10, 2) == 5.0  # Cached
+
+        with pytest.raises(ZeroDivisionError):
+            divide(10, 0)
+
+
+class TestInitLcCache:
+    """Test suite for init_lc_cache function."""
+
+    def test_init_lc_cache_disabled_by_config(self):
+        """Test that cache is not initialized when disabled in config."""
+        with patch.object(cache_module.config, "project") as mock_project:
+            mock_project.cache = False
+            with patch.object(cache_module.config, "run") as mock_run:
+                mock_run.no_cache = False
+
+                with patch.object(
+                    cache_module.importlib.util, "find_spec"
+                ) as mock_find_spec:
+                    init_lc_cache()
+
+                    # Should not check for langchain if cache is disabled
+                    mock_find_spec.assert_not_called()
+
+    def test_init_lc_cache_disabled_by_no_cache_flag(self):
+        """Test that cache is not initialized when no_cache flag is set."""
+        with patch.object(cache_module.config, "project") as mock_project:
+            mock_project.cache = True
+            with patch.object(cache_module.config, "run") as mock_run:
+                mock_run.no_cache = True
+
+                with patch.object(
+                    cache_module.importlib.util, "find_spec"
+                ) as mock_find_spec:
+                    init_lc_cache()
+
+                    # Should not check for langchain if no_cache is True
+                    mock_find_spec.assert_not_called()
+
+    def test_init_lc_cache_langchain_not_installed(self):
+        """Test behavior when langchain is not installed."""
+        with patch.object(cache_module.config, "project") as mock_project:
+            mock_project.cache = True
+            with patch.object(cache_module.config, "run") as mock_run:
+                mock_run.no_cache = False
+
+                with patch.object(
+                    cache_module.importlib.util, "find_spec", return_value=None
+                ) as mock_find_spec:
+                    # Should not raise an error
+                    init_lc_cache()
+
+                    mock_find_spec.assert_called_once_with("langchain")
+
+    def test_init_lc_cache_with_langchain_installed(self):
+        """Test cache initialization when langchain is installed."""
+        with patch.object(cache_module.config, "project") as mock_project:
+            mock_project.cache = True
+            mock_project.lc_cache_path = "/tmp/test_cache.db"
+            with patch.object(cache_module.config, "run") as mock_run:
+                mock_run.no_cache = False
+
+                mock_spec = Mock()
+                with patch.object(
+                    cache_module.importlib.util, "find_spec", return_value=mock_spec
+                ):
+                    # Mock langchain modules
+                    mock_sqlite_cache = Mock()
+                    mock_set_llm_cache = Mock()
+
+                    with patch.dict(
+                        sys.modules,
+                        {
+                            "langchain": Mock(),
+                            "langchain_community.cache": Mock(
+                                SQLiteCache=mock_sqlite_cache
+                            ),
+                            "langchain.cache": Mock(SQLiteCache=mock_sqlite_cache),
+                            "langchain_core.globals": Mock(
+                                set_llm_cache=mock_set_llm_cache
+                            ),
+                            "langchain.globals": Mock(set_llm_cache=mock_set_llm_cache),
+                        },
+                    ):
+                        with patch("os.path.exists", return_value=True):
+                            init_lc_cache()
+
+                            mock_sqlite_cache.assert_called_once_with(
+                                database_path="/tmp/test_cache.db"
+                            )
+                            mock_set_llm_cache.assert_called_once()
+
+    def test_init_lc_cache_creates_new_cache_file(self):
+        """Test that logger is called when creating new cache file."""
+        with patch.object(cache_module.config, "project") as mock_project:
+            mock_project.cache = True
+            mock_project.lc_cache_path = "/tmp/new_cache.db"
+            with patch.object(cache_module.config, "run") as mock_run:
+                mock_run.no_cache = False
+
+                mock_spec = Mock()
+                with patch.object(
+                    cache_module.importlib.util, "find_spec", return_value=mock_spec
+                ):
+                    mock_sqlite_cache = Mock()
+                    mock_set_llm_cache = Mock()
+
+                    with patch.dict(
+                        sys.modules,
+                        {
+                            "langchain": Mock(),
+                            "langchain_community.cache": Mock(
+                                SQLiteCache=mock_sqlite_cache
+                            ),
+                            "langchain.cache": Mock(SQLiteCache=mock_sqlite_cache),
+                            "langchain_core.globals": Mock(
+                                set_llm_cache=mock_set_llm_cache
+                            ),
+                            "langchain.globals": Mock(set_llm_cache=mock_set_llm_cache),
+                        },
+                    ):
+                        with patch("os.path.exists", return_value=False):
+                            with patch.object(cache_module, "logger") as mock_logger:
+                                init_lc_cache()
+
+                                mock_logger.info.assert_called_once()
+                                assert "LangChain cache created at" in str(
+                                    mock_logger.info.call_args
+                                )
+
+    def test_init_lc_cache_without_cache_path(self):
+        """Test that cache is not initialized when cache path is None."""
+        with patch.object(cache_module.config, "project") as mock_project:
+            mock_project.cache = True
+            mock_project.lc_cache_path = None
+            with patch.object(cache_module.config, "run") as mock_run:
+                mock_run.no_cache = False
+
+                mock_spec = Mock()
+                with patch.object(
+                    cache_module.importlib.util, "find_spec", return_value=mock_spec
+                ):
+                    mock_sqlite_cache = Mock()
+                    mock_set_llm_cache = Mock()
+
+                    with patch.dict(
+                        sys.modules,
+                        {
+                            "langchain": Mock(),
+                            "langchain_community.cache": Mock(
+                                SQLiteCache=mock_sqlite_cache
+                            ),
+                            "langchain.cache": Mock(SQLiteCache=mock_sqlite_cache),
+                            "langchain_core.globals": Mock(
+                                set_llm_cache=mock_set_llm_cache
+                            ),
+                            "langchain.globals": Mock(set_llm_cache=mock_set_llm_cache),
+                        },
+                    ):
+                        init_lc_cache()
+
+                        # Should not call SQLiteCache if path is None
+                        mock_sqlite_cache.assert_not_called()
+                        mock_set_llm_cache.assert_not_called()
+
+    def test_init_lc_cache_falls_back_to_legacy_langchain_modules(self):
+        """Test cache initialization falls back to legacy langchain imports."""
+        with patch.object(cache_module.config, "project") as mock_project:
+            mock_project.cache = True
+            mock_project.lc_cache_path = "/tmp/test_cache.db"
+            with patch.object(cache_module.config, "run") as mock_run:
+                mock_run.no_cache = False
+
+                mock_spec = Mock()
+                with patch.object(
+                    cache_module.importlib.util, "find_spec", return_value=mock_spec
+                ):
+                    mock_sqlite_cache = Mock()
+                    mock_set_llm_cache = Mock()
+
+                    def mock_import_module(name):
+                        if name == "langchain_core.globals":
+                            raise ImportError
+                        if name == "langchain_community.cache":
+                            raise ImportError
+                        if name == "langchain.globals":
+                            return Mock(set_llm_cache=mock_set_llm_cache)
+                        if name == "langchain.cache":
+                            return Mock(SQLiteCache=mock_sqlite_cache)
+                        raise ImportError
+
+                    with patch.object(
+                        cache_module.importlib,
+                        "import_module",
+                        side_effect=mock_import_module,
+                    ):
+                        with patch.object(
+                            cache_module.os.path, "exists", return_value=True
+                        ):
+                            init_lc_cache()
+
+                            mock_sqlite_cache.assert_called_once_with(
+                                database_path="/tmp/test_cache.db"
+                            )
+                            mock_set_llm_cache.assert_called_once()
+
+
+class TestCacheEdgeCases:
+    """Test suite for cache edge cases."""
+
+    def setup_method(self):
+        """Clear the cache before each test."""
+        cache_module._cache.clear()
+
+    def teardown_method(self):
+        """Clear the cache after each test."""
+        cache_module._cache.clear()
+
+    def test_cache_with_unhashable_arguments(self):
+        """Test that cache handles unhashable arguments gracefully."""
+
+        @cache
+        def process_list(items):
+            return sum(items)
+
+        # Lists are unhashable and will cause an error
+        with pytest.raises(TypeError):
+            process_list([1, 2, 3])
+
+    def test_cache_with_string_arguments(self):
+        """Test cache with string arguments."""
+        call_count = 0
+
+        @cache
+        def process_string(s):
+            nonlocal call_count
+            call_count += 1
+            return s.upper()
+
+        result1 = process_string("hello")
+        assert result1 == "HELLO"
+        assert call_count == 1
+
+        result2 = process_string("hello")
+        assert result2 == "HELLO"
+        assert call_count == 1  # Cache hit
+
+    def test_cache_with_tuple_arguments(self):
+        """Test cache with tuple arguments."""
+        call_count = 0
+
+        @cache
+        def process_tuple(t):
+            nonlocal call_count
+            call_count += 1
+            return sum(t)
+
+        result1 = process_tuple((1, 2, 3))
+        assert result1 == 6
+        assert call_count == 1
+
+        result2 = process_tuple((1, 2, 3))
+        assert result2 == 6
+        assert call_count == 1  # Cache hit
+
+    def test_cache_with_boolean_arguments(self):
+        """Test cache with boolean arguments."""
+        call_count = 0
+
+        @cache
+        def process_bool(flag):
+            nonlocal call_count
+            call_count += 1
+            return "yes" if flag else "no"
+
+        result1 = process_bool(True)
+        assert result1 == "yes"
+        assert call_count == 1
+
+        result2 = process_bool(True)
+        assert result2 == "yes"
+        assert call_count == 1  # Cache hit
+
+        result3 = process_bool(False)
+        assert result3 == "no"
+        assert call_count == 2
+
+    def test_cache_global_state(self):
+        """Test that cache is global across function calls."""
+
+        @cache
+        def func(x):
+            return x * 2
+
+        func(5)
+        assert len(cache_module._cache) == 1
+
+        func(10)
+        assert len(cache_module._cache) == 2
+
+        func(5)  # Cache hit
+        assert len(cache_module._cache) == 2  # No new entry

backend/tests/test_callbacks.py

@@ -0,0 +1,850 @@
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock, Mock
+
+from chainlit import config
+from chainlit.callbacks import password_auth_callback
+from chainlit.data.base import BaseDataLayer
+from chainlit.types import ThreadDict
+from chainlit.user import User
+
+
+async def test_password_auth_callback(test_config: config.ChainlitConfig):
+    @password_auth_callback
+    async def auth_func(username: str, password: str) -> User | None:
+        if username == "testuser" and password == "testpass":  # nosec B105
+            return User(identifier="testuser")
+        return None
+
+    # Test that the callback is properly registered
+    assert test_config.code.password_auth_callback is not None
+
+    # Test the wrapped function
+    result = await test_config.code.password_auth_callback("testuser", "testpass")
+    assert isinstance(result, User)
+    assert result.identifier == "testuser"
+
+    # Test with incorrect credentials
+    result = await test_config.code.password_auth_callback("wronguser", "wrongpass")
+    assert result is None
+
+
+async def test_header_auth_callback(test_config: config.ChainlitConfig):
+    from starlette.datastructures import Headers
+
+    from chainlit.callbacks import header_auth_callback
+
+    @header_auth_callback
+    async def auth_func(headers: Headers) -> User | None:
+        if headers.get("Authorization") == "Bearer valid_token":
+            return User(identifier="testuser")
+        return None
+
+    # Test that the callback is properly registered
+    assert test_config.code.header_auth_callback is not None
+
+    # Test the wrapped function with valid header
+    valid_headers = Headers({"Authorization": "Bearer valid_token"})
+    result = await test_config.code.header_auth_callback(valid_headers)
+    assert isinstance(result, User)
+    assert result.identifier == "testuser"
+
+    # Test with invalid header
+    invalid_headers = Headers({"Authorization": "Bearer invalid_token"})
+    result = await test_config.code.header_auth_callback(invalid_headers)
+    assert result is None
+
+    # Test with missing header
+    missing_headers = Headers({})
+    result = await test_config.code.header_auth_callback(missing_headers)
+    assert result is 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
+
+    # Mock the get_configured_oauth_providers function
+    with patch(
+        "chainlit.callbacks.get_configured_oauth_providers", return_value=["google"]
+    ):
+
+        @oauth_callback
+        async def auth_func(
+            provider_id: str,
+            token: str,
+            raw_user_data: dict,
+            default_app_user: User,
+            id_token: str | None = None,
+        ) -> User | None:
+            if provider_id == "google" and token == "valid_token":  # nosec B105
+                return User(identifier="oauth_user")
+            return None
+
+        # Test that the callback is properly registered
+        assert test_config.code.oauth_callback is not None
+
+        # Test the wrapped function with valid data
+        result = await test_config.code.oauth_callback(
+            "google", "valid_token", {}, User(identifier="default_user")
+        )
+        assert isinstance(result, User)
+        assert result.identifier == "oauth_user"
+
+        # Test with invalid data
+        result = await test_config.code.oauth_callback(
+            "facebook", "invalid_token", {}, User(identifier="default_user")
+        )
+        assert result is None
+
+
+async def test_on_message(mock_chainlit_context, test_config: config.ChainlitConfig):
+    from chainlit.callbacks import on_message
+    from chainlit.message import Message
+
+    async with mock_chainlit_context:
+        message_received = None
+
+        @on_message
+        async def handle_message(message: Message):
+            nonlocal message_received
+            message_received = message
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_message is not None
+
+        # Create a test message
+        test_message = Message(content="Test message", author="User")
+
+        # Call the registered callback
+        await test_config.code.on_message(test_message)
+
+        # Check that the message was received by our handler
+        assert message_received is not None
+        assert message_received.content == "Test message"
+        assert message_received.author == "User"
+
+
+async def test_on_stop(mock_chainlit_context, test_config: config.ChainlitConfig):
+    from chainlit.callbacks import on_stop
+
+    async with mock_chainlit_context:
+        stop_called = False
+
+        @on_stop
+        async def handle_stop():
+            nonlocal stop_called
+            stop_called = True
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_stop is not None
+
+        # Call the registered callback
+        await test_config.code.on_stop()
+
+        # Check that the stop_called flag was set
+        assert stop_called
+
+
+async def test_action_callback(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.action import Action
+    from chainlit.callbacks import action_callback
+
+    async with mock_chainlit_context:
+        action_handled = False
+
+        @action_callback("test_action")
+        async def handle_action(action: Action):
+            nonlocal action_handled
+            action_handled = True
+            assert action.name == "test_action"
+
+        # Test that the callback is properly registered
+        assert "test_action" in test_config.code.action_callbacks
+
+        # Call the registered callback
+        test_action = Action(name="test_action", payload={"value": "test_value"})
+        await test_config.code.action_callbacks["test_action"](test_action)
+
+        # Check that the action_handled flag was set
+        assert action_handled
+
+
+async def test_on_settings_update(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import on_settings_update
+
+    async with mock_chainlit_context:
+        settings_updated = False
+
+        @on_settings_update
+        async def handle_settings_update(settings: dict):
+            nonlocal settings_updated
+            settings_updated = True
+            assert settings == {"test_setting": "test_value"}
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_settings_update is not None
+
+        # Call the registered callback
+        await test_config.code.on_settings_update({"test_setting": "test_value"})
+
+        # Check that the settings_updated flag was set
+        assert settings_updated
+
+
+async def test_author_rename(test_config: config.ChainlitConfig):
+    from chainlit.callbacks import author_rename
+
+    @author_rename
+    async def rename_author(author: str) -> str:
+        if author == "AI":
+            return "Assistant"
+        return author
+
+    # Test that the callback is properly registered
+    assert test_config.code.author_rename is not None
+
+    # Call the registered callback
+    result = await test_config.code.author_rename("AI")
+    assert result == "Assistant"
+
+    result = await test_config.code.author_rename("Human")
+    assert result == "Human"
+
+    # Test that the callback is properly registered
+    assert test_config.code.author_rename is not None
+
+    # Call the registered callback
+    result = await test_config.code.author_rename("AI")
+    assert result == "Assistant"
+
+    result = await test_config.code.author_rename("Human")
+    assert result == "Human"
+
+
+async def test_on_app_startup(test_config: config.ChainlitConfig):
+    """Test the on_app_startup callback registration and execution for sync and async functions."""
+    from chainlit.callbacks import on_app_startup
+
+    # Test with synchronous function
+    sync_startup_called = False
+
+    @on_app_startup
+    def sync_startup():
+        nonlocal sync_startup_called
+        sync_startup_called = True
+
+    assert test_config.code.on_app_startup is not None, (
+        "Sync startup callback not registered"
+    )
+    # Call the wrapped function (which might be async due to wrap_user_function)
+    result = test_config.code.on_app_startup()
+    if asyncio.iscoroutine(result):
+        await result
+    assert sync_startup_called, "Sync startup function was not called"
+
+    # Reset for async test
+    test_config.code.on_app_startup = None  # Explicitly clear previous registration
+
+    # Test with asynchronous function
+    async_startup_called = False
+
+    @on_app_startup
+    async def async_startup():
+        nonlocal async_startup_called
+        await asyncio.sleep(0)  # Simulate async work
+        async_startup_called = True
+
+    assert test_config.code.on_app_startup is not None, (
+        "Async startup callback not registered"
+    )
+    # Call the wrapped function (which should be async)
+    result = test_config.code.on_app_startup()
+    assert asyncio.iscoroutine(result), (
+        "Async startup function did not return a coroutine"
+    )
+    await result
+    assert async_startup_called, "Async startup function was not called"
+
+
+async def test_on_app_shutdown(test_config: config.ChainlitConfig):
+    """Test the on_app_shutdown callback registration and execution for sync and async functions."""
+    from chainlit.callbacks import on_app_shutdown
+
+    # Test with synchronous function
+    sync_shutdown_called = False
+
+    @on_app_shutdown
+    def sync_shutdown():
+        nonlocal sync_shutdown_called
+        sync_shutdown_called = True
+
+    assert test_config.code.on_app_shutdown is not None, (
+        "Sync shutdown callback not registered"
+    )
+    # Call the wrapped function
+    result = test_config.code.on_app_shutdown()
+    if asyncio.iscoroutine(result):
+        await result
+    assert sync_shutdown_called, "Sync shutdown function was not called"
+
+    # Reset for async test
+    test_config.code.on_app_shutdown = None  # Explicitly clear previous registration
+
+    # Test with asynchronous function
+    async_shutdown_called = False
+
+    @on_app_shutdown
+    async def async_shutdown():
+        nonlocal async_shutdown_called
+        await asyncio.sleep(0)  # Simulate async work
+        async_shutdown_called = True
+
+    assert test_config.code.on_app_shutdown is not None, (
+        "Async shutdown callback not registered"
+    )
+    # Call the wrapped function
+    result = test_config.code.on_app_shutdown()
+    assert asyncio.iscoroutine(result), (
+        "Async shutdown function did not return a coroutine"
+    )
+    await result
+    assert async_shutdown_called, "Async shutdown function was not called"
+
+
+async def test_on_chat_start(mock_chainlit_context, test_config: config.ChainlitConfig):
+    from chainlit.callbacks import on_chat_start
+
+    async with mock_chainlit_context:
+        chat_started = False
+
+        @on_chat_start
+        async def handle_chat_start():
+            nonlocal chat_started
+            chat_started = True
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_chat_start is not None
+
+        # Call the registered callback
+        await test_config.code.on_chat_start()
+
+        # Check that the chat_started flag was set
+        assert chat_started
+
+
+async def test_on_chat_resume(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import on_chat_resume
+
+    async with mock_chainlit_context:
+        chat_resumed = False
+
+        @on_chat_resume
+        async def handle_chat_resume(thread: ThreadDict):
+            nonlocal chat_resumed
+            chat_resumed = True
+            assert thread["id"] == "test_thread_id"
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_chat_resume is not None
+
+        # Call the registered callback
+        await test_config.code.on_chat_resume(
+            {
+                "id": "test_thread_id",
+                "createdAt": "2023-01-01T00:00:00Z",
+                "name": "Test Thread",
+                "userId": "test_user_id",
+                "userIdentifier": "test_user",
+                "tags": [],
+                "metadata": {},
+                "steps": [],
+                "elements": [],
+            }
+        )
+
+        # Check that the chat_resumed flag was set
+        assert chat_resumed
+
+
+async def test_set_chat_profiles(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import set_chat_profiles
+    from chainlit.types import ChatProfile
+
+    async with mock_chainlit_context:
+
+        @set_chat_profiles
+        async def get_chat_profiles(user, language):
+            return [
+                ChatProfile(name="Test Profile", markdown_description="A test profile")
+            ]
+
+        # Test that the callback is properly registered
+        assert test_config.code.set_chat_profiles is not None
+
+        # Call the registered callback
+        result = await test_config.code.set_chat_profiles(None, None)
+
+        # Check the result
+        assert result is not None
+        assert isinstance(result, list)
+        assert len(result) == 1
+        assert isinstance(result[0], ChatProfile)
+        assert result[0].name == "Test Profile"
+        assert result[0].markdown_description == "A test profile"
+
+
+async def test_set_chat_profiles_language(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import set_chat_profiles
+    from chainlit.types import ChatProfile
+
+    async with mock_chainlit_context:
+
+        @set_chat_profiles
+        async def get_chat_profiles(user, language):
+            if language == "fr-CA":
+                return [
+                    ChatProfile(
+                        name="Profil de test", markdown_description="Un profil de test"
+                    )
+                ]
+
+            return [
+                ChatProfile(name="Test Profile", markdown_description="A test profile")
+            ]
+
+        # Test that the callback is properly registered
+        assert test_config.code.set_chat_profiles is not None
+
+        # Call the registered callback
+        result = await test_config.code.set_chat_profiles(None, "fr-CA")
+
+        # Check the result
+        assert result is not None
+        assert isinstance(result, list)
+        assert len(result) == 1
+        assert isinstance(result[0], ChatProfile)
+        assert result[0].name == "Profil de test"
+        assert result[0].markdown_description == "Un profil de test"
+
+
+async def test_set_starters(mock_chainlit_context, test_config: config.ChainlitConfig):
+    from chainlit.callbacks import set_starters
+    from chainlit.types import Starter
+
+    async with mock_chainlit_context:
+
+        @set_starters
+        async def get_starters(user):
+            return [
+                Starter(
+                    label="Test Label",
+                    message="Test Message",
+                )
+            ]
+
+        # Test that the callback is properly registered
+        assert test_config.code.set_starters is not None
+
+        # Call the registered callback
+        result = await test_config.code.set_starters(None, None)
+
+        # Check the result
+        assert result is not None
+        assert isinstance(result, list)
+        assert len(result) == 1
+        assert isinstance(result[0], Starter)
+        assert result[0].label == "Test Label"
+        assert result[0].message == "Test Message"
+
+
+async def test_set_starters_language(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import set_starters
+    from chainlit.types import Starter
+
+    async with mock_chainlit_context:
+
+        @set_starters
+        async def get_starters(user, language):
+            if language == "fr-CA":
+                return [
+                    Starter(
+                        label="Étiquette de test",
+                        message="Message de test",
+                    )
+                ]
+
+            return [
+                Starter(
+                    label="Test Label",
+                    message="Test Message",
+                )
+            ]
+
+        # Test that the callback is properly registered
+        assert test_config.code.set_starters is not None
+
+        # Call the registered callback
+        result = await test_config.code.set_starters(None, "fr-CA")
+
+        # Check the result
+        assert result is not None
+        assert isinstance(result, list)
+        assert len(result) == 1
+        assert isinstance(result[0], Starter)
+        assert result[0].label == "Étiquette de test"
+        assert result[0].message == "Message de test"
+
+
+async def test_set_starter_categories(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import set_starter_categories
+    from chainlit.types import Starter, StarterCategory
+
+    async with mock_chainlit_context:
+
+        @set_starter_categories
+        async def get_starter_categories(user, language):
+            return [
+                StarterCategory(
+                    label="Creative",
+                    icon="https://example.com/creative.png",
+                    starters=[
+                        Starter(label="Write a poem", message="Write a poem"),
+                        Starter(label="Write a story", message="Write a story"),
+                    ],
+                ),
+                StarterCategory(
+                    label="Educational",
+                    starters=[
+                        Starter(label="Explain concept", message="Explain it"),
+                    ],
+                ),
+            ]
+
+        assert test_config.code.set_starter_categories is not None
+
+        result = await test_config.code.set_starter_categories(None, None, None)
+
+        assert result is not None
+        assert isinstance(result, list)
+        assert len(result) == 2
+
+        assert result[0].label == "Creative"
+        assert result[0].icon == "https://example.com/creative.png"
+        assert len(result[0].starters) == 2
+        assert result[0].starters[0].label == "Write a poem"
+
+        assert result[1].label == "Educational"
+        assert result[1].icon is None
+        assert len(result[1].starters) == 1
+
+        category_dict = result[0].to_dict()
+        assert category_dict["label"] == "Creative"
+        assert category_dict["icon"] == "https://example.com/creative.png"
+        starters_list = category_dict["starters"]
+        assert isinstance(starters_list, list)
+        assert len(starters_list) == 2
+
+
+async def test_set_starter_categories_with_chat_profile(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import set_starter_categories
+    from chainlit.types import Starter, StarterCategory
+
+    async with mock_chainlit_context:
+
+        @set_starter_categories
+        async def get_starter_categories(user, language, chat_profile):
+            return [
+                StarterCategory(
+                    label="Profile Specific",
+                    icon="https://example.com/profile.png",
+                    starters=[
+                        Starter(
+                            label=f"Starter for {chat_profile}",
+                            message=f"Message for {chat_profile}",
+                        ),
+                    ],
+                ),
+            ]
+
+        assert test_config.code.set_starter_categories is not None
+
+        result = await test_config.code.set_starter_categories(
+            None, "en", "test-profile"
+        )
+
+        assert result is not None
+        assert isinstance(result, list)
+        assert len(result) == 1
+
+        assert result[0].label == "Profile Specific"
+        assert result[0].icon == "https://example.com/profile.png"
+        assert len(result[0].starters) == 1
+        assert result[0].starters[0].label == "Starter for test-profile"
+        assert result[0].starters[0].message == "Message for test-profile"
+
+
+async def test_on_shared_thread_view_allow(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import on_shared_thread_view
+    from chainlit.user import User
+
+    async with mock_chainlit_context:
+        # Simulate a viewer with access to certain chat profiles
+        allowed_profiles_by_user = {"viewer": {"pro", "basic"}}
+
+        @on_shared_thread_view
+        async def allow_shared_view(thread, viewer: User | None):
+            md = thread.get("metadata") or {}
+            chat_profile = (md or {}).get("chat_profile")
+            if not md.get("is_shared"):
+                return False
+            if not viewer:
+                return False
+            return chat_profile in allowed_profiles_by_user.get(
+                viewer.identifier, set()
+            )
+
+        assert test_config.code.on_shared_thread_view is not None
+
+        thread: ThreadDict = {
+            "id": "t1",
+            "createdAt": "2025-09-03T00:00:00Z",
+            "name": "Shared Thread",
+            "userId": "author_id",
+            "userIdentifier": "author",
+            "tags": [],
+            "metadata": {"is_shared": True, "chat_profile": "pro"},
+            "steps": [],
+            "elements": [],
+        }
+        viewer = User(identifier="viewer")
+
+        res = await test_config.code.on_shared_thread_view(thread, viewer)
+        assert res is True
+
+
+async def test_on_shared_thread_view_block_and_exception(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    from chainlit.callbacks import on_shared_thread_view
+    from chainlit.user import User
+
+    async with mock_chainlit_context:
+        # Case 1: Explicitly return False when profile not allowed
+        @on_shared_thread_view
+        async def deny_when_not_allowed(thread, viewer: User | None):
+            md = thread.get("metadata") or {}
+            return md.get("chat_profile") == "allowed"
+
+        assert test_config.code.on_shared_thread_view is not None
+
+        thread: ThreadDict = {
+            "id": "t2",
+            "createdAt": "2025-09-03T00:00:00Z",
+            "name": "Shared Thread",
+            "userId": "author_id",
+            "userIdentifier": "author",
+            "tags": [],
+            "metadata": {"is_shared": True, "chat_profile": "restricted"},
+            "steps": [],
+            "elements": [],
+        }
+        viewer = User(identifier="viewer")
+        res = await test_config.code.on_shared_thread_view(thread, viewer)
+        assert not res
+
+        # Case 2: Raise an exception inside callback; wrapper should swallow and result should be falsy
+        @on_shared_thread_view
+        async def raise_on_forbidden(thread, viewer: User | None):
+            md = thread.get("metadata") or {}
+            if md.get("chat_profile") == "forbidden":
+                raise ValueError("Viewer not allowed for this profile")
+            return True
+
+        assert test_config.code.on_shared_thread_view is not None
+
+        thread_err: ThreadDict = {
+            "id": "t3",
+            "createdAt": "2025-09-03T00:00:00Z",
+            "name": "Shared Thread",
+            "userId": "author_id",
+            "userIdentifier": "author",
+            "tags": [],
+            "metadata": {"is_shared": True, "chat_profile": "forbidden"},
+            "steps": [],
+            "elements": [],
+        }
+        res2 = await test_config.code.on_shared_thread_view(thread_err, viewer)
+        assert not res2
+
+
+async def test_on_chat_end(mock_chainlit_context, test_config: config.ChainlitConfig):
+    from chainlit.callbacks import on_chat_end
+
+    async with mock_chainlit_context:
+        chat_ended = False
+
+        @on_chat_end
+        async def handle_chat_end():
+            nonlocal chat_ended
+            chat_ended = True
+
+        # Test that the callback is properly registered
+        assert test_config.code.on_chat_end is not None
+
+        # Call the registered callback
+        await test_config.code.on_chat_end()
+
+        # Check that the chat_ended flag was set
+        assert chat_ended
+
+
+def test_data_layer_config(
+    mock_data_layer: AsyncMock,
+    test_config: config.ChainlitConfig,
+    mock_get_data_layer: Mock,
+):
+    """Test whether we can properly configure a data layer."""
+
+    # Test that the callback is properly registered
+    assert test_config.code.data_layer is not None
+
+    # Call the registered callback
+    result = test_config.code.data_layer()
+
+    # Check that the result is an instance of MockDataLayer
+    assert isinstance(result, BaseDataLayer)
+
+    mock_get_data_layer.assert_called_once()
+
+
+def test_chat_profile_with_config_overrides():
+    """Test that ChatProfile can be created with config_overrides."""
+    from chainlit.config import (
+        ChainlitConfigOverrides,
+        FeaturesSettings,
+        McpFeature,
+        UISettings,
+    )
+    from chainlit.types import ChatProfile
+
+    # Test creating a profile without config_overrides
+    basic_profile = ChatProfile(
+        name="Basic Profile", markdown_description="A basic profile without overrides"
+    )
+    assert basic_profile.config_overrides is None
+
+    # Test creating a profile with config_overrides
+    config_overrides = ChainlitConfigOverrides(
+        features=FeaturesSettings(mcp=McpFeature(enabled=True)),
+        ui=UISettings(
+            name="Custom App Name",
+            description="Custom description",
+            default_theme="light",
+        ),
+    )
+
+    profile_with_overrides = ChatProfile(
+        name="MCP Profile",
+        markdown_description="A profile with MCP enabled",
+        config_overrides=config_overrides,
+    )
+
+    # Verify the profile was created successfully
+    assert profile_with_overrides.name == "MCP Profile"
+    assert profile_with_overrides.config_overrides is not None
+    assert profile_with_overrides.config_overrides.features.mcp.enabled is True
+    assert profile_with_overrides.config_overrides.ui.name == "Custom App Name"
+    assert profile_with_overrides.config_overrides.ui.default_theme == "light"
+
+
+async def test_set_chat_profiles_with_config_overrides(
+    mock_chainlit_context, test_config: config.ChainlitConfig
+):
+    """Test that set_chat_profiles callback works with profiles that have config_overrides."""
+    from chainlit.callbacks import set_chat_profiles
+    from chainlit.config import (
+        ChainlitConfigOverrides,
+        FeaturesSettings,
+        McpFeature,
+        UISettings,
+    )
+    from chainlit.types import ChatProfile
+
+    async with mock_chainlit_context:
+
+        @set_chat_profiles
+        async def get_chat_profiles(user, language):
+            return [
+                ChatProfile(
+                    name="Basic Profile",
+                    markdown_description="A basic profile without overrides",
+                ),
+                ChatProfile(
+                    name="MCP Profile",
+                    markdown_description="A profile with MCP enabled",
+                    config_overrides=ChainlitConfigOverrides(
+                        features=FeaturesSettings(mcp=McpFeature(enabled=True)),
+                        ui=UISettings(name="MCP Assistant", default_theme="dark"),
+                    ),
+                ),
+                ChatProfile(
+                    name="Light Theme Profile",
+                    markdown_description="A profile with light theme",
+                    config_overrides=ChainlitConfigOverrides(
+                        ui=UISettings(name="Light Theme App", default_theme="light")
+                    ),
+                ),
+            ]
+
+        # Test that the callback is properly registered
+        assert test_config.code.set_chat_profiles is not None
+
+        # Call the registered callback
+        result = await test_config.code.set_chat_profiles(None, None)
+
+        # Check the result
+        assert result is not None
+        assert isinstance(result, list)
+        assert len(result) == 3
+
+        # Test basic profile
+        basic_profile = result[0]
+        assert basic_profile.name == "Basic Profile"
+        assert basic_profile.config_overrides is None
+
+        # Test MCP profile
+        mcp_profile = result[1]
+        assert mcp_profile.name == "MCP Profile"
+        assert mcp_profile.config_overrides is not None
+        assert mcp_profile.config_overrides.features.mcp.enabled is True
+        assert mcp_profile.config_overrides.ui.name == "MCP Assistant"
+        assert mcp_profile.config_overrides.ui.default_theme == "dark"
+
+        # Test light theme profile
+        light_profile = result[2]
+        assert light_profile.name == "Light Theme Profile"
+        assert light_profile.config_overrides is not None
+        assert light_profile.config_overrides.ui.name == "Light Theme App"
+        assert light_profile.config_overrides.ui.default_theme == "light"

backend/tests/test_chat_context.py

@@ -0,0 +1,459 @@
+import asyncio
+from contextlib import contextmanager
+from unittest.mock import Mock, patch
+
+from chainlit.chat_context import chat_context, chat_contexts
+from chainlit.context import ChainlitContext, context_var
+
+
+@contextmanager
+def mock_chainlit_context(session=None):
+    """Context manager to set up and tear down Chainlit context."""
+    # Mock the event loop since we're not in an async context
+    mock_loop = Mock(spec=asyncio.AbstractEventLoop)
+
+    with patch("asyncio.get_running_loop", return_value=mock_loop):
+        mock_context = ChainlitContext(session=session)
+        token = context_var.set(mock_context)
+        try:
+            yield mock_context
+        finally:
+            context_var.reset(token)
+
+
+class TestChatContext:
+    """Test suite for ChatContext class."""
+
+    def setup_method(self):
+        """Clear chat_contexts before each test."""
+        chat_contexts.clear()
+
+    def teardown_method(self):
+        """Clear chat_contexts after each test."""
+        chat_contexts.clear()
+
+    def test_get_without_session(self):
+        """Test get returns empty list when no session exists."""
+        with mock_chainlit_context(session=None):
+            result = chat_context.get()
+            assert result == []
+
+    def test_get_with_new_session(self):
+        """Test get creates new chat context for new session."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.get()
+
+            assert result == []
+            assert "session_123" in chat_contexts
+            assert chat_contexts["session_123"] == []
+
+    def test_get_returns_copy(self):
+        """Test get returns a copy of the chat context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.get()
+
+            assert result == [mock_message]
+            # Verify it's a copy, not the original
+            assert result is not chat_contexts["session_123"]
+
+    def test_get_with_existing_messages(self):
+        """Test get returns existing messages."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.get()
+
+            assert len(result) == 2
+            assert mock_msg1 in result
+            assert mock_msg2 in result
+
+    def test_add_without_session(self):
+        """Test add does nothing when no session exists."""
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=None):
+            result = chat_context.add(mock_message)
+
+            assert result is None
+            assert len(chat_contexts) == 0
+
+    def test_add_with_new_session(self):
+        """Test add creates new chat context and adds message."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_message)
+
+            assert result == mock_message
+            assert "session_123" in chat_contexts
+            assert mock_message in chat_contexts["session_123"]
+
+    def test_add_message_to_existing_context(self):
+        """Test add appends message to existing context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_msg2)
+
+            assert result == mock_msg2
+            assert len(chat_contexts["session_123"]) == 2
+            assert mock_msg1 in chat_contexts["session_123"]
+            assert mock_msg2 in chat_contexts["session_123"]
+
+    def test_add_duplicate_message(self):
+        """Test add does not add duplicate messages."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_message)
+
+            assert result == mock_message
+            assert len(chat_contexts["session_123"]) == 1
+
+    def test_remove_without_session(self):
+        """Test remove returns False when no session exists."""
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=None):
+            result = chat_context.remove(mock_message)
+
+            assert result is False
+
+    def test_remove_with_nonexistent_context(self):
+        """Test remove returns False when context doesn't exist."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.remove(mock_message)
+
+            assert result is False
+
+    def test_remove_nonexistent_message(self):
+        """Test remove returns False when message not in context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.remove(mock_msg2)
+
+            assert result is False
+            assert mock_msg1 in chat_contexts["session_123"]
+
+    def test_remove_existing_message(self):
+        """Test remove successfully removes message."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.remove(mock_msg1)
+
+            assert result is True
+            assert mock_msg1 not in chat_contexts["session_123"]
+            assert mock_msg2 in chat_contexts["session_123"]
+            assert len(chat_contexts["session_123"]) == 1
+
+    def test_clear_without_session(self):
+        """Test clear does nothing when no session exists."""
+        chat_contexts["session_123"] = [Mock()]
+
+        with mock_chainlit_context(session=None):
+            chat_context.clear()
+
+            # Original context should remain
+            assert "session_123" in chat_contexts
+
+    def test_clear_with_nonexistent_context(self):
+        """Test clear does nothing when context doesn't exist."""
+        mock_session = Mock()
+        mock_session.id = "session_456"
+
+        chat_contexts["session_123"] = [Mock()]
+
+        with mock_chainlit_context(session=mock_session):
+            chat_context.clear()
+
+            # Original context should remain
+            assert "session_123" in chat_contexts
+
+    def test_clear_existing_context(self):
+        """Test clear empties existing context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2]
+
+        with mock_chainlit_context(session=mock_session):
+            chat_context.clear()
+
+            assert "session_123" in chat_contexts
+            assert chat_contexts["session_123"] == []
+
+    def test_to_openai_with_assistant_message(self):
+        """Test to_openai converts assistant messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "assistant_message"
+        mock_message.content = "Hello, how can I help?"
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {
+                "role": "assistant",
+                "content": "Hello, how can I help?",
+            }
+
+    def test_to_openai_with_user_message(self):
+        """Test to_openai converts user messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "user_message"
+        mock_message.content = "What is the weather?"
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {"role": "user", "content": "What is the weather?"}
+
+    def test_to_openai_with_system_message(self):
+        """Test to_openai converts system messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "system_message"
+        mock_message.content = "You are a helpful assistant."
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {
+                "role": "system",
+                "content": "You are a helpful assistant.",
+            }
+
+    def test_to_openai_with_unknown_message_type(self):
+        """Test to_openai treats unknown types as system messages."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_message = Mock()
+        mock_message.type = "unknown_type"
+        mock_message.content = "Unknown message"
+
+        chat_contexts["session_123"] = [mock_message]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 1
+            assert result[0] == {"role": "system", "content": "Unknown message"}
+
+    def test_to_openai_with_multiple_messages(self):
+        """Test to_openai converts multiple messages correctly."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg1.type = "user_message"
+        mock_msg1.content = "Hello"
+
+        mock_msg2 = Mock()
+        mock_msg2.type = "assistant_message"
+        mock_msg2.content = "Hi there!"
+
+        mock_msg3 = Mock()
+        mock_msg3.type = "user_message"
+        mock_msg3.content = "How are you?"
+
+        chat_contexts["session_123"] = [mock_msg1, mock_msg2, mock_msg3]
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 3
+            assert result[0] == {"role": "user", "content": "Hello"}
+            assert result[1] == {"role": "assistant", "content": "Hi there!"}
+            assert result[2] == {"role": "user", "content": "How are you?"}
+
+    def test_to_openai_with_empty_context(self):
+        """Test to_openai returns empty list for empty context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        chat_contexts["session_123"] = []
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert result == []
+
+    def test_to_openai_without_session(self):
+        """Test to_openai returns empty list when no session exists."""
+        with mock_chainlit_context(session=None):
+            result = chat_context.to_openai()
+
+            assert result == []
+
+
+class TestChatContextEdgeCases:
+    """Test suite for chat_context edge cases."""
+
+    def setup_method(self):
+        """Clear chat_contexts before each test."""
+        chat_contexts.clear()
+
+    def teardown_method(self):
+        """Clear chat_contexts after each test."""
+        chat_contexts.clear()
+
+    def test_multiple_sessions_isolated(self):
+        """Test that different sessions have isolated contexts."""
+        mock_session1 = Mock()
+        mock_session1.id = "session_1"
+
+        mock_session2 = Mock()
+        mock_session2.id = "session_2"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+
+        with mock_chainlit_context(session=mock_session1):
+            chat_context.add(mock_msg1)
+
+        with mock_chainlit_context(session=mock_session2):
+            chat_context.add(mock_msg2)
+
+        assert len(chat_contexts) == 2
+        assert mock_msg1 in chat_contexts["session_1"]
+        assert mock_msg2 in chat_contexts["session_2"]
+        assert mock_msg1 not in chat_contexts["session_2"]
+        assert mock_msg2 not in chat_contexts["session_1"]
+
+    def test_add_then_remove_then_add_again(self):
+        """Test adding, removing, and re-adding the same message."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            # Add
+            chat_context.add(mock_message)
+            assert len(chat_contexts["session_123"]) == 1
+
+            # Remove
+            result = chat_context.remove(mock_message)
+            assert result is True
+            assert len(chat_contexts["session_123"]) == 0
+
+            # Add again
+            chat_context.add(mock_message)
+            assert len(chat_contexts["session_123"]) == 1
+
+    def test_clear_then_add(self):
+        """Test adding messages after clearing context."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        mock_msg1 = Mock()
+        mock_msg2 = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            chat_context.add(mock_msg1)
+            chat_context.clear()
+            chat_context.add(mock_msg2)
+
+            result = chat_context.get()
+            assert len(result) == 1
+            assert mock_msg2 in result
+            assert mock_msg1 not in result
+
+    def test_to_openai_with_mixed_message_types(self):
+        """Test to_openai with various message types in sequence."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+
+        messages = [
+            Mock(type="system_message", content="System prompt"),
+            Mock(type="user_message", content="User query"),
+            Mock(type="assistant_message", content="Assistant response"),
+            Mock(type="other_type", content="Other message"),
+        ]
+
+        chat_contexts["session_123"] = messages
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.to_openai()
+
+            assert len(result) == 4
+            assert result[0]["role"] == "system"
+            assert result[1]["role"] == "user"
+            assert result[2]["role"] == "assistant"
+            assert result[3]["role"] == "system"  # Unknown types default to system
+
+    def test_chat_context_singleton(self):
+        """Test that chat_context is a singleton instance."""
+        from chainlit.chat_context import chat_context as imported_context
+
+        assert chat_context is imported_context
+
+    def test_add_returns_message(self):
+        """Test that add returns the message for chaining."""
+        mock_session = Mock()
+        mock_session.id = "session_123"
+        mock_message = Mock()
+
+        with mock_chainlit_context(session=mock_session):
+            result = chat_context.add(mock_message)
+
+            assert result is mock_message

backend/tests/test_chat_settings.py

@@ -0,0 +1,343 @@
+import pytest
+
+from chainlit.chat_settings import ChatSettings
+from chainlit.input_widget import (
+    Checkbox,
+    NumberInput,
+    Select,
+    Slider,
+    Switch,
+    Tab,
+    TextInput,
+)
+
+
+@pytest.mark.asyncio
+class TestChatSettings:
+    """Test suite for ChatSettings class."""
+
+    async def test_chat_settings_initialization(self, mock_chainlit_context):
+        """Test ChatSettings initialization with input widgets."""
+        async with mock_chainlit_context:
+            switch = Switch(id="enable_feature", label="Enable Feature")
+            slider = Slider(id="temperature", label="Temperature", initial=0.7)
+
+            settings = ChatSettings(inputs=[switch, slider])
+
+            assert len(settings.inputs) == 2
+            assert settings.inputs[0] == switch
+            assert settings.inputs[1] == slider
+
+    async def test_chat_settings_with_empty_inputs(self, mock_chainlit_context):
+        """Test ChatSettings with empty inputs list."""
+        async with mock_chainlit_context:
+            settings = ChatSettings(inputs=[])
+
+            assert settings.inputs == []
+
+    async def test_chat_settings_settings_method(self, mock_chainlit_context):
+        """Test ChatSettings.settings() method returns initial values."""
+        async with mock_chainlit_context:
+            switch = Switch(id="enable_feature", label="Enable", initial=True)
+            slider = Slider(id="temperature", label="Temperature", initial=0.7)
+            text = TextInput(id="model", label="Model", initial="gpt-4")
+
+            settings = ChatSettings(inputs=[switch, slider, text])
+            result = settings.settings()
+
+            assert result["enable_feature"] is True
+            assert result["temperature"] == 0.7
+            assert result["model"] == "gpt-4"
+
+    async def test_chat_settings_with_tabs(self, mock_chainlit_context):
+        """Test ChatSettings with Tab containers."""
+        async with mock_chainlit_context:
+            # Create inputs for tabs
+            switch1 = Switch(id="switch1", label="Switch 1", initial=True)
+            slider1 = Slider(id="slider1", label="Slider 1", initial=5)
+
+            switch2 = Switch(id="switch2", label="Switch 2", initial=False)
+            text2 = TextInput(id="text2", label="Text 2", initial="value")
+
+            # Create tabs
+            tab1 = Tab(id="tab1", label="Tab 1", inputs=[switch1, slider1])
+            tab2 = Tab(id="tab2", label="Tab 2", inputs=[switch2, text2])
+
+            settings = ChatSettings(inputs=[tab1, tab2])
+
+            assert len(settings.inputs) == 2
+            assert isinstance(settings.inputs[0], Tab)
+            assert isinstance(settings.inputs[1], Tab)
+
+    async def test_chat_settings_settings_with_tabs(self, mock_chainlit_context):
+        """Test ChatSettings.settings() collects values from tabs."""
+        async with mock_chainlit_context:
+            switch1 = Switch(id="switch1", label="Switch 1", initial=True)
+            slider1 = Slider(id="slider1", label="Slider 1", initial=5)
+
+            switch2 = Switch(id="switch2", label="Switch 2", initial=False)
+            text2 = TextInput(id="text2", label="Text 2", initial="value")
+
+            tab1 = Tab(id="tab1", label="Tab 1", inputs=[switch1, slider1])
+            tab2 = Tab(id="tab2", label="Tab 2", inputs=[switch2, text2])
+
+            settings = ChatSettings(inputs=[tab1, tab2])
+            result = settings.settings()
+
+            # Should collect all inputs from all tabs
+            assert result["switch1"] is True
+            assert result["slider1"] == 5
+            assert result["switch2"] is False
+            assert result["text2"] == "value"
+
+    async def test_chat_settings_send(self, mock_chainlit_context):
+        """Test ChatSettings.send() method."""
+        async with mock_chainlit_context as ctx:
+            switch = Switch(id="enable", label="Enable", initial=True)
+            slider = Slider(id="temp", label="Temperature", initial=0.8)
+
+            settings = ChatSettings(inputs=[switch, slider])
+            result = await settings.send()
+
+            # Verify settings were returned
+            assert result["enable"] is True
+            assert result["temp"] == 0.8
+
+            # Verify emitter methods were called
+            ctx.emitter.set_chat_settings.assert_called_once_with(result)
+            ctx.emitter.emit.assert_called_once()
+
+            # Verify emit was called with correct arguments
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "chat_settings"
+            assert len(call_args[0][1]) == 2  # Two inputs
+
+    async def test_chat_settings_refresh(self, mock_chainlit_context):
+        """Test ChatSettings.refresh() emits UI without updating session."""
+        async with mock_chainlit_context as ctx:
+            switch = Switch(id="enable", label="Enable", initial=True)
+            slider = Slider(id="temp", label="Temperature", initial=0.8)
+
+            settings = ChatSettings(inputs=[switch, slider])
+            result = await settings.refresh()
+
+            assert result["enable"] is True
+            assert result["temp"] == 0.8
+
+            ctx.emitter.set_chat_settings.assert_not_called()
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "chat_settings"
+            assert len(call_args[0][1]) == 2
+
+    async def test_chat_settings_send_with_tabs(self, mock_chainlit_context):
+        """Test ChatSettings.send() with tabs."""
+        async with mock_chainlit_context as ctx:
+            switch = Switch(id="switch1", label="Switch", initial=True)
+            slider = Slider(id="slider1", label="Slider", initial=5)
+
+            tab = Tab(id="tab1", label="Settings", inputs=[switch, slider])
+            settings = ChatSettings(inputs=[tab])
+            result = await settings.send()
+
+            # Verify settings collected from tab
+            assert result["switch1"] is True
+            assert result["slider1"] == 5
+
+            # Verify emitter was called
+            ctx.emitter.set_chat_settings.assert_called_once()
+            ctx.emitter.emit.assert_called_once()
+
+    async def test_chat_settings_with_all_widget_types(self, mock_chainlit_context):
+        """Test ChatSettings with all widget types."""
+        async with mock_chainlit_context:
+            widgets = [
+                Switch(id="switch", label="Switch", initial=True),
+                Slider(id="slider", label="Slider", initial=5, min=0, max=10),
+                Select(
+                    id="select",
+                    label="Select",
+                    values=["a", "b", "c"],
+                    initial_index=1,
+                ),
+                TextInput(id="text", label="Text", initial="hello"),
+                NumberInput(id="number", label="Number", initial=42.0),
+                Checkbox(id="checkbox", label="Checkbox", initial=False),
+            ]
+
+            settings = ChatSettings(inputs=widgets)
+            result = settings.settings()
+
+            assert result["switch"] is True
+            assert result["slider"] == 5
+            assert result["select"] == "b"
+            assert result["text"] == "hello"
+            assert result["number"] == 42.0
+            assert result["checkbox"] is False
+
+    async def test_chat_settings_with_nested_tabs(self, mock_chainlit_context):
+        """Test ChatSettings.settings() with nested structure."""
+        async with mock_chainlit_context:
+            # Create multiple tabs with different inputs
+            tab1_inputs = [
+                Switch(id="t1_switch", label="T1 Switch", initial=True),
+                Slider(id="t1_slider", label="T1 Slider", initial=3),
+            ]
+
+            tab2_inputs = [
+                TextInput(id="t2_text", label="T2 Text", initial="test"),
+                Checkbox(id="t2_check", label="T2 Check", initial=True),
+            ]
+
+            tab1 = Tab(id="tab1", label="Tab 1", inputs=tab1_inputs)
+            tab2 = Tab(id="tab2", label="Tab 2", inputs=tab2_inputs)
+
+            settings = ChatSettings(inputs=[tab1, tab2])
+            result = settings.settings()
+
+            # All inputs from all tabs should be collected
+            assert result["t1_switch"] is True
+            assert result["t1_slider"] == 3
+            assert result["t2_text"] == "test"
+            assert result["t2_check"] is True
+            assert len(result) == 4
+
+    async def test_chat_settings_only_widgets_or_only_tabs(self, mock_chainlit_context):
+        """Test that ChatSettings accepts either all widgets or all tabs, not mixed."""
+        async with mock_chainlit_context:
+            # Test with only widgets
+            widgets = [
+                Switch(id="switch", label="Switch", initial=True),
+                Slider(id="slider", label="Slider", initial=7),
+            ]
+            settings_widgets = ChatSettings(inputs=widgets)
+            result_widgets = settings_widgets.settings()
+            assert result_widgets["switch"] is True
+            assert result_widgets["slider"] == 7
+
+            # Test with only tabs
+            tab1 = Tab(
+                id="tab1",
+                label="Tab 1",
+                inputs=[Switch(id="t1_switch", label="Switch", initial=False)],
+            )
+            tab2 = Tab(
+                id="tab2",
+                label="Tab 2",
+                inputs=[Slider(id="t2_slider", label="Slider", initial=3)],
+            )
+            settings_tabs = ChatSettings(inputs=[tab1, tab2])
+            result_tabs = settings_tabs.settings()
+            assert result_tabs["t1_switch"] is False
+            assert result_tabs["t2_slider"] == 3
+
+    async def test_chat_settings_with_none_initial_values(self, mock_chainlit_context):
+        """Test ChatSettings with widgets having None initial values."""
+        async with mock_chainlit_context:
+            text = TextInput(id="text", label="Text", initial=None)
+            number = NumberInput(id="number", label="Number", initial=None)
+            select = Select(
+                id="select", label="Select", values=["a", "b"], initial_value=None
+            )
+
+            settings = ChatSettings(inputs=[text, number, select])
+            result = settings.settings()
+
+            assert result["text"] is None
+            assert result["number"] is None
+            assert result["select"] is None
+
+
+@pytest.mark.asyncio
+class TestChatSettingsEdgeCases:
+    """Test suite for ChatSettings edge cases."""
+
+    async def test_chat_settings_empty_tabs(self, mock_chainlit_context):
+        """Test ChatSettings with empty tabs."""
+        async with mock_chainlit_context:
+            empty_tab = Tab(id="empty", label="Empty Tab", inputs=[])
+            settings = ChatSettings(inputs=[empty_tab])
+            result = settings.settings()
+
+            assert result == {}
+
+    async def test_chat_settings_duplicate_ids(self, mock_chainlit_context):
+        """Test ChatSettings behavior with duplicate IDs (last one wins)."""
+        async with mock_chainlit_context:
+            switch1 = Switch(id="duplicate", label="Switch 1", initial=True)
+            switch2 = Switch(id="duplicate", label="Switch 2", initial=False)
+
+            settings = ChatSettings(inputs=[switch1, switch2])
+            result = settings.settings()
+
+            # Last value should win
+            assert result["duplicate"] is False
+
+    async def test_chat_settings_send_returns_settings(self, mock_chainlit_context):
+        """Test that send() returns the settings dictionary."""
+        async with mock_chainlit_context:
+            switch = Switch(id="test", label="Test", initial=True)
+            settings = ChatSettings(inputs=[switch])
+
+            result = await settings.send()
+
+            assert isinstance(result, dict)
+            assert "test" in result
+            assert result["test"] is True
+
+    async def test_chat_settings_to_dict_serialization(self, mock_chainlit_context):
+        """Test that inputs are properly serialized in send()."""
+        async with mock_chainlit_context as ctx:
+            switch = Switch(id="switch", label="Switch", initial=True)
+            slider = Slider(id="slider", label="Slider", initial=5)
+
+            settings = ChatSettings(inputs=[switch, slider])
+            await settings.send()
+
+            # Check that emit was called with serialized inputs
+            call_args = ctx.emitter.emit.call_args
+            inputs_content = call_args[0][1]
+
+            assert len(inputs_content) == 2
+            assert inputs_content[0]["type"] == "switch"
+            assert inputs_content[0]["id"] == "switch"
+            assert inputs_content[1]["type"] == "slider"
+            assert inputs_content[1]["id"] == "slider"
+
+    async def test_chat_settings_with_complex_tab_structure(
+        self, mock_chainlit_context
+    ):
+        """Test ChatSettings with complex tab structure."""
+        async with mock_chainlit_context:
+            # Create a complex structure with multiple tabs
+            tab1 = Tab(
+                id="general",
+                label="General",
+                inputs=[
+                    Switch(id="enabled", label="Enabled", initial=True),
+                    TextInput(id="name", label="Name", initial="MyApp"),
+                ],
+            )
+
+            tab2 = Tab(
+                id="advanced",
+                label="Advanced",
+                inputs=[
+                    Slider(id="timeout", label="Timeout", initial=30, min=0, max=60),
+                    Select(
+                        id="mode",
+                        label="Mode",
+                        values=["dev", "prod"],
+                        initial_index=0,
+                    ),
+                ],
+            )
+
+            settings = ChatSettings(inputs=[tab1, tab2])
+            result = settings.settings()
+
+            assert result["enabled"] is True
+            assert result["name"] == "MyApp"
+            assert result["timeout"] == 30
+            assert result["mode"] == "dev"
+            assert len(result) == 4

backend/tests/test_config.py

@@ -0,0 +1,127 @@
+import json
+from pathlib import Path
+
+import pytest
+
+from chainlit import config as chainlit_config
+from chainlit.config import ChainlitConfig
+
+
+@pytest.fixture
+def translation_dir(tmp_path: Path) -> Path:
+    """Minimal translation directory with a controlled set of locale files."""
+    t_dir = tmp_path / "translations"
+    t_dir.mkdir()
+
+    files: dict[str, dict] = {
+        "en-US.json": {"greeting": "Hello"},
+        "es.json": {"greeting": "Hola"},
+        "da-DK.json": {"greeting": "Hej"},
+        "de-DE.json": {"greeting": "Hallo"},
+        "zh-CN.json": {"greeting": "你好 CN"},
+        "zh-TW.json": {"greeting": "你好 TW"},
+    }
+    for filename, content in files.items():
+        (t_dir / filename).write_text(json.dumps(content), encoding="utf-8")
+
+    return t_dir
+
+
+class TestLoadTranslation:
+    """Regression tests for the load_translation fallback chain."""
+
+    def test_exact_match_regional(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Exact regional locale (da-DK) resolves directly to its file."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("da-DK") == {"greeting": "Hej"}
+
+    def test_exact_match_base(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Exact base locale (es) resolves directly to its file."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("es") == {"greeting": "Hola"}
+
+    def test_parent_fallback(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Regional locale (es-419) falls back to base file (es.json) when no exact match."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("es-419") == {"greeting": "Hola"}
+
+    def test_regional_variant_lookup(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Base locale (da) resolves to regional file (da-DK.json) when no exact match exists."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("da") == {"greeting": "Hej"}
+
+    def test_regional_variant_lookup_de(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Base locale (de) resolves to regional file (de-DE.json) via variant lookup."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("de") == {"greeting": "Hallo"}
+
+    def test_regional_variant_sorted_deterministic(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """When multiple regional variants exist, the first sorted match (zh-CN) is returned."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("zh") == {"greeting": "你好 CN"}
+
+    def test_default_fallback_unknown_locale(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Completely unknown locale (xx) falls back to en-US."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("xx") == {"greeting": "Hello"}
+
+    def test_default_fallback_base_without_regional_variant(
+        self,
+        test_config: ChainlitConfig,
+        translation_dir: Path,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Base locale (fr) with no matching file at all falls back to en-US."""
+        monkeypatch.setattr(
+            chainlit_config, "config_translation_dir", str(translation_dir)
+        )
+        assert test_config.load_translation("fr") == {"greeting": "Hello"}

backend/tests/test_context.py

@@ -0,0 +1,55 @@
+from unittest.mock import Mock
+
+import pytest
+
+from chainlit.context import (
+    ChainlitContext,
+    ChainlitContextException,
+    get_context,
+    init_http_context,
+    init_ws_context,
+)
+from chainlit.emitter import BaseChainlitEmitter, ChainlitEmitter
+from chainlit.session import HTTPSession
+
+
+@pytest.fixture
+def mock_emitter():
+    return Mock(spec=BaseChainlitEmitter)
+
+
+async def test_chainlit_context_init_with_websocket(
+    mock_websocket_session, mock_emitter
+):
+    context = ChainlitContext(mock_websocket_session, mock_emitter)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+    assert context.session == mock_websocket_session
+
+
+async def test_chainlit_context_init_with_http(mock_http_session):
+    context = ChainlitContext(mock_http_session)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+    assert context.session == mock_http_session
+
+
+async def test_init_ws_context(mock_websocket_session):
+    context = init_ws_context(mock_websocket_session)
+    assert isinstance(context, ChainlitContext)
+    assert context.session == mock_websocket_session
+    assert isinstance(context.emitter, ChainlitEmitter)
+
+
+async def test_init_http_context():
+    context = init_http_context()
+    assert isinstance(context, ChainlitContext)
+    assert isinstance(context.session, HTTPSession)
+    assert isinstance(context.emitter, BaseChainlitEmitter)
+
+
+async def test_get_context():
+    with pytest.raises(ChainlitContextException):
+        get_context()
+
+    init_http_context()  # Initialize a context
+    context = get_context()
+    assert isinstance(context, ChainlitContext)

backend/tests/test_element.py

@@ -0,0 +1,584 @@
+import json
+import uuid
+from unittest.mock import AsyncMock
+
+import pytest
+
+from chainlit.element import (
+    Audio,
+    CustomElement,
+    Dataframe,
+    Element,
+    ElementDict,
+    File,
+    Image,
+    Pdf,
+    Task,
+    TaskList,
+    TaskStatus,
+    Text,
+    Video,
+)
+
+
+@pytest.mark.asyncio
+class TestElementBase:
+    """Test suite for the base Element class."""
+
+    async def test_element_initialization_with_url(self, mock_chainlit_context):
+        """Test Element initialization with URL."""
+        async with mock_chainlit_context:
+            element = File(name="test_file", url="https://example.com/file.pdf")
+
+            assert element.name == "test_file"
+            assert element.url == "https://example.com/file.pdf"
+            assert isinstance(element.id, str)
+            uuid.UUID(element.id)  # Verify valid UUID
+            assert element.persisted is False
+            assert element.updatable is False
+
+    async def test_element_initialization_with_content(self, mock_chainlit_context):
+        """Test Element initialization with content."""
+        async with mock_chainlit_context:
+            content = b"test content"
+            element = File(name="test_file", content=content)
+
+            assert element.name == "test_file"
+            assert element.content == content
+            assert element.url is None
+            assert element.path is None
+
+    async def test_element_initialization_with_path(self, mock_chainlit_context):
+        """Test Element initialization with path."""
+        async with mock_chainlit_context:
+            element = File(name="test_file", path="/path/to/file.txt")
+
+            assert element.name == "test_file"
+            assert element.path == "/path/to/file.txt"
+            assert element.url is None
+            assert element.content is None
+
+    async def test_element_requires_url_path_or_content(self, mock_chainlit_context):
+        """Test that Element raises error without url, path, or content."""
+        async with mock_chainlit_context:
+            with pytest.raises(ValueError, match="Must provide url, path or content"):
+                File(name="test_file")
+
+    async def test_element_to_dict(self, mock_chainlit_context):
+        """Test Element serialization to dictionary."""
+        async with mock_chainlit_context as ctx:
+            element = File(
+                name="test_file",
+                url="https://example.com/file.pdf",
+                display="inline",
+            )
+
+            element_dict = element.to_dict()
+
+            assert element_dict["name"] == "test_file"
+            assert element_dict["url"] == "https://example.com/file.pdf"
+            assert element_dict["type"] == "file"
+            assert element_dict["id"] == element.id
+            assert element_dict["threadId"] == ctx.session.thread_id
+            assert element_dict["display"] == "inline"
+
+    async def test_element_send(self, mock_chainlit_context):
+        """Test Element.send() method."""
+        async with mock_chainlit_context as ctx:
+            element = File(name="test_file", url="https://example.com/file.pdf")
+
+            await element.send(for_id="message_123")
+
+            assert element.for_id == "message_123"
+            ctx.emitter.send_element.assert_called_once()
+
+    async def test_element_remove(self, mock_chainlit_context):
+        """Test Element.remove() method."""
+        async with mock_chainlit_context as ctx:
+            element = File(name="test_file", url="https://example.com/file.pdf")
+
+            await element.remove()
+
+            ctx.emitter.emit.assert_called_once_with(
+                "remove_element", {"id": element.id}
+            )
+
+    async def test_element_display_options(self, mock_chainlit_context):
+        """Test Element display options."""
+        async with mock_chainlit_context:
+            element_inline = File(
+                name="test", url="https://example.com/file.pdf", display="inline"
+            )
+            element_side = File(
+                name="test", url="https://example.com/file.pdf", display="side"
+            )
+            element_page = File(
+                name="test", url="https://example.com/file.pdf", display="page"
+            )
+
+            assert element_inline.display == "inline"
+            assert element_side.display == "side"
+            assert element_page.display == "page"
+
+    async def test_element_from_dict_file(self, mock_chainlit_context):
+        """Test Element.from_dict() for File type."""
+        async with mock_chainlit_context:
+            element_dict: ElementDict = {
+                "id": str(uuid.uuid4()),
+                "name": "test_file",
+                "type": "file",
+                "url": "https://example.com/file.pdf",
+                "display": "inline",
+            }
+
+            element = Element.from_dict(element_dict)
+
+            assert isinstance(element, File)
+            assert element.name == "test_file"
+            assert element.url == "https://example.com/file.pdf"
+
+    async def test_element_from_dict_image(self, mock_chainlit_context):
+        """Test Element.from_dict() for Image type."""
+        async with mock_chainlit_context:
+            element_dict: ElementDict = {
+                "id": str(uuid.uuid4()),
+                "name": "test_image",
+                "type": "image",
+                "url": "https://example.com/image.png",
+                "display": "inline",
+            }
+
+            element = Element.from_dict(element_dict)
+
+            assert isinstance(element, Image)
+            assert element.name == "test_image"
+            assert element.type == "image"
+
+    async def test_element_infer_type_from_mime(self):
+        """Test Element.infer_type_from_mime() method."""
+        assert Element.infer_type_from_mime("image/png") == "image"
+        assert Element.infer_type_from_mime("image/jpeg") == "image"
+        assert Element.infer_type_from_mime("application/pdf") == "pdf"
+        assert Element.infer_type_from_mime("audio/mp3") == "audio"
+        assert Element.infer_type_from_mime("video/mp4") == "video"
+        assert Element.infer_type_from_mime("text/plain") == "file"
+        assert Element.infer_type_from_mime("application/json") == "file"
+
+
+@pytest.mark.asyncio
+class TestImageElement:
+    """Test suite for Image element."""
+
+    async def test_image_initialization(self, mock_chainlit_context):
+        """Test Image element initialization."""
+        async with mock_chainlit_context:
+            image = Image(
+                name="test_image",
+                url="https://example.com/image.png",
+                size="large",
+            )
+
+            assert image.type == "image"
+            assert image.name == "test_image"
+            assert image.size == "large"
+
+    async def test_image_size_options(self, mock_chainlit_context):
+        """Test Image size options."""
+        async with mock_chainlit_context:
+            small = Image(name="test", url="https://example.com/img.png", size="small")
+            medium = Image(
+                name="test", url="https://example.com/img.png", size="medium"
+            )
+            large = Image(name="test", url="https://example.com/img.png", size="large")
+
+            assert small.size == "small"
+            assert medium.size == "medium"
+            assert large.size == "large"
+
+
+@pytest.mark.asyncio
+class TestTextElement:
+    """Test suite for Text element."""
+
+    async def test_text_initialization(self, mock_chainlit_context):
+        """Test Text element initialization."""
+        async with mock_chainlit_context:
+            text = Text(name="test_text", content="Hello, World!", language="python")
+
+            assert text.type == "text"
+            assert text.name == "test_text"
+            assert text.content == "Hello, World!"
+            assert text.language == "python"
+
+    async def test_text_without_language(self, mock_chainlit_context):
+        """Test Text element without language."""
+        async with mock_chainlit_context:
+            text = Text(name="test_text", content="Plain text")
+
+            assert text.language is None
+
+
+@pytest.mark.asyncio
+class TestPdfElement:
+    """Test suite for Pdf element."""
+
+    async def test_pdf_initialization(self, mock_chainlit_context):
+        """Test Pdf element initialization."""
+        async with mock_chainlit_context:
+            pdf = Pdf(name="test_pdf", url="https://example.com/document.pdf", page=5)
+
+            assert pdf.type == "pdf"
+            assert pdf.name == "test_pdf"
+            assert pdf.mime == "application/pdf"
+            assert pdf.page == 5
+
+    async def test_pdf_without_page(self, mock_chainlit_context):
+        """Test Pdf element without page number."""
+        async with mock_chainlit_context:
+            pdf = Pdf(name="test_pdf", url="https://example.com/document.pdf")
+
+            assert pdf.page is None
+
+
+@pytest.mark.asyncio
+class TestAudioElement:
+    """Test suite for Audio element."""
+
+    async def test_audio_initialization(self, mock_chainlit_context):
+        """Test Audio element initialization."""
+        async with mock_chainlit_context:
+            audio = Audio(
+                name="test_audio",
+                url="https://example.com/audio.mp3",
+                auto_play=True,
+            )
+
+            assert audio.type == "audio"
+            assert audio.name == "test_audio"
+            assert audio.auto_play is True
+
+    async def test_audio_default_auto_play(self, mock_chainlit_context):
+        """Test Audio element default auto_play."""
+        async with mock_chainlit_context:
+            audio = Audio(name="test_audio", url="https://example.com/audio.mp3")
+
+            assert audio.auto_play is False
+
+
+@pytest.mark.asyncio
+class TestVideoElement:
+    """Test suite for Video element."""
+
+    async def test_video_initialization(self, mock_chainlit_context):
+        """Test Video element initialization."""
+        async with mock_chainlit_context:
+            player_config = {"youtube": {"playerVars": {"showinfo": 1}}}
+            video = Video(
+                name="test_video",
+                url="https://example.com/video.mp4",
+                size="large",
+                player_config=player_config,
+            )
+
+            assert video.type == "video"
+            assert video.name == "test_video"
+            assert video.size == "large"
+            assert video.player_config == player_config
+
+    async def test_video_without_player_config(self, mock_chainlit_context):
+        """Test Video element without player config."""
+        async with mock_chainlit_context:
+            video = Video(name="test_video", url="https://example.com/video.mp4")
+
+            assert video.player_config is None
+
+
+@pytest.mark.asyncio
+class TestFileElement:
+    """Test suite for File element."""
+
+    async def test_file_initialization(self, mock_chainlit_context):
+        """Test File element initialization."""
+        async with mock_chainlit_context:
+            file = File(name="test_file", url="https://example.com/file.txt")
+
+            assert file.type == "file"
+            assert file.name == "test_file"
+
+    async def test_file_with_content(self, mock_chainlit_context):
+        """Test File element with content."""
+        async with mock_chainlit_context:
+            content = b"File content"
+            file = File(name="test_file", content=content)
+
+            assert file.content == content
+
+
+@pytest.mark.asyncio
+class TestTaskListElement:
+    """Test suite for TaskList element."""
+
+    async def test_tasklist_initialization(self, mock_chainlit_context):
+        """Test TaskList element initialization."""
+        async with mock_chainlit_context:
+            tasklist = TaskList(name="test_tasklist")
+
+            assert tasklist.type == "tasklist"
+            assert tasklist.name == "test_tasklist"
+            assert tasklist.tasks == []
+            assert tasklist.status == "Ready"
+            assert tasklist.updatable is True
+
+    async def test_tasklist_add_task(self, mock_chainlit_context):
+        """Test adding tasks to TaskList."""
+        async with mock_chainlit_context:
+            tasklist = TaskList(name="test_tasklist")
+            task1 = Task(title="Task 1", status=TaskStatus.READY)
+            task2 = Task(title="Task 2", status=TaskStatus.RUNNING)
+
+            await tasklist.add_task(task1)
+            await tasklist.add_task(task2)
+
+            assert len(tasklist.tasks) == 2
+            assert tasklist.tasks[0].title == "Task 1"
+            assert tasklist.tasks[1].title == "Task 2"
+
+    async def test_tasklist_preprocess_content(self, mock_chainlit_context):
+        """Test TaskList content preprocessing."""
+        async with mock_chainlit_context:
+            tasklist = TaskList(name="test_tasklist", status="In Progress")
+            task = Task(title="Test Task", status=TaskStatus.DONE)
+            await tasklist.add_task(task)
+
+            await tasklist.preprocess_content()
+
+            assert isinstance(tasklist.content, str)
+            assert "Test Task" in tasklist.content
+            assert "done" in tasklist.content
+            assert "In Progress" in tasklist.content
+
+
+@pytest.mark.asyncio
+class TestTaskClass:
+    """Test suite for Task class."""
+
+    def test_task_initialization(self):
+        """Test Task initialization."""
+        task = Task(title="Test Task", status=TaskStatus.READY)
+
+        assert task.title == "Test Task"
+        assert task.status == TaskStatus.READY
+        assert task.forId is None
+
+    def test_task_with_for_id(self):
+        """Test Task with forId."""
+        task = Task(title="Test Task", status=TaskStatus.RUNNING, forId="step_123")
+
+        assert task.forId == "step_123"
+
+    def test_task_status_enum(self):
+        """Test TaskStatus enum values."""
+        assert TaskStatus.READY.value == "ready"
+        assert TaskStatus.RUNNING.value == "running"
+        assert TaskStatus.FAILED.value == "failed"
+        assert TaskStatus.DONE.value == "done"
+
+
+@pytest.mark.asyncio
+class TestCustomElement:
+    """Test suite for CustomElement."""
+
+    async def test_custom_element_initialization(self, mock_chainlit_context):
+        """Test CustomElement initialization."""
+        async with mock_chainlit_context:
+            props = {"key1": "value1", "key2": 42}
+            custom = CustomElement(name="test_custom", props=props)
+
+            assert custom.type == "custom"
+            assert custom.name == "test_custom"
+            assert custom.props == props
+            assert custom.mime == "application/json"
+            assert custom.updatable is True
+
+    async def test_custom_element_content_serialization(self, mock_chainlit_context):
+        """Test CustomElement content serialization."""
+        async with mock_chainlit_context:
+            props = {"nested": {"data": [1, 2, 3]}}
+            custom = CustomElement(name="test_custom", props=props)
+
+            assert isinstance(custom.content, str)
+            assert "nested" in custom.content
+            assert "data" in custom.content
+
+    async def test_custom_element_update(self, mock_chainlit_context):
+        """Test CustomElement update method."""
+        async with mock_chainlit_context as ctx:
+            custom = CustomElement(
+                name="test_custom",
+                props={"key": "value"},
+                url="https://example.com/custom",
+            )
+            custom.for_id = "message_123"
+
+            await custom.update()
+
+            ctx.emitter.send_element.assert_called()
+
+
+@pytest.mark.asyncio
+class TestElementEdgeCases:
+    """Test suite for Element edge cases."""
+
+    async def test_element_with_custom_id(self, mock_chainlit_context):
+        """Test Element with custom ID."""
+        async with mock_chainlit_context:
+            custom_id = str(uuid.uuid4())
+            element = File(
+                id=custom_id, name="test_file", url="https://example.com/file.txt"
+            )
+
+            assert element.id == custom_id
+
+    async def test_element_with_object_key(self, mock_chainlit_context):
+        """Test Element with object_key."""
+        async with mock_chainlit_context:
+            element = File(
+                name="test_file",
+                url="https://example.com/file.txt",
+                object_key="s3://bucket/key",
+            )
+
+            assert element.object_key == "s3://bucket/key"
+
+    async def test_element_with_chainlit_key(self, mock_chainlit_context):
+        """Test Element with chainlit_key."""
+        async with mock_chainlit_context:
+            element = File(
+                name="test_file",
+                url="https://example.com/file.txt",
+                chainlit_key="chainlit_key_123",
+            )
+
+            assert element.chainlit_key == "chainlit_key_123"
+
+    async def test_element_send_without_url_or_key_raises_error(
+        self, mock_chainlit_context
+    ):
+        """Test that send() raises error without url or chainlit_key."""
+        async with mock_chainlit_context as ctx:
+            # Mock persist_file to not set chainlit_key
+            ctx.session.persist_file = AsyncMock(return_value={"id": None})
+
+            element = File(name="test_file", content=b"test content")
+
+            with pytest.raises(ValueError, match="Must provide url or chainlit key"):
+                await element.send(for_id="message_123", persist=False)
+
+    async def test_element_from_dict_with_missing_fields(self, mock_chainlit_context):
+        """Test Element.from_dict() with minimal fields."""
+        async with mock_chainlit_context:
+            element_dict: ElementDict = {
+                "type": "file",
+                "url": "https://example.com/file.txt",
+            }
+
+            element = Element.from_dict(element_dict)
+
+            assert isinstance(element, File)
+            assert element.name == ""
+            assert element.url == "https://example.com/file.txt"
+
+    async def test_element_id_uniqueness(self, mock_chainlit_context):
+        """Test that each Element gets a unique ID."""
+        async with mock_chainlit_context:
+            element1 = File(name="file1", url="https://example.com/file1.txt")
+            element2 = File(name="file2", url="https://example.com/file2.txt")
+            element3 = File(name="file3", url="https://example.com/file3.txt")
+
+            ids = {element1.id, element2.id, element3.id}
+            assert len(ids) == 3  # All unique
+
+
+@pytest.mark.asyncio
+class TestDataframeElement:
+    """Test suite for Dataframe element."""
+
+    async def test_dataframe_with_pandas(self, mock_chainlit_context):
+        """Test Dataframe element with a pandas DataFrame."""
+        import pandas as pd
+
+        async with mock_chainlit_context:
+            df = pd.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
+            element = Dataframe(name="test_df", data=df)
+
+            assert element.type == "dataframe"
+            assert element.size == "large"
+
+            parsed = json.loads(element.content)
+            assert parsed["columns"] == ["a", "b"]
+            assert parsed["data"] == [[4, "foo"], [2, "bar"], [0, "baz"]]
+            assert parsed["index"] == [0, 1, 2]
+
+    async def test_dataframe_with_polars(self, mock_chainlit_context):
+        """Test Dataframe element with a polars DataFrame."""
+        import polars as pl
+
+        async with mock_chainlit_context:
+            df = pl.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
+            element = Dataframe(name="test_df", data=df)
+
+            assert element.type == "dataframe"
+            assert element.size == "large"
+
+            parsed = json.loads(element.content)
+            assert parsed["columns"] == ["a", "b"]
+            assert parsed["data"] == [[4, "foo"], [2, "bar"], [0, "baz"]]
+            assert parsed["index"] == [0, 1, 2]
+
+    async def test_dataframe_with_invalid_data(self, mock_chainlit_context):
+        """Test Dataframe element rejects non-DataFrame data."""
+        async with mock_chainlit_context:
+            with pytest.raises(
+                TypeError,
+                match=r"data must be a pandas\.DataFrame or polars\.DataFrame",
+            ):
+                Dataframe(name="test_df", data={"a": [1, 2]})
+
+    async def test_dataframe_polars_and_pandas_produce_equal_outputs(
+        self, mock_chainlit_context
+    ):
+        """Test that pandas and polars DataFrames produce the same JSON."""
+        import pandas as pd
+        import polars as pl
+
+        async with mock_chainlit_context:
+            pd_df = pd.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
+            pl_df = pl.DataFrame({"a": [4, 2, 0], "b": ["foo", "bar", "baz"]})
+
+            pd_element = Dataframe(name="pd_df", data=pd_df)
+            pl_element = Dataframe(name="pl_df", data=pl_df)
+
+            pd_parsed = json.loads(pd_element.content)
+            pl_parsed = json.loads(pl_element.content)
+
+            assert pd_parsed["columns"] == pl_parsed["columns"]
+            assert pd_parsed["index"] == pl_parsed["index"]
+            assert pd_parsed["data"] == pl_parsed["data"]
+
+    async def test_dataframe_polars_with_dates(self, mock_chainlit_context):
+        """Test Dataframe element with polars date columns serializes correctly."""
+        from datetime import date
+
+        import polars as pl
+
+        async with mock_chainlit_context:
+            df = pl.DataFrame(
+                {"date": [date(2026, 1, 1), date(2025, 12, 31)], "val": [1, 2]}
+            )
+            element = Dataframe(name="test_df", data=df)
+
+            parsed = json.loads(element.content)
+            assert parsed["columns"] == ["date", "val"]
+            assert len(parsed["data"]) == 2
+            assert parsed["data"][0][0] == "2026-01-01"
+            assert parsed["data"][1][0] == "2025-12-31"

backend/tests/test_emitter.py

@@ -0,0 +1,211 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from chainlit.element import ElementDict
+from chainlit.emitter import ChainlitEmitter
+from chainlit.step import StepDict
+
+
+@pytest.fixture
+def emitter(mock_websocket_session):
+    return ChainlitEmitter(mock_websocket_session)
+
+
+async def test_send_element(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    element_dict: ElementDict = {
+        "id": "test_element",
+        "threadId": None,
+        "type": "text",
+        "chainlitKey": None,
+        "url": None,
+        "objectKey": None,
+        "name": "Test Element",
+        "display": "inline",
+        "size": None,
+        "language": None,
+        "page": None,
+        "props": None,
+        "autoPlay": None,
+        "playerConfig": None,
+        "forId": None,
+        "mime": None,
+    }
+
+    await emitter.send_element(element_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("element", element_dict)
+
+
+async def test_send_step(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_step",
+        "type": "user_message",
+        "name": "Test Step",
+        "output": "This is a test step",
+    }
+
+    await emitter.send_step(step_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("new_message", step_dict)
+
+
+async def test_send_step_with_icon(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_step_with_icon",
+        "type": "tool",
+        "name": "Test Step with Icon",
+        "output": "This is a test step with an icon",
+        "metadata": {"icon": "search"},
+    }
+
+    await emitter.send_step(step_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("new_message", step_dict)
+
+
+async def test_update_step(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_step",
+        "type": "assistant_message",
+        "name": "Updated Test Step",
+        "output": "This is an updated test step",
+    }
+
+    await emitter.update_step(step_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("update_message", step_dict)
+
+
+async def test_update_step_with_icon(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_step_with_icon",
+        "type": "tool",
+        "name": "Updated Test Step with Icon",
+        "output": "This is an updated test step with an icon",
+        "metadata": {"icon": "database"},
+    }
+
+    await emitter.update_step(step_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("update_message", step_dict)
+
+
+async def test_delete_step(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_step",
+        "type": "system_message",
+        "name": "Deleted Test Step",
+        "output": "This step will be deleted",
+    }
+
+    await emitter.delete_step(step_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("delete_message", step_dict)
+
+
+async def test_send_timeout(emitter, mock_websocket_session):
+    await emitter.send_timeout("ask_timeout")
+    mock_websocket_session.emit.assert_called_once_with("ask_timeout", {})
+
+
+async def test_clear(emitter, mock_websocket_session):
+    await emitter.clear("clear_ask")
+    mock_websocket_session.emit.assert_called_once_with("clear_ask", {})
+
+
+async def test_send_token(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    await emitter.send_token("test_id", "test_token", is_sequence=True, is_input=False)
+    mock_websocket_session.emit.assert_called_once_with(
+        "stream_token",
+        {"id": "test_id", "token": "test_token", "isSequence": True, "isInput": False},
+    )
+
+
+async def test_set_chat_settings(emitter, mock_websocket_session):
+    settings = {"key": "value"}
+    emitter.set_chat_settings(settings)
+    assert emitter.session.chat_settings == settings
+
+
+async def test_update_token_count(emitter, mock_websocket_session):
+    count = 100
+    await emitter.update_token_count(count)
+    mock_websocket_session.emit.assert_called_once_with("token_usage", count)
+
+
+async def test_task_start(emitter, mock_websocket_session):
+    await emitter.task_start()
+    mock_websocket_session.emit.assert_called_once_with("task_start", {})
+
+
+async def test_task_end(emitter, mock_websocket_session):
+    await emitter.task_end()
+    mock_websocket_session.emit.assert_called_once_with("task_end", {})
+
+
+async def test_stream_start(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_stream",
+        "type": "run",
+        "name": "Test Stream",
+        "output": "This is a test stream",
+    }
+    await emitter.stream_start(step_dict)
+    mock_websocket_session.emit.assert_called_once_with("stream_start", step_dict)
+
+
+async def test_stream_start_with_icon(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_stream_with_icon",
+        "type": "tool",
+        "name": "Test Stream with Icon",
+        "output": "This is a test stream with an icon",
+        "metadata": {"icon": "cpu"},
+    }
+    await emitter.stream_start(step_dict)
+    mock_websocket_session.emit.assert_called_once_with("stream_start", step_dict)
+
+
+async def test_send_toast(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    message = "This is a test message"
+    await emitter.send_toast(message)
+    mock_websocket_session.emit.assert_called_once_with(
+        "toast", {"message": message, "type": "info"}
+    )
+
+
+async def test_send_toast_with_type(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    message = "This is a test message"
+    await emitter.send_toast(message, type="error")
+    mock_websocket_session.emit.assert_called_once_with(
+        "toast", {"message": message, "type": "error"}
+    )
+
+
+async def test_send_toast_invalid_type(emitter: ChainlitEmitter) -> None:
+    message = "This is a test message"
+    with pytest.raises(ValueError, match="Invalid toast type: invalid"):
+        await emitter.send_toast(message, type="invalid")  # type: ignore[arg-type]

backend/tests/test_input_widget.py

@@ -0,0 +1,734 @@
+import pytest
+
+from chainlit.input_widget import (
+    Checkbox,
+    MultiSelect,
+    NumberInput,
+    RadioGroup,
+    Select,
+    Slider,
+    Switch,
+    Tab,
+    Tags,
+    TextInput,
+)
+
+
+class TestInputWidgetBase:
+    """Test suite for base InputWidget validation."""
+
+    def test_input_widget_requires_id_and_label(self):
+        """Test that InputWidget requires both id and label."""
+        with pytest.raises(ValueError, match="Must provide key and label"):
+            Switch(id="", label="Test Label")
+
+        with pytest.raises(ValueError, match="Must provide key and label"):
+            Switch(id="test_id", label="")
+
+
+class TestSwitchWidget:
+    """Test suite for Switch input widget."""
+
+    def test_switch_initialization(self):
+        """Test Switch widget initialization."""
+        switch = Switch(id="test_switch", label="Enable Feature")
+
+        assert switch.id == "test_switch"
+        assert switch.label == "Enable Feature"
+        assert switch.type == "switch"
+        assert switch.initial is False
+        assert switch.disabled is False
+
+    def test_switch_with_initial_value(self):
+        """Test Switch widget with initial value."""
+        switch = Switch(id="test_switch", label="Enable Feature", initial=True)
+
+        assert switch.initial is True
+
+    def test_switch_with_tooltip_and_description(self):
+        """Test Switch widget with tooltip and description."""
+        switch = Switch(
+            id="test_switch",
+            label="Enable Feature",
+            tooltip="Toggle this feature",
+            description="This enables the advanced feature",
+        )
+
+        assert switch.tooltip == "Toggle this feature"
+        assert switch.description == "This enables the advanced feature"
+
+    def test_switch_disabled(self):
+        """Test Switch widget in disabled state."""
+        switch = Switch(id="test_switch", label="Enable Feature", disabled=True)
+
+        assert switch.disabled is True
+
+    def test_switch_to_dict(self):
+        """Test Switch widget serialization."""
+        switch = Switch(
+            id="test_switch",
+            label="Enable Feature",
+            initial=True,
+            tooltip="Toggle",
+            description="Description",
+            disabled=False,
+        )
+
+        result = switch.to_dict()
+
+        assert result["type"] == "switch"
+        assert result["id"] == "test_switch"
+        assert result["label"] == "Enable Feature"
+        assert result["initial"] is True
+        assert result["tooltip"] == "Toggle"
+        assert result["description"] == "Description"
+        assert result["disabled"] is False
+
+
+class TestSliderWidget:
+    """Test suite for Slider input widget."""
+
+    def test_slider_initialization(self):
+        """Test Slider widget initialization."""
+        slider = Slider(id="test_slider", label="Temperature")
+
+        assert slider.id == "test_slider"
+        assert slider.label == "Temperature"
+        assert slider.type == "slider"
+        assert slider.initial == 0
+        assert slider.min == 0
+        assert slider.max == 10
+        assert slider.step == 1
+
+    def test_slider_with_custom_range(self):
+        """Test Slider widget with custom range."""
+        slider = Slider(
+            id="test_slider",
+            label="Temperature",
+            initial=0.5,
+            min=0.0,
+            max=1.0,
+            step=0.1,
+        )
+
+        assert slider.initial == 0.5
+        assert slider.min == 0.0
+        assert slider.max == 1.0
+        assert slider.step == 0.1
+
+    def test_slider_to_dict(self):
+        """Test Slider widget serialization."""
+        slider = Slider(
+            id="test_slider",
+            label="Temperature",
+            initial=0.7,
+            min=0.0,
+            max=2.0,
+            step=0.1,
+            tooltip="Adjust temperature",
+        )
+
+        result = slider.to_dict()
+
+        assert result["type"] == "slider"
+        assert result["id"] == "test_slider"
+        assert result["label"] == "Temperature"
+        assert result["initial"] == 0.7
+        assert result["min"] == 0.0
+        assert result["max"] == 2.0
+        assert result["step"] == 0.1
+        assert result["tooltip"] == "Adjust temperature"
+
+
+class TestSelectWidget:
+    """Test suite for Select input widget."""
+
+    def test_select_with_values(self):
+        """Test Select widget with values list."""
+        select = Select(
+            id="test_select",
+            label="Choose Model",
+            values=["gpt-4", "gpt-3.5", "claude"],
+        )
+
+        assert select.id == "test_select"
+        assert select.label == "Choose Model"
+        assert select.type == "select"
+        assert select.items == {
+            "gpt-4": "gpt-4",
+            "gpt-3.5": "gpt-3.5",
+            "claude": "claude",
+        }
+
+    def test_select_with_items(self):
+        """Test Select widget with items dict."""
+        items = {"gpt4": "GPT-4", "gpt35": "GPT-3.5", "claude": "Claude"}
+        select = Select(id="test_select", label="Choose Model", items=items)
+
+        assert select.items == items
+
+    def test_select_with_initial_index(self):
+        """Test Select widget with initial_index."""
+        select = Select(
+            id="test_select",
+            label="Choose Model",
+            values=["gpt-4", "gpt-3.5", "claude"],
+            initial_index=1,
+        )
+
+        assert select.initial == "gpt-3.5"
+
+    def test_select_with_initial_value(self):
+        """Test Select widget with initial_value."""
+        select = Select(
+            id="test_select",
+            label="Choose Model",
+            values=["gpt-4", "gpt-3.5", "claude"],
+            initial_value="claude",
+        )
+
+        assert select.initial == "claude"
+
+    def test_select_requires_values_or_items(self):
+        """Test that Select requires either values or items."""
+        with pytest.raises(ValueError, match="Must provide values or items"):
+            Select(id="test_select", label="Choose Model")
+
+    def test_select_cannot_have_both_values_and_items(self):
+        """Test that Select cannot have both values and items."""
+        with pytest.raises(ValueError, match="only provide either values or items"):
+            Select(
+                id="test_select",
+                label="Choose Model",
+                values=["a", "b"],
+                items={"a": "A"},
+            )
+
+    def test_select_initial_index_requires_values(self):
+        """Test that initial_index requires values."""
+        with pytest.raises(
+            ValueError,
+            match="Initial_index can only be used in combination with values",
+        ):
+            Select(
+                id="test_select",
+                label="Choose Model",
+                items={"a": "A"},
+                initial_index=0,
+            )
+
+    def test_select_to_dict(self):
+        """Test Select widget serialization."""
+        select = Select(
+            id="test_select",
+            label="Choose Model",
+            values=["gpt-4", "gpt-3.5"],
+            initial_index=0,
+            tooltip="Select a model",
+        )
+
+        result = select.to_dict()
+
+        assert result["type"] == "select"
+        assert result["id"] == "test_select"
+        assert result["label"] == "Choose Model"
+        assert result["initial"] == "gpt-4"
+        assert len(result["items"]) == 2
+        assert result["items"][0] == {"label": "gpt-4", "value": "gpt-4"}
+        assert result["tooltip"] == "Select a model"
+
+
+class TestTextInputWidget:
+    """Test suite for TextInput widget."""
+
+    def test_textinput_initialization(self):
+        """Test TextInput widget initialization."""
+        text_input = TextInput(id="test_input", label="Enter Name")
+
+        assert text_input.id == "test_input"
+        assert text_input.label == "Enter Name"
+        assert text_input.type == "textinput"
+        assert text_input.initial is None
+        assert text_input.placeholder is None
+        assert text_input.multiline is False
+
+    def test_textinput_with_initial_and_placeholder(self):
+        """Test TextInput widget with initial value and placeholder."""
+        text_input = TextInput(
+            id="test_input",
+            label="Enter Name",
+            initial="John Doe",
+            placeholder="Enter your name",
+        )
+
+        assert text_input.initial == "John Doe"
+        assert text_input.placeholder == "Enter your name"
+
+    def test_textinput_multiline(self):
+        """Test TextInput widget in multiline mode."""
+        text_input = TextInput(
+            id="test_input", label="Enter Description", multiline=True
+        )
+
+        assert text_input.multiline is True
+
+    def test_textinput_to_dict(self):
+        """Test TextInput widget serialization."""
+        text_input = TextInput(
+            id="test_input",
+            label="Enter Name",
+            initial="Default",
+            placeholder="Type here",
+            multiline=True,
+            tooltip="Enter your name",
+        )
+
+        result = text_input.to_dict()
+
+        assert result["type"] == "textinput"
+        assert result["id"] == "test_input"
+        assert result["label"] == "Enter Name"
+        assert result["initial"] == "Default"
+        assert result["placeholder"] == "Type here"
+        assert result["multiline"] is True
+        assert result["tooltip"] == "Enter your name"
+
+
+class TestNumberInputWidget:
+    """Test suite for NumberInput widget."""
+
+    def test_numberinput_initialization(self):
+        """Test NumberInput widget initialization."""
+        number_input = NumberInput(id="test_number", label="Enter Age")
+
+        assert number_input.id == "test_number"
+        assert number_input.label == "Enter Age"
+        assert number_input.type == "numberinput"
+        assert number_input.initial is None
+        assert number_input.placeholder is None
+
+    def test_numberinput_with_initial(self):
+        """Test NumberInput widget with initial value."""
+        number_input = NumberInput(
+            id="test_number", label="Enter Age", initial=25.5, placeholder="Age"
+        )
+
+        assert number_input.initial == 25.5
+        assert number_input.placeholder == "Age"
+
+    def test_numberinput_to_dict(self):
+        """Test NumberInput widget serialization."""
+        number_input = NumberInput(
+            id="test_number",
+            label="Enter Age",
+            initial=30.0,
+            placeholder="Enter a number",
+            tooltip="Your age",
+        )
+
+        result = number_input.to_dict()
+
+        assert result["type"] == "numberinput"
+        assert result["id"] == "test_number"
+        assert result["label"] == "Enter Age"
+        assert result["initial"] == 30.0
+        assert result["placeholder"] == "Enter a number"
+        assert result["tooltip"] == "Your age"
+
+
+class TestTagsWidget:
+    """Test suite for Tags widget."""
+
+    def test_tags_initialization(self):
+        """Test Tags widget initialization."""
+        tags = Tags(id="test_tags", label="Add Tags")
+
+        assert tags.id == "test_tags"
+        assert tags.label == "Add Tags"
+        assert tags.type == "tags"
+        assert tags.initial == []
+        assert tags.values == []
+
+    def test_tags_with_initial_values(self):
+        """Test Tags widget with initial values."""
+        tags = Tags(
+            id="test_tags",
+            label="Add Tags",
+            initial=["python", "javascript"],
+            values=["python", "javascript", "go", "rust"],
+        )
+
+        assert tags.initial == ["python", "javascript"]
+        assert tags.values == ["python", "javascript", "go", "rust"]
+
+    def test_tags_to_dict(self):
+        """Test Tags widget serialization."""
+        tags = Tags(
+            id="test_tags",
+            label="Add Tags",
+            initial=["tag1"],
+            tooltip="Add your tags",
+        )
+
+        result = tags.to_dict()
+
+        assert result["type"] == "tags"
+        assert result["id"] == "test_tags"
+        assert result["label"] == "Add Tags"
+        assert result["initial"] == ["tag1"]
+        assert result["tooltip"] == "Add your tags"
+
+
+class TestMultiSelectWidget:
+    """Test suite for MultiSelect widget."""
+
+    def test_multiselect_with_values(self):
+        """Test MultiSelect widget with values list."""
+        multi_select = MultiSelect(
+            id="test_multiselect",
+            label="Choose Languages",
+            values=["Python", "JavaScript", "Go"],
+        )
+
+        assert multi_select.id == "test_multiselect"
+        assert multi_select.label == "Choose Languages"
+        assert multi_select.type == "multiselect"
+        assert multi_select.items == {
+            "Python": "Python",
+            "JavaScript": "JavaScript",
+            "Go": "Go",
+        }
+
+    def test_multiselect_with_items(self):
+        """Test MultiSelect widget with items dict."""
+        items = {"py": "Python", "js": "JavaScript", "go": "Go"}
+        multi_select = MultiSelect(
+            id="test_multiselect", label="Choose Languages", items=items
+        )
+
+        assert multi_select.items == items
+
+    def test_multiselect_with_initial(self):
+        """Test MultiSelect widget with initial selection."""
+        multi_select = MultiSelect(
+            id="test_multiselect",
+            label="Choose Languages",
+            values=["Python", "JavaScript", "Go"],
+            initial=["Python", "Go"],
+        )
+
+        assert multi_select.initial == ["Python", "Go"]
+
+    def test_multiselect_requires_values_or_items(self):
+        """Test that MultiSelect requires either values or items."""
+        with pytest.raises(ValueError, match="Must provide values or items"):
+            MultiSelect(id="test_multiselect", label="Choose Languages")
+
+    def test_multiselect_cannot_have_both_values_and_items(self):
+        """Test that MultiSelect cannot have both values and items."""
+        with pytest.raises(ValueError, match="only provide either values or items"):
+            MultiSelect(
+                id="test_multiselect",
+                label="Choose Languages",
+                values=["a", "b"],
+                items={"a": "A"},
+            )
+
+    def test_multiselect_to_dict(self):
+        """Test MultiSelect widget serialization."""
+        multi_select = MultiSelect(
+            id="test_multiselect",
+            label="Choose Languages",
+            values=["Python", "JavaScript"],
+            initial=["Python"],
+            tooltip="Select languages",
+        )
+
+        result = multi_select.to_dict()
+
+        assert result["type"] == "multiselect"
+        assert result["id"] == "test_multiselect"
+        assert result["label"] == "Choose Languages"
+        assert result["initial"] == ["Python"]
+        assert len(result["items"]) == 2
+        assert result["tooltip"] == "Select languages"
+
+
+class TestCheckboxWidget:
+    """Test suite for Checkbox widget."""
+
+    def test_checkbox_initialization(self):
+        """Test Checkbox widget initialization."""
+        checkbox = Checkbox(id="test_checkbox", label="Accept Terms")
+
+        assert checkbox.id == "test_checkbox"
+        assert checkbox.label == "Accept Terms"
+        assert checkbox.type == "checkbox"
+        assert checkbox.initial is False
+
+    def test_checkbox_with_initial_value(self):
+        """Test Checkbox widget with initial value."""
+        checkbox = Checkbox(id="test_checkbox", label="Accept Terms", initial=True)
+
+        assert checkbox.initial is True
+
+    def test_checkbox_to_dict(self):
+        """Test Checkbox widget serialization."""
+        checkbox = Checkbox(
+            id="test_checkbox",
+            label="Accept Terms",
+            initial=True,
+            tooltip="Check to accept",
+            description="Terms and conditions",
+        )
+
+        result = checkbox.to_dict()
+
+        assert result["type"] == "checkbox"
+        assert result["id"] == "test_checkbox"
+        assert result["label"] == "Accept Terms"
+        assert result["initial"] is True
+        assert result["tooltip"] == "Check to accept"
+        assert result["description"] == "Terms and conditions"
+
+
+class TestRadioGroupWidget:
+    """Test suite for RadioGroup widget."""
+
+    def test_radiogroup_with_values(self):
+        """Test RadioGroup widget with values list."""
+        radio = RadioGroup(
+            id="test_radio", label="Choose Size", values=["Small", "Medium", "Large"]
+        )
+
+        assert radio.id == "test_radio"
+        assert radio.label == "Choose Size"
+        assert radio.type == "radio"
+        assert radio.items == {"Small": "Small", "Medium": "Medium", "Large": "Large"}
+
+    def test_radiogroup_with_items(self):
+        """Test RadioGroup widget with items dict."""
+        items = {"s": "Small", "m": "Medium", "l": "Large"}
+        radio = RadioGroup(id="test_radio", label="Choose Size", items=items)
+
+        assert radio.items == items
+
+    def test_radiogroup_with_initial_index(self):
+        """Test RadioGroup widget with initial_index."""
+        radio = RadioGroup(
+            id="test_radio",
+            label="Choose Size",
+            values=["Small", "Medium", "Large"],
+            initial_index=1,
+        )
+
+        assert radio.initial == "Medium"
+
+    def test_radiogroup_with_initial_value(self):
+        """Test RadioGroup widget with initial_value."""
+        radio = RadioGroup(
+            id="test_radio",
+            label="Choose Size",
+            values=["Small", "Medium", "Large"],
+            initial_value="Large",
+        )
+
+        assert radio.initial == "Large"
+
+    def test_radiogroup_requires_values_or_items(self):
+        """Test that RadioGroup requires either values or items."""
+        with pytest.raises(ValueError, match="Must provide values or items"):
+            RadioGroup(id="test_radio", label="Choose Size")
+
+    def test_radiogroup_cannot_have_both_values_and_items(self):
+        """Test that RadioGroup cannot have both values and items."""
+        with pytest.raises(ValueError, match="only provide either values or items"):
+            RadioGroup(
+                id="test_radio",
+                label="Choose Size",
+                values=["a", "b"],
+                items={"a": "A"},
+            )
+
+    def test_radiogroup_initial_index_requires_values(self):
+        """Test that initial_index requires values."""
+        with pytest.raises(
+            ValueError,
+            match="Initial_index can only be used in combination with values",
+        ):
+            RadioGroup(
+                id="test_radio", label="Choose Size", items={"a": "A"}, initial_index=0
+            )
+
+    def test_radiogroup_to_dict(self):
+        """Test RadioGroup widget serialization."""
+        radio = RadioGroup(
+            id="test_radio",
+            label="Choose Size",
+            values=["Small", "Medium"],
+            initial_index=0,
+            tooltip="Select size",
+        )
+
+        result = radio.to_dict()
+
+        assert result["type"] == "radio"
+        assert result["id"] == "test_radio"
+        assert result["label"] == "Choose Size"
+        assert result["initial"] == "Small"
+        assert len(result["items"]) == 2
+        assert result["items"][0] == {"label": "Small", "value": "Small"}
+        assert result["tooltip"] == "Select size"
+
+
+class TestTabWidget:
+    """Test suite for Tab widget."""
+
+    def test_tab_initialization(self):
+        """Test Tab initialization."""
+        tab = Tab(id="test_tab", label="Settings")
+
+        assert tab.id == "test_tab"
+        assert tab.label == "Settings"
+        assert tab.inputs == []
+
+    def test_tab_with_inputs(self):
+        """Test Tab with input widgets."""
+        switch = Switch(id="switch1", label="Enable")
+        slider = Slider(id="slider1", label="Value")
+        tab = Tab(id="test_tab", label="Settings", inputs=[switch, slider])
+
+        assert len(tab.inputs) == 2
+        assert tab.inputs[0] == switch
+        assert tab.inputs[1] == slider
+
+    def test_tab_to_dict(self):
+        """Test Tab serialization."""
+        switch = Switch(id="switch1", label="Enable", initial=True)
+        slider = Slider(id="slider1", label="Value", initial=5)
+        tab = Tab(id="test_tab", label="Settings", inputs=[switch, slider])
+
+        result = tab.to_dict()
+
+        assert result["id"] == "test_tab"
+        assert result["label"] == "Settings"
+        assert len(result["inputs"]) == 2
+        assert result["inputs"][0]["type"] == "switch"
+        assert result["inputs"][0]["id"] == "switch1"
+        assert result["inputs"][1]["type"] == "slider"
+        assert result["inputs"][1]["id"] == "slider1"
+
+    def test_tab_to_dict_empty_inputs(self):
+        """Test Tab serialization with no inputs."""
+        tab = Tab(id="test_tab", label="Empty Tab")
+
+        result = tab.to_dict()
+
+        assert result["id"] == "test_tab"
+        assert result["label"] == "Empty Tab"
+        assert result["inputs"] == []
+
+
+class TestInputWidgetEdgeCases:
+    """Test suite for InputWidget edge cases."""
+
+    def test_all_widgets_have_consistent_common_fields(self):
+        """Test that all widgets support common fields."""
+        widgets = [
+            Switch(
+                id="test",
+                label="Test",
+                tooltip="Tooltip",
+                description="Description",
+                disabled=True,
+            ),
+            Slider(
+                id="test",
+                label="Test",
+                tooltip="Tooltip",
+                description="Description",
+                disabled=True,
+            ),
+            Checkbox(
+                id="test",
+                label="Test",
+                tooltip="Tooltip",
+                description="Description",
+                disabled=True,
+            ),
+            TextInput(
+                id="test",
+                label="Test",
+                tooltip="Tooltip",
+                description="Description",
+                disabled=True,
+            ),
+            NumberInput(
+                id="test",
+                label="Test",
+                tooltip="Tooltip",
+                description="Description",
+                disabled=True,
+            ),
+            Tags(
+                id="test",
+                label="Test",
+                tooltip="Tooltip",
+                description="Description",
+                disabled=True,
+            ),
+        ]
+
+        for widget in widgets:
+            assert widget.tooltip == "Tooltip"
+            assert widget.description == "Description"
+            assert widget.disabled is True
+            result = widget.to_dict()
+            assert result["tooltip"] == "Tooltip"
+            assert result["description"] == "Description"
+            assert result["disabled"] is True
+
+    def test_select_with_complex_items(self):
+        """Test Select with complex item labels and values."""
+        items = {
+            "option_1": "Option One with Spaces",
+            "option_2": "Option Two (with parentheses)",
+            "option_3": "Option Three - with dashes",
+        }
+        select = Select(id="test_select", label="Choose", items=items)
+
+        result = select.to_dict()
+        assert len(result["items"]) == 3
+        assert {"label": "option_1", "value": "Option One with Spaces"} in result[
+            "items"
+        ]
+
+    def test_multiselect_initial_with_multiple_values(self):
+        """Test MultiSelect with multiple initial values."""
+        multi_select = MultiSelect(
+            id="test",
+            label="Choose",
+            values=["A", "B", "C", "D"],
+            initial=["A", "C", "D"],
+        )
+
+        assert len(multi_select.initial) == 3
+        assert "A" in multi_select.initial
+        assert "C" in multi_select.initial
+        assert "D" in multi_select.initial
+
+    def test_slider_with_negative_range(self):
+        """Test Slider with negative range."""
+        slider = Slider(id="test", label="Test", min=-10, max=10, initial=-5, step=1)
+
+        assert slider.min == -10
+        assert slider.max == 10
+        assert slider.initial == -5
+
+    def test_textinput_empty_initial_value(self):
+        """Test TextInput with empty string as initial value."""
+        text_input = TextInput(id="test", label="Test", initial="")
+
+        assert text_input.initial == ""
+        result = text_input.to_dict()
+        assert result["initial"] == ""

backend/tests/test_markdown.py

@@ -0,0 +1,322 @@
+import os
+import tempfile
+from unittest.mock import patch
+
+import pytest
+
+from chainlit.markdown import DEFAULT_MARKDOWN_STR, get_markdown_str, init_markdown
+
+
+class TestInitMarkdown:
+    """Test suite for init_markdown function."""
+
+    def test_init_markdown_creates_file(self):
+        """Test that init_markdown creates chainlit.md if it doesn't exist."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            init_markdown(tmpdir)
+
+            chainlit_md_path = os.path.join(tmpdir, "chainlit.md")
+            assert os.path.exists(chainlit_md_path)
+
+            # Verify content is the default markdown
+            with open(chainlit_md_path, encoding="utf-8") as f:
+                content = f.read()
+            assert content == DEFAULT_MARKDOWN_STR
+
+    def test_init_markdown_does_not_overwrite_existing(self):
+        """Test that init_markdown doesn't overwrite existing chainlit.md."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            chainlit_md_path = os.path.join(tmpdir, "chainlit.md")
+            custom_content = "# My Custom Markdown"
+
+            # Create existing file
+            with open(chainlit_md_path, "w", encoding="utf-8") as f:
+                f.write(custom_content)
+
+            # Call init_markdown
+            init_markdown(tmpdir)
+
+            # Verify content is unchanged
+            with open(chainlit_md_path, encoding="utf-8") as f:
+                content = f.read()
+            assert content == custom_content
+
+    def test_init_markdown_with_nonexistent_directory(self):
+        """Test init_markdown with a directory that doesn't exist."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            nonexistent_dir = os.path.join(tmpdir, "nonexistent")
+
+            # Should raise an error when trying to create file in nonexistent dir
+            with pytest.raises(FileNotFoundError):
+                init_markdown(nonexistent_dir)
+
+    def test_init_markdown_creates_utf8_file(self):
+        """Test that init_markdown creates file with UTF-8 encoding."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            init_markdown(tmpdir)
+
+            chainlit_md_path = os.path.join(tmpdir, "chainlit.md")
+
+            # Verify UTF-8 encoding by reading with explicit encoding
+            with open(chainlit_md_path, encoding="utf-8") as f:
+                content = f.read()
+
+            # Should contain emoji characters from DEFAULT_MARKDOWN_STR
+            assert "🚀" in content
+            assert "🤖" in content
+
+
+class TestGetMarkdownStr:
+    """Test suite for get_markdown_str function."""
+
+    def test_get_markdown_str_returns_default(self):
+        """Test get_markdown_str returns default chainlit.md content."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            chainlit_md_path = os.path.join(tmpdir, "chainlit.md")
+            content = "# Default Chainlit Markdown"
+
+            with open(chainlit_md_path, "w", encoding="utf-8") as f:
+                f.write(content)
+
+            result = get_markdown_str(tmpdir, "en")
+            assert result == content
+
+    def test_get_markdown_str_returns_translated(self):
+        """Test get_markdown_str returns translated markdown when available."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create default markdown
+            default_content = "# Default English"
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(default_content)
+
+            # Create translated markdown
+            translated_content = "# Français"
+            with open(
+                os.path.join(tmpdir, "chainlit_fr.md"), "w", encoding="utf-8"
+            ) as f:
+                f.write(translated_content)
+
+            result = get_markdown_str(tmpdir, "fr")
+            assert result == translated_content
+
+    def test_get_markdown_str_falls_back_to_default(self):
+        """Test get_markdown_str falls back to default when translation missing."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            default_content = "# Default English"
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(default_content)
+
+            # Request non-existent translation
+            with patch("chainlit.markdown.logger") as mock_logger:
+                result = get_markdown_str(tmpdir, "es")
+
+                assert result == default_content
+                mock_logger.warning.assert_called_once()
+                assert "es" in str(mock_logger.warning.call_args)
+
+    def test_get_markdown_str_returns_none_when_no_file(self):
+        """Test get_markdown_str returns None when no markdown file exists."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            result = get_markdown_str(tmpdir, "en")
+            assert result is None
+
+    def test_get_markdown_str_with_utf8_content(self):
+        """Test get_markdown_str handles UTF-8 content correctly."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            content = "# Welcome 欢迎 🎉\n\nこんにちは"
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(content)
+
+            result = get_markdown_str(tmpdir, "en")
+            assert result == content
+            assert "欢迎" in result
+            assert "こんにちは" in result
+
+    def test_get_markdown_str_prevents_path_traversal(self):
+        """Test get_markdown_str prevents path traversal attacks."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create default markdown
+            default_content = "# Default"
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(default_content)
+
+            # Try to access file outside root using path traversal
+            # The is_path_inside check should prevent this
+            result = get_markdown_str(tmpdir, "../../../etc/passwd")
+
+            # Should fall back to default since traversal is blocked
+            assert result == default_content
+
+    def test_get_markdown_str_with_multiple_languages(self):
+        """Test get_markdown_str with multiple language files."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create multiple language files
+            languages = {
+                "en": "# English",
+                "fr": "# Français",
+                "es": "# Español",
+                "ja": "# 日本語",
+            }
+
+            # Create default
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(languages["en"])
+
+            # Create translations
+            for lang, content in languages.items():
+                if lang != "en":
+                    path = os.path.join(tmpdir, f"chainlit_{lang}.md")
+                    with open(path, "w", encoding="utf-8") as f:
+                        f.write(content)
+
+            # Test each language
+            for lang, expected_content in languages.items():
+                result = get_markdown_str(tmpdir, lang)
+                assert result == expected_content
+
+    def test_get_markdown_str_with_empty_file(self):
+        """Test get_markdown_str with empty markdown file."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create empty file
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write("")
+
+            result = get_markdown_str(tmpdir, "en")
+            assert result == ""
+
+    def test_get_markdown_str_with_large_file(self):
+        """Test get_markdown_str with large markdown file."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create large content
+            large_content = "# Header\n" + ("Lorem ipsum dolor sit amet.\n" * 1000)
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(large_content)
+
+            result = get_markdown_str(tmpdir, "en")
+            assert result == large_content
+            assert len(result) > 10000
+
+
+class TestDefaultMarkdownStr:
+    """Test suite for DEFAULT_MARKDOWN_STR constant."""
+
+    def test_default_markdown_str_is_string(self):
+        """Test that DEFAULT_MARKDOWN_STR is a string."""
+        assert isinstance(DEFAULT_MARKDOWN_STR, str)
+
+    def test_default_markdown_str_not_empty(self):
+        """Test that DEFAULT_MARKDOWN_STR is not empty."""
+        assert len(DEFAULT_MARKDOWN_STR) > 0
+
+    def test_default_markdown_str_contains_welcome(self):
+        """Test that DEFAULT_MARKDOWN_STR contains welcome message."""
+        assert "Welcome to Chainlit" in DEFAULT_MARKDOWN_STR
+
+    def test_default_markdown_str_contains_links(self):
+        """Test that DEFAULT_MARKDOWN_STR contains useful links."""
+        assert "Documentation" in DEFAULT_MARKDOWN_STR
+        assert "Discord" in DEFAULT_MARKDOWN_STR
+        assert "https://docs.chainlit.io" in DEFAULT_MARKDOWN_STR
+
+    def test_default_markdown_str_is_valid_markdown(self):
+        """Test that DEFAULT_MARKDOWN_STR contains valid markdown syntax."""
+        assert "#" in DEFAULT_MARKDOWN_STR  # Headers
+        assert "**" in DEFAULT_MARKDOWN_STR  # Bold
+        assert "[" in DEFAULT_MARKDOWN_STR  # Links
+        assert "](" in DEFAULT_MARKDOWN_STR  # Link syntax
+
+
+class TestMarkdownEdgeCases:
+    """Test suite for markdown edge cases."""
+
+    def test_init_markdown_with_special_characters_in_path(self):
+        """Test init_markdown with special characters in directory path."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            special_dir = os.path.join(tmpdir, "test dir with spaces")
+            os.makedirs(special_dir)
+
+            init_markdown(special_dir)
+
+            chainlit_md_path = os.path.join(special_dir, "chainlit.md")
+            assert os.path.exists(chainlit_md_path)
+
+    def test_get_markdown_str_with_symlink(self):
+        """Test get_markdown_str with symlinked markdown file."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create original file
+            original_dir = os.path.join(tmpdir, "original")
+            os.makedirs(original_dir)
+            original_file = os.path.join(original_dir, "chainlit.md")
+            content = "# Original Content"
+            with open(original_file, "w", encoding="utf-8") as f:
+                f.write(content)
+
+            # Create symlink directory
+            link_dir = os.path.join(tmpdir, "link")
+            os.makedirs(link_dir)
+            link_file = os.path.join(link_dir, "chainlit.md")
+
+            # Create symlink (skip on Windows if no permissions)
+            try:
+                os.symlink(original_file, link_file)
+                result = get_markdown_str(link_dir, "en")
+                assert result == content
+            except OSError:
+                pytest.skip("Symlink creation not supported")
+
+    def test_get_markdown_str_with_relative_path(self):
+        """Test get_markdown_str with relative path."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            content = "# Test Content"
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(content)
+
+            # Use relative path
+            original_cwd = os.getcwd()
+            try:
+                os.chdir(tmpdir)
+                result = get_markdown_str(".", "en")
+                assert result == content
+            finally:
+                os.chdir(original_cwd)
+
+    def test_get_markdown_str_language_case_sensitivity(self):
+        """Test get_markdown_str language code is used as-is in filename."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            default_content = "# Default"
+            with open(os.path.join(tmpdir, "chainlit.md"), "w", encoding="utf-8") as f:
+                f.write(default_content)
+
+            # Create a language file
+            fr_content = "# Français"
+            with open(
+                os.path.join(tmpdir, "chainlit_fr.md"), "w", encoding="utf-8"
+            ) as f:
+                f.write(fr_content)
+
+            # Test exact match - should get the file
+            result = get_markdown_str(tmpdir, "fr")
+            assert result == fr_content
+
+            # Test different case that doesn't exist - should fall back to default
+            # Note: On case-insensitive file systems (Windows), this might still find the file
+            # On case-sensitive file systems (Linux), it will fall back to default
+            with patch("chainlit.markdown.logger"):
+                result_different = get_markdown_str(tmpdir, "es")
+                assert result_different == default_content
+
+    def test_init_markdown_concurrent_calls(self):
+        """Test init_markdown with concurrent calls (race condition)."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Call init_markdown multiple times
+            init_markdown(tmpdir)
+            init_markdown(tmpdir)
+            init_markdown(tmpdir)
+
+            # Should only have one file with default content
+            chainlit_md_path = os.path.join(tmpdir, "chainlit.md")
+            assert os.path.exists(chainlit_md_path)
+
+            with open(chainlit_md_path, encoding="utf-8") as f:
+                content = f.read()
+            assert content == DEFAULT_MARKDOWN_STR

backend/tests/test_mcp.py

@@ -0,0 +1,446 @@
+import sys
+from unittest.mock import patch
+
+import pytest
+from pydantic import ValidationError
+
+from chainlit.mcp import (
+    HttpMcpConnection,
+    SseMcpConnection,
+    StdioMcpConnection,
+    validate_mcp_command,
+)
+
+
+class TestStdioMcpConnection:
+    """Test suite for StdioMcpConnection model."""
+
+    def test_stdio_connection_initialization(self):
+        """Test StdioMcpConnection initialization."""
+        connection = StdioMcpConnection(
+            name="test_server", command="python", args=["-m", "mcp_server"]
+        )
+
+        assert connection.name == "test_server"
+        assert connection.command == "python"
+        assert connection.args == ["-m", "mcp_server"]
+        assert connection.clientType == "stdio"
+
+    def test_stdio_connection_with_empty_args(self):
+        """Test StdioMcpConnection with empty args list."""
+        connection = StdioMcpConnection(name="test_server", command="node", args=[])
+
+        assert connection.args == []
+        assert connection.clientType == "stdio"
+
+    def test_stdio_connection_requires_name(self):
+        """Test that StdioMcpConnection requires name."""
+        with pytest.raises(ValidationError):
+            StdioMcpConnection(command="python", args=[])
+
+    def test_stdio_connection_requires_command(self):
+        """Test that StdioMcpConnection requires command."""
+        with pytest.raises(ValidationError):
+            StdioMcpConnection(name="test_server", args=[])
+
+    def test_stdio_connection_requires_args(self):
+        """Test that StdioMcpConnection requires args."""
+        with pytest.raises(ValidationError):
+            StdioMcpConnection(name="test_server", command="python")
+
+    def test_stdio_connection_client_type_is_literal(self):
+        """Test that clientType is always 'stdio'."""
+        connection = StdioMcpConnection(name="test_server", command="python", args=[])
+
+        assert connection.clientType == "stdio"
+
+    def test_stdio_connection_serialization(self):
+        """Test StdioMcpConnection serialization."""
+        connection = StdioMcpConnection(
+            name="test_server", command="python", args=["-m", "server"]
+        )
+
+        data = connection.model_dump()
+
+        assert data["name"] == "test_server"
+        assert data["command"] == "python"
+        assert data["args"] == ["-m", "server"]
+        assert data["clientType"] == "stdio"
+
+
+class TestSseMcpConnection:
+    """Test suite for SseMcpConnection model."""
+
+    def test_sse_connection_initialization(self):
+        """Test SseMcpConnection initialization."""
+        connection = SseMcpConnection(name="test_server", url="https://example.com/mcp")
+
+        assert connection.name == "test_server"
+        assert connection.url == "https://example.com/mcp"
+        assert connection.headers is None
+        assert connection.clientType == "sse"
+
+    def test_sse_connection_with_headers(self):
+        """Test SseMcpConnection with headers."""
+        headers = {"Authorization": "Bearer token123", "X-Custom": "value"}
+        connection = SseMcpConnection(
+            name="test_server", url="https://example.com/mcp", headers=headers
+        )
+
+        assert connection.headers == headers
+
+    def test_sse_connection_requires_name(self):
+        """Test that SseMcpConnection requires name."""
+        with pytest.raises(ValidationError):
+            SseMcpConnection(url="https://example.com/mcp")
+
+    def test_sse_connection_requires_url(self):
+        """Test that SseMcpConnection requires url."""
+        with pytest.raises(ValidationError):
+            SseMcpConnection(name="test_server")
+
+    def test_sse_connection_client_type_is_literal(self):
+        """Test that clientType is always 'sse'."""
+        connection = SseMcpConnection(name="test_server", url="https://example.com/mcp")
+
+        assert connection.clientType == "sse"
+
+    def test_sse_connection_serialization(self):
+        """Test SseMcpConnection serialization."""
+        headers = {"Authorization": "Bearer token"}
+        connection = SseMcpConnection(
+            name="test_server", url="https://example.com/mcp", headers=headers
+        )
+
+        data = connection.model_dump()
+
+        assert data["name"] == "test_server"
+        assert data["url"] == "https://example.com/mcp"
+        assert data["headers"] == headers
+        assert data["clientType"] == "sse"
+
+
+class TestHttpMcpConnection:
+    """Test suite for HttpMcpConnection model."""
+
+    def test_http_connection_initialization(self):
+        """Test HttpMcpConnection initialization."""
+        connection = HttpMcpConnection(
+            name="test_server", url="https://example.com/mcp"
+        )
+
+        assert connection.name == "test_server"
+        assert connection.url == "https://example.com/mcp"
+        assert connection.headers is None
+        assert connection.clientType == "streamable-http"
+
+    def test_http_connection_with_headers(self):
+        """Test HttpMcpConnection with headers."""
+        headers = {
+            "Authorization": "Bearer token123",
+            "Content-Type": "application/json",
+        }
+        connection = HttpMcpConnection(
+            name="test_server", url="https://example.com/mcp", headers=headers
+        )
+
+        assert connection.headers == headers
+
+    def test_http_connection_requires_name(self):
+        """Test that HttpMcpConnection requires name."""
+        with pytest.raises(ValidationError):
+            HttpMcpConnection(url="https://example.com/mcp")
+
+    def test_http_connection_requires_url(self):
+        """Test that HttpMcpConnection requires url."""
+        with pytest.raises(ValidationError):
+            HttpMcpConnection(name="test_server")
+
+    def test_http_connection_client_type_is_literal(self):
+        """Test that clientType is always 'streamable-http'."""
+        connection = HttpMcpConnection(
+            name="test_server", url="https://example.com/mcp"
+        )
+
+        assert connection.clientType == "streamable-http"
+
+    def test_http_connection_serialization(self):
+        """Test HttpMcpConnection serialization."""
+        headers = {"Authorization": "Bearer token"}
+        connection = HttpMcpConnection(
+            name="test_server", url="https://example.com/mcp", headers=headers
+        )
+
+        data = connection.model_dump()
+
+        assert data["name"] == "test_server"
+        assert data["url"] == "https://example.com/mcp"
+        assert data["headers"] == headers
+        assert data["clientType"] == "streamable-http"
+
+
+class TestValidateMcpCommand:
+    """Test suite for validate_mcp_command function."""
+
+    def test_validate_simple_command(self):
+        """Test validation of a simple command."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python", "node"]
+
+            env, executable, args = validate_mcp_command("python -m mcp_server")
+
+            assert env == {}
+            assert executable == "python"
+            assert args == ["-m", "mcp_server"]
+
+    def test_validate_command_with_path(self):
+        """Test validation of command with full path."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            env, executable, args = validate_mcp_command(
+                "/usr/bin/python -m mcp_server"
+            )
+
+            assert env == {}
+            assert executable == "/usr/bin/python"
+            assert args == ["-m", "mcp_server"]
+
+    def test_validate_command_with_windows_path(self):
+        """Test validation of command with Windows path (using forward slashes or quoted)."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python.exe"]
+
+            # Use forward slashes which work on Windows and don't get escaped by shlex
+            env, executable, args = validate_mcp_command(
+                "C:/Python/python.exe -m mcp_server"
+            )
+
+            assert env == {}
+            assert executable == "C:/Python/python.exe"
+            assert args == ["-m", "mcp_server"]
+
+    def test_validate_command_with_env_vars(self):
+        """Test validation of command with environment variables."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["node"]
+
+            env, executable, args = validate_mcp_command(
+                "MY_VAR=value NODE_ENV=production node server.js"
+            )
+
+            assert env == {"MY_VAR": "value", "NODE_ENV": "production"}
+            assert executable == "node"
+            assert args == ["server.js"]
+
+    def test_validate_command_with_env_var_with_spaces(self):
+        """Test validation of command with env var containing spaces."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            env, executable, args = validate_mcp_command(
+                'MY_VAR="value with spaces" python script.py'
+            )
+
+            assert env == {"MY_VAR": "value with spaces"}
+            assert executable == "python"
+            assert args == ["script.py"]
+
+    def test_validate_command_with_quoted_args(self):
+        """Test validation of command with quoted arguments."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            env, executable, args = validate_mcp_command(
+                'python script.py --arg "value with spaces"'
+            )
+
+            assert env == {}
+            assert executable == "python"
+            assert args == ["script.py", "--arg", "value with spaces"]
+
+    def test_validate_command_with_multiple_args(self):
+        """Test validation of command with multiple arguments."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["node"]
+
+            env, executable, args = validate_mcp_command(
+                "node server.js --port 3000 --host localhost"
+            )
+
+            assert env == {}
+            assert executable == "node"
+            assert args == ["server.js", "--port", "3000", "--host", "localhost"]
+
+    def test_validate_command_not_in_allowed_list(self):
+        """Test that validation fails for disallowed executable."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python", "node"]
+
+            with pytest.raises(ValueError, match="Only commands in"):
+                validate_mcp_command("bash script.sh")
+
+    def test_validate_empty_command(self):
+        """Test that validation fails for empty command."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            with pytest.raises(ValueError, match="Empty command string"):
+                validate_mcp_command("")
+
+    def test_validate_command_with_invalid_syntax(self):
+        """Test that validation fails for invalid command syntax."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            with pytest.raises(ValueError, match="Invalid command string"):
+                validate_mcp_command('python "unclosed quote')
+
+    def test_validate_command_with_none_allowed_executables(self):
+        """Test validation when allowed_executables is None (all allowed)."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = None
+
+            env, executable, args = validate_mcp_command("any_command --arg value")
+
+            assert env == {}
+            assert executable == "any_command"
+            assert args == ["--arg", "value"]
+
+    def test_validate_command_with_invalid_env_var_format(self):
+        """Test that validation fails for invalid env var format."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            with pytest.raises(ValueError, match="Invalid environment variable format"):
+                validate_mcp_command("INVALID_ENV python script.py")
+
+    def test_validate_command_with_complex_env_vars(self):
+        """Test validation with complex environment variables."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            env, executable, args = validate_mcp_command(
+                'API_KEY=sk-123456 BASE_URL="https://api.example.com" python app.py'
+            )
+
+            assert env == {
+                "API_KEY": "sk-123456",
+                "BASE_URL": "https://api.example.com",
+            }
+            assert executable == "python"
+            assert args == ["app.py"]
+
+    def test_validate_command_with_equals_in_arg(self):
+        """Test validation with equals sign in argument."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            env, executable, args = validate_mcp_command(
+                "python script.py --config=value"
+            )
+
+            assert env == {}
+            assert executable == "python"
+            assert args == ["script.py", "--config=value"]
+
+    def test_validate_command_preserves_arg_order(self):
+        """Test that argument order is preserved."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["node"]
+
+            _, _, args = validate_mcp_command(
+                "node app.js arg1 arg2 arg3 --flag1 --flag2"
+            )
+
+            assert args == ["app.js", "arg1", "arg2", "arg3", "--flag1", "--flag2"]
+
+
+class TestMcpConnectionEdgeCases:
+    """Test suite for MCP connection edge cases."""
+
+    def test_stdio_connection_with_complex_args(self):
+        """Test StdioMcpConnection with complex arguments."""
+        connection = StdioMcpConnection(
+            name="complex_server",
+            command="python",
+            args=[
+                "-m",
+                "mcp_server",
+                "--config",
+                "/path/to/config.json",
+                "--verbose",
+            ],
+        )
+
+        assert len(connection.args) == 5
+        assert connection.args[0] == "-m"
+        assert connection.args[3] == "/path/to/config.json"
+
+    def test_sse_connection_with_multiple_headers(self):
+        """Test SseMcpConnection with multiple headers."""
+        headers = {
+            "Authorization": "Bearer token",
+            "X-API-Key": "key123",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+        connection = SseMcpConnection(
+            name="multi_header_server", url="https://api.example.com", headers=headers
+        )
+
+        assert len(connection.headers) == 4
+        assert connection.headers["Authorization"] == "Bearer token"
+        assert connection.headers["X-API-Key"] == "key123"
+
+    def test_http_connection_with_localhost_url(self):
+        """Test HttpMcpConnection with localhost URL."""
+        connection = HttpMcpConnection(
+            name="local_server", url="http://localhost:8000/mcp"
+        )
+
+        assert connection.url == "http://localhost:8000/mcp"
+
+    def test_connection_names_can_be_descriptive(self):
+        """Test that connection names can be descriptive strings."""
+        stdio_conn = StdioMcpConnection(
+            name="My Custom MCP Server (Python)", command="python", args=[]
+        )
+        sse_conn = SseMcpConnection(
+            name="Production API Server", url="https://api.example.com"
+        )
+        http_conn = HttpMcpConnection(
+            name="Development Server - Local", url="http://localhost:3000"
+        )
+
+        assert "Python" in stdio_conn.name
+        assert "Production" in sse_conn.name
+        assert "Development" in http_conn.name
+
+    def test_validate_command_with_special_characters_in_path(self):
+        """Test validation with special characters in path."""
+        mcp_module = sys.modules["chainlit.mcp"]
+        with patch.object(mcp_module, "config") as mock_config:
+            mock_config.features.mcp.stdio.allowed_executables = ["python"]
+
+            _, executable, args = validate_mcp_command(
+                "/opt/my-app/bin/python script.py"
+            )
+
+            assert executable == "/opt/my-app/bin/python"
+            assert args == ["script.py"]

backend/tests/test_message.py

@@ -0,0 +1,757 @@
+import asyncio
+import json
+from contextlib import contextmanager
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from chainlit.action import Action
+from chainlit.context import ChainlitContext, context_var
+from chainlit.message import (
+    AskActionMessage,
+    AskElementMessage,
+    AskFileMessage,
+    AskUserMessage,
+    ErrorMessage,
+    Message,
+    MessageBase,
+)
+
+
+@contextmanager
+def mock_chainlit_context(session=None):
+    """Context manager to set up and tear down Chainlit context."""
+    mock_loop = Mock(spec=asyncio.AbstractEventLoop)
+    mock_session = session or Mock()
+    mock_session.thread_id = "thread_123"
+
+    with patch("asyncio.get_running_loop", return_value=mock_loop):
+        mock_emitter = AsyncMock()
+        mock_context = ChainlitContext(session=mock_session, emitter=mock_emitter)
+        token = context_var.set(mock_context)
+        try:
+            yield mock_context
+        finally:
+            context_var.reset(token)
+
+
+class TestMessageBase:
+    """Test suite for MessageBase class."""
+
+    def test_post_init_sets_thread_id(self):
+        """Test that __post_init__ sets thread_id from session."""
+        with mock_chainlit_context():
+            msg = Message(content="test")
+            assert msg.thread_id == "thread_123"
+
+    def test_post_init_generates_id_if_not_provided(self):
+        """Test that __post_init__ generates UUID if id not provided."""
+        with mock_chainlit_context():
+            msg = Message(content="test")
+            assert msg.id is not None
+            assert len(msg.id) == 36
+
+    def test_post_init_uses_provided_id(self):
+        """Test that __post_init__ uses provided id."""
+        with mock_chainlit_context():
+            msg = Message(content="test", id="custom_id")
+            assert msg.id == "custom_id"
+
+    def test_from_dict_creates_message(self):
+        """Test creating message from dictionary."""
+        step_dict = {
+            "id": "msg_123",
+            "parentId": "parent_123",
+            "createdAt": "2024-01-01T00:00:00Z",
+            "output": "Hello world",
+            "name": "Assistant",
+            "command": "/test",
+            "type": "user_message",
+            "language": "python",
+            "metadata": {"key": "value"},
+        }
+
+        with mock_chainlit_context():
+            msg = MessageBase.from_dict(step_dict)
+
+            assert msg.id == "msg_123"
+            assert msg.parent_id == "parent_123"
+            assert msg.created_at == "2024-01-01T00:00:00Z"
+            assert msg.content == "Hello world"
+            assert msg.author == "Assistant"
+            assert msg.command == "/test"
+            assert msg.type == "user_message"
+            assert msg.language == "python"
+            assert msg.metadata == {"key": "value"}
+
+    def test_from_dict_with_minimal_data(self):
+        """Test from_dict with minimal required fields."""
+        step_dict = {
+            "id": "msg_123",
+            "createdAt": "2024-01-01T00:00:00Z",
+            "output": "Hello",
+        }
+
+        with mock_chainlit_context():
+            with patch("chainlit.message.config") as mock_config:
+                mock_config.ui.name = "DefaultBot"
+                msg = MessageBase.from_dict(step_dict)
+
+                assert msg.id == "msg_123"
+                assert msg.content == "Hello"
+                assert msg.author == "DefaultBot"
+                assert msg.type == "assistant_message"
+
+    def test_to_dict_returns_step_dict(self):
+        """Test converting message to dictionary."""
+        with mock_chainlit_context():
+            msg = Message(
+                content="Test content",
+                author="TestBot",
+                language="python",
+                type="user_message",
+                metadata={"key": "value"},
+                tags=["tag1", "tag2"],
+                id="msg_123",
+                parent_id="parent_123",
+                command="/test",
+            )
+            msg.created_at = "2024-01-01T00:00:00Z"
+
+            result = msg.to_dict()
+
+            assert result["id"] == "msg_123"
+            assert result["threadId"] == "thread_123"
+            assert result["parentId"] == "parent_123"
+            assert result["createdAt"] == "2024-01-01T00:00:00Z"
+            assert result["command"] == "/test"
+            assert result["output"] == "Test content"
+            assert result["name"] == "TestBot"
+            assert result["type"] == "user_message"
+            assert result["language"] == "python"
+            assert result["streaming"] is False
+            assert result["isError"] is False
+            assert result["waitForAnswer"] is False
+            assert result["metadata"] == {"key": "value"}
+            assert result["tags"] == ["tag1", "tag2"]
+
+    @pytest.mark.asyncio
+    async def test_update_stops_streaming(self):
+        """Test that update stops streaming."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="test")
+            msg.streaming = True
+
+            with patch("chainlit.message.chat_context") as mock_chat_ctx:
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    result = await msg.update()
+
+                    assert msg.streaming is False
+                    assert result is True
+                    mock_chat_ctx.add.assert_called_once_with(msg)
+                    ctx.emitter.update_step.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_update_with_data_layer(self):
+        """Test update with data layer."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="test")
+            mock_data_layer = AsyncMock()
+
+            with patch("chainlit.message.chat_context"):
+                with patch(
+                    "chainlit.message.get_data_layer", return_value=mock_data_layer
+                ):
+                    with patch("asyncio.create_task") as mock_create_task:
+                        await msg.update()
+
+                        mock_create_task.assert_called_once()
+                        ctx.emitter.update_step.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_remove_from_chat_context(self):
+        """Test removing message from chat context."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="test", id="msg_123")
+
+            with patch("chainlit.message.chat_context") as mock_chat_ctx:
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    result = await msg.remove()
+
+                    assert result is True
+                    mock_chat_ctx.remove.assert_called_once_with(msg)
+                    ctx.emitter.delete_step.assert_called_once()
+
+
+class TestMessage:
+    """Test suite for Message class."""
+
+    def test_message_with_string_content(self):
+        """Test creating message with string content."""
+        with mock_chainlit_context():
+            msg = Message(content="Hello world")
+
+            assert msg.content == "Hello world"
+            assert msg.language is None
+
+    def test_message_with_dict_content(self):
+        """Test creating message with dict content."""
+        with mock_chainlit_context():
+            content_dict = {"key": "value", "number": 42}
+            msg = Message(content=content_dict)
+
+            expected = json.dumps(content_dict, indent=4, ensure_ascii=False)
+            assert msg.content == expected
+            assert msg.language == "json"
+
+    def test_message_with_non_serializable_dict(self):
+        """Test message with non-JSON-serializable dict."""
+        with mock_chainlit_context():
+
+            class NonSerializable:
+                pass
+
+            content_dict = {"obj": NonSerializable()}
+            msg = Message(content=content_dict)
+
+            assert msg.language == "text"
+            assert "NonSerializable" in msg.content
+
+    def test_message_with_non_string_content(self):
+        """Test message with non-string, non-dict content."""
+        with mock_chainlit_context():
+            msg = Message(content=12345)
+
+            assert msg.content == "12345"
+            assert msg.language == "text"
+
+    def test_message_with_custom_author(self):
+        """Test message with custom author."""
+        with mock_chainlit_context():
+            msg = Message(content="test", author="CustomBot")
+
+            assert msg.author == "CustomBot"
+
+    def test_message_with_default_author(self):
+        """Test message uses default author from config."""
+        with mock_chainlit_context():
+            with patch("chainlit.message.config") as mock_config:
+                mock_config.ui.name = "DefaultBot"
+                msg = Message(content="test")
+
+                assert msg.author == "DefaultBot"
+
+    def test_message_with_actions(self):
+        """Test message with actions."""
+        with mock_chainlit_context():
+            action1 = Mock(spec=Action)
+            action2 = Mock(spec=Action)
+            msg = Message(content="test", actions=[action1, action2])
+
+            assert len(msg.actions) == 2
+            assert action1 in msg.actions
+            assert action2 in msg.actions
+
+    def test_message_with_elements(self):
+        """Test message with elements."""
+        with mock_chainlit_context():
+            element1 = Mock()
+            element2 = Mock()
+            msg = Message(content="test", elements=[element1, element2])
+
+            assert len(msg.elements) == 2
+            assert element1 in msg.elements
+            assert element2 in msg.elements
+
+    @pytest.mark.asyncio
+    async def test_message_send(self):
+        """Test sending a message."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="test")
+
+            with patch("chainlit.message.chat_context") as mock_chat_ctx:
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    with patch("chainlit.message.config") as mock_config:
+                        mock_config.code.author_rename = None
+
+                        result = await msg.send()
+
+                        assert result == msg
+                        assert msg.created_at is not None
+                        assert msg.streaming is False
+                        mock_chat_ctx.add.assert_called_once_with(msg)
+                        ctx.emitter.send_step.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_message_send_with_author_rename(self):
+        """Test sending message with author rename."""
+        with mock_chainlit_context():
+            msg = Message(content="test", author="OldName")
+
+            async def rename_author(name):
+                return "NewName"
+
+            with patch("chainlit.message.chat_context"):
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    with patch("chainlit.message.config") as mock_config:
+                        mock_config.code.author_rename = rename_author
+
+                        await msg.send()
+
+                        assert msg.author == "NewName"
+
+    @pytest.mark.asyncio
+    async def test_message_send_with_actions_and_elements(self):
+        """Test sending message with actions and elements."""
+        with mock_chainlit_context():
+            action = AsyncMock(spec=Action)
+            element = AsyncMock()
+            msg = Message(content="test", actions=[action], elements=[element])
+
+            with patch("chainlit.message.chat_context"):
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    with patch("chainlit.message.config") as mock_config:
+                        mock_config.code.author_rename = None
+
+                        await msg.send()
+
+                        action.send.assert_called_once()
+                        element.send.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_message_update_with_actions(self):
+        """Test updating message with new actions."""
+        with mock_chainlit_context():
+            action1 = AsyncMock(spec=Action)
+            action1.forId = None
+            action2 = AsyncMock(spec=Action)
+            action2.forId = "existing_id"
+
+            msg = Message(content="test", actions=[action1, action2])
+
+            with patch("chainlit.message.chat_context"):
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    result = await msg.update()
+
+                    assert result is True
+                    action1.send.assert_called_once()
+                    action2.send.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_message_remove_actions(self):
+        """Test removing all actions from message."""
+        with mock_chainlit_context():
+            action1 = AsyncMock(spec=Action)
+            action2 = AsyncMock(spec=Action)
+            msg = Message(content="test", actions=[action1, action2])
+
+            await msg.remove_actions()
+
+            action1.remove.assert_called_once()
+            action2.remove.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_stream_token_starts_streaming(self):
+        """Test that stream_token starts streaming."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="")
+
+            await msg.stream_token("Hello")
+
+            assert msg.streaming is True
+            assert msg.content == "Hello"
+            ctx.emitter.stream_start.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_stream_token_appends_content(self):
+        """Test that stream_token appends to content."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="Hello")
+            msg.streaming = True
+
+            await msg.stream_token(" world")
+
+            assert msg.content == "Hello world"
+            ctx.emitter.send_token.assert_called_once_with(
+                id=msg.id, token=" world", is_sequence=False
+            )
+
+    @pytest.mark.asyncio
+    async def test_stream_token_with_sequence(self):
+        """Test stream_token with is_sequence=True."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="Old content")
+            msg.streaming = True
+
+            await msg.stream_token("New content", is_sequence=True)
+
+            assert msg.content == "New content"
+            ctx.emitter.send_token.assert_called_once_with(
+                id=msg.id, token="New content", is_sequence=True
+            )
+
+    @pytest.mark.asyncio
+    async def test_stream_token_ignores_empty_token(self):
+        """Test that empty tokens are ignored."""
+        with mock_chainlit_context() as ctx:
+            msg = Message(content="test")
+
+            await msg.stream_token("")
+
+            assert msg.content == "test"
+            ctx.emitter.stream_start.assert_not_called()
+
+
+class TestErrorMessage:
+    """Test suite for ErrorMessage class."""
+
+    def test_error_message_initialization(self):
+        """Test ErrorMessage initialization."""
+        with mock_chainlit_context():
+            msg = ErrorMessage(content="An error occurred")
+
+            assert msg.content == "An error occurred"
+            assert msg.author is not None
+            assert msg.type == "assistant_message"
+            assert msg.is_error is True
+            assert msg.fail_on_persist_error is False
+
+    def test_error_message_with_custom_author(self):
+        """Test ErrorMessage with custom author."""
+        with mock_chainlit_context():
+            msg = ErrorMessage(content="Error", author="ErrorBot")
+
+            assert msg.author == "ErrorBot"
+
+    def test_error_message_with_fail_on_persist(self):
+        """Test ErrorMessage with fail_on_persist_error=True."""
+        with mock_chainlit_context():
+            msg = ErrorMessage(content="Error", fail_on_persist_error=True)
+
+            assert msg.fail_on_persist_error is True
+
+    @pytest.mark.asyncio
+    async def test_error_message_send(self):
+        """Test sending error message."""
+        with mock_chainlit_context() as ctx:
+            msg = ErrorMessage(content="Error occurred")
+
+            with patch("chainlit.message.chat_context"):
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    with patch("chainlit.message.config") as mock_config:
+                        mock_config.code.author_rename = None
+
+                        result = await msg.send()
+
+                        assert result == msg
+                        ctx.emitter.send_step.assert_called_once()
+
+
+class TestAskUserMessage:
+    """Test suite for AskUserMessage class."""
+
+    def test_ask_user_message_initialization(self):
+        """Test AskUserMessage initialization."""
+        with mock_chainlit_context():
+            msg = AskUserMessage(content="What is your name?")
+
+            assert msg.content == "What is your name?"
+            assert msg.author is not None
+            assert msg.timeout == 60
+            assert msg.raise_on_timeout is False
+
+    def test_ask_user_message_with_custom_timeout(self):
+        """Test AskUserMessage with custom timeout."""
+        with mock_chainlit_context():
+            msg = AskUserMessage(content="Question?", timeout=120)
+
+            assert msg.timeout == 120
+
+    @pytest.mark.asyncio
+    async def test_ask_user_message_send(self):
+        """Test sending AskUserMessage."""
+        with mock_chainlit_context() as ctx:
+            msg = AskUserMessage(content="Question?")
+            ctx.emitter.send_ask_user = AsyncMock(return_value={"output": "Answer"})
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    result = await msg.send()
+
+                    assert result == {"output": "Answer"}
+                    assert msg.wait_for_answer is False
+                    ctx.emitter.send_ask_user.assert_called_once()
+
+
+class TestAskFileMessage:
+    """Test suite for AskFileMessage class."""
+
+    def test_ask_file_message_initialization(self):
+        """Test AskFileMessage initialization."""
+        with mock_chainlit_context():
+            with patch("chainlit.message.config") as mock_config:
+                mock_config.ui.name = "Bot"
+                msg = AskFileMessage(
+                    content="Upload a file", accept=["text/plain", "application/pdf"]
+                )
+
+                assert msg.content == "Upload a file"
+                assert msg.accept == ["text/plain", "application/pdf"]
+                assert msg.max_size_mb == 2
+                assert msg.max_files == 1
+
+    def test_ask_file_message_with_custom_limits(self):
+        """Test AskFileMessage with custom limits."""
+        with mock_chainlit_context():
+            msg = AskFileMessage(
+                content="Upload", accept=["image/*"], max_size_mb=10, max_files=5
+            )
+
+            assert msg.max_size_mb == 10
+            assert msg.max_files == 5
+
+    @pytest.mark.asyncio
+    async def test_ask_file_message_send_with_response(self):
+        """Test AskFileMessage send with file response."""
+        with mock_chainlit_context() as ctx:
+            msg = AskFileMessage(content="Upload", accept=["text/plain"])
+            file_response = [
+                {
+                    "id": "file_123",
+                    "name": "test.txt",
+                    "path": "/path/to/test.txt",
+                    "size": 1024,
+                    "type": "text/plain",
+                }
+            ]
+            ctx.emitter.send_ask_user = AsyncMock(return_value=file_response)
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    result = await msg.send()
+
+                    assert result is not None
+                    assert len(result) == 1
+                    assert result[0].id == "file_123"
+                    assert result[0].name == "test.txt"
+                    assert result[0].path == "/path/to/test.txt"
+
+    @pytest.mark.asyncio
+    async def test_ask_file_message_send_with_no_response(self):
+        """Test AskFileMessage send with no response."""
+        with mock_chainlit_context() as ctx:
+            msg = AskFileMessage(content="Upload", accept=["text/plain"])
+            ctx.emitter.send_ask_user = AsyncMock(return_value=None)
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    result = await msg.send()
+
+                    assert result is None
+
+
+class TestAskActionMessage:
+    """Test suite for AskActionMessage class."""
+
+    def test_ask_action_message_initialization(self):
+        """Test AskActionMessage initialization."""
+        with mock_chainlit_context():
+            with patch("chainlit.message.config") as mock_config:
+                mock_config.ui.name = "Bot"
+                action1 = Mock(spec=Action)
+                action2 = Mock(spec=Action)
+                msg = AskActionMessage(
+                    content="Choose an action", actions=[action1, action2]
+                )
+
+                assert msg.content == "Choose an action"
+                assert len(msg.actions) == 2
+
+    @pytest.mark.asyncio
+    async def test_ask_action_message_send_with_response(self):
+        """Test AskActionMessage send with action response."""
+        with mock_chainlit_context() as ctx:
+            action = AsyncMock(spec=Action)
+            action.id = "action_123"
+            msg = AskActionMessage(content="Choose", actions=[action])
+            ctx.emitter.send_ask_user = AsyncMock(
+                return_value={"id": "action_123", "label": "Confirm"}
+            )
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    with patch("chainlit.message.chat_context"):
+                        result = await msg.send()
+
+                        assert result == {"id": "action_123", "label": "Confirm"}
+                        assert msg.content == "**Selected:** Confirm"
+                        action.send.assert_called_once()
+                        action.remove.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_ask_action_message_send_timeout(self):
+        """Test AskActionMessage send with timeout."""
+        with mock_chainlit_context() as ctx:
+            action = AsyncMock(spec=Action)
+            action.id = "action_123"
+            msg = AskActionMessage(content="Choose", actions=[action])
+            ctx.emitter.send_ask_user = AsyncMock(return_value=None)
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    with patch("chainlit.message.chat_context"):
+                        result = await msg.send()
+
+                        assert result is None
+                        assert msg.content == "Timed out: no action was taken"
+
+
+class TestAskElementMessage:
+    """Test suite for AskElementMessage class."""
+
+    def test_ask_element_message_initialization(self):
+        """Test AskElementMessage initialization."""
+        with mock_chainlit_context():
+            with patch("chainlit.message.config") as mock_config:
+                mock_config.ui.name = "Bot"
+                element = Mock()
+                msg = AskElementMessage(content="Submit form", element=element)
+
+                assert msg.content == "Submit form"
+                assert msg.element == element
+
+    @pytest.mark.asyncio
+    async def test_ask_element_message_send_submitted(self):
+        """Test AskElementMessage send with submitted response."""
+        with mock_chainlit_context() as ctx:
+            element = AsyncMock()
+            element.id = "element_123"
+            msg = AskElementMessage(content="Submit", element=element)
+            ctx.emitter.send_ask_user = AsyncMock(
+                return_value={"submitted": True, "data": {"field": "value"}}
+            )
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    with patch("chainlit.message.chat_context"):
+                        result = await msg.send()
+
+                        assert result == {"submitted": True, "data": {"field": "value"}}
+                        assert msg.content == "Thanks for submitting"
+                        element.send.assert_called_once()
+                        element.remove.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_ask_element_message_send_cancelled(self):
+        """Test AskElementMessage send with cancelled response."""
+        with mock_chainlit_context() as ctx:
+            element = AsyncMock()
+            element.id = "element_123"
+            msg = AskElementMessage(content="Submit", element=element)
+            ctx.emitter.send_ask_user = AsyncMock(return_value={"submitted": False})
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    with patch("chainlit.message.chat_context"):
+                        result = await msg.send()
+
+                        assert result == {"submitted": False}
+                        assert msg.content == "Cancelled"
+
+    @pytest.mark.asyncio
+    async def test_ask_element_message_send_timeout(self):
+        """Test AskElementMessage send with timeout."""
+        with mock_chainlit_context() as ctx:
+            element = AsyncMock()
+            element.id = "element_123"
+            msg = AskElementMessage(content="Submit", element=element)
+            ctx.emitter.send_ask_user = AsyncMock(return_value=None)
+
+            with patch("chainlit.message.get_data_layer", return_value=None):
+                with patch("chainlit.message.config") as mock_config:
+                    mock_config.code.author_rename = None
+
+                    with patch("chainlit.message.chat_context"):
+                        result = await msg.send()
+
+                        assert result is None
+                        assert msg.content == "Timed out"
+
+
+class TestMessageEdgeCases:
+    """Test suite for message edge cases."""
+
+    def test_message_with_none_content(self):
+        """Test message handles None content."""
+        with mock_chainlit_context():
+            msg = Message(content=None)
+            assert msg.content == "None"
+
+    def test_message_language_override(self):
+        """Test that dict content sets language to json."""
+        with mock_chainlit_context():
+            msg = Message(content={"key": "value"}, language="python")
+            # Dict content always sets language to json, overriding the parameter
+            assert msg.language == "json"
+
+    @pytest.mark.asyncio
+    async def test_message_send_with_none_content(self):
+        """Test sending message with None content."""
+        with mock_chainlit_context():
+            msg = Message(content="test")
+            msg.content = None
+
+            with patch("chainlit.message.chat_context"):
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    with patch("chainlit.message.config") as mock_config:
+                        mock_config.code.author_rename = None
+
+                        await msg.send()
+
+                        assert msg.content == ""
+
+    @pytest.mark.asyncio
+    async def test_ask_message_remove_clears_ask(self):
+        """Test that AskMessage remove clears ask state."""
+        with mock_chainlit_context() as ctx:
+            msg = AskUserMessage(content="Question?")
+
+            with patch("chainlit.message.chat_context"):
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    await msg.remove()
+
+                    ctx.emitter.clear.assert_called_once_with("clear_ask")
+
+    def test_message_metadata_and_tags(self):
+        """Test message with metadata and tags."""
+        with mock_chainlit_context():
+            metadata = {"key1": "value1", "key2": 123}
+            tags = ["important", "user-query"]
+            msg = Message(content="test", metadata=metadata, tags=tags)
+
+            assert msg.metadata == metadata
+            assert msg.tags == tags
+
+    def test_message_to_dict_with_none_metadata(self):
+        """Test to_dict with None metadata."""
+        with mock_chainlit_context():
+            msg = Message(content="test")
+            msg.metadata = None
+
+            result = msg.to_dict()
+
+            assert result["metadata"] == {}

backend/tests/test_modes.py

@@ -0,0 +1,279 @@
+"""Tests for Modes system functionality."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import chainlit as cl
+from chainlit.emitter import ChainlitEmitter
+from chainlit.mode import Mode, ModeOption
+
+
+@pytest.fixture
+def mock_modes():
+    """Fixture providing sample modes for testing."""
+    return [
+        Mode(
+            id="model",
+            name="Model",
+            options=[
+                ModeOption(
+                    id="gpt-4",
+                    name="GPT-4",
+                    description="Most capable model",
+                    icon="sparkles",
+                    default=True,
+                ),
+                ModeOption(
+                    id="gpt-3.5-turbo",
+                    name="GPT-3.5 Turbo",
+                    description="Fast and efficient",
+                    icon="bolt",
+                    default=False,
+                ),
+            ],
+        ),
+        Mode(
+            id="reasoning",
+            name="Reasoning",
+            options=[
+                ModeOption(id="high", name="High", description="Maximum depth"),
+                ModeOption(
+                    id="medium", name="Medium", description="Balanced", default=True
+                ),
+                ModeOption(id="low", name="Low", description="Quick responses"),
+            ],
+        ),
+    ]
+
+
+@pytest.mark.asyncio
+class TestModeOption:
+    """Test suite for ModeOption dataclass."""
+
+    def test_mode_option_required_fields(self):
+        """Test ModeOption with required fields only."""
+        option = ModeOption(id="test", name="Test Option")
+
+        assert option.id == "test"
+        assert option.name == "Test Option"
+        assert option.description is None
+        assert option.icon is None
+        assert option.default is False
+
+    def test_mode_option_all_fields(self):
+        """Test ModeOption with all fields."""
+        option = ModeOption(
+            id="gpt-4",
+            name="GPT-4",
+            description="Most capable model",
+            icon="sparkles",
+            default=True,
+        )
+
+        assert option.id == "gpt-4"
+        assert option.name == "GPT-4"
+        assert option.description == "Most capable model"
+        assert option.icon == "sparkles"
+        assert option.default is True
+
+    def test_mode_option_to_dict(self):
+        """Test ModeOption serialization."""
+        option = ModeOption(
+            id="test",
+            name="Test",
+            description="Test desc",
+            icon="star",
+            default=True,
+        )
+
+        option_dict = option.to_dict()
+
+        assert option_dict["id"] == "test"
+        assert option_dict["name"] == "Test"
+        assert option_dict["description"] == "Test desc"
+        assert option_dict["icon"] == "star"
+        assert option_dict["default"] is True
+
+
+@pytest.mark.asyncio
+class TestMode:
+    """Test suite for Mode dataclass."""
+
+    def test_mode_creation(self, mock_modes):
+        """Test Mode creation with options."""
+        mode = mock_modes[0]
+
+        assert mode.id == "model"
+        assert mode.name == "Model"
+        assert len(mode.options) == 2
+
+    def test_mode_to_dict(self, mock_modes):
+        """Test Mode serialization."""
+        mode = mock_modes[0]
+        mode_dict = mode.to_dict()
+
+        assert mode_dict["id"] == "model"
+        assert mode_dict["name"] == "Model"
+        assert len(mode_dict["options"]) == 2
+        assert mode_dict["options"][0]["id"] == "gpt-4"
+
+    def test_mode_default_option(self, mock_modes):
+        """Test finding default option in mode."""
+        mode = mock_modes[0]
+
+        default_option = next(
+            (opt for opt in mode.options if opt.default), mode.options[0]
+        )
+
+        assert default_option.id == "gpt-4"
+
+    def test_mode_fallback_to_first(self, mock_modes):
+        """Test fallback to first option when no default set."""
+        mode = Mode(
+            id="test",
+            name="Test",
+            options=[
+                ModeOption(id="opt1", name="Option 1"),
+                ModeOption(id="opt2", name="Option 2"),
+            ],
+        )
+
+        default_option = next(
+            (opt for opt in mode.options if opt.default),
+            mode.options[0] if mode.options else None,
+        )
+
+        assert default_option is not None
+        assert default_option.id == "opt1"
+
+
+@pytest.mark.asyncio
+class TestMessageWithModes:
+    """Test suite for Message with modes field."""
+
+    async def test_message_with_modes(self, mock_chainlit_context):
+        """Test that Message can be created with modes field."""
+        async with mock_chainlit_context:
+            modes = {"model": "gpt-4", "reasoning": "high"}
+            message = cl.Message(content="Test message", modes=modes)
+
+            assert message.modes == modes
+            assert message.content == "Test message"
+
+    async def test_message_to_dict_includes_modes(self, mock_chainlit_context):
+        """Test that Message.to_dict() includes the modes field."""
+        async with mock_chainlit_context:
+            modes = {"model": "gpt-4", "reasoning": "medium"}
+            message = cl.Message(content="Test", modes=modes)
+            message_dict = message.to_dict()
+
+            assert "modes" in message_dict
+            assert message_dict["modes"] == modes
+
+    async def test_message_from_dict_with_modes(self, mock_chainlit_context):
+        """Test that Message.from_dict() correctly handles modes field."""
+        async with mock_chainlit_context:
+            message_dict = {
+                "id": "test-id",
+                "content": "Test message",
+                "modes": {"model": "gpt-3.5-turbo", "reasoning": "low"},
+                "type": "user_message",
+                "createdAt": "2024-01-01T00:00:00",
+                "output": "Test message",
+            }
+            message = cl.Message.from_dict(message_dict)
+
+            assert message.modes == {"model": "gpt-3.5-turbo", "reasoning": "low"}
+            assert message.content == "Test message"
+
+    async def test_message_without_modes(self, mock_chainlit_context):
+        """Test that Message works without modes field (backward compatibility)."""
+        async with mock_chainlit_context:
+            message = cl.Message(content="Test message")
+
+            assert message.modes is None
+            message_dict = message.to_dict()
+            assert message_dict.get("modes") is None
+
+    async def test_message_send_with_modes(self, mock_chainlit_context):
+        """Test that sending a message with modes works."""
+        async with mock_chainlit_context as ctx:
+            modes = {"model": "gpt-4", "reasoning": "high"}
+            message = cl.Message(content="Test", modes=modes)
+
+            with patch("chainlit.message.chat_context") as mock_chat_ctx:
+                with patch("chainlit.message.get_data_layer", return_value=None):
+                    with patch("chainlit.message.config") as mock_config:
+                        mock_config.code.author_rename = None
+
+                        result = await message.send()
+
+                        assert result == message
+                        assert message.modes == modes
+                        mock_chat_ctx.add.assert_called_once_with(message)
+                        ctx.emitter.send_step.assert_called_once()
+
+                        # Verify the dict sent to emitter includes modes
+                        call_args = ctx.emitter.send_step.call_args[0][0]
+                        assert call_args["modes"] == modes
+
+
+@pytest.mark.asyncio
+class TestEmitterSetModes:
+    """Test suite for emitter set_modes functionality."""
+
+    async def test_set_modes(
+        self, mock_modes, mock_websocket_session: MagicMock
+    ) -> None:
+        """Test set_modes emits correct event."""
+        emitter = ChainlitEmitter(mock_websocket_session)
+        modes_dicts = [mode.to_dict() for mode in mock_modes]
+
+        await emitter.set_modes(mock_modes)
+
+        mock_websocket_session.emit.assert_called_once_with("set_modes", modes_dicts)
+
+    async def test_set_modes_empty_list(
+        self, mock_websocket_session: MagicMock
+    ) -> None:
+        """Test set_modes with empty list."""
+        emitter = ChainlitEmitter(mock_websocket_session)
+
+        await emitter.set_modes([])
+
+        mock_websocket_session.emit.assert_called_once_with("set_modes", [])
+
+    async def test_set_modes_single_mode(
+        self, mock_websocket_session: MagicMock
+    ) -> None:
+        """Test set_modes with single mode."""
+        emitter = ChainlitEmitter(mock_websocket_session)
+        mode = Mode(
+            id="model",
+            name="Model",
+            options=[ModeOption(id="gpt-4", name="GPT-4", default=True)],
+        )
+
+        await emitter.set_modes([mode])
+
+        mock_websocket_session.emit.assert_called_once()
+        call_args = mock_websocket_session.emit.call_args
+        assert call_args[0][0] == "set_modes"
+        assert len(call_args[0][1]) == 1
+        assert call_args[0][1][0]["id"] == "model"
+
+
+@pytest.mark.asyncio
+class TestModeExports:
+    """Test that Mode and ModeOption are properly exported."""
+
+    def test_mode_exported_from_chainlit(self):
+        """Test Mode is exported from chainlit package."""
+        assert hasattr(cl, "Mode")
+        assert cl.Mode is Mode
+
+    def test_mode_option_exported_from_chainlit(self):
+        """Test ModeOption is exported from chainlit package."""
+        assert hasattr(cl, "ModeOption")
+        assert cl.ModeOption is ModeOption

backend/tests/test_session.py

@@ -0,0 +1,765 @@
+import builtins
+import json
+import tempfile
+import uuid
+from pathlib import Path
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from chainlit.session import (
+    BaseSession,
+    HTTPSession,
+    JSONEncoderIgnoreNonSerializable,
+    McpSession,
+    WebsocketSession,
+    clean_metadata,
+)
+
+
+def make_exception_group(message: str, exceptions: list[BaseException]):
+    base_exception_group = getattr(builtins, "BaseExceptionGroup", None)
+    if base_exception_group is None:
+        pytest.skip("BaseExceptionGroup is unavailable on this Python version")
+    return base_exception_group(message, exceptions)  # type: ignore[misc]
+
+
+class TestJSONEncoderIgnoreNonSerializable:
+    """Test suite for JSONEncoderIgnoreNonSerializable."""
+
+    def test_encoder_handles_serializable_objects(self):
+        """Test that encoder handles normal serializable objects."""
+        data = {
+            "string": "value",
+            "number": 42,
+            "list": [1, 2, 3],
+            "dict": {"key": "value"},
+        }
+        result = json.dumps(data, cls=JSONEncoderIgnoreNonSerializable)
+        assert json.loads(result) == data
+
+    def test_encoder_ignores_non_serializable_objects(self):
+        """Test that encoder returns None for non-serializable objects."""
+
+        class NonSerializable:
+            pass
+
+        data = {"normal": "value", "non_serializable": NonSerializable()}
+        result = json.dumps(data, cls=JSONEncoderIgnoreNonSerializable)
+        parsed = json.loads(result)
+
+        assert parsed["normal"] == "value"
+        assert parsed["non_serializable"] is None
+
+    def test_encoder_with_nested_non_serializable(self):
+        """Test encoder with nested non-serializable objects."""
+
+        class NonSerializable:
+            pass
+
+        data = {
+            "level1": {
+                "level2": {
+                    "serializable": "value",
+                    "non_serializable": NonSerializable(),
+                }
+            }
+        }
+        result = json.dumps(data, cls=JSONEncoderIgnoreNonSerializable)
+        parsed = json.loads(result)
+
+        assert parsed["level1"]["level2"]["serializable"] == "value"
+        assert parsed["level1"]["level2"]["non_serializable"] is None
+
+
+class TestCleanMetadata:
+    """Test suite for clean_metadata function."""
+
+    def test_clean_metadata_with_normal_data(self):
+        """Test clean_metadata with normal serializable data."""
+        metadata = {"key": "value", "number": 42, "list": [1, 2, 3]}
+        result = clean_metadata(metadata)
+        assert result == metadata
+
+    def test_clean_metadata_removes_non_serializable(self):
+        """Test that clean_metadata removes non-serializable objects."""
+
+        class NonSerializable:
+            pass
+
+        metadata = {"normal": "value", "non_serializable": NonSerializable()}
+        result = clean_metadata(metadata)
+
+        assert result["normal"] == "value"
+        assert result["non_serializable"] is None
+
+    def test_clean_metadata_redacts_large_data(self):
+        """Test that clean_metadata redacts data exceeding max size."""
+        # Create large metadata
+        large_data = {"data": "x" * 2000000}  # > 1MB
+        result = clean_metadata(large_data, max_size=1048576)
+
+        assert "message" in result
+        assert "exceeds the limit" in result["message"]
+
+    def test_clean_metadata_with_custom_max_size(self):
+        """Test clean_metadata with custom max size."""
+        small_data = {"data": "x" * 100}
+        result = clean_metadata(small_data, max_size=50)
+
+        # Should be redacted because it exceeds 50 bytes
+        assert "message" in result
+        assert "exceeds the limit" in result["message"]
+
+    def test_clean_metadata_preserves_unicode(self):
+        """Test that clean_metadata preserves Unicode characters."""
+        metadata = {"chinese": "你好", "emoji": "🎉", "japanese": "こんにちは"}
+        result = clean_metadata(metadata)
+
+        assert result["chinese"] == "你好"
+        assert result["emoji"] == "🎉"
+        assert result["japanese"] == "こんにちは"
+
+
+class TestBaseSession:
+    """Test suite for BaseSession class."""
+
+    def test_base_session_initialization(self):
+        """Test BaseSession initialization with required parameters."""
+        session = BaseSession(
+            id="test_id",
+            client_type="webapp",
+            thread_id=None,
+            user=None,
+            token=None,
+            user_env=None,
+        )
+
+        assert session.id == "test_id"
+        assert session.client_type == "webapp"
+        assert session.thread_id is not None  # Auto-generated UUID
+        assert session.user is None
+        assert session.token is None
+        assert session.user_env == {}
+        assert session.chat_settings == {}
+
+    def test_base_session_with_thread_id(self):
+        """Test BaseSession with provided thread_id."""
+        thread_id = str(uuid.uuid4())
+        session = BaseSession(
+            id="test_id",
+            client_type="webapp",
+            thread_id=thread_id,
+            user=None,
+            token=None,
+            user_env=None,
+        )
+
+        assert session.thread_id == thread_id
+        assert session.thread_id_to_resume == thread_id
+
+    def test_base_session_with_user_env(self):
+        """Test BaseSession with user environment variables."""
+        user_env = {"API_KEY": "secret", "ENV_VAR": "value"}
+        session = BaseSession(
+            id="test_id",
+            client_type="webapp",
+            thread_id=None,
+            user=None,
+            token=None,
+            user_env=user_env,
+        )
+
+        assert session.user_env == user_env
+
+    def test_base_session_with_chat_profile(self):
+        """Test BaseSession with chat profile."""
+        session = BaseSession(
+            id="test_id",
+            client_type="webapp",
+            thread_id=None,
+            user=None,
+            token=None,
+            user_env=None,
+            chat_profile="gpt-4",
+        )
+
+        assert session.chat_profile == "gpt-4"
+
+    def test_base_session_files_dir(self):
+        """Test BaseSession files_dir property."""
+        with patch("chainlit.config.FILES_DIRECTORY", Path("/tmp/files")):
+            session = BaseSession(
+                id="test_id",
+                client_type="webapp",
+                thread_id=None,
+                user=None,
+                token=None,
+                user_env=None,
+            )
+
+            assert session.files_dir == Path("/tmp/files/test_id")
+
+    @pytest.mark.asyncio
+    async def test_base_session_persist_file_with_content(self):
+        """Test persisting a file with content."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch("chainlit.config.FILES_DIRECTORY", Path(tmpdir)):
+                session = BaseSession(
+                    id="test_id",
+                    client_type="webapp",
+                    thread_id=None,
+                    user=None,
+                    token=None,
+                    user_env=None,
+                )
+
+                content = b"test file content"
+                result = await session.persist_file(
+                    name="test.txt",
+                    mime="text/plain",
+                    content=content,
+                )
+
+                assert "id" in result
+                assert result["id"] in session.files
+                assert session.files[result["id"]]["name"] == "test.txt"
+                assert session.files[result["id"]]["type"] == "text/plain"
+
+    @pytest.mark.asyncio
+    async def test_base_session_persist_file_with_string_content(self):
+        """Test persisting a file with string content."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch("chainlit.config.FILES_DIRECTORY", Path(tmpdir)):
+                session = BaseSession(
+                    id="test_id",
+                    client_type="webapp",
+                    thread_id=None,
+                    user=None,
+                    token=None,
+                    user_env=None,
+                )
+
+                content = "test string content"
+                result = await session.persist_file(
+                    name="test.txt",
+                    mime="text/plain",
+                    content=content,
+                )
+
+                assert "id" in result
+                file_id = result["id"]
+                assert session.files[file_id]["size"] > 0
+
+    @pytest.mark.asyncio
+    async def test_base_session_persist_file_without_path_or_content(self):
+        """Test that persist_file raises error without path or content."""
+        session = BaseSession(
+            id="test_id",
+            client_type="webapp",
+            thread_id=None,
+            user=None,
+            token=None,
+            user_env=None,
+        )
+
+        with pytest.raises(ValueError, match="Either path or content must be provided"):
+            await session.persist_file(name="test.txt", mime="text/plain")
+
+    def test_base_session_to_persistable(self):
+        """Test BaseSession to_persistable method."""
+        from chainlit.user_session import user_sessions
+
+        original_sessions = user_sessions.copy()
+        user_sessions.update({"test_id": {"key": "value"}})
+
+        try:
+            with patch("chainlit.config.config") as mock_config:
+                mock_config.project.persist_user_env = True
+
+                session = BaseSession(
+                    id="test_id",
+                    client_type="webapp",
+                    thread_id=None,
+                    user=None,
+                    token=None,
+                    user_env={"API_KEY": "secret"},
+                    chat_profile="gpt-4",
+                )
+                session.chat_settings = {"temperature": 0.7}
+
+                result = session.to_persistable()
+
+                assert result["chat_settings"] == {"temperature": 0.7}
+                assert result["chat_profile"] == "gpt-4"
+                assert result["client_type"] == "webapp"
+        finally:
+            user_sessions.clear()
+            user_sessions.update(original_sessions)
+
+    def test_base_session_to_persistable_without_persist_user_env(self):
+        """Test to_persistable removes user_env when persist_user_env is False."""
+        from chainlit.user_session import user_sessions
+
+        original_sessions = user_sessions.copy()
+        user_sessions.update({"test_id": {"env": {"KEY": "value"}}})
+
+        try:
+            with patch("chainlit.config.config") as mock_config:
+                mock_config.project.persist_user_env = False
+
+                session = BaseSession(
+                    id="test_id",
+                    client_type="webapp",
+                    thread_id=None,
+                    user=None,
+                    token=None,
+                    user_env={"API_KEY": "secret"},
+                )
+
+                result = session.to_persistable()
+
+                assert result["env"] == {}
+        finally:
+            user_sessions.clear()
+            user_sessions.update(original_sessions)
+
+
+class TestHTTPSession:
+    """Test suite for HTTPSession class."""
+
+    def test_http_session_initialization(self):
+        """Test HTTPSession initialization."""
+        session = HTTPSession(
+            id="http_id",
+            client_type="copilot",
+            thread_id=None,
+            user=None,
+            token=None,
+            user_env=None,
+        )
+
+        assert session.id == "http_id"
+        assert session.client_type == "copilot"
+        assert isinstance(session, BaseSession)
+
+    @pytest.mark.asyncio
+    async def test_http_session_delete(self):
+        """Test HTTPSession delete method."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch("chainlit.config.FILES_DIRECTORY", Path(tmpdir)):
+                session = HTTPSession(
+                    id="http_id",
+                    client_type="copilot",
+                )
+
+                # Create files directory
+                session.files_dir.mkdir(exist_ok=True)
+                test_file = session.files_dir / "test.txt"
+                test_file.write_text("test")
+
+                assert session.files_dir.exists()
+
+                await session.delete()
+
+                assert not session.files_dir.exists()
+
+
+class TestWebsocketSession:
+    """Test suite for WebsocketSession class."""
+
+    def test_websocket_session_initialization(self):
+        """Test WebsocketSession initialization."""
+        emit_mock = Mock()
+        emit_call_mock = Mock()
+
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=emit_mock,
+            emit_call=emit_call_mock,
+            user_env={},
+            client_type="webapp",
+        )
+
+        assert session.id == "ws_id"
+        assert session.socket_id == "socket_123"
+        assert session.emit == emit_mock
+        assert session.emit_call == emit_call_mock
+        assert session.restored is False
+        assert session.mcp_sessions == {}
+
+    def test_websocket_session_language_detection(self):
+        """Test WebsocketSession language detection from HTTP headers."""
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+            environ={"HTTP_ACCEPT_LANGUAGE": "fr-FR,fr;q=0.9,en;q=0.8"},
+        )
+
+        assert session.language == "fr-FR"
+
+    def test_websocket_session_default_language(self):
+        """Test WebsocketSession defaults to en-US without language header."""
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+            environ={},
+        )
+
+        assert session.language == "en-US"
+
+    def test_websocket_session_restore(self):
+        """Test WebsocketSession restore method."""
+        from chainlit.session import ws_sessions_sid
+
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="old_socket",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+        )
+
+        assert ws_sessions_sid.get("old_socket") == session
+
+        session.restore("new_socket")
+
+        assert session.socket_id == "new_socket"
+        assert session.restored is True
+        assert ws_sessions_sid.get("old_socket") is None
+        assert ws_sessions_sid.get("new_socket") == session
+
+    @pytest.mark.asyncio
+    async def test_websocket_session_delete(self):
+        """Test WebsocketSession delete method."""
+        from chainlit.session import ws_sessions_id, ws_sessions_sid
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch("chainlit.config.FILES_DIRECTORY", Path(tmpdir)):
+                session = WebsocketSession(
+                    id="ws_id",
+                    socket_id="socket_123",
+                    emit=Mock(),
+                    emit_call=Mock(),
+                    user_env={},
+                    client_type="webapp",
+                )
+
+                # Create files directory
+                session.files_dir.mkdir(exist_ok=True)
+
+                assert ws_sessions_sid.get("socket_123") == session
+                assert ws_sessions_id.get("ws_id") == session
+
+                await session.delete()
+
+                assert not session.files_dir.exists()
+                assert ws_sessions_sid.get("socket_123") is None
+                assert ws_sessions_id.get("ws_id") is None
+
+    def test_websocket_session_get(self):
+        """Test WebsocketSession.get class method."""
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+        )
+
+        retrieved = WebsocketSession.get("socket_123")
+        assert retrieved == session
+
+    def test_websocket_session_get_by_id(self):
+        """Test WebsocketSession.get_by_id class method."""
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+        )
+
+        retrieved = WebsocketSession.get_by_id("ws_id")
+        assert retrieved == session
+
+    def test_websocket_session_require_success(self):
+        """Test WebsocketSession.require with existing session."""
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+        )
+
+        retrieved = WebsocketSession.require("socket_123")
+        assert retrieved == session
+
+    def test_websocket_session_require_failure(self):
+        """Test WebsocketSession.require raises error for missing session."""
+        with pytest.raises(ValueError, match="Session not found"):
+            WebsocketSession.require("nonexistent_socket")
+
+    @pytest.mark.asyncio
+    async def test_websocket_session_flush_method_queue(self):
+        """Test WebsocketSession flush_method_queue."""
+        from collections import deque
+
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+        )
+
+        # Create a mock async method
+        mock_method = AsyncMock()
+
+        # Add items to queue
+        session.thread_queues["test_method"] = deque(
+            [
+                (mock_method, session, ("arg1",), {"kwarg1": "value1"}),
+                (mock_method, session, ("arg2",), {"kwarg2": "value2"}),
+            ]
+        )
+
+        await session.flush_method_queue()
+
+        assert mock_method.call_count == 2
+        assert len(session.thread_queues["test_method"]) == 0
+
+
+class TestSessionEdgeCases:
+    """Test suite for session edge cases."""
+
+    def test_base_session_with_all_client_types(self):
+        """Test BaseSession with different client types."""
+        client_types = ["webapp", "copilot", "teams", "slack", "discord"]
+
+        for client_type in client_types:
+            session = BaseSession(
+                id=f"test_{client_type}",
+                client_type=client_type,
+                thread_id=None,
+                user=None,
+                token=None,
+                user_env=None,
+            )
+            assert session.client_type == client_type
+
+    @pytest.mark.asyncio
+    async def test_persist_file_with_mime_extension(self):
+        """Test that persist_file adds correct file extension based on MIME type."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch("chainlit.config.FILES_DIRECTORY", Path(tmpdir)):
+                session = BaseSession(
+                    id="test_id",
+                    client_type="webapp",
+                    thread_id=None,
+                    user=None,
+                    token=None,
+                    user_env=None,
+                )
+
+                # Test with image MIME type
+                result = await session.persist_file(
+                    name="image.png",
+                    mime="image/png",
+                    content=b"fake image data",
+                )
+
+                file_id = result["id"]
+                file_path = session.files[file_id]["path"]
+                assert file_path.suffix == ".png"
+
+    def test_clean_metadata_with_empty_dict(self):
+        """Test clean_metadata with empty dictionary."""
+        result = clean_metadata({})
+        assert result == {}
+
+    def test_websocket_session_with_chat_profile(self):
+        """Test WebsocketSession with chat profile."""
+        session = WebsocketSession(
+            id="ws_id",
+            socket_id="socket_123",
+            emit=Mock(),
+            emit_call=Mock(),
+            user_env={},
+            client_type="webapp",
+            chat_profile="gpt-4",
+        )
+
+        assert session.chat_profile == "gpt-4"
+
+    @pytest.mark.asyncio
+    async def test_websocket_session_delete_with_mcp_sessions(self):
+        """Test WebsocketSession delete with MCP sessions."""
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch("chainlit.config.FILES_DIRECTORY", Path(tmpdir)):
+                session = WebsocketSession(
+                    id="ws_id",
+                    socket_id="socket_123",
+                    emit=Mock(),
+                    emit_call=Mock(),
+                    user_env={},
+                    client_type="webapp",
+                )
+
+                # Create a real McpSession with a completed task
+                import asyncio
+
+                stop = asyncio.Event()
+                stop.set()  # already stopped
+
+                async def _noop():
+                    pass
+
+                task = asyncio.create_task(_noop())
+                await task  # let it finish
+
+                mcp = McpSession(
+                    name="mcp1",
+                    client=Mock(),
+                    task=task,
+                    stop_event=stop,
+                )
+                session.mcp_sessions["mcp1"] = mcp
+
+                await session.delete()
+
+                assert "mcp1" not in session.mcp_sessions
+
+    @pytest.mark.asyncio
+    async def test_websocket_session_delete_with_hanging_mcp(self):
+        """Test that session delete handles a slow MCP session gracefully."""
+        import asyncio
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch("chainlit.config.FILES_DIRECTORY", Path(tmpdir)):
+                session = WebsocketSession(
+                    id="ws_id",
+                    socket_id="socket_123",
+                    emit=Mock(),
+                    emit_call=Mock(),
+                    user_env={},
+                    client_type="webapp",
+                )
+
+                stop = asyncio.Event()
+
+                async def _hang():
+                    await stop.wait()
+
+                task = asyncio.create_task(_hang())
+
+                mcp = McpSession(
+                    name="mcp1",
+                    client=Mock(),
+                    task=task,
+                    stop_event=stop,
+                )
+                session.mcp_sessions["mcp1"] = mcp
+
+                # delete() should close the session cleanly
+                await session.delete()
+
+                assert task.done()
+                assert "mcp1" not in session.mcp_sessions
+
+
+class TestMcpSession:
+    """Test suite for the McpSession dataclass."""
+
+    @pytest.mark.asyncio
+    async def test_close_signals_stop_and_awaits_task(self):
+        """close() sets the stop event and waits for the task."""
+        import asyncio
+
+        stop = asyncio.Event()
+
+        async def _runner():
+            await stop.wait()
+
+        task = asyncio.create_task(_runner())
+        mcp = McpSession(
+            name="test",
+            client=Mock(),
+            task=task,
+            stop_event=stop,
+        )
+
+        await mcp.close()
+
+        assert stop.is_set()
+        assert task.done()
+
+    @pytest.mark.asyncio
+    async def test_close_cancels_on_timeout(self):
+        """close() cancels a task that doesn't respond to stop_event."""
+        import asyncio
+
+        stop = asyncio.Event()
+
+        async def _stuck():
+            # Ignore stop_event entirely
+            await asyncio.sleep(3600)
+
+        task = asyncio.create_task(_stuck())
+        mcp = McpSession(
+            name="stuck",
+            client=Mock(),
+            task=task,
+            stop_event=stop,
+        )
+
+        # Temporarily reduce timeout for this test
+        import chainlit.session as session_mod
+
+        original_timeout = session_mod._CLOSE_TIMEOUT
+        session_mod._CLOSE_TIMEOUT = 0.1
+        try:
+            await mcp.close()
+        finally:
+            session_mod._CLOSE_TIMEOUT = original_timeout
+
+        assert task.done()
+
+    @pytest.mark.asyncio
+    async def test_close_idempotent(self):
+        """Calling close() twice does not raise."""
+        import asyncio
+
+        stop = asyncio.Event()
+
+        async def _runner():
+            await stop.wait()
+
+        task = asyncio.create_task(_runner())
+        mcp = McpSession(
+            name="test",
+            client=Mock(),
+            task=task,
+            stop_event=stop,
+        )
+
+        await mcp.close()
+        await mcp.close()  # second call should be safe
+
+        assert task.done()

backend/tests/test_sidebar.py

@@ -0,0 +1,303 @@
+import pytest
+
+from chainlit.element import File, Image, Text
+from chainlit.sidebar import ElementSidebar
+
+
+@pytest.mark.asyncio
+class TestElementSidebar:
+    """Test suite for ElementSidebar class."""
+
+    async def test_set_title(self, mock_chainlit_context):
+        """Test ElementSidebar.set_title() method."""
+        async with mock_chainlit_context as ctx:
+            await ElementSidebar.set_title("My Sidebar Title")
+
+            ctx.emitter.emit.assert_called_once_with(
+                "set_sidebar_title", "My Sidebar Title"
+            )
+
+    async def test_set_title_with_empty_string(self, mock_chainlit_context):
+        """Test ElementSidebar.set_title() with empty string."""
+        async with mock_chainlit_context as ctx:
+            await ElementSidebar.set_title("")
+
+            ctx.emitter.emit.assert_called_once_with("set_sidebar_title", "")
+
+    async def test_set_title_with_special_characters(self, mock_chainlit_context):
+        """Test ElementSidebar.set_title() with special characters."""
+        async with mock_chainlit_context as ctx:
+            title = "Title with 特殊字符 & symbols! 🎉"
+            await ElementSidebar.set_title(title)
+
+            ctx.emitter.emit.assert_called_once_with("set_sidebar_title", title)
+
+    async def test_set_elements_with_single_element(self, mock_chainlit_context):
+        """Test ElementSidebar.set_elements() with a single element."""
+        async with mock_chainlit_context as ctx:
+            element = File(name="test.txt", url="https://example.com/test.txt")
+
+            await ElementSidebar.set_elements([element])
+
+            # Verify element.send() was called
+            ctx.emitter.send_element.assert_called_once()
+
+            # Verify emit was called with correct structure
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "set_sidebar_elements"
+            assert "elements" in call_args[0][1]
+            assert "key" in call_args[0][1]
+            assert len(call_args[0][1]["elements"]) == 1
+            assert call_args[0][1]["key"] is None
+
+    async def test_set_elements_with_multiple_elements(self, mock_chainlit_context):
+        """Test ElementSidebar.set_elements() with multiple elements."""
+        async with mock_chainlit_context as ctx:
+            elements = [
+                File(name="file1.txt", url="https://example.com/file1.txt"),
+                Image(name="image1.png", url="https://example.com/image1.png"),
+                Text(name="text1", content="Some text content"),
+            ]
+
+            await ElementSidebar.set_elements(elements)
+
+            # Verify all elements were sent (3 send_element calls)
+            assert ctx.emitter.send_element.call_count == 3
+
+            # Verify emit was called with all elements
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "set_sidebar_elements"
+            assert len(call_args[0][1]["elements"]) == 3
+
+    async def test_set_elements_with_empty_list(self, mock_chainlit_context):
+        """Test ElementSidebar.set_elements() with empty list (closes sidebar)."""
+        async with mock_chainlit_context as ctx:
+            await ElementSidebar.set_elements([])
+
+            # No elements to send
+            ctx.emitter.send_element.assert_not_called()
+
+            # Emit should still be called with empty elements
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][0] == "set_sidebar_elements"
+            assert call_args[0][1]["elements"] == []
+            assert call_args[0][1]["key"] is None
+
+    async def test_set_elements_with_key(self, mock_chainlit_context):
+        """Test ElementSidebar.set_elements() with a key."""
+        async with mock_chainlit_context as ctx:
+            element = File(name="test.txt", url="https://example.com/test.txt")
+            key = "my_sidebar_key"
+
+            await ElementSidebar.set_elements([element], key=key)
+
+            # Verify emit was called with the key
+            ctx.emitter.emit.assert_called_once()
+            call_args = ctx.emitter.emit.call_args
+            assert call_args[0][1]["key"] == key
+
+    async def test_set_elements_with_for_id(self, mock_chainlit_context):
+        """Test ElementSidebar.set_elements() with elements that have for_id."""
+        async with mock_chainlit_context as ctx:
+            element = File(
+                name="test.txt",
+                url="https://example.com/test.txt",
+                for_id="message_123",
+            )
+
+            await ElementSidebar.set_elements([element])
+
+            # Element should be sent with its for_id
+            ctx.emitter.send_element.assert_called_once()
+
+            # Verify emit was called
+            ctx.emitter.emit.assert_called_once()
+
+    async def test_set_elements_without_for_id(self, mock_chainlit_context):
+        """Test ElementSidebar.set_elements() with elements without for_id."""
+        async with mock_chainlit_context as ctx:
+            element = File(name="test.txt", url="https://example.com/test.txt")
+
+            await ElementSidebar.set_elements([element])
+
+            # Element should be sent with empty string for_id
+            ctx.emitter.send_element.assert_called_once()
+
+            # Verify emit was called
+            ctx.emitter.emit.assert_called_once()
+
+    async def test_set_elements_persist_false(self, mock_chainlit_context):
+        """Test that set_elements() sends elements with persist=False."""
+        async with mock_chainlit_context as ctx:
+            # Mock persist_file to provide chainlit_key
+            ctx.session.persist_file.return_value = {"id": "test_key"}
+
+            element = File(name="test.txt", content=b"test content")
+
+            await ElementSidebar.set_elements([element])
+
+            # persist_file is still called to get chainlit_key, even with persist=False
+            # The persist=False affects data layer persistence, not file upload
+            ctx.session.persist_file.assert_called_once()
+
+            # Verify element was sent
+            ctx.emitter.send_element.assert_called_once()
+
+    async def test_set_elements_serialization(self, mock_chainlit_context):
+        """Test that elements are properly serialized in set_elements()."""
+        async with mock_chainlit_context as ctx:
+            file_elem = File(name="file.txt", url="https://example.com/file.txt")
+            image_elem = Image(
+                name="image.png", url="https://example.com/image.png", size="large"
+            )
+
+            await ElementSidebar.set_elements([file_elem, image_elem])
+
+            # Verify emit was called with serialized elements
+            call_args = ctx.emitter.emit.call_args
+            elements_data = call_args[0][1]["elements"]
+
+            assert len(elements_data) == 2
+            assert elements_data[0]["name"] == "file.txt"
+            assert elements_data[0]["type"] == "file"
+            assert elements_data[1]["name"] == "image.png"
+            assert elements_data[1]["type"] == "image"
+            assert elements_data[1]["size"] == "large"
+
+
+@pytest.mark.asyncio
+class TestElementSidebarEdgeCases:
+    """Test suite for ElementSidebar edge cases."""
+
+    async def test_set_title_multiple_times(self, mock_chainlit_context):
+        """Test calling set_title() multiple times."""
+        async with mock_chainlit_context as ctx:
+            await ElementSidebar.set_title("First Title")
+            await ElementSidebar.set_title("Second Title")
+            await ElementSidebar.set_title("Third Title")
+
+            assert ctx.emitter.emit.call_count == 3
+
+            # Verify last call had the third title
+            last_call = ctx.emitter.emit.call_args
+            assert last_call[0][1] == "Third Title"
+
+    async def test_set_elements_multiple_times(self, mock_chainlit_context):
+        """Test calling set_elements() multiple times."""
+        async with mock_chainlit_context as ctx:
+            element1 = File(name="file1.txt", url="https://example.com/file1.txt")
+            element2 = File(name="file2.txt", url="https://example.com/file2.txt")
+
+            await ElementSidebar.set_elements([element1])
+            await ElementSidebar.set_elements([element2])
+
+            # Should have sent both elements
+            assert ctx.emitter.send_element.call_count == 2
+
+            # Should have emitted twice
+            assert ctx.emitter.emit.call_count == 2
+
+    async def test_set_elements_with_same_key_twice(self, mock_chainlit_context):
+        """Test calling set_elements() with the same key twice."""
+        async with mock_chainlit_context as ctx:
+            element1 = File(name="file1.txt", url="https://example.com/file1.txt")
+            element2 = File(name="file2.txt", url="https://example.com/file2.txt")
+
+            await ElementSidebar.set_elements([element1], key="same_key")
+            await ElementSidebar.set_elements([element2], key="same_key")
+
+            # Both should be sent (server doesn't prevent this)
+            assert ctx.emitter.send_element.call_count == 2
+            assert ctx.emitter.emit.call_count == 2
+
+    async def test_set_elements_with_different_keys(self, mock_chainlit_context):
+        """Test calling set_elements() with different keys."""
+        async with mock_chainlit_context as ctx:
+            element1 = File(name="file1.txt", url="https://example.com/file1.txt")
+            element2 = File(name="file2.txt", url="https://example.com/file2.txt")
+
+            await ElementSidebar.set_elements([element1], key="key1")
+            await ElementSidebar.set_elements([element2], key="key2")
+
+            assert ctx.emitter.emit.call_count == 2
+
+            # Verify different keys were used
+            calls = ctx.emitter.emit.call_args_list
+            assert calls[0][0][1]["key"] == "key1"
+            assert calls[1][0][1]["key"] == "key2"
+
+    async def test_set_elements_with_large_number_of_elements(
+        self, mock_chainlit_context
+    ):
+        """Test set_elements() with many elements."""
+        async with mock_chainlit_context as ctx:
+            # Create 50 elements
+            elements = [
+                File(name=f"file{i}.txt", url=f"https://example.com/file{i}.txt")
+                for i in range(50)
+            ]
+
+            await ElementSidebar.set_elements(elements)
+
+            # All 50 elements should be sent
+            assert ctx.emitter.send_element.call_count == 50
+
+            # Verify emit was called with all 50 elements
+            call_args = ctx.emitter.emit.call_args
+            assert len(call_args[0][1]["elements"]) == 50
+
+    async def test_set_title_and_set_elements_together(self, mock_chainlit_context):
+        """Test using set_title() and set_elements() together."""
+        async with mock_chainlit_context as ctx:
+            await ElementSidebar.set_title("My Documents")
+
+            elements = [
+                File(name="doc1.pdf", url="https://example.com/doc1.pdf"),
+                File(name="doc2.pdf", url="https://example.com/doc2.pdf"),
+            ]
+            await ElementSidebar.set_elements(elements)
+
+            # Verify both methods were called
+            assert ctx.emitter.emit.call_count == 2
+
+            # Verify the calls were correct
+            calls = ctx.emitter.emit.call_args_list
+            assert calls[0][0][0] == "set_sidebar_title"
+            assert calls[0][0][1] == "My Documents"
+            assert calls[1][0][0] == "set_sidebar_elements"
+
+    async def test_set_elements_with_mixed_element_types(self, mock_chainlit_context):
+        """Test set_elements() with various element types."""
+        async with mock_chainlit_context as ctx:
+            elements = [
+                File(name="document.pdf", url="https://example.com/doc.pdf"),
+                Image(
+                    name="photo.jpg", url="https://example.com/photo.jpg", size="medium"
+                ),
+                Text(name="notes", content="Some important notes"),
+            ]
+
+            await ElementSidebar.set_elements(elements)
+
+            # Verify all different types were sent
+            assert ctx.emitter.send_element.call_count == 3
+
+            # Verify serialization includes type information
+            call_args = ctx.emitter.emit.call_args
+            elements_data = call_args[0][1]["elements"]
+
+            assert elements_data[0]["type"] == "file"
+            assert elements_data[1]["type"] == "image"
+            assert elements_data[2]["type"] == "text"
+
+    async def test_set_title_with_long_string(self, mock_chainlit_context):
+        """Test set_title() with a very long title."""
+        async with mock_chainlit_context as ctx:
+            long_title = "A" * 1000  # 1000 character title
+
+            await ElementSidebar.set_title(long_title)
+
+            ctx.emitter.emit.assert_called_once_with("set_sidebar_title", long_title)

backend/tests/test_slack_socket_mode.py

@@ -0,0 +1,54 @@
+# tests/test_slack_socket_mode.py
+import importlib
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_start_socket_mode_starts_handler(monkeypatch):
+    """
+    The function should:
+      • build an AsyncSocketModeHandler with the global slack_app
+      • use the token found in SLACK_WEBSOCKET_TOKEN
+      • await the handler.start_async() coroutine exactly once
+    """
+    token = "xapp-fake-token"
+    # minimal env required for the Slack module to initialise
+    monkeypatch.setenv("SLACK_BOT_TOKEN", "xoxb-fake-bot")
+    monkeypatch.setenv("SLACK_WEBSOCKET_TOKEN", token)
+
+    # Import the module first to avoid lazy import registry issues
+    slack_app_mod = importlib.import_module("chainlit.slack.app")
+
+    # Patch the object directly instead of using string path
+    with patch.object(
+        slack_app_mod, "AsyncSocketModeHandler", autospec=True
+    ) as handler_cls:
+        handler_instance = AsyncMock()
+        handler_cls.return_value = handler_instance
+
+        # Run: should build handler + await start_async
+        await slack_app_mod.start_socket_mode()
+
+        handler_cls.assert_called_once_with(slack_app_mod.slack_app, token)
+        handler_instance.start_async.assert_awaited_once()
+
+
+def test_slack_http_route_registered(monkeypatch):
+    """
+    When only the classic HTTP tokens are set (no websocket token),
+    the FastAPI app should expose POST /slack/events.
+    """
+    # HTTP-only environment
+    monkeypatch.setenv("SLACK_BOT_TOKEN", "xoxb-fake-bot")
+    monkeypatch.setenv("SLACK_SIGNING_SECRET", "shhh-fake-secret")
+    monkeypatch.delenv("SLACK_WEBSOCKET_TOKEN", raising=False)
+
+    # Re-import server with the fresh env so the route table is built correctly
+    server = importlib.reload(importlib.import_module("chainlit.server"))
+
+    assert any(
+        route.path == "/slack/events" and "POST" in route.methods
+        for route in server.router.routes
+    ), "Slack HTTP handler route was not registered"

backend/tests/test_socket.py

@@ -0,0 +1,630 @@
+import json
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from chainlit.session import WebsocketSession
+from chainlit.socket import (
+    _authenticate_connection,
+    _get_token,
+    _get_token_from_cookie,
+    clean_session,
+    connection_successful,
+    load_user_env,
+    persist_user_session,
+    restore_existing_session,
+    resume_thread,
+)
+
+
+class TestGetTokenFromCookie:
+    """Test suite for _get_token_from_cookie function."""
+
+    def test_get_token_from_cookie_with_valid_cookie(self):
+        """Test extracting token from valid cookie header."""
+        with patch("chainlit.socket.get_token_from_cookies") as mock_get_token:
+            mock_get_token.return_value = "test_token"
+            environ = {"HTTP_COOKIE": "session=abc123; token=test_token"}
+
+            result = _get_token_from_cookie(environ)
+
+            assert result == "test_token"
+            mock_get_token.assert_called_once()
+
+    def test_get_token_from_cookie_without_cookie(self):
+        """Test when no cookie header is present."""
+        environ = {}
+        result = _get_token_from_cookie(environ)
+        assert result is None
+
+    def test_get_token_from_cookie_with_empty_cookie(self):
+        """Test with empty cookie header."""
+        with patch("chainlit.socket.get_token_from_cookies") as mock_get_token:
+            mock_get_token.return_value = None
+            environ = {"HTTP_COOKIE": ""}
+
+            result = _get_token_from_cookie(environ)
+
+            assert result is None
+
+
+class TestGetToken:
+    """Test suite for _get_token function."""
+
+    def test_get_token_calls_get_token_from_cookie(self):
+        """Test that _get_token delegates to _get_token_from_cookie."""
+        with patch("chainlit.socket._get_token_from_cookie") as mock_get_cookie:
+            mock_get_cookie.return_value = "token_value"
+            environ = {"HTTP_COOKIE": "token=token_value"}
+
+            result = _get_token(environ)
+
+            assert result == "token_value"
+            mock_get_cookie.assert_called_once_with(environ)
+
+
+class TestAuthenticateConnection:
+    """Test suite for _authenticate_connection function."""
+
+    @pytest.mark.asyncio
+    async def test_authenticate_connection_with_valid_token(self):
+        """Test authentication with valid token."""
+        mock_user = Mock()
+        mock_user.identifier = "user123"
+
+        with patch("chainlit.socket._get_token") as mock_get_token:
+            with patch("chainlit.socket.get_current_user") as mock_get_user:
+                mock_get_token.return_value = "valid_token"
+                mock_get_user.return_value = mock_user
+
+                environ = {"HTTP_COOKIE": "token=valid_token"}
+                user, token = await _authenticate_connection(environ)
+
+                assert user == mock_user
+                assert token == "valid_token"
+                mock_get_user.assert_called_once_with(token="valid_token")
+
+    @pytest.mark.asyncio
+    async def test_authenticate_connection_without_token(self):
+        """Test authentication without token."""
+        with patch("chainlit.socket._get_token") as mock_get_token:
+            mock_get_token.return_value = None
+
+            environ = {}
+            user, token = await _authenticate_connection(environ)
+
+            assert user is None
+            assert token is None
+
+    @pytest.mark.asyncio
+    async def test_authenticate_connection_with_invalid_token(self):
+        """Test authentication with invalid token."""
+        with patch("chainlit.socket._get_token") as mock_get_token:
+            with patch("chainlit.socket.get_current_user") as mock_get_user:
+                mock_get_token.return_value = "invalid_token"
+                mock_get_user.return_value = None
+
+                environ = {"HTTP_COOKIE": "token=invalid_token"}
+                user, token = await _authenticate_connection(environ)
+
+                assert user is None
+                assert token is None
+
+
+class TestRestoreExistingSession:
+    """Test suite for restore_existing_session function."""
+
+    def test_restore_existing_session_success(self):
+        """Test restoring an existing session."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = None
+        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, 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_with_matching_user(self):
+        """Test restoring a session for its authenticated owner."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock(identifier="user123")
+        authenticated_user = Mock(identifier="user123")
+
+        with patch.object(WebsocketSession, "get_by_id") as mock_get:
+            mock_get.return_value = mock_session
+
+            result = restore_existing_session(
+                "new_sid",
+                "session_123",
+                Mock(),
+                Mock(),
+                {"HTTP_COOKIE": "token=token"},
+                user=authenticated_user,
+            )
+
+            assert result is True
+            mock_session.restore.assert_called_once_with(new_socket_id="new_sid")
+
+    def test_restore_existing_session_rejects_different_user(self):
+        """Test that a session cannot be restored by another user."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock(identifier="victim")
+        authenticated_user = Mock(identifier="attacker")
+
+        with patch.object(WebsocketSession, "get_by_id") as mock_get:
+            mock_get.return_value = mock_session
+
+            with pytest.raises(ConnectionRefusedError, match="authorization failed"):
+                restore_existing_session(
+                    "new_sid",
+                    "session_123",
+                    Mock(),
+                    Mock(),
+                    {"HTTP_COOKIE": "token=token"},
+                    user=authenticated_user,
+                )
+
+            mock_session.restore.assert_not_called()
+
+    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(), {"HTTP_COOKIE": "token=token"}
+            )
+
+            assert result is False
+
+
+class TestPersistUserSession:
+    """Test suite for persist_user_session function."""
+
+    @pytest.mark.asyncio
+    async def test_persist_user_session_with_data_layer(self):
+        """Test persisting user session with data layer."""
+        mock_data_layer = AsyncMock()
+
+        with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+            mock_get_dl.return_value = mock_data_layer
+
+            metadata = {"key": "value"}
+            await persist_user_session("thread_123", metadata)
+
+            mock_data_layer.update_thread.assert_called_once_with(
+                thread_id="thread_123", metadata=metadata
+            )
+
+    @pytest.mark.asyncio
+    async def test_persist_user_session_without_data_layer(self):
+        """Test persisting when no data layer is available."""
+        with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+            mock_get_dl.return_value = None
+
+            # Should not raise an error
+            await persist_user_session("thread_123", {"key": "value"})
+
+
+class TestResumeThread:
+    """Test suite for resume_thread function."""
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_without_data_layer(self):
+        """Test resume thread when no data layer exists."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock()
+        mock_session.thread_id_to_resume = "thread_123"
+
+        with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+            mock_get_dl.return_value = None
+
+            result = await resume_thread(mock_session)
+
+            assert result is None
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_without_user(self):
+        """Test resume thread when session has no user."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = None
+        mock_session.thread_id_to_resume = "thread_123"
+
+        result = await resume_thread(mock_session)
+
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_without_thread_id(self):
+        """Test resume thread when no thread_id_to_resume."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock()
+        mock_session.thread_id_to_resume = None
+
+        result = await resume_thread(mock_session)
+
+        assert result is None
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_thread_not_found(self):
+        """Test resume thread when thread doesn't exist."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock(identifier="user123")
+        mock_session.thread_id_to_resume = "thread_123"
+        mock_session.id = "session_123"
+
+        mock_data_layer = AsyncMock()
+        mock_data_layer.get_thread.return_value = None
+
+        with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+            mock_get_dl.return_value = mock_data_layer
+
+            result = await resume_thread(mock_session)
+
+            assert result is None
+            mock_data_layer.get_thread.assert_called_once_with(thread_id="thread_123")
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_user_not_author(self):
+        """Test resume thread when user is not the thread author."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock(identifier="user123")
+        mock_session.thread_id_to_resume = "thread_123"
+        mock_session.id = "session_123"
+
+        thread = {"userIdentifier": "different_user", "metadata": {}}
+        mock_data_layer = AsyncMock()
+        mock_data_layer.get_thread.return_value = thread
+
+        with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+            mock_get_dl.return_value = mock_data_layer
+
+            result = await resume_thread(mock_session)
+
+            assert result is None
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_success(self):
+        """Test successful thread resumption."""
+        from chainlit.user_session import user_sessions
+
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock(identifier="user123")
+        mock_session.thread_id_to_resume = "thread_123"
+        mock_session.id = "session_123"
+
+        metadata = {
+            "chat_profile": "gpt-4",
+            "chat_settings": {"temperature": 0.7},
+        }
+        thread = {"userIdentifier": "user123", "metadata": metadata}
+
+        mock_data_layer = AsyncMock()
+        mock_data_layer.get_thread.return_value = thread
+
+        original_sessions = user_sessions.copy()
+        try:
+            with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+                mock_get_dl.return_value = mock_data_layer
+
+                result = await resume_thread(mock_session)
+
+                assert result == thread
+                assert mock_session.chat_profile == "gpt-4"
+                assert mock_session.chat_settings == {"temperature": 0.7}
+                assert user_sessions.get("session_123") == metadata
+        finally:
+            user_sessions.clear()
+            user_sessions.update(original_sessions)
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_with_string_metadata(self):
+        """Test thread resumption with JSON string metadata."""
+        from chainlit.user_session import user_sessions
+
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock(identifier="user123")
+        mock_session.thread_id_to_resume = "thread_123"
+        mock_session.id = "session_123"
+
+        metadata_dict = {"chat_profile": "gpt-4"}
+        thread = {
+            "userIdentifier": "user123",
+            "metadata": json.dumps(metadata_dict),
+        }
+
+        mock_data_layer = AsyncMock()
+        mock_data_layer.get_thread.return_value = thread
+
+        original_sessions = user_sessions.copy()
+        try:
+            with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+                mock_get_dl.return_value = mock_data_layer
+
+                result = await resume_thread(mock_session)
+
+                assert result == thread
+                assert mock_session.chat_profile == "gpt-4"
+        finally:
+            user_sessions.clear()
+            user_sessions.update(original_sessions)
+
+
+class TestLoadUserEnv:
+    """Test suite for load_user_env function."""
+
+    def test_load_user_env_with_valid_json(self):
+        """Test loading valid user environment JSON."""
+        user_env = '{"API_KEY": "secret", "ENV_VAR": "value"}'
+
+        with patch("chainlit.socket.config") as mock_config:
+            mock_config.project.user_env = []
+
+            result = load_user_env(user_env)
+
+            assert result == {"API_KEY": "secret", "ENV_VAR": "value"}
+
+    def test_load_user_env_with_required_keys(self):
+        """Test loading user env with required keys."""
+        user_env = '{"API_KEY": "secret", "OTHER_KEY": "value"}'
+
+        with patch("chainlit.socket.config") as mock_config:
+            mock_config.project.user_env = ["API_KEY", "OTHER_KEY"]
+
+            result = load_user_env(user_env)
+
+            assert result == {"API_KEY": "secret", "OTHER_KEY": "value"}
+
+    def test_load_user_env_missing_required_key(self):
+        """Test error when required key is missing."""
+        user_env = '{"API_KEY": "secret"}'
+
+        with patch("chainlit.socket.config") as mock_config:
+            mock_config.project.user_env = ["API_KEY", "MISSING_KEY"]
+
+            with pytest.raises(
+                ConnectionRefusedError, match="Missing user environment variable"
+            ):
+                load_user_env(user_env)
+
+    def test_load_user_env_none_with_required_keys(self):
+        """Test error when user_env is None but keys are required."""
+        with patch("chainlit.socket.config") as mock_config:
+            mock_config.project.user_env = ["API_KEY"]
+
+            with pytest.raises(
+                ConnectionRefusedError, match="Missing user environment variables"
+            ):
+                load_user_env(None)
+
+    def test_load_user_env_none_without_required_keys(self):
+        """Test when user_env is None and no keys are required."""
+        with patch("chainlit.socket.config") as mock_config:
+            mock_config.project.user_env = []
+
+            result = load_user_env(None)
+
+            assert result == {}
+
+
+class TestCleanSession:
+    """Test suite for clean_session function."""
+
+    @pytest.mark.asyncio
+    async def test_clean_session_with_existing_session(self):
+        """Test marking session for cleanup."""
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.to_clear = False
+
+        with patch.object(WebsocketSession, "get") as mock_get:
+            mock_get.return_value = mock_session
+
+            await clean_session("socket_123")
+
+            assert mock_session.to_clear is True
+            mock_get.assert_called_once_with("socket_123")
+
+    @pytest.mark.asyncio
+    async def test_clean_session_without_session(self):
+        """Test clean_session when session doesn't exist."""
+        with patch.object(WebsocketSession, "get") as mock_get:
+            mock_get.return_value = None
+
+            # Should not raise an error
+            await clean_session("socket_123")
+
+
+class TestSocketEdgeCases:
+    """Test suite for socket edge cases."""
+
+    def test_restore_existing_session_with_none_session_id(self):
+        """Test restore with None session_id."""
+        with patch.object(WebsocketSession, "get_by_id") as mock_get:
+            mock_get.return_value = None
+
+            result = restore_existing_session(None, None, Mock(), Mock(), None)
+
+            assert result is False
+
+    @pytest.mark.asyncio
+    async def test_persist_user_session_with_empty_metadata(self):
+        """Test persisting empty metadata."""
+        mock_data_layer = AsyncMock()
+
+        with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+            mock_get_dl.return_value = mock_data_layer
+
+            await persist_user_session("thread_123", {})
+
+            mock_data_layer.update_thread.assert_called_once_with(
+                thread_id="thread_123", metadata={}
+            )
+
+    def test_load_user_env_with_empty_json(self):
+        """Test loading empty user environment."""
+        user_env = "{}"
+
+        with patch("chainlit.socket.config") as mock_config:
+            mock_config.project.user_env = []
+
+            result = load_user_env(user_env)
+
+            assert result == {}
+
+    @pytest.mark.asyncio
+    async def test_resume_thread_with_empty_metadata(self):
+        """Test resuming thread with empty metadata."""
+        from chainlit.user_session import user_sessions
+
+        mock_session = Mock(spec=WebsocketSession)
+        mock_session.user = Mock(identifier="user123")
+        mock_session.thread_id_to_resume = "thread_123"
+        mock_session.id = "session_123"
+
+        thread = {"userIdentifier": "user123", "metadata": {}}
+
+        mock_data_layer = AsyncMock()
+        mock_data_layer.get_thread.return_value = thread
+
+        original_sessions = user_sessions.copy()
+        try:
+            with patch("chainlit.socket.get_data_layer") as mock_get_dl:
+                mock_get_dl.return_value = mock_data_layer
+
+                result = await resume_thread(mock_session)
+
+                assert result == thread
+                assert user_sessions.get("session_123") == {}
+        finally:
+            user_sessions.clear()
+            user_sessions.update(original_sessions)
+
+    @pytest.mark.asyncio
+    async def test_authenticate_connection_with_exception(self):
+        """Test authentication when get_current_user raises exception."""
+        with patch("chainlit.socket._get_token") as mock_get_token:
+            with patch("chainlit.socket.get_current_user") as mock_get_user:
+                mock_get_token.return_value = "token"
+                mock_get_user.side_effect = Exception("Auth error")
+
+                environ = {"HTTP_COOKIE": "token=token"}
+
+                # Should propagate the exception
+                with pytest.raises(Exception, match="Auth error"):
+                    await _authenticate_connection(environ)
+
+
+class TestConnectionSuccessfulIdempotency:
+    """Regression tests: on_chat_start must fire exactly once per
+    WebsocketSession, even when connection_successful is dispatched multiple
+    times (Socket.IO reconnect, React StrictMode double-mount, reverse-proxy
+    WS 101 retry).
+
+    Refs chainlit#2535, chainlit#2549, chainlit#2228.
+    """
+
+    @pytest.mark.asyncio
+    async def test_on_chat_start_not_duplicated_on_reconnect(
+        self, mock_session_factory
+    ):
+        """Reconnect path (session.restored=True): on_chat_start scheduled once."""
+        on_chat_start = AsyncMock()
+
+        session = mock_session_factory(has_first_interaction=False)
+        session.restored = True
+        session.chat_started = False
+        session.current_task = None
+        session.thread_id_to_resume = None
+
+        mock_context = Mock()
+        mock_context.session = session
+        mock_context.emitter = AsyncMock()
+
+        mock_config = Mock()
+        mock_config.code.on_chat_start = on_chat_start
+        mock_config.code.on_chat_resume = None
+
+        with (
+            patch("chainlit.socket.init_ws_context", return_value=mock_context),
+            patch("chainlit.socket.config", mock_config),
+        ):
+            await connection_successful("sid-1")
+            # Simulate reconnect: same session object, chat_started now True.
+            await connection_successful("sid-1")
+
+        assert on_chat_start.call_count == 1, (
+            "on_chat_start must be scheduled exactly once per WebsocketSession"
+        )
+        assert session.chat_started is True
+
+    @pytest.mark.asyncio
+    async def test_on_chat_start_fires_once_on_fresh_session(
+        self, mock_session_factory
+    ):
+        """Normal one-connect path still greets exactly once."""
+        on_chat_start = AsyncMock()
+
+        session = mock_session_factory(has_first_interaction=False)
+        session.restored = False
+        session.chat_started = False
+        session.current_task = None
+        session.thread_id_to_resume = None
+
+        mock_context = Mock()
+        mock_context.session = session
+        mock_context.emitter = AsyncMock()
+
+        mock_config = Mock()
+        mock_config.code.on_chat_start = on_chat_start
+        mock_config.code.on_chat_resume = None
+
+        with (
+            patch("chainlit.socket.init_ws_context", return_value=mock_context),
+            patch("chainlit.socket.config", mock_config),
+        ):
+            await connection_successful("sid-1")
+
+        assert on_chat_start.call_count == 1
+        assert session.chat_started is True
+
+    @pytest.mark.asyncio
+    async def test_on_chat_start_not_duplicated_on_fresh_then_reconnect(
+        self, mock_session_factory
+    ):
+        """Fresh connect followed by a reconnect still fires exactly once."""
+        on_chat_start = AsyncMock()
+
+        session = mock_session_factory(has_first_interaction=False)
+        session.restored = False
+        session.chat_started = False
+        session.current_task = None
+        session.thread_id_to_resume = None
+
+        mock_context = Mock()
+        mock_context.session = session
+        mock_context.emitter = AsyncMock()
+
+        mock_config = Mock()
+        mock_config.code.on_chat_start = on_chat_start
+        mock_config.code.on_chat_resume = None
+
+        with (
+            patch("chainlit.socket.init_ws_context", return_value=mock_context),
+            patch("chainlit.socket.config", mock_config),
+        ):
+            await connection_successful("sid-1")
+            session.restored = True
+            await connection_successful("sid-1")
+
+        assert on_chat_start.call_count == 1

backend/tests/test_step.py

@@ -0,0 +1,616 @@
+import sys
+import uuid
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from chainlit.context import local_steps
+from chainlit.element import Element
+from chainlit.step import (
+    Step,
+    check_add_step_in_cot,
+    flatten_args_kwargs,
+    step,
+    stub_step,
+)
+
+
+@pytest.mark.asyncio
+class TestStepClass:
+    """Test suite for the Step class."""
+
+    async def test_step_initialization_with_defaults(self, mock_chainlit_context):
+        """Test Step initialization with default values."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test_step")
+
+            assert test_step.name == "test_step"
+            assert test_step.type == "undefined"
+            assert isinstance(test_step.id, str)
+            uuid.UUID(test_step.id)  # Verify valid UUID
+            assert test_step.parent_id is None
+            assert test_step.metadata == {}
+            assert test_step.tags is None
+            assert test_step.is_error is False
+            assert test_step.show_input == "json"
+            assert test_step.language is None
+            assert test_step.default_open is False
+            assert test_step.elements == []
+            assert test_step.streaming is False
+            assert test_step.persisted is False
+            assert test_step.fail_on_persist_error is False
+            assert test_step.input == ""
+            assert test_step.output == ""
+            assert test_step.created_at is not None
+            assert test_step.start is None
+            assert test_step.end is None
+
+    async def test_step_initialization_with_all_fields(self, mock_chainlit_context):
+        """Test Step initialization with all fields provided."""
+        async with mock_chainlit_context:
+            test_id = str(uuid.uuid4())
+            parent_id = str(uuid.uuid4())
+            metadata = {"key": "value"}
+            tags = ["tag1", "tag2"]
+
+            test_step = Step(
+                name="custom_step",
+                type="tool",
+                id=test_id,
+                parent_id=parent_id,
+                metadata=metadata,
+                tags=tags,
+                language="python",
+                default_open=True,
+                show_input=False,
+            )
+
+            assert test_step.name == "custom_step"
+            assert test_step.type == "tool"
+            assert test_step.id == test_id
+            assert test_step.parent_id == parent_id
+            assert test_step.metadata == metadata
+            assert test_step.tags == tags
+            assert test_step.language == "python"
+            assert test_step.default_open is True
+            assert test_step.show_input is False
+
+    async def test_step_input_setter_with_string(self, mock_chainlit_context):
+        """Test Step input setter with string content."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            test_step.input = "This is input text"
+
+            assert test_step.input == "This is input text"
+
+    async def test_step_input_setter_with_dict(self, mock_chainlit_context):
+        """Test Step input setter with dictionary content."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            input_dict = {"param1": "value1", "param2": 42}
+            test_step.input = input_dict
+
+            # Should be JSON formatted
+            assert "param1" in test_step.input
+            assert "value1" in test_step.input
+            assert isinstance(test_step.input, str)
+
+    async def test_step_output_setter_with_string(self, mock_chainlit_context):
+        """Test Step output setter with string content."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            test_step.output = "This is output text"
+
+            assert test_step.output == "This is output text"
+
+    async def test_step_output_setter_with_dict(self, mock_chainlit_context):
+        """Test Step output setter with dictionary content and language detection."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            output_dict = {"result": "success", "data": [1, 2, 3]}
+            test_step.output = output_dict
+
+            # Should be JSON formatted and language set to json
+            assert "result" in test_step.output
+            assert "success" in test_step.output
+            assert test_step.language == "json"
+
+    async def test_step_clean_content_with_bytes(self, mock_chainlit_context):
+        """Test that bytes in content are stripped."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            content_with_bytes = {
+                "text": "hello",
+                "binary": b"binary_data",
+                "nested": {"data": b"more_binary"},
+            }
+            test_step.output = content_with_bytes
+
+            assert "STRIPPED_BINARY_DATA" in test_step.output
+            assert b"binary_data" not in test_step.output.encode()
+
+    async def test_step_to_dict(self, mock_chainlit_context):
+        """Test Step serialization to dictionary."""
+        async with mock_chainlit_context as ctx:
+            test_step = Step(
+                name="test_step",
+                type="tool",
+                metadata={"key": "value"},
+                tags=["tag1"],
+            )
+            test_step.input = "test input"
+            test_step.output = "test output"
+
+            step_dict = test_step.to_dict()
+
+            assert step_dict["name"] == "test_step"
+            assert step_dict["type"] == "tool"
+            assert step_dict["id"] == test_step.id
+            assert step_dict["threadId"] == ctx.session.thread_id
+            assert step_dict["parentId"] is None
+            assert step_dict["streaming"] is False
+            assert step_dict["metadata"] == {"key": "value"}
+            assert step_dict["tags"] == ["tag1"]
+            assert step_dict["input"] == "test input"
+            assert step_dict["output"] == "test output"
+            assert step_dict["isError"] is False
+            assert step_dict["createdAt"] is not None
+            assert step_dict["start"] is None
+            assert step_dict["end"] is None
+
+    async def test_step_send(self, mock_chainlit_context):
+        """Test Step.send() method."""
+        async with mock_chainlit_context as ctx:
+            test_step = Step(name="test_step")
+
+            result = await test_step.send()
+
+            assert result == test_step
+            assert test_step.persisted is False  # No data layer configured
+            ctx.emitter.send_step.assert_called_once()
+
+    async def test_step_send_with_elements(self, mock_chainlit_context):
+        """Test Step.send() with elements."""
+        async with mock_chainlit_context:
+            mock_element = Mock(spec=Element)
+            mock_element.send = AsyncMock()
+
+            test_step = Step(name="test_step", elements=[mock_element])
+
+            await test_step.send()
+
+            mock_element.send.assert_called_once_with(for_id=test_step.id)
+
+    async def test_step_send_already_persisted(self, mock_chainlit_context):
+        """Test that send() returns early if already persisted."""
+        async with mock_chainlit_context as ctx:
+            test_step = Step(name="test_step")
+            test_step.persisted = True
+
+            result = await test_step.send()
+
+            assert result == test_step
+            ctx.emitter.send_step.assert_not_called()
+
+    async def test_step_update(self, mock_chainlit_context):
+        """Test Step.update() method."""
+        async with mock_chainlit_context as ctx:
+            test_step = Step(name="test_step")
+            test_step.streaming = True
+
+            result = await test_step.update()
+
+            assert result is True
+            assert test_step.streaming is False
+            ctx.emitter.update_step.assert_called_once()
+
+    async def test_step_remove(self, mock_chainlit_context):
+        """Test Step.remove() method."""
+        async with mock_chainlit_context as ctx:
+            test_step = Step(name="test_step")
+
+            result = await test_step.remove()
+
+            assert result is True
+            ctx.emitter.delete_step.assert_called_once()
+
+    async def test_step_stream_token_output(self, mock_chainlit_context):
+        """Test streaming tokens to output."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test_step")
+
+            await test_step.stream_token("Hello")
+            await test_step.stream_token(" ")
+            await test_step.stream_token("World")
+
+            assert test_step.output == "Hello World"
+            assert test_step.streaming is True
+
+    async def test_step_stream_token_input(self, mock_chainlit_context):
+        """Test streaming tokens to input."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test_step")
+
+            await test_step.stream_token("Input", is_input=True)
+            await test_step.stream_token(" text", is_input=True)
+
+            assert test_step.input == "Input text"
+
+    async def test_step_stream_token_sequence(self, mock_chainlit_context):
+        """Test streaming tokens with is_sequence flag."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test_step")
+
+            await test_step.stream_token("First", is_sequence=True)
+            await test_step.stream_token("Second", is_sequence=True)
+
+            # With is_sequence, it replaces instead of appending
+            assert test_step.output == "Second"
+
+    async def test_step_stream_token_empty(self, mock_chainlit_context):
+        """Test that empty tokens are ignored."""
+        async with mock_chainlit_context as ctx:
+            test_step = Step(name="test_step")
+
+            await test_step.stream_token("")
+
+            assert test_step.output == ""
+            ctx.emitter.stream_start.assert_not_called()
+
+    async def test_step_context_manager_async(self, mock_chainlit_context):
+        """Test Step as async context manager."""
+        async with mock_chainlit_context as ctx:
+            async with Step(name="context_step") as test_step:
+                assert test_step.start is not None
+                assert test_step.end is None
+
+            # After exiting context
+            assert test_step.end is not None
+            assert ctx.emitter.send_step.call_count == 1
+            assert ctx.emitter.update_step.call_count == 1
+
+    async def test_step_context_manager_with_exception(self, mock_chainlit_context):
+        """Test Step context manager handles exceptions."""
+        async with mock_chainlit_context:
+            try:
+                async with Step(name="error_step") as test_step:
+                    raise ValueError("Test error")
+            except ValueError:
+                pass
+
+            assert test_step.is_error is True
+            assert "Test error" in test_step.output
+
+    async def test_step_parent_id_from_context(self, mock_chainlit_context):
+        """Test that parent_id is set from context when nesting steps."""
+        async with mock_chainlit_context:
+            async with Step(name="parent_step") as parent:
+                async with Step(name="child_step") as child:
+                    assert child.parent_id == parent.id
+
+    async def test_step_local_steps_tracking(self, mock_chainlit_context):
+        """Test that local_steps tracks step hierarchy."""
+        async with mock_chainlit_context:
+            async with Step(name="step1") as step1:
+                steps = local_steps.get()
+                assert step1 in steps
+
+                async with Step(name="step2") as step2:
+                    steps = local_steps.get()
+                    assert step1 in steps
+                    assert step2 in steps
+
+                # After step2 exits
+                steps = local_steps.get()
+                assert step1 in steps
+                assert step2 not in steps
+
+    async def test_step_with_none_input(self, mock_chainlit_context):
+        """Test Step handles None input correctly."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            test_step.input = None
+
+            assert test_step.input == ""
+
+    async def test_step_with_none_output(self, mock_chainlit_context):
+        """Test Step handles None output correctly."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            test_step.output = None
+
+            assert test_step.output == ""
+
+    async def test_step_with_list_content(self, mock_chainlit_context):
+        """Test Step handles list content."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            test_step.output = [1, 2, 3, "four"]
+
+            assert "[" in test_step.output
+            assert "1" in test_step.output
+            assert "four" in test_step.output
+            assert test_step.language == "json"
+
+    async def test_step_with_tuple_content(self, mock_chainlit_context):
+        """Test Step handles tuple content."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            test_step.output = ("a", "b", "c")
+
+            assert test_step.output != ""
+            assert test_step.language == "json"
+
+
+@pytest.mark.asyncio
+class TestStepDecorator:
+    """Test suite for the @step decorator."""
+
+    async def test_step_decorator_async_function(self, mock_chainlit_context):
+        """Test @step decorator on async function."""
+        async with mock_chainlit_context as ctx:
+
+            @step(name="async_step", type="tool")
+            async def async_function(x: int, y: int):
+                return x + y
+
+            result = await async_function(2, 3)
+
+            assert result == 5
+            ctx.emitter.send_step.assert_called()
+
+    async def test_step_decorator_sync_function(self, mock_chainlit_context):
+        """Test @step decorator on sync function."""
+        async with mock_chainlit_context:
+
+            @step(name="sync_step", type="tool")
+            def sync_function(x: int, y: int):
+                return x + y
+
+            result = sync_function(2, 3)
+
+            assert result == 5
+
+    async def test_step_decorator_uses_function_name(self, mock_chainlit_context):
+        """Test that decorator uses function name when name not provided."""
+        async with mock_chainlit_context as ctx:
+
+            @step(type="tool")
+            async def my_custom_function():
+                return "result"
+
+            await my_custom_function()
+
+            # Check that step was created with function name
+            call_args = ctx.emitter.send_step.call_args
+            step_dict = call_args[0][0]
+            assert step_dict["name"] == "my_custom_function"
+
+    async def test_step_decorator_captures_input(self, mock_chainlit_context):
+        """Test that decorator captures function arguments as input."""
+        async with mock_chainlit_context as ctx:
+
+            @step(name="test_step")
+            async def function_with_args(a: str, b: int, c: bool = True):
+                return "done"
+
+            await function_with_args("hello", 42, c=False)
+
+            # Verify send_step was called (input is set during step execution)
+            ctx.emitter.send_step.assert_called()
+
+    async def test_step_decorator_captures_output(self, mock_chainlit_context):
+        """Test that decorator captures function return value as output."""
+        async with mock_chainlit_context as ctx:
+
+            @step(name="test_step")
+            async def function_with_return():
+                return {"status": "success", "value": 123}
+
+            await function_with_return()
+
+            call_args = ctx.emitter.update_step.call_args
+            step_dict = call_args[0][0]
+            assert "status" in step_dict["output"]
+            assert "success" in step_dict["output"]
+
+    async def test_step_decorator_handles_exception(self, mock_chainlit_context):
+        """Test that decorator handles exceptions in wrapped function."""
+        async with mock_chainlit_context as ctx:
+
+            @step(name="error_step")
+            async def function_with_error():
+                raise ValueError("Something went wrong")
+
+            try:
+                await function_with_error()
+            except ValueError:
+                pass
+
+            call_args = ctx.emitter.update_step.call_args
+            step_dict = call_args[0][0]
+            assert step_dict["isError"] is True
+            assert "Something went wrong" in step_dict["output"]
+
+    async def test_step_decorator_with_metadata(self, mock_chainlit_context):
+        """Test decorator with metadata parameter."""
+        async with mock_chainlit_context as ctx:
+            metadata = {"version": "1.0", "author": "test"}
+
+            @step(name="test_step", metadata=metadata)
+            async def function_with_metadata():
+                return "result"
+
+            await function_with_metadata()
+
+            call_args = ctx.emitter.send_step.call_args
+            step_dict = call_args[0][0]
+            assert step_dict["metadata"] == metadata
+
+    async def test_step_decorator_with_tags(self, mock_chainlit_context):
+        """Test decorator with tags parameter."""
+        async with mock_chainlit_context as ctx:
+            tags = ["important", "production"]
+
+            @step(name="test_step", tags=tags)
+            async def function_with_tags():
+                return "result"
+
+            await function_with_tags()
+
+            call_args = ctx.emitter.send_step.call_args
+            step_dict = call_args[0][0]
+            assert step_dict["tags"] == tags
+
+    async def test_step_decorator_without_parentheses(self, mock_chainlit_context):
+        """Test @step decorator without parentheses."""
+        async with mock_chainlit_context as ctx:
+
+            @step
+            async def simple_function():
+                return "result"
+
+            result = await simple_function()
+
+            assert result == "result"
+            ctx.emitter.send_step.assert_called()
+
+
+@pytest.mark.asyncio
+class TestStepHelperFunctions:
+    """Test suite for Step helper functions."""
+
+    def test_flatten_args_kwargs(self):
+        """Test flatten_args_kwargs function."""
+
+        def sample_func(a, b, c=10, d=20):
+            pass
+
+        result = flatten_args_kwargs(sample_func, (1, 2), {"d": 30})
+
+        assert result["a"] == 1
+        assert result["b"] == 2
+        assert result["c"] == 10  # default value
+        assert result["d"] == 30
+
+    def test_flatten_args_kwargs_with_all_kwargs(self):
+        """Test flatten_args_kwargs with all keyword arguments."""
+
+        def sample_func(x, y, z):
+            pass
+
+        result = flatten_args_kwargs(sample_func, (), {"x": 1, "y": 2, "z": 3})
+
+        assert result == {"x": 1, "y": 2, "z": 3}
+
+    async def test_stub_step(self, mock_chainlit_context):
+        """Test stub_step function creates minimal step dict."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test_step", type="tool")
+            test_step.parent_id = "parent_123"
+            test_step.input = "full input"
+            test_step.output = "full output"
+
+            stub = stub_step(test_step)
+
+            assert stub["name"] == "test_step"
+            assert stub["type"] == "tool"
+            assert stub["id"] == test_step.id
+            assert stub["parentId"] == "parent_123"
+            assert stub["threadId"] == test_step.thread_id
+            assert stub["input"] == ""  # Stubbed
+            assert stub["output"] == ""  # Stubbed
+
+    async def test_check_add_step_in_cot_hidden(self, mock_chainlit_context):
+        """Test check_add_step_in_cot with hidden COT."""
+        async with mock_chainlit_context:
+            step_module = sys.modules["chainlit.step"]
+            with patch.object(step_module, "config") as mock_config:
+                mock_config.ui.cot = "hidden"
+
+                # Message types should be added
+                message_step = Step(name="test", type="assistant_message")
+                assert check_add_step_in_cot(message_step) is True
+
+                # Non-message types should not be added
+                tool_step = Step(name="test", type="tool")
+                assert check_add_step_in_cot(tool_step) is False
+
+    async def test_check_add_step_in_cot_visible(self, mock_chainlit_context):
+        """Test check_add_step_in_cot with visible COT."""
+        async with mock_chainlit_context:
+            step_module = sys.modules["chainlit.step"]
+            with patch.object(step_module, "config") as mock_config:
+                mock_config.ui.cot = "visible"
+
+                # All steps should be added
+                tool_step = Step(name="test", type="tool")
+                assert check_add_step_in_cot(tool_step) is True
+
+
+@pytest.mark.asyncio
+class TestStepEdgeCases:
+    """Test suite for Step edge cases and error handling."""
+
+    async def test_step_with_non_serializable_content(self, mock_chainlit_context):
+        """Test Step handles non-JSON-serializable content."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+
+            class NonSerializable:
+                pass
+
+            test_step.output = NonSerializable()
+
+            # Should convert to string
+            assert isinstance(test_step.output, str)
+            assert test_step.language == "text"
+
+    async def test_step_with_very_long_content(self, mock_chainlit_context):
+        """Test Step handles very long content."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+            long_text = "x" * 10000
+
+            test_step.output = long_text
+
+            assert len(test_step.output) == 10000
+
+    async def test_step_multiple_updates(self, mock_chainlit_context):
+        """Test calling update() multiple times."""
+        async with mock_chainlit_context as ctx:
+            test_step = Step(name="test")
+
+            await test_step.update()
+            await test_step.update()
+            await test_step.update()
+
+            assert ctx.emitter.update_step.call_count == 3
+
+    async def test_step_id_uniqueness(self, mock_chainlit_context):
+        """Test that each Step gets a unique ID."""
+        async with mock_chainlit_context:
+            step1 = Step(name="step1")
+            step2 = Step(name="step2")
+            step3 = Step(name="step3")
+
+            ids = {step1.id, step2.id, step3.id}
+            assert len(ids) == 3  # All unique
+
+    async def test_step_with_custom_thread_id(self, mock_chainlit_context):
+        """Test Step with custom thread_id."""
+        async with mock_chainlit_context:
+            custom_thread_id = "custom_thread_123"
+            test_step = Step(name="test", thread_id=custom_thread_id)
+
+            assert test_step.thread_id == custom_thread_id
+
+    async def test_step_fail_on_persist_error_flag(self, mock_chainlit_context):
+        """Test fail_on_persist_error flag behavior."""
+        async with mock_chainlit_context:
+            test_step = Step(name="test")
+
+            assert test_step.fail_on_persist_error is False
+
+            test_step.fail_on_persist_error = True
+            assert test_step.fail_on_persist_error is True

backend/tests/test_teams_adapter.py

@@ -0,0 +1,58 @@
+import importlib
+from unittest.mock import MagicMock, call, patch
+
+
+def _reload_teams_app(monkeypatch, env_vars: dict):
+    """Reload chainlit.teams.app with the given environment variables set."""
+    for key, value in env_vars.items():
+        monkeypatch.setenv(key, value)
+    # Remove keys not in env_vars so tests are independent
+    for key in ("TEAMS_APP_ID", "TEAMS_APP_PASSWORD", "TEAMS_APP_TENANT_ID"):
+        if key not in env_vars:
+            monkeypatch.delenv(key, raising=False)
+
+    teams_mod = importlib.import_module("chainlit.teams.app")
+    return importlib.reload(teams_mod)
+
+
+def test_teams_adapter_without_tenant(monkeypatch):
+    """Omitting TEAMS_APP_TENANT_ID leaves channel_auth_tenant as None (multi-tenant)."""
+    with (
+        patch("botbuilder.core.BotFrameworkAdapterSettings") as mock_settings,
+        patch("botbuilder.core.BotFrameworkAdapter"),
+    ):
+        mock_settings.return_value = MagicMock()
+        _reload_teams_app(
+            monkeypatch,
+            {"TEAMS_APP_ID": "app-id", "TEAMS_APP_PASSWORD": "app-secret"},
+        )
+        assert mock_settings.call_count >= 1
+        assert mock_settings.call_args_list[-1] == call(
+            app_id="app-id",
+            app_password="app-secret",
+            channel_auth_tenant=None,
+        )
+
+
+def test_teams_adapter_with_tenant(monkeypatch):
+    """Setting TEAMS_APP_TENANT_ID forwards the tenant to BotFrameworkAdapterSettings."""
+    tenant_id = "00000000-0000-0000-0000-000000000001"
+    with (
+        patch("botbuilder.core.BotFrameworkAdapterSettings") as mock_settings,
+        patch("botbuilder.core.BotFrameworkAdapter"),
+    ):
+        mock_settings.return_value = MagicMock()
+        _reload_teams_app(
+            monkeypatch,
+            {
+                "TEAMS_APP_ID": "app-id",
+                "TEAMS_APP_PASSWORD": "app-secret",
+                "TEAMS_APP_TENANT_ID": tenant_id,
+            },
+        )
+        assert mock_settings.call_count >= 1
+        assert mock_settings.call_args_list[-1] == call(
+            app_id="app-id",
+            app_password="app-secret",
+            channel_auth_tenant=tenant_id,
+        )

backend/tests/test_translations.py

@@ -0,0 +1,386 @@
+from io import StringIO
+from unittest.mock import patch
+
+import pytest
+
+from chainlit.translations import compare_json_structures, lint_translation_json
+
+
+class TestCompareJsonStructures:
+    """Test suite for compare_json_structures function."""
+
+    def test_compare_identical_structures(self):
+        """Test comparing identical JSON structures."""
+        truth = {"key1": "value1", "key2": "value2"}
+        to_compare = {"key1": "value1", "key2": "value2"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_compare_with_missing_keys(self):
+        """Test when to_compare is missing keys."""
+        truth = {"key1": "value1", "key2": "value2", "key3": "value3"}
+        to_compare = {"key1": "value1"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 2
+        assert "❌ Missing key: 'key2'" in errors
+        assert "❌ Missing key: 'key3'" in errors
+
+    def test_compare_with_extra_keys(self):
+        """Test when to_compare has extra keys."""
+        truth = {"key1": "value1"}
+        to_compare = {"key1": "value1", "key2": "value2", "key3": "value3"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 2
+        assert "⚠️ Extra key: 'key2'" in errors
+        assert "⚠️ Extra key: 'key3'" in errors
+
+    def test_compare_with_both_missing_and_extra_keys(self):
+        """Test when there are both missing and extra keys."""
+        truth = {"key1": "value1", "key2": "value2"}
+        to_compare = {"key1": "value1", "key3": "value3"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 2
+        assert any("Extra key: 'key3'" in e for e in errors)
+        assert any("Missing key: 'key2'" in e for e in errors)
+
+    def test_compare_nested_structures(self):
+        """Test comparing nested JSON structures."""
+        truth = {"level1": {"level2": {"key": "value"}}}
+        to_compare = {"level1": {"level2": {"key": "value"}}}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_compare_nested_with_missing_keys(self):
+        """Test nested structures with missing keys."""
+        truth = {"level1": {"key1": "value1", "key2": "value2"}}
+        to_compare = {"level1": {"key1": "value1"}}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 1
+        assert "❌ Missing key: 'level1.key2'" in errors
+
+    def test_compare_nested_with_extra_keys(self):
+        """Test nested structures with extra keys."""
+        truth = {"level1": {"key1": "value1"}}
+        to_compare = {"level1": {"key1": "value1", "key2": "value2"}}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 1
+        assert "⚠️ Extra key: 'level1.key2'" in errors
+
+    def test_compare_deeply_nested_structures(self):
+        """Test deeply nested structures."""
+        truth = {"a": {"b": {"c": {"d": "value"}}}}
+        to_compare = {"a": {"b": {"c": {}}}}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 1
+        assert "❌ Missing key: 'a.b.c.d'" in errors
+
+    def test_compare_structure_mismatch_dict_vs_value(self):
+        """Test when one is dict and other is value."""
+        truth = {"key": {"nested": "value"}}
+        to_compare = {"key": "not_a_dict"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 1
+        assert "❌ Structure mismatch at: 'key'" in errors
+
+    def test_compare_structure_mismatch_value_vs_dict(self):
+        """Test when truth is value and to_compare is dict."""
+        truth = {"key": "value"}
+        to_compare = {"key": {"nested": "value"}}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 1
+        assert "❌ Structure mismatch at: 'key'" in errors
+
+    def test_compare_with_non_dict_input_truth(self):
+        """Test error when truth is not a dict."""
+        with pytest.raises(ValueError, match="Both inputs must be dictionaries"):
+            compare_json_structures("not_a_dict", {})
+
+    def test_compare_with_non_dict_input_to_compare(self):
+        """Test error when to_compare is not a dict."""
+        with pytest.raises(ValueError, match="Both inputs must be dictionaries"):
+            compare_json_structures({}, "not_a_dict")
+
+    def test_compare_with_both_non_dict_inputs(self):
+        """Test error when both inputs are not dicts."""
+        with pytest.raises(ValueError, match="Both inputs must be dictionaries"):
+            compare_json_structures("not_a_dict", "also_not_a_dict")
+
+    def test_compare_empty_dicts(self):
+        """Test comparing empty dictionaries."""
+        truth = {}
+        to_compare = {}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_compare_empty_truth_with_data(self):
+        """Test when truth is empty but to_compare has data."""
+        truth = {}
+        to_compare = {"key1": "value1", "key2": "value2"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 2
+        assert all("Extra key" in e for e in errors)
+
+    def test_compare_empty_to_compare_with_data(self):
+        """Test when to_compare is empty but truth has data."""
+        truth = {"key1": "value1", "key2": "value2"}
+        to_compare = {}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 2
+        assert all("Missing key" in e for e in errors)
+
+    def test_compare_with_different_value_types(self):
+        """Test that different value types at leaf nodes don't cause errors."""
+        truth = {"key1": "string", "key2": 123, "key3": True}
+        to_compare = {"key1": "different", "key2": 456, "key3": False}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        # Structure matches, so no errors (values are not compared)
+        assert errors == []
+
+    def test_compare_complex_nested_structure(self):
+        """Test complex nested structure with multiple levels."""
+        truth = {
+            "app": {
+                "title": "My App",
+                "settings": {"theme": "dark", "language": "en"},
+            },
+            "user": {"name": "John", "preferences": {"notifications": True}},
+        }
+        to_compare = {
+            "app": {
+                "title": "My App",
+                "settings": {"theme": "light"},  # Missing 'language'
+            },
+            "user": {
+                "name": "Jane",
+                "preferences": {"notifications": False, "extra": "value"},  # Extra key
+            },
+        }
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 2
+        assert any("Missing key: 'app.settings.language'" in e for e in errors)
+        assert any("Extra key: 'user.preferences.extra'" in e for e in errors)
+
+    def test_compare_with_null_values(self):
+        """Test structures with None/null values."""
+        truth = {"key1": None, "key2": "value"}
+        to_compare = {"key1": None, "key2": "value"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_compare_with_list_values(self):
+        """Test structures with list values (treated as leaf nodes)."""
+        truth = {"key1": ["a", "b", "c"], "key2": "value"}
+        to_compare = {"key1": ["x", "y"], "key2": "value"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        # Lists are leaf nodes, structure matches
+        assert errors == []
+
+    def test_compare_path_formatting(self):
+        """Test that error paths are formatted correctly."""
+        truth = {"a": {"b": {"c": "value"}}}
+        to_compare = {"a": {"b": {}}}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 1
+        assert "a.b.c" in errors[0]
+        assert not errors[0].startswith(".")
+
+
+class TestLintTranslationJson:
+    """Test suite for lint_translation_json function."""
+
+    def test_lint_with_no_errors(self):
+        """Test linting when there are no errors."""
+        truth = {"key1": "value1", "key2": "value2"}
+        to_compare = {"key1": "value1", "key2": "value2"}
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            lint_translation_json("test.json", truth, to_compare)
+            output = fake_out.getvalue()
+
+            assert "Linting test.json..." in output
+            assert "✅ No errors found in test.json" in output
+
+    def test_lint_with_errors(self):
+        """Test linting when there are errors."""
+        truth = {"key1": "value1", "key2": "value2"}
+        to_compare = {"key1": "value1", "key3": "value3"}
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            lint_translation_json("test.json", truth, to_compare)
+            output = fake_out.getvalue()
+
+            assert "Linting test.json..." in output
+            assert "Missing key: 'key2'" in output
+            assert "Extra key: 'key3'" in output
+            assert "✅ No errors found" not in output
+
+    def test_lint_with_nested_errors(self):
+        """Test linting with nested structure errors."""
+        truth = {"level1": {"key1": "value1", "key2": "value2"}}
+        to_compare = {"level1": {"key1": "value1"}}
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            lint_translation_json("nested.json", truth, to_compare)
+            output = fake_out.getvalue()
+
+            assert "Linting nested.json..." in output
+            assert "Missing key: 'level1.key2'" in output
+
+    def test_lint_with_structure_mismatch(self):
+        """Test linting with structure mismatch."""
+        truth = {"key": {"nested": "value"}}
+        to_compare = {"key": "not_nested"}
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            lint_translation_json("mismatch.json", truth, to_compare)
+            output = fake_out.getvalue()
+
+            assert "Linting mismatch.json..." in output
+            assert "Structure mismatch at: 'key'" in output
+
+    def test_lint_with_multiple_errors(self):
+        """Test linting with multiple types of errors."""
+        truth = {
+            "key1": "value1",
+            "key2": {"nested": "value"},
+            "key3": "value3",
+        }
+        to_compare = {
+            "key1": "value1",
+            "key2": "not_nested",
+            "key4": "extra",
+        }
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            lint_translation_json("multi.json", truth, to_compare)
+            output = fake_out.getvalue()
+
+            assert "Linting multi.json..." in output
+            assert "Structure mismatch" in output
+            assert "Missing key: 'key3'" in output
+            assert "Extra key: 'key4'" in output
+
+    def test_lint_output_format(self):
+        """Test that lint output is properly formatted."""
+        truth = {"key1": "value1"}
+        to_compare = {"key2": "value2"}
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            lint_translation_json("format.json", truth, to_compare)
+            output = fake_out.getvalue()
+
+            # Check that output starts with newline and linting message
+            lines = output.strip().split("\n")
+            assert "Linting format.json..." in lines[0]
+            assert len(lines) >= 2  # At least linting message + errors
+
+
+class TestTranslationsEdgeCases:
+    """Test suite for edge cases in translations module."""
+
+    def test_compare_with_numeric_keys(self):
+        """Test structures with numeric keys (as strings)."""
+        truth = {"1": "value1", "2": "value2"}
+        to_compare = {"1": "value1", "2": "value2"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_compare_with_special_characters_in_keys(self):
+        """Test keys with special characters."""
+        truth = {"key-1": "value", "key_2": "value", "key.3": "value"}
+        to_compare = {"key-1": "value", "key_2": "value", "key.3": "value"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_compare_with_unicode_keys(self):
+        """Test keys with unicode characters."""
+        truth = {"键": "value", "clé": "value", "مفتاح": "value"}
+        to_compare = {"键": "value", "clé": "value", "مفتاح": "value"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_compare_very_deeply_nested(self):
+        """Test very deeply nested structures."""
+        truth = {"a": {"b": {"c": {"d": {"e": {"f": "value"}}}}}}
+        to_compare = {"a": {"b": {"c": {"d": {"e": {}}}}}}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert len(errors) == 1
+        assert "a.b.c.d.e.f" in errors[0]
+
+    def test_compare_with_empty_string_values(self):
+        """Test structures with empty string values."""
+        truth = {"key1": "", "key2": "value"}
+        to_compare = {"key1": "", "key2": "value"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        assert errors == []
+
+    def test_lint_with_empty_filename(self):
+        """Test lint with empty filename."""
+        truth = {"key": "value"}
+        to_compare = {"key": "value"}
+
+        with patch("sys.stdout", new=StringIO()) as fake_out:
+            lint_translation_json("", truth, to_compare)
+            output = fake_out.getvalue()
+
+            assert "Linting ..." in output
+
+    def test_compare_preserves_error_order(self):
+        """Test that errors are reported in a consistent order."""
+        truth = {"a": "1", "b": "2", "c": "3"}
+        to_compare = {"d": "4", "e": "5"}
+
+        errors = compare_json_structures(truth, to_compare)
+
+        # Should have 2 extra keys and 3 missing keys
+        assert len(errors) == 5
+        extra_errors = [e for e in errors if "Extra" in e]
+        missing_errors = [e for e in errors if "Missing" in e]
+        assert len(extra_errors) == 2
+        assert len(missing_errors) == 3

backend/tests/test_user_session.py

@@ -0,0 +1,14 @@
+async def test_user_session_set_get(mock_chainlit_context, user_session):
+    async with mock_chainlit_context as context:
+        # Test setting a value
+        user_session.set("test_key", "test_value")
+
+        # Test getting the value
+        assert user_session.get("test_key") == "test_value"
+
+        # Test getting a default value for a non-existent key
+        assert user_session.get("non_existent_key", "default") == "default"
+
+        # Test getting session-related values
+        assert user_session.get("id") == context.session.id
+        assert user_session.get("env") == context.session.user_env

backend/tests/test_utils.py

@@ -0,0 +1,449 @@
+import os
+import tempfile
+from datetime import datetime, timezone
+from unittest.mock import AsyncMock, patch
+
+import click
+import pytest
+
+from chainlit.utils import (
+    check_file,
+    check_module_version,
+    make_module_getattr,
+    timestamp_utc,
+    utc_now,
+    wrap_user_function,
+)
+
+
+class TestUtcNow:
+    """Test suite for utc_now function."""
+
+    def test_utc_now_returns_string(self):
+        """Test that utc_now returns a string."""
+        result = utc_now()
+        assert isinstance(result, str)
+
+    def test_utc_now_ends_with_z(self):
+        """Test that utc_now returns ISO format with Z suffix."""
+        result = utc_now()
+        assert result.endswith("Z")
+
+    def test_utc_now_is_iso_format(self):
+        """Test that utc_now returns valid ISO format."""
+        result = utc_now()
+        # Remove the Z and parse
+        dt_str = result[:-1]
+        # Should be parseable as ISO format
+        datetime.fromisoformat(dt_str)
+
+    def test_utc_now_is_current_time(self):
+        """Test that utc_now returns approximately current time."""
+        before = datetime.now(timezone.utc).replace(tzinfo=None)
+        result = utc_now()
+        after = datetime.now(timezone.utc).replace(tzinfo=None)
+
+        # Parse the result (naive datetime)
+        result_dt = datetime.fromisoformat(result[:-1])
+
+        # Should be between before and after (with some tolerance for microseconds)
+        assert (
+            before.replace(microsecond=0) <= result_dt <= after.replace(microsecond=0)
+            or before <= result_dt <= after
+        )
+
+    def test_utc_now_multiple_calls(self):
+        """Test that multiple calls to utc_now return different values."""
+        result1 = utc_now()
+        result2 = utc_now()
+
+        # Results should be very close but might differ
+        assert isinstance(result1, str)
+        assert isinstance(result2, str)
+
+
+class TestTimestampUtc:
+    """Test suite for timestamp_utc function."""
+
+    def test_timestamp_utc_returns_string(self):
+        """Test that timestamp_utc returns a string."""
+        result = timestamp_utc(1234567890.0)
+        assert isinstance(result, str)
+
+    def test_timestamp_utc_ends_with_z(self):
+        """Test that timestamp_utc returns ISO format with Z suffix."""
+        result = timestamp_utc(1234567890.0)
+        assert result.endswith("Z")
+
+    def test_timestamp_utc_converts_correctly(self):
+        """Test that timestamp_utc converts timestamp correctly."""
+        # Known timestamp: 2009-02-13 23:31:30 UTC
+        timestamp = 1234567890.0
+        result = timestamp_utc(timestamp)
+
+        # Parse and verify
+        dt = datetime.fromisoformat(result[:-1])
+        assert dt.year == 2009
+        assert dt.month == 2
+        assert dt.day == 13
+
+    def test_timestamp_utc_with_zero(self):
+        """Test timestamp_utc with epoch (0)."""
+        result = timestamp_utc(0.0)
+        dt = datetime.fromisoformat(result[:-1])
+        assert dt.year == 1970
+        assert dt.month == 1
+        assert dt.day == 1
+
+    def test_timestamp_utc_with_fractional_seconds(self):
+        """Test timestamp_utc with fractional seconds."""
+        timestamp = 1234567890.123456
+        result = timestamp_utc(timestamp)
+
+        # Should be valid ISO format
+        dt = datetime.fromisoformat(result[:-1])
+        assert isinstance(dt, datetime)
+
+    def test_timestamp_utc_with_negative_timestamp(self):
+        """Test timestamp_utc with negative timestamp (before epoch)."""
+        # 1969-12-31 23:00:00 UTC
+        timestamp = -3600.0
+        result = timestamp_utc(timestamp)
+
+        dt = datetime.fromisoformat(result[:-1])
+        assert dt.year == 1969
+
+
+@pytest.mark.asyncio
+class TestWrapUserFunction:
+    """Test suite for wrap_user_function."""
+
+    async def test_wrap_user_function_with_sync_function(self, mock_chainlit_context):
+        """Test wrapping a synchronous function."""
+        async with mock_chainlit_context:
+
+            def user_func(a, b):
+                return a + b
+
+            wrapped = wrap_user_function(user_func)
+            result = await wrapped(5, 3)
+
+            assert result == 8
+
+    async def test_wrap_user_function_with_async_function(self, mock_chainlit_context):
+        """Test wrapping an asynchronous function."""
+        async with mock_chainlit_context:
+
+            async def user_func(x, y):
+                return x * y
+
+            wrapped = wrap_user_function(user_func)
+            result = await wrapped(4, 7)
+
+            assert result == 28
+
+    async def test_wrap_user_function_with_no_args(self, mock_chainlit_context):
+        """Test wrapping a function with no arguments."""
+        async with mock_chainlit_context:
+
+            def user_func():
+                return "hello"
+
+            wrapped = wrap_user_function(user_func)
+            result = await wrapped()
+
+            assert result == "hello"
+
+    async def test_wrap_user_function_with_task(self, mock_chainlit_context):
+        """Test wrapping a function with task management."""
+        async with mock_chainlit_context as ctx:
+            ctx.emitter.task_start = AsyncMock()
+            ctx.emitter.task_end = AsyncMock()
+
+            def user_func(value):
+                return value * 2
+
+            wrapped = wrap_user_function(user_func, with_task=True)
+            result = await wrapped(10)
+
+            assert result == 20
+            ctx.emitter.task_start.assert_called_once()
+            ctx.emitter.task_end.assert_called_once()
+
+    async def test_wrap_user_function_handles_exception(self, mock_chainlit_context):
+        """Test that wrapped function handles exceptions."""
+        async with mock_chainlit_context:
+
+            def user_func():
+                raise ValueError("Test error")
+
+            wrapped = wrap_user_function(user_func)
+            result = await wrapped()
+
+            # Should return None when exception occurs
+            assert result is None
+
+    async def test_wrap_user_function_with_task_handles_exception(
+        self, mock_chainlit_context
+    ):
+        """Test that wrapped function with task handles exceptions."""
+        async with mock_chainlit_context as ctx:
+            ctx.emitter.task_start = AsyncMock()
+            ctx.emitter.task_end = AsyncMock()
+
+            def user_func():
+                raise ValueError("Test error")
+
+            with patch("chainlit.utils.logger") as mock_logger:
+                wrapped = wrap_user_function(user_func, with_task=True)
+                result = await wrapped()
+
+                assert result is None
+                ctx.emitter.task_start.assert_called_once()
+                ctx.emitter.task_end.assert_called_once()
+                mock_logger.exception.assert_called_once()
+
+    async def test_wrap_user_function_preserves_function_metadata(
+        self, mock_chainlit_context
+    ):
+        """Test that wrapping preserves function metadata."""
+        async with mock_chainlit_context:
+
+            def user_func(a, b):
+                """Test function docstring."""
+                return a + b
+
+            wrapped = wrap_user_function(user_func)
+
+            assert wrapped.__name__ == "user_func"
+            assert wrapped.__doc__ == "Test function docstring."
+
+    async def test_wrap_user_function_with_kwargs(self, mock_chainlit_context):
+        """Test wrapping a function and calling with positional args."""
+        async with mock_chainlit_context:
+
+            def user_func(x, y, z):
+                return x + y + z
+
+            wrapped = wrap_user_function(user_func)
+            result = await wrapped(1, 2, 3)
+
+            assert result == 6
+
+
+class TestMakeModuleGetattr:
+    """Test suite for make_module_getattr."""
+
+    def test_make_module_getattr_creates_function(self):
+        """Test that make_module_getattr creates a function."""
+        registry = {"SomeClass": "some.module"}
+        getattr_func = make_module_getattr(registry)
+
+        assert callable(getattr_func)
+
+    def test_make_module_getattr_imports_module(self):
+        """Test that the created function imports modules."""
+        # Use a real module for testing
+        registry = {"datetime": "datetime"}
+        getattr_func = make_module_getattr(registry)
+
+        result = getattr_func("datetime")
+        assert result is datetime
+
+    def test_make_module_getattr_with_nested_module(self):
+        """Test with nested module path."""
+        registry = {"timezone": "datetime"}
+        getattr_func = make_module_getattr(registry)
+
+        result = getattr_func("timezone")
+        assert result is timezone
+
+
+class TestCheckModuleVersion:
+    """Test suite for check_module_version."""
+
+    def test_check_module_version_with_installed_module(self):
+        """Test checking version of an installed module."""
+        # pytest should be installed
+        result = check_module_version("pytest", "1.0.0")
+        assert result is True
+
+    def test_check_module_version_with_higher_required_version(self):
+        """Test with a required version higher than installed."""
+        # Require an impossibly high version
+        result = check_module_version("pytest", "999.0.0")
+        assert result is False
+
+    def test_check_module_version_with_nonexistent_module(self):
+        """Test with a module that doesn't exist."""
+        result = check_module_version("nonexistent_module_xyz", "1.0.0")
+        assert result is False
+
+    def test_check_module_version_exact_match(self):
+        """Test with exact version match."""
+        # Get actual pytest version
+        result = check_module_version("pytest", pytest.__version__)
+        assert result is True
+
+    def test_check_module_version_with_builtin_module(self):
+        """Test with a builtin module that has no __version__."""
+        # os module doesn't have __version__
+        with pytest.raises(AttributeError):
+            check_module_version("os", "1.0.0")
+
+
+class TestCheckFile:
+    """Test suite for check_file function."""
+
+    def test_check_file_with_valid_py_file(self):
+        """Test check_file with a valid .py file."""
+        with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as f:
+            temp_file = f.name
+
+        try:
+            # Should not raise any exception
+            check_file(temp_file)
+        finally:
+            os.unlink(temp_file)
+
+    def test_check_file_with_valid_py3_file(self):
+        """Test check_file with a valid .py3 file."""
+        with tempfile.NamedTemporaryFile(suffix=".py3", delete=False) as f:
+            temp_file = f.name
+
+        try:
+            # Should not raise any exception
+            check_file(temp_file)
+        finally:
+            os.unlink(temp_file)
+
+    def test_check_file_with_invalid_extension(self):
+        """Test check_file with invalid file extension."""
+        with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
+            temp_file = f.name
+
+        try:
+            with pytest.raises(click.BadArgumentUsage) as exc_info:
+                check_file(temp_file)
+            assert ".txt" in str(exc_info.value)
+        finally:
+            os.unlink(temp_file)
+
+    def test_check_file_with_no_extension(self):
+        """Test check_file with file that has no extension."""
+        with tempfile.NamedTemporaryFile(suffix="", delete=False) as f:
+            temp_file = f.name
+
+        try:
+            with pytest.raises(click.BadArgumentUsage) as exc_info:
+                check_file(temp_file)
+            assert "no extension" in str(exc_info.value)
+        finally:
+            os.unlink(temp_file)
+
+    def test_check_file_with_nonexistent_file(self):
+        """Test check_file with a file that doesn't exist."""
+        nonexistent_file = "/path/to/nonexistent/file.py"
+
+        with pytest.raises(click.BadParameter) as exc_info:
+            check_file(nonexistent_file)
+        assert "does not exist" in str(exc_info.value)
+
+    def test_check_file_with_json_extension(self):
+        """Test check_file with .json extension."""
+        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
+            temp_file = f.name
+
+        try:
+            with pytest.raises(click.BadArgumentUsage) as exc_info:
+                check_file(temp_file)
+            assert ".json" in str(exc_info.value)
+        finally:
+            os.unlink(temp_file)
+
+
+class TestUtilsEdgeCases:
+    """Test suite for utils edge cases."""
+
+    def test_utc_now_format_consistency(self):
+        """Test that utc_now format is consistent across calls."""
+        results = [utc_now() for _ in range(5)]
+
+        for result in results:
+            # All should have same format
+            assert result.endswith("Z")
+            assert "T" in result
+            # Should be parseable
+            datetime.fromisoformat(result[:-1])
+
+    def test_timestamp_utc_with_large_timestamp(self):
+        """Test timestamp_utc with very large timestamp (far future)."""
+        # Year 2100
+        timestamp = 4102444800.0
+        result = timestamp_utc(timestamp)
+
+        dt = datetime.fromisoformat(result[:-1])
+        assert dt.year == 2100
+
+    @pytest.mark.asyncio
+    async def test_wrap_user_function_with_multiple_exceptions(
+        self, mock_chainlit_context
+    ):
+        """Test wrapped function handles different exception types."""
+        async with mock_chainlit_context:
+            exceptions = [ValueError("error1"), TypeError("error2"), KeyError("error3")]
+
+            for exc in exceptions:
+
+                def user_func():
+                    raise exc
+
+                with patch("chainlit.utils.logger"):
+                    wrapped = wrap_user_function(user_func)
+                    result = await wrapped()
+                    assert result is None
+
+    def test_check_file_with_relative_path(self):
+        """Test check_file with relative path."""
+        # Create a temp file in current directory
+        with tempfile.NamedTemporaryFile(suffix=".py", delete=False, dir=".") as f:
+            temp_file = os.path.basename(f.name)
+
+        try:
+            # Should work with relative path
+            check_file(temp_file)
+        finally:
+            os.unlink(temp_file)
+
+    def test_check_file_with_absolute_path(self):
+        """Test check_file with absolute path."""
+        with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as f:
+            temp_file = os.path.abspath(f.name)
+
+        try:
+            # Should work with absolute path
+            check_file(temp_file)
+        finally:
+            os.unlink(temp_file)
+
+    def test_make_module_getattr_with_empty_registry(self):
+        """Test make_module_getattr with empty registry."""
+        registry = {}
+        getattr_func = make_module_getattr(registry)
+
+        with pytest.raises(KeyError):
+            getattr_func("nonexistent")
+
+    @pytest.mark.asyncio
+    async def test_wrap_user_function_with_default_args(self, mock_chainlit_context):
+        """Test wrapping function with default arguments."""
+        async with mock_chainlit_context:
+
+            def user_func(a, b=10):
+                return a + b
+
+            wrapped = wrap_user_function(user_func)
+
+            # Call with only required arg
+            result = await wrapped(5)
+            assert result == 15