Add workspace abstraction (#2982)

* Add workspace abstraction

* Remove old workspace implementation

* Extract path resolution to a helper function

* Add api key requirements to new tests

tests/conftest.py

@@ -1,3 +1,19 @@
+from pathlib import Path
+
+import pytest
 from dotenv import load_dotenv
 
+from autogpt.workspace import Workspace
+
 load_dotenv()
+
+
+@pytest.fixture()
+def workspace_root(tmp_path) -> Path:
+    return tmp_path / "home/users/monty/auto_gpt_workspace"
+
+
+@pytest.fixture()
+def workspace(workspace_root: Path) -> Workspace:
+    workspace_root = Workspace.make_workspace(workspace_root)
+    return Workspace(workspace_root, restrict_to_workspace=True)

tests/integration/goal_oriented/test_write_file.py

@@ -6,12 +6,9 @@ import vcr
 
 from autogpt.agent import Agent
 from autogpt.commands.command import CommandRegistry
-from autogpt.commands.file_operations import LOG_FILE, delete_file, read_file
-from autogpt.config import AIConfig, Config, check_openai_api_key
+from autogpt.commands.file_operations import delete_file, read_file
+from autogpt.config import AIConfig, Config
 from autogpt.memory import get_memory
-
-# from autogpt.prompt import Prompt
-from autogpt.workspace import WORKSPACE_PATH
 from tests.integration.goal_oriented.vcr_helper import before_record_request
 from tests.utils import requires_api_key
 
@@ -28,19 +25,12 @@ CFG = Config()
 
 
 @requires_api_key("OPENAI_API_KEY")
-def test_write_file() -> None:
-    # if file exist
-    file_name = "hello_world.txt"
+def test_write_file(workspace) -> None:
+    CFG.workspace_path = workspace.root
+    CFG.file_logger_path = os.path.join(workspace.root, "file_logger.txt")
 
-    file_path_to_write_into = f"{WORKSPACE_PATH}/{file_name}"
-    if os.path.exists(file_path_to_write_into):
-        os.remove(file_path_to_write_into)
-    file_logger_path = f"{WORKSPACE_PATH}/{LOG_FILE}"
-    if os.path.exists(file_logger_path):
-        os.remove(file_logger_path)
-
-    delete_file(file_name)
-    agent = create_writer_agent()
+    file_name = str(workspace.get_path("hello_world.txt"))
+    agent = create_writer_agent(workspace)
     try:
         with my_vcr.use_cassette(
             "write_file.vcr.yml",
@@ -58,14 +48,11 @@ def test_write_file() -> None:
                     assert False, "The process took longer than 45 seconds to complete."
     # catch system exit exceptions
     except SystemExit:  # the agent returns an exception when it shuts down
-        content = ""
         content = read_file(file_name)
-        os.remove(file_path_to_write_into)
-
         assert content == "Hello World", f"Expected 'Hello World', got {content}"
 
 
-def create_writer_agent():
+def create_writer_agent(workspace):
     command_registry = CommandRegistry()
     command_registry.import_commands("autogpt.commands.file_operations")
     command_registry.import_commands("autogpt.app")
@@ -96,6 +83,7 @@ def create_writer_agent():
         next_action_count=0,
         system_prompt=system_prompt,
         triggering_prompt=triggering_prompt,
+        workspace_directory=workspace.root,
     )
     CFG.set_continuous_mode(True)
     CFG.set_memory_backend("no_memory")

tests/test_image_gen.py

@@ -1,39 +1,46 @@
 import hashlib
 import os
+import shutil
 import unittest
+from pathlib import Path
 
 from PIL import Image
 
 from autogpt.commands.image_gen import generate_image, generate_image_with_sd_webui
 from autogpt.config import Config
-from autogpt.workspace import path_in_workspace
+from autogpt.workspace import Workspace
 from tests.utils import requires_api_key
 
 
 def lst(txt):
-    return txt.split(":")[1].strip()
+    return Path(txt.split(":")[1].strip())
 
 
-@unittest.skipIf(os.getenv("CI"), "Skipping image generation tests")
+@unittest.skip("Skipping image generation tests")
 class TestImageGen(unittest.TestCase):
     def setUp(self):
         self.config = Config()
+        workspace_path = os.path.join(os.path.dirname(__file__), "workspace")
+        self.workspace_path = Workspace.make_workspace(workspace_path)
+        self.config.workspace_path = workspace_path
+        self.workspace = Workspace(workspace_path, restrict_to_workspace=True)
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.workspace_path)
 
     @requires_api_key("OPENAI_API_KEY")
     def test_dalle(self):
         self.config.image_provider = "dalle"
 
         # Test using size 256
