team-10/env/Lib/site-packages/streamlit/runtime/session_manager.py

# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from abc import abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, Protocol, cast

if TYPE_CHECKING:
    from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
    from streamlit.runtime.app_session import AppSession
    from streamlit.runtime.script_data import ScriptData
    from streamlit.runtime.scriptrunner.script_cache import ScriptCache
    from streamlit.runtime.uploaded_file_manager import UploadedFileManager


class SessionClientDisconnectedError(Exception):
    """Raised by operations on a disconnected SessionClient."""


class SessionClient(Protocol):
    """Interface for sending data to a session's client."""

    @abstractmethod
    def write_forward_msg(self, msg: ForwardMsg) -> None:
        """Deliver a ForwardMsg to the client.

        If the SessionClient has been disconnected, it should raise a
        SessionClientDisconnectedError.
        """
        raise NotImplementedError


@dataclass
class ActiveSessionInfo:
    """Type containing data related to an active session.

    This type is nearly identical to SessionInfo. The difference is that when using it,
    we are guaranteed that SessionClient is not None.
    """

    client: SessionClient
    session: AppSession
    # The number of times the script has been run for this session.
    # At the moment, this is only used for testing and debugging purposes.
    script_run_count: int = 0


@dataclass
class SessionInfo:
    """Type containing data related to an AppSession."""

    client: SessionClient | None
    session: AppSession
    # The number of times the script has been run for this session.
    # At the moment, this is only used for testing and debugging purposes.
    script_run_count: int = 0

    def is_active(self) -> bool:
        return self.client is not None

    def to_active(self) -> ActiveSessionInfo:
        if not self.is_active():
            raise RuntimeError("A SessionInfo with no client cannot be active!")

        # NOTE: The cast here (rather than copying this SessionInfo's fields into a new
        # ActiveSessionInfo) is important as the Runtime expects to be able to mutate
        # what's returned from get_active_session_info to increment script_run_count.
        return cast("ActiveSessionInfo", self)


class SessionStorageError(Exception):
    """Exception class for errors raised by SessionStorage.

    The original error that causes a SessionStorageError to be (re)raised will generally
    be an I/O error specific to the concrete SessionStorage implementation.
    """


class SessionStorage(Protocol):
    @abstractmethod
    def get(self, session_id: str) -> SessionInfo | None:
        """Return the SessionInfo corresponding to session_id, or None if one does not
        exist.

        Parameters
        ----------
        session_id
            The unique ID of the session being fetched.

        Returns
        -------
        SessionInfo or None

        Raises
        ------
        SessionStorageError
            Raised if an error occurs while attempting to fetch the session. This will
            generally happen if there is an error with the underlying storage backend
            (e.g. if we lose our connection to it).
        """
        raise NotImplementedError

    @abstractmethod
    def save(self, session_info: SessionInfo) -> None:
        """Save the given session.

        Parameters
        ----------
        session_info
            The SessionInfo being saved.

        Raises
        ------
        SessionStorageError
            Raised if an error occurs while saving the given session.
        """
        raise NotImplementedError

    @abstractmethod
    def delete(self, session_id: str) -> None:
        """Mark the session corresponding to session_id for deletion and stop tracking
        it.

        Note that:
          - Calling delete on an ID corresponding to a nonexistent session is a no-op.
          - Calling delete on an ID should cause the given session to no longer be
            tracked by this SessionStorage, but exactly when and how the session's data
            is eventually cleaned up is a detail left up to the implementation.

        Parameters
        ----------
        session_id
            The unique ID of the session to delete.

        Raises
        ------
        SessionStorageError
            Raised if an error occurs while attempting to delete the session.
        """
        raise NotImplementedError

    @abstractmethod
    def list(self) -> list[SessionInfo]:
        """List all sessions tracked by this SessionStorage.

        Returns
        -------
        List[SessionInfo]

        Raises
        ------
        SessionStorageError
            Raised if an error occurs while attempting to list sessions.
        """
        raise NotImplementedError


