hte-cli 0.1.23__py3-none-any.whl → 0.2.8__py3-none-any.whl

hte_cli/api_client.py

@@ -201,6 +201,31 @@ class APIClient:
             },
         )
 
+    # =========================================================================
+    # Session-based API (new flow: hte-cli session join <session_id>)
+    # =========================================================================
+
+    def join_session(self, session_id: str) -> dict:
+        """Join an existing session created by web UI.
+
+        Returns session info including task data, benchmark, mode, etc.
+        Sets cli_connected_at on the server.
+        """
+        return self.post(f"/sessions/{session_id}/join")
+
+    def get_session_files(self, session_id: str) -> bytes:
+        """Download task files for a session as zip."""
+        return self.get_raw(f"/sessions/{session_id}/files")
+
+    def get_session_compose(self, session_id: str) -> str:
+        """Get compose.yaml content for a session."""
+        content = self.get_raw(f"/sessions/{session_id}/compose")
+        return content.decode("utf-8")
+
+    # =========================================================================
+    # Result Upload
+    # =========================================================================
+
     def upload_result(
         self,
         session_id: str,

hte_cli/cli.py

@@ -13,9 +13,7 @@ import click
 from rich.console import Console
 from rich.table import Table
 from rich.panel import Panel
-from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
-from rich.live import Live
-from rich.text import Text
+from rich.progress import Progress, SpinnerColumn, TextColumn
 
 from hte_cli import __version__, API_BASE_URL
 from hte_cli.config import Config, get_eval_logs_dir
@@ -149,13 +147,319 @@ def auth_status(ctx):
 
 
 # =============================================================================
-# Tasks Commands
+# Session Commands (New flow: session join <session_id>)
+# =============================================================================
+
+
+@cli.group()
+def session():
+    """Session management commands."""
+    pass
+
+
+@session.command("join")
+@click.argument("session_id")
+@click.option("--force-setup", is_flag=True, help="Re-run setup even if reconnecting")
+@click.pass_context
+def session_join(ctx, session_id: str, force_setup: bool):
+    """Join an existing session by ID.
+
+    This is the primary way to start working on a task:
+    1. Start the task from the web UI (creates session)
+    2. Run this command with the session ID shown in the web UI
+    3. The environment will be set up and the timer will start
+    """
+    config: Config = ctx.obj["config"]
+
+    if not config.is_authenticated():
+        console.print("[red]Not logged in. Run: hte-cli auth login[/red]")
+        sys.exit(1)
+
+    api = APIClient(config)
+
+    # Step 1: Join session
+    console.print()
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        console=console,
+    ) as progress:
+        progress.add_task("Joining session...", total=None)
+
+        try:
+            session_info = api.join_session(session_id)
+        except APIError as e:
+            if "Invalid session ID format" in str(e):
+                console.print(f"[red]{e}[/red]")
+            elif e.status_code == 404:
+                console.print("[red]Session not found. Check the session ID and try again.[/red]")
+            elif e.status_code == 400 and "paused" in str(e).lower():
+                console.print("[yellow]Session is paused. Please resume from the web UI first.[/yellow]")
+            else:
+                console.print(f"[red]Error: {e}[/red]")
+            sys.exit(1)
+
+    # Check if reconnecting (session already in_progress)
+    is_reconnect = session_info.get("status") == "in_progress"
+
+    if is_reconnect and not force_setup:
+        console.print("[yellow]Reconnecting to existing session...[/yellow]")
+        console.print()
+
+    console.print(
+        Panel(
+            f"[bold]Task:[/bold] {session_info['task_id']}\n"
+            f"[bold]Benchmark:[/bold] {session_info['benchmark']}\n"
+            f"[bold]Mode:[/bold] {session_info['mode']}",
+            title="Session Joined",
+        )
+    )
+    console.print()
+
+    # Import runner components
+    import time
+    from hte_cli.events import EventStreamer
+    from hte_cli.runner import TaskRunner
+    from hte_cli.image_utils import extract_images_from_compose, pull_image_with_progress
+
+    # Create event streamer
+    events = EventStreamer(api, session_id)
+
+    # Step 2: Download task files and compose (skip if reconnecting without force)
+    files_zip = None
+    compose_yaml = None
+
+    if not is_reconnect or force_setup:
+        console.print("[bold]Step 1:[/bold] Downloading task resources...")
+
+        with console.status("[dim]Fetching task files...[/dim]"):
+            try:
+                files_zip = api.get_session_files(session_id)
+                console.print("  [green]✓[/green] Task files downloaded")
+            except APIError as e:
+                console.print(f"  [dim]○ No task files (optional)[/dim]")
+                files_zip = None
+
+        with console.status("[dim]Fetching compose configuration...[/dim]"):
+            try:
+                compose_yaml = api.get_session_compose(session_id)
+                console.print("  [green]✓[/green] Docker compose downloaded")
+            except APIError as e:
+                console.print(f"  [dim]○ No compose file (optional)[/dim]")
+                compose_yaml = None
+
+        console.print()
+
+    # Validate compose for sandbox-required benchmarks
+    benchmark = session_info.get("benchmark", "").lower()
+    SANDBOX_REQUIRED_BENCHMARKS = {"cybench", "intercode-ctf", "cvebench", "cybergym", "nyuctf"}
+    if benchmark in SANDBOX_REQUIRED_BENCHMARKS and not compose_yaml and not is_reconnect:
+        console.print(
+            f"[red]Error: {benchmark} requires a Docker sandbox but no compose file was found.[/red]"
+        )
+        console.print()
+        console.print(
+            f"Please contact support: {SUPPORT_EMAIL}"
+        )
+        sys.exit(1)
+
+    # Build assignment dict for runner compatibility
+    assignment = {
+        "assignment_id": session_info.get("assignment_id"),
+        "session_id": session_id,
+        "task_id": session_info["task_id"],
+        "benchmark": session_info["benchmark"],
+        "mode": session_info["mode"],
+        "time_cap_seconds": session_info.get("time_cap_seconds"),
+        "task": {
+            "instructions": session_info.get("instructions", ""),
+            "metadata": session_info.get("metadata", {}),
+            "scorer_type": session_info.get("scorer_type"),
+            "intermediate_scoring": session_info.get("intermediate_scoring", False),
+        },
+    }
+
+    # Step 3: Run setup (skip if reconnecting without force)
+    setup_start_time = time.monotonic()
+    images = []
+    pulled_images = []
+    cached_images = []
+    failed_images = []
+
+    if not is_reconnect or force_setup:
+        # Extract images from compose
+        if compose_yaml:
+            images = extract_images_from_compose(compose_yaml)
+
+        # Send setup_started event
+        events.setup_started(images=images)
+
+        # Pull images if we have any
+        if images:
+            from hte_cli.image_utils import check_image_exists_locally
+
+            console.print(f"[bold]Step 2:[/bold] Pulling {len(images)} Docker image(s)...")
+            pull_start = time.monotonic()
+
+            for img in images:
+                short_name = img.split("/")[-1][:40]
+
+                # Check if already cached
+                if check_image_exists_locally(img):
+                    console.print(f"  [green]✓[/green] {short_name} [dim](cached)[/dim]")
+                    cached_images.append(img)
+                    continue
+
+                # Need to pull - show progress
+                last_status = ["connecting..."]  # Use list for closure mutability
+                with console.status(f"[yellow]↓[/yellow] {short_name} [dim]connecting...[/dim]") as status:
+                    def show_progress(image: str, line: str):
+                        # Parse docker pull output for layer progress
+                        # Lines look like: "abc123: Downloading [====>   ] 10MB/50MB"
+                        # Or: "abc123: Extracting [====>   ] 10MB/50MB"
+                        # Or: "abc123: Pull complete", "Digest: sha256:...", "Status: ..."
+                        display_text = None
+                        if ": " in line:
+                            parts = line.split(": ", 1)
+                            if len(parts) == 2:
+                                layer_id = parts[0][-8:]  # Last 8 chars of layer ID
+                                layer_status = parts[1]  # Don't truncate - keep full progress
+                                # Show progress bars for Downloading/Extracting with MB info
+                                if "Downloading" in layer_status or "Extracting" in layer_status:
+                                    # Keep progress: "[====>   ] 10.5MB/50MB"
+                                    display_text = f"{layer_id}: {layer_status[:50]}"
+                                elif "Pull complete" in layer_status:
+                                    display_text = f"{layer_id}: done"
+                                elif "Download complete" in layer_status:
+                                    display_text = f"{layer_id}: download done"
+                                elif "Already exists" in layer_status:
+                                    display_text = f"{layer_id}: cached"
+                                elif "Waiting" in layer_status:
+                                    display_text = f"{layer_id}: waiting"
+                                elif "Verifying" in layer_status:
+                                    display_text = f"{layer_id}: verifying"
+                                else:
+                                    display_text = line[:55]
+                        elif line.strip():
+                            display_text = line[:55]
+
+                        if display_text and display_text != last_status[0]:
+                            last_status[0] = display_text
+                            status.update(f"[yellow]↓[/yellow] {short_name} [dim]{display_text}[/dim]")
+
+                    success = pull_image_with_progress(img, on_progress=show_progress)
+
+                if success:
+                    console.print(f"  [green]✓[/green] {short_name} [dim](downloaded)[/dim]")
+                    pulled_images.append(img)
+                else:
+                    console.print(f"  [red]✗[/red] {short_name} [dim](failed)[/dim]")
+                    failed_images.append(img)
+
+            pull_duration = time.monotonic() - pull_start
+            events.image_pull_completed(
+                duration_seconds=pull_duration,
+                pulled=pulled_images,
+                cached=cached_images,
+                failed=failed_images,
+            )
+            console.print()
+
+        # Send setup_completed - THIS STARTS THE TIMER ON SERVER
+        total_setup = time.monotonic() - setup_start_time
+        events.setup_completed(total_seconds=total_setup)
+        console.print("[green]Environment ready! Timer started.[/green]")
+        console.print()
+    else:
+        # Reconnecting - compose should already be running
+        console.print("[dim]Skipping setup (use --force-setup to re-run)[/dim]")
+        console.print()
+
+    # Step 4: Show instructions
+    if session_info.get("instructions"):
+        console.print(Panel(session_info["instructions"], title="Task Instructions"))
+        console.print()
+
+    # Step 3: Run the task using TaskRunner
+    step_num = "3" if (not is_reconnect or force_setup) and images else "2" if (not is_reconnect or force_setup) else "1"
+    console.print(f"[bold]Step {step_num}:[/bold] Starting task environment...")
+
+    events.docker_started()
+
+    runner = TaskRunner()
+    eval_log_bytes = None
+    try:
+        with console.status("[dim]Launching Docker containers (this may take a minute)...[/dim]"):
+            result = runner.run_from_assignment(
+                assignment=assignment,
+                compose_yaml=compose_yaml,
+                files_zip=files_zip,
+            )
+        # Read eval log before cleanup
+        if result.eval_log_path and result.eval_log_path.exists():
+            eval_log_bytes = result.eval_log_path.read_bytes()
+    except KeyboardInterrupt:
+        events.docker_stopped(exit_code=130)
+        console.print()
+        console.print("[yellow]Interrupted. Session remains active - you can reconnect later.[/yellow]")
+        sys.exit(0)
+    except Exception as e:
+        events.docker_stopped(exit_code=1)
+        console.print(f"[red]Task execution failed: {e}[/red]")
+        sys.exit(1)
+    finally:
+        runner.cleanup()
+
+    events.docker_stopped(exit_code=0)
+
+    # Step 6: Upload result
+    if result and result.answer:
+        events.session_completed(
+            elapsed_seconds=result.time_seconds,
+            answer=result.answer,
+        )
+
+        console.print()
+        console.print("[green]Task completed![/green]")
+        console.print(f"Answer: {result.answer}")
+        console.print(f"Time: {result.time_seconds:.1f}s")
+
+        # Upload to server
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            console=console,
+        ) as progress:
+            progress.add_task("Uploading result...", total=None)
+            try:
+                upload_result = api.upload_result(
+                    session_id=session_id,
+                    answer=result.answer or "",
+                    client_active_seconds=result.time_seconds,
+                    eval_log_bytes=eval_log_bytes,
+                    score=result.score,
+                    score_binarized=result.score_binarized,
+                    agent_id=result.agent_id,
+                )
+            except APIError as e:
+                console.print(f"[red]Failed to upload result: {e}[/red]")
+                sys.exit(1)
+
+        if upload_result.get("score") is not None:
+            console.print(f"Score: {upload_result['score']}")
+
+        console.print()
+        console.print("[green]Done! Return to the web UI to see your results.[/green]")
+
+
+# =============================================================================
+# Tasks Commands (DEPRECATED - use 'session join' instead)
 # =============================================================================
 
 
 @cli.group()
 def tasks():
