Merge pull request #159 from BESSER-PEARL/hotfix/initial_message_not_sent

Hotfix/initial message not sent

Author: Aran30

besser/agent/core/session.py

@@ -5,6 +5,7 @@
 from asyncio import TimerHandle
 from collections import deque
 from typing import Any, TYPE_CHECKING
+from datetime import datetime
 
 from pandas import DataFrame
 from websocket import WebSocketApp
@@ -142,7 +143,7 @@ def _stop_event_thread(self) -> None:
         self._event_loop = None
         self._event_thread = None
 
-    def get_chat_history(self, n: int = None) -> list[Message]:
+    def get_chat_history(self, n: int = None, until_timestamp: datetime = None) -> list[Message]:
         """Get the history of messages between this session and its agent.
 
         Args:
@@ -154,7 +155,7 @@ def get_chat_history(self, n: int = None) -> list[Message]:
         """
         chat_history: list[Message] = []
         if self._agent.get_property(DB_MONITORING) and self._agent._monitoring_db.connected:
-            chat_df: DataFrame = self._agent._monitoring_db.select_chat(self, n=n)
+            chat_df: DataFrame = self._agent._monitoring_db.select_chat(self, n=n, until_timestamp=until_timestamp)
             for i, row in chat_df.iterrows():
                 t = get_message_type(row['type'])
                 chat_history.append(Message(t=t, content=row['content'], is_user=row['is_user'], timestamp=row['timestamp']))

besser/agent/db/monitoring_db.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Optional
 
 import json
 import pandas as pd
