From dd79b10e3c022b7a32aa5ee5acf961730aa087d3 Mon Sep 17 00:00:00 2001 From: haikdc Date: Sun, 5 Apr 2026 21:10:07 -0700 Subject: [PATCH] [Haik]: ckpt, made classes for NineRouterClient and NineRouterProcess, now gonna abstract the NineRouter dir a bit, then clean up the type specing/checking in all the classes --- .../subscriptions/NineRouter/NineRouter.py | 43 ++- .../NineRouter/NineRouterClient.py | 129 +++++---- .../NineRouterProcess/NineRouterProcess.py | 268 +++++++++--------- backend/apps/subscriptions/subscriptions.py | 3 +- 4 files changed, 216 insertions(+), 227 deletions(-) diff --git a/backend/apps/subscriptions/NineRouter/NineRouter.py b/backend/apps/subscriptions/NineRouter/NineRouter.py index cd215945..b7545182 100644 --- a/backend/apps/subscriptions/NineRouter/NineRouter.py +++ b/backend/apps/subscriptions/NineRouter/NineRouter.py @@ -2,7 +2,7 @@ All interaction with 9Router from the subscriptions subapp goes through NineRouter.get() so that mutable state (subprocess handle, background -ensure-task) lives in one place. +ensure-task, HTTP client) lives in one place. """ import asyncio @@ -10,25 +10,16 @@ from typing import ClassVar, Optional from typeguard import typechecked -from backend.apps.subscriptions.NineRouter.NineRouterProcess.NineRouterProcess import ( - is_running as _is_running, - ensure_running as _ensure_running, - stop as _stop, -) -from backend.apps.subscriptions.NineRouter.NineRouterClient import ( - get_providers as _get_providers, - get_models as _get_models, - start_oauth as _start_oauth, - poll_oauth as _poll_oauth, - exchange_oauth as _exchange_oauth, - disconnect_provider as _disconnect_provider, -) +from backend.apps.subscriptions.NineRouter.NineRouterProcess.NineRouterProcess import NineRouterProcess +from backend.apps.subscriptions.NineRouter.NineRouterClient import NineRouterClient class NineRouter: _instance: ClassVar[Optional["NineRouter"]] = None def __init__(self) -> None: + self._process: NineRouterProcess = NineRouterProcess() + self._client: NineRouterClient = NineRouterClient() self._ensure_task: Optional[asyncio.Task] = None @classmethod @@ -41,21 +32,23 @@ class NineRouter: @typechecked def is_running(self) -> bool: - return _is_running() + return self._process.is_running() @typechecked async def ensure_running(self) -> None: - await _ensure_running() + await self._process.ensure_running() @typechecked - def stop(self) -> None: - _stop() + async def stop(self) -> None: + self.cancel_ensure_task() + await self._client.aclose() + self._process.stop() @typechecked async def ensure_running_background(self) -> None: """Kick off ensure_running as a background task if not already in flight.""" if self._ensure_task is None or self._ensure_task.done(): - self._ensure_task = asyncio.create_task(_ensure_running()) + self._ensure_task = asyncio.create_task(self._process.ensure_running()) @typechecked def cancel_ensure_task(self) -> None: @@ -66,15 +59,15 @@ class NineRouter: @typechecked async def get_providers(self) -> list[dict] | dict: - return await _get_providers() + return await self._client.get_providers() @typechecked async def get_models(self) -> list[dict]: - return await _get_models() + return await self._client.get_models() @typechecked async def start_oauth(self, provider: str) -> dict: - return await _start_oauth(provider) + return await self._client.start_oauth(provider) @typechecked async def poll_oauth( @@ -84,7 +77,7 @@ class NineRouter: code_verifier: str | None = None, extra_data: dict | None = None, ) -> dict: - return await _poll_oauth(provider, device_code, code_verifier=code_verifier, extra_data=extra_data) + return await self._client.poll_oauth(provider, device_code, code_verifier=code_verifier, extra_data=extra_data) @typechecked async def exchange_oauth( @@ -95,8 +88,8 @@ class NineRouter: code_verifier: str, state: str = "", ) -> dict: - return await _exchange_oauth(provider, code, redirect_uri, code_verifier, state) + return await self._client.exchange_oauth(provider, code, redirect_uri, code_verifier, state) @typechecked async def disconnect_provider(self, provider_id: str) -> bool: - return await _disconnect_provider(provider_id) + return await self._client.disconnect_provider(provider_id) diff --git a/backend/apps/subscriptions/NineRouter/NineRouterClient.py b/backend/apps/subscriptions/NineRouter/NineRouterClient.py index 416c6594..5108b2bb 100644 --- a/backend/apps/subscriptions/NineRouter/NineRouterClient.py +++ b/backend/apps/subscriptions/NineRouter/NineRouterClient.py @@ -6,28 +6,33 @@ from typeguard import typechecked from backend.apps.subscriptions.NineRouter.constants import NINE_ROUTER_API, NINE_ROUTER_V1 from backend.ports import NINE_ROUTER_PORT -@typechecked -async def get_providers() -> list[dict] | dict: - try: - async with httpx.AsyncClient(timeout=5.0) as client: - r = await client.get(f"{NINE_ROUTER_API}/providers") + +class NineRouterClient: + def __init__(self) -> None: + self._http: httpx.AsyncClient = httpx.AsyncClient(timeout=15.0) + + async def aclose(self) -> None: + await self._http.aclose() + + @typechecked + async def get_providers(self) -> list[dict] | dict: + try: + r = await self._http.get(f"{NINE_ROUTER_API}/providers", timeout=5.0) if r.status_code == 200: return r.json() - except Exception as e: - print(f"9Router providers fetch failed: {e}") - return [] + except Exception as e: + print(f"9Router providers fetch failed: {e}") + return [] + @typechecked + async def start_oauth(self, provider: str) -> dict: + """Start OAuth flow for a provider. -@typechecked -async def start_oauth(provider: str) -> dict: - """Start OAuth flow for a provider. - - device_code providers: returns {user_code, verification_uri, device_code} - authorization_code providers: returns {authUrl, codeVerifier, state} - """ - async with httpx.AsyncClient(timeout=15.0) as client: + device_code providers: returns {user_code, verification_uri, device_code} + authorization_code providers: returns {authUrl, codeVerifier, state} + """ try: - r = await client.get(f"{NINE_ROUTER_API}/oauth/{provider}/device-code") + r = await self._http.get(f"{NINE_ROUTER_API}/oauth/{provider}/device-code") if r.status_code == 200: data: dict = r.json() return { @@ -42,7 +47,7 @@ async def start_oauth(provider: str) -> dict: pass callback_url: str = f"http://localhost:{NINE_ROUTER_PORT}/callback" - r = await client.get( + r = await self._http.get( f"{NINE_ROUTER_API}/oauth/{provider}/authorize", params={"redirect_uri": callback_url}, ) @@ -56,53 +61,49 @@ async def start_oauth(provider: str) -> dict: "redirect_uri": callback_url, } + @typechecked + async def poll_oauth( + self, + provider: str, + device_code: str, + code_verifier: str | None = None, + extra_data: dict | None = None, + ) -> dict: + body: dict = {"deviceCode": device_code} + if code_verifier: + body["codeVerifier"] = code_verifier + if extra_data: + body["extraData"] = extra_data -@typechecked -async def poll_oauth( - provider: str, - device_code: str, - code_verifier: str | None = None, - extra_data: dict | None = None, -) -> dict: - body: dict = {"deviceCode": device_code} - if code_verifier: - body["codeVerifier"] = code_verifier - if extra_data: - body["extraData"] = extra_data - - async with httpx.AsyncClient(timeout=15.0) as client: - r = await client.post(f"{NINE_ROUTER_API}/oauth/{provider}/poll", json=body) + r = await self._http.post(f"{NINE_ROUTER_API}/oauth/{provider}/poll", json=body) r.raise_for_status() return r.json() - -@typechecked -async def exchange_oauth( - provider: str, - code: str, - redirect_uri: str, - code_verifier: str, - state: str = "", -) -> dict: - payload: dict = { - "code": code, - "redirectUri": redirect_uri, - "codeVerifier": code_verifier, - "state": state, - } - print(f"exchange_oauth: provider={provider} redirect_uri={redirect_uri}") - async with httpx.AsyncClient(timeout=15.0) as client: - r = await client.post(f"{NINE_ROUTER_API}/oauth/{provider}/exchange", json=payload) + @typechecked + async def exchange_oauth( + self, + provider: str, + code: str, + redirect_uri: str, + code_verifier: str, + state: str = "", + ) -> dict: + payload: dict = { + "code": code, + "redirectUri": redirect_uri, + "codeVerifier": code_verifier, + "state": state, + } + print(f"exchange_oauth: provider={provider} redirect_uri={redirect_uri}") + r = await self._http.post(f"{NINE_ROUTER_API}/oauth/{provider}/exchange", json=payload) print(f"exchange_oauth: status={r.status_code}") r.raise_for_status() return r.json() - -@typechecked -async def get_models() -> list[dict]: - try: - async with httpx.AsyncClient(timeout=5.0) as client: - r = await client.get(f"{NINE_ROUTER_V1}/models") + @typechecked + async def get_models(self) -> list[dict]: + try: + r = await self._http.get(f"{NINE_ROUTER_V1}/models", timeout=5.0) if r.status_code == 200: data: dict = r.json() models: list = data.get("data", []) @@ -115,13 +116,11 @@ async def get_models() -> list[dict]: } for m in models ] - except Exception as e: - print(f"9Router models fetch failed: {e}") - return [] + except Exception as e: + print(f"9Router models fetch failed: {e}") + return [] - -@typechecked -async def disconnect_provider(provider_id: str) -> bool: - async with httpx.AsyncClient(timeout=10.0) as client: - r = await client.delete(f"{NINE_ROUTER_API}/providers/{provider_id}") + @typechecked + async def disconnect_provider(self, provider_id: str) -> bool: + r = await self._http.delete(f"{NINE_ROUTER_API}/providers/{provider_id}", timeout=10.0) return r.status_code == 200 diff --git a/backend/apps/subscriptions/NineRouter/NineRouterProcess/NineRouterProcess.py b/backend/apps/subscriptions/NineRouter/NineRouterProcess/NineRouterProcess.py index 8727851d..5d3043e6 100644 --- a/backend/apps/subscriptions/NineRouter/NineRouterProcess/NineRouterProcess.py +++ b/backend/apps/subscriptions/NineRouter/NineRouterProcess/NineRouterProcess.py @@ -20,145 +20,143 @@ from backend.apps.subscriptions.NineRouter.NineRouterProcess.helpers.forward_out from backend.apps.subscriptions.NineRouter.NineRouterProcess.helpers.find_9router_dir import find_9router_dir from backend.apps.subscriptions.NineRouter.NineRouterProcess.helpers.find_node import find_node -P_PROCESS: subprocess.Popen | None = None -P_THIS_DIR: str = os.path.dirname(os.path.abspath(__file__)) +_THIS_DIR: str = os.path.dirname(os.path.abspath(__file__)) +class NineRouterProcess: + def __init__(self) -> None: + self._process: Optional[subprocess.Popen] = None -@typechecked -def is_running() -> bool: - try: - r = httpx.get(f"{NINE_ROUTER_V1}/models", timeout=2.0) - return r.status_code == 200 - except Exception: - return False - - -@typechecked -async def ensure_running() -> None: - """Start 9Router if not already running.""" - global P_PROCESS - _is_packaged: bool = os.environ.get("OPENSWARM_PACKAGED") == "1" - - if is_running(): - if not _is_packaged: - try: - result = subprocess.run( - ["pgrep", "-f", "next-server"], - capture_output=True, text=True, timeout=3, - ) - if result.stdout.strip(): - print("9Router: killing stale standalone to use next dev", flush=True) - subprocess.run(["pkill", "-f", "next-server"], timeout=5) - await asyncio.sleep(2) - else: - return - except Exception: - return - else: - return - - _9router_dir: Optional[str] = find_9router_dir(P_THIS_DIR) - cmd: list[str] - cwd: str | None - env: dict[str, str] - - if _is_packaged and _9router_dir: - standalone_server: str = os.path.join(_9router_dir, "server.js") - if not os.path.exists(standalone_server): - standalone_server = os.path.join(_9router_dir, ".next", "standalone", "server.js") - if not os.path.exists(standalone_server): - print("9Router: standalone build not found in", _9router_dir, flush=True) - return - - node: Optional[str] = find_node() - if not node: - print("9Router: Node.js not found, cannot start in packaged mode", flush=True) - return - - print(f"9Router: starting (production) on port {NINE_ROUTER_PORT}...", flush=True) - cmd = [node, standalone_server] - cwd = os.path.dirname(standalone_server) - env = { - **os.environ, - "PORT": str(NINE_ROUTER_PORT), - "NEXT_PUBLIC_BASE_URL": NINE_ROUTER_URL, - "NODE_ENV": "production", - } - if node == os.environ.get("OPENSWARM_ELECTRON_PATH"): - env["ELECTRON_RUN_AS_NODE"] = "1" - - elif _9router_dir: - npx: str | None = shutil.which("npx") - if not npx: - print("9Router: npx not found, cannot auto-start", flush=True) - return - - if not os.path.isdir(os.path.join(_9router_dir, "node_modules")): - print("9Router: installing dependencies...", flush=True) - npm: str | None = shutil.which("npm") - if npm: - subprocess.run( - [npm, "install"], cwd=_9router_dir, - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=120, - ) - - print(f"9Router: starting (dev) on port {NINE_ROUTER_PORT}...", flush=True) - cmd = [npx, "next", "dev", "--webpack", "-p", str(NINE_ROUTER_PORT)] - cwd = _9router_dir - env = { - **os.environ, - "PORT": str(NINE_ROUTER_PORT), - "NEXT_PUBLIC_BASE_URL": NINE_ROUTER_URL, - } - - else: - npx = shutil.which("npx") - if not npx: - print("9Router: npx not found and no bundled 9router directory", flush=True) - return - print(f"9Router: starting (npx) on port {NINE_ROUTER_PORT}...", flush=True) - cmd = [npx, "9router", "--port", str(NINE_ROUTER_PORT), "--no-browser", "--skip-update"] - cwd = None - env = { - **os.environ, - "PORT": str(NINE_ROUTER_PORT), - "NEXT_PUBLIC_BASE_URL": NINE_ROUTER_URL, - } - - try: - P_PROCESS = subprocess.Popen( - cmd, cwd=cwd, - stdin=subprocess.DEVNULL, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - env=env, - ) - threading.Thread(target=forward_output, args=(P_PROCESS.stdout,), daemon=True).start() - - timeout: int = 20 if _is_packaged else 30 - for _ in range(timeout * 2): - await asyncio.sleep(0.5) - if is_running(): - print("9Router: started successfully", flush=True) - return - - print(f"9Router: did not start within {timeout}s", flush=True) - except Exception as e: - print(f"9Router: failed to start: {e}", flush=True) - - -@typechecked -def stop() -> None: - global P_PROCESS - if P_PROCESS: + @typechecked + def is_running(self) -> bool: try: - P_PROCESS.terminate() - P_PROCESS.wait(timeout=5) + r = httpx.get(f"{NINE_ROUTER_V1}/models", timeout=2.0) + return r.status_code == 200 except Exception: + return False + + @typechecked + async def ensure_running(self) -> None: + """Start 9Router if not already running.""" + _is_packaged: bool = os.environ.get("OPENSWARM_PACKAGED") == "1" + + if self.is_running(): + if not _is_packaged: + try: + result = subprocess.run( + ["pgrep", "-f", "next-server"], + capture_output=True, text=True, timeout=3, + ) + if result.stdout.strip(): + print("9Router: killing stale standalone to use next dev", flush=True) + subprocess.run(["pkill", "-f", "next-server"], timeout=5) + await asyncio.sleep(2) + else: + return + except Exception: + return + else: + return + + _9router_dir: Optional[str] = find_9router_dir(_THIS_DIR) + cmd: list[str] + cwd: str | None + env: dict[str, str] + + if _is_packaged and _9router_dir: + standalone_server: str = os.path.join(_9router_dir, "server.js") + if not os.path.exists(standalone_server): + standalone_server = os.path.join(_9router_dir, ".next", "standalone", "server.js") + if not os.path.exists(standalone_server): + print("9Router: standalone build not found in", _9router_dir, flush=True) + return + + node: Optional[str] = find_node() + if not node: + print("9Router: Node.js not found, cannot start in packaged mode", flush=True) + return + + print(f"9Router: starting (production) on port {NINE_ROUTER_PORT}...", flush=True) + cmd = [node, standalone_server] + cwd = os.path.dirname(standalone_server) + env = { + **os.environ, + "PORT": str(NINE_ROUTER_PORT), + "NEXT_PUBLIC_BASE_URL": NINE_ROUTER_URL, + "NODE_ENV": "production", + } + if node == os.environ.get("OPENSWARM_ELECTRON_PATH"): + env["ELECTRON_RUN_AS_NODE"] = "1" + + elif _9router_dir: + npx: str | None = shutil.which("npx") + if not npx: + print("9Router: npx not found, cannot auto-start", flush=True) + return + + if not os.path.isdir(os.path.join(_9router_dir, "node_modules")): + print("9Router: installing dependencies...", flush=True) + npm: str | None = shutil.which("npm") + if npm: + subprocess.run( + [npm, "install"], cwd=_9router_dir, + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=120, + ) + + print(f"9Router: starting (dev) on port {NINE_ROUTER_PORT}...", flush=True) + cmd = [npx, "next", "dev", "--webpack", "-p", str(NINE_ROUTER_PORT)] + cwd = _9router_dir + env = { + **os.environ, + "PORT": str(NINE_ROUTER_PORT), + "NEXT_PUBLIC_BASE_URL": NINE_ROUTER_URL, + } + + else: + npx = shutil.which("npx") + if not npx: + print("9Router: npx not found and no bundled 9router directory", flush=True) + return + print(f"9Router: starting (npx) on port {NINE_ROUTER_PORT}...", flush=True) + cmd = [npx, "9router", "--port", str(NINE_ROUTER_PORT), "--no-browser", "--skip-update"] + cwd = None + env = { + **os.environ, + "PORT": str(NINE_ROUTER_PORT), + "NEXT_PUBLIC_BASE_URL": NINE_ROUTER_URL, + } + + try: + self._process = subprocess.Popen( + cmd, cwd=cwd, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + env=env, + ) + threading.Thread(target=forward_output, args=(self._process.stdout,), daemon=True).start() + + timeout: int = 20 if _is_packaged else 30 + for _ in range(timeout * 2): + await asyncio.sleep(0.5) + if self.is_running(): + print("9Router: started successfully", flush=True) + return + + print(f"9Router: did not start within {timeout}s", flush=True) + except Exception as e: + print(f"9Router: failed to start: {e}", flush=True) + + @typechecked + def stop(self) -> None: + if self._process: try: - P_PROCESS.kill() + self._process.terminate() + self._process.wait(timeout=5) except Exception: - pass - P_PROCESS = None - print("9Router stopped") + try: + self._process.kill() + except Exception: + pass + self._process = None + print("9Router stopped") diff --git a/backend/apps/subscriptions/subscriptions.py b/backend/apps/subscriptions/subscriptions.py index 29a35009..4babf892 100644 --- a/backend/apps/subscriptions/subscriptions.py +++ b/backend/apps/subscriptions/subscriptions.py @@ -29,9 +29,8 @@ async def subscriptions_lifespan(): except Exception as e: print(f"9Router auto-start failed: {e}") yield - router.cancel_ensure_task() try: - router.stop() + await router.stop() except Exception: pass