src/crud/session.py

from dataclasses import dataclass
from logging import getLogger
from typing import Any
from typing import cast as typing_cast

from cashews import NOT_NONE
from nanoid import generate as generate_nanoid
from sqlalchemy import Select, and_, case, cast, delete, func, insert, select, update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.engine import CursorResult
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import make_transient_to_detached
from sqlalchemy.types import BigInteger, Boolean

from src import models, schemas
from src.cache.client import (
    cache,
    get_cache_namespace,
    safe_cache_delete,
    safe_cache_set,
)
from src.config import settings
from src.exceptions import (
    ConflictException,
    ObserverException,
    ResourceNotFoundException,
)
from src.utils.filter import apply_filter
from src.utils.types import GetOrCreateResult
from src.vector_store import get_external_vector_store

from .peer import get_or_create_peers, get_peer
from .workspace import get_or_create_workspace

logger = getLogger(__name__)


@dataclass
class SessionDeletionResult:
    """Result of a session deletion including cascade counts."""

    messages_deleted: int
    conclusions_deleted: int


SESSION_CACHE_KEY_TEMPLATE = "v2:workspace:{workspace_name}:session:{session_name}"
SESSION_LOCK_PREFIX = f"{get_cache_namespace()}:lock:v2"


def session_cache_key(workspace_name: str, session_name: str) -> str:
    """Generate cache key for session."""
    return (
        get_cache_namespace()
        + ":"
        + SESSION_CACHE_KEY_TEMPLATE.format(
            workspace_name=workspace_name,
            session_name=session_name,
        )
    )


@cache(
    key=SESSION_CACHE_KEY_TEMPLATE,
    ttl=f"{settings.CACHE.DEFAULT_TTL_SECONDS}s",
    prefix=get_cache_namespace(),
    condition=NOT_NONE,
)
@cache.locked(
    key=SESSION_CACHE_KEY_TEMPLATE,
    ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s",
    prefix=SESSION_LOCK_PREFIX,
)
async def _fetch_session(
    db: AsyncSession,
    workspace_name: str,
    session_name: str,
) -> dict[str, Any] | None:
    """Fetch a session from the database and return as a plain dict for safe caching."""
    obj = await db.scalar(
        select(models.Session)
        .where(models.Session.workspace_name == workspace_name)
        .where(models.Session.name == session_name)
    )
    if obj is None:
        return None
    return {
        "id": obj.id,
        "name": obj.name,
        "workspace_name": obj.workspace_name,
        "is_active": obj.is_active,
        "h_metadata": obj.h_metadata,
        "internal_metadata": obj.internal_metadata,
        "configuration": obj.configuration,
        "created_at": obj.created_at,
    }


def count_observers_in_config(
    peer_configs: dict[str, schemas.SessionPeerConfig],
) -> int:
    """
    Count the number of peers that will be observing others based on their configurations.

    Args:
        peer_configs: Dictionary of peer names to their session configurations

    Returns:
        Number of peers that will be observing others
    """
    return sum(1 for config in peer_configs.values() if config.observe_others)


async def get_sessions(
    workspace_name: str,
    filters: dict[str, Any] | None = None,
) -> Select[tuple[models.Session]]:
    """
    Get all active sessions in a workspace.
    """
    stmt = (
        select(models.Session)
        .where(models.Session.workspace_name == workspace_name)
        .where(models.Session.is_active == True)  # noqa: E712
    )

    stmt = apply_filter(stmt, models.Session, filters)

    return stmt.order_by(models.Session.created_at)


