diff --git a/src/serena/jetbrains/jetbrains_plugin_client.py b/src/serena/jetbrains/jetbrains_plugin_client.py index 88e32377..f11dc080 100644 --- a/src/serena/jetbrains/jetbrains_plugin_client.py +++ b/src/serena/jetbrains/jetbrains_plugin_client.py @@ -2,9 +2,12 @@ Client for the Serena JetBrains Plugin """ +import concurrent import json import logging import re +import threading +from concurrent.futures.thread import ThreadPoolExecutor from pathlib import Path from typing import Any, Literal, Optional, Self, TypeVar, cast @@ -15,6 +18,7 @@ from sensai.util.string import ToStringMixin import serena.jetbrains.jetbrains_types as jb from serena.jetbrains.jetbrains_types import PluginStatusDTO from serena.project import Project +from serena.util.class_decorators import singleton from serena.util.text_utils import render_html from serena.util.version import Version @@ -59,6 +63,60 @@ class ServerNotFoundError(Exception): """Raised when the plugin's service is not found.""" +@singleton +class JetBrainsPluginClientManager: + """ + Manager for JetBrainsPluginClient instances, responsible for scanning ports to find available plugin instances + """ + + NUM_PORTS_TO_SCAN = 20 + + def __init__(self) -> None: + self._clients: dict[int, "JetBrainsPluginClient"] = {} + self._lock = threading.Lock() + + def _submit_scan(self) -> list[concurrent.futures.Future["JetBrainsPluginClient"]]: + """ + Performs a port scan to find available plugin instances in parallel. + + :return: futures that will resolve to plugin clients for every port + """ + + def scan_port(port: int) -> JetBrainsPluginClient: + client = JetBrainsPluginClient(port) + with self._lock: + self._clients[port] = client + return client + + futures = [] + with ThreadPoolExecutor(max_workers=self.NUM_PORTS_TO_SCAN) as executor: + for i in range(self.NUM_PORTS_TO_SCAN): + future = executor.submit(scan_port, JetBrainsPluginClient.BASE_PORT + i) + futures.append(future) + return futures + + def find_client(self, project_root: Path) -> "JetBrainsPluginClient": + plugin_paths_found = [] + for future in self._submit_scan(): + client = future.result() + if client.matches(project_root): + return client + elif client.project_root is not None: + plugin_paths_found.append(client.project_root) + + log.warning( + "Searched for Serena JetBrains plugin service for project at %s but found no matching service. " + "Found plugin instances for the following project paths: %s", + project_root, + plugin_paths_found, + ) + raise ServerNotFoundError( + f"Found no Serena service in a JetBrains IDE instance for the project at {project_root}. " + "STOP. Do not attempt any other tools or workarounds. Ask the user to open this folder as a project in a JetBrains IDE " + "with the Serena plugin installed and running!" + ) + + class JetBrainsPluginClient(ToStringMixin): """ Python client for the Serena Backend Service. @@ -87,11 +145,11 @@ class JetBrainsPluginClient(ToStringMixin): self._session.headers.update({"Content-Type": "application/json", "Accept": "application/json"}) # connect and obtain status - self._project_root: str | None = None + self.project_root: str | None = None self._plugin_version: Version | None = None try: status_response: PluginStatusDTO = cast(jb.PluginStatusDTO, self._make_request("GET", "/status")) - self._project_root = status_response["project_root"] + self.project_root = status_response["project_root"] self._plugin_version = Version(status_response["plugin_version"]) except ConnectionError: # expected if no server is running at the port pass @@ -107,7 +165,7 @@ class JetBrainsPluginClient(ToStringMixin): cls._server_address = address def _tostring_includes(self) -> list[str]: - return ["_port", "_project_root", "_plugin_version"] + return ["_port", "project_root", "_plugin_version"] @classmethod def from_project(cls, project: Project) -> Self: @@ -118,27 +176,9 @@ class JetBrainsPluginClient(ToStringMixin): if client.matches(resolved_path): return client - plugin_paths_found = [] - for port in range(cls.BASE_PORT, cls.BASE_PORT + 20): - client = JetBrainsPluginClient(port) - if client.matches(resolved_path): - log.info("Found matching %s", client) - cls._last_port = port - return client - elif client._project_root is not None: - plugin_paths_found.append(client._project_root) - - log.warning( - "Searched for Serena JetBrains plugin service for project at %s but found no matching service. " - "Found plugin instances for the following project paths: %s", - resolved_path, - plugin_paths_found, - ) - raise ServerNotFoundError( - f"Found no Serena service in a JetBrains IDE instance for the project at {resolved_path}. " - "STOP. Do not attempt any other tools or workarounds. Ask the user to open this folder as a project in a JetBrains IDE " - "with the Serena plugin installed and running!" - ) + client = JetBrainsPluginClientManager().find_client(resolved_path) + cls._last_port = client._port + return client @staticmethod def _paths_match(resolved_serena_path: str, plugin_path: str) -> bool: @@ -184,9 +224,9 @@ class JetBrainsPluginClient(ToStringMixin): :param resolved_path: the resolved project root path from Serena's perspective :return: whether this client instance matches the given project path """ - if self._project_root is None: + if self.project_root is None: return False - return self._paths_match(str(resolved_path), self._project_root) + return self._paths_match(str(resolved_path), self.project_root) def is_version_at_least(self, *version_parts: int) -> bool: if self._plugin_version is None: