Merge pull request #95 from UiPath/fix/reload-factory-after-login

fix: reload runtime factory after login

Author: cristipufu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-dev"
-version = "0.0.68"
+version = "0.0.69"
 description = "UiPath Developer Console"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/dev/server/app.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
 from pathlib import Path
@@ -160,12 +161,37 @@ async def _config():
     from uipath.dev.server.ws.handler import router as ws_router
 
     if auth_enabled:
-        from uipath.dev.server.auth import restore_session
+        from uipath.dev.server.auth import get_auth_state, restore_session
         from uipath.dev.server.routes.auth import router as auth_router
 
         app.include_router(auth_router, prefix="/api")
         restore_session()
 
+        # Reload the runtime factory when authentication completes so the
+        # newly-written credentials are picked up by subsequent runs.
+        def _on_authenticated() -> None:
+            async def _safe_reload() -> None:
+                # Wait for active runs to finish before reloading
+                while any(
+                    r.status in ("pending", "running")
+                    for r in server.run_service.runs.values()
+                ):
+                    await asyncio.sleep(1)
+                await server.reload_factory()
+
+            def _on_reload_done(t: asyncio.Task[None]) -> None:
+                try:
+                    t.result()
+                except asyncio.CancelledError:
+                    pass
+                except Exception:
+                    logger.exception("Factory reload after login failed")
+
+            task = asyncio.create_task(_safe_reload())
+            task.add_done_callback(_on_reload_done)
+
+        get_auth_state()._on_authenticated = _on_authenticated
+
     app.include_router(entrypoints_router, prefix="/api")
     app.include_router(runs_router, prefix="/api")
     app.include_router(graph_router, prefix="/api")

src/uipath/dev/server/auth.py

@@ -18,6 +18,7 @@
 import socketserver
 import threading
 import time
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
@@ -109,6 +110,8 @@ class AuthState:
     _last_tenant: str | None = None
     _last_org: dict[str, str] = field(default_factory=dict)
     _last_environment: str | None = None
+    # Callback invoked after successful authentication
+    _on_authenticated: Callable[[], None] | None = None
     # internal
     _code_verifier: str | None = None
     _state: str | None = None
@@ -128,9 +131,30 @@ def get_auth_state() -> AuthState:
 
 
 def reset_auth_state() -> None:
-    """Reset the auth state to its initial (unauthenticated) values."""
-    global _auth
-    _auth = AuthState()
+    """Reset the auth state to its initial (unauthenticated) values.
+
+    Preserves registered callbacks (e.g. ``_on_authenticated``).
+    """
+    _auth.status = "unauthenticated"
+    _auth.environment = "cloud"
+    _auth.token_data = {}
+    _auth.tenants = []
+    _auth.organization = {}
+    _auth.uipath_url = None
+    _auth._last_tenant = None
+    _auth._last_org = {}
+    _auth._last_environment = None
+    _auth._code_verifier = None
+    _auth._state = None
+    _auth._port = None
+    _auth._callback_server = None
+    if _auth._wait_task and not _auth._wait_task.done():
+        _auth._wait_task.cancel()
+    _auth._wait_task = None
+    if _auth._token_event:
+        _auth._token_event.set()
+    _auth._token_event = None
+    _auth._loop = None
 
 
 # ---------------------------------------------------------------------------
@@ -990,6 +1014,40 @@ def select_tenant(tenant_name: str) -> dict[str, Any]:
     return {"status": "authenticated", "uipath_url": auth.uipath_url}
 
 
+def _update_env_file(env_contents: dict[str, str]) -> None:
+    """Merge *env_contents* into the CWD ``.env`` file.
+
+    New keys take priority; existing keys not in *env_contents* are preserved.
+    Comments and blank lines are kept as-is.
+    """
+    env_path = Path.cwd() / ".env"
+    lines: list[str] = []
+    seen_keys: set[str] = set()
+
+    if env_path.exists():
+        with open(env_path) as f:
+            for raw_line in f:
+                stripped = raw_line.strip()
+                if stripped.startswith("#") or "=" not in stripped:
+                    # Preserve comments and blank lines
+                    lines.append(raw_line)
+                    continue
+                key = stripped.split("=", 1)[0]
+                if key in env_contents:
+                    lines.append(f"{key}={env_contents[key]}\n")
+                else:
+                    lines.append(raw_line)
+                seen_keys.add(key)
+
+    # Append new keys that weren't already in the file
+    for key, value in env_contents.items():
+        if key not in seen_keys:
+            lines.append(f"{key}={value}\n")
+
+    with open(env_path, "w") as f:
+        f.writelines(lines)
+
+
 def _finalize_tenant(auth: AuthState, tenant_name: str) -> None:
     """Write .env and os.environ with the resolved credentials."""
     org_name = auth.organization.get("name", "")
@@ -1010,45 +1068,27 @@ def _finalize_tenant(auth: AuthState, tenant_name: str) -> None:
     auth._last_org = dict(auth.organization)
     auth._last_environment = auth.environment
 
-    # Update os.environ
-    os.environ["UIPATH_ACCESS_TOKEN"] = access_token
-    os.environ["UIPATH_URL"] = uipath_url
-    os.environ["UIPATH_TENANT_ID"] = tenant_id
-    os.environ["UIPATH_ORGANIZATION_ID"] = org_id
-
-    # Write/update .env file (preserving comments, blank lines, and ordering)
-    env_path = Path.cwd() / ".env"
-    lines: list[str] = []
-    updated_keys: set[str] = set()
-    new_values = {
-        "UIPATH_ACCESS_TOKEN": access_token,
-        "UIPATH_URL": uipath_url,
-        "UIPATH_TENANT_ID": tenant_id,
-        "UIPATH_ORGANIZATION_ID": org_id,
-    }
-
-    if env_path.exists():
-        with open(env_path) as f:
-            for raw_line in f:
-                stripped = raw_line.strip()
-                if "=" in stripped and not stripped.startswith("#"):
-                    key = stripped.split("=", 1)[0]
-                    if key in new_values:
-                        lines.append(f"{key}={new_values[key]}\n")
-                        updated_keys.add(key)
-                        continue
-                lines.append(raw_line)
-
-    # Append any keys that weren't already in the file
-    for key, value in new_values.items():
-        if key not in updated_keys:
-            lines.append(f"{key}={value}\n")
+    # Write .env using the same approach as `uipath auth`
+    _update_env_file(
+        {
+            "UIPATH_ACCESS_TOKEN": access_token,
+            "UIPATH_URL": uipath_url,
+            "UIPATH_TENANT_ID": tenant_id,
+            "UIPATH_ORGANIZATION_ID": org_id,
+        }
+    )
 
-    with open(env_path, "w") as f:
-        f.writelines(lines)
+    # Reload .env into os.environ (same as CLI root: cwd + override)
+    load_dotenv(
+        dotenv_path=os.path.join(os.getcwd(), ".env"),
+        override=True,
+    )
 
-    # Reload all .env variables into os.environ
-    load_dotenv(override=True)
+    if auth._on_authenticated:
+        try:
+            auth._on_authenticated()
+        except Exception:
+            logger.exception("Error in post-authentication callback")
 
 
 def logout() -> None: