weco 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

weco/api.py

@@ -4,7 +4,6 @@ import requests
 from rich.console import Console
 
 from weco import __pkg_version__, __base_url__
-from .constants import CODEGEN_API_TIMEOUT, STATUS_API_TIMEOUT
 from .utils import truncate_output, determine_model_for_onboarding
 
 
@@ -48,6 +47,51 @@ def handle_api_error(e: requests.exceptions.HTTPError, console: Console) -> None
     _render(detail)
 
 
+def _recover_suggest_after_transport_error(
+    console: Console, run_id: str, step: int, auth_headers: dict
+) -> Optional[Dict[str, Any]]:
+    """
+    Try to reconstruct the /suggest response after a transport error (ReadTimeout/502/RemoteDisconnected)
+    by fetching run status and using the latest nodes.
+
+    Args:
+        console: The console object to use for logging.
+        run_id: The ID of the run to recover.
+        step: The step of the solution to recover.
+        auth_headers: The authentication headers to use for the request.
+
+    Returns:
+        The recovered response if the run is in a valid state, otherwise None.
+    """
+    run_status_recovery_response = get_optimization_run_status(
+        console=console, run_id=run_id, include_history=True, auth_headers=auth_headers
+    )
+    current_step = run_status_recovery_response.get("current_step")
+    current_status = run_status_recovery_response.get("status")
+    # The run should be "running" and the current step should correspond to the solution step we are attempting to generate
+    is_valid_run_state = current_status is not None and current_status == "running"
+    is_valid_step = current_step is not None and current_step == step
+    if is_valid_run_state and is_valid_step:
+        nodes = run_status_recovery_response.get("nodes") or []
+        # We need at least 2 nodes to reconstruct the expected response i.e., the last two nodes
+        if len(nodes) >= 2:
+            nodes_sorted_ascending = sorted(nodes, key=lambda n: n["step"])
+            latest_node = nodes_sorted_ascending[-1]
+            penultimate_node = nodes_sorted_ascending[-2]
+            # If the server finished generating the next candidate, it should be exactly this step
+            if latest_node and latest_node["step"] == step:
+                # Try to reconstruct the expected response from the /suggest endpoint using the run status info
+                return {
+                    "run_id": run_id,
+                    "previous_solution_metric_value": penultimate_node.get("metric_value"),
+                    "solution_id": latest_node.get("solution_id"),
+                    "code": latest_node.get("code"),
+                    "plan": latest_node.get("plan"),
+                    "is_done": False,
+                }
+    return None
+
+
 def start_optimization_run(
     console: Console,
     source_code: str,
@@ -64,7 +108,7 @@ def start_optimization_run(
     save_logs: bool = False,
     log_dir: str = ".runs",
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = (10, 3650),
 ) -> Optional[Dict[str, Any]]:
     """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
@@ -107,7 +151,7 @@ def start_optimization_run(
 
 
 def resume_optimization_run(
-    console: Console, run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = STATUS_API_TIMEOUT
+    console: Console, run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = (5, 10)
 ) -> Optional[Dict[str, Any]]:
     """Request the backend to resume an interrupted run."""
     with console.status("[bold green]Resuming run..."):
@@ -136,7 +180,7 @@ def evaluate_feedback_then_suggest_next_solution(
     execution_output: str,
     additional_instructions: str = None,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = (10, 3650),
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
@@ -158,44 +202,38 @@ def evaluate_feedback_then_suggest_next_solution(
             result["code"] = ""
         return result
     except requests.exceptions.ReadTimeout as e:
-        # NOTE: This can occur when:
-        # 1. The server is busy and the request times out
-        # 2. When intermediaries drop the connection so even after the server responds,
-        # the client doesn't receive the response and times out
-
-        # Here we ONLY try to recover in the latter case i.e., server completed request but
-        # client didn't receive the response and timed out
-        run_status_recovery_response = get_optimization_run_status(
-            console=console, run_id=run_id, include_history=True, auth_headers=auth_headers
+        # ReadTimeout can mean either:
+        # 1) the server truly didn't finish before the client's read timeout, or
+        # 2) the server finished but an intermediary (proxy/LB) dropped the response.
+        # We only try to recover case (2): fetch run status to confirm the step completed and reconstruct the response.
+        recovered = _recover_suggest_after_transport_error(
+            console=console, run_id=run_id, step=step, auth_headers=auth_headers
         )
-        current_step = run_status_recovery_response.get("current_step")
-        current_status = run_status_recovery_response.get("status")
-        # The run should be "running" and the current step should correspond to the solution step we are attempting to generate
-        is_valid_run_state = current_status is not None and current_status == "running"
-        is_valid_step = current_step is not None and current_step == step
-        if is_valid_run_state and is_valid_step:
-            nodes = run_status_recovery_response.get("nodes") or []
-            # We need at least 2 nodes to reconstruct the expected response i.e., the last two nodes
-            if len(nodes) >= 2:
-                nodes_sorted_ascending = sorted(nodes, key=lambda n: n["step"])
-                latest_node = nodes_sorted_ascending[-1]
-                penultimate_node = nodes_sorted_ascending[-2]
-                # If the server finished generating the next candidate, it should be exactly this step
-                if latest_node and latest_node["step"] == step:
-                    # Try to reconstruct the expected response from the /suggest endpoint using the run status info
-                    reconstructed_expected_response = {
-                        "run_id": run_id,
-                        "previous_solution_metric_value": penultimate_node.get("metric_value"),
-                        "solution_id": latest_node.get("solution_id"),
-                        "code": latest_node.get("code"),
-                        "plan": latest_node.get("plan"),
-                        "is_done": False,
-                    }
-                    return reconstructed_expected_response
-        # If we couldn't recover, raise the timeout error so the run can be resumed by the user
+        if recovered is not None:
+            return recovered
+        # If we cannot confirm completion, bubble up the timeout so the caller can resume later.
         raise requests.exceptions.ReadTimeout(e)
     except requests.exceptions.HTTPError as e:
-        # Allow caller to handle suggest errors, maybe retry or terminate
+        # Treat only 502 Bad Gateway as a transient transport/gateway issue (akin to a dropped response).
+        # For 502, attempt the status-based recovery method used for ReadTimeout errors; otherwise render the HTTP error normally.
+        if (resp := getattr(e, "response", None)) is not None and resp.status_code == 502:
+            recovered = _recover_suggest_after_transport_error(
+                console=console, run_id=run_id, step=step, auth_headers=auth_headers
+            )
+            if recovered is not None:
+                return recovered
+        # Surface non-502 HTTP errors to the user.
+        handle_api_error(e, console)
+        raise
+    except requests.exceptions.ConnectionError as e:
+        # Covers connection resets with no HTTP response (e.g., RemoteDisconnected).
+        # Treat as a potential "response lost after completion": try status-based recovery first similar to how ReadTimeout errors are handled.
+        recovered = _recover_suggest_after_transport_error(
+            console=console, run_id=run_id, step=step, auth_headers=auth_headers
+        )
+        if recovered is not None:
+            return recovered
+        # Surface the connection error to the user.
         handle_api_error(e, console)
         raise
     except Exception as e:
@@ -208,7 +246,7 @@ def get_optimization_run_status(
     run_id: str,
     include_history: bool = False,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = STATUS_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = (5, 10),
 ) -> Dict[str, Any]:
     """Get the current status of the optimization run."""
     try:
@@ -239,7 +277,7 @@ def get_optimization_run_status(
         raise
 
 
-def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = (10, 10)) -> bool:
+def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: Union[int, Tuple[int, int]] = (5, 10)) -> bool:
     """Send a heartbeat signal to the backend."""
     try:
         response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
@@ -262,7 +300,7 @@ def report_termination(
     reason: str,
     details: Optional[str] = None,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = (10, 30),
+    timeout: Union[int, Tuple[int, int]] = (5, 10),
 ) -> bool:
     """Report the termination reason to the backend."""
     try:
@@ -285,7 +323,7 @@ def get_optimization_suggestions_from_codebase(
     gitingest_tree: str,
     gitingest_content_str: str,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = (10, 3650),
 ) -> Optional[List[Dict[str, Any]]]:
     """Analyze codebase and get optimization suggestions using the model-agnostic backend API."""
     try:
@@ -320,7 +358,7 @@ def generate_evaluation_script_and_metrics(
     description: str,
     gitingest_content_str: str,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = (10, 3650),
 ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
     """Generate evaluation script and determine metrics using the model-agnostic backend API."""
     try:
@@ -356,7 +394,7 @@ def analyze_evaluation_environment(
     gitingest_tree: str,
     gitingest_content_str: str,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = (10, 3650),
 ) -> Optional[Dict[str, Any]]:
     """Analyze existing evaluation scripts and environment using the model-agnostic backend API."""
     try:
@@ -392,7 +430,7 @@ def analyze_script_execution_requirements(
     script_path: str,
     target_file: str,
     auth_headers: dict = {},
-    timeout: Union[int, Tuple[int, int]] = CODEGEN_API_TIMEOUT,
+    timeout: Union[int, Tuple[int, int]] = (10, 3650),
 ) -> Optional[str]:
     """Analyze script to determine proper execution command using the model-agnostic backend API."""
     try:

weco/constants.py

@@ -3,10 +3,6 @@
 Constants for the Weco CLI package.
 """
 
-# API timeout configuration (connect_timeout, read_timeout) in seconds
-CODEGEN_API_TIMEOUT = (30, 3650)
-STATUS_API_TIMEOUT = (10, 180)
-
 # Output truncation configuration
 TRUNCATION_THRESHOLD = 51000  # Maximum length before truncation
 TRUNCATION_KEEP_LENGTH = 25000  # Characters to keep from beginning and end

weco/optimizer.py

@@ -168,7 +168,6 @@ def execute_optimization(
                 reason=f"user_terminated_{signal_name.lower()}",
                 details=f"Process terminated by signal {signal_name} ({signum}).",
                 auth_headers=current_auth_headers_for_heartbeat,
-                timeout=3,
             )
             console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {current_run_id_for_heartbeat}[/]")
 
@@ -547,7 +546,6 @@ def resume_optimization(run_id: str, console: Optional[Console] = None) -> bool:
                 reason=f"user_terminated_{signal_name.lower()}",
                 details=f"Process terminated by signal {signal_name} ({signum}).",
                 auth_headers=current_auth_headers_for_heartbeat,
-                timeout=3,
             )
             console.print(f"\n[cyan]To resume this run, use:[/] [bold cyan]weco resume {current_run_id_for_heartbeat}[/]")
         sys.exit(0)

