refactor skyvern current state json file (#2489)

skyvern/forge/sdk/artifact/storage/base.py

@@ -94,7 +94,7 @@ class BaseStorage(ABC):
         pass
 
     @abstractmethod
-    async def get_streaming_file(self, organization_id: str, file_name: str, use_default: bool = True) -> bytes | None:
+    async def get_streaming_file(self, organization_id: str, file_name: str) -> bytes | None:
         pass
 
     @abstractmethod

skyvern/forge/sdk/artifact/storage/local.py

@@ -116,10 +116,10 @@ class LocalStorage(BaseStorage):
     async def save_streaming_file(self, organization_id: str, file_name: str) -> None:
         return
 
-    async def get_streaming_file(self, organization_id: str, file_name: str, use_default: bool = True) -> bytes | None:
+    async def get_streaming_file(self, organization_id: str, file_name: str) -> bytes | None:
-        file_path = Path(f"{get_skyvern_temp_dir()}/skyvern_screenshot.png")
+        # make the directory if it doesn't exist
-        if not use_default:
+        Path(f"{get_skyvern_temp_dir()}/{organization_id}").mkdir(parents=True, exist_ok=True)
-            file_path = Path(f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}")
+        file_path = Path(f"{get_skyvern_temp_dir()}/{organization_id}/{file_name}")
         try:
             with open(file_path, "rb") as f:
                 return f.read()

skyvern/utils/files.py

@@ -0,0 +1,29 @@
+import json
+import os
+
+import aiofiles
+
+from skyvern.forge.sdk.api.files import get_skyvern_temp_dir
+
+
+def get_skyvern_state_file_path() -> str:
+    return f"{get_skyvern_temp_dir()}/current.json"
+
+
+async def initialize_skyvern_state_file(
+    task_id: str | None = None, workflow_run_id: str | None = None, organization_id: str | None = None
+) -> None:
+    # create the file if it doesn't exist
+    async with aiofiles.open(get_skyvern_state_file_path(), "w") as json_file:
+        await json_file.write(
+            json.dumps({"task_id": task_id, "workflow_run_id": workflow_run_id, "organization_id": organization_id})
+        )
+
+
+def get_json_from_file(file_path: str) -> dict[str, str]:
+    # check if file exists
+    if not os.path.exists(file_path):
+        return {}
+
+    with open(file_path, "r") as json_file:
+        return json.load(json_file)