from __future__ import annotations

from threading import Lock

from app.core.api.domain.models.agent_session import AgentSession


class InMemorySessionStore:
    def __init__(self) -> None:
        self._items: dict[str, AgentSession] = {}
        self._lock = Lock()

    def save(self, session: AgentSession) -> AgentSession:
        with self._lock:
            self._items[session.session_id] = session
        return session

    def get(self, session_id: str) -> AgentSession | None:
        with self._lock:
            return self._items.get(session_id)
