AgentSessionManager: Session Lifecycle Management

AgentSessionManager Implementation - Session Lifecycle Management

As Claude Control grew more complex, it became necessary to manage session lifecycles in a structured way. I implemented an AgentSessionManager that defines clear state transitions for sessions — creation, activation, pause, resume, and termination.

Session State Machine

Session state transitions are modeled as a state machine.

from enum import Enum
from typing import Dict, Set, Optional, Callable, Awaitable
from datetime import datetime

class SessionState(str, Enum):
    CREATED = "created"
    INITIALIZING = "initializing"
    ACTIVE = "active"
    BUSY = "busy"
    PAUSED = "paused"
    ERROR = "error"
    TERMINATING = "terminating"
    TERMINATED = "terminated"

# 허용된 상태 전이
VALID_TRANSITIONS: Dict[SessionState, Set[SessionState]] = {
    SessionState.CREATED: {SessionState.INITIALIZING},
    SessionState.INITIALIZING: {SessionState.ACTIVE, SessionState.ERROR},
    SessionState.ACTIVE: {SessionState.BUSY, SessionState.PAUSED,
                          SessionState.TERMINATING, SessionState.ERROR},
    SessionState.BUSY: {SessionState.ACTIVE, SessionState.ERROR},
    SessionState.PAUSED: {SessionState.ACTIVE, SessionState.TERMINATING},
    SessionState.ERROR: {SessionState.ACTIVE, SessionState.TERMINATING},
    SessionState.TERMINATING: {SessionState.TERMINATED},
}

AgentSessionManager Core Implementation

class AgentSessionManager:
    def __init__(self, redis_store, process_manager, logger):
        self.store = redis_store
        self.process_mgr = process_manager
        self.logger = logger
        self.sessions: Dict[str, SessionInfo] = {}
        self._hooks: Dict[str, List[Callable]] = {}

    async def create_session(self, config: SessionConfig) -> SessionInfo:
        session = SessionInfo(
            id=config.session_id,
            state=SessionState.CREATED,
            workdir=config.workdir,
            created_at=datetime.now(),
        )
        self.sessions[session.id] = session
        await self._transition(session.id, SessionState.INITIALIZING)

        try:
            await self.process_mgr.start_process(session.id, config)
            await self._load_mcp_servers(session.id, config)
            await self._apply_prompt(session.id, config)
            await self._transition(session.id, SessionState.ACTIVE)
        except Exception as e:
            await self._transition(session.id, SessionState.ERROR)
            raise

        return session

    async def _transition(self, session_id: str, new_state: SessionState):
        session = self.sessions[session_id]
        valid_next = VALID_TRANSITIONS.get(session.state, set())
        if new_state not in valid_next:
            raise ValueError(
                f"Invalid transition: {session.state} -> {new_state}"
            )
        old_state = session.state
        session.state = new_state
        session.updated_at = datetime.now()

        await self.store.save_session(session_id, {
            "state": new_state.value,
            "updated_at": session.updated_at.isoformat(),
        })
        await self._fire_hooks(session_id, old_state, new_state)

Lifecycle Hooks

Callbacks can be registered to fire on state transitions.

    def on_transition(self, from_state: SessionState,
                      to_state: SessionState, callback: Callable):
        key = f"{from_state.value}:{to_state.value}"
        if key not in self._hooks:
            self._hooks[key] = []
        self._hooks[key].append(callback)

    async def _fire_hooks(self, session_id, old_state, new_state):
        key = f"{old_state.value}:{new_state.value}"
        for hook in self._hooks.get(key, []):
            await hook(session_id)

Bulk Session Management

    async def pause_all(self):
        for sid, session in self.sessions.items():
            if session.state == SessionState.ACTIVE:
                await self._transition(sid, SessionState.PAUSED)

    async def terminate_all(self):
        tasks = []
        for sid, session in self.sessions.items():
            if session.state not in (SessionState.TERMINATED,
                                      SessionState.TERMINATING):
                tasks.append(self.terminate_session(sid))
        await asyncio.gather(*tasks, return_exceptions=True)

Retrospective

Adopting the state machine pattern made it immediately clear which operations are valid in which states. Previously, there was a bug where a termination request arriving while a session was in the BUSY state would leave the process as a zombie. Enforcing explicit state transitions eliminated that class of problem entirely.