from __future__ import annotations

from typing import Dict, Optional


class RuntimeRegistry:
    def __init__(self) -> None:
        self._states: Dict[str, "AgentRuntimeState"] = {}

    def register(self, agent_id: str, state: "AgentRuntimeState") -> None:
        self._states[agent_id] = state

    def get(self, agent_id: str) -> Optional["AgentRuntimeState"]:
        return self._states.get(agent_id)

    def unregister(self, agent_id: str) -> Optional["AgentRuntimeState"]:
        return self._states.pop(agent_id, None)

    def list_agents(self) -> list[str]:
        return list(self._states.keys())

    def clear(self) -> None:
        self._states.clear()