diff --git a/pyproject.toml b/pyproject.toml index c1415a6..415a7da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,11 +1,11 @@ [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" dependencies = [ - "uipath-core>=0.5.28, <0.6.0", + "uipath-core>=0.5.31, <0.6.0", "vaderSentiment>=3.3.2, <4.0", "chardet>=5.2.0, <8.0", ] 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..2462284 100644 --- a/src/uipath/runtime/workspace/__init__.py +++ b/src/uipath/runtime/workspace/__init__.py @@ -1,5 +1,6 @@ """Workspace persistence primitives for runtime implementations.""" +from uipath.runtime.workspace.conversational import ConversationalWorkspaceRuntime from uipath.runtime.workspace.hydration import ( HydrationPolicy, HydrationRuntime, @@ -13,6 +14,7 @@ __all__ = [ "AttachmentRegistryEntry", + "ConversationalWorkspaceRuntime", "HydrationPolicy", "HydrationRuntime", "Workspace", diff --git a/src/uipath/runtime/workspace/cas_uri.py b/src/uipath/runtime/workspace/cas_uri.py new file mode 100644 index 0000000..8a73d53 --- /dev/null +++ b/src/uipath/runtime/workspace/cas_uri.py @@ -0,0 +1,22 @@ +"""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) -> str | None: + """Return the normalized attachment ID from an Orchestrator CAS URI.""" + if not uri.startswith(CAS_FILE_URI_PREFIX): + return None + + attachment_id = uri.removeprefix(CAS_FILE_URI_PREFIX) + try: + return str(uuid.UUID(attachment_id)) + except ValueError: + return None diff --git a/src/uipath/runtime/workspace/conversational.py b/src/uipath/runtime/workspace/conversational.py new file mode 100644 index 0000000..f277968 --- /dev/null +++ b/src/uipath/runtime/workspace/conversational.py @@ -0,0 +1,296 @@ +"""Attachment-backed workspace persistence for conversational agents.""" + +import logging +import mimetypes +from collections.abc import Iterator, Mapping, Sequence +from pathlib import Path +from typing import Any, AsyncGenerator + +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.registry_store import WorkspaceRegistryStore + +logger = logging.getLogger(__name__) + +WORKSPACE_FILE_KIND = "workspace" +DEFAULT_MIME_TYPE = "application/octet-stream" + +_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 + + +def _last_assistant_message( + messages: Sequence[object], +) -> UiPathConversationMessage | None: + for message in reversed(messages): + if isinstance(message, UiPathConversationMessage): + if message.role == "assistant": + return message + elif isinstance(message, Mapping) and message.get("role") == "assistant": + return UiPathConversationMessage.model_validate(message) + return None + + +def _attachment_keys_from_history( + input: Mapping[str, object] | None, +) -> dict[str, str]: + """Collect workspace attachment keys from the last assistant message.""" + messages = (input or {}).get("messages") + if not isinstance(messages, list): + return {} + + last_assistant = _last_assistant_message(messages) + if last_assistant is None: + return {} + + attachment_keys_by_path: dict[str, str] = {} + for part in last_assistant.content_parts: + if not isinstance(part.data, UiPathExternalValue): + continue + is_workspace_file = ( + part.metadata.get("fileKind") == WORKSPACE_FILE_KIND + if part.metadata is not None + else 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 + attachment_keys_by_path[part.name] = attachment_id + return attachment_keys_by_path + + +def _conversation_message_from_event( + event: UiPathRuntimeEvent, +) -> tuple[UiPathRuntimeMessageEvent, UiPathConversationMessageEvent] | None: + if not isinstance(event, UiPathRuntimeMessageEvent): + return None + if not isinstance(event.payload, UiPathConversationMessageEvent): + return None + return event, event.payload + + +class _AssistantMessageEndBuffer: + """Delay assistant end events, including end-only events emitted after resume.""" + + def __init__(self) -> None: + self._open_message_roles: dict[str, str] = {} + self._pending_end: UiPathRuntimeMessageEvent | None = None + + def process(self, event: UiPathRuntimeEvent) -> Iterator[UiPathRuntimeEvent]: + conversation_event = _conversation_message_from_event(event) + if conversation_event is None: + yield event + return + message_event, message = conversation_event + + if self._pending_end is not None: + pending_end = self._pending_end + self._pending_end = None + yield pending_end + + if message.start is not None: + self._open_message_roles[message.message_id] = message.start.role + + message_role = self._open_message_roles.get(message.message_id) + is_assistant_end = message.end is not None and ( + message_role is None or message_role == "assistant" + ) + if message.end is not None: + self._open_message_roles.pop(message.message_id, None) + if is_assistant_end: + self._pending_end = message_event + else: + yield message_event + + def get_pending_end(self) -> UiPathRuntimeMessageEvent | None: + pending_end = self._pending_end + self._pending_end = None + return pending_end + + +class ConversationalWorkspaceRuntime: + """Persists a workspace in marked file parts on assistant messages.""" + + def __init__( + self, + delegate: UiPathRuntimeProtocol, + *, + hydrator: WorkspaceHydrator, + registry_store: WorkspaceRegistryStore | None = None, + ): + """Initialize the wrapper with its delegate and hydrator.""" + self.delegate = delegate + self.hydrator = hydrator + self.registry_store = registry_store + self._registry: dict[str, dict[str, Any]] = {} + self._hydrated = False + + async def execute( + self, + input: dict[str, Any] | None = None, + options: UiPathExecuteOptions | None = None, + ) -> UiPathRuntimeResult: + """Execute by draining the stream.""" + result: UiPathRuntimeResult | None = None + 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: + raise RuntimeError("Delegate stream completed without a runtime result") + return result + + 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) + + end_buffer = _AssistantMessageEndBuffer() + final_result: UiPathRuntimeResult | None = None + + async for event in self.delegate.stream(input, options=options): + if isinstance(event, UiPathRuntimeResult): + final_result = event + continue + + for event_to_emit in end_buffer.process(event): + yield event_to_emit + + pending_end = end_buffer.get_pending_end() + + if ( + final_result is not None + and final_result.status == UiPathRuntimeStatus.SUCCESSFUL + ): + if pending_end is None: + logger.error( + "Skipping workspace persistence because the delegate did not " + "complete an assistant message" + ) + else: + for file_event in await self._dehydrate(pending_end.payload.message_id): + yield file_event + + if pending_end is not None: + yield pending_end + if final_result is not None: + 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: + """Release resources owned by this wrapper.""" + + async def _hydrate(self, input: Mapping[str, object] | None) -> None: + """Rebuild the workspace from the last assistant message's file parts.""" + if self._hydrated: + return + + persisted_registry = ( + await self.registry_store.try_load() if self.registry_store else None + ) + attachment_keys_by_path = ( + _attachment_keys_from_history(input) if persisted_registry is None else {} + ) + logger.info( + "Conversational workspace hydrate: %d file part(s) from history", + len(attachment_keys_by_path), + ) + if persisted_registry is not None: + self._registry = await self.hydrator.hydrate_from_registry( + persisted_registry + ) + elif attachment_keys_by_path: + self._registry = await self.hydrator.hydrate_from_attachments( + attachment_keys_by_path + ) + self._hydrated = True + + 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) + if self.registry_store is not None: + await self.registry_store.save(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())): + 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(), + ), + ) + ) + ) + return events diff --git a/src/uipath/runtime/workspace/hydration.py b/src/uipath/runtime/workspace/hydration.py index 2b3d8e1..8dbaf3b 100644 --- a/src/uipath/runtime/workspace/hydration.py +++ b/src/uipath/runtime/workspace/hydration.py @@ -96,7 +96,7 @@ async def dispose(self) -> None: async def _hydrate(self) -> None: registry = await self.registry_store.load() - hydrated = await self.hydrator.hydrate(registry) + hydrated = await self.hydrator.hydrate_from_registry(registry) if hydrated != registry: await self.registry_store.save(hydrated) diff --git a/src/uipath/runtime/workspace/hydrator.py b/src/uipath/runtime/workspace/hydrator.py index fba298e..51c4e40 100644 --- a/src/uipath/runtime/workspace/hydrator.py +++ b/src/uipath/runtime/workspace/hydrator.py @@ -50,7 +50,14 @@ async def hydrate( self, registry: dict[str, dict[str, Any]], ) -> dict[str, dict[str, Any]]: - """Download registry files into the workspace. + """Restore a registry using the original public API.""" + return await self.hydrate_from_registry(registry) + + async def hydrate_from_registry( + self, + registry: dict[str, dict[str, Any]], + ) -> dict[str, dict[str, Any]]: + """Restore registry attachments into the workspace. Files with matching SHA-256 are left untouched. """ @@ -68,6 +75,30 @@ async def hydrate( ) return self._dump_registry(normalized) + async def hydrate_from_attachments( + self, + attachment_keys_by_path: dict[str, str], + ) -> dict[str, dict[str, Any]]: + """Restore attachments into the workspace and build their registry.""" + registry: dict[str, AttachmentRegistryEntry] = {} + for virtual_path, attachment_key in attachment_keys_by_path.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]], @@ -173,7 +204,6 @@ def _normalize_registry( ) -> dict[str, AttachmentRegistryEntry]: normalized: dict[str, AttachmentRegistryEntry] = {} for virtual_path, entry in registry.items(): - self._resolve_workspace_path(virtual_path) normalized[virtual_path] = AttachmentRegistryEntry.model_validate(entry) return normalized diff --git a/src/uipath/runtime/workspace/registry_store.py b/src/uipath/runtime/workspace/registry_store.py index d6e1eb4..a494dba 100644 --- a/src/uipath/runtime/workspace/registry_store.py +++ b/src/uipath/runtime/workspace/registry_store.py @@ -25,11 +25,11 @@ def __init__( self.namespace = namespace self.key = key - async def load(self) -> dict[str, dict[str, Any]]: - """Load registry entries keyed by workspace-relative path.""" + async def try_load(self) -> dict[str, dict[str, Any]] | None: + """Load the registry, or return ``None`` when it has not been saved.""" value = await self.storage.get_value(self.runtime_id, self.namespace, self.key) if value is None: - return {} + return None if not isinstance(value, dict): raise TypeError("Workspace registry payload must be a dictionary.") @@ -39,6 +39,11 @@ async def load(self) -> dict[str, dict[str, Any]]: registry[path] = entry return registry + async def load(self) -> dict[str, dict[str, Any]]: + """Load registry entries, defaulting to an empty registry.""" + registry = await self.try_load() + return registry if registry is not None else {} + async def save(self, registry: dict[str, dict[str, Any]]) -> None: """Persist registry entries.""" await self.storage.set_value( 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 new file mode 100644 index 0000000..7bd0df2 --- /dev/null +++ b/tests/workspace/test_conversational_workspace.py @@ -0,0 +1,704 @@ +from __future__ import annotations + +import mimetypes +import uuid +from pathlib import Path +from typing import Any, AsyncGenerator + +import pytest +from pytest_mock import MockerFixture +from uipath.core.chat import ( + UiPathConversationContentPartEvent, + UiPathConversationContentPartStartEvent, + UiPathConversationMessageEndEvent, + UiPathConversationMessageEvent, + UiPathConversationMessageStartEvent, + UiPathExternalValue, +) + +from tests.workspace.test_workspace_hydration import ( + FakeAttachments, + FakeJobs, + MemoryStorage, +) +from uipath.runtime import ( + ConversationalWorkspaceRuntime, + UiPathExecuteOptions, + UiPathRuntimeResult, + UiPathRuntimeStatus, + UiPathStreamOptions, + Workspace, + WorkspaceHydrator, + WorkspaceRegistryStore, +) +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 +from uipath.runtime.workspace.conversational import _guess_mime_type, logger + + +@pytest.mark.parametrize( + ("filename", "expected_mime_type"), + [ + ("plan.md", "text/markdown"), + ("config.yaml", "application/yaml"), + ("config.yml", "application/yaml"), + ], +) +def test_workspace_mime_types_are_stable( + filename: str, expected_mime_type: str +) -> None: + assert _guess_mime_type(filename) == expected_mime_type + + +@pytest.mark.parametrize( + ("filename", "expected_mime_type"), + [ + ("plan.md", "text/markdown"), + ("config.yaml", "application/yaml"), + ("config.yml", "application/yaml"), + ], +) +def test_workspace_mime_types_do_not_depend_on_host_database( + monkeypatch: pytest.MonkeyPatch, + filename: str, + expected_mime_type: str, +) -> None: + monkeypatch.setattr( + mimetypes, + "guess_type", + lambda _: ("application/x-host-specific", None), + ) + + assert _guess_mime_type(filename) == expected_mime_type + + +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": part.get("contentPartId", f"ws-{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) + ], + } + + +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 + self.received_options: UiPathStreamOptions | 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 + 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 + 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, + registry_store: WorkspaceRegistryStore | 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, + registry_store=registry_store, + ), + 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) + 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 + assert parts[1].end.metadata is None + + 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_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() + 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"} + 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})] + + assert (workspace.path / "plan/todo.md").read_text(encoding="utf-8") == "step 1" + assert not (workspace.path / "stale.md").exists() + + assert attachments.uploads == 0 + assert (uuid.UUID(job_key), old_key) in jobs.links + + 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_saved_empty_registry_does_not_restore_files_from_history( + tmp_path: Path, +) -> None: + attachments = FakeAttachments() + old_key = uuid.uuid4() + attachments.files[old_key] = ("plan.md", b"deleted") + history = [ + history_assistant_message( + "a1", [{"name": "plan.md", "uri": build_cas_uri(str(old_key))}] + ) + ] + store = WorkspaceRegistryStore(MemoryStorage(), "runtime-1") + await store.save({}) + delegate = ScriptedRuntime( + events=[message_start("m1"), message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + ) + runtime, workspace = make_runtime( + tmp_path, + delegate, + attachments, + registry_store=store, + ) + + events = [event async for event in runtime.stream({"messages": history})] + + assert not (workspace.path / "plan.md").exists() + assert attachments.downloads == 0 + assert file_part_events(events) == [] + + +@pytest.mark.asyncio +async def test_changed_file_is_reuploaded(tmp_path: Path) -> None: + attachments = FakeAttachments() + store = WorkspaceRegistryStore(MemoryStorage(), "runtime-1") + 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, + registry_store=store, + ) + + 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) + registry = await store.load() + assert registry["todo.md"]["attachment_key"] != 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 + 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.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 + + +@pytest.mark.asyncio +async def test_end_only_resumed_assistant_message_persists_workspace( + tmp_path: Path, +) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[message_end("m1")], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + workspace_path=tmp_path / "workspace", + files={"notes.txt": "resumed"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + events = [event async for event in runtime.stream({"messages": []})] + + assert attachments.uploads == 1 + assert len(file_part_events(events)) == 2 + assert events[-2] == message_end("m1") + assert isinstance(events[-1], UiPathRuntimeResult) + + +@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 + 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, +) -> 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=[message_start("m1"), message_end("m1")], + 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_explicit_non_workspace_metadata_takes_precedence_over_part_id( + tmp_path: Path, +) -> None: + attachments = FakeAttachments() + attachment_key = uuid.uuid4() + attachments.files[attachment_key] = ("report.pdf", b"report") + history = [ + history_assistant_message( + "a1", + [ + { + "contentPartId": "ws-a1-0", + "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_logs_and_skips_persistence( + tmp_path: Path, + mocker: MockerFixture, +) -> None: + attachments = FakeAttachments() + delegate = ScriptedRuntime( + events=[], + result=UiPathRuntimeResult(status=UiPathRuntimeStatus.SUCCESSFUL), + workspace_path=tmp_path / "workspace", + files={"plan.md": "plan"}, + ) + runtime, _ = make_runtime(tmp_path, delegate, attachments) + + log_error = mocker.patch.object(logger, "error") + events = [event async for event in runtime.stream({"messages": []})] + + assert isinstance(events[-1], UiPathRuntimeResult) + assert attachments.uploads == 0 + log_error.assert_called_once_with( + "Skipping workspace persistence because the delegate did not " + "complete an assistant message" + ) + + +@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_does_not_dispose_injected_dependencies(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 + ), + ) + + await runtime.dispose() + + assert not delegate.disposed + assert workspace.path.exists() + + await delegate.dispose() + await workspace.dispose() + + assert delegate.disposed + assert not workspace.path.exists() diff --git a/tests/workspace/test_workspace_hydration.py b/tests/workspace/test_workspace_hydration.py index c6a2b29..5710041 100644 --- a/tests/workspace/test_workspace_hydration.py +++ b/tests/workspace/test_workspace_hydration.py @@ -6,8 +6,13 @@ from typing import Any, AsyncGenerator import pytest +from uipath.core.chat import ( + UiPathConversationMessageEndEvent, + UiPathConversationMessageEvent, +) from uipath.runtime import ( + ConversationalWorkspaceRuntime, HydrationPolicy, HydrationRuntime, UiPathExecuteOptions, @@ -18,7 +23,11 @@ WorkspaceHydrator, WorkspaceRegistryStore, ) -from uipath.runtime.events import UiPathRuntimeEvent, UiPathRuntimeStateEvent +from uipath.runtime.events import ( + UiPathRuntimeEvent, + UiPathRuntimeMessageEvent, + UiPathRuntimeStateEvent, +) from uipath.runtime.schema import UiPathRuntimeSchema @@ -95,6 +104,18 @@ async def link_attachment_async( self.links.append((job_key, attachment_key)) +@pytest.mark.asyncio +async def test_registry_store_distinguishes_missing_from_saved_empty() -> None: + store = WorkspaceRegistryStore(MemoryStorage(), "runtime-1") + + assert await store.try_load() is None + + await store.save({}) + + assert await store.try_load() == {} + assert await store.load() == {} + + class WritingRuntime: def __init__(self, workspace_path: Path, status: UiPathRuntimeStatus) -> None: self.workspace_path = workspace_path @@ -125,6 +146,35 @@ 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" + (self.workspace_path / "notes.txt").write_text( + "hello after resume", encoding="utf-8" + ) + return UiPathRuntimeResult(status=self.status, output={"ok": True}) + + async def stream( + self, + input: dict[str, Any] | None = None, + options: UiPathStreamOptions | None = None, + ) -> AsyncGenerator[UiPathRuntimeEvent, None]: + message_id = "assistant-1" + 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, @@ -183,7 +233,9 @@ async def test_successful_completion_persists_when_policy_allows( @pytest.mark.asyncio -async def test_hydrate_downloads_missing_registry_files(tmp_path: Path) -> None: +async def test_hydrate_compatibility_api_downloads_registry_files( + tmp_path: Path, +) -> None: workspace = Workspace.create(tmp_path / "workspace") attachments = FakeAttachments() key = uuid.uuid4() @@ -233,7 +285,62 @@ async def test_stream_persists_on_suspend(tmp_path: Path) -> None: @pytest.mark.asyncio -async def test_hydrate_skips_unchanged_local_file(tmp_path: Path) -> None: +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, + 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, + 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 == 2 + assert len(resumed_events) == 4 + + +@pytest.mark.asyncio +async def test_hydrate_from_registry_skips_unchanged_local_file( + tmp_path: Path, +) -> None: workspace = Workspace.create(tmp_path / "workspace") attachments = FakeAttachments() (workspace.path / "notes.txt").write_text("same", encoding="utf-8") @@ -252,7 +359,7 @@ async def test_hydrate_skips_unchanged_local_file(tmp_path: Path) -> None: } } - await hydrator.hydrate(registry) + await hydrator.hydrate_from_registry(registry) assert attachments.downloads == 0 diff --git a/uv.lock b/uv.lock index 1aa86de..8538484 100644 --- a/uv.lock +++ b/uv.lock @@ -1139,21 +1139,21 @@ wheels = [ [[package]] name = "uipath-core" -version = "0.5.28" +version = "0.5.31" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4c/f9/8d2f1d98cbebbcf059cf4561f38f34ad4cd58423e4f15cad22bd297a2563/uipath_core-0.5.28.tar.gz", hash = "sha256:942987f6b612c64f93d612ad7b242276ed75f129fdd8f25bc71c24ec8887e388", size = 130578, upload-time = "2026-06-30T14:04:48.841Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1d/b2/518fcb7d466a39c675194c492e8accbffdc7fd6b280d89b21ed11700a7bd/uipath_core-0.5.31.tar.gz", hash = "sha256:6abb443582d9a8ab6a504791296d49f524fb430f66d8dfaf6d6d1849958700bb", size = 130979, upload-time = "2026-07-15T06:56:03.867Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/1e/385bb166232a57ebe938cc57ad2717f350bc922bb5d2ce31af84306b7569/uipath_core-0.5.28-py3-none-any.whl", hash = "sha256:b952a46a21710073cbc16d6d5684e9aa645c107f57a636b778cfb94aa81a1e48", size = 54980, upload-time = "2026-06-30T14:04:47.374Z" }, + { url = "https://files.pythonhosted.org/packages/1a/37/47a4e12bc7bdcfab4bd4e8e1f2a5c083bd59fca42fd431026a3f14c642a3/uipath_core-0.5.31-py3-none-any.whl", hash = "sha256:3dff5ecb236bf46af683c34d79bb2cf458a942ec041e1cbf7b4825ae246ff00c", size = 55040, upload-time = "2026-07-15T06:56:02.446Z" }, ] [[package]] name = "uipath-runtime" -version = "0.12.3" +version = "0.12.4" source = { editable = "." } dependencies = [ { name = "chardet" }, @@ -1179,7 +1179,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "chardet", specifier = ">=5.2.0,<8.0" }, - { name = "uipath-core", specifier = ">=0.5.28,<0.6.0" }, + { name = "uipath-core", specifier = ">=0.5.31,<0.6.0" }, { name = "vadersentiment", specifier = ">=3.3.2,<4.0" }, ]