-        result = lst(generate_image("astronaut riding a horse", 256))
-        image_path = path_in_workspace(result)
+        image_path = lst(generate_image("astronaut riding a horse", 256))
         self.assertTrue(image_path.exists())
         with Image.open(image_path) as img:
             self.assertEqual(img.size, (256, 256))
         image_path.unlink()
 
         # Test using size 512
-        result = lst(generate_image("astronaut riding a horse", 512))
-        image_path = path_in_workspace(result)
+        image_path = lst(generate_image("astronaut riding a horse", 512))
         with Image.open(image_path) as img:
             self.assertEqual(img.size, (512, 512))
         image_path.unlink()
@@ -44,8 +51,7 @@ class TestImageGen(unittest.TestCase):
 
         # Test usin SD 1.4 model and size 512
         self.config.huggingface_image_model = "CompVis/stable-diffusion-v1-4"
-        result = lst(generate_image("astronaut riding a horse", 512))
-        image_path = path_in_workspace(result)
+        image_path = lst(generate_image("astronaut riding a horse", 512))
         self.assertTrue(image_path.exists())
         with Image.open(image_path) as img:
             self.assertEqual(img.size, (512, 512))
@@ -53,8 +59,7 @@ class TestImageGen(unittest.TestCase):
 
         # Test using SD 2.1 768 model and size 768
         self.config.huggingface_image_model = "stabilityai/stable-diffusion-2-1"
-        result = lst(generate_image("astronaut riding a horse", 768))
-        image_path = path_in_workspace(result)
+        image_path = lst(generate_image("astronaut riding a horse", 768))
         with Image.open(image_path) as img:
             self.assertEqual(img.size, (768, 768))
         image_path.unlink()
@@ -64,8 +69,7 @@ class TestImageGen(unittest.TestCase):
         return
 
         # Test using size 128
-        result = lst(generate_image_with_sd_webui("astronaut riding a horse", 128))
-        image_path = path_in_workspace(result)
+        image_path = lst(generate_image_with_sd_webui("astronaut riding a horse", 128))
         self.assertTrue(image_path.exists())
         with Image.open(image_path) as img:
             self.assertEqual(img.size, (128, 128))

tests/test_workspace.py

@@ -0,0 +1,86 @@
+from pathlib import Path
+
+import pytest
+
+from autogpt.workspace import Workspace
+
+_WORKSPACE_ROOT = Path("home/users/monty/auto_gpt_workspace")
+
+_ACCESSIBLE_PATHS = [
+    Path("."),
+    Path("test_file.txt"),
+    Path("test_folder"),
+    Path("test_folder/test_file.txt"),
+    Path("test_folder/.."),
+    Path("test_folder/../test_file.txt"),
+    Path("test_folder/../test_folder"),
+    Path("test_folder/../test_folder/test_file.txt"),
+]
+
+_INACCESSIBLE_PATHS = [
+    # Takes us out of the workspace
+    Path(".."),
+    Path("../test_file.txt"),
+    Path("../not_auto_gpt_workspace"),
+    Path("../not_auto_gpt_workspace/test_file.txt"),
+    Path("test_folder/../.."),
+    Path("test_folder/../../test_file.txt"),
+    Path("test_folder/../../not_auto_gpt_workspace"),
+    Path("test_folder/../../not_auto_gpt_workspace/test_file.txt"),
+    # Contains null bytes
+    Path("\x00"),
+    Path("\x00test_file.txt"),
+    Path("test_folder/\x00"),
+    Path("test_folder/\x00test_file.txt"),
+    # Absolute paths
+    Path("/"),
+    Path("/test_file.txt"),
+    Path("/home"),
+]
+
+
+@pytest.fixture()
+def workspace_root(tmp_path):
+    return tmp_path / _WORKSPACE_ROOT
+
+
+@pytest.fixture(params=_ACCESSIBLE_PATHS)
+def accessible_path(request):
+    return request.param
+
+
+@pytest.fixture(params=_INACCESSIBLE_PATHS)
+def inaccessible_path(request):
+    return request.param
+
+
+def test_sanitize_path_accessible(accessible_path, workspace_root):
+    full_path = Workspace._sanitize_path(
+        accessible_path,
+        root=workspace_root,
+        restrict_to_root=True,
+    )
+    assert full_path.is_absolute()
+    assert full_path.is_relative_to(workspace_root)
+
+
+def test_sanitize_path_inaccessible(inaccessible_path, workspace_root):
+    with pytest.raises(ValueError):
+        Workspace._sanitize_path(
+            inaccessible_path,
+            root=workspace_root,
+            restrict_to_root=True,
+        )
+
+
+def test_get_path_accessible(accessible_path, workspace_root):
+    workspace = Workspace(workspace_root, True)
+    full_path = workspace.get_path(accessible_path)
+    assert full_path.is_absolute()
+    assert full_path.is_relative_to(workspace_root)
+
+
+def test_get_path_inaccessible(inaccessible_path, workspace_root):
+    workspace = Workspace(workspace_root, True)
+    with pytest.raises(ValueError):
+        workspace.get_path(inaccessible_path)