-    """Task commands."""
+    """Task commands (deprecated - use 'session join' instead)."""
     pass
 
 
@@ -225,7 +529,23 @@ def tasks_list(ctx):
 @click.argument("task_id", required=False)
 @click.pass_context
 def tasks_run(ctx, task_id: str | None):
-    """Run a task (default: highest priority pending task)."""
+    """[DEPRECATED] Run a task - use 'session join' instead."""
+    console.print()
+    console.print("[red]This command is deprecated.[/red]")
+    console.print()
+    console.print("The new workflow is:")
+    console.print("  1. Start the task from the web UI: https://cyber-task-horizons.com")
+    console.print("  2. Run the command shown: [bold]hte-cli session join <session_id>[/bold]")
+    console.print()
+    console.print("This ensures accurate timing by starting the timer only when")
+    console.print("the environment is ready, not including Docker setup time.")
+    console.print()
+    sys.exit(1)
+
+
+# Keep the old implementation as _tasks_run_legacy for testing if needed
+def _tasks_run_legacy(ctx, task_id: str | None):
+    """Legacy implementation of tasks run (for testing only)."""
     config: Config = ctx.obj["config"]
 
     if not config.is_authenticated():
@@ -397,12 +717,18 @@ def tasks_run(ctx, task_id: str | None):
         return
 
     # Step 5: Pre-pull Docker images with progress
-    from hte_cli.runner import prepull_compose_images, extract_images_from_compose
+    from hte_cli.image_utils import extract_images_from_compose
     import re
+    import time
+
+    setup_start_time = time.monotonic()
+    images: list[str] = []
+    results: list[tuple[str, bool, str]] = []
 
     if compose_yaml:
         images = extract_images_from_compose(compose_yaml)
         if images:
+            events.setup_started(images)
             console.print()
             console.print(f"[bold]Preparing Docker environment ({len(images)} images)...[/bold]")
 
@@ -448,34 +774,48 @@ def tasks_run(ctx, task_id: str | None):
             def get_progress_summary(image: str) -> str:
                 """Get a human-readable progress summary for an image with MB counts."""
                 if image not in image_layers or not image_layers[image]:
-                    return "Starting..."
+                    return "connecting..."
 
                 layers = image_layers[image]
                 total_layers = len(layers)
-                complete = sum(1 for s, _, _ in layers.values() if "complete" in s.lower())
 
-                # Sum up download progress
+                # Count layers in different states
+                complete = 0
+                downloading = 0
+                waiting = 0
                 total_downloaded_mb = 0
                 total_size_mb = 0
+
                 for status, downloaded, total in layers.values():
-                    if "complete" in status.lower():
-                        # Completed layers: use total as both downloaded and total
+                    status_lower = status.lower()
+                    if "complete" in status_lower:
+                        complete += 1
                         total_downloaded_mb += total
                         total_size_mb += total