async def get_or_create_session(
    db: AsyncSession,
    session: schemas.SessionCreate,
    workspace_name: str,
    *,
    _retry: bool = False,
) -> GetOrCreateResult[models.Session]:
    """
    Get an active session in a workspace or create it if it does not exist.

    If the session already exists, provided metadata replaces the current
    metadata, provided configuration keys are merged into the existing
    configuration, and any provided peers are ensured to be members of the
    session. If the session does not exist, the workspace and peers are created
    as needed before the session is created.

    Args:
        db: Database session
        session: Session creation payload, including optional metadata,
            configuration, and session-peer configuration
        workspace_name: Name of the workspace
        _retry: Whether to retry after a concurrent create conflict

    Returns:
        GetOrCreateResult containing the session and whether it was created

    Raises:
        ValueError: If session.name is empty
        ResourceNotFoundException: If the named session exists but is inactive
        ObserverException: If adding peers would exceed the observer limit
        ConflictException: If concurrent creation prevents fetching or creating
            the session
    """

    if not session.name:
        raise ValueError("Session name must be provided")

    session_data = await _fetch_session(db, workspace_name, session.name)

    # Reconstruct and merge cached dict into session if it exists
    honcho_session: models.Session | None = None
    if session_data is not None:
        obj = models.Session(**session_data)
        make_transient_to_detached(obj)
        honcho_session = await db.merge(obj, load=False)

        # Reject operations on inactive sessions (marked for deletion)
        if not honcho_session.is_active:
            raise ResourceNotFoundException(
                f"Session {session.name} not found in workspace {workspace_name}"
            )

    # Track if we need to update cache and if session was created
    needs_cache_update = False
    created = False
    ws_result = None
    peers_result = None

    # Check if session already exists
    if honcho_session is None:
        if session.peer_names:
            # Count peers that will be observing others
            observer_count = count_observers_in_config(session.peer_names)
            if observer_count > settings.SESSION_OBSERVERS_LIMIT:
                raise ObserverException(session.name, observer_count)

        # Get or create workspace to ensure it exists
        ws_result = await get_or_create_workspace(
            db,
            schemas.WorkspaceCreate(name=workspace_name),
        )

        # Create honcho session
        honcho_session = models.Session(
            workspace_name=workspace_name,
            name=session.name,
            h_metadata=session.metadata or {},
            configuration=session.configuration.model_dump(exclude_none=True)
            if session.configuration
            else {},
        )
        try:
            async with db.begin_nested():
                db.add(honcho_session)
            needs_cache_update = True
            created = True

        except IntegrityError:
            logger.debug(
                "Race condition detected for session: %s, retrying get", session.name
            )
            if _retry:
                raise ConflictException(
                    f"Unable to create or get session: {session.name}"
                ) from None
            return await get_or_create_session(db, session, workspace_name, _retry=True)
    else:
        # Update existing session with metadata and feature flags if provided
        if (
            session.metadata is not None
            and honcho_session.h_metadata != session.metadata
        ):
            honcho_session.h_metadata = session.metadata
            needs_cache_update = True
        if session.configuration is not None:
            # Merge configuration instead of replacing to preserve existing keys
            existing_config = (honcho_session.configuration or {}).copy()
            incoming_config = session.configuration.model_dump(exclude_none=True)
            merged_config = {**existing_config, **incoming_config}
            if honcho_session.configuration != merged_config:
                honcho_session.configuration = merged_config
                needs_cache_update = True

    # Add all peers to session
    if session.peer_names:
        peers_result = await get_or_create_peers(
            db,
            workspace_name=workspace_name,
            peers=[
                schemas.PeerCreate(name=peer_name) for peer_name in session.peer_names
            ],
        )
        await _get_or_add_peers_to_session(
            db,
            workspace_name=workspace_name,
            session_name=session.name,
            peer_names=session.peer_names,
            fetch_after_upsert=False,
        )

    await db.commit()

    # Run deferred cache operations from workspace/peer creation
    if ws_result is not None:
        await ws_result.post_commit()
    if peers_result is not None:
        await peers_result.post_commit()

    # Only update cache if session data changed or was newly created
    if needs_cache_update:
        cache_key = session_cache_key(workspace_name, session.name)
        await safe_cache_set(
            cache_key,
            {
                "id": honcho_session.id,
                "name": honcho_session.name,
                "workspace_name": honcho_session.workspace_name,
                "is_active": honcho_session.is_active,
                "h_metadata": honcho_session.h_metadata,
                "internal_metadata": honcho_session.internal_metadata,
                "configuration": honcho_session.configuration,
                "created_at": honcho_session.created_at,
            },
            expire=settings.CACHE.DEFAULT_TTL_SECONDS,
        )
        logger.debug(
            "Session %s cache updated in workspace %s", session.name, workspace_name
        )

    return GetOrCreateResult(honcho_session, created=created)


async def get_session(
    db: AsyncSession,
    session_name: str,
    workspace_name: str,
    *,
    include_inactive: bool = False,
) -> models.Session:
    """
    Get a session in a workspace.

    Args:
        db: Database session
        session_name: Name of the session
        workspace_name: Name of the workspace
        include_inactive: If True, return sessions even if they are marked for deletion.
            This should only be used for internal operations like the deletion task.

    Returns:
        The session

    Raises:
        ResourceNotFoundException: If the session does not exist or is inactive
    """
    data = await _fetch_session(db, workspace_name, session_name)

    if data is None:
        raise ResourceNotFoundException(
            f"Session {session_name} not found in workspace {workspace_name}"
        )

    # Check if session is active (unless include_inactive is True)
    # Check on the dict before constructing the ORM object
    if not include_inactive and not data["is_active"]:
        raise ResourceNotFoundException(
            f"Session {session_name} not found in workspace {workspace_name}"
        )

    # Reconstruct ORM object from cached dict and merge into session
    obj = models.Session(**data)
    make_transient_to_detached(obj)
    session = await db.merge(obj, load=False)

    return session


async def update_session(
    db: AsyncSession,
    session: schemas.SessionUpdate,
    workspace_name: str,
    session_name: str,
) -> models.Session:
    """
    Get or create a session, then apply metadata and configuration updates.

    Provided metadata replaces the current metadata when present. Provided
    configuration keys are merged into the existing configuration instead of
    replacing it wholesale.

    Args:
        db: Database session
        session: Session update schema
        workspace_name: Name of the workspace
        session_name: Name of the session

    Returns:
        The updated session

    Raises:
        ResourceNotFoundException: If the named session exists but is inactive
        ConflictException: If concurrent creation prevents fetching or creating
            the session
    """
    honcho_session: models.Session = (
        await get_or_create_session(
            db, schemas.SessionCreate(name=session_name), workspace_name=workspace_name
        )
    ).resource

    # Track if anything changed
    needs_update = False

    if session.metadata is not None and honcho_session.h_metadata != session.metadata:
        honcho_session.h_metadata = session.metadata
        needs_update = True

    if session.configuration is not None:
        # Merge configuration instead of replacing to preserve existing keys
        base_config = (honcho_session.configuration or {}).copy()
        merged_config = {
            **base_config,
            **session.configuration.model_dump(exclude_none=True),
        }
        if honcho_session.configuration != merged_config:
            honcho_session.configuration = merged_config
            needs_update = True

    if not needs_update:
        logger.debug(
            "Session %s unchanged in workspace %s, skipping update",
            session_name,
            workspace_name,
        )
        return honcho_session

    await db.commit()

    # Only invalidate if we actually updated
    cache_key = session_cache_key(workspace_name, session_name)
    await safe_cache_delete(cache_key)

    logger.debug("Session %s updated successfully", session_name)
    return honcho_session


async def _batch_delete_matching(
    db: AsyncSession,
    model: Any,
    filter_conditions: list[Any],
    batch_size: int = 5000,
) -> int:
    """
    Delete records in batches that match the given filter conditions.

    Args:
        db: Database session
        model: SQLAlchemy model class
        filter_conditions: List of SQLAlchemy filter conditions
        batch_size: Number of records to delete per batch

    Returns:
        Total number of records deleted
    """
    total_deleted = 0
    primary_key_column = model.__table__.primary_key.columns.values()[0]

    while True:
        subquery = (
            select(primary_key_column).where(and_(*filter_conditions)).limit(batch_size)
        )
        delete_stmt = delete(model).where(primary_key_column.in_(subquery))
        delete_result = typing_cast(CursorResult[Any], await db.execute(delete_stmt))
        batch_deleted = delete_result.rowcount or 0
        total_deleted += batch_deleted

        if batch_deleted == 0:
            break

    return total_deleted


