SIA/sia/auto_approver.py

from threading import Thread, Event
from typing import Callable, TypeAlias, TypedDict
from .web_agent import WebAgent, WebAgentState

class AutoApproverConfig(TypedDict):
    context_enabled: bool
    response_enabled: bool
    context_timeout: float
    response_timeout: float

ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]

class AutoApprover:
    """
    Handles automatic approvals for WebAgent states after configurable timeouts.
    """
    
    def __init__(self, agent: WebAgent):
        """
        Initialize auto approver with a WebAgent instance.
        """
        self.agent = agent
        self._context_timeout = 5.0
        self._response_timeout = 10.0
        self._context_enabled = False
        self._response_enabled = False
        self._stop_event = Event()
        self._context_thread: Thread | None = None
        self._response_thread: Thread | None = None
        self._config_change_handlers: list[ConfigChangeHandler] = []
        
        self.agent.add_state_change_handler(self._handle_state_change)

    @property
    def config(self) -> AutoApproverConfig:
        return AutoApproverConfig(
            context_enabled=self._context_enabled,
            response_enabled=self._response_enabled,
            context_timeout=self._context_timeout,
            response_timeout=self._response_timeout
        )

    def set_config(self, config: AutoApproverConfig) -> None:
        self._context_enabled = config['context_enabled']
        self._response_enabled = config['response_enabled'] 
        self._context_timeout = config['context_timeout']
        self._response_timeout = config['response_timeout']
        
        if self._context_enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
            self._start_context_thread()
        else:
            self._stop_context_thread()

        if self._response_enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
            self._start_response_thread()
        else:
            self._stop_response_thread()
        
        self._notify_config_change()

    def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
        self._config_change_handlers.append(handler)

    def _notify_config_change(self) -> None:
        current_config = self.config
        for handler in self._config_change_handlers:
            handler(current_config)

    @property
    def context_timeout(self) -> float:
        return self._context_timeout
    
    @context_timeout.setter 
    def context_timeout(self, timeout: float) -> None:
        if self._context_enabled:
            raise ValueError("Cannot change timeout while auto-approval is enabled")
        self._context_timeout = timeout
        self._notify_config_change()

    @property
    def response_timeout(self) -> float:
        return self._response_timeout
    
    @response_timeout.setter
    def response_timeout(self, timeout: float) -> None:
        if self._response_enabled:
            raise ValueError("Cannot change timeout while auto-approval is enabled")
        self._response_timeout = timeout
        self._notify_config_change()

    @property
    def context_enabled(self) -> bool:
        return self._context_enabled

    @context_enabled.setter
    def context_enabled(self, enabled: bool) -> None:
        if enabled == self._context_enabled:
            return
        self._context_enabled = enabled
        self._stop_context_thread()
        if enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
            self._start_context_thread()
        self._notify_config_change()

    @property 
    def response_enabled(self) -> bool:
        return self._response_enabled

    @response_enabled.setter
    def response_enabled(self, enabled: bool) -> None:
        if enabled == self._response_enabled:
            return
        self._response_enabled = enabled
        self._stop_response_thread()
        if enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
            self._start_response_thread()
        self._notify_config_change()

    def _handle_state_change(self, new_state: WebAgentState) -> None:
        if new_state == WebAgentState.CONTEXT_APPROVAL:
            self._start_context_thread()
        elif self._context_enabled:
            self._stop_context_thread()

        if new_state == WebAgentState.RESPONSE_APPROVAL:
            self._start_response_thread()
        elif self._response_enabled:
            self._stop_response_thread()

    def _stop_context_thread(self) -> None:
        if self._context_thread:
            self._stop_event.set()
            self._context_thread = None
            self._stop_event.clear()

    def _stop_response_thread(self) -> None:
        if self._response_thread:
            self._stop_event.set()
            self._response_thread = None
            self._stop_event.clear()

    def _start_context_thread(self) -> None:
        self._context_thread = Thread(target=self._context_approval_thread)
        self._context_thread.start()

    def _start_response_thread(self) -> None:
        self._response_thread = Thread(target=self._response_approval_thread)
        self._response_thread.start()

    def _context_approval_thread(self) -> None:
        if self._stop_event.wait(self._context_timeout):
            return
        if (self._context_enabled and 
            self.agent.state == WebAgentState.CONTEXT_APPROVAL):
            self.agent.approve_context()

    def _response_approval_thread(self) -> None:
        if self._stop_event.wait(self._response_timeout):
            return
        if (self._response_enabled and 
            self.agent.state == WebAgentState.RESPONSE_APPROVAL):
            self.agent.approve_response()