feat: expose managed workspace during runtime execution

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-runtime"
-version = "0.12.5"
+version = "0.12.6"
 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

@@ -50,6 +50,7 @@
     Workspace,
     WorkspaceHydrator,
     WorkspaceRegistryStore,
+    get_workspace_path,
 )
 
 __all__ = [
@@ -84,6 +85,7 @@
     "AttachmentRegistryEntry",
     "HydrationPolicy",
     "HydrationRuntime",
+    "get_workspace_path",
     "Workspace",
     "WorkspaceHydrator",
     "WorkspaceRegistryStore",

src/uipath/runtime/errors/codes.py

@@ -16,6 +16,7 @@ class UiPathErrorCode(str, Enum):
     MODULE_EXECUTION_ERROR = "MODULE_EXECUTION_ERROR"
     FUNCTION_EXECUTION_ERROR = "FUNCTION_EXECUTION_ERROR"
     EXECUTION_ERROR = "EXECUTION_ERROR"
+    MANAGED_WORKSPACE_UNAVAILABLE = "MANAGED_WORKSPACE_UNAVAILABLE"
 
     # Input validation errors
     INVALID_INPUT_FILE_EXTENSION = "INVALID_INPUT_FILE_EXTENSION"

src/uipath/runtime/factory.py

@@ -23,6 +23,9 @@ class UiPathRuntimeFactorySettings(BaseModel):
 
     agent_framework: str | None = None
 
+    managed_workspace: bool = False
+    """Whether this runtime can use host-managed persistent workspace storage."""
+
 
 class UiPathRuntimeFactoryProtocol(
     UiPathDisposableProtocol,

src/uipath/runtime/workspace/__init__.py

@@ -1,5 +1,6 @@
 """Workspace persistence primitives for runtime implementations."""
 
+from uipath.runtime.workspace.context import get_workspace_path
 from uipath.runtime.workspace.hydration import (
     HydrationPolicy,
     HydrationRuntime,
@@ -15,6 +16,7 @@
     "AttachmentRegistryEntry",
     "HydrationPolicy",
     "HydrationRuntime",
+    "get_workspace_path",
     "Workspace",
     "WorkspaceHydrator",
     "WorkspaceRegistryStore",

src/uipath/runtime/workspace/context.py

@@ -0,0 +1,52 @@
+"""Execution-scoped access to a managed workspace."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from pathlib import Path
+
+from uipath.runtime.errors import (
+    UiPathErrorCategory,
+    UiPathErrorCode,
+    UiPathRuntimeError,
+)
+
+_workspace_path: ContextVar[Path | None] = ContextVar(
+    "uipath_workspace_path", default=None
+)
+
+
+def get_workspace_path() -> Path:
+    """Return the workspace available to the current runtime execution.
+
+    A workspace is available only while a runtime managed by
+    :class:`HydrationRuntime` is executing. Files created beneath this path are
+    restored before a resumed execution and persisted when the runtime
+    suspends, according to its hydration policy.
+
+    Raises:
+        UiPathRuntimeError: If called outside a managed runtime execution.
+    """
+    workspace_path = _workspace_path.get()
+    if workspace_path is None:
+        raise UiPathRuntimeError(
+            code=UiPathErrorCode.MANAGED_WORKSPACE_UNAVAILABLE,
+            title="Managed Workspace Unavailable",
+            detail=(
+                "No managed workspace is available outside a runtime execution. "
+                "Call get_workspace_path() only from a graph node or tool."
+            ),
+            category=UiPathErrorCategory.USER,
+            include_traceback=False,
+        )
+    return workspace_path
+
+
+@contextmanager
+def _workspace_execution(path: Path) -> Iterator[None]:
+    """Make ``path`` available for one managed runtime execution."""
+    token = _workspace_path.set(path)
+    try:
+        yield
+    finally:
+        _workspace_path.reset(token)

src/uipath/runtime/workspace/hydration.py

@@ -11,6 +11,7 @@
 from uipath.runtime.events import UiPathRuntimeEvent
 from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus
 from uipath.runtime.schema import UiPathRuntimeSchema
+from uipath.runtime.workspace.context import _workspace_execution
 from uipath.runtime.workspace.hydrator import WorkspaceHydrator
 from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore
 from uipath.runtime.workspace.workspace import Workspace
@@ -85,39 +86,41 @@ async def execute(
         options: UiPathExecuteOptions | None = None,
     ) -> UiPathRuntimeResult:
         """Hydrate, execute, then persist files according to policy."""
-        await self._hydrate()
-        try:
-            result = await self.delegate.execute(input, options=options)
-        except Exception:
-            if self.policy == HydrationPolicy.ALWAYS:
-                await self._persist()
-            raise
-        await self._dehydrate(result)
-        return result
+        with _workspace_execution(self.workspace.path):
+            await self._hydrate()
+            try:
+                result = await self.delegate.execute(input, options=options)
+            except Exception:
+                if self.policy == HydrationPolicy.ALWAYS:
+                    await self._persist()
+                raise
+            await self._dehydrate(result)
+            return result
 
     async def stream(
         self,
         input: dict[str, Any] | None = None,
         options: UiPathStreamOptions | None = None,
     ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
         """Hydrate, stream delegate events, then persist files according to policy."""
-        await self._hydrate()
-        final_result: UiPathRuntimeResult | None = None
-
-        try:
-            async for event in self.delegate.stream(input, options=options):
-                if isinstance(event, UiPathRuntimeResult):
-                    final_result = event
-                else:
-                    yield event
-        except Exception:
-            if self.policy == HydrationPolicy.ALWAYS:
-                await self._persist()
-            raise
-
-        if final_result is not None:
-            await self._dehydrate(final_result)
-            yield final_result
+        with _workspace_execution(self.workspace.path):
+            await self._hydrate()
+            final_result: UiPathRuntimeResult | None = None
+
+            try:
+                async for event in self.delegate.stream(input, options=options):
+                    if isinstance(event, UiPathRuntimeResult):
+                        final_result = event
+                    else:
+                        yield event
+            except Exception:
+                if self.policy == HydrationPolicy.ALWAYS:
+                    await self._persist()
+                raise
+
+            if final_result is not None:
+                await self._dehydrate(final_result)
+                yield final_result
 
     async def get_schema(self) -> UiPathRuntimeSchema:
         """Passthrough schema from delegate runtime."""

src/uipath/runtime/workspace/workspace.py

@@ -1,5 +1,6 @@
 """Disk workspace lifecycle helpers."""
 
+import asyncio
 import shutil
 import tempfile
 from pathlib import Path
@@ -44,4 +45,4 @@ def create(
     async def dispose(self) -> None:
         """Remove the workspace directory when configured to do so."""
         if self.cleanup and self.path.exists():
-            shutil.rmtree(self.path)
+            await asyncio.to_thread(shutil.rmtree, self.path)

tests/workspace/test_workspace_hydration.py

@@ -17,7 +17,9 @@
     Workspace,
     WorkspaceHydrator,
     WorkspaceRegistryStore,
+    get_workspace_path,
 )
+from uipath.runtime.errors import UiPathErrorCategory, UiPathRuntimeError
 from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeStateEvent
 from uipath.runtime.schema import UiPathRuntimeSchema
 
@@ -139,6 +141,25 @@ async def get_schema(self) -> UiPathRuntimeSchema:
         )
 
 
+class ContextAwareRuntime(WritingRuntime):
+    async def execute(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathExecuteOptions | None = None,
+    ) -> UiPathRuntimeResult:
+        assert get_workspace_path() == self.workspace_path
+        return await super().execute(input, options)
+
+    async def stream(
+        self,
+        input: dict[str, Any] | None = None,
+        options: UiPathStreamOptions | None = None,
+    ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
+        assert get_workspace_path() == self.workspace_path
+        async for event in super().stream(input, options):
+            yield event
+
+
 @pytest.mark.asyncio
 async def test_dehydrate_uploads_changed_files_and_saves_registry(
     tmp_path: Path,
@@ -172,6 +193,57 @@ async def test_dehydrate_uploads_changed_files_and_saves_registry(
     assert len(jobs.links) == 1
 
 
+def test_get_workspace_path_requires_a_managed_execution() -> None:
+    with pytest.raises(UiPathRuntimeError, match="No managed workspace") as error:
+        get_workspace_path()
+
+    assert error.value.error_info.code == "Python.MANAGED_WORKSPACE_UNAVAILABLE"
+    assert error.value.error_info.category == UiPathErrorCategory.USER
+
+
+@pytest.mark.asyncio
+async def test_workspace_path_is_available_only_during_execution(
+    tmp_path: Path,
+) -> None:
+    workspace = Workspace.create(tmp_path / "workspace")
+    runtime = HydrationRuntime(
+        ContextAwareRuntime(workspace.path, UiPathRuntimeStatus.SUCCESSFUL),
+        workspace=workspace,
+        hydrator=WorkspaceHydrator(
+            workspace_path=workspace.path,
+            attachments=FakeAttachments(),
+        ),
+        registry_store=WorkspaceRegistryStore(MemoryStorage(), "runtime-1"),
+    )
+
+    await runtime.execute({})
+
+    with pytest.raises(UiPathRuntimeError, match="No managed workspace"):
+        get_workspace_path()
+
+
+@pytest.mark.asyncio
+async def test_workspace_path_is_available_during_streaming(
+    tmp_path: Path,
+) -> None:
+    workspace = Workspace.create(tmp_path / "workspace")
+    runtime = HydrationRuntime(
+        ContextAwareRuntime(workspace.path, UiPathRuntimeStatus.SUCCESSFUL),
+        workspace=workspace,
+        hydrator=WorkspaceHydrator(
+            workspace_path=workspace.path,
+            attachments=FakeAttachments(),
+        ),
+        registry_store=WorkspaceRegistryStore(MemoryStorage(), "runtime-1"),
+    )
+
+    events = [event async for event in runtime.stream({})]
+
+    assert isinstance(events[-1], UiPathRuntimeResult)
+    with pytest.raises(UiPathRuntimeError, match="No managed workspace"):
+        get_workspace_path()
+
+
 @pytest.mark.asyncio
 async def test_hydrator_factory_is_deferred_and_cached(tmp_path: Path) -> None:
     workspace = Workspace.create(tmp_path / "workspace")