feat(workspace): add ManagedWorkspace to build workspace persistence [PC-5359]

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-runtime"
-version = "0.13.6"
+version = "0.13.7"
 description = "Runtime abstractions and interfaces for building agents and automation scripts in the UiPath ecosystem"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/runtime/__init__.py

@@ -48,9 +48,11 @@
     ConversationalWorkspaceRuntime,
     HydrationPolicy,
     HydrationRuntime,
+    ManagedWorkspace,
     Workspace,
     WorkspaceHydrator,
     WorkspaceRegistryStore,
+    WorkspaceServicesProtocol,
     get_workspace_path,
 )
 
@@ -88,8 +90,10 @@
     "ConversationalWorkspaceRuntime",
     "HydrationPolicy",
     "HydrationRuntime",
+    "ManagedWorkspace",
     "get_workspace_path",
     "Workspace",
     "WorkspaceHydrator",
     "WorkspaceRegistryStore",
+    "WorkspaceServicesProtocol",
 ]

src/uipath/runtime/workspace/__init__.py

@@ -10,6 +10,10 @@
     AttachmentRegistryEntry,
     WorkspaceHydrator,
 )
+from uipath.runtime.workspace.managed import (
+    ManagedWorkspace,
+    WorkspaceServicesProtocol,
+)
 from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore
 from uipath.runtime.workspace.workspace import Workspace
 
@@ -18,8 +22,10 @@
     "ConversationalWorkspaceRuntime",
     "HydrationPolicy",
     "HydrationRuntime",
+    "ManagedWorkspace",
     "get_workspace_path",
     "Workspace",
     "WorkspaceHydrator",
     "WorkspaceRegistryStore",
+    "WorkspaceServicesProtocol",
 ]

src/uipath/runtime/workspace/hydration.py

@@ -8,6 +8,7 @@
     UiPathRuntimeProtocol,
     UiPathStreamOptions,
 )
+from uipath.runtime.debug.breakpoint import UiPathBreakpointResult
 from uipath.runtime.events import UiPathRuntimeEvent
 from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus
 from uipath.runtime.schema import UiPathRuntimeSchema
@@ -43,6 +44,7 @@ def __init__(
         registry_store: WorkspaceRegistryStore,
         hydrator_factory: None = None,
         policy: HydrationPolicy = HydrationPolicy.SUSPEND_ONLY,
+        runtime_id: str | None = None,
     ) -> None: ...
 
     @overload
