refactor: scope execution source via ExecutionSourceContext manager

Per review feedback, set the execution source with a context manager
that releases the ContextVar token on exit (like ResourceOverwritesContext
and the runtime ctx), instead of a bare setter. CLI handlers enter
'with ExecutionSourceContext(ctx.execution_source), ctx:'; dev wraps the
console/server run. Adds a dev terminal test covering the scoped run.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: Valentina Bojan

packages/uipath-platform/src/uipath/platform/common/__init__.py

@@ -16,7 +16,7 @@
 )
 from ._config import UiPathApiConfig, UiPathConfig
 from ._endpoints_manager import EndpointManager
-from ._execution_context import UiPathExecutionContext, set_execution_source
+from ._execution_context import ExecutionSourceContext, UiPathExecutionContext
 from ._external_application_service import ExternalApplicationService
 from ._folder_context import FolderContext, header_folder
 from ._http_config import get_ca_bundle_path, get_httpx_client_kwargs
@@ -61,7 +61,7 @@
     "BaseService",
     "UiPathApiConfig",
     "UiPathExecutionContext",
-    "set_execution_source",
+    "ExecutionSourceContext",
     "ExternalApplicationService",
     "FolderContext",
     "TokenData",

packages/uipath-platform/src/uipath/platform/common/_execution_context.py

@@ -1,4 +1,4 @@
-from contextvars import ContextVar
+from contextvars import ContextVar, Token
 from os import environ as env
 from typing import Optional
 
@@ -9,15 +9,28 @@
 )
 
 
-def set_execution_source(execution_source: Optional[str]) -> None:
-    """Set the execution source for the current run.
+class ExecutionSourceContext:
+    """Scope the execution source for the duration of a run.
 
     Carries the source (e.g. ``runtime``/``playground``/``eval``) via a context
-    variable instead of a process-global env var. The CLI sets this from
+    variable instead of a process-global env var, and releases it on exit so it
+    stays correctly scoped in concurrent runs. The CLI enters this with
     ``UiPathRuntimeContext.execution_source`` so platform clients can read it
     via :attr:`UiPathExecutionContext.execution_source`.
     """
-    _execution_source.set(execution_source)
+
+    def __init__(self, execution_source: Optional[str]) -> None:
+        self._execution_source = execution_source
+        self._token: Optional[Token[Optional[str]]] = None
+
+    def __enter__(self) -> "ExecutionSourceContext":
+        self._token = _execution_source.set(self._execution_source)
+        return self
+
+    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
+        if self._token is not None:
+            _execution_source.reset(self._token)
+            self._token = None
 
 
 class UiPathExecutionContext:

packages/uipath-platform/tests/common/test_execution_context.py

@@ -1,32 +1,26 @@
-import pytest
-
-from uipath.platform.common import UiPathExecutionContext, set_execution_source
-
-
-@pytest.fixture(autouse=True)
-def _reset_execution_source() -> object:
-    set_execution_source(None)
-    yield
-    set_execution_source(None)
+from uipath.platform.common import ExecutionSourceContext, UiPathExecutionContext
 
 
 def test_execution_source_none_by_default() -> None:
     assert UiPathExecutionContext().execution_source is None
 
 
-def test_execution_source_reflects_set_value() -> None:
+def test_execution_source_set_within_context() -> None:
     ctx = UiPathExecutionContext()
 
-    set_execution_source("runtime")
+    with ExecutionSourceContext("runtime"):
+        assert ctx.execution_source == "runtime"
 
-    assert ctx.execution_source == "runtime"
+    assert ctx.execution_source is None
 
 
-def test_execution_source_can_be_cleared() -> None:
+def test_execution_source_context_restores_previous_value() -> None:
     ctx = UiPathExecutionContext()
 
-    set_execution_source("eval")
-    assert ctx.execution_source == "eval"
+    with ExecutionSourceContext("eval"):
+        assert ctx.execution_source == "eval"
+        with ExecutionSourceContext("playground"):
+            assert ctx.execution_source == "playground"
+        assert ctx.execution_source == "eval"
 
-    set_execution_source(None)
     assert ctx.execution_source is None

packages/uipath-platform/tests/services/test_guardrails_service.py

@@ -10,7 +10,7 @@
 )
 
 from uipath.platform import UiPathApiConfig, UiPathExecutionContext
-from uipath.platform.common import set_execution_source
+from uipath.platform.common import ExecutionSourceContext
 from uipath.platform.guardrails import (
     BuiltInValidatorGuardrail,
     EnumListParameterValue,
@@ -19,13 +19,6 @@
 )
 
 
-@pytest.fixture(autouse=True)
-def _reset_execution_source() -> object:
-    set_execution_source(None)
-    yield
-    set_execution_source(None)
-
-
 @pytest.fixture
 def service(
     config: UiPathApiConfig,
@@ -405,8 +398,8 @@ def capture_request(request):
                 validator_parameters=[],
             )
 
-            set_execution_source("runtime")
-            service.evaluate_guardrail("test input", pii_guardrail)
+            with ExecutionSourceContext("runtime"):
+                service.evaluate_guardrail("test input", pii_guardrail)
 
             assert captured_request is not None
             headers = dict(captured_request.headers)

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

@@ -12,9 +12,9 @@
 from uipath.eval.mocks import UiPathMockRuntime
 from uipath.eval.mocks._mock_runtime import load_simulation_config
 from uipath.platform.common import (
+    ExecutionSourceContext,
     ResourceOverwritesContext,
     UiPathConfig,
-    set_execution_source,
 )
 from uipath.runtime import (
     UiPathExecuteOptions,
@@ -134,8 +134,7 @@ async def execute_debug_runtime():
                     trace_manager=trace_manager,
                     command="debug",
                 )
-                set_execution_source(ctx.execution_source)
-                with ctx:
+                with ExecutionSourceContext(ctx.execution_source), ctx:
                     factory: UiPathRuntimeFactoryProtocol | None = None
 
                     try:

packages/uipath/src/uipath/_cli/cli_dev.py

@@ -7,19 +7,18 @@
 from uipath._cli._utils._debug import setup_debugging
 from uipath._cli.middlewares import Middlewares
 from uipath.core.tracing import UiPathTraceManager
-from uipath.platform.common import set_execution_source
+from uipath.platform.common import ExecutionSourceContext
 from uipath.runtime import UiPathRuntimeContext, UiPathRuntimeFactoryRegistry
 
 from ._telemetry import track_command
 
 console = ConsoleLogger()
 
 
-def _create_dev_factory(trace_manager: UiPathTraceManager):
-    """Build the dev runtime factory and propagate the execution source."""
+def _create_dev_context_and_factory(trace_manager: UiPathTraceManager):
+    """Build the dev runtime context and its factory."""
     context = UiPathRuntimeContext(trace_manager=trace_manager, command="dev")
-    set_execution_source(context.execution_source)
-    return UiPathRuntimeFactoryRegistry.get(context=context)
+    return context, UiPathRuntimeFactoryRegistry.get(context=context)
 
 
 def _check_dev_dependency(interface: str) -> None:
@@ -88,13 +87,14 @@ async def run_terminal() -> None:
             factory = None
             try:
                 trace_manager = UiPathTraceManager()
-                factory = _create_dev_factory(trace_manager)
+                context, factory = _create_dev_context_and_factory(trace_manager)
 
                 app = UiPathDeveloperConsole(
                     runtime_factory=factory, trace_manager=trace_manager
                 )
 
-                await app.run_async()
+                with ExecutionSourceContext(context.execution_source):
+                    await app.run_async()
 
             except KeyboardInterrupt:
                 console.info("Debug session interrupted by user")
@@ -128,7 +128,7 @@ def signal_handler(sig, frame):
 
             try:
                 trace_manager = UiPathTraceManager()
-                factory = _create_dev_factory(trace_manager)
+                context, factory = _create_dev_context_and_factory(trace_manager)
 
                 app = UiPathDeveloperServer(
                     runtime_factory=factory,
@@ -140,13 +140,17 @@ def signal_handler(sig, frame):
                     ),
                 )
 
-                server_task = asyncio.create_task(app.run_async())
-                shutdown_task = asyncio.create_task(shutdown_event.wait())
-
-                # Wait for either server to complete or shutdown signal
-                done, pending = await asyncio.wait(
-                    {server_task, shutdown_task}, return_when=asyncio.FIRST_COMPLETED
-                )
+                # Enter the execution source context before creating the server
+                # task so request tasks spawned during the run inherit it.
+                with ExecutionSourceContext(context.execution_source):
+                    server_task = asyncio.create_task(app.run_async())
+                    shutdown_task = asyncio.create_task(shutdown_event.wait())
+
+                    # Wait for either server to complete or shutdown signal
+                    done, pending = await asyncio.wait(
+                        {server_task, shutdown_task},
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
 
                 for task in pending:
                     task.cancel()

packages/uipath/src/uipath/_cli/cli_eval.py

@@ -22,9 +22,9 @@
 from uipath.eval.runtime import UiPathEvalContext, evaluate
 from uipath.platform.chat import set_llm_concurrency
 from uipath.platform.common import (
+    ExecutionSourceContext,
     ResourceOverwritesContext,
     UiPathConfig,
-    set_execution_source,
 )
 from uipath.runtime import (
     UiPathRuntimeContext,
@@ -319,8 +319,7 @@ async def execute_eval():
                     command="eval",
                     resume=resume,
                 )
-                set_execution_source(ctx.execution_source)
-                with ctx:
+                with ExecutionSourceContext(ctx.execution_source), ctx:
                     # Set job_id in eval context for single runtime runs
                     eval_context.job_id = ctx.job_id
 

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

@@ -10,9 +10,9 @@
 from uipath.core.tracing import UiPathTraceManager
 from uipath.eval.mocks import SimulationConfig, UiPathMockRuntime, build_mocking_context
 from uipath.platform.common import (
+    ExecutionSourceContext,
     ResourceOverwritesContext,
     UiPathConfig,
-    set_execution_source,
 )
 from uipath.runtime import (
     UiPathExecuteOptions,
@@ -205,8 +205,6 @@ async def execute() -> None:
                     trace_manager=trace_manager,
                 )
 
-                set_execution_source(ctx.execution_source)
-
                 if ctx.trace_file:
                     trace_manager.add_span_exporter(
                         JsonLinesFileExporter(ctx.trace_file)
@@ -215,7 +213,7 @@ async def execute() -> None:
                 async with ResourceOverwritesContext(
                     lambda: read_resource_overwrites_from_file(ctx.runtime_dir)
                 ):
-                    with ctx:
+                    with ExecutionSourceContext(ctx.execution_source), ctx:
                         base_runtime: UiPathRuntimeProtocol | None = None
                         runtime: UiPathRuntimeProtocol | None = None
                         chat_runtime: UiPathRuntimeProtocol | None = None

packages/uipath/tests/cli/test_dev.py

@@ -1,32 +1,70 @@
-from unittest.mock import MagicMock
+import sys
+import types
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
+from click.testing import CliRunner
 
-from uipath._cli import cli_dev
+from uipath._cli import cli, cli_dev
+from uipath._cli.middlewares import MiddlewareResult
+from uipath.platform.common import UiPathExecutionContext
 
 
-def test_create_dev_factory_propagates_execution_source(
+def test_create_dev_context_and_factory_uses_dev_command(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """_create_dev_factory sets the 'playground' source and returns the factory."""
-    captured: dict[str, object] = {}
-
-    def fake_set_execution_source(value: object) -> None:
-        captured["source"] = value
-
+    """The helper builds a 'dev' context (source 'playground') and its factory."""
     sentinel_factory = MagicMock(name="factory")
+    captured: dict[str, object] = {}
 
     def fake_get(context: object) -> object:
         captured["command"] = context.command  # type: ignore[attr-defined]
         return sentinel_factory
 
-    monkeypatch.setattr(cli_dev, "set_execution_source", fake_set_execution_source)
     monkeypatch.setattr(
         "uipath._cli.cli_dev.UiPathRuntimeFactoryRegistry.get", fake_get
     )
 
-    factory = cli_dev._create_dev_factory(None)  # type: ignore[arg-type]
+    context, factory = cli_dev._create_dev_context_and_factory(None)  # type: ignore[arg-type]
 
     assert factory is sentinel_factory
     assert captured["command"] == "dev"
-    assert captured["source"] == "playground"
+    assert context.execution_source == "playground"
+
+
+def test_dev_terminal_sets_execution_source_during_run(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Running `dev terminal` scopes the execution source to 'playground'."""
+    seen: dict[str, object] = {}
+
+    async def fake_run_async() -> None:
+        seen["source"] = UiPathExecutionContext().execution_source
+
+    fake_console = MagicMock()
+    fake_console.run_async = fake_run_async
+
+    fake_module = types.ModuleType("uipath.dev")
+    fake_module.UiPathDeveloperConsole = MagicMock(return_value=fake_console)  # type: ignore[attr-defined]
+    monkeypatch.setitem(sys.modules, "uipath.dev", fake_module)
+
+    mock_factory = MagicMock()
+    mock_factory.dispose = AsyncMock()
+
+    monkeypatch.setattr(cli_dev, "_check_dev_dependency", lambda interface: None)
+    monkeypatch.setattr(cli_dev, "setup_debugging", lambda debug, port: True)
+    monkeypatch.setattr(
+        "uipath._cli.cli_dev.Middlewares.next",
+        lambda *a, **k: MiddlewareResult(should_continue=True),
+    )
+    monkeypatch.setattr(
+        "uipath._cli.cli_dev.UiPathRuntimeFactoryRegistry.get",
+        lambda context: mock_factory,
+    )
+
+    result = CliRunner().invoke(cli, ["dev", "terminal"])
+
+    assert result.exit_code == 0, result.output
+    assert seen["source"] == "playground"
+    # token released once the run completes
+    assert UiPathExecutionContext().execution_source is None