diff --git a/test_env/fake_env.py b/test_env/fake_env.py
index 98543c4..ec3ea3a 100644
--- a/test_env/fake_env.py
+++ b/test_env/fake_env.py
@@ -83,7 +83,7 @@ class DesktopEnv:
     def _stop_video_recording(self):
         pass
     
-    def step(self, action) -> Tuple:
+    def step(self, action, *args, **kargs) -> Tuple:
         self._step_no += 1
         self.action_history.append(action)
 
@@ -116,7 +116,7 @@ class DesktopEnv:
                 }
                 info.update({"tool_result": tool_result})
         
-        return (terminated, info)
+        return (self._get_obs(), 0, terminated, info)
     
     def close(self):
         self._step_no = 0