diff --git a/surfsense_backend/alembic/versions/65_add_message_author_id.py b/surfsense_backend/alembic/versions/65_add_message_author_id.py index 8d891db81..2253c2d4e 100644 --- a/surfsense_backend/alembic/versions/65_add_message_author_id.py +++ b/surfsense_backend/alembic/versions/65_add_message_author_id.py @@ -37,10 +37,5 @@ def upgrade() -> None: def downgrade() -> None: """Remove author_id column from new_chat_messages table.""" - op.execute( - """ - DROP INDEX IF EXISTS ix_new_chat_messages_author_id; - ALTER TABLE new_chat_messages - DROP COLUMN IF EXISTS author_id; - """ - ) + op.execute("DROP INDEX IF EXISTS ix_new_chat_messages_author_id") + op.execute("ALTER TABLE new_chat_messages DROP COLUMN IF EXISTS author_id") diff --git a/surfsense_backend/alembic/versions/68_add_chat_comments_table.py b/surfsense_backend/alembic/versions/68_add_chat_comments_table.py new file mode 100644 index 000000000..01c11f3e1 --- /dev/null +++ b/surfsense_backend/alembic/versions/68_add_chat_comments_table.py @@ -0,0 +1,52 @@ +"""Add chat_comments table for comments on AI responses + +Revision ID: 68 +Revises: 67 +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "68" +down_revision: str | None = "67" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Create chat_comments table.""" + op.execute( + """ + CREATE TABLE IF NOT EXISTS chat_comments ( + id SERIAL PRIMARY KEY, + message_id INTEGER NOT NULL REFERENCES new_chat_messages(id) ON DELETE CASCADE, + parent_id INTEGER REFERENCES chat_comments(id) ON DELETE CASCADE, + author_id UUID REFERENCES "user"(id) ON DELETE SET NULL, + content TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """ + ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_comments_message_id ON chat_comments(message_id)" + ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_comments_parent_id ON chat_comments(parent_id)" + ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_comments_author_id ON chat_comments(author_id)" + ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_comments_created_at ON chat_comments(created_at)" + ) + + +def downgrade() -> None: + """Drop chat_comments table.""" + op.execute( + """ + DROP TABLE IF EXISTS chat_comments; + """ + ) diff --git a/surfsense_backend/alembic/versions/69_add_chat_comment_mentions_table.py b/surfsense_backend/alembic/versions/69_add_chat_comment_mentions_table.py new file mode 100644 index 000000000..c8d5c9c9a --- /dev/null +++ b/surfsense_backend/alembic/versions/69_add_chat_comment_mentions_table.py @@ -0,0 +1,41 @@ +"""Add chat_comment_mentions table for @mentions in comments + +Revision ID: 69 +Revises: 68 +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "69" +down_revision: str | None = "68" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Create chat_comment_mentions table.""" + op.execute( + """ + CREATE TABLE IF NOT EXISTS chat_comment_mentions ( + id SERIAL PRIMARY KEY, + comment_id INTEGER NOT NULL REFERENCES chat_comments(id) ON DELETE CASCADE, + mentioned_user_id UUID NOT NULL REFERENCES "user"(id) ON DELETE CASCADE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE (comment_id, mentioned_user_id) + ) + """ + ) + op.execute( + "CREATE INDEX IF NOT EXISTS idx_chat_comment_mentions_comment_id ON chat_comment_mentions(comment_id)" + ) + + +def downgrade() -> None: + """Drop chat_comment_mentions table.""" + op.execute( + """ + DROP TABLE IF EXISTS chat_comment_mentions; + """ + ) diff --git a/surfsense_backend/alembic/versions/70_add_comments_permissions_to_roles.py b/surfsense_backend/alembic/versions/70_add_comments_permissions_to_roles.py new file mode 100644 index 000000000..ec7df8b56 --- /dev/null +++ b/surfsense_backend/alembic/versions/70_add_comments_permissions_to_roles.py @@ -0,0 +1,94 @@ +"""Add comments permissions to existing roles + +Revision ID: 70 +Revises: 69 +Create Date: 2024-01-16 + +""" + +from sqlalchemy import text + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "70" +down_revision = "69" +branch_labels = None +depends_on = None + + +def upgrade(): + connection = op.get_bind() + + # Add comments:create to Admin, Editor, Viewer roles (if not already present) + connection.execute( + text( + """ + UPDATE search_space_roles + SET permissions = array_append(permissions, 'comments:create') + WHERE name IN ('Admin', 'Editor', 'Viewer') + AND NOT ('comments:create' = ANY(permissions)) + """ + ) + ) + + # Add comments:read to Admin, Editor, Viewer roles (if not already present) + connection.execute( + text( + """ + UPDATE search_space_roles + SET permissions = array_append(permissions, 'comments:read') + WHERE name IN ('Admin', 'Editor', 'Viewer') + AND NOT ('comments:read' = ANY(permissions)) + """ + ) + ) + + # Add comments:delete to Admin roles only (if not already present) + connection.execute( + text( + """ + UPDATE search_space_roles + SET permissions = array_append(permissions, 'comments:delete') + WHERE name = 'Admin' + AND NOT ('comments:delete' = ANY(permissions)) + """ + ) + ) + + +def downgrade(): + connection = op.get_bind() + + # Remove comments:create from Admin, Editor, Viewer roles + connection.execute( + text( + """ + UPDATE search_space_roles + SET permissions = array_remove(permissions, 'comments:create') + WHERE name IN ('Admin', 'Editor', 'Viewer') + """ + ) + ) + + # Remove comments:read from Admin, Editor, Viewer roles + connection.execute( + text( + """ + UPDATE search_space_roles + SET permissions = array_remove(permissions, 'comments:read') + WHERE name IN ('Admin', 'Editor', 'Viewer') + """ + ) + ) + + # Remove comments:delete from Admin roles only + connection.execute( + text( + """ + UPDATE search_space_roles + SET permissions = array_remove(permissions, 'comments:delete') + WHERE name = 'Admin' + """ + ) + ) diff --git a/surfsense_backend/alembic/versions/71_add_comments_electric_replication.py b/surfsense_backend/alembic/versions/71_add_comments_electric_replication.py new file mode 100644 index 000000000..d99843584 --- /dev/null +++ b/surfsense_backend/alembic/versions/71_add_comments_electric_replication.py @@ -0,0 +1,59 @@ +"""Add Electric SQL replication for chat_comment_mentions table + +Revision ID: 71 +Revises: 70 + +Enables Electric SQL replication for the chat_comment_mentions table to support +real-time live updates for mentions. +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "71" +down_revision: str | None = "70" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Enable Electric SQL replication for chat_comment_mentions table.""" + op.execute("ALTER TABLE chat_comment_mentions REPLICA IDENTITY FULL;") + + op.execute( + """ + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_publication_tables + WHERE pubname = 'electric_publication_default' + AND tablename = 'chat_comment_mentions' + ) THEN + ALTER PUBLICATION electric_publication_default ADD TABLE chat_comment_mentions; + END IF; + END + $$; + """ + ) + + +def downgrade() -> None: + """Remove chat_comment_mentions from Electric SQL replication.""" + op.execute( + """ + DO $$ + BEGIN + IF EXISTS ( + SELECT 1 FROM pg_publication_tables + WHERE pubname = 'electric_publication_default' + AND tablename = 'chat_comment_mentions' + ) THEN + ALTER PUBLICATION electric_publication_default DROP TABLE chat_comment_mentions; + END IF; + END + $$; + """ + ) + + op.execute("ALTER TABLE chat_comment_mentions REPLICA IDENTITY DEFAULT;") diff --git a/surfsense_backend/app/db.py b/surfsense_backend/app/db.py index ee0d0724d..2b514483a 100644 --- a/surfsense_backend/app/db.py +++ b/surfsense_backend/app/db.py @@ -152,6 +152,11 @@ class Permission(str, Enum): CHATS_UPDATE = "chats:update" CHATS_DELETE = "chats:delete" + # Comments + COMMENTS_CREATE = "comments:create" + COMMENTS_READ = "comments:read" + COMMENTS_DELETE = "comments:delete" + # LLM Configs LLM_CONFIGS_CREATE = "llm_configs:create" LLM_CONFIGS_READ = "llm_configs:read" @@ -209,6 +214,10 @@ DEFAULT_ROLE_PERMISSIONS = { Permission.CHATS_READ.value, Permission.CHATS_UPDATE.value, Permission.CHATS_DELETE.value, + # Comments + Permission.COMMENTS_CREATE.value, + Permission.COMMENTS_READ.value, + Permission.COMMENTS_DELETE.value, # LLM Configs Permission.LLM_CONFIGS_CREATE.value, Permission.LLM_CONFIGS_READ.value, @@ -252,6 +261,9 @@ DEFAULT_ROLE_PERMISSIONS = { Permission.CHATS_READ.value, Permission.CHATS_UPDATE.value, Permission.CHATS_DELETE.value, + # Comments (no delete) + Permission.COMMENTS_CREATE.value, + Permission.COMMENTS_READ.value, # LLM Configs (read only) Permission.LLM_CONFIGS_READ.value, Permission.LLM_CONFIGS_CREATE.value, @@ -279,6 +291,9 @@ DEFAULT_ROLE_PERMISSIONS = { Permission.DOCUMENTS_READ.value, # Chats (read only) Permission.CHATS_READ.value, + # Comments (no delete) + Permission.COMMENTS_CREATE.value, + Permission.COMMENTS_READ.value, # LLM Configs (read only) Permission.LLM_CONFIGS_READ.value, # Podcasts (read only) @@ -424,6 +439,84 @@ class NewChatMessage(BaseModel, TimestampMixin): # Relationships thread = relationship("NewChatThread", back_populates="messages") author = relationship("User") + comments = relationship( + "ChatComment", + back_populates="message", + cascade="all, delete-orphan", + ) + + +class ChatComment(BaseModel, TimestampMixin): + """ + Comment model for comments on AI chat responses. + Supports one level of nesting (replies to comments, but no replies to replies). + """ + + __tablename__ = "chat_comments" + + message_id = Column( + Integer, + ForeignKey("new_chat_messages.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + parent_id = Column( + Integer, + ForeignKey("chat_comments.id", ondelete="CASCADE"), + nullable=True, + index=True, + ) + author_id = Column( + UUID(as_uuid=True), + ForeignKey("user.id", ondelete="SET NULL"), + nullable=True, + index=True, + ) + content = Column(Text, nullable=False) + updated_at = Column( + TIMESTAMP(timezone=True), + nullable=False, + default=lambda: datetime.now(UTC), + onupdate=lambda: datetime.now(UTC), + index=True, + ) + + # Relationships + message = relationship("NewChatMessage", back_populates="comments") + author = relationship("User") + parent = relationship( + "ChatComment", remote_side="ChatComment.id", backref="replies" + ) + mentions = relationship( + "ChatCommentMention", + back_populates="comment", + cascade="all, delete-orphan", + ) + + +class ChatCommentMention(BaseModel, TimestampMixin): + """ + Tracks @mentions in chat comments for notification purposes. + """ + + __tablename__ = "chat_comment_mentions" + + comment_id = Column( + Integer, + ForeignKey("chat_comments.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + mentioned_user_id = Column( + UUID(as_uuid=True), + ForeignKey("user.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + + # Relationships + comment = relationship("ChatComment", back_populates="mentions") + mentioned_user = relationship("User") class Document(BaseModel, TimestampMixin): diff --git a/surfsense_backend/app/routes/__init__.py b/surfsense_backend/app/routes/__init__.py index 84ce86451..d3091af59 100644 --- a/surfsense_backend/app/routes/__init__.py +++ b/surfsense_backend/app/routes/__init__.py @@ -3,6 +3,7 @@ from fastapi import APIRouter from .airtable_add_connector_route import ( router as airtable_add_connector_router, ) +from .chat_comments_routes import router as chat_comments_router from .circleback_webhook_route import router as circleback_webhook_router from .clickup_add_connector_route import router as clickup_add_connector_router from .confluence_add_connector_route import router as confluence_add_connector_router @@ -43,6 +44,7 @@ router.include_router(editor_router) router.include_router(documents_router) router.include_router(notes_router) router.include_router(new_chat_router) # Chat with assistant-ui persistence +router.include_router(chat_comments_router) router.include_router(podcasts_router) # Podcast task status and audio router.include_router(search_source_connectors_router) router.include_router(google_calendar_add_connector_router) diff --git a/surfsense_backend/app/routes/chat_comments_routes.py b/surfsense_backend/app/routes/chat_comments_routes.py new file mode 100644 index 000000000..1c21c0f4a --- /dev/null +++ b/surfsense_backend/app/routes/chat_comments_routes.py @@ -0,0 +1,95 @@ +""" +Routes for chat comments and mentions. +""" + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db import User, get_async_session +from app.schemas.chat_comments import ( + CommentCreateRequest, + CommentListResponse, + CommentReplyResponse, + CommentResponse, + CommentUpdateRequest, + MentionListResponse, +) +from app.services.chat_comments_service import ( + create_comment, + create_reply, + delete_comment, + get_comments_for_message, + get_user_mentions, + update_comment, +) +from app.users import current_active_user + +router = APIRouter() + + +@router.get("/messages/{message_id}/comments", response_model=CommentListResponse) +async def list_comments( + message_id: int, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """List all comments for a message with their replies.""" + return await get_comments_for_message(session, message_id, user) + + +@router.post("/messages/{message_id}/comments", response_model=CommentResponse) +async def add_comment( + message_id: int, + request: CommentCreateRequest, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """Create a top-level comment on an AI response.""" + return await create_comment(session, message_id, request.content, user) + + +@router.post("/comments/{comment_id}/replies", response_model=CommentReplyResponse) +async def add_reply( + comment_id: int, + request: CommentCreateRequest, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """Reply to an existing comment.""" + return await create_reply(session, comment_id, request.content, user) + + +@router.put("/comments/{comment_id}", response_model=CommentReplyResponse) +async def edit_comment( + comment_id: int, + request: CommentUpdateRequest, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """Update a comment's content (author only).""" + return await update_comment(session, comment_id, request.content, user) + + +@router.delete("/comments/{comment_id}") +async def remove_comment( + comment_id: int, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """Delete a comment (author or user with COMMENTS_DELETE permission).""" + return await delete_comment(session, comment_id, user) + + +# ============================================================================= +# Mention Routes +# ============================================================================= + + +@router.get("/mentions", response_model=MentionListResponse) +async def list_mentions( + search_space_id: int | None = None, + session: AsyncSession = Depends(get_async_session), + user: User = Depends(current_active_user), +): + """List mentions for the current user.""" + return await get_user_mentions(session, user, search_space_id) diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py index e4dc5714a..7a5224ba6 100644 --- a/surfsense_backend/app/routes/new_chat_routes.py +++ b/surfsense_backend/app/routes/new_chat_routes.py @@ -19,13 +19,14 @@ from datetime import UTC, datetime from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile from fastapi.responses import StreamingResponse -from sqlalchemy import or_ +from sqlalchemy import func, or_ from sqlalchemy.exc import IntegrityError, OperationalError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from sqlalchemy.orm import selectinload from app.db import ( + ChatComment, ChatVisibility, NewChatMessage, NewChatMessageRole, @@ -508,7 +509,19 @@ async def get_thread_full( # Check thread-level access based on visibility await check_thread_access(session, thread, user) - return thread + # Check if thread has any comments + comment_count = await session.scalar( + select(func.count()) + .select_from(ChatComment) + .join(NewChatMessage, ChatComment.message_id == NewChatMessage.id) + .where(NewChatMessage.thread_id == thread.id) + ) + + return { + **thread.__dict__, + "messages": thread.messages, + "has_comments": (comment_count or 0) > 0, + } except HTTPException: raise diff --git a/surfsense_backend/app/routes/rbac_routes.py b/surfsense_backend/app/routes/rbac_routes.py index e90970b29..84e95f7ca 100644 --- a/surfsense_backend/app/routes/rbac_routes.py +++ b/surfsense_backend/app/routes/rbac_routes.py @@ -452,6 +452,8 @@ async def list_members( "created_at": membership.created_at, "role": membership.role, "user_email": member_user.email if member_user else None, + "user_display_name": member_user.display_name if member_user else None, + "user_avatar_url": member_user.avatar_url if member_user else None, } response.append(membership_dict) diff --git a/surfsense_backend/app/schemas/chat_comments.py b/surfsense_backend/app/schemas/chat_comments.py new file mode 100644 index 000000000..b87ee58a4 --- /dev/null +++ b/surfsense_backend/app/schemas/chat_comments.py @@ -0,0 +1,129 @@ +""" +Pydantic schemas for chat comments and mentions. +""" + +from datetime import datetime +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field + +# ============================================================================= +# Request Schemas +# ============================================================================= + + +class CommentCreateRequest(BaseModel): + """Schema for creating a comment or reply.""" + + content: str = Field(..., min_length=1, max_length=5000) + + +class CommentUpdateRequest(BaseModel): + """Schema for updating a comment.""" + + content: str = Field(..., min_length=1, max_length=5000) + + +# ============================================================================= +# Author Schema +# ============================================================================= + + +class AuthorResponse(BaseModel): + """Author information for comments.""" + + id: UUID + display_name: str | None = None + avatar_url: str | None = None + email: str + + model_config = ConfigDict(from_attributes=True) + + +# ============================================================================= +# Comment Schemas +# ============================================================================= + + +class CommentReplyResponse(BaseModel): + """Schema for a comment reply (no nested replies).""" + + id: int + content: str + content_rendered: str + author: AuthorResponse | None = None + created_at: datetime + updated_at: datetime + is_edited: bool + can_edit: bool = False + can_delete: bool = False + + model_config = ConfigDict(from_attributes=True) + + +class CommentResponse(BaseModel): + """Schema for a top-level comment with replies.""" + + id: int + message_id: int + content: str + content_rendered: str + author: AuthorResponse | None = None + created_at: datetime + updated_at: datetime + is_edited: bool + can_edit: bool = False + can_delete: bool = False + reply_count: int + replies: list[CommentReplyResponse] = [] + + model_config = ConfigDict(from_attributes=True) + + +class CommentListResponse(BaseModel): + """Response for listing comments on a message.""" + + comments: list[CommentResponse] + total_count: int + + +# ============================================================================= +# Mention Schemas +# ============================================================================= + + +class MentionContextResponse(BaseModel): + """Context information for where a mention occurred.""" + + thread_id: int + thread_title: str + message_id: int + search_space_id: int + search_space_name: str + + +class MentionCommentResponse(BaseModel): + """Abbreviated comment info for mention display.""" + + id: int + content_preview: str + author: AuthorResponse | None = None + created_at: datetime + + +class MentionResponse(BaseModel): + """Schema for a mention notification.""" + + id: int + created_at: datetime + comment: MentionCommentResponse + context: MentionContextResponse + + model_config = ConfigDict(from_attributes=True) + + +class MentionListResponse(BaseModel): + """Response for listing user's mentions.""" + + mentions: list[MentionResponse] + total_count: int diff --git a/surfsense_backend/app/schemas/new_chat.py b/surfsense_backend/app/schemas/new_chat.py index 3734b0470..24e779b50 100644 --- a/surfsense_backend/app/schemas/new_chat.py +++ b/surfsense_backend/app/schemas/new_chat.py @@ -105,6 +105,7 @@ class NewChatThreadWithMessages(NewChatThreadRead): """Schema for reading a thread with its messages.""" messages: list[NewChatMessageRead] = [] + has_comments: bool = False # ============================================================================= diff --git a/surfsense_backend/app/schemas/rbac_schemas.py b/surfsense_backend/app/schemas/rbac_schemas.py index 736d40807..a51f3bc28 100644 --- a/surfsense_backend/app/schemas/rbac_schemas.py +++ b/surfsense_backend/app/schemas/rbac_schemas.py @@ -73,8 +73,10 @@ class MembershipRead(BaseModel): created_at: datetime # Nested role info role: RoleRead | None = None - # User email (populated separately) + # User details (populated separately) user_email: str | None = None + user_display_name: str | None = None + user_avatar_url: str | None = None class Config: from_attributes = True diff --git a/surfsense_backend/app/services/chat_comments_service.py b/surfsense_backend/app/services/chat_comments_service.py new file mode 100644 index 000000000..fa26bf6d5 --- /dev/null +++ b/surfsense_backend/app/services/chat_comments_service.py @@ -0,0 +1,733 @@ +""" +Service layer for chat comments and mentions. +""" + +from uuid import UUID + +from fastapi import HTTPException +from sqlalchemy import delete, select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import selectinload + +from app.db import ( + ChatComment, + ChatCommentMention, + NewChatMessage, + NewChatMessageRole, + NewChatThread, + Permission, + SearchSpaceMembership, + User, + has_permission, +) +from app.schemas.chat_comments import ( + AuthorResponse, + CommentListResponse, + CommentReplyResponse, + CommentResponse, + MentionCommentResponse, + MentionContextResponse, + MentionListResponse, + MentionResponse, +) +from app.services.notification_service import NotificationService +from app.utils.chat_comments import parse_mentions, render_mentions +from app.utils.rbac import check_permission, get_user_permissions + + +async def get_user_names_for_mentions( + session: AsyncSession, + user_ids: set[UUID], +) -> dict[UUID, str]: + """ + Fetch display names for a set of user IDs. + + Args: + session: Database session + user_ids: Set of user UUIDs to look up + + Returns: + Dictionary mapping user UUID to display name + """ + if not user_ids: + return {} + + result = await session.execute( + select(User.id, User.display_name).filter(User.id.in_(user_ids)) + ) + return {row.id: row.display_name or "Unknown" for row in result.all()} + + +async def process_mentions( + session: AsyncSession, + comment_id: int, + content: str, + search_space_id: int, +) -> dict[UUID, int]: + """ + Parse mentions from content, validate users are members, and insert mention records. + + Args: + session: Database session + comment_id: ID of the comment containing mentions + content: Comment text with @[uuid] mentions + search_space_id: ID of the search space for membership validation + + Returns: + Dictionary mapping mentioned user UUID to their mention record ID + """ + mentioned_uuids = parse_mentions(content) + if not mentioned_uuids: + return {} + + # Get valid members from the mentioned UUIDs + result = await session.execute( + select(SearchSpaceMembership.user_id).filter( + SearchSpaceMembership.search_space_id == search_space_id, + SearchSpaceMembership.user_id.in_(mentioned_uuids), + ) + ) + valid_member_ids = result.scalars().all() + + # Insert mention records for valid members and collect their IDs + mentions_map: dict[UUID, int] = {} + for user_id in valid_member_ids: + mention = ChatCommentMention( + comment_id=comment_id, + mentioned_user_id=user_id, + ) + session.add(mention) + await session.flush() + mentions_map[user_id] = mention.id + + return mentions_map + + +async def get_comments_for_message( + session: AsyncSession, + message_id: int, + user: User, +) -> CommentListResponse: + """ + Get all comments for a message with their replies. + + Args: + session: Database session + message_id: ID of the message to get comments for + user: The current authenticated user + + Returns: + CommentListResponse with all top-level comments and their replies + + Raises: + HTTPException: If message not found or user lacks COMMENTS_READ permission + """ + result = await session.execute( + select(NewChatMessage) + .options(selectinload(NewChatMessage.thread)) + .filter(NewChatMessage.id == message_id) + ) + message = result.scalars().first() + + if not message: + raise HTTPException(status_code=404, detail="Message not found") + + search_space_id = message.thread.search_space_id + + # Check permission to read comments + await check_permission( + session, + user, + search_space_id, + Permission.COMMENTS_READ.value, + "You don't have permission to read comments in this search space", + ) + + # Get user permissions for can_delete computation + user_permissions = await get_user_permissions(session, user.id, search_space_id) + can_delete_any = has_permission(user_permissions, Permission.COMMENTS_DELETE.value) + + # Get top-level comments (parent_id IS NULL) with their authors and replies + result = await session.execute( + select(ChatComment) + .options( + selectinload(ChatComment.author), + selectinload(ChatComment.replies).selectinload(ChatComment.author), + ) + .filter( + ChatComment.message_id == message_id, + ChatComment.parent_id.is_(None), + ) + .order_by(ChatComment.created_at) + ) + top_level_comments = result.scalars().all() + + # Collect all mentioned UUIDs from comments and replies for rendering + all_mentioned_uuids: set[UUID] = set() + for comment in top_level_comments: + all_mentioned_uuids.update(parse_mentions(comment.content)) + for reply in comment.replies: + all_mentioned_uuids.update(parse_mentions(reply.content)) + + # Fetch display names for mentioned users + user_names = await get_user_names_for_mentions(session, all_mentioned_uuids) + + comments = [] + for comment in top_level_comments: + author = None + if comment.author: + author = AuthorResponse( + id=comment.author.id, + display_name=comment.author.display_name, + avatar_url=comment.author.avatar_url, + email=comment.author.email, + ) + + replies = [] + for reply in sorted(comment.replies, key=lambda r: r.created_at): + reply_author = None + if reply.author: + reply_author = AuthorResponse( + id=reply.author.id, + display_name=reply.author.display_name, + avatar_url=reply.author.avatar_url, + email=reply.author.email, + ) + + is_reply_author = reply.author_id == user.id if reply.author_id else False + replies.append( + CommentReplyResponse( + id=reply.id, + content=reply.content, + content_rendered=render_mentions(reply.content, user_names), + author=reply_author, + created_at=reply.created_at, + updated_at=reply.updated_at, + is_edited=reply.updated_at > reply.created_at, + can_edit=is_reply_author, + can_delete=is_reply_author or can_delete_any, + ) + ) + + is_comment_author = comment.author_id == user.id if comment.author_id else False + comments.append( + CommentResponse( + id=comment.id, + message_id=comment.message_id, + content=comment.content, + content_rendered=render_mentions(comment.content, user_names), + author=author, + created_at=comment.created_at, + updated_at=comment.updated_at, + is_edited=comment.updated_at > comment.created_at, + can_edit=is_comment_author, + can_delete=is_comment_author or can_delete_any, + reply_count=len(replies), + replies=replies, + ) + ) + + return CommentListResponse( + comments=comments, + total_count=len(comments), + ) + + +async def create_comment( + session: AsyncSession, + message_id: int, + content: str, + user: User, +) -> CommentResponse: + """ + Create a top-level comment on an AI response. + + Args: + session: Database session + message_id: ID of the message to comment on + content: Comment text content + user: The current authenticated user + + Returns: + CommentResponse for the created comment + + Raises: + HTTPException: If message not found, not AI response, or user lacks COMMENTS_CREATE permission + """ + result = await session.execute( + select(NewChatMessage) + .options(selectinload(NewChatMessage.thread)) + .filter(NewChatMessage.id == message_id) + ) + message = result.scalars().first() + + if not message: + raise HTTPException(status_code=404, detail="Message not found") + + # Validate message is an AI response + if message.role != NewChatMessageRole.ASSISTANT: + raise HTTPException( + status_code=400, + detail="Comments can only be added to AI responses", + ) + + search_space_id = message.thread.search_space_id + + # Check permission to create comments + user_permissions = await get_user_permissions(session, user.id, search_space_id) + if not has_permission(user_permissions, Permission.COMMENTS_CREATE.value): + raise HTTPException( + status_code=403, + detail="You don't have permission to create comments in this search space", + ) + + comment = ChatComment( + message_id=message_id, + author_id=user.id, + content=content, + ) + session.add(comment) + await session.flush() + + # Process mentions - returns map of user_id -> mention_id + mentions_map = await process_mentions(session, comment.id, content, search_space_id) + + await session.commit() + await session.refresh(comment) + + # Fetch user names for rendering mentions (reuse mentions_map keys) + user_names = await get_user_names_for_mentions(session, set(mentions_map.keys())) + + # Create notifications for mentioned users (excluding author) + thread = message.thread + author_name = user.display_name or user.email + content_preview = render_mentions(content, user_names) + for mentioned_user_id, mention_id in mentions_map.items(): + if mentioned_user_id == user.id: + continue # Don't notify yourself + await NotificationService.mention.notify_new_mention( + session=session, + mentioned_user_id=mentioned_user_id, + mention_id=mention_id, + comment_id=comment.id, + message_id=message_id, + thread_id=thread.id, + thread_title=thread.title or "Untitled thread", + author_id=str(user.id), + author_name=author_name, + content_preview=content_preview[:200], + search_space_id=search_space_id, + ) + + author = AuthorResponse( + id=user.id, + display_name=user.display_name, + avatar_url=user.avatar_url, + email=user.email, + ) + + return CommentResponse( + id=comment.id, + message_id=comment.message_id, + content=comment.content, + content_rendered=render_mentions(content, user_names), + author=author, + created_at=comment.created_at, + updated_at=comment.updated_at, + is_edited=False, + can_edit=True, + can_delete=True, + reply_count=0, + replies=[], + ) + + +async def create_reply( + session: AsyncSession, + comment_id: int, + content: str, + user: User, +) -> CommentReplyResponse: + """ + Create a reply to an existing comment. + + Args: + session: Database session + comment_id: ID of the parent comment to reply to + content: Reply text content + user: The current authenticated user + + Returns: + CommentReplyResponse for the created reply + + Raises: + HTTPException: If comment not found, is already a reply, or user lacks COMMENTS_CREATE permission + """ + # Get parent comment with its message and thread + result = await session.execute( + select(ChatComment) + .options(selectinload(ChatComment.message).selectinload(NewChatMessage.thread)) + .filter(ChatComment.id == comment_id) + ) + parent_comment = result.scalars().first() + + if not parent_comment: + raise HTTPException(status_code=404, detail="Comment not found") + + # Validate parent is a top-level comment (cannot reply to a reply) + if parent_comment.parent_id is not None: + raise HTTPException( + status_code=400, + detail="Cannot reply to a reply", + ) + + search_space_id = parent_comment.message.thread.search_space_id + + # Check permission to create comments + user_permissions = await get_user_permissions(session, user.id, search_space_id) + if not has_permission(user_permissions, Permission.COMMENTS_CREATE.value): + raise HTTPException( + status_code=403, + detail="You don't have permission to create comments in this search space", + ) + + reply = ChatComment( + message_id=parent_comment.message_id, + parent_id=comment_id, + author_id=user.id, + content=content, + ) + session.add(reply) + await session.flush() + + # Process mentions - returns map of user_id -> mention_id + mentions_map = await process_mentions(session, reply.id, content, search_space_id) + + await session.commit() + await session.refresh(reply) + + # Fetch user names for rendering mentions (reuse mentions_map keys) + user_names = await get_user_names_for_mentions(session, set(mentions_map.keys())) + + # Create notifications for mentioned users (excluding author) + thread = parent_comment.message.thread + author_name = user.display_name or user.email + content_preview = render_mentions(content, user_names) + for mentioned_user_id, mention_id in mentions_map.items(): + if mentioned_user_id == user.id: + continue # Don't notify yourself + await NotificationService.mention.notify_new_mention( + session=session, + mentioned_user_id=mentioned_user_id, + mention_id=mention_id, + comment_id=reply.id, + message_id=parent_comment.message_id, + thread_id=thread.id, + thread_title=thread.title or "Untitled thread", + author_id=str(user.id), + author_name=author_name, + content_preview=content_preview[:200], + search_space_id=search_space_id, + ) + + author = AuthorResponse( + id=user.id, + display_name=user.display_name, + avatar_url=user.avatar_url, + email=user.email, + ) + + return CommentReplyResponse( + id=reply.id, + content=reply.content, + content_rendered=render_mentions(content, user_names), + author=author, + created_at=reply.created_at, + updated_at=reply.updated_at, + is_edited=False, + can_edit=True, + can_delete=True, + ) + + +async def update_comment( + session: AsyncSession, + comment_id: int, + content: str, + user: User, +) -> CommentReplyResponse: + """ + Update a comment's content (author only). + + Args: + session: Database session + comment_id: ID of the comment to update + content: New comment text content + user: The current authenticated user + + Returns: + CommentReplyResponse for the updated comment + + Raises: + HTTPException: If comment not found or user is not the author + """ + result = await session.execute( + select(ChatComment) + .options( + selectinload(ChatComment.author), + selectinload(ChatComment.message).selectinload(NewChatMessage.thread), + ) + .filter(ChatComment.id == comment_id) + ) + comment = result.scalars().first() + + if not comment: + raise HTTPException(status_code=404, detail="Comment not found") + + if comment.author_id != user.id: + raise HTTPException( + status_code=403, + detail="You can only edit your own comments", + ) + + search_space_id = comment.message.thread.search_space_id + + # Get existing mentioned user IDs + existing_result = await session.execute( + select(ChatCommentMention.mentioned_user_id).filter( + ChatCommentMention.comment_id == comment_id + ) + ) + existing_mention_ids = set(existing_result.scalars().all()) + + # Parse new mentions from updated content + new_mention_uuids = set(parse_mentions(content)) + + # Validate new mentions are search space members + if new_mention_uuids: + valid_result = await session.execute( + select(SearchSpaceMembership.user_id).filter( + SearchSpaceMembership.search_space_id == search_space_id, + SearchSpaceMembership.user_id.in_(new_mention_uuids), + ) + ) + valid_new_mentions = set(valid_result.scalars().all()) + else: + valid_new_mentions = set() + + # Compute diff: removed, kept (preserve read status), added + mentions_to_remove = existing_mention_ids - valid_new_mentions + mentions_to_add = valid_new_mentions - existing_mention_ids + + # Delete removed mentions + if mentions_to_remove: + await session.execute( + delete(ChatCommentMention).where( + ChatCommentMention.comment_id == comment_id, + ChatCommentMention.mentioned_user_id.in_(mentions_to_remove), + ) + ) + + # Add new mentions and collect their IDs for notifications + new_mentions_map: dict[UUID, int] = {} + for user_id in mentions_to_add: + mention = ChatCommentMention( + comment_id=comment_id, + mentioned_user_id=user_id, + ) + session.add(mention) + await session.flush() + new_mentions_map[user_id] = mention.id + + comment.content = content + + await session.commit() + await session.refresh(comment) + + # Fetch user names for rendering mentions + user_names = await get_user_names_for_mentions(session, valid_new_mentions) + + # Create notifications for newly added mentions (excluding author) + if new_mentions_map: + thread = comment.message.thread + author_name = user.display_name or user.email + content_preview = render_mentions(content, user_names) + for mentioned_user_id, mention_id in new_mentions_map.items(): + if mentioned_user_id == user.id: + continue # Don't notify yourself + await NotificationService.mention.notify_new_mention( + session=session, + mentioned_user_id=mentioned_user_id, + mention_id=mention_id, + comment_id=comment_id, + message_id=comment.message_id, + thread_id=thread.id, + thread_title=thread.title or "Untitled thread", + author_id=str(user.id), + author_name=author_name, + content_preview=content_preview[:200], + search_space_id=search_space_id, + ) + + author = AuthorResponse( + id=user.id, + display_name=user.display_name, + avatar_url=user.avatar_url, + email=user.email, + ) + + return CommentReplyResponse( + id=comment.id, + content=comment.content, + content_rendered=render_mentions(content, user_names), + author=author, + created_at=comment.created_at, + updated_at=comment.updated_at, + is_edited=comment.updated_at > comment.created_at, + can_edit=True, + can_delete=True, + ) + + +async def delete_comment( + session: AsyncSession, + comment_id: int, + user: User, +) -> dict: + """ + Delete a comment (author or user with COMMENTS_DELETE permission). + + Args: + session: Database session + comment_id: ID of the comment to delete + user: The current authenticated user + + Returns: + Dict with deletion confirmation + + Raises: + HTTPException: If comment not found or user lacks permission to delete + """ + result = await session.execute( + select(ChatComment) + .options(selectinload(ChatComment.message).selectinload(NewChatMessage.thread)) + .filter(ChatComment.id == comment_id) + ) + comment = result.scalars().first() + + if not comment: + raise HTTPException(status_code=404, detail="Comment not found") + + is_author = comment.author_id == user.id + + # Check if user has COMMENTS_DELETE permission + search_space_id = comment.message.thread.search_space_id + user_permissions = await get_user_permissions(session, user.id, search_space_id) + can_delete_any = has_permission(user_permissions, Permission.COMMENTS_DELETE.value) + + if not is_author and not can_delete_any: + raise HTTPException( + status_code=403, + detail="You do not have permission to delete this comment", + ) + + await session.delete(comment) + await session.commit() + + return {"message": "Comment deleted successfully", "comment_id": comment_id} + + +async def get_user_mentions( + session: AsyncSession, + user: User, + search_space_id: int | None = None, +) -> MentionListResponse: + """ + Get mentions for the current user, optionally filtered by search space. + + Args: + session: Database session + user: The current authenticated user + search_space_id: Optional search space ID to filter mentions + + Returns: + MentionListResponse with mentions and total count + """ + # Build query with joins for filtering by search_space_id + query = ( + select(ChatCommentMention) + .join(ChatComment, ChatCommentMention.comment_id == ChatComment.id) + .join(NewChatMessage, ChatComment.message_id == NewChatMessage.id) + .join(NewChatThread, NewChatMessage.thread_id == NewChatThread.id) + .options( + selectinload(ChatCommentMention.comment).selectinload(ChatComment.author), + selectinload(ChatCommentMention.comment).selectinload(ChatComment.message), + ) + .filter(ChatCommentMention.mentioned_user_id == user.id) + .order_by(ChatCommentMention.created_at.desc()) + ) + + if search_space_id is not None: + query = query.filter(NewChatThread.search_space_id == search_space_id) + + result = await session.execute(query) + mention_records = result.scalars().all() + + # Fetch search space info for context (single query for all unique search spaces) + thread_ids = {m.comment.message.thread_id for m in mention_records} + if thread_ids: + thread_result = await session.execute( + select(NewChatThread) + .options(selectinload(NewChatThread.search_space)) + .filter(NewChatThread.id.in_(thread_ids)) + ) + threads_map = {t.id: t for t in thread_result.scalars().all()} + else: + threads_map = {} + + mentions = [] + for mention in mention_records: + comment = mention.comment + message = comment.message + thread = threads_map.get(message.thread_id) + search_space = thread.search_space if thread else None + + author = None + if comment.author: + author = AuthorResponse( + id=comment.author.id, + display_name=comment.author.display_name, + avatar_url=comment.author.avatar_url, + email=comment.author.email, + ) + + content_preview = ( + comment.content[:100] + "..." + if len(comment.content) > 100 + else comment.content + ) + + mentions.append( + MentionResponse( + id=mention.id, + created_at=mention.created_at, + comment=MentionCommentResponse( + id=comment.id, + content_preview=content_preview, + author=author, + created_at=comment.created_at, + ), + context=MentionContextResponse( + thread_id=thread.id if thread else 0, + thread_title=thread.title or "Untitled" if thread else "Unknown", + message_id=message.id, + search_space_id=search_space.id if search_space else 0, + search_space_name=search_space.name if search_space else "Unknown", + ), + ) + ) + + return MentionListResponse( + mentions=mentions, + total_count=len(mentions), + ) diff --git a/surfsense_backend/app/services/notification_service.py b/surfsense_backend/app/services/notification_service.py index 5e8d2aa8b..97e0f9457 100644 --- a/surfsense_backend/app/services/notification_service.py +++ b/surfsense_backend/app/services/notification_service.py @@ -617,12 +617,83 @@ class DocumentProcessingNotificationHandler(BaseNotificationHandler): ) +class MentionNotificationHandler(BaseNotificationHandler): + """Handler for new mention notifications.""" + + def __init__(self): + super().__init__("new_mention") + + async def notify_new_mention( + self, + session: AsyncSession, + mentioned_user_id: UUID, + mention_id: int, + comment_id: int, + message_id: int, + thread_id: int, + thread_title: str, + author_id: str, + author_name: str, + content_preview: str, + search_space_id: int, + ) -> Notification: + """ + Create notification when a user is @mentioned in a comment. + + Args: + session: Database session + mentioned_user_id: User who was mentioned + mention_id: ID of the mention record + comment_id: ID of the comment containing the mention + message_id: ID of the message being commented on + thread_id: ID of the chat thread + thread_title: Title of the chat thread + author_id: ID of the comment author + author_name: Display name of the comment author + content_preview: First ~100 chars of the comment + search_space_id: Search space ID + + Returns: + Notification: The created notification + """ + title = f"{author_name} mentioned you" + message = content_preview[:100] + ("..." if len(content_preview) > 100 else "") + + metadata = { + "mention_id": mention_id, + "comment_id": comment_id, + "message_id": message_id, + "thread_id": thread_id, + "thread_title": thread_title, + "author_id": author_id, + "author_name": author_name, + "content_preview": content_preview[:200], + } + + notification = Notification( + user_id=mentioned_user_id, + search_space_id=search_space_id, + type=self.notification_type, + title=title, + message=message, + notification_metadata=metadata, + ) + session.add(notification) + await session.commit() + await session.refresh(notification) + logger.info( + f"Created new_mention notification {notification.id} for user {mentioned_user_id}" + ) + return notification + + class NotificationService: """Service for creating and managing notifications that sync via Electric SQL.""" # Handler instances connector_indexing = ConnectorIndexingNotificationHandler() document_processing = DocumentProcessingNotificationHandler() + mention = MentionNotificationHandler() @staticmethod async def create_notification( diff --git a/surfsense_backend/app/utils/chat_comments.py b/surfsense_backend/app/utils/chat_comments.py new file mode 100644 index 000000000..5349466f9 --- /dev/null +++ b/surfsense_backend/app/utils/chat_comments.py @@ -0,0 +1,64 @@ +""" +Utility functions for chat comments, including mention parsing. +""" + +import re +from uuid import UUID + +# Pattern to match @[uuid] mentions in comment content +MENTION_PATTERN = re.compile(r"@\[([0-9a-fA-F-]{36})\]") + + +def parse_mentions(content: str) -> list[UUID]: + """ + Extract user UUIDs from @[uuid] mentions in content. + + Args: + content: Comment text that may contain @[uuid] mentions + + Returns: + List of unique user UUIDs found in the content + """ + matches = MENTION_PATTERN.findall(content) + unique_uuids = [] + seen = set() + + for match in matches: + try: + uuid = UUID(match) + if uuid not in seen: + seen.add(uuid) + unique_uuids.append(uuid) + except ValueError: + # Invalid UUID format, skip + continue + + return unique_uuids + + +def render_mentions(content: str, user_names: dict[UUID, str]) -> str: + """ + Replace @[uuid] mentions with @{DisplayName} in content. + + Uses curly braces as delimiters for unambiguous frontend parsing. + + Args: + content: Comment text with @[uuid] mentions + user_names: Dict mapping user UUIDs to display names + + Returns: + Content with mentions rendered as @{DisplayName} + """ + + def replace_mention(match: re.Match) -> str: + try: + uuid = UUID(match.group(1)) + name = user_names.get(uuid) + if name: + return f"@{{{name}}}" + # Keep original format if user not found + return match.group(0) + except ValueError: + return match.group(0) + + return MENTION_PATTERN.sub(replace_mention, content) diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 827646dd2..43c33ba5a 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -8,10 +8,11 @@ import { } from "@assistant-ui/react"; import { useQueryClient } from "@tanstack/react-query"; import { useAtomValue, useSetAtom } from "jotai"; -import { useParams } from "next/navigation"; +import { useParams, useSearchParams } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; import { z } from "zod"; +import { currentThreadAtom } from "@/atoms/chat/current-thread.atom"; import { type MentionedDocumentInfo, mentionedDocumentIdsAtom, @@ -251,6 +252,7 @@ export default function NewChatPage() { const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom); const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom); const hydratePlanState = useSetAtom(hydratePlanStateAtom); + const setCurrentThreadState = useSetAtom(currentThreadAtom); // Get current user for author info in shared chats const { data: currentUser } = useAtomValue(currentUserAtom); @@ -365,6 +367,48 @@ export default function NewChatPage() { initializeThread(); }, [initializeThread]); + // Handle scroll to comment from URL query params (e.g., from notification click) + const searchParams = useSearchParams(); + const targetCommentId = searchParams.get("commentId"); + + useEffect(() => { + if (!targetCommentId || isInitializing || messages.length === 0) return; + + const tryScroll = () => { + const el = document.querySelector(`[data-comment-id="${targetCommentId}"]`); + if (el) { + el.scrollIntoView({ behavior: "smooth", block: "center" }); + return true; + } + return false; + }; + + // Try immediately + if (tryScroll()) return; + + // Retry every 200ms for up to 10 seconds + const intervalId = setInterval(() => { + if (tryScroll()) clearInterval(intervalId); + }, 200); + + const timeoutId = setTimeout(() => clearInterval(intervalId), 10000); + + return () => { + clearInterval(intervalId); + clearTimeout(timeoutId); + }; + }, [targetCommentId, isInitializing, messages.length]); + + // Sync current thread state to atom + useEffect(() => { + setCurrentThreadState({ + id: currentThread?.id ?? null, + visibility: currentThread?.visibility ?? null, + hasComments: currentThread?.has_comments ?? false, + addingCommentToMessageId: null, + }); + }, [currentThread, setCurrentThreadState]); + // Cancel ongoing request const cancelRun = useCallback(async () => { if (abortControllerRef.current) { @@ -842,10 +886,32 @@ export default function NewChatPage() { // Persist assistant message (with thinking steps for restoration on refresh) const finalContent = buildContentForPersistence(); if (contentParts.length > 0) { - appendMessage(currentThreadId, { - role: "assistant", - content: finalContent, - }).catch((err) => console.error("Failed to persist assistant message:", err)); + try { + const savedMessage = await appendMessage(currentThreadId, { + role: "assistant", + content: finalContent, + }); + + // Update message ID from temporary to database ID so comments work immediately + const newMsgId = `msg-${savedMessage.id}`; + setMessages((prev) => + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + ); + + // Also update thinking steps map with new ID + setMessageThinkingSteps((prev) => { + const steps = prev.get(assistantMsgId); + if (steps) { + const newMap = new Map(prev); + newMap.delete(assistantMsgId); + newMap.set(newMsgId, steps); + return newMap; + } + return prev; + }); + } catch (err) { + console.error("Failed to persist assistant message:", err); + } // Track successful response trackChatResponseReceived(searchSpaceId, currentThreadId); @@ -860,10 +926,20 @@ export default function NewChatPage() { ); if (hasContent && currentThreadId) { const partialContent = buildContentForPersistence(); - appendMessage(currentThreadId, { - role: "assistant", - content: partialContent, - }).catch((err) => console.error("Failed to persist partial assistant message:", err)); + try { + const savedMessage = await appendMessage(currentThreadId, { + role: "assistant", + content: partialContent, + }); + + // Update message ID from temporary to database ID + const newMsgId = `msg-${savedMessage.id}`; + setMessages((prev) => + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + ); + } catch (err) { + console.error("Failed to persist partial assistant message:", err); + } } return; } diff --git a/surfsense_web/app/globals.css b/surfsense_web/app/globals.css index 4f4ab6de1..7324ffeb3 100644 --- a/surfsense_web/app/globals.css +++ b/surfsense_web/app/globals.css @@ -157,5 +157,33 @@ button { cursor: pointer; } +/* Custom scrollbar styles */ +.scrollbar-thin { + scrollbar-width: thin; + scrollbar-color: hsl(var(--muted-foreground) / 0.2) transparent; +} + +.scrollbar-thin:hover { + scrollbar-color: hsl(var(--muted-foreground) / 0.4) transparent; +} + +/* Webkit scrollbar styles */ +.scrollbar-thin::-webkit-scrollbar { + width: 6px; +} + +.scrollbar-thin::-webkit-scrollbar-track { + background: transparent; +} + +.scrollbar-thin::-webkit-scrollbar-thumb { + background-color: hsl(var(--muted-foreground) / 0.2); + border-radius: 3px; +} + +.scrollbar-thin::-webkit-scrollbar-thumb:hover { + background-color: hsl(var(--muted-foreground) / 0.4); +} + @source '../node_modules/@llamaindex/chat-ui/**/*.{ts,tsx}'; @source '../node_modules/streamdown/dist/*.js'; diff --git a/surfsense_web/atoms/chat-comments/comments-mutation.atoms.ts b/surfsense_web/atoms/chat-comments/comments-mutation.atoms.ts new file mode 100644 index 000000000..e6a9767ca --- /dev/null +++ b/surfsense_web/atoms/chat-comments/comments-mutation.atoms.ts @@ -0,0 +1,72 @@ +import { atomWithMutation } from "jotai-tanstack-query"; +import { toast } from "sonner"; +import type { + CreateCommentRequest, + CreateReplyRequest, + DeleteCommentRequest, + UpdateCommentRequest, +} from "@/contracts/types/chat-comments.types"; +import { chatCommentsApiService } from "@/lib/apis/chat-comments-api.service"; +import { cacheKeys } from "@/lib/query-client/cache-keys"; +import { queryClient } from "@/lib/query-client/client"; + +export const createCommentMutationAtom = atomWithMutation(() => ({ + mutationFn: async (request: CreateCommentRequest) => { + return chatCommentsApiService.createComment(request); + }, + onSuccess: (_, variables) => { + queryClient.invalidateQueries({ + queryKey: cacheKeys.comments.byMessage(variables.message_id), + }); + }, + onError: (error: Error) => { + console.error("Error creating comment:", error); + toast.error("Failed to create comment"); + }, +})); + +export const createReplyMutationAtom = atomWithMutation(() => ({ + mutationFn: async (request: CreateReplyRequest & { message_id: number }) => { + return chatCommentsApiService.createReply(request); + }, + onSuccess: (_, variables) => { + queryClient.invalidateQueries({ + queryKey: cacheKeys.comments.byMessage(variables.message_id), + }); + }, + onError: (error: Error) => { + console.error("Error creating reply:", error); + toast.error("Failed to create reply"); + }, +})); + +export const updateCommentMutationAtom = atomWithMutation(() => ({ + mutationFn: async (request: UpdateCommentRequest & { message_id: number }) => { + return chatCommentsApiService.updateComment(request); + }, + onSuccess: (_, variables) => { + queryClient.invalidateQueries({ + queryKey: cacheKeys.comments.byMessage(variables.message_id), + }); + }, + onError: (error: Error) => { + console.error("Error updating comment:", error); + toast.error("Failed to update comment"); + }, +})); + +export const deleteCommentMutationAtom = atomWithMutation(() => ({ + mutationFn: async (request: DeleteCommentRequest & { message_id: number }) => { + return chatCommentsApiService.deleteComment(request); + }, + onSuccess: (_, variables) => { + queryClient.invalidateQueries({ + queryKey: cacheKeys.comments.byMessage(variables.message_id), + }); + toast.success("Comment deleted"); + }, + onError: (error: Error) => { + console.error("Error deleting comment:", error); + toast.error("Failed to delete comment"); + }, +})); diff --git a/surfsense_web/atoms/chat/current-thread.atom.ts b/surfsense_web/atoms/chat/current-thread.atom.ts new file mode 100644 index 000000000..1231887f8 --- /dev/null +++ b/surfsense_web/atoms/chat/current-thread.atom.ts @@ -0,0 +1,52 @@ +import { atom } from "jotai"; +import type { ChatVisibility } from "@/lib/chat/thread-persistence"; + +// TODO: Update `hasComments` to true when the first comment is created on a thread. +// Currently it only updates on thread load. The gutter still works because +// `addingCommentToMessageId` keeps it open, but the state is technically stale. + +// TODO: Reset `addingCommentToMessageId` to null after a comment is successfully created. +// Currently it stays set until navigation or clicking another message's bubble. +// Not causing issues since panel visibility is driven by per-message comment count. + +// TODO: Consider calling `resetCurrentThreadAtom` when unmounting the chat page +// for explicit cleanup, though React navigation handles this implicitly. + +interface CurrentThreadState { + id: number | null; + visibility: ChatVisibility | null; + hasComments: boolean; + addingCommentToMessageId: number | null; +} + +const initialState: CurrentThreadState = { + id: null, + visibility: null, + hasComments: false, + addingCommentToMessageId: null, +}; + +export const currentThreadAtom = atom(initialState); + +export const commentsEnabledAtom = atom( + (get) => get(currentThreadAtom).visibility === "SEARCH_SPACE" +); + +export const showCommentsGutterAtom = atom((get) => { + const thread = get(currentThreadAtom); + return ( + thread.visibility === "SEARCH_SPACE" && + (thread.hasComments || thread.addingCommentToMessageId !== null) + ); +}); + +export const addingCommentToMessageIdAtom = atom( + (get) => get(currentThreadAtom).addingCommentToMessageId, + (get, set, messageId: number | null) => { + set(currentThreadAtom, { ...get(currentThreadAtom), addingCommentToMessageId: messageId }); + } +); + +export const resetCurrentThreadAtom = atom(null, (_, set) => { + set(currentThreadAtom, initialState); +}); diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index 83c573e2f..b2ca9d729 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -5,9 +5,15 @@ import { MessagePrimitive, useAssistantState, } from "@assistant-ui/react"; +import { useAtom, useAtomValue } from "jotai"; import { CheckIcon, CopyIcon, DownloadIcon, RefreshCwIcon } from "lucide-react"; import type { FC } from "react"; -import { useContext } from "react"; +import { useContext, useEffect, useRef, useState } from "react"; +import { + addingCommentToMessageIdAtom, + commentsEnabledAtom, +} from "@/atoms/chat/current-thread.atom"; +import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { BranchPicker } from "@/components/assistant-ui/branch-picker"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { @@ -16,6 +22,9 @@ import { } from "@/components/assistant-ui/thinking-steps"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container"; +import { CommentTrigger } from "@/components/chat-comments/comment-trigger/comment-trigger"; +import { useComments } from "@/hooks/use-comments"; export const MessageError: FC = () => { return ( @@ -76,13 +85,89 @@ const AssistantMessageInner: FC = () => { ); }; +function parseMessageId(assistantUiMessageId: string | undefined): number | null { + if (!assistantUiMessageId) return null; + const match = assistantUiMessageId.match(/^msg-(\d+)$/); + return match ? Number.parseInt(match[1], 10) : null; +} + export const AssistantMessage: FC = () => { + const [messageHeight, setMessageHeight] = useState(undefined); + const messageRef = useRef(null); + const messageId = useAssistantState(({ message }) => message?.id); + const searchSpaceId = useAtomValue(activeSearchSpaceIdAtom); + const dbMessageId = parseMessageId(messageId); + const commentsEnabled = useAtomValue(commentsEnabledAtom); + const [addingCommentToMessageId, setAddingCommentToMessageId] = useAtom( + addingCommentToMessageIdAtom + ); + + const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); + const isLastMessage = useAssistantState(({ message }) => message?.isLast ?? false); + const isMessageStreaming = isThreadRunning && isLastMessage; + + const { data: commentsData } = useComments({ + messageId: dbMessageId ?? 0, + enabled: !!dbMessageId, + }); + + const hasComments = (commentsData?.total_count ?? 0) > 0; + const isAddingComment = dbMessageId !== null && addingCommentToMessageId === dbMessageId; + const showCommentPanel = hasComments || isAddingComment; + + const handleToggleAddComment = () => { + if (!dbMessageId) return; + setAddingCommentToMessageId(isAddingComment ? null : dbMessageId); + }; + + useEffect(() => { + if (!messageRef.current) return; + const el = messageRef.current; + const update = () => setMessageHeight(el.offsetHeight); + update(); + const observer = new ResizeObserver(update); + observer.observe(el); + return () => observer.disconnect(); + }, []); + return ( + + {searchSpaceId && commentsEnabled && !isMessageStreaming && ( +
+
+ {!hasComments && ( + + )} + + {showCommentPanel && dbMessageId && ( +
+ +
+ )} +
+
+ )} ); }; diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 18f45f810..02731ae6e 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -26,6 +26,7 @@ import { import { useParams } from "next/navigation"; import { type FC, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react"; import { createPortal } from "react-dom"; +import { showCommentsGutterAtom } from "@/atoms/chat/current-thread.atom"; import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, @@ -36,6 +37,7 @@ import { newLLMConfigsAtom, } from "@/atoms/new-llm-config/new-llm-config-query.atoms"; import { currentUserAtom } from "@/atoms/user/user-query.atoms"; +import { AssistantMessage } from "@/components/assistant-ui/assistant-message"; import { ComposerAddAttachment, ComposerAttachments } from "@/components/assistant-ui/attachment"; import { ConnectorIndicator } from "@/components/assistant-ui/connector-popup"; import { @@ -59,57 +61,63 @@ import { Button } from "@/components/ui/button"; import type { Document } from "@/contracts/types/document.types"; import { cn } from "@/lib/utils"; -/** - * Props for the Thread component - */ interface ThreadProps { messageThinkingSteps?: Map; - /** Optional header component to render at the top of the viewport (sticky) */ header?: React.ReactNode; } export const Thread: FC = ({ messageThinkingSteps = new Map(), header }) => { return ( - - - {/* Optional sticky header for model selector etc. */} - {header &&
{header}
} - - thread.isEmpty}> - - - - - - - - !thread.isEmpty}> -
- -
- - - - + ); }; +const ThreadContent: FC<{ header?: React.ReactNode }> = ({ header }) => { + const showGutter = useAtomValue(showCommentsGutterAtom); + + return ( + + + {header &&
{header}
} + + thread.isEmpty}> + + + + + + + + !thread.isEmpty}> +
+ +
+ + + + + ); +}; + const ThreadScrollToBottom: FC = () => { return ( @@ -579,17 +587,6 @@ const AssistantMessageInner: FC = () => { ); }; -const AssistantMessage: FC = () => { - return ( - - - - ); -}; - const AssistantActionBar: FC = () => { return ( b.displayName.length - a.displayName.length); + + for (const mention of sortedMentions) { + const displayPattern = new RegExp( + `@${escapeRegExp(mention.displayName)}(?=\\s|$|[.,!?;:])`, + "g" + ); + const dataFormat = `@[${mention.id}]`; + result = result.replace(displayPattern, dataFormat); + } + + return result; +} + +function escapeRegExp(string: string): string { + return string.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); +} + +function findMentionTrigger( + text: string, + cursorPos: number, + insertedMentions: InsertedMention[] +): { isActive: boolean; query: string; startIndex: number } { + const textBeforeCursor = text.slice(0, cursorPos); + + const mentionMatch = textBeforeCursor.match(/(?:^|[\s])@([^\s]*)$/); + + if (!mentionMatch) { + return { isActive: false, query: "", startIndex: 0 }; + } + + const fullMatch = mentionMatch[0]; + const query = mentionMatch[1]; + const atIndex = cursorPos - query.length - 1; + + if (atIndex > 0) { + const charBefore = text[atIndex - 1]; + if (charBefore && !/[\s]/.test(charBefore)) { + return { isActive: false, query: "", startIndex: 0 }; + } + } + + const textFromAt = text.slice(atIndex); + + for (const mention of insertedMentions) { + const mentionPattern = `@${mention.displayName}`; + + if (textFromAt.startsWith(mentionPattern)) { + const charAfterMention = text[atIndex + mentionPattern.length]; + if (!charAfterMention || /[\s.,!?;:]/.test(charAfterMention)) { + if (cursorPos <= atIndex + mentionPattern.length) { + return { isActive: false, query: "", startIndex: 0 }; + } + } + } + } + + if (query.length > 50) { + return { isActive: false, query: "", startIndex: 0 }; + } + + return { isActive: true, query, startIndex: atIndex }; +} + +export function CommentComposer({ + members, + membersLoading = false, + placeholder = "Write a comment...", + submitLabel = "Send", + isSubmitting = false, + onSubmit, + onCancel, + autoFocus = false, + initialValue = "", +}: CommentComposerProps) { + const [displayContent, setDisplayContent] = useState(initialValue); + const [insertedMentions, setInsertedMentions] = useState([]); + const [mentionsInitialized, setMentionsInitialized] = useState(false); + const [mentionState, setMentionState] = useState({ + isActive: false, + query: "", + startIndex: 0, + }); + const [highlightedIndex, setHighlightedIndex] = useState(0); + const textareaRef = useRef(null); + + const filteredMembers = mentionState.query + ? members.filter( + (member) => + member.displayName?.toLowerCase().includes(mentionState.query.toLowerCase()) || + member.email.toLowerCase().includes(mentionState.query.toLowerCase()) + ) + : members; + + const closeMentionPicker = useCallback(() => { + setMentionState({ isActive: false, query: "", startIndex: 0 }); + setHighlightedIndex(0); + }, []); + + const insertMention = useCallback( + (member: MemberOption) => { + const displayName = member.displayName || member.email.split("@")[0]; + const before = displayContent.slice(0, mentionState.startIndex); + const cursorPos = textareaRef.current?.selectionStart ?? displayContent.length; + const after = displayContent.slice(cursorPos); + const mentionText = `@${displayName} `; + const newContent = before + mentionText + after; + + setDisplayContent(newContent); + setInsertedMentions((prev) => { + const exists = prev.some((m) => m.id === member.id && m.displayName === displayName); + if (exists) return prev; + return [...prev, { id: member.id, displayName }]; + }); + closeMentionPicker(); + + requestAnimationFrame(() => { + if (textareaRef.current) { + const cursorPos = before.length + mentionText.length; + textareaRef.current.focus(); + textareaRef.current.setSelectionRange(cursorPos, cursorPos); + } + }); + }, + [displayContent, mentionState.startIndex, closeMentionPicker] + ); + + const handleInputChange = (e: React.ChangeEvent) => { + const value = e.target.value; + const cursorPos = e.target.selectionStart; + setDisplayContent(value); + + const triggerResult = findMentionTrigger(value, cursorPos, insertedMentions); + + if (triggerResult.isActive) { + setMentionState(triggerResult); + setHighlightedIndex(0); + } else if (mentionState.isActive) { + closeMentionPicker(); + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (!mentionState.isActive) { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + handleSubmit(); + } + return; + } + + switch (e.key) { + case "ArrowDown": + case "Tab": + if (!e.shiftKey) { + e.preventDefault(); + setHighlightedIndex((prev) => (prev < filteredMembers.length - 1 ? prev + 1 : 0)); + } else if (e.key === "Tab") { + e.preventDefault(); + setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : filteredMembers.length - 1)); + } + break; + case "ArrowUp": + e.preventDefault(); + setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : filteredMembers.length - 1)); + break; + case "Enter": + e.preventDefault(); + if (filteredMembers[highlightedIndex]) { + insertMention(filteredMembers[highlightedIndex]); + } + break; + case "Escape": + e.preventDefault(); + closeMentionPicker(); + break; + } + }; + + const handleSubmit = () => { + const trimmed = displayContent.trim(); + if (!trimmed || isSubmitting) return; + + const dataContent = convertDisplayToData(trimmed, insertedMentions); + onSubmit(dataContent); + setDisplayContent(""); + setInsertedMentions([]); + }; + + // Pre-populate insertedMentions from initialValue when members are loaded + useEffect(() => { + if (mentionsInitialized || !initialValue || members.length === 0) return; + + const mentionPattern = /@([^\s@]+(?:\s+[^\s@]+)*?)(?=\s|$|[.,!?;:]|@)/g; + const foundMentions: InsertedMention[] = []; + let match: RegExpExecArray | null; + + while ((match = mentionPattern.exec(initialValue)) !== null) { + const displayName = match[1]; + const member = members.find( + (m) => m.displayName === displayName || m.email.split("@")[0] === displayName + ); + if (member) { + const exists = foundMentions.some((m) => m.id === member.id); + if (!exists) { + foundMentions.push({ id: member.id, displayName }); + } + } + } + + if (foundMentions.length > 0) { + setInsertedMentions(foundMentions); + } + setMentionsInitialized(true); + }, [initialValue, members, mentionsInitialized]); + + useEffect(() => { + if (autoFocus && textareaRef.current) { + textareaRef.current.focus(); + } + }, [autoFocus]); + + const canSubmit = displayContent.trim().length > 0 && !isSubmitting; + + return ( +
+ !open && closeMentionPicker()} + modal={false} + > + +