From a922f1dd7a41c303296faebfd380b20ccc45ff30 Mon Sep 17 00:00:00 2001 From: radu-mocanu Date: Mon, 6 Jul 2026 17:49:01 +0300 Subject: [PATCH 1/4] feat: workspace persistence across conversational exchanges via content-parts --- pyproject.toml | 2 +- src/uipath/runtime/__init__.py | 2 + src/uipath/runtime/workspace/__init__.py | 5 + src/uipath/runtime/workspace/cas_uri.py | 33 ++ .../runtime/workspace/conversational.py | 259 +++++++++++ src/uipath/runtime/workspace/hydrator.py | 33 ++ .../test_conversational_workspace.py | 426 ++++++++++++++++++ uv.lock | 2 +- 8 files changed, 760 insertions(+), 2 deletions(-) create mode 100644 src/uipath/runtime/workspace/cas_uri.py create mode 100644 src/uipath/runtime/workspace/conversational.py create mode 100644 tests/workspace/test_conversational_workspace.py diff --git a/pyproject.toml b/pyproject.toml index c1415a6..d02f960 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-runtime" -version = "0.12.3" +version = "0.12.4" 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 afc2719..7ecbb32 100644 --- a/src/uipath/runtime/__init__.py +++ b/src/uipath/runtime/__init__.py @@ -45,6 +45,7 @@ from uipath.runtime.storage import UiPathRuntimeStorageProtocol from uipath.runtime.workspace import ( AttachmentRegistryEntry, + ConversationalWorkspaceRuntime, HydrationPolicy, HydrationRuntime, Workspace, @@ -82,6 +83,7 @@ "UiPathChatProtocol", "UiPathChatRuntime", "AttachmentRegistryEntry", + "ConversationalWorkspaceRuntime", "HydrationPolicy", "HydrationRuntime", "Workspace", diff --git a/src/uipath/runtime/workspace/__init__.py b/src/uipath/runtime/workspace/__init__.py index c8dc23b..d68da44 100644 --- a/src/uipath/runtime/workspace/__init__.py +++ b/src/uipath/runtime/workspace/__init__.py @@ -1,5 +1,7 @@ """Workspace persistence primitives for runtime implementations.""" +from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id +from uipath.runtime.workspace.conversational import ConversationalWorkspaceRuntime from uipath.runtime.workspace.hydration import ( HydrationPolicy, HydrationRuntime, @@ -13,9 +15,12 @@ __all__ = [ "AttachmentRegistryEntry", + "ConversationalWorkspaceRuntime", "HydrationPolicy", "HydrationRuntime", "Workspace", "WorkspaceHydrator", "WorkspaceRegistryStore", + "build_cas_uri", + "parse_attachment_id", ] diff --git a/src/uipath/runtime/workspace/cas_uri.py b/src/uipath/runtime/workspace/cas_uri.py new file mode 100644 index 0000000..c52bb5f --- /dev/null +++ b/src/uipath/runtime/workspace/cas_uri.py @@ -0,0 +1,33 @@ +"""CAS file URI helpers for workspace attachments.""" + +import uuid + +CAS_FILE_URI_PREFIX = "urn:uipath:cas:file:orchestrator:" + + +def build_cas_uri(attachment_key: str) -> str: + """Build a CAS file URI for an Orchestrator attachment key.""" + return f"{CAS_FILE_URI_PREFIX}{attachment_key}" + + +def parse_attachment_id(uri: str | None) -> str | None: + """Parse the attachment ID from a CAS file URI. + + Extracts the UUID from URIs like: + "urn:uipath:cas:file:orchestrator:a940a416-b97b-4146-3089-08de5f4d0a87" + + Returns the normalized (lowercase) attachment ID, or None when the URI + does not end in a valid UUID. + """ + if not uri: + return None + + # The UUID is the last segment after the final colon + parts = uri.rsplit(":", 1) + if len(parts) != 2 or not parts[1]: + return None + + try: + return str(uuid.UUID(parts[1])) + except (ValueError, AttributeError): + return None diff --git a/src/uipath/runtime/workspace/conversational.py b/src/uipath/runtime/workspace/conversational.py new file mode 100644 index 0000000..4f853fc --- /dev/null +++ b/src/uipath/runtime/workspace/conversational.py @@ -0,0 +1,259 @@ +"""Workspace persistence across conversational exchanges via content-parts. + +Conversational agents run one job per exchange, so the workspace cannot be +carried in per-run storage. Instead, the previous exchange publishes each +workspace file as a job attachment referenced by a file content-part on its +last assistant message, and the next exchange rebuilds the workspace from +those content-parts found in the conversation history. +""" + +import logging +import mimetypes +from pathlib import Path +from typing import Any, AsyncGenerator + +from pydantic import ValidationError +from uipath.core.chat import ( + UiPathConversationContentPartEndEvent, + UiPathConversationContentPartEvent, + UiPathConversationContentPartStartEvent, + UiPathConversationMessage, + UiPathConversationMessageEvent, + UiPathExternalValue, +) + +from uipath.runtime.base import ( + UiPathExecuteOptions, + UiPathRuntimeProtocol, + UiPathStreamOptions, +) +from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeMessageEvent +from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus +from uipath.runtime.schema import UiPathRuntimeSchema +from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id +from uipath.runtime.workspace.hydrator import WorkspaceHydrator +from uipath.runtime.workspace.workspace import Workspace + +logger = logging.getLogger(__name__) + +WORKSPACE_FILE_KIND = "workspace" +DEFAULT_MIME_TYPE = "application/octet-stream" + +# stdlib mimetypes lacks these until newer Python versions +_EXTRA_MIME_TYPES = { + ".md": "text/markdown", + ".yaml": "application/yaml", + ".yml": "application/yaml", +} + + +def _guess_mime_type(virtual_path: str) -> str: + suffix = Path(virtual_path).suffix.lower() + if suffix in _EXTRA_MIME_TYPES: + return _EXTRA_MIME_TYPES[suffix] + return mimetypes.guess_type(virtual_path)[0] or DEFAULT_MIME_TYPE + + +class ConversationalWorkspaceRuntime: + """Persists the workspace across exchanges through file content-parts. + + Wraps the runtime whose stream emits ``UiPathRuntimeMessageEvent``s and must + sit below ``UiPathChatRuntime`` so the injected content-part events flow + through the chat bridge like any other message event. + + Before delegating, the workspace is hydrated from the file content-parts of + the last assistant message in the conversation history. On a successful + terminal result, workspace files are uploaded (or relinked when unchanged) + as job attachments and one content-part per file is appended to the final + assistant message, before its ``endMessage`` event. + """ + + def __init__( + self, + delegate: UiPathRuntimeProtocol, + *, + hydrator: WorkspaceHydrator, + workspace: Workspace, + ): + """Initialize the wrapper with the hydrator and workspace to persist.""" + self.delegate = delegate + self.hydrator = hydrator + self.workspace = workspace + self._registry: dict[str, dict[str, Any]] = {} + self._hydrated = False + # message ids whose startMessage declared role=assistant; message starts + # and ends can arrive in different stream passes (interrupt/resume) + self._assistant_message_ids: set[str] = set() + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + """Execute by draining the stream.""" + result: UiPathRuntimeResult | None = None + async for event in self.stream( + input, + options=UiPathStreamOptions(resume=options.resume if options else False), + ): + if isinstance(event, UiPathRuntimeResult): + result = event + return result or UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL) + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + """Hydrate, stream delegate events, and emit workspace content-parts.""" + await self._hydrate(input) + + # Withhold the latest assistant endMessage event so workspace + # content-parts can be injected before it if it turns out to be the + # last assistant message of the exchange. + pending_end: UiPathRuntimeMessageEvent | None = None + final_result: UiPathRuntimeResult | None = None + + async for event in self.delegate.stream(input, options=options): + if isinstance(event, UiPathRuntimeResult): + final_result = event + continue + + if isinstance(event, UiPathRuntimeMessageEvent) and isinstance( + event.payload, UiPathConversationMessageEvent + ): + payload = event.payload + if pending_end is not None: + yield pending_end + pending_end = None + if payload.start and payload.start.role == "assistant": + self._assistant_message_ids.add(payload.message_id) + if payload.end and payload.message_id in self._assistant_message_ids: + pending_end = event + continue + + yield event + + if final_result is None: + if pending_end is not None: + yield pending_end + return + + if ( + final_result.status == UiPathRuntimeStatus.SUCCESSFUL + and pending_end is not None + ): + try: + for file_event in await self._dehydrate(pending_end.payload.message_id): + yield file_event + except Exception: + logger.exception("Failed to persist workspace files as content-parts") + + if pending_end is not None: + yield pending_end + 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, input: dict[str, Any] | None) -> None: + """Rebuild the workspace from the last assistant message's file parts.""" + if self._hydrated: + # resume passes re-enter stream() with resume payloads, not history + return + + messages = (input or {}).get("messages") + if not isinstance(messages, list): + return + self._hydrated = True + + last_assistant = self._last_assistant_message(messages) + if last_assistant is None: + return + + files: dict[str, str] = {} + for part in last_assistant.content_parts: + if not isinstance(part.data, UiPathExternalValue): + continue + attachment_id = parse_attachment_id(part.data.uri) + if attachment_id is None or not part.name: + continue + # Interim rule (until CAS persists content-part metaData): every + # file content-part on an assistant message is a workspace file. + files[part.name] = attachment_id + + if files: + self._registry = await self.hydrator.hydrate_files(files) + + async def _dehydrate(self, message_id: str) -> list[UiPathRuntimeMessageEvent]: + """Upload/relink workspace files and build their content-part events.""" + self._registry = await self.hydrator.dehydrate(self._registry) + + events: list[UiPathRuntimeMessageEvent] = [] + for index, (virtual_path, entry) in enumerate(sorted(self._registry.items())): + content_part_id = f"ws-{message_id}-{index}" + mime_type = _guess_mime_type(virtual_path) + metadata = { + "fileKind": WORKSPACE_FILE_KIND, + "sha256": entry["sha256"], + } + events.append( + UiPathRuntimeMessageEvent( + payload=UiPathConversationMessageEvent( + message_id=message_id, + content_part=UiPathConversationContentPartEvent( + content_part_id=content_part_id, + start=UiPathConversationContentPartStartEvent( + mime_type=mime_type, + name=virtual_path, + external_value=UiPathExternalValue( + uri=build_cas_uri(entry["attachment_key"]), + byte_count=entry["size"], + ), + metadata=metadata, + ), + ), + ) + ) + ) + events.append( + UiPathRuntimeMessageEvent( + payload=UiPathConversationMessageEvent( + message_id=message_id, + content_part=UiPathConversationContentPartEvent( + content_part_id=content_part_id, + end=UiPathConversationContentPartEndEvent( + metadata=metadata, + ), + ), + ) + ) + ) + return events + + @staticmethod + def _last_assistant_message( + messages: list[Any], + ) -> UiPathConversationMessage | None: + last: UiPathConversationMessage | None = None + for message in messages: + if isinstance(message, UiPathConversationMessage): + parsed = message + elif isinstance(message, dict): + try: + parsed = UiPathConversationMessage.model_validate(message) + except ValidationError: + continue + else: + continue + if parsed.role == "assistant": + last = parsed + return last diff --git a/src/uipath/runtime/workspace/hydrator.py b/src/uipath/runtime/workspace/hydrator.py index fba298e..961a8fa 100644 --- a/src/uipath/runtime/workspace/hydrator.py +++ b/src/uipath/runtime/workspace/hydrator.py @@ -68,6 +68,39 @@ async def hydrate( ) return self._dump_registry(normalized) + async def hydrate_files( + self, + files: dict[str, str], + ) -> dict[str, dict[str, Any]]: + """Download attachments into the workspace and build a registry from them. + + Args: + files: Mapping of virtual path to attachment key. + + Unlike ``hydrate``, no prior registry (and thus no SHA-256) is known, so + every file is downloaded and its digest/size are computed from the + downloaded bytes. The returned registry can be fed to ``dehydrate`` so + unchanged files are relinked instead of re-uploaded. + """ + registry: dict[str, AttachmentRegistryEntry] = {} + for virtual_path, attachment_key in files.items(): + target = self._resolve_workspace_path(virtual_path) + target.parent.mkdir(parents=True, exist_ok=True) + await self.attachments.download_async( + key=UUID(attachment_key), + destination_path=str(target), + folder_key=self.folder_key, + folder_path=self.folder_path, + ) + registry[virtual_path] = AttachmentRegistryEntry( + attachment_key=attachment_key, + sha256=self._sha256(target), + size=target.stat().st_size, + uploaded_at=datetime.now(timezone.utc).isoformat(), + attachment_name=self._attachment_name_for_virtual_path(virtual_path), + ) + return self._dump_registry(registry) + async def dehydrate( self, registry: dict[str, dict[str, Any]], diff --git a/tests/workspace/test_conversational_workspace.py b/tests/workspace/test_conversational_workspace.py new file mode 100644 index 0000000..f60db42 --- /dev/null +++ b/tests/workspace/test_conversational_workspace.py @@ -0,0 +1,426 @@ +from __future__ import annotations + +import uuid +from pathlib import Path +from typing import Any, AsyncGenerator + +import pytest +from uipath.core.chat import ( + UiPathConversationContentPartEvent, + UiPathConversationContentPartStartEvent, + UiPathConversationMessageEndEvent, + UiPathConversationMessageEvent, + UiPathConversationMessageStartEvent, + UiPathExternalValue, +) + +from tests.workspace.test_workspace_hydration import FakeAttachments, FakeJobs +from uipath.runtime import ( + ConversationalWorkspaceRuntime, + UiPathExecuteOptions, + UiPathRuntimeResult, + UiPathRuntimeStatus, + UiPathStreamOptions, + Workspace, + WorkspaceHydrator, +) +from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeMessageEvent +from uipath.runtime.schema import UiPathRuntimeSchema +from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id + + +def message_start( + message_id: str, role: str = "assistant" +) -> UiPathRuntimeMessageEvent: + return UiPathRuntimeMessageEvent( + payload=UiPathConversationMessageEvent( + message_id=message_id, + start=UiPathConversationMessageStartEvent(role=role), + ) + ) + + +def message_end(message_id: str) -> UiPathRuntimeMessageEvent: + return UiPathRuntimeMessageEvent( + payload=UiPathConversationMessageEvent( + message_id=message_id, + end=UiPathConversationMessageEndEvent(), + ) + ) + + +def history_assistant_message( + message_id: str, parts: list[dict[str, Any]] +) -> dict[str, Any]: + return { + "messageId": message_id, + "role": "assistant", + "contentParts": [ + { + "contentPartId": f"{message_id}-{i}", + "mimeType": part.get("mimeType", "application/octet-stream"), + "name": part.get("name"), + "data": {"uri": part["uri"]}, + } + for i, part in enumerate(parts) + ], + } + + +class ScriptedRuntime: + """Delegate emitting a fixed event script then a result.""" + + def __init__( + self, + events: list[UiPathRuntimeEvent], + result: UiPathRuntimeResult, + workspace_path: Path | None = None, + files: dict[str, str] | None = None, + ) -> None: + self.events = events + self.result = result + self.workspace_path = workspace_path + self.files = files or {} + self.disposed = False + self.received_input: dict[str, Any] | None = None + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + raise NotImplementedError + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + self.received_input = input + if self.workspace_path is not None: + for virtual_path, content in self.files.items(): + target = self.workspace_path / virtual_path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(content, encoding="utf-8") + for event in self.events: + yield event + yield self.result + + async def get_schema(self) -> UiPathRuntimeSchema: + raise NotImplementedError + + async def dispose(self) -> None: + self.disposed = True + + +def make_runtime( + tmp_path: Path, + delegate: ScriptedRuntime, + attachments: FakeAttachments, + jobs: FakeJobs | None = None, + job_key: str | None = None, +) -> tuple[ConversationalWorkspaceRuntime, Workspace]: + workspace = Workspace.create(tmp_path / "workspace") + hydrator = WorkspaceHydrator( + workspace_path=workspace.path, + attachments=attachments, + jobs=jobs, + current_job_key=job_key, + ) + return ( + ConversationalWorkspaceRuntime( + delegate, hydrator=hydrator, workspace=workspace + ), + workspace, + ) + + +def file_part_events( + events: list[UiPathRuntimeEvent], +) -> list[UiPathConversationContentPartEvent]: + parts = [] + for event in events: + if isinstance(event, UiPathRuntimeMessageEvent) and isinstance( + event.payload, UiPathConversationMessageEvent + ): + part = event.payload.content_part + if part and part.content_part_id.startswith("ws-"): + parts.append(part) + return parts + + +@pytest.mark.asyncio +async def test_emits_file_content_parts_before_final_message_end( + tmp_path: Path, +) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + workspace_path=tmp_path / "workspace", + files={"plan/todo.md": "step 1"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": []})] + + parts = file_part_events(events) + # one start + one end event per workspace file + assert len(parts) == 2 + start = parts[0].start + assert isinstance(start, UiPathConversationContentPartStartEvent) + assert start.name == "plan/todo.md" + assert start.mime_type == "text/markdown" + assert start.metadata is not None + assert start.metadata["fileKind"] == "workspace" + assert start.metadata["sha256"] + assert isinstance(start.external_value, UiPathExternalValue) + assert parse_attachment_id(start.external_value.uri) is not None + assert parts[1].end is not None + + # ordering: file parts on m1 come before m1's endMessage, then the result + end_index = next( + i + for i, e in enumerate(events) + if isinstance(e, UiPathRuntimeMessageEvent) + and isinstance(e.payload, UiPathConversationMessageEvent) + and e.payload.end is not None + ) + part_indices = [ + i + for i, e in enumerate(events) + if isinstance(e, UiPathRuntimeMessageEvent) + and isinstance(e.payload, UiPathConversationMessageEvent) + and e.payload.content_part is not None + ] + assert all(i < end_index for i in part_indices) + assert isinstance(events[-1], UiPathRuntimeResult) + assert attachments.uploads == 1 + + +@pytest.mark.asyncio +async def test_only_last_assistant_message_gets_file_parts(tmp_path: Path) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[ + message_start("m1"), + message_end("m1"), + message_start("m2"), + message_end("m2"), + ], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + workspace_path=tmp_path / "workspace", + files={"notes.txt": "hello"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": []})] + + part_message_ids = { + e.payload.message_id + for e in events + if isinstance(e, UiPathRuntimeMessageEvent) + and isinstance(e.payload, UiPathConversationMessageEvent) + and e.payload.content_part is not None + } + assert part_message_ids == {"m2"} + # m1's end was flushed before m2 started + order = [ + (e.payload.message_id, "end" if e.payload.end else "other") + for e in events + if isinstance(e, UiPathRuntimeMessageEvent) + and isinstance(e.payload, UiPathConversationMessageEvent) + ] + assert order.index(("m1", "end")) < order.index(("m2", "other")) + + +@pytest.mark.asyncio +async def test_hydrates_workspace_from_last_assistant_history_message( + tmp_path: Path, +) -> None: + attachments = FakeAttachments() + jobs = FakeJobs() + job_key = str(uuid.uuid4()) + old_key = uuid.uuid4() + stale_key = uuid.uuid4() + attachments.files[old_key] = ("todo.md", b"step 1") + attachments.files[stale_key] = ("stale.md", b"old") + history = [ + {"messageId": "u1", "role": "user", "contentParts": []}, + history_assistant_message( + "a1", [{"name": "stale.md", "uri": build_cas_uri(str(stale_key))}] + ), + {"messageId": "u2", "role": "user", "contentParts": []}, + history_assistant_message( + "a2", + [{"name": "plan/todo.md", "uri": build_cas_uri(str(old_key))}], + ), + ] + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, workspace = make_runtime( + tmp_path, delegate, attachments, jobs=jobs, job_key=job_key + ) + + events = [event async for event in runtime.stream({"messages": history})] + + # only the LAST assistant message's parts are hydrated + assert (workspace.path / "plan/todo.md").read_text(encoding="utf-8") == "step 1" + assert not (workspace.path / "stale.md").exists() + + # unchanged file is relinked to the current job, not re-uploaded + assert attachments.uploads == 0 + assert (uuid.UUID(job_key), old_key) in jobs.links + + # and re-emitted with the SAME attachment id + starts = [p.start for p in file_part_events(events) if p.start] + assert len(starts) == 1 + assert starts[0].external_value is not None + assert parse_attachment_id(starts[0].external_value.uri) == str(old_key) + + +@pytest.mark.asyncio +async def test_changed_file_is_reuploaded(tmp_path: Path) -> None: + attachments = FakeAttachments() + old_key = uuid.uuid4() + attachments.files[old_key] = ("todo.md", b"step 1") + history = [ + history_assistant_message( + "a1", [{"name": "todo.md", "uri": build_cas_uri(str(old_key))}] + ) + ] + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + workspace_path=tmp_path / "workspace", + files={"todo.md": "step 1 done, step 2"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": history})] + + assert attachments.uploads == 1 + starts = [p.start for p in file_part_events(events) if p.start] + assert starts[0].external_value is not None + assert parse_attachment_id(starts[0].external_value.uri) != str(old_key) + + +@pytest.mark.asyncio +async def test_no_emission_on_faulted_result(tmp_path: Path) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.FAULTED), + workspace_path=tmp_path / "workspace", + files={"notes.txt": "hello"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": []})] + + assert file_part_events(events) == [] + assert attachments.uploads == 0 + # the buffered end message is still flushed + ends = [ + e + for e in events + if isinstance(e, UiPathRuntimeMessageEvent) + and isinstance(e.payload, UiPathConversationMessageEvent) + and e.payload.end is not None + ] + assert len(ends) == 1 + assert isinstance(events[-1], UiPathRuntimeResult) + + +@pytest.mark.asyncio +async def test_non_assistant_message_end_is_not_buffered(tmp_path: Path) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[message_start("m1", role="user"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + workspace_path=tmp_path / "workspace", + files={"notes.txt": "hello"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": []})] + + assert file_part_events(events) == [] + assert attachments.uploads == 0 + + +@pytest.mark.asyncio +async def test_hydrate_only_happens_once_across_resume_passes( + tmp_path: Path, +) -> None: + attachments = FakeAttachments() + key = uuid.uuid4() + attachments.files[key] = ("todo.md", b"step 1") + history = [ + history_assistant_message( + "a1", [{"name": "todo.md", "uri": build_cas_uri(str(key))}] + ) + ] + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + async for _ in runtime.stream({"messages": history}): + pass + downloads_after_first = attachments.downloads + # resume pass: input is a resume map, not history + async for _ in runtime.stream({"interrupt-1": {"approved": True}}): + pass + + assert attachments.downloads == downloads_after_first == 1 + + +@pytest.mark.asyncio +async def test_history_parts_without_name_or_invalid_uri_are_skipped( + tmp_path: Path, +) -> None: + attachments = FakeAttachments() + history = [ + history_assistant_message( + "a1", + [ + {"name": None, "uri": build_cas_uri(str(uuid.uuid4()))}, + {"name": "bad.md", "uri": "urn:uipath:cas:file:orchestrator:nope"}, + ], + ) + ] + delegate = ScriptedRuntime( + events=[], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": history})] + + assert attachments.downloads == 0 + assert isinstance(events[-1], UiPathRuntimeResult) + + +@pytest.mark.asyncio +async def test_dispose_disposes_delegate_and_workspace(tmp_path: Path) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[], result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL) + ) + workspace = Workspace.create(tmp_path / "workspace", cleanup=True) + runtime = ConversationalWorkspaceRuntime( + delegate, + hydrator=WorkspaceHydrator( + workspace_path=workspace.path, attachments=attachments + ), + workspace=workspace, + ) + + await runtime.dispose() + + assert delegate.disposed + assert not workspace.path.exists() diff --git a/uv.lock b/uv.lock index 1aa86de..685e071 100644 --- a/uv.lock +++ b/uv.lock @@ -1153,7 +1153,7 @@ wheels = [ [[package]] name = "uipath-runtime" -version = "0.12.3" +version = "0.12.4" source = { editable = "." } dependencies = [ { name = "chardet" }, From 2ca571cde48ee6e629d38fdf542cca67086070b4 Mon Sep 17 00:00:00 2001 From: radu-mocanu Date: Tue, 14 Jul 2026 14:54:55 +0300 Subject: [PATCH 2/4] fix: harden conversational workspace persistence --- src/uipath/runtime/workspace/cas_uri.py | 18 +--- .../runtime/workspace/conversational.py | 91 +++++++++---------- tests/workspace/test_cas_uri.py | 16 ++++ .../test_conversational_workspace.py | 81 ++++++++++++++++- 4 files changed, 143 insertions(+), 63 deletions(-) create mode 100644 tests/workspace/test_cas_uri.py diff --git a/src/uipath/runtime/workspace/cas_uri.py b/src/uipath/runtime/workspace/cas_uri.py index c52bb5f..909cc0f 100644 --- a/src/uipath/runtime/workspace/cas_uri.py +++ b/src/uipath/runtime/workspace/cas_uri.py @@ -11,23 +11,15 @@ def build_cas_uri(attachment_key: str) -> str: def parse_attachment_id(uri: str | None) -> str | None: - """Parse the attachment ID from a CAS file URI. - - Extracts the UUID from URIs like: - "urn:uipath:cas:file:orchestrator:a940a416-b97b-4146-3089-08de5f4d0a87" - - Returns the normalized (lowercase) attachment ID, or None when the URI - does not end in a valid UUID. - """ - if not uri: + """Return the normalized attachment ID from an Orchestrator CAS URI.""" + if not uri or not uri.startswith(CAS_FILE_URI_PREFIX): return None - # The UUID is the last segment after the final colon - parts = uri.rsplit(":", 1) - if len(parts) != 2 or not parts[1]: + attachment_id = uri.removeprefix(CAS_FILE_URI_PREFIX) + if not attachment_id or ":" in attachment_id: return None try: - return str(uuid.UUID(parts[1])) + return str(uuid.UUID(attachment_id)) except (ValueError, AttributeError): return None diff --git a/src/uipath/runtime/workspace/conversational.py b/src/uipath/runtime/workspace/conversational.py index 4f853fc..0ca7035 100644 --- a/src/uipath/runtime/workspace/conversational.py +++ b/src/uipath/runtime/workspace/conversational.py @@ -1,11 +1,4 @@ -"""Workspace persistence across conversational exchanges via content-parts. - -Conversational agents run one job per exchange, so the workspace cannot be -carried in per-run storage. Instead, the previous exchange publishes each -workspace file as a job attachment referenced by a file content-part on its -last assistant message, and the next exchange rebuilds the workspace from -those content-parts found in the conversation history. -""" +"""Attachment-backed workspace persistence for conversational agents.""" import logging import mimetypes @@ -55,18 +48,7 @@ def _guess_mime_type(virtual_path: str) -> str: class ConversationalWorkspaceRuntime: - """Persists the workspace across exchanges through file content-parts. - - Wraps the runtime whose stream emits ``UiPathRuntimeMessageEvent``s and must - sit below ``UiPathChatRuntime`` so the injected content-part events flow - through the chat bridge like any other message event. - - Before delegating, the workspace is hydrated from the file content-parts of - the last assistant message in the conversation history. On a successful - terminal result, workspace files are uploaded (or relinked when unchanged) - as job attachments and one content-part per file is appended to the final - assistant message, before its ``endMessage`` event. - """ + """Persists a workspace in marked file parts on assistant messages.""" def __init__( self, @@ -81,8 +63,6 @@ def __init__( self.workspace = workspace self._registry: dict[str, dict[str, Any]] = {} self._hydrated = False - # message ids whose startMessage declared role=assistant; message starts - # and ends can arrive in different stream passes (interrupt/resume) self._assistant_message_ids: set[str] = set() async def execute( @@ -108,9 +88,7 @@ async def stream( """Hydrate, stream delegate events, and emit workspace content-parts.""" await self._hydrate(input) - # Withhold the latest assistant endMessage event so workspace - # content-parts can be injected before it if it turns out to be the - # last assistant message of the exchange. + # Workspace parts must precede the final assistant end event. pending_end: UiPathRuntimeMessageEvent | None = None final_result: UiPathRuntimeResult | None = None @@ -139,15 +117,13 @@ async def stream( yield pending_end return - if ( - final_result.status == UiPathRuntimeStatus.SUCCESSFUL - and pending_end is not None - ): - try: - for file_event in await self._dehydrate(pending_end.payload.message_id): - yield file_event - except Exception: - logger.exception("Failed to persist workspace files as content-parts") + if final_result.status == UiPathRuntimeStatus.SUCCESSFUL: + if pending_end is None: + raise RuntimeError( + "Conversational workspace requires a completed assistant message" + ) + for file_event in await self._dehydrate(pending_end.payload.message_id): + yield file_event if pending_end is not None: yield pending_end @@ -167,35 +143,58 @@ async def dispose(self) -> None: async def _hydrate(self, input: dict[str, Any] | None) -> None: """Rebuild the workspace from the last assistant message's file parts.""" if self._hydrated: - # resume passes re-enter stream() with resume payloads, not history return + self._hydrated = True + files = self._files_from_history(input) + logger.info( + "Conversational workspace hydrate: %d file part(s) from history", + len(files), + ) + if files: + self._registry = await self.hydrator.hydrate_files(files) + + @staticmethod + def _files_from_history(input: dict[str, Any] | None) -> dict[str, str]: + """Collect workspace files from the last assistant message's file parts.""" messages = (input or {}).get("messages") if not isinstance(messages, list): - return - self._hydrated = True + return {} - last_assistant = self._last_assistant_message(messages) + last_assistant = ConversationalWorkspaceRuntime._last_assistant_message( + messages + ) if last_assistant is None: - return + return {} files: dict[str, str] = {} for part in last_assistant.content_parts: if not isinstance(part.data, UiPathExternalValue): continue + metadata = getattr(part, "metadata", None) + has_workspace_metadata = ( + isinstance(metadata, dict) + and metadata.get("fileKind") == WORKSPACE_FILE_KIND + ) + is_workspace_file = has_workspace_metadata or part.content_part_id.startswith( + "ws-" + ) + if not is_workspace_file: + continue attachment_id = parse_attachment_id(part.data.uri) if attachment_id is None or not part.name: continue - # Interim rule (until CAS persists content-part metaData): every - # file content-part on an assistant message is a workspace file. files[part.name] = attachment_id - - if files: - self._registry = await self.hydrator.hydrate_files(files) + return files async def _dehydrate(self, message_id: str) -> list[UiPathRuntimeMessageEvent]: """Upload/relink workspace files and build their content-part events.""" self._registry = await self.hydrator.dehydrate(self._registry) + logger.info( + "Conversational workspace dehydrate: emitting %d workspace file(s) on message %s", + len(self._registry), + message_id, + ) events: list[UiPathRuntimeMessageEvent] = [] for index, (virtual_path, entry) in enumerate(sorted(self._registry.items())): @@ -230,9 +229,7 @@ async def _dehydrate(self, message_id: str) -> list[UiPathRuntimeMessageEvent]: message_id=message_id, content_part=UiPathConversationContentPartEvent( content_part_id=content_part_id, - end=UiPathConversationContentPartEndEvent( - metadata=metadata, - ), + end=UiPathConversationContentPartEndEvent(), ), ) ) diff --git a/tests/workspace/test_cas_uri.py b/tests/workspace/test_cas_uri.py new file mode 100644 index 0000000..ba60db3 --- /dev/null +++ b/tests/workspace/test_cas_uri.py @@ -0,0 +1,16 @@ +import uuid + +from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id + + +def test_parse_attachment_id_accepts_orchestrator_cas_uri() -> None: + attachment_id = uuid.uuid4() + + assert parse_attachment_id(build_cas_uri(str(attachment_id))) == str(attachment_id) + + +def test_parse_attachment_id_rejects_other_uri_namespaces() -> None: + attachment_id = uuid.uuid4() + + assert parse_attachment_id(f"urn:other:{attachment_id}") is None + assert parse_attachment_id(f"https://example.test/{attachment_id}") is None diff --git a/tests/workspace/test_conversational_workspace.py b/tests/workspace/test_conversational_workspace.py index f60db42..5ca4e81 100644 --- a/tests/workspace/test_conversational_workspace.py +++ b/tests/workspace/test_conversational_workspace.py @@ -57,10 +57,11 @@ def history_assistant_message( "role": "assistant", "contentParts": [ { - "contentPartId": f"{message_id}-{i}", + "contentPartId": part.get("contentPartId", f"{message_id}-{i}"), "mimeType": part.get("mimeType", "application/octet-stream"), "name": part.get("name"), "data": {"uri": part["uri"]}, + "metaData": part.get("metaData", {"fileKind": "workspace"}), } for i, part in enumerate(parts) ], @@ -177,6 +178,7 @@ async def test_emits_file_content_parts_before_final_message_end( assert isinstance(start.external_value, UiPathExternalValue) assert parse_attachment_id(start.external_value.uri) is not None assert parts[1].end is not None + assert parts[1].end.metadata is None # ordering: file parts on m1 come before m1's endMessage, then the result end_index = next( @@ -339,7 +341,7 @@ async def test_non_assistant_message_end_is_not_buffered(tmp_path: Path) -> None attachments = FakeAttachments() delegate = ScriptedRuntime( events=[message_start("m1", role="user"), message_end("m1")], - result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.FAULTED), workspace_path=tmp_path / "workspace", files={"notes.txt": "hello"}, ) @@ -394,7 +396,7 @@ async def test_history_parts_without_name_or_invalid_uri_are_skipped( ) ] delegate = ScriptedRuntime( - events=[], + events=[message_start("m1"), message_end("m1")], result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), ) runtime, _ = make_runtime(tmp_path, delegate, attachments) @@ -405,6 +407,79 @@ async def test_history_parts_without_name_or_invalid_uri_are_skipped( assert isinstance(events[-1], UiPathRuntimeResult) +@pytest.mark.asyncio +async def test_ordinary_assistant_attachments_are_not_hydrated(tmp_path: Path) -> None: + attachments = FakeAttachments() + attachment_key = uuid.uuid4() + attachments.files[attachment_key] = ("report.pdf", b"report") + history = [ + history_assistant_message( + "a1", + [ + { + "name": "report.pdf", + "uri": build_cas_uri(str(attachment_key)), + "metaData": {"fileKind": "attachment"}, + } + ], + ) + ] + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, workspace = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": history})] + + assert attachments.downloads == 0 + assert not (workspace.path / "report.pdf").exists() + assert isinstance(events[-1], UiPathRuntimeResult) + + +@pytest.mark.asyncio +async def test_workspace_part_id_supports_older_content_models(tmp_path: Path) -> None: + attachments = FakeAttachments() + attachment_key = uuid.uuid4() + attachments.files[attachment_key] = ("plan.md", b"plan") + history = [ + history_assistant_message( + "a1", + [ + { + "contentPartId": "ws-a1-0", + "name": "plan.md", + "uri": build_cas_uri(str(attachment_key)), + "metaData": None, + } + ], + ) + ] + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, workspace = make_runtime(tmp_path, delegate, attachments) + + async for _ in runtime.stream({"messages": history}): + pass + + assert (workspace.path / "plan.md").read_bytes() == b"plan" + + +@pytest.mark.asyncio +async def test_success_without_assistant_message_fails(tmp_path: Path) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[], result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL) + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + with pytest.raises(RuntimeError, match="completed assistant message"): + async for _ in runtime.stream({"messages": []}): + pass + + @pytest.mark.asyncio async def test_dispose_disposes_delegate_and_workspace(tmp_path: Path) -> None: attachments = FakeAttachments() From 4484b84fe6638b31d6c6969c8183434c74f96b1c Mon Sep 17 00:00:00 2001 From: radu-mocanu Date: Tue, 14 Jul 2026 15:02:57 +0300 Subject: [PATCH 3/4] fix: propagate conversational workspace persistence failures --- .../runtime/workspace/conversational.py | 8 ++++--- .../test_conversational_workspace.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/uipath/runtime/workspace/conversational.py b/src/uipath/runtime/workspace/conversational.py index 0ca7035..6bfb51f 100644 --- a/src/uipath/runtime/workspace/conversational.py +++ b/src/uipath/runtime/workspace/conversational.py @@ -78,7 +78,9 @@ async def execute( ): if isinstance(event, UiPathRuntimeResult): result = event - return result or UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL) + if result is None: + raise RuntimeError("Delegate stream completed without a runtime result") + return result async def stream( self, @@ -176,8 +178,8 @@ def _files_from_history(input: dict[str, Any] | None) -> dict[str, str]: isinstance(metadata, dict) and metadata.get("fileKind") == WORKSPACE_FILE_KIND ) - is_workspace_file = has_workspace_metadata or part.content_part_id.startswith( - "ws-" + is_workspace_file = ( + has_workspace_metadata or part.content_part_id.startswith("ws-") ) if not is_workspace_file: continue diff --git a/tests/workspace/test_conversational_workspace.py b/tests/workspace/test_conversational_workspace.py index 5ca4e81..dd40f8a 100644 --- a/tests/workspace/test_conversational_workspace.py +++ b/tests/workspace/test_conversational_workspace.py @@ -480,6 +480,29 @@ async def test_success_without_assistant_message_fails(tmp_path: Path) -> None: pass +@pytest.mark.asyncio +async def test_attachment_upload_failure_propagates( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + workspace_path=tmp_path / "workspace", + files={"plan.md": "plan"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + async def fail_upload(**_: Any) -> uuid.UUID: + raise RuntimeError("upload failed") + + monkeypatch.setattr(attachments, "upload_async", fail_upload) + + with pytest.raises(RuntimeError, match="upload failed"): + async for _ in runtime.stream({"messages": []}): + pass + + @pytest.mark.asyncio async def test_dispose_disposes_delegate_and_workspace(tmp_path: Path) -> None: attachments = FakeAttachments() From 359f08fc839bc8ef3c4e06c57489b32204f683d3 Mon Sep 17 00:00:00 2001 From: radu-mocanu Date: Tue, 14 Jul 2026 15:38:11 +0300 Subject: [PATCH 4/4] fix: preserve conversational workspace state --- .../runtime/workspace/conversational.py | 27 +++-- .../test_conversational_workspace.py | 71 ++++++++++++-- tests/workspace/test_workspace_hydration.py | 98 ++++++++++++++++++- 3 files changed, 179 insertions(+), 17 deletions(-) diff --git a/src/uipath/runtime/workspace/conversational.py b/src/uipath/runtime/workspace/conversational.py index 6bfb51f..5a7e377 100644 --- a/src/uipath/runtime/workspace/conversational.py +++ b/src/uipath/runtime/workspace/conversational.py @@ -25,6 +25,7 @@ from uipath.runtime.schema import UiPathRuntimeSchema from uipath.runtime.workspace.cas_uri import build_cas_uri, parse_attachment_id from uipath.runtime.workspace.hydrator import WorkspaceHydrator +from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore from uipath.runtime.workspace.workspace import Workspace logger = logging.getLogger(__name__) @@ -56,11 +57,16 @@ def __init__( *, hydrator: WorkspaceHydrator, workspace: Workspace, + runtime_id: str | None = None, + registry_store: WorkspaceRegistryStore | None = None, ): """Initialize the wrapper with the hydrator and workspace to persist.""" self.delegate = delegate self.hydrator = hydrator self.workspace = workspace + self.registry_store = registry_store + if runtime_id is not None: + self.runtime_id = runtime_id self._registry: dict[str, dict[str, Any]] = {} self._hydrated = False self._assistant_message_ids: set[str] = set() @@ -72,10 +78,12 @@ async def execute( ) -> UiPathRuntimeResult: """Execute by draining the stream.""" result: UiPathRuntimeResult | None = None - async for event in self.stream( - input, - options=UiPathStreamOptions(resume=options.resume if options else False), - ): + stream_options = ( + UiPathStreamOptions.model_validate(options.model_dump()) + if options is not None + else None + ) + async for event in self.stream(input, options=stream_options): if isinstance(event, UiPathRuntimeResult): result = event if result is None: @@ -146,15 +154,20 @@ async def _hydrate(self, input: dict[str, Any] | None) -> None: """Rebuild the workspace from the last assistant message's file parts.""" if self._hydrated: return - self._hydrated = True - files = self._files_from_history(input) + persisted_registry = ( + await self.registry_store.load() if self.registry_store else {} + ) + files = self._files_from_history(input) if not persisted_registry else {} logger.info( "Conversational workspace hydrate: %d file part(s) from history", len(files), ) - if files: + if persisted_registry: + self._registry = await self.hydrator.hydrate(persisted_registry) + elif files: self._registry = await self.hydrator.hydrate_files(files) + self._hydrated = True @staticmethod def _files_from_history(input: dict[str, Any] | None) -> dict[str, str]: diff --git a/tests/workspace/test_conversational_workspace.py b/tests/workspace/test_conversational_workspace.py index dd40f8a..54f5c6c 100644 --- a/tests/workspace/test_conversational_workspace.py +++ b/tests/workspace/test_conversational_workspace.py @@ -57,7 +57,7 @@ def history_assistant_message( "role": "assistant", "contentParts": [ { - "contentPartId": part.get("contentPartId", f"{message_id}-{i}"), + "contentPartId": part.get("contentPartId", f"ws-{message_id}-{i}"), "mimeType": part.get("mimeType", "application/octet-stream"), "name": part.get("name"), "data": {"uri": part["uri"]}, @@ -84,6 +84,7 @@ def __init__( self.files = files or {} self.disposed = False self.received_input: dict[str, Any] | None = None + self.received_options: UiPathStreamOptions | None = None async def execute( self, @@ -98,6 +99,7 @@ async def stream( options: UiPathStreamOptions | None = None, ) -> AsyncGenerator[UiPathRuntimeEvent, None]: self.received_input = input + self.received_options = options if self.workspace_path is not None: for virtual_path, content in self.files.items(): target = self.workspace_path / virtual_path @@ -166,7 +168,6 @@ async def test_emits_file_content_parts_before_final_message_end( events = [event async for event in runtime.stream({"messages": []})] parts = file_part_events(events) - # one start + one end event per workspace file assert len(parts) == 2 start = parts[0].start assert isinstance(start, UiPathConversationContentPartStartEvent) @@ -180,7 +181,6 @@ async def test_emits_file_content_parts_before_final_message_end( assert parts[1].end is not None assert parts[1].end.metadata is None - # ordering: file parts on m1 come before m1's endMessage, then the result end_index = next( i for i, e in enumerate(events) @@ -200,6 +200,30 @@ async def test_emits_file_content_parts_before_final_message_end( assert attachments.uploads == 1 +@pytest.mark.asyncio +async def test_execute_preserves_all_options(tmp_path: Path) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + options = UiPathExecuteOptions.model_validate( + { + "resume": True, + "breakpoints": ["review"], + "protocol_extension": "value", + } + ) + + await runtime.execute({"messages": []}, options=options) + + assert delegate.received_options is not None + assert delegate.received_options.resume is True + assert delegate.received_options.breakpoints == ["review"] + assert delegate.received_options.model_extra == {"protocol_extension": "value"} + + @pytest.mark.asyncio async def test_only_last_assistant_message_gets_file_parts(tmp_path: Path) -> None: attachments = FakeAttachments() @@ -226,7 +250,6 @@ async def test_only_last_assistant_message_gets_file_parts(tmp_path: Path) -> No and e.payload.content_part is not None } assert part_message_ids == {"m2"} - # m1's end was flushed before m2 started order = [ (e.payload.message_id, "end" if e.payload.end else "other") for e in events @@ -268,15 +291,12 @@ async def test_hydrates_workspace_from_last_assistant_history_message( events = [event async for event in runtime.stream({"messages": history})] - # only the LAST assistant message's parts are hydrated assert (workspace.path / "plan/todo.md").read_text(encoding="utf-8") == "step 1" assert not (workspace.path / "stale.md").exists() - # unchanged file is relinked to the current job, not re-uploaded assert attachments.uploads == 0 assert (uuid.UUID(job_key), old_key) in jobs.links - # and re-emitted with the SAME attachment id starts = [p.start for p in file_part_events(events) if p.start] assert len(starts) == 1 assert starts[0].external_value is not None @@ -324,7 +344,6 @@ async def test_no_emission_on_faulted_result(tmp_path: Path) -> None: assert file_part_events(events) == [] assert attachments.uploads == 0 - # the buffered end message is still flushed ends = [ e for e in events @@ -374,13 +393,46 @@ async def test_hydrate_only_happens_once_across_resume_passes( async for _ in runtime.stream({"messages": history}): pass downloads_after_first = attachments.downloads - # resume pass: input is a resume map, not history async for _ in runtime.stream({"interrupt-1": {"approved": True}}): pass assert attachments.downloads == downloads_after_first == 1 +@pytest.mark.asyncio +async def test_failed_hydration_can_be_retried( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + attachments = FakeAttachments() + attachment_key = uuid.uuid4() + attachments.files[attachment_key] = ("todo.md", b"step 1") + history = [ + history_assistant_message( + "a1", [{"name": "todo.md", "uri": build_cas_uri(str(attachment_key))}] + ) + ] + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, workspace = make_runtime(tmp_path, delegate, attachments) + download = attachments.download_async + + async def fail_download(**_: Any) -> str: + raise RuntimeError("download failed") + + monkeypatch.setattr(attachments, "download_async", fail_download) + with pytest.raises(RuntimeError, match="download failed"): + async for _ in runtime.stream({"messages": history}): + pass + + monkeypatch.setattr(attachments, "download_async", download) + async for _ in runtime.stream({"messages": history}): + pass + + assert (workspace.path / "todo.md").read_bytes() == b"step 1" + + @pytest.mark.asyncio async def test_history_parts_without_name_or_invalid_uri_are_skipped( tmp_path: Path, @@ -417,6 +469,7 @@ async def test_ordinary_assistant_attachments_are_not_hydrated(tmp_path: Path) - "a1", [ { + "contentPartId": "attachment-a1-0", "name": "report.pdf", "uri": build_cas_uri(str(attachment_key)), "metaData": {"fileKind": "attachment"}, diff --git a/tests/workspace/test_workspace_hydration.py b/tests/workspace/test_workspace_hydration.py index c6a2b29..2d9e4af 100644 --- a/tests/workspace/test_workspace_hydration.py +++ b/tests/workspace/test_workspace_hydration.py @@ -6,8 +6,14 @@ from typing import Any, AsyncGenerator import pytest +from uipath.core.chat import ( + UiPathConversationMessageEndEvent, + UiPathConversationMessageEvent, + UiPathConversationMessageStartEvent, +) from uipath.runtime import ( + ConversationalWorkspaceRuntime, HydrationPolicy, HydrationRuntime, UiPathExecuteOptions, @@ -18,7 +24,11 @@ WorkspaceHydrator, WorkspaceRegistryStore, ) -from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeStateEvent +from uipath.runtime.events import ( + UiPathRuntimeEvent, + UiPathRuntimeMessageEvent, + UiPathRuntimeStateEvent, +) from uipath.runtime.schema import UiPathRuntimeSchema @@ -125,6 +135,38 @@ async def dispose(self) -> None: self.disposed = True +class ReadingRuntime(WritingRuntime): + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + assert (self.workspace_path / "notes.txt").read_text( + encoding="utf-8" + ) == "hello" + 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]: + message_id = "assistant-1" + yield UiPathRuntimeMessageEvent( + payload=UiPathConversationMessageEvent( + message_id=message_id, + start=UiPathConversationMessageStartEvent(role="assistant"), + ) + ) + yield UiPathRuntimeMessageEvent( + payload=UiPathConversationMessageEvent( + message_id=message_id, + end=UiPathConversationMessageEndEvent(), + ) + ) + yield await self.execute(input, options) + + @pytest.mark.asyncio async def test_dehydrate_uploads_changed_files_and_saves_registry( tmp_path: Path, @@ -232,6 +274,60 @@ async def test_stream_persists_on_suspend(tmp_path: Path) -> None: assert "notes.txt" in await runtime.registry_store.load() +@pytest.mark.asyncio +async def test_conversational_suspend_restores_in_new_workspace(tmp_path: Path) -> None: + attachments = FakeAttachments() + storage = MemoryStorage() + first_workspace = Workspace.create(tmp_path / "first", cleanup=False) + first_hydrator = WorkspaceHydrator( + workspace_path=first_workspace.path, + attachments=attachments, + ) + suspended_runtime = HydrationRuntime( + WritingRuntime(first_workspace.path, UiPathRuntimeStatus.SUSPENDED), + workspace=first_workspace, + hydrator=first_hydrator, + registry_store=WorkspaceRegistryStore(storage, "runtime-1"), + ) + conversation_runtime = ConversationalWorkspaceRuntime( + suspended_runtime, + workspace=first_workspace, + hydrator=first_hydrator, + ) + + events = [event async for event in conversation_runtime.stream({"messages": []})] + + assert isinstance(events[-1], UiPathRuntimeResult) + assert attachments.uploads == 1 + + second_workspace = Workspace.create(tmp_path / "second", cleanup=False) + second_hydrator = WorkspaceHydrator( + workspace_path=second_workspace.path, + attachments=attachments, + ) + second_registry_store = WorkspaceRegistryStore(storage, "runtime-1") + resumed_runtime = HydrationRuntime( + ReadingRuntime(second_workspace.path, UiPathRuntimeStatus.SUCCESSFUL), + workspace=second_workspace, + hydrator=second_hydrator, + registry_store=second_registry_store, + ) + resumed_conversation_runtime = ConversationalWorkspaceRuntime( + resumed_runtime, + workspace=second_workspace, + hydrator=second_hydrator, + registry_store=second_registry_store, + ) + + resumed_events = [ + event async for event in resumed_conversation_runtime.stream({"messages": []}) + ] + + assert isinstance(resumed_events[-1], UiPathRuntimeResult) + assert attachments.downloads == 1 + assert attachments.uploads == 1 + + @pytest.mark.asyncio async def test_hydrate_skips_unchanged_local_file(tmp_path: Path) -> None: workspace = Workspace.create(tmp_path / "workspace")