feat: multi-session terminal backend core (PR 1)
- Add TerminalSessionModel DB table with instance_id FK, name, status, created_at, last_activity_at, closed_at columns - Add Alembic migration for terminal_sessions table - Refactor TerminalManager to use composite key (instance_id, session_id) supporting up to 5 concurrent sessions per instance - Add create_session, get_session, get_sessions_for_instance, close_session - Preserve get_or_create_session for backward compatibility (default session) - Fix attach_websocket to only close sockets within same session - Add name (auto-generated 'Session N') and status tracking to TerminalSession - Add 7 unit tests for multi-session logic Quality gates: pytest (7 new passed, 174 total passed, 51 pre-existing failures)
This commit is contained in:
@@ -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")
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user