Fix Python code format check issues

This commit is contained in:
yaoyao-open-source committed 2026-09-10 17:54:18 -07:00
1 parent 70be3c37e5
commit 3047c9c20a
15 files changed
+3155 -3039

No files matched your search

+29 -6
View File
@@ -79,11 +79,15 @@ class AccessibilityClient:
# 4. Verify connection # 4. Verify connection
for attempt in range(5): for attempt in range(5):
if self.ping(): if self.ping():
logger.info(f"ArtemisAccessibilityHelper connected successfully on {self._device_id}") logger.info(
f"ArtemisAccessibilityHelper connected successfully on {self._device_id}"
)
return True return True
time.sleep(0.5) time.sleep(0.5)
logger.warning(f"Failed to connect to ArtemisAccessibilityHelper on port {self._local_port}") logger.warning(
f"Failed to connect to ArtemisAccessibilityHelper on port {self._local_port}"
)
return False return False
def ping(self) -> bool: def ping(self) -> bool:
@@ -129,7 +133,11 @@ class AccessibilityClient:
req = urllib.request.Request(f"{self._base_url}/snapshot") req = urllib.request.Request(f"{self._base_url}/snapshot")
with urllib.request.urlopen(req, timeout=6.0) as resp: with urllib.request.urlopen(req, timeout=6.0) as resp:
data = json.loads(resp.read().decode()) data = json.loads(resp.read().decode())
if data.get("success") and data.get("has_screenshot") and data.get("screenshot_base64"): if (
data.get("success")
and data.get("has_screenshot")
and data.get("screenshot_base64")
):
return data return data
except Exception as e: except Exception as e:
logger.debug(f"Atomic snapshot request failed (falling back to dual path): {e}") logger.debug(f"Atomic snapshot request failed (falling back to dual path): {e}")
@@ -226,11 +234,26 @@ class AccessibilityClient:
def _enable_accessibility_service(self) -> None: def _enable_accessibility_service(self) -> None:
"""Silently enable the accessibility service via adb settings without UI prompts.""" """Silently enable the accessibility service via adb settings without UI prompts."""
res = self._run_adb(["shell", "settings", "get", "secure", "enabled_accessibility_services"]) res = self._run_adb(
["shell", "settings", "get", "secure", "enabled_accessibility_services"]
)
current_services = res.stdout.strip() current_services = res.stdout.strip()
if SERVICE_NAME not in current_services: if SERVICE_NAME not in current_services:
new_services = f"{current_services}:{SERVICE_NAME}" if current_services and current_services != "null" else SERVICE_NAME new_services = (
self._run_adb(["shell", "settings", "put", "secure", "enabled_accessibility_services", new_services]) f"{current_services}:{SERVICE_NAME}"
if current_services and current_services != "null"
else SERVICE_NAME
)
self._run_adb(
[
"shell",
"settings",
"put",
"secure",
"enabled_accessibility_services",
new_services,
]
)
self._run_adb(["shell", "settings", "put", "secure", "accessibility_enabled", "1"]) self._run_adb(["shell", "settings", "put", "secure", "accessibility_enabled", "1"])
def _setup_port_forward(self) -> None: def _setup_port_forward(self) -> None:
+15 -4
View File
@@ -163,16 +163,27 @@ def generate_banner(
center_x = px + iw / 2.0 center_x = px + iw / 2.0
center_y = py + ih / 2.0 center_y = py + ih / 2.0
print(f"Banner successfully generated: {output_path} ({cw}x{ch})") print(f"Banner successfully generated: {output_path} ({cw}x{ch})")
print(f"Ink dimensions: {iw}x{ih}px, Exact center: ({center_x:.1f}, {center_y:.1f}) [Canvas center: ({cw/2:.1f}, {ch/2:.1f}), Offset X: {offset_x:+d}px]") print(
f"Ink dimensions: {iw}x{ih}px, Exact center: ({center_x:.1f}, {center_y:.1f}) [Canvas center: ({cw / 2:.1f}, {ch / 2:.1f}), Offset X: {offset_x:+d}px]"
)
if __name__ == "__main__": if __name__ == "__main__":
import argparse import argparse
parser = argparse.ArgumentParser(description="ARTEMIS Official Banner Generator") parser = argparse.ArgumentParser(description="ARTEMIS Official Banner Generator")
parser.add_argument("--offset-x", type=int, default=160, help="Horizontal pixel offset to the right from mathematical center (default: 160)") parser.add_argument(
parser.add_argument("--tracking", type=int, default=64, help="Character tracking spacing (default: 64)") "--offset-x",
parser.add_argument("--font-size", type=int, default=98, help="Font size in pixels (default: 98)") type=int,
default=160,
help="Horizontal pixel offset to the right from mathematical center (default: 160)",
)
parser.add_argument(
"--tracking", type=int, default=64, help="Character tracking spacing (default: 64)"
)
parser.add_argument(
"--font-size", type=int, default=98, help="Font size in pixels (default: 98)"
)
args = parser.parse_args() args = parser.parse_args()
assets_dir = os.path.dirname(os.path.abspath(__file__)) assets_dir = os.path.dirname(os.path.abspath(__file__))
@@ -23,7 +23,9 @@ security = HTTPBearer(auto_error=False)
def create_access_token(user_id: str, extra_claims: dict | None = None) -> str: def create_access_token(user_id: str, extra_claims: dict | None = None) -> str:
"""Generate a signed JWT access token.""" """Generate a signed JWT access token."""
expire = datetime.now(timezone.utc) + timedelta(minutes=settings.JWT_ACCESS_TOKEN_EXPIRE_MINUTES) expire = datetime.now(timezone.utc) + timedelta(
minutes=settings.JWT_ACCESS_TOKEN_EXPIRE_MINUTES
)
to_encode = { to_encode = {
"sub": user_id, "sub": user_id,
"exp": expire, "exp": expire,
@@ -48,7 +50,9 @@ def decode_access_token(token: str) -> dict:
) )
async def get_current_user(credentials: HTTPAuthorizationCredentials | None = Security(security)) -> str: async def get_current_user(
credentials: HTTPAuthorizationCredentials | None = Security(security),
) -> str:
"""FastAPI dependency to extract and verify the current user_id.""" """FastAPI dependency to extract and verify the current user_id."""
if not credentials: if not credentials:
raise HTTPException( raise HTTPException(
@@ -45,7 +45,9 @@ class OTPService:
} }
# Dispatch via SMS / Email provider (mocked logger in template) # Dispatch via SMS / Email provider (mocked logger in template)
logger.info(f"[OTP Service] Generated OTP for {ident}: [{code}] (Expires at: {expires_at.isoformat()})") logger.info(
f"[OTP Service] Generated OTP for {ident}: [{code}] (Expires at: {expires_at.isoformat()})"
)
return code return code
def verify_otp(self, identifier: str, code: str) -> bool: def verify_otp(self, identifier: str, code: str) -> bool:
+3 -1
View File
@@ -43,7 +43,9 @@ class Settings(BaseSettings):
ARTEMIS_MEMORY_LIMIT: str = "2g" ARTEMIS_MEMORY_LIMIT: str = "2g"
# Cloud Orchestrator & Cuttlefish Host Settings # Cloud Orchestrator & Cuttlefish Host Settings
CLOUD_ORCHESTRATOR_URL: str = os.getenv("CLOUD_ORCHESTRATOR_URL", "http://cloud-orchestrator:2081") CLOUD_ORCHESTRATOR_URL: str = os.getenv(
"CLOUD_ORCHESTRATOR_URL", "http://cloud-orchestrator:2081"
)
CUTTLEFISH_HOST_GATEWAY: str = os.getenv("CUTTLEFISH_HOST_GATEWAY", "cloud-orchestrator") CUTTLEFISH_HOST_GATEWAY: str = os.getenv("CUTTLEFISH_HOST_GATEWAY", "cloud-orchestrator")
CUTTLEFISH_START_ADB_PORT: int = 6520 CUTTLEFISH_START_ADB_PORT: int = 6520
CUTTLEFISH_START_WEBRTC_PORT: int = 8443 CUTTLEFISH_START_WEBRTC_PORT: int = 8443
+1
View File
@@ -76,6 +76,7 @@ async def serve_spa_root():
if __name__ == "__main__": if __name__ == "__main__":
import uvicorn import uvicorn
uvicorn.run( uvicorn.run(
"app.main:app", "app.main:app",
host=settings.HOST, host=settings.HOST,
@@ -83,7 +83,9 @@ async def terminate_session(
user_id: str = Depends(get_current_user), user_id: str = Depends(get_current_user),
): ):
"""Terminate the session, destroy Artemis container and Cuttlefish AVD.""" """Terminate the session, destroy Artemis container and Cuttlefish AVD."""
success = await session_manager.terminate_session(session_id, user_id=user_id, reason="User terminated") success = await session_manager.terminate_session(
session_id, user_id=user_id, reason="User terminated"
)
if not success: if not success:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, status_code=status.HTTP_404_NOT_FOUND,
@@ -16,7 +16,9 @@ from pydantic import BaseModel, Field
class OTPRequest(BaseModel): class OTPRequest(BaseModel):
identifier: str = Field(..., description="User phone number or email address", example="user@example.com") identifier: str = Field(
..., description="User phone number or email address", example="user@example.com"
)
class OTPRequestResponse(BaseModel): class OTPRequestResponse(BaseModel):
@@ -26,7 +28,9 @@ class OTPRequestResponse(BaseModel):
class OTPVerify(BaseModel): class OTPVerify(BaseModel):
identifier: str = Field(..., description="User phone number or email address", example="user@example.com") identifier: str = Field(
..., description="User phone number or email address", example="user@example.com"
)
code: str = Field(..., description="6-digit verification code", example="123456") code: str = Field(..., description="6-digit verification code", example="123456")
@@ -27,9 +27,15 @@ class SessionStatus(str, Enum):
class CreateSessionRequest(BaseModel): class CreateSessionRequest(BaseModel):
device_preset: str = Field(default="pixel_7_pro", description="Target Cuttlefish device configuration") device_preset: str = Field(
initial_task: str | None = Field(default=None, description="Optional initial prompt/task for Artemis") default="pixel_7_pro", description="Target Cuttlefish device configuration"
ttl_minutes: int = Field(default=60, ge=5, le=480, description="Session time to live in minutes") )
initial_task: str | None = Field(
default=None, description="Optional initial prompt/task for Artemis"
)
ttl_minutes: int = Field(
default=60, ge=5, le=480, description="Session time to live in minutes"
)
class SessionRecord(BaseModel): class SessionRecord(BaseModel):
@@ -23,6 +23,7 @@ logger = logging.getLogger("artemis.bigquery")
try: try:
from google.cloud import bigquery from google.cloud import bigquery
from google.cloud.exceptions import NotFound from google.cloud.exceptions import NotFound
BQ_AVAILABLE = True BQ_AVAILABLE = True
except ImportError: except ImportError:
BQ_AVAILABLE = False BQ_AVAILABLE = False
@@ -34,7 +35,9 @@ class BigQueryMappingService:
def __init__(self): def __init__(self):
self._local_cache: dict[str, SessionRecord] = {} self._local_cache: dict[str, SessionRecord] = {}
self._client: Optional["bigquery.Client"] = None self._client: Optional["bigquery.Client"] = None
self._table_ref: str = f"{settings.GCP_PROJECT_ID}.{settings.BQ_DATASET}.{settings.BQ_TABLE}" self._table_ref: str = (
f"{settings.GCP_PROJECT_ID}.{settings.BQ_DATASET}.{settings.BQ_TABLE}"
)
if BQ_AVAILABLE and not settings.USE_LOCAL_FALLBACK_DB: if BQ_AVAILABLE and not settings.USE_LOCAL_FALLBACK_DB:
try: try:
@@ -42,7 +45,9 @@ class BigQueryMappingService:
self._ensure_dataset_and_table() self._ensure_dataset_and_table()
logger.info(f"[BigQuery] Connected to BigQuery table: {self._table_ref}") logger.info(f"[BigQuery] Connected to BigQuery table: {self._table_ref}")
except Exception as e: except Exception as e:
logger.warning(f"[BigQuery] Failed to initialize BigQuery client: {e}. Falling back to in-memory store.") logger.warning(
f"[BigQuery] Failed to initialize BigQuery client: {e}. Falling back to in-memory store."
)
self._client = None self._client = None
else: else:
logger.info("[BigQuery] Using in-memory fallback database for local/test execution.") logger.info("[BigQuery] Using in-memory fallback database for local/test execution.")
@@ -110,7 +115,9 @@ class BigQueryMappingService:
} }
errors = self._client.insert_rows_json(self._table_ref, [row]) errors = self._client.insert_rows_json(self._table_ref, [row])
if errors: if errors:
logger.error(f"[BigQuery] Insert errors for session {record.session_id}: {errors}") logger.error(
f"[BigQuery] Insert errors for session {record.session_id}: {errors}"
)
except Exception as e: except Exception as e:
logger.error(f"[BigQuery] Streaming insert failed: {e}") logger.error(f"[BigQuery] Streaming insert failed: {e}")
@@ -121,18 +128,22 @@ class BigQueryMappingService:
async def list_user_sessions(self, user_id: str) -> list[SessionRecord]: async def list_user_sessions(self, user_id: str) -> list[SessionRecord]:
"""List active/valid sessions belonging to a specific user.""" """List active/valid sessions belonging to a specific user."""
return [ return [
rec for rec in self._local_cache.values() rec
for rec in self._local_cache.values()
if rec.user_id == user_id and rec.status != SessionStatus.TERMINATED if rec.user_id == user_id and rec.status != SessionStatus.TERMINATED
] ]
async def list_all_active_sessions(self) -> list[SessionRecord]: async def list_all_active_sessions(self) -> list[SessionRecord]:
"""List all non-terminated sessions for maintenance/reaping.""" """List all non-terminated sessions for maintenance/reaping."""
return [ return [
rec for rec in self._local_cache.values() rec
for rec in self._local_cache.values()
if rec.status in (SessionStatus.PROVISIONING, SessionStatus.ACTIVE, SessionStatus.IDLE) if rec.status in (SessionStatus.PROVISIONING, SessionStatus.ACTIVE, SessionStatus.IDLE)
] ]
async def update_status(self, session_id: str, status: SessionStatus, error_message: str | None = None) -> None: async def update_status(
self, session_id: str, status: SessionStatus, error_message: str | None = None
) -> None:
"""Update the status of an existing session record.""" """Update the status of an existing session record."""
record = self._local_cache.get(session_id) record = self._local_cache.get(session_id)
if record: if record:
@@ -47,7 +47,9 @@ class CloudOrchestratorService:
try: try:
resp = await client.get(f"{settings.CLOUD_ORCHESTRATOR_URL}/cvds") resp = await client.get(f"{settings.CLOUD_ORCHESTRATOR_URL}/cvds")
if resp.status_code != 200: if resp.status_code != 200:
logger.error(f"[Cloud Orchestrator] Failed to fetch CVD list: HTTP {resp.status_code} ({resp.text})") logger.error(
f"[Cloud Orchestrator] Failed to fetch CVD list: HTTP {resp.status_code} ({resp.text})"
)
return [] return []
all_cvds = resp.json().get("cvds", []) all_cvds = resp.json().get("cvds", [])
@@ -69,9 +71,13 @@ class CloudOrchestratorService:
logger.warning(f"[Cloud Orchestrator] Error listing CVDs from orchestrator: {e}") logger.warning(f"[Cloud Orchestrator] Error listing CVDs from orchestrator: {e}")
return [] return []
async def create_new_cuttlefish_emulator(self, instance_name: str, preset: str = "pixel_7_pro") -> dict: async def create_new_cuttlefish_emulator(
self, instance_name: str, preset: str = "pixel_7_pro"
) -> dict:
"""Trigger creation of a new CVD via Host Orchestrator and poll until ready.""" """Trigger creation of a new CVD via Host Orchestrator and poll until ready."""
logger.info(f"[Cloud Orchestrator] Provisioning new Cuttlefish emulator '{instance_name}' via Host Orchestrator...") logger.info(
f"[Cloud Orchestrator] Provisioning new Cuttlefish emulator '{instance_name}' via Host Orchestrator..."
)
async with httpx.AsyncClient(timeout=10.0) as client: async with httpx.AsyncClient(timeout=10.0) as client:
create_resp = await client.post( create_resp = await client.post(
f"{settings.CLOUD_ORCHESTRATOR_URL}/cvds", f"{settings.CLOUD_ORCHESTRATOR_URL}/cvds",
@@ -94,19 +100,27 @@ class CloudOrchestratorService:
) )
if create_resp.status_code not in (200, 201): if create_resp.status_code not in (200, 201):
raise RuntimeError(f"Host Orchestrator rejected CVD creation (HTTP {create_resp.status_code}): {create_resp.text}") raise RuntimeError(
f"Host Orchestrator rejected CVD creation (HTTP {create_resp.status_code}): {create_resp.text}"
)
op = create_resp.json() op = create_resp.json()
op_name = op.get("name") op_name = op.get("name")
logger.info(f"[Cloud Orchestrator] Creation operation started: {op_name}. Waiting for device boot...") logger.info(
f"[Cloud Orchestrator] Creation operation started: {op_name}. Waiting for device boot..."
)
# Poll operation until done # Poll operation until done
for attempt in range(60): # Poll up to 120 seconds for attempt in range(60): # Poll up to 120 seconds
await asyncio.sleep(2.0) await asyncio.sleep(2.0)
try: try:
op_resp = await client.get(f"{settings.CLOUD_ORCHESTRATOR_URL}/operations/{op_name}") op_resp = await client.get(
f"{settings.CLOUD_ORCHESTRATOR_URL}/operations/{op_name}"
)
if op_resp.status_code == 200 and op_resp.json().get("done", False): if op_resp.status_code == 200 and op_resp.json().get("done", False):
logger.info(f"[Cloud Orchestrator] Operation {op_name} successfully finished!") logger.info(
f"[Cloud Orchestrator] Operation {op_name} successfully finished!"
)
break break
except Exception as e: except Exception as e:
logger.debug(f"[Cloud Orchestrator] Polling operation {op_name}: {e}") logger.debug(f"[Cloud Orchestrator] Polling operation {op_name}: {e}")
@@ -119,9 +133,13 @@ class CloudOrchestratorService:
logger.info(f"[Cloud Orchestrator] Newly provisioned CVD found: {cvd}") logger.info(f"[Cloud Orchestrator] Newly provisioned CVD found: {cvd}")
return cvd return cvd
raise RuntimeError(f"CVD '{instance_name}' creation operation finished, but device was not found in running CVD list") raise RuntimeError(
f"CVD '{instance_name}' creation operation finished, but device was not found in running CVD list"
)
async def create_cuttlefish_instance(self, session_id: str, preset: str = "pixel_7_pro") -> CuttlefishInstanceInfo: async def create_cuttlefish_instance(
self, session_id: str, preset: str = "pixel_7_pro"
) -> CuttlefishInstanceInfo:
""" """
Main entrypoint: Main entrypoint:
1. Find an available running emulator not bound to any active Artemis session. 1. Find an available running emulator not bound to any active Artemis session.
@@ -140,7 +158,9 @@ class CloudOrchestratorService:
chosen_cvd = None chosen_cvd = None
if available: if available:
chosen_cvd = available[0] chosen_cvd = available[0]
logger.info(f"[Cloud Orchestrator] Reusing existing free emulator '{chosen_cvd.get('name')}' for session {session_id}.") logger.info(
f"[Cloud Orchestrator] Reusing existing free emulator '{chosen_cvd.get('name')}' for session {session_id}."
)
else: else:
# 2. If no emulator available, create one # 2. If no emulator available, create one
instance_name = f"cvd-{session_id[:8]}" instance_name = f"cvd-{session_id[:8]}"
@@ -158,7 +178,9 @@ class CloudOrchestratorService:
pass pass
if not adb_port: if not adb_port:
raise RuntimeError(f"Could not dynamically determine ADB port for emulator '{emulator_name}': {chosen_cvd}") raise RuntimeError(
f"Could not dynamically determine ADB port for emulator '{emulator_name}': {chosen_cvd}"
)
adb_port = int(adb_port) adb_port = int(adb_port)
logger.info(f"[Cloud Orchestrator] Dynamic ADB port for '{emulator_name}': {adb_port}") logger.info(f"[Cloud Orchestrator] Dynamic ADB port for '{emulator_name}': {adb_port}")
@@ -174,7 +196,9 @@ class CloudOrchestratorService:
self._allocated_instances[session_id] = info self._allocated_instances[session_id] = info
self._emulator_bindings[emulator_name] = session_id self._emulator_bindings[emulator_name] = session_id
logger.info(f"[Cloud Orchestrator] Successfully bound emulator '{emulator_name}' (port {adb_port}) to session {session_id}.") logger.info(
f"[Cloud Orchestrator] Successfully bound emulator '{emulator_name}' (port {adb_port}) to session {session_id}."
)
return info return info
async def destroy_cuttlefish_instance(self, session_id: str) -> None: async def destroy_cuttlefish_instance(self, session_id: str) -> None:
@@ -182,11 +206,15 @@ class CloudOrchestratorService:
async with self._lock: async with self._lock:
info = self._allocated_instances.pop(session_id, None) info = self._allocated_instances.pop(session_id, None)
if not info: if not info:
logger.info(f"[Cloud Orchestrator] No active binding found for session {session_id}") logger.info(
f"[Cloud Orchestrator] No active binding found for session {session_id}"
)
return return
self._emulator_bindings.pop(info.instance_id, None) self._emulator_bindings.pop(info.instance_id, None)
logger.info(f"[Cloud Orchestrator] Unbound emulator '{info.instance_id}' from session {session_id}. Emulator returned to pool.") logger.info(
f"[Cloud Orchestrator] Unbound emulator '{info.instance_id}' from session {session_id}. Emulator returned to pool."
)
cloud_orchestrator_service = CloudOrchestratorService() cloud_orchestrator_service = CloudOrchestratorService()
@@ -23,6 +23,7 @@ logger = logging.getLogger("artemis.docker")
try: try:
import docker import docker
from docker.errors import DockerException, NotFound from docker.errors import DockerException, NotFound
DOCKER_AVAILABLE = True DOCKER_AVAILABLE = True
except ImportError: except ImportError:
DOCKER_AVAILABLE = False DOCKER_AVAILABLE = False
@@ -36,10 +37,14 @@ class DockerManagerService:
if DOCKER_AVAILABLE: if DOCKER_AVAILABLE:
try: try:
self._client = docker.DockerClient(base_url=settings.DOCKER_SOCKET_PATH) self._client = docker.DockerClient(base_url=settings.DOCKER_SOCKET_PATH)
logger.info(f"[Docker Service] Connected to Docker socket at {settings.DOCKER_SOCKET_PATH}") logger.info(
f"[Docker Service] Connected to Docker socket at {settings.DOCKER_SOCKET_PATH}"
)
self._ensure_network() self._ensure_network()
except Exception as e: except Exception as e:
logger.warning(f"[Docker Service] Could not connect to Docker socket: {e}. Running in simulation mode.") logger.warning(
f"[Docker Service] Could not connect to Docker socket: {e}. Running in simulation mode."
)
self._client = None self._client = None
else: else:
logger.warning("[Docker Service] Docker SDK not available; running in simulation mode.") logger.warning("[Docker Service] Docker SDK not available; running in simulation mode.")
@@ -66,7 +71,9 @@ class DockerManagerService:
container_name = f"{settings.ARTEMIS_CONTAINER_PREFIX}-{session_id}" container_name = f"{settings.ARTEMIS_CONTAINER_PREFIX}-{session_id}"
adb_target = f"{cuttlefish_host}:{cuttlefish_adb_port}" adb_target = f"{cuttlefish_host}:{cuttlefish_adb_port}"
logger.info(f"[Docker Service] Launching container '{container_name}' targeting ADB device '{adb_target}'...") logger.info(
f"[Docker Service] Launching container '{container_name}' targeting ADB device '{adb_target}'..."
)
if not self._client: if not self._client:
# Simulated mode for environments without raw Docker socket access # Simulated mode for environments without raw Docker socket access
@@ -55,11 +55,21 @@ class SessionManagerService:
for s in active_sessions: for s in active_sessions:
# Check if session has exceeded expiration or missed heartbeats # Check if session has exceeded expiration or missed heartbeats
if s.expires_at < now: if s.expires_at < now:
logger.info(f"[Session Reaper] Session {s.session_id} has expired (TTL elapsed). Tearing down...") logger.info(
await self.terminate_session(s.session_id, user_id=s.user_id, reason="TTL expired") f"[Session Reaper] Session {s.session_id} has expired (TTL elapsed). Tearing down..."
elif (now - s.last_heartbeat_at).total_seconds() > settings.SESSION_HEARTBEAT_TIMEOUT_SECONDS: )
logger.info(f"[Session Reaper] Session {s.session_id} heartbeat timed out. Tearing down...") await self.terminate_session(
await self.terminate_session(s.session_id, user_id=s.user_id, reason="Heartbeat timeout") s.session_id, user_id=s.user_id, reason="TTL expired"
)
elif (
now - s.last_heartbeat_at
).total_seconds() > settings.SESSION_HEARTBEAT_TIMEOUT_SECONDS:
logger.info(
f"[Session Reaper] Session {s.session_id} heartbeat timed out. Tearing down..."
)
await self.terminate_session(
s.session_id, user_id=s.user_id, reason="Heartbeat timeout"
)
except asyncio.CancelledError: except asyncio.CancelledError:
break break
except Exception as e: except Exception as e:
@@ -71,7 +81,9 @@ class SessionManagerService:
now = datetime.now(timezone.utc) now = datetime.now(timezone.utc)
expires_at = now + timedelta(minutes=request.ttl_minutes) expires_at = now + timedelta(minutes=request.ttl_minutes)
logger.info(f"[Session Manager] Starting session creation for user={user_id}, session_id={session_id}...") logger.info(
f"[Session Manager] Starting session creation for user={user_id}, session_id={session_id}..."
)
# Initialize record with PROVISIONING status # Initialize record with PROVISIONING status
record = SessionRecord( record = SessionRecord(
@@ -126,7 +138,9 @@ class SessionManagerService:
if not record: if not record:
return None return None
if record.user_id != user_id: if record.user_id != user_id:
logger.warning(f"[Session Manager] Tenant mismatch: user {user_id} attempted access to session {session_id}") logger.warning(
f"[Session Manager] Tenant mismatch: user {user_id} attempted access to session {session_id}"
)
return None return None
return self._build_session_response(record) return self._build_session_response(record)
@@ -162,7 +176,9 @@ class SessionManagerService:
expires_at=record.expires_at, expires_at=record.expires_at,
) )
async def terminate_session(self, session_id: str, user_id: str | None = None, reason: str = "User requested") -> bool: async def terminate_session(
self, session_id: str, user_id: str | None = None, reason: str = "User requested"
) -> bool:
"""Gracefully terminate container, emulator, and update records.""" """Gracefully terminate container, emulator, and update records."""
record = await bigquery_service.get_session(session_id) record = await bigquery_service.get_session(session_id)
if not record: if not record:
@@ -4,7 +4,6 @@ from artemis.clients.accessibility_client import AccessibilityClient, DEFAULT_PO
class TestAccessibilityClient(unittest.TestCase): class TestAccessibilityClient(unittest.TestCase):
def setUp(self): def setUp(self):
self.client = AccessibilityClient(device_id="dummy_device", local_port=DEFAULT_PORT) self.client = AccessibilityClient(device_id="dummy_device", local_port=DEFAULT_PORT)
@@ -38,7 +37,7 @@ class TestAccessibilityClient(unittest.TestCase):
@patch("artemis.clients.accessibility_client.urllib.request.urlopen") @patch("artemis.clients.accessibility_client.urllib.request.urlopen")
def test_get_hierarchy_xml(self, mock_urlopen): def test_get_hierarchy_xml(self, mock_urlopen):
mock_resp = MagicMock() mock_resp = MagicMock()
mock_resp.read.return_value = b"<hierarchy rotation=\"0\"><node text=\"Save\" /></hierarchy>" mock_resp.read.return_value = b'<hierarchy rotation="0"><node text="Save" /></hierarchy>'
mock_resp.__enter__.return_value = mock_resp mock_resp.__enter__.return_value = mock_resp
mock_urlopen.return_value = mock_resp mock_urlopen.return_value = mock_resp
@@ -52,7 +51,7 @@ class TestAccessibilityClient(unittest.TestCase):
"success": True, "success": True,
"has_screenshot": True, "has_screenshot": True,
"screenshot_base64": "dummy_b64", "screenshot_base64": "dummy_b64",
"xml": "<hierarchy rotation=\"0\"><node text=\"OK\" /></hierarchy>", "xml": '<hierarchy rotation="0"><node text="OK" /></hierarchy>',
"elements": [{"text": "OK"}], "elements": [{"text": "OK"}],
"width": 1080, "width": 1080,
"height": 2400, "height": 2400,
@@ -61,7 +60,7 @@ class TestAccessibilityClient(unittest.TestCase):
self.assertEqual(screen_data.width, 1080) self.assertEqual(screen_data.width, 1080)
self.assertEqual(screen_data.height, 2400) self.assertEqual(screen_data.height, 2400)
self.assertEqual(screen_data.base64, "dummy_b64") self.assertEqual(screen_data.base64, "dummy_b64")
self.assertIn("rotation=\"0\"", screen_data.hierarchy_xml) self.assertIn('rotation="0"', screen_data.hierarchy_xml)
self.assertEqual(len(screen_data.elements), 1) self.assertEqual(len(screen_data.elements), 1)
@patch.object(AccessibilityClient, "get_atomic_snapshot", return_value=None) @patch.object(AccessibilityClient, "get_atomic_snapshot", return_value=None)
@@ -69,10 +68,11 @@ class TestAccessibilityClient(unittest.TestCase):
@patch.object(AccessibilityClient, "get_hierarchy") @patch.object(AccessibilityClient, "get_hierarchy")
def test_get_screen_data_fallback(self, mock_get_hierarchy, mock_get_screenshot, _mock_atomic): def test_get_screen_data_fallback(self, mock_get_hierarchy, mock_get_screenshot, _mock_atomic):
from PIL import Image from PIL import Image
mock_get_screenshot.return_value = Image.new("RGB", (1080, 2400)) mock_get_screenshot.return_value = Image.new("RGB", (1080, 2400))
mock_get_hierarchy.return_value = { mock_get_hierarchy.return_value = {
"success": True, "success": True,
"xml": "<hierarchy rotation=\"0\"><node text=\"Fallback\" /></hierarchy>", "xml": '<hierarchy rotation="0"><node text="Fallback" /></hierarchy>',
"elements": [{"text": "Fallback"}], "elements": [{"text": "Fallback"}],
} }
@@ -85,4 +85,3 @@ class TestAccessibilityClient(unittest.TestCase):
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
Generated
+2982 -2982
View File
File diff suppressed because it is too large. Load diff