@@ -55,6 +57,7 @@ def __init__(
         registry_store: WorkspaceRegistryStore,
         hydrator_factory: Callable[[], WorkspaceHydrator],
         policy: HydrationPolicy = HydrationPolicy.SUSPEND_ONLY,
+        runtime_id: str | None = None,
     ) -> None: ...
 
     def __init__(
@@ -66,8 +69,13 @@ def __init__(
         registry_store: WorkspaceRegistryStore,
         hydrator_factory: Callable[[], WorkspaceHydrator] | None = None,
         policy: HydrationPolicy = HydrationPolicy.SUSPEND_ONLY,
+        runtime_id: str | None = None,
     ) -> None:
-        """Initialize the hydration wrapper."""
+        """Initialize the hydration wrapper.
+
+        ``runtime_id`` is not used here; it is exposed for outer wrappers that
+        read the identity of the runtime they wrap.
+        """
         if (hydrator is None) == (hydrator_factory is None):
             raise ValueError("Provide exactly one of hydrator or hydrator_factory")
 
@@ -77,6 +85,8 @@ def __init__(
         self._hydrator_factory = hydrator_factory
         self.registry_store = registry_store
         self.policy = policy
+        self.runtime_id = runtime_id
+        self._paused_at_breakpoint = False
 
     def _get_hydrator(self) -> WorkspaceHydrator:
         if self.hydrator is None:
@@ -99,7 +109,7 @@ async def execute(
                 if self.policy == HydrationPolicy.ALWAYS:
                     await self._persist()
                 raise
-            await self._dehydrate(result)
+            await self._dehydrate(result, options)
             return result
 
     async def stream(
@@ -135,7 +145,7 @@ async def stream(
         else:
             if final_result is not None:
                 with _bind_workspace_execution(execution):
-                    await self._dehydrate(final_result)
+                    await self._dehydrate(final_result, options)
         finally:
             try:
                 if delegate_stream is not None:
@@ -159,12 +169,24 @@ async def dispose(self) -> None:
             await self.workspace.dispose()
 
     async def _hydrate(self) -> None:
+        if self._paused_at_breakpoint:
+            self._paused_at_breakpoint = False
+            return
         registry = await self.registry_store.load()
         hydrated = await self._get_hydrator().hydrate_from_registry(registry)
         if hydrated != registry:
             await self.registry_store.save(hydrated)
 
-    async def _dehydrate(self, result: UiPathRuntimeResult) -> None:
+    async def _dehydrate(
+        self, result: UiPathRuntimeResult, options: UiPathExecuteOptions | None
+    ) -> None:
+        if (
+            isinstance(result, UiPathBreakpointResult)
+            and options is not None
+            and options.breakpoints
+        ):
+            self._paused_at_breakpoint = True
+            return
         if self._should_dehydrate(result):
             await self._persist()
 

src/uipath/runtime/workspace/managed.py

@@ -0,0 +1,112 @@
+"""Shared wiring that persists a runtime's workspace through job attachments."""
+
+from typing import Callable, Protocol
+
+from uipath.core.workspace import AttachmentsProtocol, JobsProtocol
+
+from uipath.runtime.base import UiPathRuntimeProtocol
+from uipath.runtime.workspace.conversational import ConversationalWorkspaceRuntime
+from uipath.runtime.workspace.hydration import HydrationPolicy, HydrationRuntime
+from uipath.runtime.workspace.hydrator import WorkspaceHydrator
+from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore
+from uipath.runtime.workspace.workspace import Workspace
+
+
+class WorkspaceServicesProtocol(Protocol):
+    """Platform services the workspace hydrator talks to, such as ``UiPath()``."""
+
+    @property
+    def attachments(self) -> AttachmentsProtocol:
+        """Uploads and downloads workspace files."""
+        ...
+
+    @property
+    def jobs(self) -> JobsProtocol:
+        """Links uploaded files to the current job."""
+        ...
+
+
+class ManagedWorkspace:
+    """Builds the workspace wrappers for one runtime.
+
+    The wrappers returned by :meth:`wrap` and :meth:`wrap_conversational` share
+    one workspace, one hydrator and one registry store. The platform services
+    are created on first use: :meth:`wrap` defers them to the first run, while
+    :meth:`wrap_conversational` needs them when it is called.
+    """
+
+    def __init__(
+        self,
+        *,
+        services: Callable[[], WorkspaceServicesProtocol],
+        runtime_id: str,
+        job_key: str | None,
+        folder_key: str | None = None,
+        workspace: Workspace | None = None,
+        registry_store: WorkspaceRegistryStore | None = None,
+    ) -> None:
+        """Initialize the managed workspace.
+
+        Args:
+            services: Returns the attachments and jobs services, for example ``UiPath``.
+            runtime_id: Runtime id the wrapped runtime was created with.
+            job_key: Job that uploaded files are linked to, or ``None`` outside a job.
+            folder_key: Orchestrator folder of the attachments.
+            workspace: Folder owned by the caller. When omitted, a temporary
+                folder is created; disposing the runtime returned by :meth:`wrap`
+                removes it, so call :meth:`wrap` once per instance.
+            registry_store: Where the uploaded files are recorded. Required by
+                :meth:`wrap`; optional for :meth:`wrap_conversational`.
+        """
+        self.workspace = workspace if workspace is not None else Workspace.create()
+        self.registry_store = registry_store
+        self.runtime_id = runtime_id
+        self._services = services
+        self._job_key = job_key
+        self._folder_key = folder_key
+        self._hydrator: WorkspaceHydrator | None = None
+
+    def hydrator(self) -> WorkspaceHydrator:
+        """Return the shared hydrator, creating the platform services on first use."""
+        if self._hydrator is None:
+            services = self._services()
+            self._hydrator = WorkspaceHydrator(
+                workspace_path=self.workspace.path,
+                attachments=services.attachments,
+                jobs=services.jobs,
+                current_job_key=self._job_key,
+                folder_key=self._folder_key,
+            )
+        return self._hydrator
+
+    def wrap(
+        self,
+        runtime: UiPathRuntimeProtocol,
+        *,
+        policy: HydrationPolicy = HydrationPolicy.SUSPEND_ONLY,
+    ) -> HydrationRuntime:
+        """Restore the workspace before each run and persist it according to policy."""
+        if self.registry_store is None:
+            raise ValueError("Workspace hydration needs a registry store")
+        return HydrationRuntime(
+            runtime,
+            workspace=self.workspace,
+            hydrator_factory=self.hydrator,
+            registry_store=self.registry_store,
+            policy=policy,
+            runtime_id=self.runtime_id,
+        )
+
+    def wrap_conversational(
+        self, runtime: UiPathRuntimeProtocol
+    ) -> ConversationalWorkspaceRuntime:
+        """Carry the workspace from one conversation exchange to the next.
+
+        The returned runtime never removes the workspace folder. Without a
+        caller-owned ``workspace``, also call :meth:`wrap`, whose runtime removes it.
+        """
+        return ConversationalWorkspaceRuntime(
+            runtime,
+            hydrator=self.hydrator(),
+            registry_store=self.registry_store,
+        )