-                    elif total > 0:
+                    elif "downloading" in status_lower:
+                        downloading += 1
                         total_downloaded_mb += downloaded
                         total_size_mb += total
+                    elif "waiting" in status_lower:
+                        waiting += 1
 
+                # Choose the most informative display
                 if complete == total_layers and total_layers > 0:
                     if total_size_mb > 0:
-                        return f"Done ({total_size_mb:.0f}MB)"
-                    return f"Done ({total_layers} layers)"
+                        return f"done ({total_size_mb:.0f}MB)"
+                    return f"done ({total_layers} layers)"
                 elif total_size_mb > 0:
-                    return f"{total_downloaded_mb:.0f}MB / {total_size_mb:.0f}MB"
+                    # Show MB progress when available
+                    pct = int(100 * total_downloaded_mb / total_size_mb) if total_size_mb > 0 else 0
+                    return f"{total_downloaded_mb:.0f}/{total_size_mb:.0f}MB ({pct}%)"
+                elif downloading > 0:
+                    return f"downloading ({complete}/{total_layers} done)"
                 elif complete > 0:
-                    return f"Pulling ({complete}/{total_layers} layers)"
+                    return f"extracting ({complete}/{total_layers} done)"
+                elif waiting > 0:
+                    return f"queued ({total_layers} layers)"
                 else:
