feat: workspace persistence across conversational exchanges via content-parts

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-runtime"
-version = "0.12.1"
+version = "0.12.2"
 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"

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

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

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

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

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]],