async def delete_session(
    db: AsyncSession, workspace_name: str, session_name: str
) -> SessionDeletionResult:
    """
    Delete a session and all associated data (hard delete).

    This performs cascading deletes for all session-related data including:
    - Active queue sessions
    - Queue items
    - Message embeddings (batched)
    - Documents (theory-of-mind data, batched)
    - Messages (batched)
    - Session peer associations
    - The session itself

    Args:
        db: Database session
        workspace_name: Name of the workspace
        session_name: Name of the session

    Returns:
        SessionDeletionResult containing cascade counts

    Raises:
        ResourceNotFoundException: If the session does not exist
    """
    honcho_session = await get_session(
        db, session_name, workspace_name, include_inactive=True
    )

    # Perform cascading deletes in order
    # Order is important to avoid foreign key constraint violations
    try:
        # Delete ActiveQueueSession entries
        # Work unit keys have format: {task_type}:{workspace_name}:{session_name}:{...}
        await db.execute(
            delete(models.ActiveQueueSession).where(
                and_(
                    func.split_part(models.ActiveQueueSession.work_unit_key, ":", 2)
                    == workspace_name,
                    func.split_part(models.ActiveQueueSession.work_unit_key, ":", 3)
                    == session_name,
                )
            )
        )

        # Delete QueueItem entries
        await db.execute(
            delete(models.QueueItem).where(
                models.QueueItem.session_id == honcho_session.id
            )
        )

        # Delete message vectors from vector store before deleting DB records
        # Fetch all MessageEmbedding records to build vector IDs with {message_id}_{chunk_index}
        embedding_result = await db.execute(
            select(models.MessageEmbedding).where(
                models.MessageEmbedding.session_name == session_name,
                models.MessageEmbedding.workspace_name == workspace_name,
            )
        )
        embeddings = list(embedding_result.scalars().all())
        external_vector_store = get_external_vector_store()

        # Only delete from external vector store if one exists
        if external_vector_store is not None and embeddings:
            # Compute chunk_index for each embedding based on message_id ordering
            message_chunks: dict[str, list[models.MessageEmbedding]] = {}
            for emb in embeddings:
                message_chunks.setdefault(emb.message_id, []).append(emb)

            # Sort each message's chunks by id and build vector IDs
            vector_ids: list[str] = []
            for chunks in message_chunks.values():
                chunks.sort(key=lambda e: e.id)
                for chunk_idx, chunk in enumerate(chunks):
                    vector_ids.append(f"{chunk.message_id}_{chunk_idx}")

            # Try to delete from external vector store (best effort)
            try:
                namespace = external_vector_store.get_vector_namespace(
                    "message", workspace_name
                )
                await external_vector_store.delete_many(namespace, vector_ids)
                logger.debug(
                    f"Deleted {len(vector_ids)} message vectors for session {session_name}"
                )
            except Exception as e:
                # Log warning but continue - workspace deletion will clean up eventually
                logger.warning(
                    f"Failed to delete message vectors for session {session_name}: {e}"
                )

        # Delete MessageEmbedding entries in batches
        await _batch_delete_matching(
            db,
            models.MessageEmbedding,
            [
                models.MessageEmbedding.session_name == session_name,
                models.MessageEmbedding.workspace_name == workspace_name,
            ],
            batch_size=5000,
        )

        # Delete document vectors from vector store before deleting DB records
        # Fetch all Document records to get IDs and namespaces
        doc_result = await db.execute(
            select(
                models.Document.id,
                models.Document.observer,
                models.Document.observed,
            ).where(
                models.Document.session_name == session_name,
                models.Document.workspace_name == workspace_name,
            )
        )
        documents = doc_result.all()

        # Only delete from external vector store if one exists
        if external_vector_store is not None and documents:
            # Group document IDs by namespace (observer/observed)
            docs_by_namespace: dict[str, list[str]] = {}
            for doc in documents:
                namespace = external_vector_store.get_vector_namespace(
                    "document",
                    workspace_name,
                    doc.observer,
                    doc.observed,
                )
                docs_by_namespace.setdefault(namespace, []).append(doc.id)

            # Try to delete from external vector store (best effort, per namespace)
            for namespace, doc_ids in docs_by_namespace.items():
                try:
                    await external_vector_store.delete_many(namespace, doc_ids)
                    logger.debug(
                        f"Deleted {len(doc_ids)} document vectors from {namespace}"
                    )
                except Exception as e:
                    # Log warning but continue - workspace deletion will clean up eventually
                    logger.warning(
                        f"Failed to delete document vectors from {namespace}: {e}"
                    )

        # Delete Document entries associated with this session in batches
        conclusions_deleted = await _batch_delete_matching(
            db,
            models.Document,
            [
                models.Document.session_name == session_name,
                models.Document.workspace_name == workspace_name,
            ],
            batch_size=5000,
        )

        # Delete Message entries in batches
        messages_deleted = await _batch_delete_matching(
            db,
            models.Message,
            [
                models.Message.session_name == session_name,
                models.Message.workspace_name == workspace_name,
            ],
            batch_size=5000,
        )

        # Delete SessionPeer associations
        await db.execute(
            delete(models.SessionPeer).where(
                models.SessionPeer.session_name == session_name,
                models.SessionPeer.workspace_name == workspace_name,
            )
        )

        # Finally, delete the session itself
        await db.delete(honcho_session)
        await db.commit()

        # Invalidate session cache
        await safe_cache_delete(session_cache_key(workspace_name, session_name))

        logger.debug("Session %s and all associated data deleted", session_name)
    except Exception as e:
        logger.error("Failed to delete session %s: %s", session_name, e)
        await db.rollback()
        raise e

    return SessionDeletionResult(
        messages_deleted=messages_deleted,
        conclusions_deleted=conclusions_deleted,
    )


