From 59508bb57abc16a8bcecb0b55fbf901ff17845cc Mon Sep 17 00:00:00 2001 From: radu-mocanu Date: Thu, 23 Jul 2026 15:35:31 +0300 Subject: [PATCH] feat: expose managed workspace during runtime execution --- pyproject.toml | 2 +- src/uipath/runtime/__init__.py | 2 + src/uipath/runtime/errors/codes.py | 1 + src/uipath/runtime/factory.py | 3 + src/uipath/runtime/workspace/__init__.py | 2 + src/uipath/runtime/workspace/context.py | 52 +++++++++++++++ src/uipath/runtime/workspace/hydration.py | 55 ++++++++-------- src/uipath/runtime/workspace/workspace.py | 3 +- tests/workspace/test_workspace_hydration.py | 72 +++++++++++++++++++++ uv.lock | 2 +- 10 files changed, 165 insertions(+), 29 deletions(-) create mode 100644 src/uipath/runtime/workspace/context.py diff --git a/pyproject.toml b/pyproject.toml index a7b90c5f..2158ad6e 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/uipath/runtime/__init__.py b/src/uipath/runtime/__init__.py index afc27198..0dedd231 100644 --- a/src/uipath/runtime/__init__.py +++ b/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", diff --git a/src/uipath/runtime/errors/codes.py b/src/uipath/runtime/errors/codes.py index f198fdc1..460d0fa1 100644 --- a/src/uipath/runtime/errors/codes.py +++ b/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" diff --git a/src/uipath/runtime/factory.py b/src/uipath/runtime/factory.py index 166d2abe..cdfbf128 100644 --- a/src/uipath/runtime/factory.py +++ b/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, diff --git a/src/uipath/runtime/workspace/__init__.py b/src/uipath/runtime/workspace/__init__.py index c8dc23bc..cdf46206 100644 --- a/src/uipath/runtime/workspace/__init__.py +++ b/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", diff --git a/src/uipath/runtime/workspace/context.py b/src/uipath/runtime/workspace/context.py new file mode 100644 index 00000000..59abb014 --- /dev/null +++ b/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) diff --git a/src/uipath/runtime/workspace/hydration.py b/src/uipath/runtime/workspace/hydration.py index e8e979a6..ce2884db 100644 --- a/src/uipath/runtime/workspace/hydration.py +++ b/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,15 +86,16 @@ 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, @@ -101,23 +103,24 @@ async def stream( 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.""" diff --git a/src/uipath/runtime/workspace/workspace.py b/src/uipath/runtime/workspace/workspace.py index 6a1fb97f..6647b059 100644 --- a/src/uipath/runtime/workspace/workspace.py +++ b/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) diff --git a/tests/workspace/test_workspace_hydration.py b/tests/workspace/test_workspace_hydration.py index b37e1817..c4b0e9e0 100644 --- a/tests/workspace/test_workspace_hydration.py +++ b/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") diff --git a/uv.lock b/uv.lock index eccd2fe1..276fad66 100644 --- a/uv.lock +++ b/uv.lock @@ -1153,7 +1153,7 @@ wheels = [ [[package]] name = "uipath-runtime" -version = "0.12.5" +version = "0.12.6" source = { editable = "." } dependencies = [ { name = "chardet" },