weco 0.2.19__py3-none-any.whl → 0.2.20__py3-none-any.whl

weco/api.py

@@ -17,7 +17,7 @@ def handle_api_error(e: requests.exceptions.HTTPError, console: rich.console.Con
     # sys.exit(1)
 
 
-def start_optimization_session(
+def start_optimization_run(
     console: rich.console.Console,
     source_code: str,
     evaluation_command: str,
@@ -29,14 +29,14 @@ def start_optimization_session(
     search_policy_config: Dict[str, Any],
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
-    auth_headers: dict = {},  # Add auth_headers
+    auth_headers: dict = {},
     timeout: int = 800,
 ) -> Dict[str, Any]:
-    """Start the optimization session."""
+    """Start the optimization run."""
     with console.status("[bold green]Starting Optimization..."):
         try:
             response = requests.post(
-                f"{__base_url__}/sessions",  # Path is relative to base_url
+                f"{__base_url__}/runs",
                 json={
                     "source_code": source_code,
                     "additional_instructions": additional_instructions,
@@ -49,37 +49,37 @@ def start_optimization_session(
                     },
                     "metadata": {"client_name": "cli", "client_version": __pkg_version__, **api_keys},
                 },
-                headers=auth_headers,  # Add headers
+                headers=auth_headers,
                 timeout=timeout,
             )
             response.raise_for_status()
             return response.json()
         except requests.exceptions.HTTPError as e:
             handle_api_error(e, console)
-            sys.exit(1)  # Exit if starting session fails
+            sys.exit(1)
         except requests.exceptions.RequestException as e:
-            console.print(f"[bold red]Network Error starting session: {e}[/]")
+            console.print(f"[bold red]Network Error starting run: {e}[/]")
             sys.exit(1)
 
 
 def evaluate_feedback_then_suggest_next_solution(
-    session_id: str,
+    run_id: str,
     execution_output: str,
     additional_instructions: str = None,
     api_keys: Dict[str, Any] = {},
-    auth_headers: dict = {},  # Add auth_headers
+    auth_headers: dict = {},
     timeout: int = 800,
 ) -> Dict[str, Any]:
     """Evaluate the feedback and suggest the next solution."""
     try:
         response = requests.post(
-            f"{__base_url__}/sessions/{session_id}/suggest",  # Path is relative to base_url
+            f"{__base_url__}/runs/{run_id}/suggest",
             json={
                 "execution_output": execution_output,
                 "additional_instructions": additional_instructions,
                 "metadata": {**api_keys},
             },
-            headers=auth_headers,  # Add headers
+            headers=auth_headers,
             timeout=timeout,
         )
         response.raise_for_status()
@@ -93,16 +93,13 @@ def evaluate_feedback_then_suggest_next_solution(
         raise  # Re-raise the exception
 
 
-def get_optimization_session_status(
-    session_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
+def get_optimization_run_status(
+    run_id: str, include_history: bool = False, auth_headers: dict = {}, timeout: int = 800
 ) -> Dict[str, Any]:
-    """Get the current status of the optimization session."""
+    """Get the current status of the optimization run."""
     try:
         response = requests.get(
-            f"{__base_url__}/sessions/{session_id}",  # Path is relative to base_url
-            params={"include_history": include_history},
-            headers=auth_headers,
-            timeout=timeout,
+            f"{__base_url__}/runs/{run_id}", params={"include_history": include_history}, headers=auth_headers, timeout=timeout
         )
         response.raise_for_status()
         return response.json()
@@ -114,42 +111,30 @@ def get_optimization_session_status(
         raise  # Re-raise
 
 
-def send_heartbeat(
-    session_id: str,
-    auth_headers: dict = {},
-    timeout: int = 10,  # Shorter timeout for non-critical heartbeat
-) -> bool:
+def send_heartbeat(run_id: str, auth_headers: dict = {}, timeout: int = 10) -> bool:
     """Send a heartbeat signal to the backend."""
     try:
-        response = requests.put(f"{__base_url__}/sessions/{session_id}/heartbeat", headers=auth_headers, timeout=timeout)
-        response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
+        response = requests.put(f"{__base_url__}/runs/{run_id}/heartbeat", headers=auth_headers, timeout=timeout)
+        response.raise_for_status()
         return True
     except requests.exceptions.HTTPError as e:
-        # Log non-critical errors like 409 Conflict (session not running)
         if e.response.status_code == 409:
-            print(f"Heartbeat ignored: Session {session_id} is not running.", file=sys.stderr)
+            print(f"Heartbeat ignored: Run {run_id} is not running.", file=sys.stderr)
         else:
-            print(f"Heartbeat failed for session {session_id}: HTTP {e.response.status_code}", file=sys.stderr)
-        # Don't exit, just report failure
+            print(f"Heartbeat failed for run {run_id}: HTTP {e.response.status_code}", file=sys.stderr)
         return False
     except requests.exceptions.RequestException as e:
-        # Network errors are also non-fatal for heartbeats
-        print(f"Heartbeat network error for session {session_id}: {e}", file=sys.stderr)
+        print(f"Heartbeat network error for run {run_id}: {e}", file=sys.stderr)
         return False
 
 
 def report_termination(
-    session_id: str,
-    status_update: str,
-    reason: str,
-    details: Optional[str] = None,
-    auth_headers: dict = {},
-    timeout: int = 30,  # Reasonably longer timeout for important termination message
+    run_id: str, status_update: str, reason: str, details: Optional[str] = None, auth_headers: dict = {}, timeout: int = 30
 ) -> bool:
     """Report the termination reason to the backend."""
     try:
         response = requests.post(
-            f"{__base_url__}/sessions/{session_id}/terminate",
+            f"{__base_url__}/runs/{run_id}/terminate",
             json={"status_update": status_update, "termination_reason": reason, "termination_details": details},
             headers=auth_headers,
             timeout=timeout,
@@ -157,6 +142,5 @@ def report_termination(
         response.raise_for_status()
         return True
     except requests.exceptions.RequestException as e:
-        # Log failure, but don't prevent CLI exit
-        print(f"Warning: Failed to report termination to backend for session {session_id}: {e}", file=sys.stderr)
+        print(f"Warning: Failed to report termination to backend for run {run_id}: {e}", file=sys.stderr)
         return False

weco/cli.py

@@ -14,9 +14,9 @@ from rich.panel import Panel
 from rich.traceback import install
 from rich.prompt import Prompt
 from .api import (
-    start_optimization_session,
+    start_optimization_run,
     evaluate_feedback_then_suggest_next_solution,
-    get_optimization_session_status,
+    get_optimization_run_status,
     handle_api_error,
     send_heartbeat,
     report_termination,
@@ -51,15 +51,15 @@ console = Console()
 # --- Global variable for heartbeat thread ---
 heartbeat_thread = None
 stop_heartbeat_event = threading.Event()
-current_session_id_for_heartbeat = None
+current_run_id_for_heartbeat = None
 current_auth_headers_for_heartbeat = {}
 
 
 # --- Heartbeat Sender Class ---
 class HeartbeatSender(threading.Thread):
-    def __init__(self, session_id: str, auth_headers: dict, stop_event: threading.Event, interval: int = 30):
+    def __init__(self, run_id: str, auth_headers: dict, stop_event: threading.Event, interval: int = 30):
         super().__init__(daemon=True)  # Daemon thread exits when main thread exits
-        self.session_id = session_id
+        self.run_id = run_id
         self.auth_headers = auth_headers
         self.interval = interval
         self.stop_event = stop_event
@@ -67,7 +67,7 @@ class HeartbeatSender(threading.Thread):
     def run(self):
         try:
             while not self.stop_event.is_set():
-                if not send_heartbeat(self.session_id, self.auth_headers):
+                if not send_heartbeat(self.run_id, self.auth_headers):
                     # send_heartbeat itself prints errors to stderr if it returns False
                     # No explicit HeartbeatSender log needed here unless more detail is desired for a False return
                     pass  # Continue trying as per original logic
@@ -79,9 +79,7 @@ class HeartbeatSender(threading.Thread):
 
         except Exception as e:
             # Catch any unexpected error in the loop to prevent silent thread death
-            print(
-                f"[ERROR HeartbeatSender] Unhandled exception in run loop for session {self.session_id}: {e}", file=sys.stderr
-            )
+            print(f"[ERROR HeartbeatSender] Unhandled exception in run loop for run {self.run_id}: {e}", file=sys.stderr)
             traceback.print_exc(file=sys.stderr)
             # The loop will break due to the exception, and thread will terminate via finally.
 
@@ -97,13 +95,14 @@ def signal_handler(signum, frame):
         heartbeat_thread.join(timeout=2)  # Give it a moment to stop
 
     # Report termination (best effort)
-    if current_session_id_for_heartbeat:
+    if current_run_id_for_heartbeat:
         report_termination(
-            session_id=current_session_id_for_heartbeat,
+            run_id=current_run_id_for_heartbeat,
             status_update="terminated",
             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,
         )
 
     # Exit gracefully
@@ -158,7 +157,7 @@ def perform_login(console: Console):
 
                 try:
                     token_response = requests.post(
-                        f"{__base_url__}/auth/device/token",  # REMOVED /v1 prefix
+                        f"{__base_url__}/auth/device/token",
                         json={"grant_type": "urn:ietf:params:oauth:grant-type:device_code", "device_code": device_code},
                     )
 
@@ -172,12 +171,10 @@ def perform_login(console: Console):
                             # Unexpected 202 response format
                             console.print(f"\n[bold red]Error:[/] Received unexpected 202 response: {token_data}")
                             return False
-
                     # Check for standard OAuth2 errors (often 400 Bad Request)
                     elif token_response.status_code == 400:
                         token_data = token_response.json()
                         error_code = token_data.get("error", "unknown_error")
-                        # NOTE: Removed "authorization_pending" check from here
                         if error_code == "slow_down":
                             interval += 5  # Increase polling interval if instructed
                             live_status.update(f"Waiting... (slowing down polling to {interval}s)")
@@ -195,7 +192,6 @@ def perform_login(console: Console):
 
                     # Check for other non-200/non-202/non-400 HTTP errors
                     token_response.raise_for_status()
-
                     # If successful (200 OK and no 'error' field)
                     token_data = token_response.json()
                     if "access_token" in token_data:
@@ -206,19 +202,17 @@ def perform_login(console: Console):
                     else:
                         # Unexpected successful response format
                         console.print("\n[bold red]Error:[/] Received unexpected response from server during polling.")
-                        print(token_data)  # Log for debugging
+                        print(token_data)
                         return False
-
                 except requests.exceptions.RequestException as e:
                     # Handle network errors during polling gracefully
                     live_status.update("Waiting... (network error, retrying)")
                     console.print(f"\n[bold yellow]Warning:[/] Network error during polling: {e}. Retrying...")
-                    # Optional: implement backoff strategy
                     time.sleep(interval * 2)  # Simple backoff
-
-    except requests.exceptions.HTTPError as e:  # Catch HTTPError specifically for handle_api_error
+    except requests.exceptions.HTTPError as e:
         handle_api_error(e, console)
-    except requests.exceptions.RequestException as e:  # Catch other request errors
+    except requests.exceptions.RequestException as e:
+        # Catch other request errors
         console.print(f"\n[bold red]Network Error:[/] {e}")
         return False
     except Exception as e:
@@ -233,17 +227,15 @@ def main() -> None:
     signal.signal(signal.SIGTERM, signal_handler)
 
     # --- Perform Update Check ---
-    # Import __pkg_version__ here to avoid circular import issues if it's also used in modules imported by cli.py
     from . import __pkg_version__
 
-    check_for_cli_updates(__pkg_version__)  # Call the imported function
+    check_for_cli_updates(__pkg_version__)
 
     # --- Argument Parsing ---
     parser = argparse.ArgumentParser(
         description="[bold cyan]Weco CLI[/]", formatter_class=argparse.RawDescriptionHelpFormatter
     )
-    # Add subparsers for commands like 'run' and 'logout'
-    subparsers = parser.add_subparsers(dest="command", help="Available commands", required=True)  # Make command required
+    subparsers = parser.add_subparsers(dest="command", help="Available commands", required=True)
 
     # --- Run Command ---
     run_parser = subparsers.add_parser(
@@ -298,33 +290,27 @@ def main() -> None:
         help="Description of additional instruction or path to a file containing additional instructions. Defaults to None.",
     )
 
-    # --- Logout Command ---
     _ = subparsers.add_parser("logout", help="Log out from Weco and clear saved API key.")
-
     args = parser.parse_args()
 
-    # --- Handle Logout Command ---
     if args.command == "logout":
         clear_api_key()
         sys.exit(0)
-
-    # --- Handle Run Command ---
     elif args.command == "run":
-        global heartbeat_thread, current_session_id_for_heartbeat, current_auth_headers_for_heartbeat  # Allow modification of globals
-
-        session_id = None  # Initialize session_id
-        optimization_completed_normally = False  # Flag for finally block
-        # --- Check Authentication ---
+        global heartbeat_thread, current_run_id_for_heartbeat, current_auth_headers_for_heartbeat  # Allow modification of globals
+        run_id = None  # Initialize run_id (we receive this from the API after starting the run)
+        optimization_completed_normally = False
+        user_stop_requested_flag = False
         weco_api_key = load_weco_api_key()
         llm_api_keys = read_api_keys_from_env()  # Read keys from client environment
 
+        # --- Login/Authentication Handling ---
         if not weco_api_key:
             login_choice = Prompt.ask(
                 "Log in to Weco to save run history or use anonymously? ([bold]L[/]ogin / [bold]S[/]kip)",
                 choices=["l", "s"],
                 default="s",
             ).lower()
-
             if login_choice == "l":
                 console.print("[cyan]Starting login process...[/]")
                 if not perform_login(console):
@@ -334,7 +320,6 @@ def main() -> None:
                 if not weco_api_key:
                     console.print("[bold red]Error: Login completed but failed to retrieve API key.[/]")
                     sys.exit(1)
-
             elif login_choice == "s":
                 console.print("[yellow]Proceeding anonymously. LLM API keys must be provided via environment variables.[/]")
                 if not llm_api_keys:
@@ -343,7 +328,6 @@ def main() -> None:
                     )
                     sys.exit(1)
 
-        # --- Prepare API Call Arguments ---
         auth_headers = {}
         if weco_api_key:
             auth_headers["Authorization"] = f"Bearer {weco_api_key}"
@@ -351,7 +335,7 @@ def main() -> None:
 
         # --- Main Run Logic ---
         try:
-            # --- Configuration Loading ---
+            # --- Read Command Line Arguments ---
             evaluation_command = args.eval_command
             metric_name = args.metric
             maximize = args.goal in ["maximize", "max"]
@@ -375,13 +359,9 @@ def main() -> None:
                 "debug_prob": 0.5,
                 "max_debug_depth": max(1, math.ceil(0.1 * steps)),
             }
-            # API request timeout
             timeout = 800
-            # Read additional instructions
             additional_instructions = read_additional_instructions(additional_instructions=args.additional_instructions)
-            # Read source code path
             source_fp = pathlib.Path(args.source)
-            # Read source code content
             source_code = read_from_path(fp=source_fp, is_json=False)
 
             # --- Panel Initialization ---
@@ -395,8 +375,8 @@ def main() -> None:
             layout = create_optimization_layout()
             end_optimization_layout = create_end_optimization_layout()
 
-            # --- Start Optimization Session ---
-            session_response = start_optimization_session(
+            # --- Start Optimization Run ---
+            run_response = start_optimization_run(
                 console=console,
                 source_code=source_code,
                 evaluation_command=evaluation_command,
@@ -411,42 +391,39 @@ def main() -> None:
                 auth_headers=auth_headers,
                 timeout=timeout,
             )
-            session_id = session_response["session_id"]
-            current_session_id_for_heartbeat = session_id  # Store for signal handler/finally
+            run_id = run_response["run_id"]
+            current_run_id_for_heartbeat = run_id
 
             # --- Start Heartbeat Thread ---
-            stop_heartbeat_event.clear()  # Ensure event is clear before starting
-            heartbeat_thread = HeartbeatSender(session_id, auth_headers, stop_heartbeat_event)
+            stop_heartbeat_event.clear()
+            heartbeat_thread = HeartbeatSender(run_id, auth_headers, stop_heartbeat_event)
             heartbeat_thread.start()
 
             # --- Live Update Loop ---
             refresh_rate = 4
             with Live(layout, refresh_per_second=refresh_rate, screen=True) as live:
-                # Define the runs directory (.runs/<session-id>)
-                runs_dir = pathlib.Path(args.log_dir) / session_id
+                # Define the runs directory (.runs/<run-id>) to store logs and results
+                runs_dir = pathlib.Path(args.log_dir) / run_id
                 runs_dir.mkdir(parents=True, exist_ok=True)
-
                 # Write the initial code string to the logs
-                write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=session_response["code"])
-
+                write_to_path(fp=runs_dir / f"step_0{source_fp.suffix}", content=run_response["code"])
                 # Write the initial code string to the source file path
-                write_to_path(fp=source_fp, content=session_response["code"])
+                write_to_path(fp=source_fp, content=run_response["code"])
 
                 # Update the panels with the initial solution
-                summary_panel.set_session_id(session_id=session_id)  # Add session id now that we have it
+                summary_panel.set_run_id(run_id=run_id)  # Add run id now that we have it
                 # Set the step of the progress bar
                 summary_panel.set_step(step=0)
                 # Update the token counts
-                summary_panel.update_token_counts(usage=session_response["usage"])
-                # Update the plan
-                plan_panel.update(plan=session_response["plan"])
+                summary_panel.update_token_counts(usage=run_response["usage"])
+                plan_panel.update(plan=run_response["plan"])
                 # Build the metric tree
                 tree_panel.build_metric_tree(
                     nodes=[
                         {
-                            "solution_id": session_response["solution_id"],
+                            "solution_id": run_response["solution_id"],
                             "parent_id": None,
-                            "code": session_response["code"],
+                            "code": run_response["code"],
                             "step": 0,
                             "metric_value": None,
                             "is_buggy": False,
@@ -454,15 +431,11 @@ def main() -> None:
                     ]
                 )
                 # Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
-                tree_panel.set_unevaluated_node(node_id=session_response["solution_id"])
+                tree_panel.set_unevaluated_node(node_id=run_response["solution_id"])
                 # Update the solution panels with the initial solution and get the panel displays
                 solution_panels.update(
                     current_node=Node(
-                        id=session_response["solution_id"],
-                        parent_id=None,
-                        code=session_response["code"],
-                        metric=None,
-                        is_buggy=False,
+                        id=run_response["solution_id"], parent_id=None, code=run_response["code"], metric=None, is_buggy=False
                     ),
                     best_node=None,
                 )
@@ -482,12 +455,8 @@ def main() -> None:
                     transition_delay=0.1,
                 )
 
-                # # Send initial heartbeat immediately after starting
-                # send_heartbeat(session_id, auth_headers)
-
                 # Run evaluation on the initial solution
                 term_out = run_evaluation(eval_command=args.eval_command)
-
                 # Update the evaluation output panel
                 eval_output_panel.update(output=term_out)
                 smooth_update(
@@ -503,40 +472,50 @@ def main() -> None:
                     current_additional_instructions = read_additional_instructions(
                         additional_instructions=args.additional_instructions
                     )
+                    if run_id:
+                        try:
+                            current_status_response = get_optimization_run_status(
+                                run_id=run_id, include_history=False, timeout=30, auth_headers=auth_headers
+                            )
+                            current_run_status_val = current_status_response.get("status")
+                            if current_run_status_val == "stopping":
+                                console.print("\n[bold yellow]Stop request received. Terminating run gracefully...[/]")
+                                user_stop_requested_flag = True
+                                break
+                        except requests.exceptions.RequestException as e:
+                            console.print(
+                                f"\n[bold red]Warning: Could not check run status: {e}. Continuing optimization...[/]"
+                            )
+                        except Exception as e:
+                            console.print(
+                                f"\n[bold red]Warning: Error checking run status: {e}. Continuing optimization...[/]"
+                            )
 
                     # Send feedback and get next suggestion
                     eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
-                        session_id=session_id,
+                        run_id=run_id,
                         execution_output=term_out,
-                        additional_instructions=current_additional_instructions,  # Pass current instructions
-                        api_keys=llm_api_keys,  # Pass client LLM keys
-                        auth_headers=auth_headers,  # Pass Weco key if logged in
+                        additional_instructions=current_additional_instructions,
+                        api_keys=llm_api_keys,
+                        auth_headers=auth_headers,
                         timeout=timeout,
                     )
-
-                    # Save next solution (.runs/<session-id>/step_<step>.<extension>)
+                    # Save next solution (.runs/<run-id>/step_<step>.<extension>)
                     write_to_path(
                         fp=runs_dir / f"step_{step}{source_fp.suffix}", content=eval_and_next_solution_response["code"]
                     )
-
                     # Write the next solution to the source file
                     write_to_path(fp=source_fp, content=eval_and_next_solution_response["code"])
-
-                    # Get the optimization session status for
-                    # the best solution, its score, and the history to plot the tree
-                    status_response = get_optimization_session_status(
-                        session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
+                    status_response = get_optimization_run_status(
+                        run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
                     )
-
-                    # Update the step of the progress bar
+                    # Update the step of the progress bar, token counts, plan and metric tree
                     summary_panel.set_step(step=step)
-                    # Update the token counts
                     summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
-                    # Update the plan
                     plan_panel.update(plan=eval_and_next_solution_response["plan"])
-                    # Build the metric tree
-                    tree_panel.build_metric_tree(nodes=status_response["history"])
-                    # Set the current solution as unevaluated since we haven't run the evaluation function and fed it back to the model yet
+
+                    nodes_list_from_status = status_response.get("nodes")
+                    tree_panel.build_metric_tree(nodes=nodes_list_from_status if nodes_list_from_status is not None else [])
                     tree_panel.set_unevaluated_node(node_id=eval_and_next_solution_response["solution_id"])
 
                     # Update the solution panels with the next solution and best solution (and score)
@@ -552,27 +531,25 @@ def main() -> None:
                     else:
                         best_solution_node = None
 
-                    # Create a node for the current solution
                     current_solution_node = None
-                    for node in status_response["history"]:
-                        if node["solution_id"] == eval_and_next_solution_response["solution_id"]:
-                            current_solution_node = Node(
-                                id=node["solution_id"],
-                                parent_id=node["parent_id"],
-                                code=node["code"],
-                                metric=node["metric_value"],
-                                is_buggy=node["is_buggy"],
-                            )
+                    if status_response.get("nodes"):
+                        for node_data in status_response["nodes"]:
+                            if node_data["solution_id"] == eval_and_next_solution_response["solution_id"]:
+                                current_solution_node = Node(
+                                    id=node_data["solution_id"],
+                                    parent_id=node_data["parent_id"],
+                                    code=node_data["code"],
+                                    metric=node_data["metric_value"],
+                                    is_buggy=node_data["is_buggy"],
+                                )
                     if current_solution_node is None:
-                        raise ValueError("Current solution node not found in history")
+                        raise ValueError("Current solution node not found in nodes list from status response")
+
                     # Update the solution panels with the current and best solution
                     solution_panels.update(current_node=current_solution_node, best_node=best_solution_node)
                     current_solution_panel, best_solution_panel = solution_panels.get_display(current_step=step)
-
                     # Clear evaluation output since we are running a evaluation on a new solution
                     eval_output_panel.clear()
-
-                    # Update displays with smooth transitions
                     smooth_update(
                         live=live,
                         layout=layout,
@@ -586,171 +563,152 @@ def main() -> None:
                         ],
                         transition_delay=0.08,  # Slightly longer delay for more noticeable transitions
                     )
-
-                    # Run evaluation on the current solution
                     term_out = run_evaluation(eval_command=args.eval_command)
                     eval_output_panel.update(output=term_out)
-
-                    # Update evaluation output with a smooth transition
                     smooth_update(
                         live=live,
                         layout=layout,
                         sections_to_update=[("eval_output", eval_output_panel.get_display())],
-                        transition_delay=0.1,  # Slightly longer delay for evaluation results
+                        transition_delay=0.1,
                     )
 
-                # Re-read instructions from the original source (file path or string) BEFORE each suggest call
-                current_additional_instructions = read_additional_instructions(
-                    additional_instructions=args.additional_instructions
-                )
-
-                # Final evaluation report
-                eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
-                    session_id=session_id,
-                    execution_output=term_out,
-                    additional_instructions=current_additional_instructions,
-                    api_keys=llm_api_keys,
-                    timeout=timeout,
-                    auth_headers=auth_headers,
-                )
-
-                # Update the progress bar
-                summary_panel.set_step(step=steps)
-                # Update the token counts
-                summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
-                # No need to update the plan panel since we have finished the optimization
-                # Get the optimization session status for
-                # the best solution, its score, and the history to plot the tree
-                status_response = get_optimization_session_status(
-                    session_id=session_id, include_history=True, timeout=timeout, auth_headers=auth_headers
-                )
-                # Build the metric tree
-                tree_panel.build_metric_tree(nodes=status_response["history"])
-                # No need to set any solution to unevaluated since we have finished the optimization
-                # and all solutions have been evaluated
-                # No neeed to update the current solution panel since we have finished the optimization
-                # We only need to update the best solution panel
-                # Figure out if we have a best solution so far
-                if status_response["best_result"] is not None:
-                    best_solution_node = Node(
-                        id=status_response["best_result"]["solution_id"],
-                        parent_id=status_response["best_result"]["parent_id"],
-                        code=status_response["best_result"]["code"],
-                        metric=status_response["best_result"]["metric_value"],
-                        is_buggy=status_response["best_result"]["is_buggy"],
+                if not user_stop_requested_flag:
+                    # Re-read instructions from the original source (file path or string) BEFORE each suggest call
+                    current_additional_instructions = read_additional_instructions(
+                        additional_instructions=args.additional_instructions
                     )
-                else:
-                    best_solution_node = None
-                solution_panels.update(current_node=None, best_node=best_solution_node)
-                _, best_solution_panel = solution_panels.get_display(current_step=steps)
-
-                # Update the end optimization layout
-                final_message = (
-                    f"{summary_panel.metric_name.capitalize()} {'maximized' if summary_panel.maximize else 'minimized'}! Best solution {summary_panel.metric_name.lower()} = [green]{status_response['best_result']['metric_value']}[/] 🏆"
-                    if best_solution_node is not None and best_solution_node.metric is not None
-                    else "[red] No valid solution found.[/]"
-                )
-                end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
-                end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
-                end_optimization_layout["best_solution"].update(best_solution_panel)
-
-                # Save optimization results
-                # If the best solution does not exist or is has not been measured at the end of the optimization
-                # save the original solution as the best solution
-                if best_solution_node is not None:
-                    best_solution_code = best_solution_node.code
-                    best_solution_score = best_solution_node.metric
-                else:
-                    best_solution_code = None
-                    best_solution_score = None
-
-                if best_solution_code is None or best_solution_score is None:
-                    best_solution_content = f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_dir / f'step_0{source_fp.suffix}', is_json=False)}"
-                else:
-                    # Format score for the comment
-                    best_score_str = (
-                        format_number(best_solution_score)
-                        if best_solution_score is not None and isinstance(best_solution_score, (int, float))
-                        else "N/A"
+                    # Evaluate the final solution thats been generated
+                    eval_and_next_solution_response = evaluate_feedback_then_suggest_next_solution(
+                        run_id=run_id,
+                        execution_output=term_out,
+                        additional_instructions=current_additional_instructions,
+                        api_keys=llm_api_keys,
+                        timeout=timeout,
+                        auth_headers=auth_headers,
                     )
-                    best_solution_content = (
-                        f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
+                    summary_panel.set_step(step=steps)
+                    summary_panel.update_token_counts(usage=eval_and_next_solution_response["usage"])
+                    status_response = get_optimization_run_status(
+                        run_id=run_id, include_history=True, timeout=timeout, auth_headers=auth_headers
                     )
+                    # No need to update the plan panel since we have finished the optimization
+                    # Get the optimization run status for
+                    # the best solution, its score, and the history to plot the tree
+                    nodes_list_from_status_final = status_response.get("nodes")
+                    tree_panel.build_metric_tree(
+                        nodes=nodes_list_from_status_final if nodes_list_from_status_final is not None else []
+                    )
+                    # No need to set any solution to unevaluated since we have finished the optimization
+                    # and all solutions have been evaluated
+                    # No neeed to update the current solution panel since we have finished the optimization
+                    # We only need to update the best solution panel
+                    # Figure out if we have a best solution so far
+                    if status_response["best_result"] is not None:
+                        best_solution_node = Node(
+                            id=status_response["best_result"]["solution_id"],
+                            parent_id=status_response["best_result"]["parent_id"],
+                            code=status_response["best_result"]["code"],
+                            metric=status_response["best_result"]["metric_value"],
+                            is_buggy=status_response["best_result"]["is_buggy"],
+                        )
+                    else:
+                        best_solution_node = None
+                    solution_panels.update(current_node=None, best_node=best_solution_node)
+                    _, best_solution_panel = solution_panels.get_display(current_step=steps)
+                    # Update the end optimization layout
+                    final_message = (
+                        f"{summary_panel.metric_name.capitalize()} {'maximized' if summary_panel.maximize else 'minimized'}! Best solution {summary_panel.metric_name.lower()} = [green]{status_response['best_result']['metric_value']}[/] 🏆"
+                        if best_solution_node is not None and best_solution_node.metric is not None
+                        else "[red] No valid solution found.[/]"
+                    )
+                    end_optimization_layout["summary"].update(summary_panel.get_display(final_message=final_message))
+                    end_optimization_layout["tree"].update(tree_panel.get_display(is_done=True))
+                    end_optimization_layout["best_solution"].update(best_solution_panel)
+
+                    # Save optimization results
+                    # If the best solution does not exist or is has not been measured at the end of the optimization
+                    # save the original solution as the best solution
+                    if best_solution_node is not None:
+                        best_solution_code = best_solution_node.code
+                        best_solution_score = best_solution_node.metric
+                    else:
+                        best_solution_code = None
+                        best_solution_score = None
 
-                # Save best solution to .runs/<session-id>/best.<extension>
-                write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
-
-                # write the best solution to the source file
-                write_to_path(fp=source_fp, content=best_solution_content)
-
-                # Mark as completed normally for the finally block
-                optimization_completed_normally = True
-
-            console.print(end_optimization_layout)
+                    if best_solution_code is None or best_solution_score is None:
+                        best_solution_content = f"# Weco could not find a better solution\n\n{read_from_path(fp=runs_dir / f'step_0{source_fp.suffix}', is_json=False)}"
+                    else:
+                        # Format score for the comment
+                        best_score_str = (
+                            format_number(best_solution_score)
+                            if best_solution_score is not None and isinstance(best_solution_score, (int, float))
+                            else "N/A"
+                        )
+                        best_solution_content = (
+                            f"# Best solution from Weco with a score of {best_score_str}\n\n{best_solution_code}"
+                        )
+                    # Save best solution to .runs/<run-id>/best.<extension>
+                    write_to_path(fp=runs_dir / f"best{source_fp.suffix}", content=best_solution_content)
+                    # write the best solution to the source file
+                    write_to_path(fp=source_fp, content=best_solution_content)
+                    # Mark as completed normally for the finally block
+                    optimization_completed_normally = True
+                    console.print(end_optimization_layout)
 
         except Exception as e:
             # Catch errors during the main optimization loop or setup
             try:
-                error_message = e.response.json()["detail"]  # Try to get API error detail
+                error_message = e.response.json()["detail"]
             except Exception:
-                error_message = str(e)  # Otherwise, use the exception string
+                error_message = str(e)
             console.print(Panel(f"[bold red]Error: {error_message}", title="[bold red]Optimization Error", border_style="red"))
-            # Print traceback for debugging if needed (can be noisy)
-            # console.print_exception(show_locals=False)
-
             # Ensure optimization_completed_normally is False
             optimization_completed_normally = False
-
-            # Prepare details for termination report
             error_details = traceback.format_exc()
-
-            # Exit code will be handled by finally block or sys.exit below
-            exit_code = 1  # Indicate error
-            # No sys.exit here, let finally block run
-
+            exit_code = 1
         finally:
             # This block runs whether the try block completed normally or raised an exception
-
             # Stop heartbeat thread
             stop_heartbeat_event.set()
             if heartbeat_thread and heartbeat_thread.is_alive():
-                heartbeat_thread.join(timeout=2)  # Give it a moment to stop
+                heartbeat_thread.join(timeout=2)
 
-            # Report final status if a session was started
-            if session_id:
-                final_status = "unknown"
-                final_reason = "unknown_termination"
+            # Report final status if a run was started
+            if run_id:
+                final_status_update = "unknown"
+                final_reason_code = "unknown_termination"
                 final_details = None
-
                 if optimization_completed_normally:
-                    final_status = "completed"
-                    final_reason = "completed_successfully"
+                    final_status_update = "completed"
+                    final_reason_code = "completed_successfully"
+                elif user_stop_requested_flag:
+                    final_status_update = "terminated"
+                    final_reason_code = "user_requested_stop"
+                    final_details = "Run stopped by user request via dashboard."
                 else:
-                    # If an exception was caught and we have details
+                    final_status_update = "error"
+                    final_reason_code = "error_cli_internal"
                     if "error_details" in locals():
-                        final_status = "error"
-                        final_reason = "error_cli_internal"
-                        final_details = error_details
-                    # else: # Should have been handled by signal handler if terminated by user
-                    # Keep default 'unknown' if we somehow end up here without error/completion/signal
-
+                        final_details = locals()["error_details"]
+                    elif "e" in locals() and isinstance(locals()["e"], Exception):
+                        final_details = traceback.format_exc()
+                    else:
+                        final_details = "CLI terminated unexpectedly without a specific exception captured."
+                # Keep default 'unknown' if we somehow end up here without error/completion/signal
                 # Avoid reporting if terminated by signal handler (already reported)
                 # Check a flag or rely on status not being 'unknown'
-                if final_status != "unknown":
+                if final_status_update != "unknown":
                     report_termination(
-                        session_id=session_id,
-                        status_update=final_status,
-                        reason=final_reason,
+                        run_id=run_id,
+                        status_update=final_status_update,
+                        reason=final_reason_code,
                         details=final_details,
-                        auth_headers=auth_headers,
+                        auth_headers=current_auth_headers_for_heartbeat,
                     )
-
-            # Ensure proper exit code if an error occurred
-            if not optimization_completed_normally and "exit_code" in locals() and exit_code != 0:
-                sys.exit(exit_code)
-            elif not optimization_completed_normally:
-                # Generic error exit if no specific code was set but try block failed
-                sys.exit(1)
-            else:
-                # Normal exit
+            if optimization_completed_normally:
                 sys.exit(0)
+            elif user_stop_requested_flag:
+                console.print("[yellow]Run terminated by user request.[/]")
+                sys.exit(0)
+            else:
+                sys.exit(locals().get("exit_code", 1))

weco/panels.py

@@ -11,9 +11,9 @@ from .__init__ import __dashboard_url__
 
 
 class SummaryPanel:
-    """Holds a summary of the optimization session."""
+    """Holds a summary of the optimization run."""
 
-    def __init__(self, maximize: bool, metric_name: str, total_steps: int, model: str, runs_dir: str, session_id: str = None):
+    def __init__(self, maximize: bool, metric_name: str, total_steps: int, model: str, runs_dir: str, run_id: str = None):
         self.maximize = maximize
         self.metric_name = metric_name
         self.goal = ("Maximizing" if self.maximize else "Minimizing") + f" {self.metric_name}..."
@@ -22,7 +22,7 @@ class SummaryPanel:
         self.total_steps = total_steps
         self.model = model
         self.runs_dir = runs_dir
-        self.session_id = session_id if session_id is not None else "N/A"
+        self.run_id = run_id if run_id is not None else "N/A"
         self.dashboard_url = "N/A"
         self.progress = Progress(
             TextColumn("[progress.description]{task.description}"),
@@ -34,14 +34,14 @@ class SummaryPanel:
         )
         self.task_id = self.progress.add_task("", total=total_steps)
 
-    def set_session_id(self, session_id: str):
-        """Set the session ID."""
-        self.session_id = session_id
-        self.set_dashboard_url(session_id=session_id)
+    def set_run_id(self, run_id: str):
+        """Set the run ID."""
+        self.run_id = run_id
+        self.set_dashboard_url(run_id=run_id)
 
-    def set_dashboard_url(self, session_id: str):
+    def set_dashboard_url(self, run_id: str):
         """Set the dashboard URL."""
-        self.dashboard_url = f"{__dashboard_url__}/runs/{session_id}"
+        self.dashboard_url = f"{__dashboard_url__}/runs/{run_id}"
 
     def set_step(self, step: int):
         """Set the current step."""
@@ -70,7 +70,7 @@ class SummaryPanel:
         summary_table.add_row(f"[bold cyan]Model:[/] {self.model}")
         summary_table.add_row("")
         # Log directory
-        summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{self.runs_dir}/{self.session_id}[/]")
+        summary_table.add_row(f"[bold cyan]Logs:[/] [blue underline]{self.runs_dir}/{self.run_id}[/]")
         summary_table.add_row("")
         # Dashboard link
         summary_table.add_row(f"[bold cyan]Dashboard:[/] [blue underline]{self.dashboard_url}[/]")
@@ -175,6 +175,9 @@ class MetricTreePanel:
 
     def build_metric_tree(self, nodes: List[dict]):
         """Build the tree from the list of nodes."""
+        # Defensive: treat None as empty list
+        if nodes is None:
+            nodes = []
         # First clear then tree
         self.metric_tree.clear()
 

{weco-0.2.19.dist-info → weco-0.2.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: weco
-Version: 0.2.19
+Version: 0.2.20
 Summary: Documentation for `weco`, a CLI for using Weco AI's code optimizer.
 Author-email: Weco AI Team <contact@weco.ai>
 License: MIT
@@ -29,6 +29,9 @@ Dynamic: license-file
 [![docs](https://img.shields.io/website?url=https://docs.weco.ai/&label=docs)](https://docs.weco.ai/)
 [![PyPI version](https://badge.fury.io/py/weco.svg)](https://badge.fury.io/py/weco)
 [![AIDE](https://img.shields.io/badge/AI--Driven_Exploration-arXiv-orange?style=flat-square&logo=arxiv)](https://arxiv.org/abs/2502.13138)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WecoAI/weco-cli/blob/main/examples/hello-kernel-world/colab_notebook_walkthrough.ipynb)
+
+`pip install weco`
 
 </div>
 

weco-0.2.20.dist-info/RECORD

@@ -0,0 +1,12 @@
+weco/__init__.py,sha256=npWmRgLxfVK69GdyxIujnI87xqmPCBrZWxxAxL_QQOc,478
+weco/api.py,sha256=xHCyQPto1Lv9QysiOFwVf5NnWDh6LBCNfPLyq-L7nys,5873
+weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
+weco/cli.py,sha256=e4h5bxeg2n95AlYXanfxLbcURWchjWTES2Kwx5AjKn0,36115
+weco/panels.py,sha256=lsTHTh-XdYMH3ZV_WBteEcIt2hTWGGtqfUjGlYRHl70,13598
+weco/utils.py,sha256=LVTBo3dduJmhlbotcYoUW2nLx6IRtKs4eDFR52Qltcg,5244
+weco-0.2.20.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
+weco-0.2.20.dist-info/METADATA,sha256=rK-Y9Q0zwaKUBS0bNZfUNvL82RXUiijVcIbW3i_IKKk,10955
+weco-0.2.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+weco-0.2.20.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
+weco-0.2.20.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
+weco-0.2.20.dist-info/RECORD,,

{weco-0.2.19.dist-info → weco-0.2.20.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.7.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

weco-0.2.19.dist-info/RECORD

@@ -1,12 +0,0 @@
-weco/__init__.py,sha256=npWmRgLxfVK69GdyxIujnI87xqmPCBrZWxxAxL_QQOc,478
-weco/api.py,sha256=lJJ0j0-bABiQXDlRb43fCo7ky0N_HwfZgFdMktRKQ90,6635
-weco/auth.py,sha256=IPfiLthcNRkPyM8pWHTyDLvikw83sigacpY1PmeA03Y,2343
-weco/cli.py,sha256=eI468fxpMTfGPL-aX6EMYxh0NuaRxpaLVF_Jj2DiFhU,36383
-weco/panels.py,sha256=pM_YGnmcXM_1CBcxo_EAzOV3g_4NFdLS4MqDqx7THbA,13563
-weco/utils.py,sha256=LVTBo3dduJmhlbotcYoUW2nLx6IRtKs4eDFR52Qltcg,5244
-weco-0.2.19.dist-info/licenses/LICENSE,sha256=p_GQqJBvuZgkLNboYKyH-5dhpTDlKs2wq2TVM55WrWE,1065
-weco-0.2.19.dist-info/METADATA,sha256=3VBVsCqr7p332A10KsLr168GvOIKcCOWWfGDv8ViF7I,10729
-weco-0.2.19.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-weco-0.2.19.dist-info/entry_points.txt,sha256=ixJ2uClALbCpBvnIR6BXMNck8SHAab8eVkM9pIUowcs,39
-weco-0.2.19.dist-info/top_level.txt,sha256=F0N7v6e2zBSlsorFv-arAq2yDxQbzX3KVO8GxYhPUeE,5
-weco-0.2.19.dist-info/RECORD,,