openhands/server/session/manager.py

import asyncio
import json
import time
from dataclasses import dataclass, field
from typing import Generic, Iterable, TypeVar
from uuid import uuid4

import socketio

from openhands.core.config import AppConfig
from openhands.core.exceptions import AgentRuntimeUnavailableError
from openhands.core.logger import openhands_logger as logger
from openhands.core.schema.agent import AgentState
from openhands.events.stream import EventStream, session_exists
from openhands.server.session.agent_session import WAIT_TIME_BEFORE_CLOSE
from openhands.server.session.conversation import Conversation
from openhands.server.session.session import ROOM_KEY, Session
from openhands.server.settings import Settings
from openhands.storage.files import FileStore
from openhands.utils.async_utils import wait_all
from openhands.utils.shutdown_listener import should_continue

_REDIS_POLL_TIMEOUT = 1.5
_CHECK_ALIVE_INTERVAL = 15

_CLEANUP_INTERVAL = 15
MAX_RUNNING_CONVERSATIONS = 3
T = TypeVar('T')


@dataclass
class _ClusterQuery(Generic[T]):
    query_id: str
    request_ids: set[str] | None
    result: T
    flag: asyncio.Event = field(default_factory=asyncio.Event)


@dataclass
class SessionManager:
    sio: socketio.AsyncServer
    config: AppConfig
    file_store: FileStore
    _local_agent_loops_by_sid: dict[str, Session] = field(default_factory=dict)
    _local_connection_id_to_session_id: dict[str, str] = field(default_factory=dict)
    _last_alive_timestamps: dict[str, float] = field(default_factory=dict)
    _redis_listen_task: asyncio.Task | None = None
    _running_sid_queries: dict[str, _ClusterQuery[set[str]]] = field(
        default_factory=dict
    )
    _active_conversations: dict[str, tuple[Conversation, int]] = field(
        default_factory=dict
    )
    _detached_conversations: dict[str, tuple[Conversation, float]] = field(
        default_factory=dict
    )
    _conversations_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    _cleanup_task: asyncio.Task | None = None
    _connection_queries: dict[str, _ClusterQuery[dict[str, str]]] = field(
        default_factory=dict
    )

    async def __aenter__(self):
        redis_client = self._get_redis_client()
        if redis_client:
            self._redis_listen_task = asyncio.create_task(self._redis_subscribe())
        self._cleanup_task = asyncio.create_task(self._cleanup_stale())
        return self

    async def __aexit__(self, exc_type, exc_value, traceback):
        if self._redis_listen_task:
            self._redis_listen_task.cancel()
            self._redis_listen_task = None
        if self._cleanup_task:
            self._cleanup_task.cancel()
            self._cleanup_task = None

    def _get_redis_client(self):
        redis_client = getattr(self.sio.manager, 'redis', None)
        return redis_client

    async def _redis_subscribe(self):
        """
        We use a redis backchannel to send actions between server nodes
        """
        logger.debug('_redis_subscribe')
        redis_client = self._get_redis_client()
        pubsub = redis_client.pubsub()
        await pubsub.subscribe('session_msg')
        while should_continue():
            try:
                message = await pubsub.get_message(
                    ignore_subscribe_messages=True, timeout=5
                )
                if message:
                    await self._process_message(message)
            except asyncio.CancelledError:
                return
            except Exception as e:
                try:
                    asyncio.get_running_loop()
                    logger.error(f'error_reading_from_redis:{str(e)}')
                except RuntimeError:
                    return  # Loop has been shut down

    async def _process_message(self, message: dict):
        data = json.loads(message['data'])
        logger.debug(f'got_published_message:{message}')
        message_type = data['message_type']
        if message_type == 'event':
            sid = data['sid']
            session = self._local_agent_loops_by_sid.get(sid)
            if session:
                await session.dispatch(data['data'])
        elif message_type == 'running_agent_loops_query':
            # Another node in the cluster is asking if the current node is running the session given.
            query_id = data['query_id']
            sids = self._get_running_agent_loops_locally(
                data.get('user_id'), data.get('filter_to_sids')
            )
            if sids:
                await self._get_redis_client().publish(
                    'session_msg',
                    json.dumps(
                        {
                            'query_id': query_id,
                            'sids': list(sids),
                            'message_type': 'running_agent_loops_response',
                        }
                    ),
                )
        elif message_type == 'running_agent_loops_response':
            query_id = data['query_id']
            for sid in data['sids']:
                self._last_alive_timestamps[sid] = time.time()
            running_query = self._running_sid_queries.get(query_id)
            if running_query:
                running_query.result.update(data['sids'])
                if running_query.request_ids is not None and len(
                    running_query.request_ids
                ) == len(running_query.result):
                    running_query.flag.set()
        elif message_type == 'connections_query':
            # Another node in the cluster is asking if the current node is connected to a session
            query_id = data['query_id']
            connections = self._get_connections_locally(
                data.get('user_id'), data.get('filter_to_sids')
            )
            if connections:
                await self._get_redis_client().publish(
                    'session_msg',
                    json.dumps(
                        {
                            'query_id': query_id,
                            'connections': connections,
                            'message_type': 'connections_response',
                        }
                    ),
                )
        elif message_type == 'connections_response':
            query_id = data['query_id']
            connection_query = self._connection_queries.get(query_id)
            if connection_query:
                connection_query.result.update(**data['connections'])
                if connection_query.request_ids is not None and len(
                    connection_query.request_ids
                ) == len(connection_query.result):
                    connection_query.flag.set()
        elif message_type == 'close_session':
            sid = data['sid']
            if sid in self._local_agent_loops_by_sid:
                await self._close_session(sid)
        elif message_type == 'session_closing':
            # Session closing event - We only get this in the event of graceful shutdown,
            # which can't be guaranteed - nodes can simply vanish unexpectedly!
            sid = data['sid']
            logger.debug(f'session_closing:{sid}')
            # Create a list of items to process to avoid modifying dict during iteration
            items = list(self._local_connection_id_to_session_id.items())
            for connection_id, local_sid in items:
                if sid == local_sid:
                    logger.warning(
                        'local_connection_to_closing_session:{connection_id}:{sid}'
                    )
                    await self.sio.disconnect(connection_id)

    async def attach_to_conversation(self, sid: str) -> Conversation | None:
        start_time = time.time()
        if not await session_exists(sid, self.file_store):
            return None

        async with self._conversations_lock:
            # Check if we have an active conversation we can reuse
            if sid in self._active_conversations:
                conversation, count = self._active_conversations[sid]
                self._active_conversations[sid] = (conversation, count + 1)
                logger.info(f'Reusing active conversation {sid}')
                return conversation

            # Check if we have a detached conversation we can reuse
            if sid in self._detached_conversations:
                conversation, _ = self._detached_conversations.pop(sid)
                self._active_conversations[sid] = (conversation, 1)
                logger.info(f'Reusing detached conversation {sid}')
                return conversation

            # Create new conversation if none exists
            c = Conversation(sid, file_store=self.file_store, config=self.config)
            try:
                await c.connect()
            except AgentRuntimeUnavailableError as e:
                logger.error(f'Error connecting to conversation {c.sid}: {e}')
                await c.disconnect()
                return None
            end_time = time.time()
            logger.info(
                f'Conversation {c.sid} connected in {end_time - start_time} seconds'
            )
            self._active_conversations[sid] = (c, 1)
            return c

    async def join_conversation(
        self, sid: str, connection_id: str, settings: Settings, user_id: str | None
    ):
        logger.info(f'join_conversation:{sid}:{connection_id}')
        await self.sio.enter_room(connection_id, ROOM_KEY.format(sid=sid))
        self._local_connection_id_to_session_id[connection_id] = sid
        event_stream = await self._get_event_stream(sid)
        if not event_stream:
            return await self.maybe_start_agent_loop(sid, settings, user_id)
        return event_stream

    async def detach_from_conversation(self, conversation: Conversation):
        sid = conversation.sid
        async with self._conversations_lock:
            if sid in self._active_conversations:
                conv, count = self._active_conversations[sid]
                if count > 1:
                    self._active_conversations[sid] = (conv, count - 1)
                    return
                else:
                    self._active_conversations.pop(sid)
                    self._detached_conversations[sid] = (conversation, time.time())

    async def _cleanup_stale(self):
        while should_continue():
            if self._get_redis_client():
                # Debug info for HA envs
                logger.info(
                    f'Attached conversations: {len(self._active_conversations)}'
                )
                logger.info(
                    f'Detached conversations: {len(self._detached_conversations)}'
                )
                logger.info(
                    f'Running agent loops: {len(self._local_agent_loops_by_sid)}'
                )
                logger.info(
                    f'Local connections: {len(self._local_connection_id_to_session_id)}'
                )
            try:
                async with self._conversations_lock:
                    # Create a list of items to process to avoid modifying dict during iteration
                    items = list(self._detached_conversations.items())
                    for sid, (conversation, detach_time) in items:
                        await conversation.disconnect()
                        self._detached_conversations.pop(sid, None)

                close_threshold = time.time() - self.config.sandbox.close_delay
                running_loops = list(self._local_agent_loops_by_sid.items())
                running_loops.sort(key=lambda item: item[1].last_active_ts)
                sid_to_close: list[str] = []
                for sid, session in running_loops:
                    state = session.agent_session.get_state()
                    if session.last_active_ts < close_threshold and state not in [
                        AgentState.RUNNING,
                        None,
                    ]:
                        sid_to_close.append(sid)

                connections = self._get_connections_locally(
                    filter_to_sids=set(sid_to_close)
                )
                connected_sids = {sid for _, sid in connections.items()}
                sid_to_close = [
                    sid for sid in sid_to_close if sid not in connected_sids
                ]

                if sid_to_close:
                    connections = await self._get_connections_remotely(
                        filter_to_sids=set(sid_to_close)
                    )
                    connected_sids = {sid for _, sid in connections.items()}
                    sid_to_close = [
                        sid for sid in sid_to_close if sid not in connected_sids
                    ]

                await wait_all(self._close_session(sid) for sid in sid_to_close)
                await asyncio.sleep(_CLEANUP_INTERVAL)
            except asyncio.CancelledError:
                async with self._conversations_lock:
                    for conversation, _ in self._detached_conversations.values():
                        await conversation.disconnect()
                    self._detached_conversations.clear()
                await wait_all(
                    (
                        self._close_session(sid)
                        for sid in self._local_agent_loops_by_sid
                    ),
                    timeout=WAIT_TIME_BEFORE_CLOSE,
                )
                return
            except Exception:
                logger.warning('error_cleaning_stale', exc_info=True, stack_info=True)
                await asyncio.sleep(_CLEANUP_INTERVAL)

    async def is_agent_loop_running(self, sid: str) -> bool:
        sids = await self.get_running_agent_loops(filter_to_sids={sid})
        return bool(sids)

    async def get_running_agent_loops(
        self, user_id: str | None = None, filter_to_sids: set[str] | None = None
    ) -> set[str]:
        """Get the running session ids. If a user is supplied, then the results are limited to session ids for that user. If a set of filter_to_sids is supplied, then results are limited to these ids of interest."""
        sids = self._get_running_agent_loops_locally(user_id, filter_to_sids)
        remote_sids = await self._get_running_agent_loops_remotely(
            user_id, filter_to_sids
        )
        return sids.union(remote_sids)

    def _get_running_agent_loops_locally(
        self, user_id: str | None = None, filter_to_sids: set[str] | None = None
    ) -> set[str]:
        items: Iterable[tuple[str, Session]] = self._local_agent_loops_by_sid.items()
        if filter_to_sids is not None:
            items = (item for item in items if item[0] in filter_to_sids)
        if user_id:
            items = (item for item in items if item[1].user_id == user_id)
        sids = {sid for sid, _ in items}
        return sids

    async def _get_running_agent_loops_remotely(
        self,
        user_id: str | None = None,
        filter_to_sids: set[str] | None = None,
    ) -> set[str]:
        """As the rest of the cluster if a session is running. Wait a for a short timeout for a reply"""
        redis_client = self._get_redis_client()
        if not redis_client:
            return set()

        flag = asyncio.Event()
        query_id = str(uuid4())
        query = _ClusterQuery[set[str]](
            query_id=query_id, request_ids=filter_to_sids, result=set()
        )
        self._running_sid_queries[query_id] = query
        try:
            logger.debug(
                f'publish:_get_running_agent_loops_remotely_query:{user_id}:{filter_to_sids}'
            )
            data: dict = {
                'query_id': query_id,
                'message_type': 'running_agent_loops_query',
            }
            if user_id:
                data['user_id'] = user_id
            if filter_to_sids:
                data['filter_to_sids'] = list(filter_to_sids)
            await redis_client.publish('session_msg', json.dumps(data))
            async with asyncio.timeout(_REDIS_POLL_TIMEOUT):
                await flag.wait()

            return query.result
        except TimeoutError:
            # Nobody replied in time
            return query.result
        finally:
            self._running_sid_queries.pop(query_id, None)

    async def get_connections(
        self, user_id: str | None = None, filter_to_sids: set[str] | None = None
    ) -> dict[str, str]:
        connection_ids = self._get_connections_locally(user_id, filter_to_sids)
        remote_connection_ids = await self._get_connections_remotely(
            user_id, filter_to_sids
        )
        connection_ids.update(**remote_connection_ids)
        return connection_ids

    def _get_connections_locally(
        self, user_id: str | None = None, filter_to_sids: set[str] | None = None
    ) -> dict[str, str]:
        connections = dict(**self._local_connection_id_to_session_id)
        if filter_to_sids is not None:
            connections = {
                connection_id: sid
                for connection_id, sid in connections.items()
                if sid in filter_to_sids
            }
        if user_id:
            for connection_id, sid in list(connections.items()):
                session = self._local_agent_loops_by_sid.get(sid)
                if not session or session.user_id != user_id:
                    connections.pop(connection_id)
        return connections

    async def _get_connections_remotely(
        self, user_id: str | None = None, filter_to_sids: set[str] | None = None
    ) -> dict[str, str]:
        redis_client = self._get_redis_client()
        if not redis_client:
            return {}

        flag = asyncio.Event()
        query_id = str(uuid4())
        query = _ClusterQuery[dict[str, str]](
            query_id=query_id, request_ids=filter_to_sids, result={}
        )
        self._connection_queries[query_id] = query
        try:
            logger.debug(
                f'publish:get_connections_remotely_query:{user_id}:{filter_to_sids}'
            )
            data: dict = {
                'query_id': query_id,
                'message_type': 'connections_query',
            }
            if user_id:
                data['user_id'] = user_id
            if filter_to_sids:
                data['filter_to_sids'] = list(filter_to_sids)
            await redis_client.publish('session_msg', json.dumps(data))
            async with asyncio.timeout(_REDIS_POLL_TIMEOUT):
                await flag.wait()

            return query.result
        except TimeoutError:
            # Nobody replied in time
            return query.result
        finally:
            self._connection_queries.pop(query_id, None)

    async def maybe_start_agent_loop(
        self,
        sid: str,
        settings: Settings,
        user_id: str | None,
        initial_user_msg: str | None = None,
    ) -> EventStream:
        logger.info(f'maybe_start_agent_loop:{sid}')
        session: Session | None = None
        if not await self.is_agent_loop_running(sid):
            logger.info(f'start_agent_loop:{sid}')

            response_ids = await self.get_running_agent_loops(user_id)
            if len(response_ids) >= MAX_RUNNING_CONVERSATIONS:
                logger.info('too_many_sessions_for:{user_id}')
                await self.close_session(next(iter(response_ids)))

            session = Session(
                sid=sid,
                file_store=self.file_store,
                config=self.config,
                sio=self.sio,
                user_id=user_id,
            )
            self._local_agent_loops_by_sid[sid] = session
            asyncio.create_task(session.initialize_agent(settings, initial_user_msg))

        event_stream = await self._get_event_stream(sid)
        if not event_stream:
            logger.error(f'No event stream after starting agent loop: {sid}')
            raise RuntimeError(f'no_event_stream:{sid}')
        return event_stream

    async def _get_event_stream(self, sid: str) -> EventStream | None:
        logger.info(f'_get_event_stream:{sid}')
        session = self._local_agent_loops_by_sid.get(sid)
        if session:
            logger.info(f'found_local_agent_loop:{sid}')
            return session.agent_session.event_stream

        if await self._get_running_agent_loops_remotely(filter_to_sids={sid}):
            logger.info(f'found_remote_agent_loop:{sid}')
            return EventStream(sid, self.file_store)

        return None

    async def send_to_event_stream(self, connection_id: str, data: dict):
        # If there is a local session running, send to that
        sid = self._local_connection_id_to_session_id.get(connection_id)
        if not sid:
            raise RuntimeError(f'no_connected_session:{connection_id}')

        session = self._local_agent_loops_by_sid.get(sid)
        if session:
            await session.dispatch(data)
            return

        redis_client = self._get_redis_client()
        if redis_client:
            # If we have a recent report that the session is alive in another pod
            last_alive_at = self._last_alive_timestamps.get(sid) or 0
            next_alive_check = last_alive_at + _CHECK_ALIVE_INTERVAL
            if (
                next_alive_check > time.time()
                or await self._get_running_agent_loops_remotely(filter_to_sids={sid})
            ):
                # Send the event to the other pod
                await redis_client.publish(
                    'session_msg',
                    json.dumps(
                        {
                            'sid': sid,
                            'message_type': 'event',
                            'data': data,
                        }
                    ),
                )
                return

        raise RuntimeError(f'no_connected_session:{connection_id}:{sid}')

    async def disconnect_from_session(self, connection_id: str):
        sid = self._local_connection_id_to_session_id.pop(connection_id, None)
        logger.info(f'disconnect_from_session:{connection_id}:{sid}')
        if not sid:
            # This can occur if the init action was never run.
            logger.warning(f'disconnect_from_uninitialized_session:{connection_id}')
            return

    async def close_session(self, sid: str):
        session = self._local_agent_loops_by_sid.get(sid)
        if session:
            await self._close_session(sid)

        redis_client = self._get_redis_client()
        if redis_client:
            await redis_client.publish(
                'session_msg',
                json.dumps({'sid': sid, 'message_type': 'close_session'}),
            )

    async def _close_session(self, sid: str):
        logger.info(f'_close_session:{sid}')

        # Clear up local variables
        connection_ids_to_remove = list(
            connection_id
            for connection_id, conn_sid in self._local_connection_id_to_session_id.items()
            if sid == conn_sid
        )
        logger.info(f'removing connections: {connection_ids_to_remove}')
        for connnnection_id in connection_ids_to_remove:
            self._local_connection_id_to_session_id.pop(connnnection_id, None)

        session = self._local_agent_loops_by_sid.pop(sid, None)
        if not session:
            logger.warning(f'no_session_to_close:{sid}')
            return

        logger.info(f'closing_session:{session.sid}')
        # We alert the cluster in case they are interested
        try:
            redis_client = self._get_redis_client()
            if redis_client:
                await redis_client.publish(
                    'session_msg',
                    json.dumps({'sid': session.sid, 'message_type': 'session_closing'}),
                )
        except Exception:
            logger.info(
                'error_publishing_close_session_event', exc_info=True, stack_info=True
            )

        await session.close()
        logger.info(f'closed_session:{session.sid}')