async def clone_session(
    db: AsyncSession,
    workspace_name: str,
    original_session_name: str,
    cutoff_message_id: str | None = None,
) -> models.Session:
    """
    Clone a session and its data. If cutoff_message_id is provided,
    only clone messages up to and including that message.

    The following data is copied to the new session:
    - Session metadata
    - Session configuration
    - All messages (or up to cutoff_message_id) with their content, metadata, and peer associations
    - Session-peer associations with their configurations (observe_me, observe_others)

    The new session gets a unique ID (nanoid) and fresh timestamps.

    Args:
        db: SQLAlchemy session
        workspace_name: Name of the workspace the target session is in
        original_session_name: Name of the session to clone
        cutoff_message_id: Optional ID of the last message to include in the clone

    Returns:
        The newly created session
    """
    # Get the original session (must be active)
    stmt = (
        select(models.Session)
        .where(models.Session.workspace_name == workspace_name)
        .where(models.Session.name == original_session_name)
        .where(models.Session.is_active == True)  # noqa: E712
    )
    result = await db.execute(stmt)
    original_session = result.scalar_one_or_none()
    if original_session is None:
        raise ResourceNotFoundException("Original session not found")

    # If cutoff_message_id is provided, verify it belongs to the session
    cutoff_message = None
    if cutoff_message_id is not None:
        stmt = select(models.Message).where(
            models.Message.public_id == cutoff_message_id,
            models.Message.session_name == original_session_name,
        )
        cutoff_message = await db.scalar(stmt)
        if not cutoff_message:
            raise ValueError(
                "Message not found or doesn't belong to the specified session"
            )

    # Create new session
    new_session = models.Session(
        workspace_name=workspace_name,
        name=generate_nanoid(),
        h_metadata=original_session.h_metadata,
        configuration=original_session.configuration,
    )
    db.add(new_session)
    await db.flush()  # Flush to get the new session ID

    # Build query for messages to clone
    stmt = select(models.Message).where(
        models.Message.session_name == original_session_name
    )
    if cutoff_message_id is not None and cutoff_message is not None:
        stmt = stmt.where(models.Message.id <= cast(cutoff_message.id, BigInteger))
    stmt = stmt.order_by(models.Message.id)

    # Fetch messages to clone
    messages_to_clone_scalars = await db.scalars(stmt)
    messages_to_clone = messages_to_clone_scalars.all()

    if not messages_to_clone:
        return new_session

    # Prepare bulk insert data
    new_messages = [
        {
            "session_name": new_session.name,
            "content": message.content,
            "h_metadata": message.h_metadata,
            "workspace_name": workspace_name,
            "peer_name": message.peer_name,
            "seq_in_session": message.seq_in_session,
        }
        for message in messages_to_clone
    ]

    insert_stmt = insert(models.Message).returning(models.Message)
    result = await db.execute(insert_stmt, new_messages)

    # Clone peers from original session to new session (including their configurations)
    stmt = select(models.SessionPeer).where(
        models.SessionPeer.session_name == original_session_name
    )
    result = await db.execute(stmt)
    session_peers = result.scalars().all()
    for session_peer in session_peers:
        new_session_peer = models.SessionPeer(
            session_name=new_session.name,
            peer_name=session_peer.peer_name,
            workspace_name=workspace_name,
            configuration=session_peer.configuration,
        )
        db.add(new_session_peer)

    await db.commit()
    logger.debug("Session %s cloned successfully", original_session_name)

    # Cache will be populated on next read - read-through pattern
    return new_session