weco/utils.py

@@ -2,13 +2,13 @@ from typing import Any, Dict, List, Tuple, Union
 import json
 import time
 import subprocess
+import psutil
 from rich.layout import Layout
 from rich.live import Live
 from rich.panel import Panel
 import pathlib
 import requests
 from packaging.version import parse as parse_version
-
 from .constants import TRUNCATION_THRESHOLD, TRUNCATION_KEEP_LENGTH, DEFAULT_MODEL, SUPPORTED_FILE_EXTENSIONS
 
 
@@ -108,22 +108,52 @@ def truncate_output(output: str) -> str:
 
 def run_evaluation(eval_command: str, timeout: int | None = None) -> str:
     """Run the evaluation command on the code and return the output."""
+    process = subprocess.Popen(
+        eval_command, shell=True, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
+    )
 
-    # Run the eval command as is
     try:
-        result = subprocess.run(eval_command, shell=True, capture_output=True, text=True, check=False, timeout=timeout)
-        # Combine stdout and stderr for complete output
-        output = result.stderr if result.stderr else ""
-        if result.stdout:
-            if len(output) > 0:
-                output += "\n"
-            output += result.stdout
-        return output  # Return full output, no truncation
+        # NOTE: Process tree cleanup only happens on timeout. Normal completion relies on the OS/shell to clean up child processes, which works for typical evaluation scripts.
+        output, _ = process.communicate(timeout=timeout)
+        return output
+
     except subprocess.TimeoutExpired:
