agent-lab/src/agent_lab/application/session_store.py

from typing import Any, Protocol

from agent_lab.domain.messages import ChatMessage, TokenUsage


class SessionStore(Protocol):
    def create_session(
        self,
        *,
        title: str,
        config: dict[str, Any],
        session_id: str | None = None,
    ) -> str:
        ...

    def ensure_session(
        self,
        session_id: str | None,
        *,
        title: str,
        config: dict[str, Any],
    ) -> str:
        ...

    def list_sessions(self) -> list[dict[str, Any]]:
        ...

    def get_session(self, session_id: str) -> dict[str, Any] | None:
        ...

    def next_turn_index(self, session_id: str) -> int:
        ...

    def start_turn(
        self,
        session_id: str,
        *,
        turn_index: int,
        user_message: str,
    ) -> None:
        ...

    def complete_turn(self, session_id: str, *, turn_index: int) -> None:
        ...

    def append_message(
        self,
        session_id: str,
        *,
        turn_index: int,
        message: ChatMessage,
    ) -> None:
        ...

    def append_messages(
        self,
        session_id: str,
        *,
        turn_index: int,
        messages: list[ChatMessage],
    ) -> None:
        ...

    def append_audit(
        self,
        session_id: str,
        *,
        event: str,
        details: dict[str, Any],
        turn_index: int | None = None,
        round_index: int | None = None,
    ) -> None:
        ...

    def append_usage(
        self,
        session_id: str,
        *,
        turn_index: int,
        round_index: int,
        usage: TokenUsage,
        ttft_ms: int | None,
        elapsed_ms: int,
    ) -> None:
        ...

    def list_messages(self, session_id: str) -> list[dict[str, Any]]:
        ...

    def list_audit_logs(self, session_id: str) -> list[dict[str, Any]]:
        ...

    def usage_summary(self, session_id: str) -> dict[str, Any]:
        ...