Add extensive typing to controller directory (#7731)

Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: Ray Myers <ray.myers@gmail.com>
Co-authored-by: Xingyao Wang <xingyao@all-hands.dev>

openhands/controller/agent_controller.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import copy
 import os
@@ -190,7 +192,7 @@ class AgentController:
             self.event_stream.add_event(system_message, EventSource.AGENT)
             logger.debug(f'System message added to event stream: {system_message}')
 
-    async def close(self, set_stop_state=True) -> None:
+    async def close(self, set_stop_state: bool = True) -> None:
         """Closes the agent controller, canceling any ongoing tasks and unsubscribing from the event stream.
 
         Note that it's fairly important that this closes properly, otherwise the state is incomplete.
@@ -242,18 +244,18 @@ class AgentController:
         extra_merged = {'session_id': self.id, **extra}
         getattr(logger, level)(message, extra=extra_merged, stacklevel=2)
 
-    def update_state_before_step(self):
+    def update_state_before_step(self) -> None:
         self.state.iteration += 1
         self.state.local_iteration += 1
 
-    async def update_state_after_step(self):
+    async def update_state_after_step(self) -> None:
         # update metrics especially for cost. Use deepcopy to avoid it being modified by agent._reset()
         self.state.local_metrics = copy.deepcopy(self.agent.llm.metrics)
 
     async def _react_to_exception(
         self,
         e: Exception,
-    ):
+    ) -> None:
         """React to an exception by setting the agent state to error and sending a status message."""
         # Store the error reason before setting the agent state
         self.state.last_error = f'{type(e).__name__}: {str(e)}'
@@ -293,7 +295,10 @@ class AgentController:
         # Set the agent state to ERROR after storing the reason
         await self.set_agent_state_to(AgentState.ERROR)
 
-    async def _step_with_exception_handling(self):
+    def step(self) -> None:
+        asyncio.create_task(self._step_with_exception_handling())
+
+    async def _step_with_exception_handling(self) -> None:
         try:
             await self._step()
         except Exception as e:
@@ -1277,7 +1282,7 @@ class AgentController:
             extra={'msg_type': 'METRICS'},
         )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         pending_action_info = '<none>'
         if (
             hasattr(self, '_pending_action_info')
@@ -1300,7 +1305,7 @@ class AgentController:
             f'_pending_action={pending_action_info})'
         )
 
-    def _is_awaiting_observation(self):
+    def _is_awaiting_observation(self) -> bool:
         events = self.event_stream.get_events(reverse=True)
         for event in events:
             if isinstance(event, AgentStateChangedObservation):