Free unit when disconnected agent is in waiting (#347)

* Free unit when disconnected agent is in waiting

* Adding unit reservation clear

* Version bump

* Formatting fix

mephisto/VERSION

@@ -1 +1 @@
-0.3.0
+0.3.1

mephisto/abstractions/blueprints/parlai_chat/parlai_chat_task_runner.py

@@ -169,8 +169,10 @@ def run_onboarding(self, agent: "OnboardingAgent") -> None:
         opt: Dict[str, Any] = self.shared_state.onboarding_world_opt
         parlai_agent = MephistoAgentWrapper(agent)
         try:
-            world = self.parlai_world_module.make_onboarding_world(  
-                opt, parlai_agent, initialization_data=self.get_init_data_for_agent(agent)
+            world = self.parlai_world_module.make_onboarding_world(
+                opt,
+                parlai_agent,
+                initialization_data=self.get_init_data_for_agent(agent),
             )  # type: ignore
         except TypeError:
             # make_world doesn't ask for initialization_data

mephisto/data_model/agent.py

@@ -178,13 +178,18 @@ def update_status(self, new_status: str) -> None:
                 f"Updating a final status, was {self.db_status} "
                 f"and want to set to {new_status}"
             )
+        old_status = self.db_status
         self.db.update_agent(self.db_id, status=new_status)
         self.db_status = new_status
         self.has_updated_status.set()
         if new_status in [AgentState.STATUS_RETURNED, AgentState.STATUS_DISCONNECT]:
             # Disconnect statuses should free any pending acts
             self.has_action.set()
             self.did_submit.set()
+            if old_status == AgentState.STATUS_WAITING:
+                # Waiting agents' unit can be reassigned, as no work
+                # has been done yet.
+                self.get_unit().clear_assigned_agent()
 
     @staticmethod
     def _register_agent(

mephisto/data_model/unit.py

@@ -175,6 +175,7 @@ def get_requester(self) -> "Requester":
     def clear_assigned_agent(self) -> None:
         """Clear the agent that is assigned to this unit"""
         self.db.clear_unit_agent_assignment(self.db_id)
+        self.get_task_run().clear_reservation(self)
         self.agent_id = None
         self.__agent = None