fix: reuse storage across runs

Author: cosmyo

packages/uipath-openai-agents/src/uipath_openai_agents/runtime/factory.py

@@ -44,14 +44,50 @@ def __init__(
         self._agent_loaders: dict[str, OpenAiAgentLoader] = {}
         self._agent_lock = asyncio.Lock()
 
+        self._storage: SqliteAgentStorage | None = None
+        self._storage_lock = asyncio.Lock()
+
         self._setup_instrumentation()
 
     def _setup_instrumentation(self) -> None:
         """Setup tracing and instrumentation."""
         OpenAIAgentsInstrumentor().instrument()
         UiPathSpanUtils.register_current_span_provider(get_current_span_wrapper)
 
-    def _get_storage_path(self, runtime_id: str) -> str | None:
+    async def _get_or_create_storage(self) -> SqliteAgentStorage | None:
+        """Get or create the shared storage instance.
+
+        Returns:
+            Shared storage instance, or None if storage is disabled
+        """
+        async with self._storage_lock:
+            if self._storage is None:
+                storage_path = self._get_storage_path()
+                if storage_path:
+                    self._storage = SqliteAgentStorage(storage_path)
+                    await self._storage.setup()
+            return self._storage
+
+    def _get_storage_path(self) -> str | None:
+        """Get the storage path for agent state.
+
+        Returns:
+            Path to SQLite database for storage, or None if storage is disabled
+        """
+        if self.context.runtime_dir and self.context.state_file:
+            path = os.path.join(self.context.runtime_dir, self.context.state_file)
+
+            if not self.context.resume and self.context.job_id is None:
+                # If not resuming and no job id, delete the previous state file
+                if os.path.exists(path):
+                    os.remove(path)
+
+            os.makedirs(self.context.runtime_dir, exist_ok=True)
+            return path
+
+        return None
+
+    def _get_storage_path_legacy(self, runtime_id: str) -> str | None:
         """
         Get the storage path for agent session state.
 
@@ -230,13 +266,9 @@ async def _create_runtime_instance(
         Returns:
             Configured runtime instance
         """
-        storage_path = self._get_storage_path(runtime_id)
-
-        # Create storage instance if storage path is available
-        storage: SqliteAgentStorage | None = None
-        if storage_path:
-            storage = SqliteAgentStorage(storage_path)
-            await storage.setup()
+        # Get shared storage instance
+        storage = await self._get_or_create_storage()
+        storage_path = storage.storage_path if storage else None
 
         # Get the loaded object from the agent loader for schema inference
         loaded_object = None
@@ -282,3 +314,8 @@ async def dispose(self) -> None:
 
         self._agent_loaders.clear()
         self._agent_cache.clear()
+
+        # Dispose shared storage
+        if self._storage:
+            await self._storage.dispose()
+            self._storage = None

packages/uipath-openai-agents/src/uipath_openai_agents/runtime/runtime.py

@@ -619,12 +619,6 @@ async def get_schema(self) -> UiPathRuntimeSchema:
 
     async def dispose(self) -> None:
         """Cleanup runtime resources."""
-        # Clean up storage first (before event loop closes)
-        if self.storage:
-            try:
-                await self.storage.dispose()
-            except Exception:
-                pass  # Best effort cleanup
-
-        self._session = None
-        self.storage = None
+        # Storage is shared across runtimes and managed by the factory
+        # Do not dispose it here
+        pass