diff --git a/CODEOWNERS b/CODEOWNERS index 03c07761..8410903c 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -1 +1 @@ -* @UiPath/team-coded-agents +* @UiPath/agents @cristipufu diff --git a/pyproject.toml b/pyproject.toml index 2158ad6e..a0eb682c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-runtime" -version = "0.12.6" +version = "0.12.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" diff --git a/src/uipath/runtime/workspace/context.py b/src/uipath/runtime/workspace/context.py index 59abb014..6e554095 100644 --- a/src/uipath/runtime/workspace/context.py +++ b/src/uipath/runtime/workspace/context.py @@ -3,6 +3,7 @@ from collections.abc import Iterator from contextlib import contextmanager from contextvars import ContextVar +from dataclasses import dataclass from pathlib import Path from uipath.runtime.errors import ( @@ -11,8 +12,17 @@ UiPathRuntimeError, ) -_workspace_path: ContextVar[Path | None] = ContextVar( - "uipath_workspace_path", default=None + +@dataclass +class _WorkspaceExecution: + """Revocable workspace state shared with tasks created during execution.""" + + path: Path + active: bool = True + + +_workspace_execution_state: ContextVar[_WorkspaceExecution | None] = ContextVar( + "uipath_workspace_execution", default=None ) @@ -27,26 +37,49 @@ def get_workspace_path() -> Path: Raises: UiPathRuntimeError: If called outside a managed runtime execution. """ - workspace_path = _workspace_path.get() - if workspace_path is None: + execution = _workspace_execution_state.get() + if execution is None or not execution.active: 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." + "No managed workspace is available in the current code path. " + "Call get_workspace_path() only from code invoked by a managed runtime, " + "not from module-level initialization or from the caller that invokes " + "or consumes the execution." ), category=UiPathErrorCategory.USER, include_traceback=False, ) - return workspace_path + return execution.path @contextmanager -def _workspace_execution(path: Path) -> Iterator[None]: - """Make ``path`` available for one managed runtime execution.""" - token = _workspace_path.set(path) +def _bind_workspace_execution(execution: _WorkspaceExecution) -> Iterator[None]: + """Make an active execution available in the current context.""" + token = _workspace_execution_state.set(execution) try: yield finally: - _workspace_path.reset(token) + _workspace_execution_state.reset(token) + + +def _create_workspace_execution(path: Path) -> _WorkspaceExecution: + """Create state for a workspace execution.""" + return _WorkspaceExecution(path) + + +def _revoke_workspace_execution(execution: _WorkspaceExecution) -> None: + """Prevent all contexts that inherited ``execution`` from accessing it.""" + execution.active = False + + +@contextmanager +def _workspace_execution(path: Path) -> Iterator[None]: + """Make ``path`` available for one non-streaming execution.""" + execution = _create_workspace_execution(path) + try: + with _bind_workspace_execution(execution): + yield + finally: + _revoke_workspace_execution(execution) diff --git a/src/uipath/runtime/workspace/hydration.py b/src/uipath/runtime/workspace/hydration.py index ce2884db..b3fa1298 100644 --- a/src/uipath/runtime/workspace/hydration.py +++ b/src/uipath/runtime/workspace/hydration.py @@ -11,7 +11,12 @@ 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.context import ( + _bind_workspace_execution, + _create_workspace_execution, + _revoke_workspace_execution, + _workspace_execution, +) from uipath.runtime.workspace.hydrator import WorkspaceHydrator from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore from uipath.runtime.workspace.workspace import Workspace @@ -103,24 +108,44 @@ async def stream( options: UiPathStreamOptions | None = None, ) -> AsyncGenerator[UiPathRuntimeEvent, None]: """Hydrate, stream delegate events, then persist files according to policy.""" - with _workspace_execution(self.workspace.path): - await self._hydrate() - final_result: UiPathRuntimeResult | None = None + execution = _create_workspace_execution(self.workspace.path) + delegate_stream: AsyncGenerator[UiPathRuntimeEvent, None] | None = None + final_result: UiPathRuntimeResult | None = None + hydrated = False - 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: + try: + with _bind_workspace_execution(execution): + await self._hydrate() + hydrated = True + delegate_stream = self.delegate.stream(input, options=options) + + event = await anext(delegate_stream, None) + while event is not None: + if isinstance(event, UiPathRuntimeResult): + final_result = event + else: + yield event + with _bind_workspace_execution(execution): + event = await anext(delegate_stream, None) + except Exception: + if self.policy == HydrationPolicy.ALWAYS and hydrated: + with _bind_workspace_execution(execution): await self._persist() - raise - + raise + else: if final_result is not None: - await self._dehydrate(final_result) - yield final_result + with _bind_workspace_execution(execution): + await self._dehydrate(final_result) + finally: + try: + if delegate_stream is not None: + with _bind_workspace_execution(execution): + await delegate_stream.aclose() + finally: + _revoke_workspace_execution(execution) + + if final_result is not None: + 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 6647b059..ffd12a3e 100644 --- a/src/uipath/runtime/workspace/workspace.py +++ b/src/uipath/runtime/workspace/workspace.py @@ -44,5 +44,7 @@ def create( async def dispose(self) -> None: """Remove the workspace directory when configured to do so.""" - if self.cleanup and self.path.exists(): - await asyncio.to_thread(shutil.rmtree, self.path) + if not self.cleanup: + return + + 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 c4b0e9e0..82567e2b 100644 --- a/tests/workspace/test_workspace_hydration.py +++ b/tests/workspace/test_workspace_hydration.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import shutil import uuid from pathlib import Path @@ -72,6 +73,18 @@ async def download_async( return name +class FailingDownloadAttachments(FakeAttachments): + async def download_async( + self, + *, + key: uuid.UUID, + destination_path: str, + folder_key: str | None = None, + folder_path: str | None = None, + ) -> str: + raise RuntimeError("download failed") + + class FakeJobs: def __init__(self) -> None: self.attachments: dict[str, list[str]] = {} @@ -160,6 +173,68 @@ async def stream( yield event +class ChildTaskRuntime(WritingRuntime): + def __init__(self, workspace_path: Path, release: asyncio.Event) -> None: + super().__init__(workspace_path, UiPathRuntimeStatus.SUCCESSFUL) + self.release = release + self.task: asyncio.Task[Path] | None = None + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + async def access_workspace_after_execution() -> Path: + await self.release.wait() + return get_workspace_path() + + self.task = asyncio.create_task(access_workspace_after_execution()) + return await super().execute(input, options) + + +class StreamingChildTaskRuntime(WritingRuntime): + def __init__(self, workspace_path: Path, release: asyncio.Event) -> None: + super().__init__(workspace_path, UiPathRuntimeStatus.SUCCESSFUL) + self.release = release + self.task: asyncio.Task[Path] | None = None + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + async def access_workspace_during_stream() -> Path: + await self.release.wait() + return get_workspace_path() + + self.task = asyncio.create_task(access_workspace_during_stream()) + yield UiPathRuntimeStateEvent(payload={"started": True}) + await self.release.wait() + assert self.task is not None + assert await self.task == self.workspace_path + yield UiPathRuntimeResult(status=self.status, output={"ok": True}) + + +class FinalChildTaskRuntime(WritingRuntime): + def __init__(self, workspace_path: Path, release: asyncio.Event) -> None: + super().__init__(workspace_path, UiPathRuntimeStatus.SUCCESSFUL) + self.release = release + self.task: asyncio.Task[Path] | None = None + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + async def access_workspace_after_final_result() -> Path: + await self.release.wait() + return get_workspace_path() + + self.task = asyncio.create_task(access_workspace_after_final_result()) + yield UiPathRuntimeStateEvent(payload={"started": True}) + yield UiPathRuntimeResult(status=self.status, output={"ok": True}) + + @pytest.mark.asyncio async def test_dehydrate_uploads_changed_files_and_saves_registry( tmp_path: Path, @@ -198,6 +273,7 @@ def test_get_workspace_path_requires_a_managed_execution() -> None: get_workspace_path() assert error.value.error_info.code == "Python.MANAGED_WORKSPACE_UNAVAILABLE" + assert "not from module-level initialization" in str(error.value) assert error.value.error_info.category == UiPathErrorCategory.USER @@ -244,6 +320,221 @@ async def test_workspace_path_is_available_during_streaming( get_workspace_path() +@pytest.mark.asyncio +async def test_stream_does_not_expose_workspace_to_the_consumer( + 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"), + ) + + stream = runtime.stream({}) + assert isinstance(await anext(stream), UiPathRuntimeStateEvent) + with pytest.raises(UiPathRuntimeError, match="No managed workspace"): + get_workspace_path() + + await stream.aclose() + with pytest.raises(UiPathRuntimeError, match="No managed workspace"): + get_workspace_path() + + +@pytest.mark.asyncio +async def test_breaking_after_the_final_stream_result_resets_workspace_access( + 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"), + ) + + stream = runtime.stream({}) + async for event in stream: + if isinstance(event, UiPathRuntimeResult): + break + + await stream.aclose() + + with pytest.raises(UiPathRuntimeError, match="No managed workspace"): + get_workspace_path() + + +@pytest.mark.asyncio +async def test_stream_supports_task_wrapped_iterator_resumes( + 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"), + ) + + stream = runtime.stream({}) + first = await asyncio.wait_for(anext(stream), timeout=1) + result = await asyncio.wait_for(anext(stream), timeout=1) + + assert isinstance(first, UiPathRuntimeStateEvent) + assert isinstance(result, UiPathRuntimeResult) + await stream.aclose() + + +@pytest.mark.asyncio +async def test_stream_child_tasks_keep_workspace_access_between_events( + tmp_path: Path, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + release = asyncio.Event() + delegate = StreamingChildTaskRuntime(workspace.path, release) + runtime = HydrationRuntime( + delegate, + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=FakeAttachments(), + ), + registry_store=WorkspaceRegistryStore(MemoryStorage(), "runtime-1"), + ) + + stream = runtime.stream({}) + assert isinstance(await anext(stream), UiPathRuntimeStateEvent) + release.set() + assert isinstance(await anext(stream), UiPathRuntimeResult) + await stream.aclose() + + +@pytest.mark.asyncio +async def test_stream_does_not_retry_failed_final_persistence( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + runtime = HydrationRuntime( + WritingRuntime(workspace.path, UiPathRuntimeStatus.SUCCESSFUL), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=FakeAttachments(), + ), + registry_store=WorkspaceRegistryStore(MemoryStorage(), "runtime-1"), + policy=HydrationPolicy.ALWAYS, + ) + persist_calls = 0 + + async def fail_persist() -> None: + nonlocal persist_calls + persist_calls += 1 + raise RuntimeError("persist failed") + + monkeypatch.setattr(runtime, "_persist", fail_persist) + + with pytest.raises(RuntimeError, match="persist failed"): + [event async for event in runtime.stream({})] + + assert persist_calls == 1 + + +@pytest.mark.asyncio +async def test_stream_hydration_failure_preserves_the_workspace_registry( + tmp_path: Path, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + storage = MemoryStorage() + store = WorkspaceRegistryStore(storage, "runtime-1") + registry = { + "notes.txt": { + "attachment_key": str(uuid.uuid4()), + "sha256": "0" * 64, + "size": 1, + "uploaded_at": "2026-01-01T00:00:00+00:00", + } + } + await store.save(registry) + runtime = HydrationRuntime( + WritingRuntime(workspace.path, UiPathRuntimeStatus.SUCCESSFUL), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=FailingDownloadAttachments(), + ), + registry_store=store, + policy=HydrationPolicy.ALWAYS, + ) + + with pytest.raises(RuntimeError, match="download failed"): + [event async for event in runtime.stream({})] + + assert await store.load() == registry + + +@pytest.mark.asyncio +async def test_stream_revokes_workspace_before_yielding_final_result( + tmp_path: Path, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + release = asyncio.Event() + delegate = FinalChildTaskRuntime(workspace.path, release) + runtime = HydrationRuntime( + delegate, + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=FakeAttachments(), + ), + registry_store=WorkspaceRegistryStore(MemoryStorage(), "runtime-1"), + ) + + stream = runtime.stream({}) + assert isinstance(await anext(stream), UiPathRuntimeStateEvent) + assert isinstance(await anext(stream), UiPathRuntimeResult) + release.set() + + assert delegate.task is not None + with pytest.raises(UiPathRuntimeError, match="No managed workspace"): + await delegate.task + + +@pytest.mark.asyncio +async def test_child_tasks_cannot_access_workspace_after_execution( + tmp_path: Path, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + release = asyncio.Event() + delegate = ChildTaskRuntime(workspace.path, release) + runtime = HydrationRuntime( + delegate, + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=FakeAttachments(), + ), + registry_store=WorkspaceRegistryStore(MemoryStorage(), "runtime-1"), + ) + + await runtime.execute({}) + release.set() + + assert delegate.task is not None + with pytest.raises(UiPathRuntimeError, match="No managed workspace"): + await delegate.task + + @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 276fad66..ec486156 100644 --- a/uv.lock +++ b/uv.lock @@ -1153,7 +1153,7 @@ wheels = [ [[package]] name = "uipath-runtime" -version = "0.12.6" +version = "0.12.7" source = { editable = "." } dependencies = [ { name = "chardet" },