Repository: Chainlit/chainlit (language: python)

Issue:
Avatar image is broken when reading threads from history
**Describe the bug**

The PR https://github.com/Chainlit/chainlit/pull/1326 introduces a bug when it expects the `get_avatar` to match string (https://github.com/Chainlit/chainlit/pull/1326/files#diff-9a73e67e163edc3fae6839bf2bcc8c6ab650624eff6668502388a7a86b7ece78R964) . 

The Frontend app can make requests to the avatar that do not match the Regex (https://github.com/Chainlit/chainlit/blob/main/frontend/src/components/molecules/messages/components/Avatar.tsx#L31) this creates a bug for existing applications where the `config.ui.name` contains space (e.g. `My Assistant`). 


**To Reproduce**
Set your config.ui.name to `My Assistant`

1. Create a thread
2. Publish messages
3. Create a new thread
4. Try to resume the previous thread
5. See the 400 error and the avatar not showing

**Expected behavior**
Avatar should be loaded for threads in history.





Discussion / hints:
<!-- Answer -->
To fix the bug where the `get_avatar` function fails to load avatars correctly for threads in history when `config.ui.name` contains spaces, you need to ensure that the `author` parameter in the URL is properly encoded. This can be done using `encodeURIComponent` to handle spaces and other special characters.

Here is the modified code for the `MessageAvatar` component:

```tsx
import { useContext, useMemo } from 'react';

import { Tooltip } from '@mui/material';
import Avatar from '@mui/material/Avatar';

import {
  ChainlitContext,
  useChatSession,
  useConfig
} from '@chainlit/react-client';

interface Props {
  author?: string;
  hide?: boolean;
}

const MessageAvatar = ({ author, hide }: Props) => {
  const apiClient = useContext(ChainlitContext);
  const { chatProfile } = useChatSession();
  const { config } = useConfig();

  const selectedChatProfile = useMemo(() => {
    return config?.chatProfiles.find((profile) => profile.name === chatProfile);
  }, [config, chatProfile]);

  const avatarUrl = useMemo(() => {
    const isAssistant = !author || author === config?.ui.name;
    if (isAssistant && selectedChatProfile?.icon) {
      return selectedChatProfile.icon;
    }
    return apiClient?.buildEndpoint(`/avatars/${encodeURIComponent(author || 'default')}`);
  }, [apiClient, selectedChatProfile, config, author]);

  return (
    <span className={`message-avatar`}>
      <Tooltip title={author}>
        <Avatar
          sx={{
            width: '1.6rem',
            height: '1.6rem',
            bgcolor: 'transparent'
          }}
          src={hide ? undefined : avatarUrl}
        />
      </Tooltip>
    </span>
  );
};

export { MessageAvatar };
```

By using `encodeURIComponent` on the `author` parameter, you ensure that any spaces or special characters are correctly encoded, preventing the 400 error <sup>[[1]](https://github.com/Chainlit/chainlit/blob/main/frontend/src/components/molecules/messages/components/Avatar.tsx)</sup>.

Additionally, the `get_avatar` function in the backend handles the `avatar_id` parameter by validating it, normalizing it, and searching for a matching avatar file. Here is the implementation:

```python
@router.get("/avatars/{avatar_id:str}")
async def get_avatar(avatar_id: str):
    """Get the avatar for the user based on the avatar_id."""
    if not re.match(r"^[a-zA-Z0-9_-]+$", avatar_id):
        raise HTTPException(status_code=400, detail="Invalid avatar_id")

    if avatar_id == "default":
        avatar_id = config.ui.name

    avatar_id = avatar_id.strip().lower().replace(" ", "_")

    base_path = Path(APP_ROOT) / "public" / "avatars"
    avatar_pattern = f"{avatar_id}.*"

    matching_files = base_path.glob(avatar_pattern)

    if avatar_path := next(matching_files, None):
        if not is_path_inside(avatar_path, base_path):
            raise HTTPException(status_code=400, detail="Invalid filename")

        media_type, _ = mimetypes.guess_type(str(avatar_path))

        return FileResponse(avatar_path, media_type=media_type)

    return await get_favicon()
```

This function ensures that the `avatar_id` is properly formatted and validated before attempting to retrieve the avatar file <sup>[[2]](https://github.com/Chainlit/chainlit/blob/main/backend/chainlit/server.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*

@LucasMagnum Love to hear whether the fix suggested by our bot does the trick. If it does, we'd love a PR for that, preferably including a(n update to a) E2E test demonstrating the issue and it's resolution.

@dosu Any suggestion as to which is the most applicable place to write such a (regression) test, preferably without adding a new E2E (because they're slow)? What could such a test look like?
<!-- Answer -->
The most applicable place to write a regression test for the `get_avatar` function without adding a new E2E test is in the `backend/tests/test_server.py` file. This file already contains tests for the `get_avatar` function, so adding a regression test here would be consistent with the existing structure.

Here is an example of what such a regression test could look like:

```python
def test_get_avatar_regression(test_client: TestClient, monkeypatch: pytest.MonkeyPatch):
    """Regression test for get_avatar function."""
    # Setup: Create a custom avatar
    custom_avatar_path = os.path.join(
        APP_ROOT, "public", "avatars", "regression_avatar.png"
    )
    os.makedirs(os.path.dirname(custom_avatar_path), exist_ok=True)
    with open(custom_avatar_path, "wb") as f:
        f.write(b"regression test image data")

    # Test: Fetch the custom avatar
    response = test_client.get("/avatars/regression_avatar")
    assert response.status_code == 200
    assert response.headers["content-type"].startswith("image/")
    assert response.content == b"regression test image data"

    # Clean up
    os.remove(custom_avatar_path)
```

This test ensures that the `get_avatar` function correctly handles a custom avatar and returns the expected image data <sup>[[1]](https://github.com/Chainlit/chainlit/blob/main/backend/tests/test_server.py)</sup><sup>[[2]](https://github.com/Chainlit/chainlit/blob/main/backend/chainlit/server.py)</sup>.


<!-- Dosu Comment Footer -->
*To continue talking to [Dosu](https://dosu.dev), mention @dosu.*

> The most applicable place to write a regression test for the `get_avatar` function without adding a new E2E test is in the `backend/tests/test_server.py` file. This file already contains tests for the `get_avatar` function, so adding a regression test here would be consistent with the existing structure.
> 
> Here is an example of what such a regression test could look like:
> 
> ```python
> def test_get_avatar_regression(test_client: TestClient, monkeypatch: pytest.MonkeyPatch):
>     """Regression test for get_avatar function."""
>     # Setup: Create a custom avatar
>     custom_avatar_path = os.path.join(
>         APP_ROOT, "public", "avatars", "regression_avatar.png"
>     )
>     os.makedirs(os.path.dirname(custom_avatar_path), exist_ok=True)
>     with open(custom_avatar_path, "wb") as f:
>         f.write(b"regression test image data")
> 
>     # Test: Fetch the custom avatar
>     response = test_client.get("/avatars/regression_avatar")
>     assert response.status_code == 200
>     assert response.headers["content-type"].startswith("image/")
>     assert response.content == b"regression test image data"
> 
>     # Clean up
>     os.remove(custom_avatar_path)
> ```
> 
> This test ensures that the `get_avatar` function correctly handles a custom avatar and returns the expected image data [[1]](https://github.com/Chainlit/chainlit/blob/main/backend/tests/test_server.py)[[2]](https://github.com/Chainlit/chainlit/blob/main/backend/chainlit/server.py).
> 
> _To continue talking to [Dosu](https://dosu.dev), mention @dosu._

But our specific test case concerns a space in `config.ui.name` -- and the requirement of URL encoding in the frontend, right?
Hi @dokterbob, was there a reason for preventing avatar names from containing spaces with the regex as such? I had to monkey patch the `get_avatar()` function when updating to 1.200 from 1.1.404 to modify the regex to allow spaces.

Changing the regex to this:
```python
if not re.match(r"^[a-zA-Z0-9_. -]+$", avatar_id):
    raise HTTPException(status_code=400, detail="Invalid avatar_id")
```

Completely fixed the issue for me.

I've just been able to replicate the issue in our deployment as well. Full stack trace etc: https://synergy-os.sentry.io/share/issue/b23e6a9655a1421e9c3273e2d4a13342/

> Hi @dokterbob, was there a reason for preventing avatar names from containing spaces with the regex as such? I had to monkey patch the `get_avatar()` function when updating to 1.200 from 1.1.404 to modify the regex to allow spaces.
> 
> Changing the regex to this:
> 
> ```python
> if not re.match(r"^[a-zA-Z0-9_. -]+$", avatar_id):
>     raise HTTPException(status_code=400, detail="Invalid avatar_id")
> ```
> 
> Completely fixed the issue for me.

To ensure proper security it is essential to always sanitise any sort of user input. That is, any unreasonable/unexpected/possibly failing or erroneous input should be rejected outright.

Allowing spaces in URL filenames IMHO definitely counts amongst these, i.e. the encoding of various types of unicode spaces can expose attack vectors. It is of course unfortunate if it breaks user experience and we should take care to address it, but not at the expense of increasing security exposure.

Perhaps the solution to this could be to [slugify file names](https://github.com/django/django/blob/stable/5.1.x/django/utils/text.py#L452)? Ref: https://docs.djangoproject.com/en/5.1/ref/utils/#django.utils.text.slugify

In that case, having a `config.ui.name` of 'Funny Banana` would give an avatar of `funny_banana.*` (limiting file extensions here would be another one...).
I've given this some thought, I think indeed it would be best to 'slugify' the URL on the client side (not on the server side). I'll need to analyse more carefully what avatar's specifically are supposed to do, happy for people thinking along towards a spec for a solution.
@dokterbob will the fix for this be able to make it into the next release?
@dokterbob Are you saying that avatars (Assistants and Steps) **should not** be able to have spaces? I sure hope not because nearly all of my Avatars have spaces in the name, as do most people who have provided screenshots.

This issue is blocking many from upgrading past 1.1.404.
> https://github.com/Chainlit/chainlit/blob/main/frontend/src/components/molecules/messages/components/Avatar.tsx#L31

What I'm saying is a clean fix is to move filename sanitation (`avatar_id.strip().lower().replace(" ", "_")`) to the client side and enforce 'clean' URL's. 

I realise we need to fix this sooner though, I will make a PR with @AidanShipperley's fix to allow spaces in avatars.

Fix patch:
diff --git a/backend/chainlit/server.py b/backend/chainlit/server.py
index 597830ee43..d7d0b4db87 100644
--- a/backend/chainlit/server.py
+++ b/backend/chainlit/server.py
@@ -961,7 +961,7 @@ async def get_logo(theme: Optional[Theme] = Query(Theme.light)):
 @router.get("/avatars/{avatar_id:str}")
 async def get_avatar(avatar_id: str):
     """Get the avatar for the user based on the avatar_id."""
-    if not re.match(r"^[a-zA-Z0-9_-]+$", avatar_id):
+    if not re.match(r"^[a-zA-Z0-9_ -]+$", avatar_id):
         raise HTTPException(status_code=400, detail="Invalid avatar_id")
 
     if avatar_id == "default":


Test patch:
diff --git a/backend/tests/test_server.py b/backend/tests/test_server.py
index 6e98bca782..91b152030f 100644
--- a/backend/tests/test_server.py
+++ b/backend/tests/test_server.py
@@ -174,6 +174,24 @@ def test_get_avatar_custom(test_client: TestClient, monkeypatch: pytest.MonkeyPa
     os.remove(custom_avatar_path)
 
 
+def test_get_avatar_with_spaces(
+    test_client: TestClient, monkeypatch: pytest.MonkeyPatch
+):
+    """Test with custom avatar."""
+    custom_avatar_path = os.path.join(APP_ROOT, "public", "avatars", "my_assistant.png")
+    os.makedirs(os.path.dirname(custom_avatar_path), exist_ok=True)
+    with open(custom_avatar_path, "wb") as f:
+        f.write(b"fake image data")
+
+    response = test_client.get("/avatars/My Assistant")
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("image/")
+    assert response.content == b"fake image data"
+
+    # Clean up
+    os.remove(custom_avatar_path)
+
+
 def test_get_avatar_non_existent_favicon(
     test_client: TestClient, monkeypatch: pytest.MonkeyPatch
 ):