@@ -432,13 +432,18 @@ def get_last_state_of_session(self, agent_name: str, platform_name: str, session
         result_transition = self.conn.execute(stmt_transition).first()
         return result_transition[0] if result_transition else None
 
-    def select_chat(self, session: Session, n: int) -> pd.DataFrame:
-        """Retrieves a conversation history from the chat table of the database.
+    def select_chat(
+        self,
+        session: Session,
+        n: Optional[int] = None,
+        until_timestamp: Optional[datetime] = None,
+    ) -> pd.DataFrame:
+        """Retrieve chat messages for a session, optionally capped by count or timestamp.
 
         Args:
-            session (Session): the session to get from the database
-            n (int or None): the number of messages to get (from the most recents). If none is provided, gets all the
-                messages
+            session (Session): Session whose messages should be returned.
+            n (int | None): Optional cap for the most recent messages; None returns all.
+            until_timestamp (datetime | None): Optional inclusive upper bound on message timestamps.
         Returns:
             pandas.DataFrame: the session record, should be a 1 row DataFrame
 
@@ -448,6 +453,9 @@ def select_chat(self, session: Session, n: int) -> pd.DataFrame:
         stmt = (select(table).where(
             table.c.session_id == int(session_entry['id'][0])
         ))
+        if until_timestamp is not None:
+            stmt = stmt.where(table.c.timestamp <= until_timestamp)
+
         if n:
             stmt = stmt.order_by(desc(table.c.id)).limit(n)
         return pd.read_sql_query(stmt, self.conn).sort_values(by='id')

besser/agent/platforms/payload.py

@@ -96,18 +96,16 @@ def decode(payload_str):
         payload_dict = json.loads(payload_str)
         payload_action = payload_dict['action']
         payload_message = payload_dict['message']
-        user_id = payload_dict.get('user_id')
         history = payload_dict.get('history', False)
 
         for action in PayloadAction:
             if action.value == payload_action:
-                return Payload(action, payload_message, user_id=user_id, history=history)
+                return Payload(action, payload_message, history=history)
         return None
 
-    def __init__(self, action: PayloadAction, message: str or dict = None, user_id: str = None, history: bool = False):
+    def __init__(self, action: PayloadAction, message: str or dict = None, history: bool = False):
         self.action: str = action.value
         self.message: str or dict = message
-        self.user_id: str = user_id
         self.history: bool = history
 
 
@@ -135,7 +133,6 @@ def default(self, obj):
             payload_dict = {
                 'action': obj.action,
                 'message': obj.message,
-                'user_id': getattr(obj, 'user_id', None),
                 'history': getattr(obj, 'history', None),
             }
             return payload_dict

besser/agent/platforms/websocket/streamlit_ui/chat.py

@@ -73,7 +73,6 @@ def send_option():
                 payload = Payload(
                     action=PayloadAction.USER_MESSAGE,
                     message=option,
-                    user_id=st.session_state.get("username", "Guest"),
                 )
                 ws = ensure_websocket_connection()
                 if not ws:
@@ -140,7 +139,6 @@ def load_chat():
                 payload = Payload(
                     action=PayloadAction.FETCH_USER_MESSAGES,
                     message=None,
-                    user_id=username,
                 )
                 try:
                     ws.send(json.dumps(payload, cls=PayloadEncoder))

besser/agent/platforms/websocket/streamlit_ui/initialization.py

@@ -42,15 +42,27 @@ def _resolve_host_port():
 
 def _start_websocket(host: str, port: str):
     try:
-        ws = websocket.WebSocketApp(
-            f"ws://{host}:{port}/",
-            on_open=on_open,
-            on_message=on_message,
-            on_error=on_error,
-            on_close=on_close,
-            on_ping=on_ping,
-            on_pong=on_pong,
-        )
+        if st.session_state.get("username"):
+            ws = websocket.WebSocketApp(
+                f"ws://{host}:{port}/",
+                header={"X-User-ID": st.session_state["username"]},
+                on_open=on_open,
+                on_message=on_message,
+                on_error=on_error,
+                on_close=on_close,
+                on_ping=on_ping,
+                on_pong=on_pong,
+            )
+        else:
+            ws = websocket.WebSocketApp(
+                f"ws://{host}:{port}/",
+                on_open=on_open,
+                on_message=on_message,
+                on_error=on_error,
+                on_close=on_close,
+                on_ping=on_ping,
+                on_pong=on_pong,
+            )
         websocket_thread = threading.Thread(target=ws.run_forever)
         add_script_run_ctx(websocket_thread)
         websocket_thread.start()

besser/agent/platforms/websocket/streamlit_ui/message_input.py

@@ -23,12 +23,7 @@ def submit_text():
             st.write(user_input)
         message = Message(t=MessageType.STR, content=user_input, is_user=True, timestamp=datetime.now())
         st.session_state.history.append(message)
-        if st.session_state.get("authenticated", False):
-            payload = Payload(action=PayloadAction.USER_MESSAGE,
-                            message=user_input, user_id=st.session_state.get("username"))
-        else:
-            payload = Payload(action=PayloadAction.USER_MESSAGE,
-                            message=user_input)
+        payload = Payload(action=PayloadAction.USER_MESSAGE, message=user_input)
         try:
             ws = st.session_state[WEBSOCKET]
             ws.send(json.dumps(payload, cls=PayloadEncoder))

besser/agent/platforms/websocket/websocket_platform.py

@@ -6,6 +6,7 @@
 import os
 import time
 from datetime import datetime
+from urllib.parse import parse_qs, urlsplit
 
 import numpy as np
 import subprocess
@@ -36,6 +37,28 @@
     DB_STREAMLIT
 )
 
+def _extract_user_id_from_request(request) -> str | None:
+    if not request:
+        return None
+    for attr in ("path", "raw_path", "uri"):
+        value = getattr(request, attr, None)
+        if not value:
+            continue
+        if isinstance(value, bytes):
+            # Prefer UTF-8 for URL paths; fall back to latin-1 to remain robust to non-UTF-8 bytes.
+            try:
+                value = value.decode("utf-8")
+            except UnicodeDecodeError:
+                value = value.decode("latin-1", errors="replace")
+        query = urlsplit(value).query
+        if not query:
+            continue
+        params = parse_qs(query)
+        user_values = params.get("user_id")
+        if user_values:
+            return user_values[0]
+    return None
+
 if TYPE_CHECKING:
     from besser.agent.core.agent import Agent
 
