diff --git a/pyproject.toml b/pyproject.toml index 74c3909..4acf110 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/uipath/dev/server/app.py b/src/uipath/dev/server/app.py index 9b2f8e8..2da16f6 100644 --- a/src/uipath/dev/server/app.py +++ b/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") diff --git a/src/uipath/dev/server/auth.py b/src/uipath/dev/server/auth.py index 4159342..fd6dd6d 100644 --- a/src/uipath/dev/server/auth.py +++ b/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: diff --git a/uv.lock b/uv.lock index 14d0d8b..61fcbf4 100644 --- a/uv.lock +++ b/uv.lock @@ -2266,7 +2266,7 @@ wheels = [ [[package]] name = "uipath-dev" -version = "0.0.68" +version = "0.0.69" source = { editable = "." } dependencies = [ { name = "fastapi" },