async def remove_peers_from_session(
    db: AsyncSession,
    workspace_name: str,
    session_name: str,
    peer_names: set[str],
) -> bool:
    """
    Remove specified peers from a session.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        session_name: Name of the session
        peer_names: Set of peer names to remove from the session

    Returns:
        True if peers were removed successfully

    Raises:
        ResourceNotFoundException: If the session does not exist
    """
    # Verify session exists
    await get_session(db, session_name, workspace_name)

    # Soft delete specified session peers by setting left_at timestamp
    update_stmt = (
        update(models.SessionPeer)
        .where(
            models.SessionPeer.session_name == session_name,
            models.SessionPeer.workspace_name == workspace_name,
            models.SessionPeer.peer_name.in_(peer_names),
            models.SessionPeer.left_at.is_(None),  # Only update active peers
        )
        .values(left_at=func.now())
    )
    await db.execute(update_stmt)

    await db.commit()
    return True


async def get_peers_from_session(
    workspace_name: str,
    session_name: str,
) -> Select[tuple[models.Peer]]:
    """
    Get all peers from a session.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        session_name: Name of the session

    Returns:
        Paginated list of Peer objects in the session
    """
    # Get all active peers in the session (where left_at is NULL)
    return (
        select(models.Peer)
        .join(
            models.SessionPeer,
            and_(
                models.Peer.name == models.SessionPeer.peer_name,
                models.Peer.workspace_name == models.SessionPeer.workspace_name,
            ),
        )
        .where(models.SessionPeer.session_name == session_name)
        .where(models.Peer.workspace_name == workspace_name)
        .where(models.SessionPeer.left_at.is_(None))  # Only active peers
    )


async def get_session_peer_configuration(
    workspace_name: str,
    session_name: str,
) -> Select[tuple[str, dict[str, Any], dict[str, Any], bool]]:
    """
    Get configuration from both SessionPeer and Peer tables for all peers in a session.
    NOTE: does not filter for active peers. Will return peers that have left the session.

    Args:
        workspace_name: Name of the workspace
        session_name: Name of the session

    Returns:
        Select statement returning peer_name, peer_configuration, session_peer_configuration,
        and a boolean indicating if the peer is currently in the session
    """
    stmt: Select[tuple[str, dict[str, Any], dict[str, Any], bool]] = (
        select(
            models.Peer.name.label("peer_name"),
            models.Peer.configuration.label("peer_configuration"),
            models.SessionPeer.configuration.label("session_peer_configuration"),
            (models.SessionPeer.left_at.is_(None)).label("is_active"),
        )
        .join(
            models.SessionPeer,
            and_(
                models.Peer.name == models.SessionPeer.peer_name,
                models.Peer.workspace_name == models.SessionPeer.workspace_name,
            ),
        )
        .where(models.SessionPeer.session_name == session_name)
        .where(models.Peer.workspace_name == workspace_name)
        .where(models.SessionPeer.workspace_name == workspace_name)
    )

    return stmt


