feat(cli): keep coded agents' workspace files across suspend and resume [PC-5359]

packages/uipath/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.15.0"
+version = "2.15.1"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath/src/uipath/_cli/cli_debug.py

@@ -10,6 +10,7 @@
 from uipath._cli._utils._debug import setup_debugging
 from uipath._cli._utils._studio_project import StudioClient
 from uipath._cli._utils._tracing import create_trace_manager
+from uipath._utils._managed_workspace import wrap_with_managed_workspace
 from uipath.eval.mocks import SimulationConfig, UiPathMockRuntime, build_mocking_context
 from uipath.eval.mocks._mock_runtime import load_simulation_config
 from uipath.platform.common import (
@@ -228,6 +229,15 @@ async def execute_debug_runtime():
                                     runtime_id=governance_runtime_id,
                                 )
 
+                            runtime = await wrap_with_managed_workspace(
+                                runtime,
+                                factory=factory,
+                                settings=factory_settings,
+                                runtime_id=governance_runtime_id,
+                                job_id=ctx.job_id,
+                                conversation_id=ctx.conversation_id,
+                            )
+
                             delegate = runtime
                             if ctx.conversation_id and ctx.exchange_id:
                                 chat_bridge: UiPathChatProtocol = get_chat_bridge(

packages/uipath/src/uipath/_cli/cli_run.py

@@ -9,6 +9,7 @@
 from uipath._cli._utils._common import read_resource_overwrites_from_file
 from uipath._cli._utils._debug import setup_debugging
 from uipath._cli._utils._tracing import create_trace_manager
+from uipath._utils._managed_workspace import wrap_with_managed_workspace
 from uipath.eval.mocks import SimulationConfig, UiPathMockRuntime, build_mocking_context
 from uipath.platform.common import (
     ExecutionSourceContext,
@@ -306,6 +307,15 @@ async def execute() -> None:
                                     runtime_id=governance_runtime_id,
                                 )
 
+                            base_runtime = await wrap_with_managed_workspace(
+                                base_runtime,
+                                factory=factory,
+                                settings=factory_settings,
+                                runtime_id=governance_runtime_id,
+                                job_id=ctx.job_id,
+                                conversation_id=ctx.conversation_id,
+                            )
+
                             runtime = base_runtime
 
                             if simulation_config:

packages/uipath/src/uipath/_utils/_managed_workspace.py

@@ -0,0 +1,73 @@
+"""Host-managed workspace persistence for runtimes that opt in.
+
+A factory opts in through ``UiPathRuntimeFactorySettings.managed_workspace``.
+The host then gives each job run a workspace folder, exposed to agent code
+through ``uipath.runtime.get_workspace_path()``, and persists the folder as
+job attachments across suspend and resume.
+"""
+
+from __future__ import annotations
+
+from uipath.platform import UiPath
+from uipath.platform.common import UiPathConfig
+from uipath.runtime import (
+    HydrationPolicy,
+    ManagedWorkspace,
+    UiPathRuntimeFactoryProtocol,
+    UiPathRuntimeFactorySettings,
+    UiPathRuntimeProtocol,
+    WorkspaceRegistryStore,
+)
+
+__all__ = ["wrap_with_managed_workspace"]
+
+
+async def wrap_with_managed_workspace(
+    runtime: UiPathRuntimeProtocol,
+    *,
+    factory: UiPathRuntimeFactoryProtocol,
+    settings: UiPathRuntimeFactorySettings | None,
+    runtime_id: str,
+    job_id: str | None,
+    conversation_id: str | None = None,
+) -> UiPathRuntimeProtocol:
+    """Wrap a runtime so its workspace survives suspend and resume.
+
+    The runtime is returned unchanged unless the factory opts in, the run is a
+    job, and the factory has storage for the workspace registry. Files are
+    uploaded when the run suspends, and also when a conversational exchange
+    ends successfully, because the next exchange resumes the same job.
+
+    Args:
+        runtime: Runtime returned by the factory.
+        factory: Factory that created the runtime; provides registry storage.
+        settings: The factory's settings.
+        runtime_id: Runtime id the runtime was created with.
+        job_id: Orchestrator job key, or ``None`` outside a job.
+        conversation_id: Conversation id for conversational runs.
+
+    Returns:
+        The wrapped runtime, or ``runtime`` when no workspace applies.
+    """
+    if settings is None or not settings.managed_workspace or not job_id:
+        return runtime
+
+    storage = await factory.get_storage()
+    if storage is None:
+        return runtime
+
+    managed = ManagedWorkspace(
+        services=UiPath,
+        runtime_id=runtime_id,
+        job_key=job_id,
+        folder_key=UiPathConfig.folder_key,
+        registry_store=WorkspaceRegistryStore(storage, runtime_id),
+    )
+    return managed.wrap(
+        runtime,
+        policy=(
+            HydrationPolicy.SUSPEND_OR_SUCCESS
+            if conversation_id
+            else HydrationPolicy.SUSPEND_ONLY
+        ),
+    )

packages/uipath/src/uipath/eval/runtime/runtime.py

@@ -27,12 +27,14 @@
 )
 from pydantic import BaseModel, ValidationError
 
