diff --git a/apps/api/alembic/versions/2026_05_28_add_terminal_sessions_table.py b/apps/api/alembic/versions/2026_05_28_add_terminal_sessions_table.py new file mode 100644 index 0000000..4b4054a --- /dev/null +++ b/apps/api/alembic/versions/2026_05_28_add_terminal_sessions_table.py @@ -0,0 +1,61 @@ +"""add terminal_sessions table + +Revision ID: 2026_05_28_add_terminal_sessions +Revises: 20260527_160017_add_pi_agent +Create Date: 2026-05-28 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + +# revision identifiers, used by Alembic. +revision: str = "2026_05_28_add_terminal_sessions" +down_revision: Union[str, None] = "20260527_160017_add_pi_agent" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "terminal_sessions", + sa.Column("id", sa.UUID(), nullable=False), + sa.Column("instance_id", sa.UUID(), nullable=False), + sa.Column("name", sa.String(length=255), nullable=True), + sa.Column("status", sa.String(length=50), nullable=False), + sa.Column("last_activity_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("closed_at", sa.DateTime(timezone=True), nullable=True), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + onupdate=sa.text("now()"), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["instance_id"], ["tool_instances.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_terminal_sessions_instance_id"), + "terminal_sessions", + ["instance_id"], + unique=False, + ) + + +def downgrade() -> None: + op.drop_index( + op.f("ix_terminal_sessions_instance_id"), + table_name="terminal_sessions", + ) + op.drop_table("terminal_sessions") diff --git a/apps/api/src/main.py b/apps/api/src/main.py index ae7b9d7..98d8e4a 100644 --- a/apps/api/src/main.py +++ b/apps/api/src/main.py @@ -24,6 +24,7 @@ from src.api.tool_types import router as tool_types_router from src.api.user_config import router as user_config_router from src.api.users import router as users_router from src.config import Settings +from src.models.terminal_session import TerminalSessionModel # noqa: F401 – Alembic model discovery from src.database import init_database from src.logging_config import ( ExceptionLoggingMiddleware, @@ -66,7 +67,9 @@ def _sanitize_validation_errors(errors): "type": error.get("type"), "loc": error.get("loc"), "msg": error.get("msg"), - "input": str(error.get("input")) if error.get("input") is not None else None, + "input": str(error.get("input")) + if error.get("input") is not None + else None, } # Convert ctx to safe format ctx = error.get("ctx") @@ -110,10 +113,12 @@ async def on_startup(): if not db_ready: logger.error("Database initialization failed. Shutting down.") import sys + sys.exit(1) logger.info("Startup complete.") + app.include_router(health_router) app.include_router(auth_router) app.include_router(dashboard_router) diff --git a/apps/api/src/models/__init__.py b/apps/api/src/models/__init__.py index 4426212..790543d 100644 --- a/apps/api/src/models/__init__.py +++ b/apps/api/src/models/__init__.py @@ -4,9 +4,23 @@ from src.models.config_profile import ConfigProfile, ConfigProfileInclude from src.models.git_repository import GitRepository from src.models.project import Project from src.models.ssh_key import SSHKey +from src.models.terminal_session import TerminalSessionModel from src.models.tool_instance import ToolInstance from src.models.tool_type import ToolType from src.models.user import User from src.models.user_config import UserConfig -__all__ = ["Base", "ConfigFolder", "ConfigProfile", "ConfigProfileInclude", "GitRepository", "Project", "SSHKey", "ToolInstance", "ToolType", "User", "UserConfig"] +__all__ = [ + "Base", + "ConfigFolder", + "ConfigProfile", + "ConfigProfileInclude", + "GitRepository", + "Project", + "SSHKey", + "TerminalSessionModel", + "ToolInstance", + "ToolType", + "User", + "UserConfig", +] diff --git a/apps/api/src/models/terminal_session.py b/apps/api/src/models/terminal_session.py new file mode 100644 index 0000000..20e6f3a --- /dev/null +++ b/apps/api/src/models/terminal_session.py @@ -0,0 +1,37 @@ +"""Terminal session database model.""" + +import uuid +from datetime import datetime + +from sqlalchemy import DateTime, ForeignKey, String +from sqlalchemy import Uuid as UUID +from sqlalchemy.orm import Mapped, mapped_column + +from src.models.base import Base, TimestampMixin, UUIDPrimaryKeyMixin + + +class TerminalSessionModel(UUIDPrimaryKeyMixin, TimestampMixin, Base): + """Database model for terminal session metadata.""" + + __tablename__ = "terminal_sessions" + + instance_id: Mapped[uuid.UUID] = mapped_column( + UUID(), + ForeignKey("tool_instances.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + name: Mapped[str | None] = mapped_column(String(255), nullable=True) + status: Mapped[str] = mapped_column( + String(50), + nullable=False, + default="active", + ) + last_activity_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) + closed_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), + nullable=True, + ) diff --git a/apps/api/src/services/terminal_manager.py b/apps/api/src/services/terminal_manager.py index f83eac1..4db04d2 100644 --- a/apps/api/src/services/terminal_manager.py +++ b/apps/api/src/services/terminal_manager.py @@ -3,20 +3,37 @@ import asyncio import logging import uuid +from datetime import datetime, timezone from fastapi import WebSocket +from src.database import SessionLocal +from src.models.terminal_session import TerminalSessionModel from src.services.terminal_session import TerminalSession logger = logging.getLogger(__name__) +class MaxSessionsExceededError(Exception): + """Raised when the maximum number of terminal sessions per instance is reached.""" + + def __init__(self, instance_id: str, max_sessions: int = 5) -> None: + self.instance_id = instance_id + self.max_sessions = max_sessions + super().__init__( + f"Maximum of {max_sessions} terminal sessions reached for instance {instance_id}" + ) + + class TerminalManager: """Manages active terminal sessions with persistence support.""" + # Maximum sessions per tool instance + MAX_SESSIONS_PER_INSTANCE = 5 + def __init__(self) -> None: - # Track sessions by instance_id for persistence - self._sessions: dict[str, TerminalSession] = {} + # Track sessions by (instance_id, session_id) for multi-session support + self._sessions: dict[tuple[str, str], TerminalSession] = {} self._idle_check_task: asyncio.Task | None = None self._start_idle_check() @@ -42,16 +59,131 @@ class TerminalManager: async def _cleanup_idle_sessions(self) -> None: """Clean up sessions that have been idle for too long.""" - idle_sessions = [] - for instance_id, session in list(self._sessions.items()): + idle_keys = [] + for (instance_id, session_id), session in list(self._sessions.items()): if session.is_idle(): - idle_sessions.append(instance_id) - - for instance_id in idle_sessions: - logger.info("Cleaning up idle terminal session for instance %s", instance_id) - session = self._sessions.pop(instance_id, None) + idle_keys.append((instance_id, session_id)) + + for key in idle_keys: + instance_id, session_id = key + logger.info( + "Cleaning up idle terminal session %s for instance %s", + session_id, + instance_id, + ) + session = self._sessions.pop(key, None) if session: await session.close() + # Update DB status fire-and-forget + asyncio.create_task(self._mark_closed_in_db(session_id)) + + async def _insert_db_session_row( + self, + session_id: str, + instance_id: uuid.UUID, + name: str, + ) -> None: + """Insert a TerminalSessionModel row into the database.""" + try: + async with SessionLocal() as db_session: + db_row = TerminalSessionModel( + id=uuid.UUID(session_id), + instance_id=instance_id, + name=name, + status="active", + created_at=datetime.now(timezone.utc), + last_activity_at=datetime.now(timezone.utc), + ) + db_session.add(db_row) + await db_session.commit() + logger.debug( + "Inserted terminal session row %s for instance %s", + session_id, + instance_id, + ) + except Exception as exc: + logger.error("Failed to insert terminal session row: %s", exc) + + async def _mark_closed_in_db(self, session_id: str) -> None: + """Mark a terminal session as closed in the database.""" + try: + async with SessionLocal() as db_session: + db_row = await db_session.get( + TerminalSessionModel, uuid.UUID(session_id) + ) + if db_row: + db_row.status = "closed" + db_row.closed_at = datetime.now(timezone.utc) + await db_session.commit() + logger.debug( + "Marked terminal session %s as closed in DB", session_id + ) + except Exception as exc: + logger.error("Failed to mark terminal session as closed in DB: %s", exc) + + def _count_sessions_for_instance(self, instance_id_str: str) -> int: + """Count active in-memory sessions for a given instance.""" + return sum(1 for (iid, _sid) in self._sessions if iid == instance_id_str) + + async def create_session( + self, + instance_id: uuid.UUID, + container_id: str, + startup_command: str | None = None, + name: str | None = None, + ) -> TerminalSession: + """Create a new terminal session for an instance. + + Enforces a maximum of MAX_SESSIONS_PER_INSTANCE sessions per instance. + Inserts a DB row fire-and-forget. + + Args: + instance_id: UUID of the tool instance. + container_id: Docker container ID. + startup_command: Optional startup command to run. + name: Optional session name (auto-generated if omitted). + + Returns: + The newly created TerminalSession. + + Raises: + MaxSessionsExceededError: If the instance already has max sessions. + """ + instance_id_str = str(instance_id) + + if ( + self._count_sessions_for_instance(instance_id_str) + >= self.MAX_SESSIONS_PER_INSTANCE + ): + raise MaxSessionsExceededError( + instance_id_str, self.MAX_SESSIONS_PER_INSTANCE + ) + + session_id = str(uuid.uuid4()) + session = TerminalSession( + session_id=session_id, + instance_id=instance_id, + container_id=container_id, + startup_command=startup_command, + name=name, + ) + await session.start(startup_command=startup_command) + + key = (instance_id_str, session_id) + self._sessions[key] = session + + # Fire-and-forget DB insert + asyncio.create_task( + self._insert_db_session_row(session_id, instance_id, session.name) + ) + + logger.info( + "Created terminal session %s for instance %s (name=%s)", + session_id, + instance_id, + session.name, + ) + return session async def get_or_create_session( self, @@ -59,61 +191,128 @@ class TerminalManager: container_id: str, startup_command: str | None = None, ) -> TerminalSession: - """Get existing session or create a new one.""" + """Get existing session or create a new one. + + Backward-compatible alias that uses 'default' as the session_id. + """ # Ensure idle check is running (lazy start) self._start_idle_check() - + instance_id_str = str(instance_id) - - # Check for existing session - if instance_id_str in self._sessions: - session = self._sessions[instance_id_str] - + key = (instance_id_str, "default") + + # Check for existing default session + if key in self._sessions: + session = self._sessions[key] + # Check if session is still alive if session.is_alive(): - logger.debug("Reattaching to existing terminal session for instance %s", instance_id) + logger.debug( + "Reattaching to existing terminal session for instance %s", + instance_id, + ) return session else: # Session died, clean it up - logger.debug("Existing session for instance %s is dead, cleaning up", instance_id) + logger.debug( + "Existing session for instance %s is dead, cleaning up", + instance_id, + ) await session.close() - del self._sessions[instance_id_str] - - # Create new session - logger.info("Creating new terminal session for instance %s", instance_id) + del self._sessions[key] + + # Create new default session + logger.info( + "Creating new default terminal session for instance %s", instance_id + ) session_id = str(uuid.uuid4()) - session = TerminalSession(session_id, instance_id, container_id, startup_command=startup_command) + session = TerminalSession( + session_id=session_id, + instance_id=instance_id, + container_id=container_id, + startup_command=startup_command, + name="Session 1", + ) await session.start(startup_command=startup_command) - self._sessions[instance_id_str] = session - + self._sessions[key] = session + + # Fire-and-forget DB insert + asyncio.create_task( + self._insert_db_session_row(session_id, instance_id, session.name) + ) + return session + def get_session( + self, + instance_id: str, + session_id: str, + ) -> TerminalSession | None: + """Lookup a session by composite key.""" + return self._sessions.get((instance_id, session_id)) + + def get_sessions_for_instance( + self, + instance_id: str, + ) -> list[TerminalSession]: + """Return all in-memory sessions for a given instance.""" + return [ + session + for (iid, _sid), session in self._sessions.items() + if iid == instance_id + ] + + async def close_session( + self, + instance_id: str, + session_id: str, + ) -> None: + """Close a specific session and update its DB status.""" + key = (instance_id, session_id) + session = self._sessions.pop(key, None) + if session: + await session.close() + # Fire-and-forget DB update + asyncio.create_task(self._mark_closed_in_db(session_id)) + logger.info( + "Closed terminal session %s for instance %s", + session_id, + instance_id, + ) + async def attach_websocket( self, session: TerminalSession, websocket: WebSocket, ) -> None: - """Attach a WebSocket to an existing session.""" - # Handle concurrent connections - close existing ones + """Attach a WebSocket to an existing session. + + Closes existing WebSocket connections only for this specific session. + """ + # Handle concurrent connections - close existing ones within the same session if session.has_websockets(): - logger.debug("Closing existing WebSocket connections for instance %s", session.instance_id) + logger.debug( + "Closing existing WebSocket connections for session %s (instance %s)", + session.session_id, + session.instance_id, + ) for ws in list(session._websockets): try: await ws.close(code=4000, reason="New connection established") except Exception: - pass + pass # noqa: S110 session._websockets.clear() - + # Attach new WebSocket session.attach_websocket(websocket) - + # Replay buffer buffer = session.get_buffer() if buffer: try: await websocket.send_bytes(buffer) except Exception: - pass + pass # noqa: S110 async def detach_websocket( self, @@ -128,23 +327,53 @@ class TerminalManager: instance_id: uuid.UUID, container_id: str, startup_command: str | None = None, + session_id: str | None = None, ) -> TerminalSession: - """Reset a session by killing it and creating a new one.""" + """Reset a session by killing it and creating a new one. + + Args: + instance_id: UUID of the tool instance. + container_id: Docker container ID. + startup_command: Optional startup command. + session_id: Specific session to reset. If None, resets the default session. + + Returns: + The newly created TerminalSession. + """ instance_id_str = str(instance_id) - + target_session_id = session_id or "default" + key = (instance_id_str, target_session_id) + # Close existing session if any - if instance_id_str in self._sessions: - logger.debug("Resetting terminal session for instance %s", instance_id) - old_session = self._sessions.pop(instance_id_str) + if key in self._sessions: + logger.debug( + "Resetting terminal session %s for instance %s", + target_session_id, + instance_id, + ) + old_session = self._sessions.pop(key) await old_session.close() - - # Create new session - session_id = str(uuid.uuid4()) - session = TerminalSession(session_id, instance_id, container_id, startup_command=startup_command) - await session.start(startup_command=startup_command) - self._sessions[instance_id_str] = session - - return session + # Fire-and-forget DB update for old session + asyncio.create_task(self._mark_closed_in_db(old_session.session_id)) + + # Create new session preserving the same session_id slot + new_session_id = str(uuid.uuid4()) + new_session = TerminalSession( + session_id=new_session_id, + instance_id=instance_id, + container_id=container_id, + startup_command=startup_command, + name="Session 1" if target_session_id == "default" else None, + ) + await new_session.start(startup_command=startup_command) + self._sessions[key] = new_session + + # Fire-and-forget DB insert + asyncio.create_task( + self._insert_db_session_row(new_session_id, instance_id, new_session.name) + ) + + return new_session async def close_all(self) -> None: """Close all active sessions.""" @@ -152,7 +381,7 @@ class TerminalManager: self._sessions.clear() for session in sessions: await session.close() - + if self._idle_check_task and not self._idle_check_task.done(): self._idle_check_task.cancel() diff --git a/apps/api/src/services/terminal_session.py b/apps/api/src/services/terminal_session.py index e27c999..1297a4a 100644 --- a/apps/api/src/services/terminal_session.py +++ b/apps/api/src/services/terminal_session.py @@ -18,18 +18,28 @@ logger = logging.getLogger(__name__) class TerminalSession: """Manages a single terminal session connected to a docker container. - + Supports persistent sessions that survive WebSocket disconnections. Multiple WebSocket connections can attach/detach from the same session. """ # Circular buffer size (10KB) BUFFER_SIZE = 10 * 1024 - + # Idle timeout in seconds (30 minutes) IDLE_TIMEOUT = 30 * 60 - def __init__(self, session_id: str, instance_id: uuid.UUID, container_id: str, startup_command: str | None = None) -> None: + # Session number counter per instance_id for auto-naming + _instance_counters: dict[str, int] = {} + + def __init__( + self, + session_id: str, + instance_id: uuid.UUID, + container_id: str, + startup_command: str | None = None, + name: str | None = None, + ) -> None: self.session_id = session_id self.instance_id = instance_id self.container_id = container_id @@ -38,37 +48,52 @@ class TerminalSession: self._closed = False self._master_fd: int | None = None self._slave_fd: int | None = None - + # Circular buffer for output replay self._output_buffer: deque[bytes] = deque(maxlen=self.BUFFER_SIZE) self._buffer_size = 0 - + # WebSocket connections self._websockets: set[Any] = set() - + # Activity tracking self.last_activity = time.time() - + # Terminal size self._cols = 80 self._rows = 24 + # Session metadata + self.name = name or self._generate_name(str(instance_id)) + self.status: str = "active" + + @classmethod + def _generate_name(cls, instance_id: str) -> str: + """Generate an auto-incremented session name for the instance.""" + count = cls._instance_counters.get(instance_id, 0) + 1 + cls._instance_counters[instance_id] = count + return f"Session {count}" + async def start(self, startup_command: str | None = None) -> None: """Start the docker exec process with a shell using a PTY.""" # Create a pseudo-terminal on the host self._master_fd, self._slave_fd = pty.openpty() - + # Set the terminal size initially self._set_terminal_size(self._cols, self._rows) - logger.debug(f"Starting terminal session {self.session_id} for container {self.container_id} with initial size {self._cols}x{self._rows}") - + logger.debug( + f"Starting terminal session {self.session_id} for container {self.container_id} with initial size {self._cols}x{self._rows}" + ) + # Build the shell command if startup_command: shell_cmd = f'bash -c "{startup_command}" || true; exec bash -il' - logger.debug(f"Using startup command for session {self.session_id}: {startup_command}") + logger.debug( + f"Using startup command for session {self.session_id}: {startup_command}" + ) else: shell_cmd = "bash -il" - + # Start docker exec with the slave fd as stdin/stdout/stderr # Using -it because the slave fd IS a TTY self.process = await asyncio.create_subprocess_exec( @@ -85,11 +110,11 @@ class TerminalSession: stdout=self._slave_fd, stderr=self._slave_fd, ) - + # Close slave fd in parent process os.close(self._slave_fd) self._slave_fd = None - + self.last_activity = time.time() def _set_terminal_size(self, cols: int, rows: int) -> None: @@ -99,7 +124,7 @@ class TerminalSession: return # TIOCSWINSZ = 0x5414 on Linux TIOCSWINSZ = 0x5414 - size = struct.pack('HHHH', rows, cols, 0, 0) + size = struct.pack("HHHH", rows, cols, 0, 0) try: fcntl.ioctl(self._master_fd, TIOCSWINSZ, size) logger.debug(f"Resized PTY to {cols}x{rows} (fd={self._master_fd})") @@ -127,7 +152,7 @@ class TerminalSession: """Add data to circular buffer, maintaining size limit.""" self._output_buffer.append(data) self._buffer_size += len(data) - + # Trim if exceeds max size while self._buffer_size > self.BUFFER_SIZE and self._output_buffer: removed = self._output_buffer.popleft() @@ -152,16 +177,16 @@ class TerminalSession: if self._closed: logger.warning("Cannot resize: session is closed") return - + # Only resize if dimensions actually changed if cols == self._cols and rows == self._rows: return - + self._cols = cols self._rows = rows logger.debug(f"resize() called for session {self.session_id}: {cols}x{rows}") self._set_terminal_size(cols, rows) - + # Docker exec -it creates its own PTY inside the container, # so host PTY resize doesn't propagate to the container shell. # Send SIGWINCH to the docker exec process on the host. @@ -170,14 +195,19 @@ class TerminalSession: if self.process and self.process.pid: try: os.kill(self.process.pid, signal.SIGWINCH) - logger.debug(f"Sent SIGWINCH to docker exec process {self.process.pid} for session {self.session_id}") + logger.debug( + f"Sent SIGWINCH to docker exec process {self.process.pid} for session {self.session_id}" + ) except ProcessLookupError: - logger.warning(f"docker exec process {self.process.pid} not found for session {self.session_id}") + logger.warning( + f"docker exec process {self.process.pid} not found for session {self.session_id}" + ) except Exception as e: logger.warning(f"Failed to send SIGWINCH: {e}") async def reset(self) -> None: """Reset the session by killing the process and clearing state.""" + self.status = "resetting" await self.close() self._closed = False self._output_buffer.clear() @@ -186,18 +216,20 @@ class TerminalSession: self.process = None self._master_fd = None self._slave_fd = None + self.status = "active" async def close(self) -> None: """Close the session and cleanup.""" if self._closed: return self._closed = True + self.status = "closed" if self._master_fd is not None: try: os.close(self._master_fd) except OSError: - pass + pass # noqa: S110 self._master_fd = None if self.process is not None: @@ -240,7 +272,7 @@ class TerminalSession: await ws.send_bytes(data) except Exception: dead_sockets.add(ws) - + # Clean up dead sockets for ws in dead_sockets: self._websockets.discard(ws) diff --git a/apps/api/tests/services/test_terminal_manager_multi.py b/apps/api/tests/services/test_terminal_manager_multi.py new file mode 100644 index 0000000..ad3dbf3 --- /dev/null +++ b/apps/api/tests/services/test_terminal_manager_multi.py @@ -0,0 +1,203 @@ +"""Unit tests for TerminalManager multi-session support.""" + +import asyncio +import uuid +from unittest.mock import AsyncMock, patch + +import pytest + +from src.services.terminal_manager import MaxSessionsExceededError, TerminalManager +from src.services.terminal_session import TerminalSession + + +@pytest.fixture +def manager() -> TerminalManager: + """Provide a fresh TerminalManager instance for each test.""" + tm = TerminalManager() + # Cancel the background idle check to avoid side effects + if tm._idle_check_task and not tm._idle_check_task.done(): + tm._idle_check_task.cancel() + return tm + + +@pytest.fixture +def mock_terminal_session(monkeypatch) -> None: + """Monkeypatch TerminalSession.start and is_alive for unit tests.""" + + async def fake_start(self, startup_command=None): + self.last_activity = __import__("time").time() + + monkeypatch.setattr(TerminalSession, "start", fake_start) + monkeypatch.setattr(TerminalSession, "is_alive", lambda self: True) + + +@pytest.fixture +def instance_id() -> uuid.UUID: + return uuid.uuid4() + + +class FakeWebSocket: + """Minimal fake WebSocket for testing attach/detach behavior.""" + + def __init__(self, name: str = "ws") -> None: + self.name = name + self.closed = False + self.close_code: int | None = None + self.close_reason: str | None = None + self._sent: list[bytes] = [] + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.closed = True + self.close_code = code + self.close_reason = reason + + async def send_bytes(self, data: bytes) -> None: + self._sent.append(data) + + +@pytest.mark.asyncio +async def test_create_session_increases_count( + manager: TerminalManager, + mock_terminal_session, + instance_id: uuid.UUID, +) -> None: + """Creating sessions increments the per-instance count.""" + assert len(manager.get_sessions_for_instance(str(instance_id))) == 0 + + session1 = await manager.create_session(instance_id, "container-1") + assert len(manager.get_sessions_for_instance(str(instance_id))) == 1 + assert session1.session_id in [ + s.session_id for s in manager.get_sessions_for_instance(str(instance_id)) + ] + + session2 = await manager.create_session(instance_id, "container-1") + assert len(manager.get_sessions_for_instance(str(instance_id))) == 2 + + # Verify sessions are distinct + assert session1.session_id != session2.session_id + + +@pytest.mark.asyncio +async def test_create_session_enforces_max_5( + manager: TerminalManager, + mock_terminal_session, + instance_id: uuid.UUID, +) -> None: + """The 6th session creation raises MaxSessionsExceededError.""" + for i in range(5): + await manager.create_session(instance_id, f"container-{i}") + + assert len(manager.get_sessions_for_instance(str(instance_id))) == 5 + + with pytest.raises(MaxSessionsExceededError): + await manager.create_session(instance_id, "container-overflow") + + +@pytest.mark.asyncio +async def test_get_sessions_for_instance_filters_by_instance( + manager: TerminalManager, + mock_terminal_session, +) -> None: + """get_sessions_for_instance returns only sessions for the requested instance.""" + instance_a = uuid.uuid4() + instance_b = uuid.uuid4() + + await manager.create_session(instance_a, "container-a") + await manager.create_session(instance_a, "container-a2") + await manager.create_session(instance_b, "container-b") + + assert len(manager.get_sessions_for_instance(str(instance_a))) == 2 + assert len(manager.get_sessions_for_instance(str(instance_b))) == 1 + + +@pytest.mark.asyncio +async def test_close_session_removes_from_dict( + manager: TerminalManager, + mock_terminal_session, + instance_id: uuid.UUID, +) -> None: + """close_session removes the key from _sessions and marks DB closed.""" + session = await manager.create_session(instance_id, "container-1") + session_id = session.session_id + + assert manager.get_session(str(instance_id), session_id) is not None + + with patch.object(manager, "_mark_closed_in_db", new=AsyncMock()) as mock_mark: + await manager.close_session(str(instance_id), session_id) + # Give the fire-and-forget task a chance to be scheduled + await asyncio.sleep(0) + + assert manager.get_session(str(instance_id), session_id) is None + mock_mark.assert_called_once_with(session_id) + + +@pytest.mark.asyncio +async def test_attach_websocket_only_closes_same_session( + manager: TerminalManager, + mock_terminal_session, + instance_id: uuid.UUID, +) -> None: + """Attaching to session A must not close WebSockets on session B.""" + session_a = await manager.create_session(instance_id, "container-1") + session_b = await manager.create_session(instance_id, "container-1") + + ws_a1 = FakeWebSocket("ws-a1") + ws_b1 = FakeWebSocket("ws-b1") + + # Manually attach websockets (simulate prior connections) + session_a.attach_websocket(ws_a1) + session_b.attach_websocket(ws_b1) + + # Now attach a new websocket to session_a + ws_a2 = FakeWebSocket("ws-a2") + await manager.attach_websocket(session_a, ws_a2) + + # ws_a1 should have been closed because it's on the same session + assert ws_a1.closed is True + + # ws_b1 should NOT have been closed because it's on a different session + assert ws_b1.closed is False + + # ws_a2 should be attached and receive buffer + assert ws_a2 in session_a._websockets + + +@pytest.mark.asyncio +async def test_default_session_keyed_separately( + manager: TerminalManager, + mock_terminal_session, + instance_id: uuid.UUID, +) -> None: + """Default session uses 'default' session_id and does not collide with named sessions.""" + default_session = await manager.get_or_create_session(instance_id, "container-1") + explicit_session = await manager.create_session(instance_id, "container-1") + + # Both should exist + assert manager.get_session(str(instance_id), "default") is default_session + assert ( + manager.get_session(str(instance_id), explicit_session.session_id) + is explicit_session + ) + + # They should be different objects + assert default_session.session_id != explicit_session.session_id + + +@pytest.mark.asyncio +async def test_idle_cleanup_updates_db_status( + manager: TerminalManager, + mock_terminal_session, + instance_id: uuid.UUID, +) -> None: + """Idle cleanup removes sessions from dict and calls DB update.""" + session = await manager.create_session(instance_id, "container-1") + session_id = session.session_id + + # Make session appear idle (no websockets, old last_activity) + session.last_activity = 0 + + with patch.object(manager, "_mark_closed_in_db", new=AsyncMock()) as mock_mark: + await manager._cleanup_idle_sessions() + + assert manager.get_session(str(instance_id), session_id) is None + mock_mark.assert_called_once_with(session_id)