From dd6e9e6adebb9ac7f1f0f69d14699e113afc5901 Mon Sep 17 00:00:00 2001 From: tudormatei Date: Wed, 4 Nov 2026 18:33:50 +0200 Subject: [PATCH] feat(workspace): add ManagedWorkspace to build workspace persistence Coded agents lose the files they write when a job suspends and resumes. The wiring that builds the hydrator and the hydration and conversational wrappers lived only in uipath-agents. ManagedWorkspace builds them around one hydrator, so the uipath CLI and uipath-agents can share it. HydrationRuntime takes an optional runtime_id for outer wrappers, so uipath-agents can drop the subclass that only added that attribute. A debugger pause no longer uploads the workspace, and the call that continues skips the restore, so edits made before the pause stay. This applies only when the call asked for breakpoints; any other suspend uploads as before. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 2 +- src/uipath/runtime/__init__.py | 4 + src/uipath/runtime/workspace/__init__.py | 6 + src/uipath/runtime/workspace/hydration.py | 30 ++- src/uipath/runtime/workspace/managed.py | 112 ++++++++++ tests/workspace/test_managed_workspace.py | 181 +++++++++++++++++ tests/workspace/test_workspace_hydration.py | 214 +++++++++++++++++++- uv.lock | 2 +- 8 files changed, 544 insertions(+), 7 deletions(-) create mode 100644 src/uipath/runtime/workspace/managed.py create mode 100644 tests/workspace/test_managed_workspace.py diff --git a/pyproject.toml b/pyproject.toml index a6db49a..5cb6c0c 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/uipath/runtime/__init__.py b/src/uipath/runtime/__init__.py index 8807ae0..a11b3cb 100644 --- a/src/uipath/runtime/__init__.py +++ b/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", ] diff --git a/src/uipath/runtime/workspace/__init__.py b/src/uipath/runtime/workspace/__init__.py index 03e8e97..4d653f5 100644 --- a/src/uipath/runtime/workspace/__init__.py +++ b/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", ] diff --git a/src/uipath/runtime/workspace/hydration.py b/src/uipath/runtime/workspace/hydration.py index d1c0480..fdeba6c 100644 --- a/src/uipath/runtime/workspace/hydration.py +++ b/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() diff --git a/src/uipath/runtime/workspace/managed.py b/src/uipath/runtime/workspace/managed.py new file mode 100644 index 0000000..de76fc8 --- /dev/null +++ b/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, + ) diff --git a/tests/workspace/test_managed_workspace.py b/tests/workspace/test_managed_workspace.py new file mode 100644 index 0000000..0a8396c --- /dev/null +++ b/tests/workspace/test_managed_workspace.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import uuid +from pathlib import Path +from typing import Any, AsyncGenerator + +import pytest + +from tests.workspace.test_workspace_hydration import ( + FakeAttachments, + FakeJobs, + MemoryStorage, +) +from uipath.runtime import ( + ManagedWorkspace, + UiPathExecuteOptions, + UiPathRuntimeResult, + UiPathRuntimeStatus, + UiPathStreamOptions, + Workspace, + WorkspaceRegistryStore, + get_workspace_path, +) +from uipath.runtime.events import ( + UiPathRuntimeConversationMetaEvent, + UiPathRuntimeEvent, +) +from uipath.runtime.schema import UiPathRuntimeSchema + +JOB_KEY = str(uuid.uuid4()) + + +class FakeServices: + def __init__(self) -> None: + self.attachments = FakeAttachments() + self.jobs = FakeJobs() + self.built = 0 + + def __call__(self) -> FakeServices: + self.built += 1 + return self + + +class ReportRuntime: + """Writes or reads ``report.md`` in the managed workspace.""" + + def __init__( + self, + *, + write: str | None = None, + status: UiPathRuntimeStatus = UiPathRuntimeStatus.SUCCESSFUL, + ) -> None: + self.write = write + self.status = status + self.workspace_path: Path | None = None + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + self.workspace_path = get_workspace_path() + report = self.workspace_path / "report.md" + if self.write is not None: + report.write_text(self.write) + content = report.read_text() if report.exists() else None + return UiPathRuntimeResult(status=self.status, output={"report": content}) + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + yield await self.execute(input, options) + + async def get_schema(self) -> UiPathRuntimeSchema: + raise NotImplementedError + + async def dispose(self) -> None: + pass + + +def _managed( + services: FakeServices, + storage: MemoryStorage, + workspace: Workspace | None = None, +) -> ManagedWorkspace: + return ManagedWorkspace( + services=services, + runtime_id="runtime-1", + job_key=JOB_KEY, + workspace=workspace, + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + ) + + +@pytest.mark.asyncio +async def test_files_written_before_a_suspend_are_back_after_resume() -> None: + services, storage = FakeServices(), MemoryStorage() + + suspended = _managed(services, storage).wrap( + ReportRuntime(write="draft", status=UiPathRuntimeStatus.SUSPENDED) + ) + await suspended.execute() + await suspended.dispose() + resumed = _managed(services, storage).wrap(ReportRuntime()) + result = await resumed.execute() + await resumed.dispose() + + assert result.output == {"report": "draft"} + assert [job for job, _ in services.jobs.links] == [uuid.UUID(JOB_KEY)] + + +@pytest.mark.asyncio +async def test_services_are_built_only_when_the_run_starts() -> None: + services = FakeServices() + + runtime = _managed(services, MemoryStorage()).wrap(ReportRuntime()) + built_before_run = services.built + await runtime.execute() + await runtime.execute() + + assert built_before_run == 0 + assert services.built == 1 + + +@pytest.mark.asyncio +async def test_conversation_and_hydration_wrappers_share_one_hydrator() -> None: + services = FakeServices() + managed = _managed(services, MemoryStorage()) + + runtime = managed.wrap_conversational(managed.wrap(ReportRuntime(write="notes"))) + events = [event async for event in runtime.stream({"messages": []})] + + snapshots = [ + event.payload + for event in events + if isinstance(event, UiPathRuntimeConversationMetaEvent) + ] + assert [f["path"] for f in snapshots[0]["workspaceFiles"]] == ["report.md"] + assert services.built == 1 + assert services.attachments.uploads == 1 + + +@pytest.mark.asyncio +async def test_creates_and_removes_its_own_workspace() -> None: + runtime = _managed(FakeServices(), MemoryStorage()) + wrapped = runtime.wrap(ReportRuntime()) + + assert runtime.workspace.path.exists() + await wrapped.dispose() + assert not runtime.workspace.path.exists() + + +@pytest.mark.asyncio +async def test_keeps_a_workspace_owned_by_the_caller(tmp_path: Path) -> None: + owned = Workspace(tmp_path / "owned", cleanup=False) + owned.path.mkdir() + inner = ReportRuntime(write="draft") + + wrapped = _managed(FakeServices(), MemoryStorage(), workspace=owned).wrap(inner) + await wrapped.execute() + await wrapped.dispose() + + assert inner.workspace_path == owned.path + assert (owned.path / "report.md").read_text() == "draft" + + +def test_wrapped_runtime_carries_the_runtime_id() -> None: + wrapped = _managed(FakeServices(), MemoryStorage()).wrap(ReportRuntime()) + + assert wrapped.runtime_id == "runtime-1" + + +def test_hydration_needs_a_registry_store() -> None: + managed = ManagedWorkspace( + services=FakeServices(), runtime_id="runtime-1", job_key=JOB_KEY + ) + + with pytest.raises(ValueError, match="registry store"): + managed.wrap(ReportRuntime()) diff --git a/tests/workspace/test_workspace_hydration.py b/tests/workspace/test_workspace_hydration.py index 8d59838..06d5990 100644 --- a/tests/workspace/test_workspace_hydration.py +++ b/tests/workspace/test_workspace_hydration.py @@ -4,7 +4,7 @@ import shutil import uuid from pathlib import Path -from typing import Any, AsyncGenerator +from typing import Any, AsyncGenerator, Callable import pytest from uipath.core.chat import ( @@ -25,6 +25,7 @@ WorkspaceRegistryStore, get_workspace_path, ) +from uipath.runtime.debug import UiPathBreakpointResult from uipath.runtime.errors import UiPathErrorCategory, UiPathRuntimeError from uipath.runtime.events import ( UiPathRuntimeConversationMetaEvent, @@ -1149,3 +1150,214 @@ async def test_always_policy_persists_on_stream_failure(tmp_path: Path) -> None: pass assert attachments.uploads == 1 + + +class _ScriptedDebugRuntime: + """Runs one scripted step per call, the way a debugger drives a paused graph.""" + + def __init__(self, steps: list[Callable[[Path], UiPathRuntimeResult]]) -> None: + self.steps = steps + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + return self.steps.pop(0)(get_workspace_path()) + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + yield await self.execute(input, options) + + async def get_schema(self) -> UiPathRuntimeSchema: + raise NotImplementedError + + async def dispose(self) -> None: + pass + + +def _paused() -> UiPathBreakpointResult: + return UiPathBreakpointResult( + breakpoint_node="agent", + breakpoint_type="before", + current_state={}, + next_nodes=["agent"], + ) + + +def _write( + content: str, result: UiPathRuntimeResult +) -> Callable[[Path], UiPathRuntimeResult]: + def step(workspace: Path) -> UiPathRuntimeResult: + (workspace / "report.md").write_text(content) + return result + + return step + + +def _report(workspace: Path) -> UiPathRuntimeResult: + report = workspace / "report.md" + return UiPathRuntimeResult( + output={"report": report.read_text() if report.exists() else None} + ) + + +def _scripted( + tmp_path: Path, + steps: list[Callable[[Path], UiPathRuntimeResult]], + attachments: FakeAttachments, + storage: MemoryStorage, + policy: HydrationPolicy = HydrationPolicy.SUSPEND_ONLY, + name: str = "workspace", +) -> HydrationRuntime: + workspace = Workspace.create(tmp_path / name) + return HydrationRuntime( + _ScriptedDebugRuntime(steps), + workspace=workspace, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, attachments=attachments + ), + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + policy=policy, + ) + + +_DEBUGGING = UiPathStreamOptions(breakpoints="*") + + +async def _run( + runtime: HydrationRuntime, + *, + use_stream: bool, + options: UiPathStreamOptions | None = _DEBUGGING, +) -> UiPathRuntimeResult: + if not use_stream: + return await runtime.execute({}, options=options) + events = [event async for event in runtime.stream({}, options=options)] + result = events[-1] + assert isinstance(result, UiPathRuntimeResult) + return result + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", list(HydrationPolicy)) +async def test_breakpoint_pause_uploads_nothing( + tmp_path: Path, policy: HydrationPolicy +) -> None: + attachments = FakeAttachments() + runtime = _scripted( + tmp_path, + [_write("draft", _paused()), _write("draft", _paused())], + attachments, + MemoryStorage(), + policy=policy, + ) + + executed = await runtime.execute({}, options=_DEBUGGING) + streamed = await _run(runtime, use_stream=True) + + assert isinstance(executed, UiPathBreakpointResult) + assert isinstance(streamed, UiPathBreakpointResult) + assert attachments.uploads == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_stream", [False, True]) +async def test_continue_after_breakpoint_keeps_edits_made_before_it( + tmp_path: Path, use_stream: bool +) -> None: + suspended = UiPathRuntimeResult(status=UiPathRuntimeStatus.SUSPENDED) + runtime = _scripted( + tmp_path, + [_write("v1", suspended), _write("v2", _paused()), _report], + FakeAttachments(), + MemoryStorage(), + ) + + await _run(runtime, use_stream=use_stream) + await _run(runtime, use_stream=use_stream) + continued = await _run(runtime, use_stream=use_stream) + + assert continued.output == {"report": "v2"} + + +@pytest.mark.asyncio +async def test_continue_after_breakpoint_keeps_deleted_files_deleted( + tmp_path: Path, +) -> None: + def delete_then_pause(workspace: Path) -> UiPathRuntimeResult: + (workspace / "report.md").unlink() + return _paused() + + suspended = UiPathRuntimeResult(status=UiPathRuntimeStatus.SUSPENDED) + runtime = _scripted( + tmp_path, + [_write("v1", suspended), delete_then_pause, _report], + FakeAttachments(), + MemoryStorage(), + ) + + for _ in range(2): + await runtime.execute({}, options=_DEBUGGING) + continued = await runtime.execute({}, options=_DEBUGGING) + + assert continued.output == {"report": None} + + +@pytest.mark.asyncio +async def test_suspend_after_a_breakpoint_is_restored_on_resume( + tmp_path: Path, +) -> None: + attachments, storage = FakeAttachments(), MemoryStorage() + suspended = UiPathRuntimeResult(status=UiPathRuntimeStatus.SUSPENDED) + debug_run = _scripted( + tmp_path, + [_write("v1", suspended), _write("v2", _paused()), _write("v3", suspended)], + attachments, + storage, + ) + for _ in range(3): + await debug_run.execute({}, options=_DEBUGGING) + + resumed = _scripted(tmp_path, [_report], attachments, storage, name="resumed") + result = await resumed.execute({}) + + assert result.output == {"report": "v3"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_stream", [False, True]) +async def test_static_interrupt_without_debugger_is_restored_on_resume( + tmp_path: Path, use_stream: bool +) -> None: + attachments, storage = FakeAttachments(), MemoryStorage() + job = _scripted(tmp_path, [_write("draft", _paused())], attachments, storage) + await _run(job, use_stream=use_stream, options=None) + + resumed = _scripted(tmp_path, [_report], attachments, storage, name="resumed") + result = await _run(resumed, use_stream=use_stream, options=None) + + assert result.output == {"report": "draft"} + + +@pytest.mark.asyncio +async def test_only_the_first_call_after_a_breakpoint_skips_hydration( + tmp_path: Path, +) -> None: + suspended = UiPathRuntimeResult(status=UiPathRuntimeStatus.SUSPENDED) + runtime = _scripted( + tmp_path, + [_write("v1", _paused()), _write("v2", suspended), _report], + FakeAttachments(), + MemoryStorage(), + ) + for _ in range(2): + await runtime.execute({}, options=_DEBUGGING) + + (runtime.workspace.path / "report.md").unlink() + result = await runtime.execute({}, options=_DEBUGGING) + + assert result.output == {"report": "v2"} diff --git a/uv.lock b/uv.lock index f4a278a..d585534 100644 --- a/uv.lock +++ b/uv.lock @@ -1153,7 +1153,7 @@ wheels = [ [[package]] name = "uipath-runtime" -version = "0.13.6" +version = "0.13.7" source = { editable = "." } dependencies = [ { name = "chardet" },