@@ -101,33 +124,34 @@ def message_handler(conn: ServerConnection) -> None:
                 conn (ServerConnection): the user connection
             """
             session: Session = None
+            current_time = datetime.now()
+            request = getattr(conn, "request", None)
+            headers = getattr(request, "headers", {}) if request else {}
+            header_user = headers.get("X-User-ID") if hasattr(headers, "get") else None
+            query_user = _extract_user_id_from_request(request)
+            session_key = header_user or query_user or str(conn.id)
+            self._connections[str(session_key)] = conn
+            session = self._agent.get_or_create_session(session_key, self)
             try:
+
                 for payload_str in conn:
                     if not self.running:
                         raise ConnectionClosedError(None, None)
                     payload: Payload = Payload.decode(payload_str)
-                    if session is None:
-                        if payload.user_id:
-                            session = self._agent.get_or_create_session(payload.user_id, self)
-                            self._connections[str(payload.user_id)] = conn
-                        else:
-                            session = self._agent.get_or_create_session(str(conn.id), self)
-                            self._connections[str(conn.id)] = conn
+
                     if payload.action == PayloadAction.FETCH_USER_MESSAGES.value:
                         try:
-                            chat_history = session.get_chat_history()
+                            chat_history = session.get_chat_history(until_timestamp=current_time)
                             for message in chat_history:
                                 history_payload = None
                                 if message.is_user:
                                     history_payload = Payload(action=PayloadAction.USER_MESSAGE,
                                                               message=message.content,
-                                                              user_id=session.id,
                                                               history=True
                                                               )
                                 else:
                                     history_payload = Payload(action=PayloadAction.AGENT_REPLY_STR,
                                                               message=message.content,
-                                                              user_id=session.id,
                                                               history=True
                                                               )
                                 self._send(session.id, history_payload)
@@ -271,7 +295,6 @@ def reply(self, session: Session, message: str) -> None:
         session.save_message(Message(t=MessageType.STR, content=message, is_user=False, timestamp=datetime.now()))
         payload = Payload(action=PayloadAction.AGENT_REPLY_STR,
                           message=message,
-                          user_id=session.id
                           )
         payload.message = self._agent.process(session=session, message=payload.message, is_user_message=False)
         self._send(session.id, payload)

docs/source/release_notes.rst

@@ -3,6 +3,7 @@ Release Notes
 
 .. toctree::
 
+   release_notes/v4.2.2
    release_notes/v4.2.1
    release_notes/v4.2.0
    release_notes/v4.1.0

docs/source/release_notes/v4.2.2.rst

@@ -0,0 +1,7 @@
+Version 4.2.2
+=============
+
+Hotfix
+------
+
+- Restored initial agent messages when persistence is enabled by updating the user identification protocol to accept the user ID from either the WebSocket header or query parameters.

docs/source/wiki/platforms/websocket_platform.rst

@@ -130,19 +130,37 @@ The WebSocket platform allows the following kinds of user messages:
 
 Enabling persistent user sessions
 ---------------------------------
-When building your own UI on top of the WebSocket API, you need to implement a user authentication mechanism to enable session persistence. 
-BAF relies on user identifiers to map incoming connections to the correct sessions.
-Thus, if your UI does not authenticate users, each new connection will be treated as a new user session.
-Once your platform authenticates users, you simply need to set a unique identifier for each user in the payload sent to the agent via WebSocket:
+When building your own UI on top of the WebSocket API, implement user authentication so every connection can be tied to a stable identifier. BAF maps connections to sessions using either the ``X-User-ID`` header or the ``user_id`` query parameter from the handshake URL. If both are provided, the header takes precedence. If neither identifier is available, the platform treats the connection as a new anonymous user.
+
+Once your client authenticates users, include the identifier in the WebSocket handshake headers:
+
+.. code:: python
+
+    ws = websocket.WebSocketApp(
+        f"ws://{host}:{port}/",
+        header={"X-User-ID": user_id},
+        on_open=on_open,
+        on_message=on_message,
+        on_error=on_error,
+        on_close=on_close,
+        on_ping=on_ping,
+        on_pong=on_pong,
+    )
+
+If you cannot control the HTTP headers, you can also attach the identifier as a query parameter on the WebSocket URL:
 
 .. code:: python
 
-    payload_dict = {
-        'action': obj.action,
-        'message': obj.message,
-        'user_id': getattr(obj, 'user_id', None),
-        'history': getattr(obj, 'history', None),
-    }
+    ws = websocket.WebSocketApp(
+        f"ws://{host}:{port}/?user_id={user_id}",
+        on_open=on_open,
+        on_message=on_message,
+        on_error=on_error,
+        on_close=on_close,
+        on_ping=on_ping,
+        on_pong=on_pong,
+    )
+
 
 On the agent's side, you'll need to start the monitoring database and set the ``persist_sessions=True`` parameter when initializing the agent: