fix(frontend): prevent chat message loss during websocket disconnections or page refresh (#13380)

enterprise/migrations/versions/101_add_pending_messages_table.py

@@ -0,0 +1,39 @@
+"""Add pending_messages table for server-side message queuing
+
+Revision ID: 101
+Revises: 100
+Create Date: 2025-03-15 00:00:00.000000
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = '101'
+down_revision: Union[str, None] = '100'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Create pending_messages table for storing messages before conversation is ready.
+
+    Messages are stored temporarily until the conversation becomes ready, then
+    delivered and deleted regardless of success or failure.
+    """
+    op.create_table(
+        'pending_messages',
+        sa.Column('id', sa.String(), primary_key=True),
+        sa.Column('conversation_id', sa.String(), nullable=False, index=True),
+        sa.Column('role', sa.String(20), nullable=False, server_default='user'),
+        sa.Column('content', sa.JSON, nullable=False),
+        sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),
+    )
+
+
+def downgrade() -> None:
+    """Remove pending_messages table."""
+    op.drop_table('pending_messages')

enterprise/server/utils/saas_pending_message_injector.py

@@ -0,0 +1,172 @@
+"""Enterprise injector for PendingMessageService with SAAS filtering."""
+
+from typing import AsyncGenerator
+from uuid import UUID
+
+from fastapi import Request
+from sqlalchemy import select
+from storage.stored_conversation_metadata_saas import StoredConversationMetadataSaas
+from storage.user import User
+
+from openhands.agent_server.models import ImageContent, TextContent
+from openhands.app_server.errors import AuthError
+from openhands.app_server.pending_messages.pending_message_models import (
+    PendingMessageResponse,
+)
+from openhands.app_server.pending_messages.pending_message_service import (
+    PendingMessageService,
+    PendingMessageServiceInjector,
+    SQLPendingMessageService,
+)
+from openhands.app_server.services.injector import InjectorState
+from openhands.app_server.user.specifiy_user_context import ADMIN
+from openhands.app_server.user.user_context import UserContext
+
+
+class SaasSQLPendingMessageService(SQLPendingMessageService):
+    """Extended SQLPendingMessageService with user and organization-based filtering.
+
+    This enterprise version ensures that:
+    - Users can only queue messages for conversations they own
+    - Organization isolation is enforced for multi-tenant deployments
+    """
+
+    def __init__(self, db_session, user_context: UserContext):
+        super().__init__(db_session=db_session)
+        self.user_context = user_context
+
+    async def _get_current_user(self) -> User | None:
+        """Get the current user using the existing db_session.
+
+        Returns:
+            User object or None if no user_id is available
+        """
+        user_id_str = await self.user_context.get_user_id()
+        if not user_id_str:
+            return None
+
+        user_id_uuid = UUID(user_id_str)
+        result = await self.db_session.execute(
+            select(User).where(User.id == user_id_uuid)
+        )
+        return result.scalars().first()
+
+    async def _validate_conversation_ownership(self, conversation_id: str) -> None:
+        """Validate that the current user owns the conversation.
+
+        This ensures multi-tenant isolation by checking:
+        - The conversation belongs to the current user
+        - The conversation belongs to the user's current organization
+
+        Args:
+            conversation_id: The conversation ID to validate (can be task-id or UUID)
+
+        Raises:
+            AuthError: If user doesn't own the conversation or authentication fails
+        """
+        # For internal operations (e.g., processing pending messages during startup)
+        # we need a mode that bypasses filtering. The ADMIN context enables this.
+        if self.user_context == ADMIN:
+            return
+
+        user_id_str = await self.user_context.get_user_id()
+        if not user_id_str:
+            raise AuthError('User authentication required')
+
+        user_id_uuid = UUID(user_id_str)
+
+        # Check conversation ownership via SAAS metadata
+        query = select(StoredConversationMetadataSaas).where(
+            StoredConversationMetadataSaas.conversation_id == conversation_id
+        )
+        result = await self.db_session.execute(query)
+        saas_metadata = result.scalar_one_or_none()
+
+        # If no SAAS metadata exists, the conversation might be a new task-id
+        # that hasn't been linked to a conversation yet. Allow access in this case
+        # as the message will be validated when the conversation is created.
+        if saas_metadata is None:
+            return
+
+        # Verify user ownership
+        if saas_metadata.user_id != user_id_uuid:
+            raise AuthError('You do not have access to this conversation')
+
+        # Verify organization ownership if applicable
+        user = await self._get_current_user()
+        if user and user.current_org_id is not None:
+            if saas_metadata.org_id != user.current_org_id:
+                raise AuthError('Conversation belongs to a different organization')
+
+    async def add_message(
+        self,
+        conversation_id: str,
+        content: list[TextContent | ImageContent],
+        role: str = 'user',
+    ) -> PendingMessageResponse:
+        """Queue a message with ownership validation.
+
+        Args:
+            conversation_id: The conversation ID to queue the message for
+            content: Message content
+            role: Message role (default: 'user')
+
+        Returns:
+            PendingMessageResponse with the queued message info
+
+        Raises:
+            AuthError: If user doesn't own the conversation
+        """
+        await self._validate_conversation_ownership(conversation_id)
+        return await super().add_message(conversation_id, content, role)
+
+    async def get_pending_messages(self, conversation_id: str):
+        """Get pending messages with ownership validation.
+
+        Args:
+            conversation_id: The conversation ID to get messages for
+
+        Returns:
+            List of pending messages
+
+        Raises:
+            AuthError: If user doesn't own the conversation
+        """
+        await self._validate_conversation_ownership(conversation_id)
+        return await super().get_pending_messages(conversation_id)
+
+    async def count_pending_messages(self, conversation_id: str) -> int:
+        """Count pending messages with ownership validation.
+
+        Args:
+            conversation_id: The conversation ID to count messages for
+
+        Returns:
+            Number of pending messages
+
+        Raises:
+            AuthError: If user doesn't own the conversation
+        """
+        await self._validate_conversation_ownership(conversation_id)
+        return await super().count_pending_messages(conversation_id)
+
+
+class SaasPendingMessageServiceInjector(PendingMessageServiceInjector):
+    """Enterprise injector for PendingMessageService with SAAS filtering."""
+
+    async def inject(
+        self, state: InjectorState, request: Request | None = None
+    ) -> AsyncGenerator[PendingMessageService, None]:
+        from openhands.app_server.config import (
+            get_db_session,
+            get_user_context,
+        )
+
+        async with (
+            get_user_context(state, request) as user_context,
+            get_db_session(state, request) as db_session,
+        ):
+            service = SaasSQLPendingMessageService(
+                db_session=db_session, user_context=user_context
+            )
+            yield service