-                    return f"Preparing ({total_layers} layers)"
+                    return f"preparing ({total_layers} layers)"
 
             def on_image_progress(image: str, line: str):
                 """Track layer-level progress with size info."""
@@ -492,7 +832,7 @@ def tasks_run(ctx, task_id: str | None):
                 short_name = img.split("/")[-1] if "/" in img else img
 
                 # Check if cached first
-                from hte_cli.runner import check_image_exists_locally, pull_image_with_progress
+                from hte_cli.image_utils import check_image_exists_locally, pull_image_with_progress
 
                 if check_image_exists_locally(img):
                     console.print(f"  [green]✓[/green] {short_name} [dim](cached)[/dim]")
@@ -502,7 +842,10 @@ def tasks_run(ctx, task_id: str | None):
                 # Need to pull - use Rich Status for live updates
                 image_layers[img] = {}
 
-                with console.status(f"[yellow]↓[/yellow] {short_name} [dim]connecting...[/dim]") as status:
+                with console.status(
+                    f"[yellow]↓[/yellow] {short_name} [dim]connecting...[/dim]"
+                ) as status:
+
                     def show_progress(image: str, line: str):
                         on_image_progress(image, line)
                         summary = get_progress_summary(image)
@@ -518,14 +861,27 @@ def tasks_run(ctx, task_id: str | None):
                     console.print(f"  [red]✗[/red] {short_name} [dim](failed)[/dim]")
                     results.append((img, False, "failed"))
 
-            failed = sum(1 for _, ok, _ in results if not ok)
-            if failed > 0:
+            failed_count = sum(1 for _, ok, _ in results if not ok)
+            if failed_count > 0:
                 console.print(
-                    f"[yellow]Warning: {failed} image(s) failed to pull. "
+                    f"[yellow]Warning: {failed_count} image(s) failed to pull. "
                     "Task may fail to start.[/yellow]"
                 )
             console.print()
 
+    # Record image pull timing
+    if images:
+        pull_duration = time.monotonic() - setup_start_time
+        pulled = [img for img, ok, status in results if ok and status == "pulled"]
+        cached = [img for img, ok, status in results if ok and status == "cached"]
+        failed = [img for img, ok, status in results if not ok]
+        events.image_pull_completed(
+            duration_seconds=pull_duration,
+            pulled=pulled,
+            cached=cached,
+            failed=failed,
+        )
+
     # Step 6: Run Inspect's human_cli
     runner = TaskRunner()
     console.print("[bold]Starting task environment...[/bold]")