async def set_peers_for_session(
    db: AsyncSession,
    workspace_name: str,
    session_name: str,
    peer_names: dict[str, schemas.SessionPeerConfig],
) -> list[models.SessionPeer]:
    """
    Set peers for a session, overwriting any existing peers.
    If peers don't exist, they will be created.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        session_name: Name of the session
        peer_names: Set of peer names to set for the session

    Returns:
        List of SessionPeer objects for all peers in the session

    Raises:
        ResourceNotFoundException: If the session does not exist
    """
    # Validate observer limit before making any changes
    observer_count = count_observers_in_config(peer_names)
    if observer_count > settings.SESSION_OBSERVERS_LIMIT:
        raise ObserverException(session_name, observer_count)

    # Verify session exists
    stmt = (
        select(models.Session)
        .where(models.Session.workspace_name == workspace_name)
        .where(models.Session.name == session_name)
    )
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()

    if session is None:
        raise ResourceNotFoundException(
            f"Session {session_name} not found in workspace {workspace_name}"
        )

    # Soft delete specified session peers by setting left_at timestamp
    update_stmt = (
        update(models.SessionPeer)
        .where(
            models.SessionPeer.session_name == session_name,
            models.SessionPeer.workspace_name == workspace_name,
            models.SessionPeer.left_at.is_(None),  # Only update active peers
        )
        .values(left_at=func.now())
    )
    result = await db.execute(update_stmt)

    # Get or create peers
    peers_result = await get_or_create_peers(
        db,
        workspace_name=workspace_name,
        peers=[schemas.PeerCreate(name=peer_name) for peer_name in peer_names],
    )

    # Add new peers to session
    peers = await _get_or_add_peers_to_session(
        db,
        workspace_name=workspace_name,
        session_name=session_name,
        peer_names=peer_names,
    )

    await db.commit()
    await peers_result.post_commit()
    return peers


async def _get_or_add_peers_to_session(
    db: AsyncSession,
    workspace_name: str,
    session_name: str,
    peer_names: dict[str, schemas.SessionPeerConfig],
    *,
    fetch_after_upsert: bool = True,
) -> list[models.SessionPeer]:
    """
    Upsert session-peer memberships for a session and optionally fetch the
    active memberships afterward.

    New peers are inserted, peers that previously left the session are rejoined,
    and already-active peers keep their existing session-level configuration.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        session_name: Name of the session
        peer_names: Mapping of peer names to session-level configuration
        fetch_after_upsert: If True, query and return the active session peers
            after the upsert. If False, skip that read and return an empty list.

    Returns:
        Active SessionPeer objects after the upsert, or an empty list when the
        post-upsert fetch is skipped

    Raises:
        ObserverException: If adding peers would exceed the observer limit
    """
    # If no peers to add, skip the insert and just return existing active session peers
    if not peer_names:
        if not fetch_after_upsert:
            return []
        select_stmt = select(models.SessionPeer).where(
            models.SessionPeer.session_name == session_name,
            models.SessionPeer.workspace_name == workspace_name,
            models.SessionPeer.left_at.is_(None),  # Only active peers
        )
        result = await db.execute(select_stmt)
        return list(result.scalars().all())

    # Only validate observer limit if we're adding peers with observe_others=True
    new_observer_count = count_observers_in_config(peer_names)

    if new_observer_count > 0:
        # Use a single efficient query to count existing observers not being updated
        # This uses PostgreSQL's JSONB operators to check the observe_others field directly
        existing_observers_stmt = select(func.count()).where(
            models.SessionPeer.session_name == session_name,
            models.SessionPeer.workspace_name == workspace_name,
            models.SessionPeer.left_at.is_(None),  # Only active peers
            models.SessionPeer.peer_name.notin_(
                peer_names.keys()
            ),  # Exclude peers being updated
            models.SessionPeer.configuration["observe_others"].astext.cast(
                Boolean
            ),  # Only observers
        )
        result = await db.execute(existing_observers_stmt)
        existing_observer_count = result.scalar() or 0

        total_observers = existing_observer_count + new_observer_count

        if total_observers > settings.SESSION_OBSERVERS_LIMIT:
            raise ObserverException(session_name, total_observers)

    # Use upsert to handle both new peers and rejoining peers
    stmt = pg_insert(models.SessionPeer).values(
        [
            {
                "session_name": session_name,
                "peer_name": peer_name,
                "workspace_name": workspace_name,
                "joined_at": func.now(),
                "left_at": None,
                "configuration": configuration.model_dump(),
            }
            for peer_name, configuration in peer_names.items()
        ]
    )

    # On conflict, update joined_at and clear left_at (rejoin scenario)
    # If left_at is not None (peer has left the session): Use the new configuration (stmt.excluded.configuration)
    # If left_at is None (peer is still active): Keep the existing configuration (models.SessionPeer.configuration)
    stmt = stmt.on_conflict_do_update(
        index_elements=["session_name", "peer_name", "workspace_name"],
        set_={
            "joined_at": func.now(),
            "left_at": None,
            "configuration": case(
                (models.SessionPeer.left_at.is_not(None), stmt.excluded.configuration),
                else_=models.SessionPeer.configuration,
            ),
        },
    )
    await db.execute(stmt)

    if not fetch_after_upsert:
        return []

    # Return all active session peers after the upsert
    select_stmt = select(models.SessionPeer).where(
        models.SessionPeer.session_name == session_name,
        models.SessionPeer.workspace_name == workspace_name,
        models.SessionPeer.left_at.is_(None),  # Only active peers
    )
    result = await db.execute(select_stmt)
    return list(result.scalars().all())


async def get_peer_config(
    db: AsyncSession,
    workspace_name: str,
    session_name: str,
    peer_id: str,
) -> schemas.SessionPeerConfig:
    """
    Get the configuration for a peer in a session.


    Args:
        db: Database session
        workspace_name: Name of the workspace
        session_name: Name of the session
        peer_id: Name of the peer


    Returns:
        Configuration for the peer

    Raises:
        ResourceNotFoundException: If the session or peer does not exist
    """
    # Get row from session_peer table
    stmt = select(models.SessionPeer).where(
        models.SessionPeer.workspace_name == workspace_name,
        models.SessionPeer.session_name == session_name,
        models.SessionPeer.peer_name == peer_id,
    )
    result = await db.execute(stmt)
    session_peer = result.scalar_one_or_none()

    if session_peer is None:
        raise ResourceNotFoundException(
            f"Session peer {peer_id} not found in session {session_name} in workspace {workspace_name}"
        )

    return schemas.SessionPeerConfig(**session_peer.configuration)


async def set_peer_config(
    db: AsyncSession,
    workspace_name: str,
    session_name: str,
    peer_name: str,
    config: schemas.SessionPeerConfig,
) -> None:
    """
    Set the configuration for a specific peer in a session.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        session_name: Name of the session
        peer_name: Name of the peer
        config: The peer configuration to set

    Raises:
        ObserverException: If the update would exceed the observer limit
    """
    # First, get the session and peer to ensure they exist
    await get_session(db, session_name, workspace_name)
    await get_peer(db, workspace_name, schemas.PeerCreate(name=peer_name))

    # Check if a SessionPeer entry already exists
    stmt = (
        select(models.SessionPeer)
        .where(models.SessionPeer.session_name == session_name)
        .where(models.SessionPeer.peer_name == peer_name)
        .where(models.SessionPeer.workspace_name == workspace_name)
    )
    result = await db.execute(stmt)
    session_peer = result.scalar_one_or_none()

    # Check if this update would exceed observer limits
    if config.observe_others:
        # Check if peer is already an observer
        is_currently_observer = (
            session_peer.configuration.get("observe_others", False)
            if session_peer and session_peer.configuration
            else False
        )

        # Only need to check limit if peer is becoming a new observer
        if not is_currently_observer:
            # Use a single efficient query to count existing observers
            existing_observers_stmt = select(func.count()).where(
                models.SessionPeer.session_name == session_name,
                models.SessionPeer.workspace_name == workspace_name,
                models.SessionPeer.left_at.is_(None),  # Only active peers
                models.SessionPeer.peer_name
                != peer_name,  # Exclude the peer being updated
                models.SessionPeer.configuration["observe_others"].astext.cast(
                    Boolean
                ),  # Only observers
            )
            result = await db.execute(existing_observers_stmt)
            observer_count = result.scalar() or 0

            # Add one for this peer becoming an observer
            observer_count += 1

            if observer_count > settings.SESSION_OBSERVERS_LIMIT:
                raise ObserverException(session_name, observer_count)

    update_data = config.model_dump(exclude_none=True)

    if session_peer:
        # Update existing configuration
        if session_peer.configuration:
            # Create a new dictionary and update it to ensure SQLAlchemy tracks the change
            new_config = session_peer.configuration.copy()
            new_config.update(update_data)
            session_peer.configuration = new_config
        else:
            session_peer.configuration = update_data
    else:
        # Create a new SessionPeer entry
        session_peer = models.SessionPeer(
            session_name=session_name,
            peer_name=peer_name,
            workspace_name=workspace_name,
            configuration=update_data,
        )
        db.add(session_peer)

    await db.commit()