add psutil

This commit is contained in:
Dillon DuPont
2025-11-01 20:33:22 -04:00
parent be4c7e45aa
commit 1447950bcf
3 changed files with 81 additions and 4 deletions

View File

@@ -8,6 +8,7 @@ from pathlib import Path
from typing import Optional, Dict, Any
from urllib import request
from urllib.error import HTTPError, URLError
import psutil
# Map child PID -> listening port
_pid_to_port: Dict[int, int] = {}
@@ -30,6 +31,31 @@ def _post_json(url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
return {"error": "url_error", "reason": str(e.reason)}
def _detect_port_for_pid(pid: int) -> int:
"""Detect a listening local TCP port for the given PID using psutil.
Fails fast if psutil is unavailable or if no suitable port is found.
"""
if psutil is None:
raise RuntimeError("psutil is required for PID->port detection. Please install psutil.")
# Scan system-wide connections and filter by PID
for c in psutil.net_connections(kind="tcp"):
if getattr(c, "pid", None) != pid:
continue
laddr = getattr(c, "laddr", None)
status = str(getattr(c, "status", ""))
if not laddr or not isinstance(laddr, tuple) or len(laddr) < 2:
continue
lip, lport = laddr[0], int(laddr[1])
if status.upper() != "LISTEN":
continue
if lip in ("127.0.0.1", "::1", "0.0.0.0", "::"):
return lport
raise RuntimeError(f"Could not detect listening port for pid {pid}")
def launch_window(
url: Optional[str] = None,
*,
@@ -101,7 +127,7 @@ def get_element_rect(pid: int, selector: str, *, space: str = "window"):
Returns a dict like {"x": float, "y": float, "width": float, "height": float} or None if not found.
"""
if pid not in _pid_to_port:
raise RuntimeError(f"Unknown pid {pid}; no registered bench-ui window")
_pid_to_port[pid] = _detect_port_for_pid(pid)
port = _pid_to_port[pid]
url = f"http://127.0.0.1:{port}/rect"
last: Dict[str, Any] = {}
@@ -129,7 +155,7 @@ def execute_javascript(pid: int, javascript: str):
Retries briefly while the window is still becoming ready.
"""
if pid not in _pid_to_port:
raise RuntimeError(f"Unknown pid {pid}; no registered bench-ui window")
_pid_to_port[pid] = _detect_port_for_pid(pid)
port = _pid_to_port[pid]
url = f"http://127.0.0.1:{port}/eval"
last: Dict[str, Any] = {}

View File

@@ -4,8 +4,8 @@ build-backend = "pdm.backend"
[project]
name = "cua-bench-ui"
version = "0.1.0"
description = "Lightweight webUI window launcher for CUA bench using pywebview"
version = "0.3.0"
description = "Lightweight webUI window controller for CUA bench using pywebview"
readme = "README.md"
authors = [
{ name = "TryCua", email = "gh@trycua.com" }
@@ -13,6 +13,7 @@ authors = [
dependencies = [
"pywebview>=5.3",
"aiohttp>=3.9.0",
"psutil>=5.9",
]
requires-python = ">=3.12"

View File

@@ -0,0 +1,50 @@
import time
import psutil
import pytest
from bench_ui import launch_window, execute_javascript
from bench_ui.api import _pid_to_port
HTML = """
<!doctype html>
<html>
<head>
<meta charset="utf-8" />
<title>Bench UI Test</title>
</head>
<body>
<div id="t">hello-world</div>
</body>
</html>
"""
def test_execute_js_after_clearing_port_mapping():
# Skip if pywebview backend is unavailable on this machine
pywebview = pytest.importorskip("webview")
pid = launch_window(html=HTML, title="Bench UI Test", width=400, height=300)
try:
# Give a brief moment for window to render and server to start
time.sleep(1.0)
# Sanity: mapping should exist initially
assert pid in _pid_to_port
# Clear the cached mapping to simulate a fresh process lookup
del _pid_to_port[pid]
# Now execute JS; this should succeed by detecting the port via psutil
result = execute_javascript(pid, "document.querySelector('#t')?.textContent")
assert result == "hello-world"
finally:
# Best-effort cleanup of the child process
try:
p = psutil.Process(pid)
p.terminate()
try:
p.wait(timeout=3)
except psutil.TimeoutExpired:
p.kill()
except Exception:
pass