src/localscript_api/sessions.py - master - Watashicuvu/agentic-tools

src/localscript_api/sessions.py
65 строк2 KB
Якуб
some fixes of api server
13 апр 2026, 11:24
13 апр 2026, 11:246827c5f
"""Управление сессиями генерации для интерактивных clarification."""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger(__name__)


@dataclass
class Session:
    """Одна генерационная сессия."""

    id: str
    # Очередь ответов пользователя (вопрос → ответ)
    answer_queue: asyncio.Queue[str] = field(default_factory=asyncio.Queue)
    # Событие: новый ответ доступен
    answer_event: asyncio.Event = field(default_factory=asyncio.Event)
    # Текущий вопрос (если задан)
    current_question: Optional[str] = None
    # Флаг завершения
    finished: bool = False


class SessionManager:
    """Хранит активные сессии генерации."""

    def __init__(self):
        self._sessions: dict[str, Session] = {}

    def create(self, session_id: str) -> Session:
        """Создать новую сессию."""
        session = Session(id=session_id)
        self._sessions[session_id] = session
        logger.info(f"Session created: {session_id}")
        return session

    def get(self, session_id: str) -> Optional[Session]:
        """Получить сессию по ID."""
        return self._sessions.get(session_id)

    def finish(self, session_id: str):
        """Завершить и удалить сессию."""
        if session_id in self._sessions:
            self._sessions[session_id].finished = True
            del self._sessions[session_id]
            logger.info(f"Session finished: {session_id}")

    async def wait_for_answer(
        self, session: Session, timeout: float = 120.0
    ) -> Optional[str]:
        """Ждать ответ от клиента.

        Returns:
            Текст ответа или None по таймауту.
        """
        try:
            async with asyncio.timeout(timeout):
                return await session.answer_queue.get()
        except TimeoutError:
            logger.warning(f"Session {session.id} timed out waiting for answer")
            return None