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" },