+        # Kill process tree
+        try:
+            parent = psutil.Process(process.pid)
+            children = parent.children(recursive=True)
+
+            # Terminate gracefully
+            for child in children:
+                try:
+                    child.terminate()
+                except psutil.NoSuchProcess:
+                    pass
+            try:
+                parent.terminate()
+            except psutil.NoSuchProcess:
+                pass
+
+            # Wait, then force kill survivors
+            _, alive = psutil.wait_procs(children + [parent], timeout=1)
+            for proc in alive:
+                try:
+                    proc.kill()
+                except psutil.NoSuchProcess:
+                    pass
+
+        except psutil.NoSuchProcess:
+            pass
+
+        # Drain pipes
+        try:
+            process.communicate(timeout=1)
+        except (subprocess.TimeoutExpired, ValueError, OSError):
+            pass
+
         return f"Evaluation timed out after {'an unspecified duration' if timeout is None else f'{timeout} seconds'}."
 
 
-# Update Check Function
 def check_for_cli_updates():
     """Checks PyPI for a newer version of the weco package and notifies the user."""
     try:

{weco-0.3.0.dist-info → weco-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.3.0
+Version: 0.3.2
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: 
@@ -219,6 +219,7 @@ Requires-Dist: packaging
 Requires-Dist: gitingest
 Requires-Dist: fastapi
 Requires-Dist: slowapi
+Requires-Dist: psutil
 Provides-Extra: dev
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: build; extra == "dev"

weco-0.3.2.dist-info/RECORD

@@ -0,0 +1,16 @@
+weco/__init__.py,sha256=ClO0uT6GKOA0iSptvP0xbtdycf0VpoPTq37jHtvlhtw,303
+weco/api.py,sha256=dUjzuOKKvayzZ_1B4j40eK9Ofk264jsc6vOR1afsszY,18523
+weco/auth.py,sha256=O31Hoj-Loi8DWJJG2LfeWgUMuNqAUeGDpd2ZGjA9Ah0,9997
+weco/chatbot.py,sha256=EIK2WaOul9gn_yHLThjsZV7RnE8t3XQPwgRkO5tybSU,38415
+weco/cli.py,sha256=579f6jf-ZWuFAmNXDisRY7zWr7vw2YZQuC_QX8-qxx0,11460
+weco/constants.py,sha256=V6yFugTznKm5EC2_jr4I_whd7sqI80HiPggRn0az580,406
+weco/credits.py,sha256=C08x-TRcLg3ccfKqMGNRY7zBn7t3r7LZ119bxgfztaI,7629
+weco/optimizer.py,sha256=mJU8_0bo_6dS2PEj1E3dQHvNH9V4e8NSLNE55tmvspw,42291
+weco/panels.py,sha256=fnGPtmvxpx21AuBCtCFu1f_BpSxybNr2lhjIIKIutrY,16133
+weco/utils.py,sha256=erDDrA_g3KSlel6YEAGALlV_k8ftT-VQnPT1BrmzK8k,7021
+weco-0.3.2.dist-info/licenses/LICENSE,sha256=9LUfoGHjLPtak2zps2kL2tm65HAZIICx_FbLaRuS4KU,11337
+weco-0.3.2.dist-info/METADATA,sha256=TDJIvT1vw3VFrjEj9o8VkLuxis2MWWAL0pnDYqpFfak,31878
+weco-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+weco-0.3.2.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.3.2.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.3.2.dist-info/RECORD,,

weco-0.3.0.dist-info/RECORD

@@ -1,16 +0,0 @@
-weco/__init__.py,sha256=ClO0uT6GKOA0iSptvP0xbtdycf0VpoPTq37jHtvlhtw,303
-weco/api.py,sha256=JQ-zqRUMVuWP1eq9ojihMCHlX7ziQjrDCyJjgbqT5e4,16909
-weco/auth.py,sha256=O31Hoj-Loi8DWJJG2LfeWgUMuNqAUeGDpd2ZGjA9Ah0,9997
-weco/chatbot.py,sha256=EIK2WaOul9gn_yHLThjsZV7RnE8t3XQPwgRkO5tybSU,38415
-weco/cli.py,sha256=579f6jf-ZWuFAmNXDisRY7zWr7vw2YZQuC_QX8-qxx0,11460
-weco/constants.py,sha256=_0a50O4nRdkEwMhNIOi04bTrGhc06ZJeotjGae6t-l4,542
-weco/credits.py,sha256=C08x-TRcLg3ccfKqMGNRY7zBn7t3r7LZ119bxgfztaI,7629
-weco/optimizer.py,sha256=Tjnnx3Xq5ZcMggW8IMHhjsb3RwVrtdh7Z4rbJV0VjAs,42345
-weco/panels.py,sha256=fnGPtmvxpx21AuBCtCFu1f_BpSxybNr2lhjIIKIutrY,16133
-weco/utils.py,sha256=TT57S0YGMuMWPFNsn0tcexNHZd-kBEjDeiOLWxANiQU,6117
-weco-0.3.0.dist-info/licenses/LICENSE,sha256=9LUfoGHjLPtak2zps2kL2tm65HAZIICx_FbLaRuS4KU,11337
-weco-0.3.0.dist-info/METADATA,sha256=OOML7kA7wFKR6HWNt7D7xXa6TBySS8iMa_K6GwA7TNw,31856
-weco-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-weco-0.3.0.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.3.0.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.3.0.dist-info/RECORD,,