From e853dbf77a54e3eb9a55e0b99ecf27c40002a608 Mon Sep 17 00:00:00 2001 From: radu-mocanu Date: Tue, 23 Jun 2026 19:26:44 +0300 Subject: [PATCH] feat: add workspace hydration primitives --- pyproject.toml | 4 +- src/uipath/runtime/__init__.py | 14 + src/uipath/runtime/workspace/__init__.py | 21 + src/uipath/runtime/workspace/hydration.py | 120 +++++ src/uipath/runtime/workspace/hydrator.py | 190 +++++++ .../runtime/workspace/registry_store.py | 49 ++ src/uipath/runtime/workspace/workspace.py | 47 ++ tests/workspace/test_workspace_hydration.py | 483 ++++++++++++++++++ uv.lock | 12 +- 9 files changed, 932 insertions(+), 8 deletions(-) create mode 100644 src/uipath/runtime/workspace/__init__.py create mode 100644 src/uipath/runtime/workspace/hydration.py create mode 100644 src/uipath/runtime/workspace/hydrator.py create mode 100644 src/uipath/runtime/workspace/registry_store.py create mode 100644 src/uipath/runtime/workspace/workspace.py create mode 100644 tests/workspace/test_workspace_hydration.py diff --git a/pyproject.toml b/pyproject.toml index a7b626c..c1427ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,11 +1,11 @@ [project] name = "uipath-runtime" -version = "0.11.2" +version = "0.11.3" 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" dependencies = [ - "uipath-core>=0.5.17, <0.6.0" + "uipath-core>=0.5.22,<0.6.0", ] classifiers = [ "Intended Audience :: Developers", diff --git a/src/uipath/runtime/__init__.py b/src/uipath/runtime/__init__.py index 1eed011..afc2719 100644 --- a/src/uipath/runtime/__init__.py +++ b/src/uipath/runtime/__init__.py @@ -43,6 +43,14 @@ ) from uipath.runtime.schema import UiPathRuntimeSchema from uipath.runtime.storage import UiPathRuntimeStorageProtocol +from uipath.runtime.workspace import ( + AttachmentRegistryEntry, + HydrationPolicy, + HydrationRuntime, + Workspace, + WorkspaceHydrator, + WorkspaceRegistryStore, +) __all__ = [ "UiPathExecuteOptions", @@ -73,4 +81,10 @@ "UiPathResumeTriggerName", "UiPathChatProtocol", "UiPathChatRuntime", + "AttachmentRegistryEntry", + "HydrationPolicy", + "HydrationRuntime", + "Workspace", + "WorkspaceHydrator", + "WorkspaceRegistryStore", ] diff --git a/src/uipath/runtime/workspace/__init__.py b/src/uipath/runtime/workspace/__init__.py new file mode 100644 index 0000000..c8dc23b --- /dev/null +++ b/src/uipath/runtime/workspace/__init__.py @@ -0,0 +1,21 @@ +"""Workspace persistence primitives for runtime implementations.""" + +from uipath.runtime.workspace.hydration import ( + HydrationPolicy, + HydrationRuntime, +) +from uipath.runtime.workspace.hydrator import ( + AttachmentRegistryEntry, + WorkspaceHydrator, +) +from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore +from uipath.runtime.workspace.workspace import Workspace + +__all__ = [ + "AttachmentRegistryEntry", + "HydrationPolicy", + "HydrationRuntime", + "Workspace", + "WorkspaceHydrator", + "WorkspaceRegistryStore", +] diff --git a/src/uipath/runtime/workspace/hydration.py b/src/uipath/runtime/workspace/hydration.py new file mode 100644 index 0000000..2b3d8e1 --- /dev/null +++ b/src/uipath/runtime/workspace/hydration.py @@ -0,0 +1,120 @@ +"""Runtime wrapper for workspace hydration.""" + +from enum import Enum +from typing import Any, AsyncGenerator + +from uipath.runtime.base import ( + UiPathExecuteOptions, + UiPathRuntimeProtocol, + UiPathStreamOptions, +) +from uipath.runtime.events import UiPathRuntimeEvent +from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus +from uipath.runtime.schema import UiPathRuntimeSchema +from uipath.runtime.workspace.hydrator import WorkspaceHydrator +from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore +from uipath.runtime.workspace.workspace import Workspace + + +class HydrationPolicy(str, Enum): + """Controls when workspace changes are persisted.""" + + SUSPEND_ONLY = "suspend_only" + SUSPEND_OR_SUCCESS = "suspend_or_success" + ALWAYS = "always" + + +class HydrationRuntime: + """Wraps a runtime with hydrate-before and dehydrate-after behavior.""" + + def __init__( + self, + delegate: UiPathRuntimeProtocol, + *, + workspace: Workspace, + hydrator: WorkspaceHydrator, + registry_store: WorkspaceRegistryStore, + policy: HydrationPolicy = HydrationPolicy.SUSPEND_ONLY, + ): + """Initialize the hydration wrapper.""" + self.delegate = delegate + self.workspace = workspace + self.hydrator = hydrator + self.registry_store = registry_store + self.policy = policy + + async def execute( + self, + input: dict[str, Any] | None = None, + 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 + + 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 + + async def get_schema(self) -> UiPathRuntimeSchema: + """Passthrough schema from delegate runtime.""" + return await self.delegate.get_schema() + + async def dispose(self) -> None: + """Dispose delegate and workspace.""" + try: + await self.delegate.dispose() + finally: + await self.workspace.dispose() + + async def _hydrate(self) -> None: + registry = await self.registry_store.load() + hydrated = await self.hydrator.hydrate(registry) + if hydrated != registry: + await self.registry_store.save(hydrated) + + async def _dehydrate(self, result: UiPathRuntimeResult) -> None: + if self._should_dehydrate(result): + await self._persist() + + async def _persist(self) -> None: + registry = await self.registry_store.load() + dehydrated = await self.hydrator.dehydrate(registry) + await self.registry_store.save(dehydrated) + + def _should_dehydrate(self, result: UiPathRuntimeResult) -> bool: + if self.policy == HydrationPolicy.ALWAYS: + return True + if result.status == UiPathRuntimeStatus.SUSPENDED: + return True + return ( + self.policy == HydrationPolicy.SUSPEND_OR_SUCCESS + and result.status == UiPathRuntimeStatus.SUCCESSFUL + ) diff --git a/src/uipath/runtime/workspace/hydrator.py b/src/uipath/runtime/workspace/hydrator.py new file mode 100644 index 0000000..fba298e --- /dev/null +++ b/src/uipath/runtime/workspace/hydrator.py @@ -0,0 +1,190 @@ +"""Attachment-backed workspace hydration.""" + +import hashlib +import os +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from uuid import UUID + +from pydantic import BaseModel +from uipath.core.workspace import AttachmentsProtocol, JobsProtocol + + +class AttachmentRegistryEntry(BaseModel): + """Registry entry for one workspace file attachment.""" + + attachment_key: str + sha256: str + size: int + uploaded_at: str + attachment_name: str | None = None + + +class WorkspaceHydrator: + """Hydrates/dehydrates a workspace directory through job attachments.""" + + DEFAULT_ATTACHMENT_PREFIX = ".uipath-workspace/" + + def __init__( + self, + *, + workspace_path: str | Path, + attachments: AttachmentsProtocol, + jobs: JobsProtocol | None = None, + current_job_key: str | None = None, + folder_key: str | None = None, + folder_path: str | None = None, + attachment_prefix: str = DEFAULT_ATTACHMENT_PREFIX, + ): + """Initialize the hydrator.""" + self.workspace_path = Path(workspace_path) + self.attachments = attachments + self.jobs = jobs + self.current_job_key = current_job_key + self.folder_key = folder_key + self.folder_path = folder_path + self.attachment_prefix = attachment_prefix.strip("/") + + async def hydrate( + self, + registry: dict[str, dict[str, Any]], + ) -> dict[str, dict[str, Any]]: + """Download registry files into the workspace. + + Files with matching SHA-256 are left untouched. + """ + normalized = self._normalize_registry(registry) + for virtual_path, entry in normalized.items(): + target = self._resolve_workspace_path(virtual_path) + if target.exists() and self._sha256(target) == entry.sha256: + continue + target.parent.mkdir(parents=True, exist_ok=True) + await self.attachments.download_async( + key=UUID(entry.attachment_key), + destination_path=str(target), + folder_key=self.folder_key, + folder_path=self.folder_path, + ) + return self._dump_registry(normalized) + + async def dehydrate( + self, + registry: dict[str, dict[str, Any]], + ) -> dict[str, dict[str, Any]]: + """Upload new/changed workspace files and return the merged registry. + + The result is rebuilt from the files currently present, so entries for + files deleted locally are dropped and not restored on the next hydrate. + """ + previous = self._normalize_registry(registry) + current: dict[str, AttachmentRegistryEntry] = {} + for file_path in self._iter_files(): + virtual_path = self._virtual_path(file_path) + digest = self._sha256(file_path) + size = file_path.stat().st_size + existing = previous.get(virtual_path) + + if existing and existing.sha256 == digest and existing.size == size: + current[virtual_path] = existing + if self.current_job_key: + await self.link_attachment(existing.attachment_key) + continue + + attachment_name = self._attachment_name_for_virtual_path(virtual_path) + attachment_key = await self.attachments.upload_async( + name=attachment_name, + source_path=str(file_path), + folder_key=self.folder_key, + folder_path=self.folder_path, + ) + entry = AttachmentRegistryEntry( + attachment_key=str(attachment_key), + sha256=digest, + size=size, + uploaded_at=datetime.now(timezone.utc).isoformat(), + attachment_name=attachment_name, + ) + current[virtual_path] = entry + + if self.current_job_key: + await self.link_attachment(entry.attachment_key) + + return self._dump_registry(current) + + async def link_attachment(self, attachment_key: str) -> None: + """Link an already uploaded attachment to the current job.""" + if self.jobs is None or not self.current_job_key: + return + + await self.jobs.link_attachment_async( + job_key=UUID(self.current_job_key), + attachment_key=UUID(attachment_key), + folder_key=self.folder_key, + folder_path=self.folder_path, + ) + + def _iter_files(self) -> list[Path]: + if not self.workspace_path.exists(): + return [] + + files: list[Path] = [] + for root, _, names in os.walk(self.workspace_path): + for name in names: + path = Path(root) / name + if path.is_file(): + files.append(path) + return sorted(files) + + def _resolve_workspace_path(self, virtual_path: str) -> Path: + path = (self.workspace_path / virtual_path).resolve() + workspace = self.workspace_path.resolve() + if workspace != path and workspace not in path.parents: + raise ValueError(f"Workspace path escapes root: {virtual_path}") + return path + + def _virtual_path(self, path: Path) -> str: + return path.relative_to(self.workspace_path).as_posix() + + @staticmethod + def _escape(value: str) -> str: + # json-pointer escaping; escape ~ before / so they can't collide + return value.replace("~", "~0").replace("/", "~1") + + @staticmethod + def _unescape(value: str) -> str: + return value.replace("~1", "/").replace("~0", "~") + + def _attachment_name_for_virtual_path(self, virtual_path: str) -> str: + # the platform mishandles "/" in attachment names, so keep it slash-free + return self._escape(f"{self.attachment_prefix}/{virtual_path}") + + def _virtual_path_from_attachment_name(self, name: str) -> str | None: + decoded = self._unescape(name) + prefix = f"{self.attachment_prefix}/" + if not decoded.startswith(prefix): + return None + virtual_path = decoded[len(prefix) :] + self._resolve_workspace_path(virtual_path) + return virtual_path + + def _normalize_registry( + self, registry: dict[str, dict[str, Any]] + ) -> dict[str, AttachmentRegistryEntry]: + normalized: dict[str, AttachmentRegistryEntry] = {} + for virtual_path, entry in registry.items(): + self._resolve_workspace_path(virtual_path) + normalized[virtual_path] = AttachmentRegistryEntry.model_validate(entry) + return normalized + + def _dump_registry( + self, registry: dict[str, AttachmentRegistryEntry] + ) -> dict[str, dict[str, Any]]: + return {path: entry.model_dump() for path, entry in registry.items()} + + def _sha256(self, path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/src/uipath/runtime/workspace/registry_store.py b/src/uipath/runtime/workspace/registry_store.py new file mode 100644 index 0000000..d6e1eb4 --- /dev/null +++ b/src/uipath/runtime/workspace/registry_store.py @@ -0,0 +1,49 @@ +"""Storage-backed workspace attachment registry.""" + +from typing import Any + +from uipath.runtime.storage import UiPathRuntimeStorageProtocol + + +class WorkspaceRegistryStore: + """Stores workspace file attachment metadata in runtime storage.""" + + DEFAULT_NAMESPACE = "workspace" + DEFAULT_KEY = "attachments" + + def __init__( + self, + storage: UiPathRuntimeStorageProtocol, + runtime_id: str, + *, + namespace: str = DEFAULT_NAMESPACE, + key: str = DEFAULT_KEY, + ): + """Initialize the registry store.""" + self.storage = storage + self.runtime_id = runtime_id + self.namespace = namespace + self.key = key + + async def load(self) -> dict[str, dict[str, Any]]: + """Load registry entries keyed by workspace-relative path.""" + value = await self.storage.get_value(self.runtime_id, self.namespace, self.key) + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError("Workspace registry payload must be a dictionary.") + + registry: dict[str, dict[str, Any]] = {} + for path, entry in value.items(): + if isinstance(path, str) and isinstance(entry, dict): + registry[path] = entry + return registry + + async def save(self, registry: dict[str, dict[str, Any]]) -> None: + """Persist registry entries.""" + await self.storage.set_value( + self.runtime_id, + self.namespace, + self.key, + registry, + ) diff --git a/src/uipath/runtime/workspace/workspace.py b/src/uipath/runtime/workspace/workspace.py new file mode 100644 index 0000000..6a1fb97 --- /dev/null +++ b/src/uipath/runtime/workspace/workspace.py @@ -0,0 +1,47 @@ +"""Disk workspace lifecycle helpers.""" + +import shutil +import tempfile +from pathlib import Path + + +class Workspace: + """Owns a directory used by a runtime to persist files between executions.""" + + def __init__(self, path: Path, *, cleanup: bool = True): + """Initialize the workspace. + + Args: + path: Workspace directory. + cleanup: Whether dispose should remove the directory. + """ + self.path = path + self.cleanup = cleanup + + @classmethod + def create( + cls, + path: str | Path | None = None, + *, + prefix: str = "uipath-workspace-", + cleanup: bool | None = None, + ) -> "Workspace": + """Create a workspace directory. + + ``cleanup`` defaults to True for an owned temp dir and False for a + supplied path; pass ``cleanup=True`` for a temp path you want disposed. + """ + if path is None: + workspace_path = Path(tempfile.mkdtemp(prefix=prefix)) + should_cleanup = True if cleanup is None else cleanup + else: + workspace_path = Path(path) + workspace_path.mkdir(parents=True, exist_ok=True) + should_cleanup = False if cleanup is None else cleanup + + return cls(workspace_path, cleanup=should_cleanup) + + 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) diff --git a/tests/workspace/test_workspace_hydration.py b/tests/workspace/test_workspace_hydration.py new file mode 100644 index 0000000..c6a2b29 --- /dev/null +++ b/tests/workspace/test_workspace_hydration.py @@ -0,0 +1,483 @@ +from __future__ import annotations + +import shutil +import uuid +from pathlib import Path +from typing import Any, AsyncGenerator + +import pytest + +from uipath.runtime import ( + HydrationPolicy, + HydrationRuntime, + UiPathExecuteOptions, + UiPathRuntimeResult, + UiPathRuntimeStatus, + UiPathStreamOptions, + Workspace, + WorkspaceHydrator, + WorkspaceRegistryStore, +) +from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeStateEvent +from uipath.runtime.schema import UiPathRuntimeSchema + + +class MemoryStorage: + def __init__(self) -> None: + self.values: dict[tuple[str, str, str], Any] = {} + + async def set_value( + self, runtime_id: str, namespace: str, key: str, value: Any + ) -> None: + self.values[(runtime_id, namespace, key)] = value + + async def get_value(self, runtime_id: str, namespace: str, key: str) -> Any: + return self.values.get((runtime_id, namespace, key)) + + +class FakeAttachments: + def __init__(self) -> None: + self.files: dict[uuid.UUID, tuple[str, bytes]] = {} + self.uploads = 0 + self.downloads = 0 + + async def upload_async( + self, + *, + name: str, + content: str | bytes | None = None, + source_path: str | None = None, + folder_key: str | None = None, + folder_path: str | None = None, + ) -> uuid.UUID: + assert source_path is not None + key = uuid.uuid4() + self.files[key] = (name, Path(source_path).read_bytes()) + self.uploads += 1 + return key + + async def download_async( + self, + *, + key: uuid.UUID, + destination_path: str, + folder_key: str | None = None, + folder_path: str | None = None, + ) -> str: + name, content = self.files[key] + Path(destination_path).write_bytes(content) + self.downloads += 1 + return name + + +class FakeJobs: + def __init__(self) -> None: + self.attachments: dict[str, list[str]] = {} + self.links: list[tuple[uuid.UUID, uuid.UUID]] = [] + + async def list_attachments_async( + self, + *, + job_key: uuid.UUID, + folder_key: str | None = None, + folder_path: str | None = None, + ) -> list[str]: + return self.attachments.get(str(job_key), []) + + async def link_attachment_async( + self, + *, + job_key: uuid.UUID, + attachment_key: uuid.UUID, + folder_key: str | None = None, + folder_path: str | None = None, + ) -> None: + self.links.append((job_key, attachment_key)) + + +class WritingRuntime: + def __init__(self, workspace_path: Path, status: UiPathRuntimeStatus) -> None: + self.workspace_path = workspace_path + self.status = status + self.disposed = False + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + (self.workspace_path / "notes.txt").write_text("hello", encoding="utf-8") + return UiPathRuntimeResult(status=self.status, output={"ok": True}) + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + yield UiPathRuntimeStateEvent(payload={"started": True}) + result = await self.execute(input, options) + yield result + + async def get_schema(self) -> UiPathRuntimeSchema: + raise NotImplementedError + + async def dispose(self) -> None: + self.disposed = True + + +@pytest.mark.asyncio +async def test_dehydrate_uploads_changed_files_and_saves_registry( + tmp_path: Path, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + jobs = FakeJobs() + current_job = uuid.uuid4() + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + jobs=jobs, + current_job_key=str(current_job), + ) + storage = MemoryStorage() + store = WorkspaceRegistryStore(storage, "runtime-1") + runtime = HydrationRuntime( + WritingRuntime(workspace.path, UiPathRuntimeStatus.SUSPENDED), + workspace=workspace, + hydrator=hydrator, + registry_store=store, + ) + + result = await runtime.execute({}) + + registry = await store.load() + assert result.status == UiPathRuntimeStatus.SUSPENDED + assert list(registry) == ["notes.txt"] + assert registry["notes.txt"]["attachment_name"] == ".uipath-workspace~1notes.txt" + assert attachments.uploads == 1 + assert len(jobs.links) == 1 + + +@pytest.mark.asyncio +async def test_successful_completion_persists_when_policy_allows( + tmp_path: Path, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + storage = MemoryStorage() + runtime = HydrationRuntime( + WritingRuntime(workspace.path, UiPathRuntimeStatus.SUCCESSFUL), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ), + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + policy=HydrationPolicy.SUSPEND_OR_SUCCESS, + ) + + await runtime.execute({}) + + assert attachments.uploads == 1 + assert "notes.txt" in await runtime.registry_store.load() + + +@pytest.mark.asyncio +async def test_hydrate_downloads_missing_registry_files(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + key = uuid.uuid4() + attachments.files[key] = (".uipath-workspace/notes.txt", b"from attachment") + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ) + registry = { + "notes.txt": { + "attachment_key": str(key), + "sha256": "different", + "size": 15, + "uploaded_at": "2026-01-01T00:00:00+00:00", + "attachment_name": ".uipath-workspace/notes.txt", + } + } + + await hydrator.hydrate(registry) + + assert (workspace.path / "notes.txt").read_text( + encoding="utf-8" + ) == "from attachment" + assert attachments.downloads == 1 + + +@pytest.mark.asyncio +async def test_stream_persists_on_suspend(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + storage = MemoryStorage() + runtime = HydrationRuntime( + WritingRuntime(workspace.path, UiPathRuntimeStatus.SUSPENDED), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ), + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + ) + + events = [event async for event in runtime.stream({})] + + assert isinstance(events[-1], UiPathRuntimeResult) + assert attachments.uploads == 1 + assert "notes.txt" in await runtime.registry_store.load() + + +@pytest.mark.asyncio +async def test_hydrate_skips_unchanged_local_file(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + (workspace.path / "notes.txt").write_text("same", encoding="utf-8") + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ) + digest = hydrator._sha256(workspace.path / "notes.txt") + registry = { + "notes.txt": { + "attachment_key": str(uuid.uuid4()), + "sha256": digest, + "size": 4, + "uploaded_at": "", + "attachment_name": ".uipath-workspace~1notes.txt", + } + } + + await hydrator.hydrate(registry) + + assert attachments.downloads == 0 + + +@pytest.mark.asyncio +async def test_workspace_dispose_removes_temp_dir(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace", cleanup=True) + (workspace.path / "file.txt").write_text("x", encoding="utf-8") + + await workspace.dispose() + + assert not workspace.path.exists() + + +@pytest.mark.asyncio +async def test_workspace_dispose_keeps_owned_path_by_default(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace") + (workspace.path / "file.txt").write_text("x", encoding="utf-8") + + await workspace.dispose() + + assert workspace.path.exists() + shutil.rmtree(workspace.path) + + +@pytest.mark.asyncio +async def test_create_temp_workspace_is_cleaned_up_by_default() -> None: + workspace = Workspace.create() + + assert workspace.path.exists() + + await workspace.dispose() + + assert not workspace.path.exists() + + +@pytest.mark.asyncio +async def test_dehydrate_relinks_unchanged_file_without_reupload( + tmp_path: Path, +) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + jobs = FakeJobs() + current_job = uuid.uuid4() + key = uuid.uuid4() + (workspace.path / "notes.txt").write_text("same", encoding="utf-8") + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + jobs=jobs, + current_job_key=str(current_job), + ) + prior = { + "notes.txt": { + "attachment_key": str(key), + "sha256": hydrator._sha256(workspace.path / "notes.txt"), + "size": 4, + "uploaded_at": "", + "attachment_name": ".uipath-workspace~1notes.txt", + } + } + + result = await hydrator.dehydrate(prior) + + assert attachments.uploads == 0 + assert result["notes.txt"]["attachment_key"] == str(key) + assert jobs.links == [(current_job, key)] + + +@pytest.mark.asyncio +async def test_attachment_names_are_single_segment_for_nested_files( + tmp_path: Path, +) -> None: + """Attachment names must stay slash-free (a "/" breaks the blob round-trip).""" + workspace = Workspace.create(tmp_path / "workspace") + (workspace.path / "plan").mkdir(parents=True, exist_ok=True) + (workspace.path / "plan" / "todo.md").write_text("step 1", encoding="utf-8") + attachments = FakeAttachments() + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ) + + registry = await hydrator.dehydrate({}) + + assert "plan/todo.md" in registry + attachment_name = registry["plan/todo.md"]["attachment_name"] + assert "/" not in attachment_name + assert all("/" not in name for name, _ in attachments.files.values()) + assert ( + hydrator._virtual_path_from_attachment_name(attachment_name) == "plan/todo.md" + ) + + +@pytest.mark.asyncio +async def test_attachment_name_round_trips_special_characters(tmp_path: Path) -> None: + """The encoding must be reversible even for paths containing the escape char.""" + workspace = Workspace.create(tmp_path / "workspace") + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=FakeAttachments(), + ) + + for virtual_path in ["a~~b.txt", "plan/todo.md", "we~ird/a~~b/file.txt"]: + name = hydrator._attachment_name_for_virtual_path(virtual_path) + assert "/" not in name + assert hydrator._virtual_path_from_attachment_name(name) == virtual_path + + +@pytest.mark.asyncio +async def test_dehydrate_drops_files_deleted_locally(tmp_path: Path) -> None: + """A registry entry whose file no longer exists is not carried forward.""" + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ) + (workspace.path / "present.txt").write_text("hi", encoding="utf-8") + prior = { + "gone.txt": { + "attachment_key": str(uuid.uuid4()), + "sha256": "stale", + "size": 1, + "uploaded_at": "2026-01-01T00:00:00+00:00", + "attachment_name": ".uipath-workspace~1gone.txt", + } + } + + result = await hydrator.dehydrate(prior) + + assert "present.txt" in result + assert "gone.txt" not in result + + +class _WriteThenRaiseRuntime: + def __init__(self, workspace_path: Path) -> None: + self.workspace_path = workspace_path + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + (self.workspace_path / "partial.txt").write_text("wip", encoding="utf-8") + raise RuntimeError("boom") + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + await self.execute(input, options) + yield UiPathRuntimeStateEvent(payload={}) + + async def get_schema(self) -> UiPathRuntimeSchema: + raise NotImplementedError + + async def dispose(self) -> None: + pass + + +@pytest.mark.asyncio +async def test_always_policy_persists_on_failure(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + storage = MemoryStorage() + runtime = HydrationRuntime( + _WriteThenRaiseRuntime(workspace.path), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ), + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + policy=HydrationPolicy.ALWAYS, + ) + + with pytest.raises(RuntimeError): + await runtime.execute({}) + + assert attachments.uploads == 1 + assert "partial.txt" in await runtime.registry_store.load() + + +@pytest.mark.asyncio +async def test_suspend_only_policy_skips_persist_on_failure(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + storage = MemoryStorage() + runtime = HydrationRuntime( + _WriteThenRaiseRuntime(workspace.path), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ), + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + policy=HydrationPolicy.SUSPEND_ONLY, + ) + + with pytest.raises(RuntimeError): + await runtime.execute({}) + + assert attachments.uploads == 0 + + +@pytest.mark.asyncio +async def test_always_policy_persists_on_stream_failure(tmp_path: Path) -> None: + workspace = Workspace.create(tmp_path / "workspace") + attachments = FakeAttachments() + storage = MemoryStorage() + runtime = HydrationRuntime( + _WriteThenRaiseRuntime(workspace.path), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + ), + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + policy=HydrationPolicy.ALWAYS, + ) + + with pytest.raises(RuntimeError): + async for _ in runtime.stream({}): + pass + + assert attachments.uploads == 1 diff --git a/uv.lock b/uv.lock index f05b1f2..6b5443c 100644 --- a/uv.lock +++ b/uv.lock @@ -3,7 +3,7 @@ revision = 3 requires-python = ">=3.11" [options] -exclude-newer = "2026-06-17T19:38:11.87942Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P2D" [options.exclude-newer-package] @@ -998,21 +998,21 @@ wheels = [ [[package]] name = "uipath-core" -version = "0.5.17" +version = "0.5.22" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e3/80/a626eb3136a6765e0af06c9d5080ac0843c2a72f17b7a2170f1f45da40dd/uipath_core-0.5.17.tar.gz", hash = "sha256:13565e1eba9f059a8221494dfb3239257ddf7f265fc7057199ffe03ed066300a", size = 119023, upload-time = "2026-05-28T21:34:10.903Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e8/e0/1cdf0537ae1db831b066604e0e83132a2dd559371ac6e5d56e96b9039163/uipath_core-0.5.22.tar.gz", hash = "sha256:01ae7c3770369469acf5cef31908e8b878a5b1123f2d930f8537ea2d97d7d621", size = 136212, upload-time = "2026-06-23T16:18:43.081Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/74/cf/f4b481970621e2a9aec869302773fa2c7d346aef294a553429626369633f/uipath_core-0.5.17-py3-none-any.whl", hash = "sha256:6e088eec5130bc492ac176ab85d4924d7d4cb07ee290ed7e6a46984e9de8c12b", size = 44957, upload-time = "2026-05-28T21:34:09.534Z" }, + { url = "https://files.pythonhosted.org/packages/fb/97/2258d51969ec71b1056d67f39d612eac2d7c6e9458d3b3c9a0b10f42e730/uipath_core-0.5.22-py3-none-any.whl", hash = "sha256:60df655b207e02a6d3bfae8c61e1fc9bc0bf11576f7ead07b8b38f23d13fc4d6", size = 58222, upload-time = "2026-06-23T16:18:41.536Z" }, ] [[package]] name = "uipath-runtime" -version = "0.11.2" +version = "0.11.3" source = { editable = "." } dependencies = [ { name = "uipath-core" }, @@ -1034,7 +1034,7 @@ dev = [ ] [package.metadata] -requires-dist = [{ name = "uipath-core", specifier = ">=0.5.17,<0.6.0" }] +requires-dist = [{ name = "uipath-core", specifier = ">=0.5.22,<0.6.0" }] [package.metadata.requires-dev] dev = [