@@ -534,6 +890,10 @@ def tasks_run(ctx, task_id: str | None):
 
     events.docker_started()
 
+    # Record total setup time (image pulls + compose up)
+    total_setup = time.monotonic() - setup_start_time
+    events.setup_completed(total_seconds=total_setup)
+
     eval_log_bytes = None
     local_eval_path = None
     try:
@@ -597,13 +957,14 @@ def tasks_run(ctx, task_id: str | None):
         except Exception:
             pass  # Not a CyberGym task or malformed zip
 
-    # Show upload size info
-    upload_size_kb = 0
-    if eval_log_bytes:
-        upload_size_kb = len(eval_log_bytes) / 1024
-        size_mb = upload_size_kb / 1024
-        if size_mb > 50:
-            console.print(f"[yellow]Warning: Large eval log ({size_mb:.1f} MB)[/yellow]")
+    # Show upload size info and track timing
+    upload_size_bytes = len(eval_log_bytes) if eval_log_bytes else 0
+    upload_size_kb = upload_size_bytes / 1024
+    if upload_size_kb / 1024 > 50:
+        console.print(f"[yellow]Warning: Large eval log ({upload_size_kb / 1024:.1f} MB)[/yellow]")
+
+    events.upload_started(size_bytes=upload_size_bytes)
+    upload_start_time = time.monotonic()
 
     with Progress(
         SpinnerColumn(),
@@ -630,6 +991,10 @@ def tasks_run(ctx, task_id: str | None):
             console.print("[yellow]Your result was saved locally but not uploaded.[/yellow]")
             sys.exit(1)
 
+    # Record upload completion
+    upload_duration = time.monotonic() - upload_start_time
+    events.upload_completed(duration_seconds=upload_duration, size_bytes=upload_size_bytes)
+
     console.print()
     console.print("[green]Result uploaded successfully![/green]")
 

hte_cli/events.py

@@ -24,6 +24,12 @@ class EventStreamer:
         "docker_started",
         "docker_stopped",
         "session_completed",
+        # Overhead tracking events
+        "setup_started",
+        "image_pull_completed",
+        "setup_completed",
+        "upload_started",
+        "upload_completed",
     }
 
     def __init__(self, api: APIClient, session_id: str):
@@ -126,3 +132,45 @@ class EventStreamer:
         if answer is not None:
             data["answer_submitted"] = True
         return self.send("session_completed", data or None)
+
+    # Overhead tracking events
+
+    def setup_started(self, images: list[str]) -> bool:
+        """Record start of setup phase (before image pulls)."""
+        return self.send("setup_started", {"images": images})
+
+    def image_pull_completed(
+        self,
+        duration_seconds: float,
+        pulled: list[str],
+        cached: list[str],
+        failed: list[str],
+    ) -> bool:
+        """Record image pull results with timing."""
+        return self.send(
+            "image_pull_completed",
+            {
+                "duration_seconds": duration_seconds,
+                "pulled": pulled,
+                "cached": cached,
+                "failed": failed,
+            },
+        )
+
+    def setup_completed(self, total_seconds: float) -> bool:
+        """Record end of setup phase (environment ready for work)."""
+        return self.send("setup_completed", {"total_seconds": total_seconds})
+
+    def upload_started(self, size_bytes: int) -> bool:
+        """Record start of result upload."""
+        return self.send("upload_started", {"size_bytes": size_bytes})
+
+    def upload_completed(self, duration_seconds: float, size_bytes: int) -> bool:
+        """Record end of result upload with timing."""
+        return self.send(
+            "upload_completed",
+            {
+                "duration_seconds": duration_seconds,
+                "size_bytes": size_bytes,
+            },
+        )