+from uipath._utils._managed_workspace import wrap_with_managed_workspace
 from uipath.core.events import EventBus
 from uipath.core.tracing import UiPathTraceManager
 from uipath.runtime import (
     UiPathExecuteOptions,
     UiPathExecutionRuntime,
     UiPathRuntimeFactoryProtocol,
+    UiPathRuntimeFactorySettings,
     UiPathRuntimeResult,
     UiPathRuntimeStatus,
     UiPathRuntimeStorageProtocol,
@@ -348,12 +350,14 @@ def __init__(
         self.coverage = coverage.Coverage(branch=True)
 
         self._storage: UiPathRuntimeStorageProtocol | None = None
+        self._factory_settings: UiPathRuntimeFactorySettings | None = None
 
     async def __aenter__(self) -> UiPathEvalRuntime:
         """Async context manager entry - initialize storage and start coverage if enabled."""
         if self.context.report_coverage:
             self.coverage.start()
         self._storage = await self.factory.get_storage()
+        self._factory_settings = await self.factory.get_settings()
 
         return self
 
@@ -399,7 +403,7 @@ async def initiate_evaluation(
                 f"Please run with a single evaluation using --eval-ids to specify one evaluation."
             )
 
-        factory_settings = await self.factory.get_settings()
+        factory_settings = self._factory_settings
         agent_type = factory_settings.agent_type if factory_settings else None
 
         await self.event_bus.publish(
@@ -1018,6 +1022,13 @@ async def execute_runtime(
                 entrypoint=self.context.entrypoint or "",
                 runtime_id=runtime_id,
             )
+            eval_runtime = await wrap_with_managed_workspace(
+                eval_runtime,
+                factory=self.factory,
+                settings=self._factory_settings,
+                runtime_id=runtime_id,
+                job_id=self.context.job_id,
+            )
 
             delegate = eval_runtime
             if mocking_context:

packages/uipath/tests/cli/conftest.py

@@ -3,6 +3,7 @@
 import pytest
 
 from tests.cli.utils.project_details import ProjectDetails
+from tests.cli.utils.workspace_fakes import FakeSdk
 
 
 @pytest.fixture
@@ -29,3 +30,10 @@ def project_details() -> ProjectDetails:
         with open("tests/cli/mocks/pyproject.toml", "r") as file:
             data = file.read()
     return ProjectDetails.from_toml(data)
+
+
+@pytest.fixture
+def sdk(monkeypatch: pytest.MonkeyPatch) -> FakeSdk:
+    fake = FakeSdk()
+    monkeypatch.setattr("uipath._utils._managed_workspace.UiPath", lambda: fake)
+    return fake

packages/uipath/tests/cli/eval/test_eval_runtime_managed_workspace.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import uuid
+from pathlib import Path
+from typing import Any
+
+from tests.cli.utils.workspace_fakes import (
+    OPTED_IN,
+    REGISTRY,
+    FakeSdk,
+    FileRuntime,
+    MemoryStorage,
+    WorkspaceFactory,
+)
+from uipath.core.events import EventBus
+from uipath.core.tracing import UiPathTraceManager
+from uipath.eval.helpers import EvalHelpers
+from uipath.eval.runtime import UiPathEvalContext, evaluate
+from uipath.eval.runtime.events import EvaluationEvents
+from uipath.runtime import UiPathRuntimeFactorySettings, UiPathRuntimeStatus
+
+JOB_ID = str(uuid.uuid4())
+EVAL_SET = Path(__file__).parent / "evals" / "eval-sets" / "default.json"
+
+
+async def _eval_context(items: int = 1) -> UiPathEvalContext:
+    evaluation_set, _ = EvalHelpers.load_eval_set(str(EVAL_SET))
+    first = evaluation_set.evaluations[0]
+    evaluation_set.evaluations = [first] + [
+        first.model_copy(update={"id": str(uuid.uuid4())}) for _ in range(items - 1)
+    ]
+    context = UiPathEvalContext()
+    context.execution_id = str(uuid.uuid4())
+    context.job_id = JOB_ID
+    context.evaluation_set = evaluation_set
+    context.runtime_schema = await FileRuntime().get_schema()
+    context.evaluators = await EvalHelpers.load_evaluators(
+        str(EVAL_SET), evaluation_set, agent_model=None
+    )
+    return context
+
+
+def _suspending_factory(
+    storage: MemoryStorage,
+    settings: UiPathRuntimeFactorySettings | None = OPTED_IN,
+) -> WorkspaceFactory:
+    return WorkspaceFactory(
+        FileRuntime(write="draft", status=UiPathRuntimeStatus.SUSPENDED),
+        storage=storage,
+        settings=settings,
+    )
+
+
+async def test_eval_uploads_each_suspended_item_workspace(sdk: FakeSdk) -> None:
+    context = await _eval_context(items=2)
+    storage = MemoryStorage()
+
+    await evaluate(
+        _suspending_factory(storage), UiPathTraceManager(), context, EventBus()
+    )
+
+    item_ids = [item.id for item in context.evaluation_set.evaluations]
+    assert list(sdk.attachments.blobs.values()) == [b"draft"] * len(item_ids)
+    assert {key[0] for key in storage.values if key[1:] == REGISTRY} == set(item_ids)
+
+
+async def test_eval_set_run_carries_the_factory_agent_type(sdk: FakeSdk) -> None:
+    created: list[Any] = []
+    event_bus = EventBus()
+    event_bus.subscribe(EvaluationEvents.CREATE_EVAL_SET_RUN, created.append)
+    settings = UiPathRuntimeFactorySettings(
+        managed_workspace=True, agent_type="uipath_coded"
+    )
+
+    await evaluate(
+        _suspending_factory(MemoryStorage(), settings),
+        UiPathTraceManager(),
+        await _eval_context(),
+        event_bus,
+    )
+    await event_bus.wait_for_all()
+
+    assert [event.agent_type for event in created] == ["uipath_coded"]

packages/uipath/tests/cli/test_managed_workspace.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from uuid import UUID, uuid4
+
+import pytest
+
+from tests.cli.utils.workspace_fakes import (
+    OPTED_IN,
+    FakeSdk,
+    FileRuntime,
+    MemoryStorage,
+    WorkspaceFactory,
+)
+from uipath._utils._managed_workspace import wrap_with_managed_workspace
+from uipath.platform.constants import ENV_FOLDER_KEY
+from uipath.runtime import (
+    UiPathRuntimeFactorySettings,
+    UiPathRuntimeResult,
+    UiPathRuntimeStatus,
+)
+
+JOB_ID = str(uuid4())
+
+
+async def _run(runtime: FileRuntime, storage: MemoryStorage) -> UiPathRuntimeResult:
+    wrapped = await wrap_with_managed_workspace(
+        runtime,
+        factory=WorkspaceFactory(runtime, storage=storage),
+        settings=OPTED_IN,
+        runtime_id="runtime-1",
+        job_id=JOB_ID,
+    )
+    try:
+        return await wrapped.execute()
+    finally:
+        await wrapped.dispose()
+
+
+async def test_uploads_are_linked_to_the_job(sdk: FakeSdk) -> None:
+    await _run(
+        FileRuntime(write="draft", status=UiPathRuntimeStatus.SUSPENDED),
+        MemoryStorage(),
+    )
+
+    assert [job for job, _ in sdk.jobs.links] == [UUID(JOB_ID)]
+    assert list(sdk.attachments.blobs.values()) == [b"draft"]
+
+
+async def test_attachment_calls_use_the_job_folder(
+    sdk: FakeSdk, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv(ENV_FOLDER_KEY, "folder-1")
+    storage = MemoryStorage()
+
+    await _run(
+        FileRuntime(write="draft", status=UiPathRuntimeStatus.SUSPENDED), storage
+    )
+    await _run(FileRuntime(), storage)
+
+    assert sdk.attachments.folder_keys
+    assert sdk.jobs.folder_keys
+    assert set(sdk.attachments.folder_keys + sdk.jobs.folder_keys) == {"folder-1"}
+
+
+async def test_successful_run_outside_a_conversation_uploads_nothing(
+    sdk: FakeSdk,
+) -> None:
+    await _run(FileRuntime(write="final"), MemoryStorage())
+
+    assert sdk.attachments.blobs == {}
+
+
+@pytest.mark.parametrize(
+    ("settings", "job_id", "storage"),
+    [
+        pytest.param(None, JOB_ID, MemoryStorage(), id="factory-without-settings"),
+        pytest.param(
+            UiPathRuntimeFactorySettings(),
+            JOB_ID,
+            MemoryStorage(),
+            id="factory-not-opted-in",
+        ),
+        pytest.param(OPTED_IN, None, MemoryStorage(), id="not-a-job"),
+        pytest.param(OPTED_IN, JOB_ID, None, id="factory-without-storage"),
+    ],
+)
+async def test_runtime_is_left_unwrapped(
+    settings: UiPathRuntimeFactorySettings | None,
+    job_id: str | None,
+    storage: MemoryStorage | None,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    def no_sdk() -> FakeSdk:
+        raise AssertionError("the SDK must not be built for an unwrapped runtime")
+
+    monkeypatch.setattr("uipath._utils._managed_workspace.UiPath", no_sdk)
+    runtime = FileRuntime()
+
+    wrapped = await wrap_with_managed_workspace(
+        runtime,
+        factory=WorkspaceFactory(runtime, storage=storage),
+        settings=settings,
+        runtime_id="runtime-1",
+        job_id=job_id,
+    )
+
+    assert wrapped is runtime