# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026) # # 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 import threading import time from collections import defaultdict from typing import TYPE_CHECKING, Final, cast from streamlit import config from streamlit.logger import get_logger from streamlit.runtime.app_session import AppSession from streamlit.runtime.session_manager import ( ActiveSessionInfo, SessionClient, SessionInfo, SessionManager, SessionStorage, ) from streamlit.runtime.stats import ( ACTIVE_SESSIONS_FAMILY, SESSION_DURATION_FAMILY, SESSION_EVENTS_FAMILY, USER_SESSION_EVENTS_FAMILY, CounterStat, GaugeStat, Stat, StatsProvider, ) if TYPE_CHECKING: from collections.abc import Callable, Mapping, Sequence from streamlit.runtime.script_data import ScriptData from streamlit.runtime.scriptrunner.script_cache import ScriptCache from streamlit.runtime.scriptrunner_utils.script_run_context import ( OnScriptErrorHandler, UserInfoType, ) from streamlit.runtime.uploaded_file_manager import UploadedFileManager _LOGGER: Final = get_logger(__name__) _EVENT_TYPE_CONNECT: Final = "connect" _EVENT_TYPE_RECONNECT: Final = "reconnect" _EVENT_TYPE_DISCONNECT: Final = "disconnect" _EVENT_TYPE_CLOSE: Final = "close" class WebsocketSessionManager(SessionManager, StatsProvider): """A SessionManager used to manage sessions with lifecycles tied to those of a browser tab's websocket connection. WebsocketSessionManagers differentiate between "active" and "inactive" sessions. Active sessions are those with a currently active websocket connection. Inactive sessions are sessions without. Eventual cleanup of inactive sessions is a detail left to the specific SessionStorage that a WebsocketSessionManager is instantiated with. """ @property def stats_families(self) -> Sequence[str]: # USER_SESSION_EVENTS_FAMILY is advertised unconditionally so that the # StatsManager (which snapshots this property once at registration time) # always routes ?families=user_session_events requests to this provider. # Emission is gated separately in get_stats so the endpoint output is # unchanged while the feature is disabled. return ( SESSION_EVENTS_FAMILY, SESSION_DURATION_FAMILY, ACTIVE_SESSIONS_FAMILY, USER_SESSION_EVENTS_FAMILY, ) def __init__( self, session_storage: SessionStorage, uploaded_file_manager: UploadedFileManager, script_cache: ScriptCache, message_enqueued_callback: Callable[[], None] | None, on_script_error: OnScriptErrorHandler | None = None, ) -> None: self._session_storage = session_storage self._uploaded_file_mgr = uploaded_file_manager self._script_cache = script_cache self._message_enqueued_callback = message_enqueued_callback self._on_script_error = on_script_error # Mapping of AppSession.id -> ActiveSessionInfo. self._active_session_info_by_id: dict[str, ActiveSessionInfo] = {} # Session event counters for metrics self._stats_lock = threading.Lock() self._connect_count: int = 0 self._reconnect_count: int = 0 self._disconnect_count: int = 0 # Session duration tracking self._session_connect_times: dict[str, float] = {} self._total_session_duration_seconds: float = 0 # Per-user session-event counters, keyed by a canonical tuple of # (attr_name, attr_value) label pairs -> {event_type: count}. Only # populated when server.unsafeMetricsUserAttributes is non-empty. Never # pruned (bounded by the number of distinct users seen by the process). self._user_event_counts: defaultdict[ tuple[tuple[str, str], ...], defaultdict[str, int] ] = defaultdict(lambda: defaultdict(int)) # Identity cached per session id while a session is connected, so the # terminal disconnect/close event can be attributed to the connect-time # user. The entry is dropped on the first disconnect/close (and # refreshed on reconnect), so the cache stays bounded by the set of # currently-connected sessions and cannot leak when a disconnected # session is later evicted from storage without an explicit close. self._session_user_labels: dict[str, tuple[tuple[str, str], ...]] = {} def connect_session( self, client: SessionClient, script_data: ScriptData, user_info: UserInfoType, existing_session_id: str | None = None, session_id_override: str | None = None, ) -> str: if existing_session_id and session_id_override: # pragma: no cover - defensive raise RuntimeError( "Only one of existing_session_id and session_id_override should be truthy. " "This should never happen." ) active_session_info = ( self._active_session_info_by_id.get(existing_session_id) if existing_session_id else None ) if active_session_info is not None: if active_session_info.session.matches_user_info(user_info): _LOGGER.warning( "Session with id %s is already connected! Reconnecting to existing session.", existing_session_id, ) # Disconnecting the still-active session moves it into storage, # which makes it available to the storage-based reconnect path # below. The new client then reconnects to the same session # (preserving its state) instead of a brand-new one being # created. self.disconnect_session(active_session_info.session.id) else: # The reconnect id points at a session that is still active # under a different user's identity. Validate identity *before* # disconnecting so we never tear down another user's active # session: treat the id as invalid, hand out a fresh session, # and leave the original session and its client untouched. _LOGGER.warning( "Ignoring reconnect to active session id %s: user identity " "mismatch. Connecting to a new session instead.", existing_session_id, ) existing_session_id = None session_info = ( existing_session_id and existing_session_id not in self._active_session_info_by_id and self._session_storage.get(existing_session_id) ) # Only reconnect to a stored session when the reconnecting identity # matches the identity that originally owned it. A mismatch is treated # like an invalid reconnect id: fall through to create a fresh session # and leave the stored session untouched. This is the defense for # normally-disconnected sessions (the active-session path above handles # the not-yet-disconnected case). if isinstance( session_info, SessionInfo ) and not session_info.session.matches_user_info(user_info): session_info = None if isinstance(session_info, SessionInfo): existing_session = session_info.session existing_session.register_file_watchers() self._active_session_info_by_id[existing_session.id] = ActiveSessionInfo( client, existing_session, session_info.script_run_count, ) self._session_storage.delete(existing_session.id) with self._stats_lock: self._reconnect_count += 1 self._session_connect_times[existing_session.id] = time.monotonic() self._record_user_event( existing_session.id, _EVENT_TYPE_RECONNECT, user_info ) return existing_session.id session = AppSession( script_data=script_data, uploaded_file_manager=self._uploaded_file_mgr, script_cache=self._script_cache, message_enqueued_callback=self._message_enqueued_callback, user_info=user_info, session_id_override=session_id_override, on_script_error=self._on_script_error, ) _LOGGER.debug( "Created new session for client %s. Session ID: %s", id(client), session.id ) if ( session.id in self._active_session_info_by_id ): # pragma: no cover - defensive raise RuntimeError( f"session.id '{session.id}' registered multiple times. " "This should never happen." ) self._active_session_info_by_id[session.id] = ActiveSessionInfo(client, session) with self._stats_lock: self._connect_count += 1 self._session_connect_times[session.id] = time.monotonic() self._record_user_event(session.id, _EVENT_TYPE_CONNECT, user_info) return session.id def disconnect_session(self, session_id: str) -> None: if session_id in self._active_session_info_by_id: active_session_info = self._active_session_info_by_id[session_id] session = active_session_info.session session.request_script_stop() session.disconnect_file_watchers() session.clear_session_caches() self._session_storage.save( SessionInfo( client=None, session=session, script_run_count=active_session_info.script_run_count, ) ) del self._active_session_info_by_id[session_id] with self._stats_lock: self._disconnect_count += 1 self._accumulate_session_duration(session_id) self._record_cached_user_event(session_id, _EVENT_TYPE_DISCONNECT) if not self._active_session_info_by_id: # Avoid stale cached scripts when all file watchers and sessions are disconnected self._script_cache.clear() def get_active_session_info(self, session_id: str) -> ActiveSessionInfo | None: return self._active_session_info_by_id.get(session_id) def is_active_session(self, session_id: str) -> bool: return session_id in self._active_session_info_by_id def list_active_sessions(self) -> list[ActiveSessionInfo]: return list(self._active_session_info_by_id.values()) def close_session(self, session_id: str) -> None: if session_id in self._active_session_info_by_id: active_session_info = self._active_session_info_by_id[session_id] del self._active_session_info_by_id[session_id] active_session_info.session.shutdown() # Count disconnect for active sessions being closed directly with self._stats_lock: self._disconnect_count += 1 self._accumulate_session_duration(session_id) self._record_cached_user_event(session_id, _EVENT_TYPE_CLOSE) if not self._active_session_info_by_id: # Avoid stale cached scripts when all file watchers and sessions are disconnected self._script_cache.clear() return # For sessions in storage, the disconnect was already counted when # disconnect_session was called earlier. session_info = self._session_storage.get(session_id) if session_info: self._session_storage.delete(session_id) session_info.session.shutdown() with self._stats_lock: self._accumulate_session_duration(session_id) self._record_cached_user_event(session_id, _EVENT_TYPE_CLOSE) def _accumulate_session_duration(self, session_id: str) -> None: """Accumulate the session duration for a closed session. This method must be called while holding self._stats_lock. """ connect_time = self._session_connect_times.pop(session_id, None) if connect_time is not None: duration = time.monotonic() - connect_time self._total_session_duration_seconds += duration def _user_labels( self, user_info: UserInfoType ) -> tuple[tuple[str, str], ...] | None: """Resolve canonical (name, value) label pairs from user_info. Returns None when the feature is disabled (the option is empty). Missing/None attributes become "" (but other falsy values such as ``False`` are preserved). This method must be called while holding self._stats_lock. """ attrs = config.get_option("server.unsafeMetricsUserAttributes") if not attrs: return None return tuple( sorted( (name, "" if (value := user_info.get(name)) is None else str(value)) for name in attrs ) ) def _record_user_event( self, session_id: str, event_type: str, user_info: UserInfoType ) -> None: """Record a per-user connect/reconnect event and cache the identity. Best-effort: a telemetry failure must never break the session lifecycle. This method must be called while holding self._stats_lock. """ try: labels = self._user_labels(user_info) if labels is not None: self._user_event_counts[labels][event_type] += 1 self._session_user_labels[session_id] = labels except Exception: _LOGGER.debug( "Failed to record per-user %s event", event_type, exc_info=True ) def _record_cached_user_event(self, session_id: str, event_type: str) -> None: """Record a per-user disconnect/close event using the cached identity. The cached identity is always removed (even when the feature is disabled), so the cache stays bounded by the currently-connected sessions and cannot leak when a disconnected session is later evicted from storage without an explicit close. The option is re-checked so a runtime-disable stops emitting new events. Best-effort: a telemetry failure must never break the session lifecycle. Must be called while holding self._stats_lock. """ try: labels = self._session_user_labels.pop(session_id, None) if labels is not None and config.get_option( "server.unsafeMetricsUserAttributes" ): self._user_event_counts[labels][event_type] += 1 except Exception: _LOGGER.debug( "Failed to record per-user %s event", event_type, exc_info=True ) def get_session_info(self, session_id: str) -> SessionInfo | None: session_info = self.get_active_session_info(session_id) if session_info: return cast("SessionInfo", session_info) return self._session_storage.get(session_id) def list_sessions(self) -> list[SessionInfo]: return ( cast("list[SessionInfo]", self.list_active_sessions()) + self._session_storage.list() ) def get_stats( self, family_names: Sequence[str] | None = None ) -> Mapping[str, Sequence[Stat]]: """Return session-related metrics. Returns session event counters (connections, reconnections, disconnections) and the current count of active sessions. """ result: dict[str, list[Stat]] = {} if family_names is None or SESSION_EVENTS_FAMILY in family_names: with self._stats_lock: connect_count = self._connect_count reconnect_count = self._reconnect_count disconnect_count = self._disconnect_count result[SESSION_EVENTS_FAMILY] = [ CounterStat( family_name=SESSION_EVENTS_FAMILY, value=connect_count, labels={"type": _EVENT_TYPE_CONNECT}, help="Total count of session events by type.", ), CounterStat( family_name=SESSION_EVENTS_FAMILY, value=reconnect_count, labels={"type": _EVENT_TYPE_RECONNECT}, help="Total count of session events by type.", ), CounterStat( family_name=SESSION_EVENTS_FAMILY, value=disconnect_count, labels={"type": _EVENT_TYPE_DISCONNECT}, help="Total count of session events by type.", ), ] if family_names is None or SESSION_DURATION_FAMILY in family_names: with self._stats_lock: total_duration = int(self._total_session_duration_seconds) result[SESSION_DURATION_FAMILY] = [ CounterStat( family_name=SESSION_DURATION_FAMILY, value=total_duration, unit="seconds", help="Total time spent in active sessions, in seconds.", ), ] if family_names is None or ACTIVE_SESSIONS_FAMILY in family_names: result[ACTIVE_SESSIONS_FAMILY] = [ GaugeStat( family_name=ACTIVE_SESSIONS_FAMILY, value=len(self._active_session_info_by_id), help="Current number of active sessions.", ), ] if config.get_option("server.unsafeMetricsUserAttributes") and ( family_names is None or USER_SESSION_EVENTS_FAMILY in family_names ): with self._stats_lock: snapshot = { labels: dict(events) for labels, events in self._user_event_counts.items() } result[USER_SESSION_EVENTS_FAMILY] = [ CounterStat( family_name=USER_SESSION_EVENTS_FAMILY, value=count, # `type` must win over any user attribute; unpack user # labels first so the discriminator is never shadowed. labels={**dict(labels), "type": event_type}, help="Total count of session events by type and user.", ) for labels, events in snapshot.items() for event_type, count in events.items() ] return result