tests/unit/test_file_operations.py

@@ -4,7 +4,6 @@ import unittest
 from pathlib import Path
 
 from autogpt.commands.file_operations import (
-    LOG_FILE_PATH,
     append_to_file,
     check_duplicate_operation,
     delete_file,
@@ -15,7 +14,7 @@ from autogpt.commands.file_operations import (
     write_to_file,
 )
 from autogpt.config import Config
-from autogpt.workspace import path_in_workspace
+from autogpt.workspace import Workspace
 
 
 class TestFileOperations(unittest.TestCase):
@@ -24,24 +23,24 @@ class TestFileOperations(unittest.TestCase):
     """
 
     def setUp(self):
-        self.test_file = "test_file.txt"
+        self.config = Config()
+        workspace_path = os.path.join(os.path.dirname(__file__), "workspace")
+        self.workspace_path = Workspace.make_workspace(workspace_path)
+        self.config.workspace_path = workspace_path
+        self.config.file_logger_path = os.path.join(workspace_path, "file_logger.txt")
+        self.workspace = Workspace(workspace_path, restrict_to_workspace=True)
+
+        self.test_file = str(self.workspace.get_path("test_file.txt"))
         self.test_file2 = "test_file2.txt"
-        self.test_directory = "test_directory"
+        self.test_directory = str(self.workspace.get_path("test_directory"))
         self.file_content = "This is a test file.\n"
         self.file_logger_logs = "file_logger.txt"
 
-        with open(path_in_workspace(self.test_file), "w") as f:
+        with open(self.test_file, "w") as f:
             f.write(self.file_content)
 
-        if os.path.exists(LOG_FILE_PATH):
-            os.remove(LOG_FILE_PATH)
-
-    def tearDown(self):
-        if os.path.exists(path_in_workspace(self.test_file)):
-            os.remove(path_in_workspace(self.test_file))
-
-        if os.path.exists(self.test_directory):
-            shutil.rmtree(self.test_directory)
+    def tearDown(self) -> None:
+        shutil.rmtree(self.workspace_path)
 
     def test_check_duplicate_operation(self):
         log_operation("write", self.test_file)
@@ -53,9 +52,9 @@ class TestFileOperations(unittest.TestCase):
             os.remove(self.file_logger_logs)
 
         log_operation("log_test", self.test_file)
-        with open(LOG_FILE_PATH, "r") as f:
+        with open(self.config.file_logger_path, "r") as f:
             content = f.read()
-        self.assertIn("log_test: test_file.txt", content)
+        self.assertIn(f"log_test: {self.test_file}", content)
 
     # Test splitting a file into chunks
     def test_split_file(self):
@@ -71,80 +70,59 @@ class TestFileOperations(unittest.TestCase):
     def test_write_to_file(self):
         new_content = "This is new content.\n"
         write_to_file(self.test_file, new_content)
-        with open(path_in_workspace(self.test_file), "r") as f:
+        with open(self.test_file, "r") as f:
             content = f.read()
         self.assertEqual(content, new_content)
 
     def test_append_to_file(self):
-        with open(path_in_workspace(self.test_file), "r") as f:
+        with open(self.test_file, "r") as f:
             content_before = f.read()
 
         append_text = "This is appended text.\n"
         append_to_file(self.test_file, append_text)
-        with open(path_in_workspace(self.test_file), "r") as f:
+        with open(self.test_file, "r") as f:
             content = f.read()
 
         self.assertEqual(content, content_before + append_text)
 
     def test_delete_file(self):
         delete_file(self.test_file)
-        self.assertFalse(os.path.exists(path_in_workspace(self.test_file)))
+        self.assertFalse(os.path.exists(self.test_file))
 
     def test_search_files(self):
         # Case 1: Create files A and B, search for A, and ensure we don't return A and B
-        file_a = "file_a.txt"
-        file_b = "file_b.txt"
+        file_a = self.workspace.get_path("file_a.txt")
+        file_b = self.workspace.get_path("file_b.txt")
 
-        with open(path_in_workspace(file_a), "w") as f:
+        with open(file_a, "w") as f:
             f.write("This is file A.")
 
-        with open(path_in_workspace(file_b), "w") as f:
+        with open(file_b, "w") as f:
             f.write("This is file B.")
 
         # Create a subdirectory and place a copy of file_a in it
-        if not os.path.exists(path_in_workspace(self.test_directory)):
-            os.makedirs(path_in_workspace(self.test_directory))
+        if not os.path.exists(self.test_directory):
+            os.makedirs(self.test_directory)
 
-        with open(
-            path_in_workspace(os.path.join(self.test_directory, file_a)), "w"
-        ) as f:
+        with open(os.path.join(self.test_directory, file_a.name), "w") as f:
             f.write("This is file A in the subdirectory.")
 
-        files = search_files(path_in_workspace(""))
-        self.assertIn(file_a, files)
-        self.assertIn(file_b, files)
-        self.assertIn(os.path.join(self.test_directory, file_a), files)
+        files = search_files(str(self.workspace.root))
+        self.assertIn(file_a.name, files)
+        self.assertIn(file_b.name, files)
+        self.assertIn(f"{Path(self.test_directory).name}/{file_a.name}", files)
 
         # Clean up
-        os.remove(path_in_workspace(file_a))
-        os.remove(path_in_workspace(file_b))
-        os.remove(path_in_workspace(os.path.join(self.test_directory, file_a)))
-        os.rmdir(path_in_workspace(self.test_directory))
+        os.remove(file_a)
+        os.remove(file_b)
+        os.remove(os.path.join(self.test_directory, file_a.name))
+        os.rmdir(self.test_directory)
 
         # Case 2: Search for a file that does not exist and make sure we don't throw
         non_existent_file = "non_existent_file.txt"
         files = search_files("")
         self.assertNotIn(non_existent_file, files)
 
-    # Test to ensure we cannot read files out of workspace
-    def test_restrict_workspace(self):
-        CFG = Config()
-        with open(self.test_file2, "w+") as f:
-            f.write("test text")
-
-        CFG.restrict_to_workspace = True
-
-        # Get the absolute path of self.test_file2
-        test_file2_abs_path = os.path.abspath(self.test_file2)
-
-        with self.assertRaises(ValueError):
-            read_file(test_file2_abs_path)
-
-        CFG.restrict_to_workspace = False
-        read_file(test_file2_abs_path)
-
-        os.remove(test_file2_abs_path)
-
 
 if __name__ == "__main__":
     unittest.main()

tests/unit/test_setup.py

@@ -35,6 +35,7 @@ class TestAutoGPT(unittest.TestCase):
         self.assertGreaterEqual(len(ai_config.ai_goals), 1)
         self.assertLessEqual(len(ai_config.ai_goals), 5)
 
+    @requires_api_key("OPENAI_API_KEY")
     def test_generate_aiconfig_automatic_fallback(self):
         user_inputs = [
             "T&GF£OIBECC()!*",
@@ -52,6 +53,7 @@ class TestAutoGPT(unittest.TestCase):
         self.assertEqual(ai_config.ai_role, "an AI designed to browse bake a cake.")
         self.assertEqual(ai_config.ai_goals, ["Purchase ingredients", "Bake a cake"])
 
+    @requires_api_key("OPENAI_API_KEY")
     def test_prompt_user_manual_mode(self):
         user_inputs = [
             "--manual",

tests/utils.py

@@ -1,3 +1,4 @@
+import functools
 import os
 
 import pytest
@@ -5,6 +6,7 @@ import pytest
 
 def requires_api_key(env_var):
     def decorator(func):
+        @functools.wraps(func)
         def wrapper(*args, **kwargs):
             if not os.environ.get(env_var):
                 pytest.skip(