class SessionManager(Protocol):
    """SessionManagers are responsible for encapsulating all session lifecycle behavior
    that the Streamlit Runtime may care about.

    A SessionManager must define the following required methods:
      - __init__
      - connect_session
      - close_session
      - get_session_info
      - list_sessions

    SessionManager implementations may also choose to define the notions of active and
    inactive sessions. The precise definitions of active/inactive are left to the
    concrete implementation. SessionManagers that wish to differentiate between active
    and inactive sessions should have the required methods listed above operate on *all*
    sessions. Additionally, they should define the following methods for working with
    active sessions:
      - disconnect_session
      - get_active_session_info
      - is_active_session
      - list_active_sessions

    When active session-related methods are left undefined, their default
    implementations are the naturally corresponding required methods.

    The Runtime, unless there's a good reason to do otherwise, should generally work
    with the active-session versions of a SessionManager's methods. There isn't currently
    a need for us to be able to operate on inactive sessions stored in SessionStorage
    outside of the SessionManager itself. However, it's highly likely that we'll
    eventually have to do so, which is why the abstractions allow for this now.

    Notes
    -----
    Threading: All SessionManager methods are *not* threadsafe -- they must be called
    from the runtime's eventloop thread.
    """

    @abstractmethod
    def __init__(
        self,
        session_storage: SessionStorage,
        uploaded_file_manager: UploadedFileManager,
        script_cache: ScriptCache,
        message_enqueued_callback: Callable[[], None] | None,
    ) -> None:
        """Initialize a SessionManager with the given SessionStorage.

        Parameters
        ----------
        session_storage
            The SessionStorage instance backing this SessionManager.

        uploaded_file_manager
            Used to manage files uploaded by users via the Streamlit web client.

        script_cache
            ScriptCache instance. Caches user script bytecode.

        message_enqueued_callback
            A callback invoked after a message is enqueued to be sent to a web client.
        """
        raise NotImplementedError

    @abstractmethod
    def connect_session(
        self,
        client: SessionClient,
        script_data: ScriptData,
        user_info: dict[str, str | bool | None],
        existing_session_id: str | None = None,
        session_id_override: str | None = None,
    ) -> str:
        """Create a new session or connect to an existing one.

        Parameters
        ----------
        client
            A concrete SessionClient implementation for communicating with
            the session's client.
        script_data
            Contains parameters related to running a script.
        user_info
            A dict that contains information about the session's user. For now,
            it only (optionally) contains the user's email address.

            {
                "email": "example@example.com"
            }
        existing_session_id
            The ID of an existing session to reconnect to. If one is not provided, a new
            session is created. Note that whether a SessionManager supports reconnecting
            to an existing session is left up to the concrete SessionManager
            implementation. Those that do not support reconnection should simply ignore
            this argument.
        session_id_override
            The ID to assign to a new session being created with this method. Setting
            this can be useful when the service that a Streamlit Runtime is running in
            wants to tie the lifecycle of a Streamlit session to some other session-like
            object that it manages. Only one of existing_session_id and
            session_id_override should be set.

        Returns
        -------
        str
            The session's unique string ID.
        """
        raise NotImplementedError

    @abstractmethod
    def close_session(self, session_id: str) -> None:
        """Close and completely delete the session with the given id.

        This function may be called multiple times for the same session,
        which is not an error. (Subsequent calls just no-op.)

        Parameters
        ----------
        session_id
            The session's unique ID.
        """
        raise NotImplementedError

    @abstractmethod
    def get_session_info(self, session_id: str) -> SessionInfo | None:
        """Return the SessionInfo for the given id, or None if no such session
        exists.

        Parameters
        ----------
        session_id
            The session's unique ID.

        Returns
        -------
        SessionInfo or None
        """
        raise NotImplementedError

    @abstractmethod
    def list_sessions(self) -> list[SessionInfo]:
        """Return the SessionInfo for all sessions managed by this SessionManager.

        Returns
        -------
        List[SessionInfo]
        """
        raise NotImplementedError

    def num_sessions(self) -> int:
        """Return the number of sessions tracked by this SessionManager.

        Subclasses of SessionManager shouldn't provide their own implementation of this
        method without a *very* good reason.

        Returns
        -------
        int
        """
        return len(self.list_sessions())

    # NOTE: The following methods only need to be overwritten when a concrete
    # SessionManager implementation has a notion of active vs inactive sessions.
    # If left unimplemented in a subclass, the default implementations of these methods
    # call corresponding SessionManager methods in a natural way.

    def disconnect_session(self, session_id: str) -> None:
        """Disconnect the given session.

        This method should be idempotent.

        Parameters
        ----------
        session_id
            The session's unique ID.
        """
        self.close_session(session_id)

    def get_active_session_info(self, session_id: str) -> ActiveSessionInfo | None:
        """Return the ActiveSessionInfo for the given id, or None if either no such
        session exists or the session is not active.

        Parameters
        ----------
        session_id
            The active session's unique ID.

        Returns
        -------
        ActiveSessionInfo or None
        """
        session = self.get_session_info(session_id)
        if session is None or not session.is_active():
            return None
        return session.to_active()

    def is_active_session(self, session_id: str) -> bool:
        """Return True if the given session exists and is active, False otherwise.

        Returns
        -------
        bool
        """
        return self.get_active_session_info(session_id) is not None

    def list_active_sessions(self) -> list[ActiveSessionInfo]:
        """Return the session info for all active sessions tracked by this SessionManager.

        Returns
        -------
        List[ActiveSessionInfo]
        """
        return [s.to_active() for s in self.list_sessions()]

    def num_active_sessions(self) -> int:
        """Return the number of active sessions tracked by this SessionManager.

        Subclasses of SessionManager shouldn't provide their own implementation of this
        method without a *very* good reason.

        Returns
        -------
        int
        """
        return len(self.list_active_sessions())