"""Thread-safe dashboard state publication and management. This module provides a thread-safe state store for the dashboard that supports atomic updates, partial state changes, and event notifications for state subscribers. """ from __future__ import annotations import datetime import logging from collections.abc import Callable from dataclasses import replace from threading import RLock from typing import Any, Protocol from .models import ( ComponentStatus, DashboardSnapshot, ExperimentSession, ExperimentState, HealthSnapshot, ParameterChangeRecord, ParameterSchema, PendingParameterChange, ) logger = logging.getLogger(__name__) # ============================================================================ # Type Definitions # ============================================================================ StateChangeCallback = Callable[[DashboardSnapshot, DashboardSnapshot], None] """Callback function type for state change notifications. Args: old_snapshot: The previous dashboard state. new_snapshot: The new dashboard state. """ # ============================================================================ # Protocols # ============================================================================ class StateSource(Protocol): """Protocol for objects that can provide a dashboard snapshot.""" def snapshot(self) -> DashboardSnapshot: """Return the current dashboard snapshot.""" ... # ============================================================================ # State Store # ============================================================================ class DashboardStateStore: """Publish and retrieve immutable dashboard snapshots safely. This class provides thread-safe access to the current dashboard state with support for atomic updates, partial updates, and event notifications. Features: - Thread-safe with RLock - Atomic full-state updates - Partial updates using a builder pattern - Event callbacks for state changes - Optional state history for debugging - Compatible with the simple interface expected by server.py Example: >>> store = DashboardStateStore() >>> store.update(system=SystemMetrics(neurons=1000)) >>> current = store.snapshot() >>> store.on_change(lambda old, new: print(f"State changed: {new.status}")) """ def __init__( self, initial: DashboardSnapshot | None = None, *, max_history: int = 0, notify_on_update: bool = True, ) -> None: """Initialize the state store. Args: initial: Optional initial snapshot. If None, a default is used. max_history: Maximum number of historical states to keep. 0 disables history tracking. notify_on_update: Whether to notify callbacks on updates. """ self._lock = RLock() self._snapshot = initial or DashboardSnapshot() self._version = 0 self._callbacks: list[StateChangeCallback] = [] self._history: list[DashboardSnapshot] = [] self._max_history = max_history self._notify_on_update = notify_on_update self._update_count = 0 # Store the initial snapshot in history if enabled if max_history > 0: self._history.append(self._snapshot) # ========================================================================= # Core Operations # ========================================================================= def publish(self, snapshot: DashboardSnapshot) -> None: """Replace the currently visible snapshot atomically. Args: snapshot: The new dashboard snapshot. """ with self._lock: old = self._snapshot self._snapshot = snapshot self._version += 1 self._update_count += 1 # Track history if enabled if self._max_history > 0: self._history.append(snapshot) if len(self._history) > self._max_history: self._history = self._history[-self._max_history :] # Notify callbacks if self._notify_on_update: self._notify_callbacks(old, snapshot) def snapshot(self) -> DashboardSnapshot: """Return the latest immutable snapshot. Returns: The current dashboard snapshot. """ with self._lock: return self._snapshot def update(self, **kwargs: Any) -> DashboardSnapshot: """Update specific fields of the dashboard state. This is a convenience method that creates a new snapshot with updated fields. It uses the dataclass replace() pattern. Args: **kwargs: Fields to update. Supported keys correspond to the fields of DashboardSnapshot. Returns: The new snapshot that was published. Example: >>> store.update(status="running", system=SystemMetrics(neurons=1000)) """ with self._lock: old = self._snapshot # Create a new snapshot with updated fields new_snapshot = self._update_snapshot(old, **kwargs) # Publish the new snapshot self._snapshot = new_snapshot self._version += 1 self._update_count += 1 # Track history if enabled if self._max_history > 0: self._history.append(new_snapshot) if len(self._history) > self._max_history: self._history = self._history[-self._max_history :] # Notify callbacks if self._notify_on_update: self._notify_callbacks(old, new_snapshot) return new_snapshot def _update_snapshot( self, snapshot: DashboardSnapshot, **kwargs: Any ) -> DashboardSnapshot: """Create a new snapshot with updated fields.""" # Get all fields of DashboardSnapshot field_names = [ "system", "learning", "storage", "self_organization", "homeostasis", "structural", "spikes", "network", "language_organ", "knowledge_intake", "signal_metrics", "experiment", "embodiment", "runtime", "components", "parameters", "pending_changes", "change_history", "experiment_state", "health", "status", "version", ] # Build replacement dict replacements: dict[str, Any] = {} for name in field_names: if name in kwargs: replacements[name] = kwargs[name] # Update the snapshot using dataclass replace return replace(snapshot, **replacements) def update_component(self, status: ComponentStatus) -> DashboardSnapshot: """Update or insert a single component status. Args: status: The new component status. Returns: The new snapshot. """ with self._lock: current = self._snapshot components = dict(current.components or {}) components[status.component] = status return self.update(components=components) def update_parameter(self, parameter: ParameterSchema) -> DashboardSnapshot: """Update or insert a single parameter schema entry. Args: parameter: The new parameter schema entry. Returns: The new snapshot. """ with self._lock: current = self._snapshot parameters = dict(current.parameters or {}) parameters[parameter.name] = parameter return self.update(parameters=parameters) def set_health(self, health: HealthSnapshot) -> DashboardSnapshot: """Set the aggregated health snapshot. Args: health: The new health snapshot. Returns: The new snapshot. """ return self.update(health=health) def set_pending_change( self, change: PendingParameterChange, ) -> DashboardSnapshot: """Record or update a pending parameter change. Args: change: The pending change to store. Returns: The new snapshot. """ with self._lock: current = self._snapshot pending = dict(current.pending_changes or {}) pending[change.name] = change return self.update(pending_changes=pending) def remove_pending_change(self, name: str) -> DashboardSnapshot: """Remove a pending parameter change. Args: name: Parameter name whose pending change should be removed. Returns: The new snapshot. """ with self._lock: current = self._snapshot pending = dict(current.pending_changes or {}) pending.pop(name, None) return self.update(pending_changes=pending) def clear_pending_changes(self) -> DashboardSnapshot: """Clear all pending parameter changes. Returns: The new snapshot. """ return self.update(pending_changes={}) def append_change_history( self, record: ParameterChangeRecord, ) -> DashboardSnapshot: """Append a change record to the history. Args: record: The record to append. Returns: The new snapshot. """ with self._lock: current = self._snapshot history = list(current.change_history) history.append(record) return self.update(change_history=tuple(history)) def set_experiment_mode( self, mode: str, ) -> DashboardSnapshot: """Set the current experiment mode. Args: mode: One of 'operator', 'experiment', 'debug'. Returns: The new snapshot. """ with self._lock: current = self._snapshot new_state = ExperimentState( current_mode=mode, active_session=current.experiment_state.active_session, sessions=current.experiment_state.sessions, ) return self.update(experiment_state=new_state) def start_experiment_session( self, session: ExperimentSession, ) -> DashboardSnapshot: """Start a new experiment/debug session. Any currently active session is closed automatically. Args: session: The session to start. Returns: The new snapshot. """ with self._lock: current = self._snapshot old_state = current.experiment_state sessions = list(old_state.sessions) # Close any active session first if old_state.active_session is not None: now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat() closed = ExperimentSession( session_id=old_state.active_session.session_id, mode=old_state.active_session.mode, hypothesis=old_state.active_session.hypothesis, notes=old_state.active_session.notes, start_tick=old_state.active_session.start_tick, end_tick=session.start_tick, start_time=old_state.active_session.start_time, end_time=now, config_snapshot=old_state.active_session.config_snapshot, active=False, ) sessions.append(closed) sessions.append(session) new_state = ExperimentState( current_mode=session.mode, active_session=session, sessions=tuple(sessions), ) return self.update(experiment_state=new_state) def stop_experiment_session( self, end_tick: int = 0, ) -> DashboardSnapshot: """Stop the active experiment/debug session. Args: end_tick: Simulation tick at session stop. Returns: The new snapshot. """ with self._lock: current = self._snapshot old_state = current.experiment_state active = old_state.active_session if active is None: return current now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat() sessions = [s for s in old_state.sessions if not s.active] closed = ExperimentSession( session_id=active.session_id, mode=active.mode, hypothesis=active.hypothesis, notes=active.notes, start_tick=active.start_tick, end_tick=end_tick, start_time=active.start_time, end_time=now, config_snapshot=active.config_snapshot, active=False, ) sessions.append(closed) new_state = ExperimentState( current_mode=old_state.current_mode, active_session=None, sessions=tuple(sessions), ) return self.update(experiment_state=new_state) def add_experiment_note( self, note: str, ) -> DashboardSnapshot: """Append a note to the active experiment session. Args: note: Note text to append. Returns: The new snapshot. """ with self._lock: current = self._snapshot old_state = current.experiment_state active = old_state.active_session if active is None: return current now = datetime.datetime.now(tz=datetime.timezone.utc).isoformat() updated = ExperimentSession( session_id=active.session_id, mode=active.mode, hypothesis=active.hypothesis, notes=(*active.notes, f"[{now}] {note}"), start_tick=active.start_tick, end_tick=active.end_tick, start_time=active.start_time, end_time=active.end_time, config_snapshot=active.config_snapshot, active=active.active, ) sessions = [ s if s.session_id != active.session_id else updated for s in old_state.sessions ] # Ensure active session is present in sessions list if not any(s.session_id == active.session_id for s in sessions): sessions.append(updated) new_state = ExperimentState( current_mode=old_state.current_mode, active_session=updated, sessions=tuple(sessions), ) return self.update(experiment_state=new_state) # ========================================================================= # Event System # ========================================================================= def on_change(self, callback: StateChangeCallback) -> None: """Register a callback for state changes. Args: callback: Function to call when the state changes. Receives old and new snapshots as arguments. Raises: ValueError: If the callback is already registered. """ with self._lock: if callback in self._callbacks: raise ValueError("Callback already registered") self._callbacks.append(callback) def remove_callback(self, callback: StateChangeCallback) -> bool: """Remove a previously registered callback. Args: callback: The callback to remove. Returns: True if the callback was removed, False if not found. """ with self._lock: try: self._callbacks.remove(callback) return True except ValueError: return False def clear_callbacks(self) -> None: """Remove all registered callbacks.""" with self._lock: self._callbacks.clear() def _notify_callbacks(self, old: DashboardSnapshot, new: DashboardSnapshot) -> None: """Notify all registered callbacks of a state change.""" # Make a copy of the callback list to avoid issues if callbacks # modify the list during iteration. with self._lock: callbacks = list(self._callbacks) for callback in callbacks: try: callback(old, new) except Exception as e: logger.warning(f"State change callback failed: {e}") # ========================================================================= # Query Methods # ========================================================================= def version(self) -> int: """Return the current state version number. The version increments on every state change. """ with self._lock: return self._version def update_count(self) -> int: """Return the total number of state updates.""" with self._lock: return self._update_count def get_history(self, limit: int | None = None) -> list[DashboardSnapshot]: """Return the state history. Args: limit: Maximum number of historical states to return. If None, returns all history. Returns: List of historical snapshots (newest first). """ with self._lock: if limit is None or limit <= 0: return list(reversed(self._history)) return list(reversed(self._history[-limit:])) def get_version_at(self, version: int) -> DashboardSnapshot | None: """Get the state at a specific version. Args: version: Version number to retrieve. Returns: The snapshot at that version, or None if not found. """ with self._lock: if version < 0 or version >= len(self._history): return None return self._history[version] # ========================================================================= # Convenience Methods # ========================================================================= def set_status(self, status: str) -> DashboardSnapshot: """Update the dashboard status. Args: status: New status string (e.g., "idle", "running", "paused"). Returns: The new snapshot. """ return self.update(status=status) def set_version(self, version: str) -> DashboardSnapshot: """Update the system version. Args: version: New version string. Returns: The new snapshot. """ return self.update(version=version) def update_metrics( self, system: Any = None, learning: Any = None, storage: Any = None, homeostasis: Any = None, structural: Any = None, spikes: Any = None, network: Any = None, language_organ: Any = None, knowledge_intake: Any = None, experiment: Any = None, embodiment: Any = None, ) -> DashboardSnapshot: """Update multiple metric fields at once. This is a convenience method for updating all metric-related fields in a single call. Args: system: SystemMetrics instance. learning: LearningMetrics instance. storage: StorageMetrics instance. homeostasis: HomeostasisMetrics instance. structural: StructuralMetrics instance. spikes: SpikeMetrics instance. network: NetworkMetrics instance. language_organ: LanguageOrganMetrics instance. knowledge_intake: KnowledgeIntakeMetrics instance. experiment: ExperimentMetrics instance. embodiment: EmbodimentMetrics instance. Returns: The new snapshot. """ updates: dict[str, Any] = {} if system is not None: updates["system"] = system if learning is not None: updates["learning"] = learning if storage is not None: updates["storage"] = storage if homeostasis is not None: updates["homeostasis"] = homeostasis if structural is not None: updates["structural"] = structural if spikes is not None: updates["spikes"] = spikes if network is not None: updates["network"] = network if language_organ is not None: updates["language_organ"] = language_organ if knowledge_intake is not None: updates["knowledge_intake"] = knowledge_intake if experiment is not None: updates["experiment"] = experiment if embodiment is not None: updates["embodiment"] = embodiment return self.update(**updates) # ========================================================================= # Testing Helpers # ========================================================================= def reset(self, initial: DashboardSnapshot | None = None) -> None: """Reset the state store to a clean state. This is primarily useful for testing. Args: initial: Optional initial snapshot. If None, uses a default. """ with self._lock: self._snapshot = initial or DashboardSnapshot() self._version = 0 self._update_count = 0 self._history = [] self._callbacks.clear() if self._max_history > 0: self._history.append(self._snapshot) # ============================================================================ # Factory Functions # ============================================================================ def create_state_store( initial: DashboardSnapshot | None = None, *, with_history: bool = True, history_limit: int = 100, notify: bool = True, ) -> DashboardStateStore: """Create a configured DashboardStateStore. Args: initial: Optional initial snapshot. with_history: Whether to enable history tracking. history_limit: Maximum number of historical states to keep. notify: Whether to notify callbacks on updates. Returns: Configured DashboardStateStore instance. Example: >>> store = create_state_store(with_history=True, history_limit=50) """ return DashboardStateStore( initial=initial, max_history=history_limit if with_history else 0, notify_on_update=notify, ) # ============================================================================ # Global State Manager (Optional) # ============================================================================ class StateManager: """Global state manager for the dashboard. This is a singleton-like manager that provides access to the dashboard state store. It can be used as a central point for state management across different parts of the application. """ _instance: StateManager | None = None _store: DashboardStateStore | None = None _config_dict: dict[str, Any] | None = None def __new__(cls) -> StateManager: if cls._instance is None: cls._instance = super().__new__(cls) return cls._instance def initialize( self, initial: DashboardSnapshot | None = None, *, with_history: bool = True, history_limit: int = 100, ) -> None: """Initialize the global state store. Args: initial: Optional initial snapshot. with_history: Whether to enable history tracking. history_limit: Maximum number of historical states to keep. """ if self._store is None: self._store = create_state_store( initial=initial, with_history=with_history, history_limit=history_limit, ) def set_config(self, config_dict: dict[str, Any] | None) -> None: """Store the runtime configuration for health derivation.""" self._config_dict = config_dict @property def store(self) -> DashboardStateStore: """Get the global state store.""" if self._store is None: # Auto-initialize with default settings self._store = create_state_store() return self._store def publish(self, snapshot: DashboardSnapshot) -> None: """Publish a new snapshot.""" self.store.publish(snapshot) def snapshot(self) -> DashboardSnapshot: """Get the current snapshot.""" return self.store.snapshot() def update(self, **kwargs: Any) -> DashboardSnapshot: """Update the current state.""" return self.store.update(**kwargs) # ============================================================================ # Global Instance # ============================================================================ # Convenience global instance for simple use cases _state_manager = StateManager() def get_state_store() -> DashboardStateStore: """Get the global state store instance. Returns: The global DashboardStateStore instance. """ return _state_manager.store def publish_state(snapshot: DashboardSnapshot) -> None: """Publish a state update to the global store. The snapshot is automatically enriched with derived components, parameters and health so that every consumer sees a complete view. """ from .health_builder import enrich_snapshot _state_manager.publish( enrich_snapshot( snapshot, _state_manager._config_dict, # pyright: ignore[reportPrivateUsage] ) ) def set_dashboard_config(config_dict: dict[str, Any] | None) -> None: """Store the runtime configuration used for health derivation.""" _state_manager.set_config(config_dict) def get_current_state() -> DashboardSnapshot: """Get the current state from the global store.""" return _state_manager.snapshot() # ============================================================================ # Example Usage # ============================================================================ if __name__ == "__main__": # Basic usage store = DashboardStateStore() # Register a callback def on_state_change(old: DashboardSnapshot, new: DashboardSnapshot) -> None: print(f"State changed: {old.status} -> {new.status}") store.on_change(on_state_change) # Update state store.set_status("running") store.set_version("0.6.0") # Get current state current = store.snapshot() print(f"Current status: {current.status}") print(f"Version: {current.version}")