fix: harden managed workspace lifecycle

CODEOWNERS

@@ -1 +1 @@
-* @UiPath/team-coded-agents
+* @UiPath/agents @cristipufu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-runtime"
-version = "0.12.6"
+version = "0.12.7"
 description = "Runtime abstractions and interfaces for building agents and automation scripts in the UiPath ecosystem"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/runtime/workspace/context.py

@@ -3,6 +3,7 @@
 from collections.abc import Iterator
 from contextlib import contextmanager
 from contextvars import ContextVar
+from dataclasses import dataclass
 from pathlib import Path
 
 from uipath.runtime.errors import (
@@ -11,8 +12,17 @@
     UiPathRuntimeError,
 )
 
-_workspace_path: ContextVar[Path | None] = ContextVar(
-    "uipath_workspace_path", default=None
+
+@dataclass
+class _WorkspaceExecution:
+    """Revocable workspace state shared with tasks created during execution."""
+
+    path: Path
+    active: bool = True
+
+
+_workspace_execution_state: ContextVar[_WorkspaceExecution | None] = ContextVar(
+    "uipath_workspace_execution", default=None
 )
 
 
@@ -27,26 +37,49 @@ def get_workspace_path() -> Path:
     Raises:
         UiPathRuntimeError: If called outside a managed runtime execution.
     """
-    workspace_path = _workspace_path.get()
-    if workspace_path is None:
+    execution = _workspace_execution_state.get()
+    if execution is None or not execution.active:
         raise UiPathRuntimeError(
             code=UiPathErrorCode.MANAGED_WORKSPACE_UNAVAILABLE,
             title="Managed Workspace Unavailable",
             detail=(
-                "No managed workspace is available outside a runtime execution. "
-                "Call get_workspace_path() only from a graph node or tool."
+                "No managed workspace is available in the current code path. "
+                "Call get_workspace_path() only from code invoked by a managed runtime, "
+                "not from module-level initialization or from the caller that invokes "
+                "or consumes the execution."
             ),
             category=UiPathErrorCategory.USER,
             include_traceback=False,
         )
-    return workspace_path
+    return execution.path
 
 
 @contextmanager
-def _workspace_execution(path: Path) -> Iterator[None]:
-    """Make ``path`` available for one managed runtime execution."""
-    token = _workspace_path.set(path)
+def _bind_workspace_execution(execution: _WorkspaceExecution) -> Iterator[None]:
+    """Make an active execution available in the current context."""
+    token = _workspace_execution_state.set(execution)
     try:
         yield
     finally:
-        _workspace_path.reset(token)
+        _workspace_execution_state.reset(token)
+
+
+def _create_workspace_execution(path: Path) -> _WorkspaceExecution:
+    """Create state for a workspace execution."""
+    return _WorkspaceExecution(path)
+
+
+def _revoke_workspace_execution(execution: _WorkspaceExecution) -> None:
+    """Prevent all contexts that inherited ``execution`` from accessing it."""
+    execution.active = False
+
+
+@contextmanager
+def _workspace_execution(path: Path) -> Iterator[None]:
+    """Make ``path`` available for one non-streaming execution."""
+    execution = _create_workspace_execution(path)
+    try:
+        with _bind_workspace_execution(execution):
+            yield
+    finally:
+        _revoke_workspace_execution(execution)

src/uipath/runtime/workspace/hydration.py

@@ -11,7 +11,12 @@
 from uipath.runtime.events import UiPathRuntimeEvent
 from uipath.runtime.result import UiPathRuntimeResult, UiPathRuntimeStatus
 from uipath.runtime.schema import UiPathRuntimeSchema
-from uipath.runtime.workspace.context import _workspace_execution
+from uipath.runtime.workspace.context import (
+    _bind_workspace_execution,
+    _create_workspace_execution,
+    _revoke_workspace_execution,
+    _workspace_execution,
+)
 from uipath.runtime.workspace.hydrator import WorkspaceHydrator
 from uipath.runtime.workspace.registry_store import WorkspaceRegistryStore
 from uipath.runtime.workspace.workspace import Workspace
@@ -103,24 +108,44 @@ async def stream(
         options: UiPathStreamOptions | None = None,
     ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
         """Hydrate, stream delegate events, then persist files according to policy."""
-        with _workspace_execution(self.workspace.path):
-            await self._hydrate()
-            final_result: UiPathRuntimeResult | None = None
+        execution = _create_workspace_execution(self.workspace.path)
+        delegate_stream: AsyncGenerator[UiPathRuntimeEvent, None] | None = None
+        final_result: UiPathRuntimeResult | None = None
+        hydrated = False
 
-            try:
-                async for event in self.delegate.stream(input, options=options):
-                    if isinstance(event, UiPathRuntimeResult):
-                        final_result = event
-                    else:
-                        yield event
-            except Exception:
-                if self.policy == HydrationPolicy.ALWAYS:
+        try:
+            with _bind_workspace_execution(execution):
+                await self._hydrate()
+                hydrated = True
+                delegate_stream = self.delegate.stream(input, options=options)
+
+                event = await anext(delegate_stream, None)
+            while event is not None:
+                if isinstance(event, UiPathRuntimeResult):
+                    final_result = event
+                else:
+                    yield event
+                with _bind_workspace_execution(execution):
+                    event = await anext(delegate_stream, None)
+        except Exception:
+            if self.policy == HydrationPolicy.ALWAYS and hydrated:
+                with _bind_workspace_execution(execution):
                     await self._persist()
-                raise
-
+            raise
+        else:
             if final_result is not None:
-                await self._dehydrate(final_result)
-                yield final_result
+                with _bind_workspace_execution(execution):
+                    await self._dehydrate(final_result)
+        finally:
+            try:
+                if delegate_stream is not None:
+                    with _bind_workspace_execution(execution):
+                        await delegate_stream.aclose()
+            finally:
+                _revoke_workspace_execution(execution)
+
+        if final_result is not None:
+            yield final_result
 
     async def get_schema(self) -> UiPathRuntimeSchema:
         """Passthrough schema from delegate runtime."""

src/uipath/runtime/workspace/workspace.py

@@ -44,5 +44,7 @@ def create(
 
     async def dispose(self) -> None:
         """Remove the workspace directory when configured to do so."""
-        if self.cleanup and self.path.exists():
-            await asyncio.to_thread(shutil.rmtree, self.path)
+        if not self.cleanup:
+            return
+
+        await asyncio.to_thread(shutil.rmtree, self.path)