cloudshell.py

"""
CloudShell Integration Module for AU Agent
==========================================

Internally integrates Google Cloud Shell access into the AU Agent Flask
backend. Replaces the standalone Gradio UI (`app_google_cloud.py`) with
programmatic access so that the download/trim step of the pipeline can
be delegated to Cloud Shell (where YouTube is reachable, where there is
ample free disk, and where `server.py` performs yt-dlp + ffmpeg + Kaggle
upload in one shot).

Flow (v1.4.0+ — interactive session, per user spec):
  1. `restore_credentials()` runs once at startup. It looks for
     `gcloud-backup.zip` in known locations, extracts it, restores the
     `~/.ssh` keys and `~/.config/gcloud` credentials.
  2. `download_and_trim_via_cloudshell()` opens an INTERACTIVE Cloud
     Shell session by running `gcloud cloud-shell ssh --authorize-session`
     (no `--command=...` flag). We wait for the shell to be ready, then
     SEND the `python3 server.py ...` command via the PTY stdin.
  3. Output (yt-dlp progress, ffmpeg stats, Kaggle upload messages,
     success/failure markers) is streamed line-by-line to the
     `log_callback`, which is wired to `db_add_log` so the web UI's
     "Live Output" panel updates in real time.
  4. Special detection:
       - "Try again 1 hour later" → Cloud Shell rate-limit hit. We emit
         a `[RATE_LIMIT] ...` marker to the log so the web UI can show
         a toast + persistent banner telling the user to wait 1 hour
         or reset the VM.
       - "SUCCESS: Task pipeline completed" → success marker.
  5. The Cloud Shell session is closed IMMEDIATELY when:
       - The python command finishes (success or failure)
       - The rate-limit message is detected
       - A fatal error occurs
       - The timeout is hit
     We close by sending `exit\n` to the PTY stdin and then closing
     the master fd, which tears down the SSH tunnel.

This module is import-safe: if gcloud is not installed or credentials
are missing, every call returns `(False, error_msg)` instead of raising,
so the pipeline can mark step 1 as failed and stop (which is the only
critical step).
"""

from __future__ import annotations

import os
import re
import shlex
import shutil
import select
import pty
import signal
import time
import zipfile
import logging
import subprocess
from pathlib import Path
from typing import Callable, Optional, Tuple

log = logging.getLogger("au-agent.cloudshell")

# ── Binary & path constants ──────────────────────────────────────────────────
GCLOUD_BIN = "/opt/google-cloud-sdk/bin/gcloud"

# Places we look for the user's gcloud-backup.zip. The user said the zip
# is "already uploaded in the HF space project dir" so we just probe a
# handful of well-known locations and pick the first one that looks real.
# Order matters: most-likely locations first.
#
# HF Spaces repo-root files end up at `/app/<filename>` inside the
# container because the Dockerfile has `COPY . /app/`. So
# `/app/gcloud-backup.zip` is the canonical location and is probed first.
GCLOUD_BACKUP_ZIP_CANDIDATES = [
    "/app/gcloud-backup.zip",
    "/app/data/gcloud-backup.zip",
    "/gcloud-backup.zip",
    os.path.join(os.getcwd(), "gcloud-backup.zip"),
    "/home/user/gcloud-backup.zip",
    "/tmp/gcloud-backup.zip",
]

GCLOUD_EXTRACT_DIR = "/tmp/gcloud-backup-extracted"
SSH_DIR = os.path.expanduser("~/.ssh")
GCLOUD_CONFIG_DIR = os.path.expanduser("~/.config/gcloud")

# ── Markers we look for in the streamed output ──────────────────────────────
# server.py on Cloud Shell prints one of these on success. We treat ANY of
# them as the authoritative success signal — gcloud's own returncode is
# unreliable for SSH tunnels (sometimes None when the tunnel tears down).
#
# NOTE: The exact marker string varies between server.py versions. The
# original spec used "SUCCESS: Task pipeline completed", but the user's
# actual server.py prints "🏁 SUCCESS: Task completed." (no "pipeline").
# We match a flexible set so both work.
SUCCESS_MARKERS = [
    "SUCCESS: Task pipeline completed",
    "SUCCESS: Task completed",
    "🏁 SUCCESS",
    "🏁 SUCCESS: Task",
]

# Cloud Shell prints one of these when the user has hit the per-session
# or per-hour rate limit. When we see any of them, we abort immediately
# and tell the web UI to show a "try again in 1 hour / reset VM" banner.
RATE_LIMIT_MARKERS = [
    "Try again 1 hour later",
    "Try again in 1 hour",
    "Rate limit exceeded",
    "Quota exceeded",
    "Too many requests",
]

# Build a clean env for every gcloud invocation. gcloud SDK needs HOME,
# USER, LOGNAME, TERM and PATH pointing at its own bin directory.
def _build_env() -> dict:
    return {
        **os.environ,
        "HOME": os.environ.get("HOME", "/home/user"),
        "USER": os.environ.get("USER", "user"),
        "LOGNAME": os.environ.get("LOGNAME", "user"),
        "TERM": "xterm-256color",
        "PATH": "/opt/google-cloud-sdk/bin:"
                + os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin"),
    }


# Module-level state: True once credentials have been restored successfully.
CREDENTIALS_READY: bool = False

# v2.14.0: Track the currently running SSH subprocess so it can be killed
# when the user clicks "Stop Task". Set by _run_interactive_session() and
# cleared when the session ends.
_running_ssh_proc = None
_running_ssh_pid = None


# ── Credential restoration ───────────────────────────────────────────────────

def find_backup_zip() -> Optional[str]:
    """Return the first existing gcloud-backup.zip path, or None."""
    for path in GCLOUD_BACKUP_ZIP_CANDIDATES:
        try:
            if os.path.isfile(path) and os.path.getsize(path) > 100:
                return path
        except OSError:
            continue
    return None


def restore_credentials(force: bool = False) -> Tuple[bool, str]:
    """
    Extract gcloud-backup.zip and restore SSH keys + gcloud credentials.

    Idempotent: if `CREDENTIALS_READY` is already True and `force` is
    False, returns immediately.

    Returns (success, message).
    """
    global CREDENTIALS_READY

    if CREDENTIALS_READY and not force:
        return True, "Credentials already restored"

    log.info("[CloudShell] Restoring credentials...")

    zip_path = find_backup_zip()
    if not zip_path:
        msg = ("gcloud-backup.zip not found. Looked in: "
               + ", ".join(GCLOUD_BACKUP_ZIP_CANDIDATES))
        log.error("[CloudShell] %s", msg)
        return False, msg

    try:
        if os.path.exists(GCLOUD_EXTRACT_DIR):
            shutil.rmtree(GCLOUD_EXTRACT_DIR, ignore_errors=True)
        os.makedirs(GCLOUD_EXTRACT_DIR, exist_ok=True)

        with zipfile.ZipFile(zip_path, "r") as z:
            z.extractall(GCLOUD_EXTRACT_DIR)
        log.info("[CloudShell] Extracted %s", zip_path)

        backup_root = GCLOUD_EXTRACT_DIR
        if os.path.isdir(os.path.join(GCLOUD_EXTRACT_DIR, "gcloud-backup")):
            backup_root = os.path.join(GCLOUD_EXTRACT_DIR, "gcloud-backup")

        ssh_src = os.path.join(backup_root, ".ssh")
        if os.path.isdir(ssh_src):
            os.makedirs(SSH_DIR, exist_ok=True)
            shutil.copytree(ssh_src, SSH_DIR, dirs_exist_ok=True)
            os.system('chmod 700 "%s" && chmod 600 "%s"/* 2>/dev/null' % (SSH_DIR, SSH_DIR))
            log.info("[CloudShell] SSH keys restored to %s", SSH_DIR)
        else:
            log.warning("[CloudShell] No .ssh folder inside backup; "
                        "cloud-shell ssh may prompt for keygen on first run")

        gcloud_src = os.path.join(backup_root, "gcloud")
        if os.path.isdir(gcloud_src):
            os.makedirs(os.path.dirname(GCLOUD_CONFIG_DIR), exist_ok=True)
            shutil.copytree(gcloud_src, GCLOUD_CONFIG_DIR, dirs_exist_ok=True)
            log.info("[CloudShell] gcloud config restored to %s", GCLOUD_CONFIG_DIR)
        else:
            log.warning("[CloudShell] No gcloud folder inside backup")

        env = _build_env()

        try:
            result = subprocess.run(
                [GCLOUD_BIN, "auth", "list"],
                capture_output=True, text=True, env=env, timeout=30,
            )
            auth_output = (result.stdout or "") + (result.stderr or "")
            log.info("[CloudShell] auth list:\n%s", auth_output.strip())
        except FileNotFoundError:
            msg = f"gcloud binary not found at {GCLOUD_BIN} -- SDK not installed"
            log.error("[CloudShell] %s", msg)
            return False, msg
        except subprocess.TimeoutExpired:
            log.warning("[CloudShell] auth list timed out, continuing anyway")

        CREDENTIALS_READY = True
        return True, "Credentials restored successfully from " + zip_path

    except Exception as e:
        log.error("[CloudShell] Credential restore error: %s", e, exc_info=True)
        return False, f"Credential restore error: {e}"


# ── Cloud Shell reachability ─────────────────────────────────────────────────

def check_cloud_shell(timeout_sec: int = 60) -> Tuple[bool, str]:
    """
    Quick sanity probe: run `echo OK` on Cloud Shell.
    Returns (success, output).
    """
    if not CREDENTIALS_READY:
        ok, msg = restore_credentials()
        if not ok:
            return False, msg

    env = _build_env()
    try:
        result = subprocess.run(
            [GCLOUD_BIN, "cloud-shell", "ssh",
             "--authorize-session", "--command=echo OK && whoami"],
            capture_output=True, text=True, env=env, timeout=timeout_sec,
        )
        out = (result.stdout or "") + (result.stderr or "")
        ok = result.returncode == 0 and "OK" in (result.stdout or "")
        return ok, out.strip()
    except FileNotFoundError:
        return False, f"gcloud binary not found at {GCLOUD_BIN}"
    except subprocess.TimeoutExpired:
        return False, "cloud-shell ssh timed out"
    except Exception as e:
        return False, f"cloud-shell ssh error: {e}"


# ── Reset Cloud Shell VM ─────────────────────────────────────────────────────

def reset_cloud_shell_vm(timeout_sec: int = 60) -> Tuple[bool, str]:
    """
    Force-reset the user's Cloud Shell VM.

    Cloud Shell doesn't have a clean "reset" CLI command, but we can
    achieve the same effect by killing all processes owned by the user
    on the remote VM. Cloud Shell auto-restarts the VM within ~30 seconds
    of all sessions ending.

    We do this by sending `pkill -9 -u $USER` over `gcloud cloud-shell ssh`.
    If the user is hitting a rate-limit / stuck session, this forcibly
    tears it down so the next task starts fresh.
    """
    if not CREDENTIALS_READY:
        ok, msg = restore_credentials()
        if not ok:
            return False, msg

    env = _build_env()
    try:
        # pkill -9 -u $USER kills everything owned by the user, including
        # any lingering server.py processes and the SSH daemon. Cloud Shell
        # detects the disconnect and restarts the VM.
        result = subprocess.run(
            [GCLOUD_BIN, "cloud-shell", "ssh",
             "--authorize-session",
             "--command=pkill -9 -u $USER; echo RESET_DONE"],
            capture_output=True, text=True, env=env, timeout=timeout_sec,
        )
        out = (result.stdout or "") + (result.stderr or "")
        # The pkill will kill our own SSH session, so we expect a non-zero
        # exit code or no output. That's fine -- the kill happened.
        if "RESET_DONE" in out or result.returncode in (0, 255, -1, None):
            return True, "Reset signal sent. Cloud Shell VM will restart within ~30 seconds."
        return False, f"Reset may have failed: {out.strip()[-300:]}"
    except subprocess.TimeoutExpired:
        return False, "Reset timed out (the kill probably still happened)"
    except Exception as e:
        return False, f"Reset error: {e}"


# ── Output streaming helpers ─────────────────────────────────────────────────

_ANSI_PATTERN_1 = re.compile(r'\x1b\[[0-9;]*[mKHJABCDsu]')
_ANSI_PATTERN_2 = re.compile(r'\x1b\][^\x07]*\x07')


def _strip_ansi(text: str) -> str:
    """Remove ANSI escape sequences so the web log is clean."""
    text = _ANSI_PATTERN_1.sub('', text)
    text = _ANSI_PATTERN_2.sub('', text)
    return text


def _emit_buffer(buffer: str, callback: Callable[[str], None]) -> str:
    """
    Send complete lines from `buffer` to `callback`. Returns the leftover
    partial line so we can keep buffering it until the next chunk arrives.

    Splits on BOTH \\n and \\r. This is critical because progress bars
    (yt-dlp's [DOWNLOAD], ffmpeg's [FFMPEG], kaggle's [KAG_UPLOAD]) use
    \\r (carriage return) to overwrite the line in place in a terminal.
    If we only split on \\n, progress updates get stuck in the buffer
    and never appear in the web UI until a \\n finally arrives — which
    might be at 100% completion or never.

    By splitting on \\r as well, each progress update becomes its own
    log line in the web UI. The user sees the download progress in
    real-time instead of a frozen log.
    """
    if not buffer:
        return buffer
    # Split on any combination of \r and \n. This handles:
    #   \n       (Unix line ending)
    #   \r\n     (Windows line ending)
    #   \r       (progress bar overwrite — the common case for yt-dlp/ffmpeg)
    #   \r\r...  (multiple progress updates batched in one chunk)
    parts = re.split(r'[\r\n]+', buffer)
    for complete in parts[:-1]:
        stripped = complete.strip()
        if stripped:
            callback(stripped)
    return parts[-1]


def _detect_rate_limit(text: str) -> bool:
    """Return True if any rate-limit marker is present in `text`."""
    lowered = text.lower()
    for marker in RATE_LIMIT_MARKERS:
        if marker.lower() in lowered:
            return True
    return False


def _detect_success(text: str) -> bool:
    """Return True if any success marker is present in `text`."""
    for marker in SUCCESS_MARKERS:
        if marker in text:
            return True
    return False


# ── Run a command on Cloud Shell (interactive session) ──────────────────────

def run_command(
    task_id: str,
    command: str,
    log_callback: Callable[[str], None],
    timeout: int = 1800,
) -> Tuple[bool, str]:
    """
    Execute `command` on Google Cloud Shell.

    Per the user's spec, the flow is:
      1. Initialize an interactive Cloud Shell session:
         `gcloud cloud-shell ssh --authorize-session`  (no --command flag)
      2. Wait for the shell to be ready (we look for a prompt or just
         wait a few seconds for boot output to settle).
      3. Send `command` to the live shell via the PTY stdin (write the
         command followed by a newline).
      4. Stream stdout/stderr line-by-line to `log_callback`.
      5. Watch for special markers:
           - Any of SUCCESS_MARKERS → success, close session
           - Any RATE_LIMIT_MARKERS → emit [RATE_LIMIT] log line,
             close session, return False
      6. Close the session IMMEDIATELY by sending `exit\n` and closing
         the PTY. We do NOT keep the shell alive between tasks.

    Args:
        task_id: Used for logging only.
        command: Shell command to execute on Cloud Shell.
        log_callback: Receives each output line (str, no trailing newline).
        timeout: Max seconds to wait. Default 30 min.

    Returns:
        (success, full_output) where:
          success = True if any SUCCESS_MARKER was seen in output.
          full_output = the complete captured stdout+stderr (ANSI stripped).
    """
    if not CREDENTIALS_READY:
        ok, msg = restore_credentials()
        if not ok:
            log_callback(f"[CloudShell] ERROR: {msg}")
            return False, msg

    env = _build_env()

    # Step 1: Initialize the interactive Cloud Shell session.
    # We do NOT pass --command here. We open a live PTY and send the
    # command ourselves once the shell is ready.
    full_cmd = [
        GCLOUD_BIN, "cloud-shell", "ssh", "--authorize-session",
    ]

    log.info("[%s][CloudShell] Initializing interactive session...", task_id)
    log_callback("[CloudShell] >>> Initializing Cloud Shell session (gcloud cloud-shell ssh --authorize-session)...")
    log_callback("[CloudShell] ... waiting for Cloud Shell VM to boot ...")

    try:
        master_fd, slave_fd = pty.openpty()

        proc = subprocess.Popen(
            full_cmd,
            stdin=slave_fd,
            stdout=slave_fd,
            stderr=slave_fd,
            close_fds=True,
            env=env,
            preexec_fn=os.setsid,
        )
        os.close(slave_fd)

        # v2.14.0: Store the proc globally so kill_running_session() can terminate it
        global _running_ssh_proc, _running_ssh_pid
        _running_ssh_proc = proc
        _running_ssh_pid = proc.pid

        output_buffer: list[str] = []
        partial_line = ""
        start_time = time.time()
        last_keepalive = start_time
        session_ready = False
        command_sent = False
        command_done = False
        rate_limited = False
        boot_settle_deadline = start_time + 90  # max 90s for Cloud Shell VM boot

        while True:
            # Process ended?
            if proc.poll() is not None:
                # Drain remaining output
                try:
                    while True:
                        r, _, _ = select.select([master_fd], [], [], 0.2)
                        if not r:
                            break
                        try:
                            data = os.read(master_fd, 4096).decode("utf-8", errors="replace")
                        except OSError:
                            break
                        if not data:
                            break
                        clean = _strip_ansi(data)
                        output_buffer.append(clean)
                        partial_line = _emit_buffer(partial_line + clean, log_callback)
                except OSError:
                    pass
                if partial_line.strip():
                    log_callback(partial_line.rstrip("\r"))
                    partial_line = ""
                break

            # Timeout?
            if time.time() - start_time > timeout:
                proc.kill()
                log_callback(f"[CloudShell] ERROR: Timed out after {timeout}s")
                # Try to close gracefully
                try:
                    os.write(master_fd, b"exit\n")
                except OSError:
                    pass
                try:
                    os.close(master_fd)
                except OSError:
                    pass
                return False, "Timed out after %ds" % timeout

            # Read whatever's available (0.5s timeout so we can also do
            # periodic actions like sending the command or heartbeats).
            try:
                r, _, _ = select.select([master_fd], [], [], 0.5)
                if r:
                    try:
                        data = os.read(master_fd, 4096).decode("utf-8", errors="replace")
                    except OSError:
                        break
                    if not data:
                        break
                    clean = _strip_ansi(data)
                    output_buffer.append(clean)
                    partial_line = _emit_buffer(partial_line + clean, log_callback)
                    last_keepalive = time.time()

                    # Check for rate-limit markers in the FULL output so far.
                    # We have to check the accumulated buffer because the
                    # marker may span multiple chunks.
                    full_so_far = "".join(output_buffer)
                    if _detect_rate_limit(full_so_far):
                        rate_limited = True
                        log_callback("[RATE_LIMIT] Cloud Shell rate limit detected.")
                        log_callback("[RATE_LIMIT] Please try again after 1 hour, or reset the Cloud Shell VM via the web UI (Settings → Reset Cloud Shell).")
                        # Break out -- we'll close the session below.
                        break

                    # If we haven't sent the command yet, look for evidence
                    # that the shell is ready. Cloud Shell prints a welcome
                    # banner and then drops to a $ prompt. We treat "user@"
                    # or "$ " or "# " as readiness. Also fall back to a
                    # 8-second timer in case the prompt is not detected.
                    if not session_ready:
                        if ("user@" in full_so_far
                                or "$ " in full_so_far[-200:]
                                or "# " in full_so_far[-200:]
                                or "Welcome to Cloud Shell" in full_so_far
                                or "Your Cloud Shell session is ready" in full_so_far
                                or time.time() - start_time > 8):
                            session_ready = True
                            log_callback("[CloudShell] ... Cloud Shell session ready.")

                    # Send the actual command once the shell is ready.
                    if session_ready and not command_sent:
                        # Small delay so the prompt is fully rendered.
                        time.sleep(0.3)
                        log_callback(f"[CloudShell] >>> Sending command: {command}")
                        try:
                            os.write(master_fd, (command + "\n").encode())
                        except OSError as e:
                            log_callback(f"[CloudShell] ERROR writing command: {e}")
                            break
                        command_sent = True
                        log_callback("[CloudShell] ... command sent, monitoring output ...")

                    # Check for success marker (any of SUCCESS_MARKERS)
                    if command_sent and _detect_success(full_so_far):
                        command_done = True
                        log_callback("[CloudShell] <<< SUCCESS marker detected.")
                        break

                else:
                    # No data this iteration.
                    # If we've been silent for >30s during command execution,
                    # emit a heartbeat so the user knows we're still alive.
                    if time.time() - last_keepalive > 30:
                        elapsed = int(time.time() - start_time)
                        log_callback(f"[CloudShell] ... still running ({elapsed}s)")
                        last_keepalive = time.time()

                    # If we've been waiting for the shell to boot for too
                    # long, abort with a clear error.
                    if not session_ready and time.time() > boot_settle_deadline:
                        log_callback("[CloudShell] ERROR: Cloud Shell VM did not become ready within 90s.")
                        log_callback("[CloudShell] This usually means a stuck session. Try resetting the VM via Settings → Reset Cloud Shell.")
                        break
            except OSError:
                break

        # ── Close the session IMMEDIATELY ────────────────────────────────────
        # Per user spec: "make sure close the cloud shell terminal instantly
        # once task done or fail". We send `exit` and close the PTY master.
        if command_done or rate_limited:
            try:
                os.write(master_fd, b"exit\n")
            except OSError:
                pass
            # Give the shell a moment to process the exit
            time.sleep(0.3)

        try:
            os.close(master_fd)
        except OSError:
            pass

        # Flush any trailing partial line
        if partial_line.strip():
            log_callback(partial_line.rstrip("\r"))
            partial_line = ""

        # Wait for the process to fully exit so proc.returncode is set.
        try:
            rc = proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()
            try:
                rc = proc.wait(timeout=5)
            except Exception:
                rc = None

        full_output = "".join(output_buffer)
        rc_display = rc if rc is not None else "None (SSH tunnel closed)"

        # Rate-limit takes precedence — it's a distinct failure mode that
        # the web UI handles specially (toast + banner).
        if rate_limited:
            log.warning("[%s][CloudShell] Rate limit hit (rc=%s)", task_id, rc_display)
            log_callback("[CloudShell] <<< Session closed (rate limit).")
            return False, full_output

        # Success marker (any of SUCCESS_MARKERS)
        if _detect_success(full_output):
            log.info("[%s][CloudShell] Command succeeded (marker found, rc=%s)", task_id, rc_display)
            log_callback(f"[CloudShell] <<< Command completed successfully (rc={rc_display}).")
            log_callback("[CloudShell] <<< Session closed.")
            return True, full_output

        # rc 0 but no marker — treat as success but warn
        if rc == 0:
            log.info("[%s][CloudShell] Command exited 0 (no SUCCESS marker)", task_id)
            log_callback("[CloudShell] <<< Command exited 0 (no SUCCESS marker found).")
            log_callback("[CloudShell] <<< Session closed.")
            return True, full_output

        # Anything else is a failure
        log.warning("[%s][CloudShell] Command failed (rc=%s)", task_id, rc_display)
        log_callback(f"[CloudShell] <<< Command exited with code {rc_display}.")
        log_callback("[CloudShell] <<< Session closed.")
        return False, full_output

    except FileNotFoundError:
        msg = f"gcloud binary not found at {GCLOUD_BIN}"
        log_callback(f"[CloudShell] ERROR: {msg}")
        return False, msg
    except Exception as e:
        msg = f"CloudShell execution error: {e}"
        log.error("[%s][CloudShell] %s", task_id, msg, exc_info=True)
        log_callback(f"[CloudShell] ERROR: {msg}")
        return False, msg


# ── High-level helper: download + trim via Cloud Shell ──────────────────────

def download_and_trim_via_cloudshell(
    task_id: str,
    url: str,
    start_time: str,
    end_time: str,
    log_callback: Callable[[str], None],
    timeout: int = 1800,
    env_extra: Optional[dict] = None,
) -> Tuple[bool, str]:
    """
    Convenience wrapper: builds the `python3 server.py ...` command exactly
    as the user specified and runs it on Cloud Shell via an interactive
    session.

        python3 server.py --id "<task_id>" --url "<url>" --start "<start>" --end "<end>"

    `server.py` is assumed to already live in the Cloud Shell user home
    (the user said the zip is uploaded in the HF project dir; on the Cloud
    Shell side they have presumably placed server.py there manually, or
    it's bundled in gcloud-backup.zip).

    Kaggle credentials are NOT passed through — the user confirmed that
    Kaggle is already configured on the Cloud Shell side (via
    `~/.kaggle/kaggle.json`). The `env_extra` parameter is kept for
    future use but not populated by default.

    Args:
        task_id: Used as the --id arg AND for log routing.
        url: YouTube URL.
        start_time: HH:MM:SS trim start.
        end_time: HH:MM:SS trim end.
        log_callback: Receives each output line.
        timeout: Max seconds.
        env_extra: Optional dict of env vars to export before server.py.
            Unused by default (Kaggle is set up on Cloud Shell itself).

    Returns:
        (success, output)
    """
    # Build the actual server.py command with shell-quoted args.
    server_cmd = "python3 server.py " + " ".join([
        "--id",     shlex.quote(task_id),
        "--url",    shlex.quote(url),
        "--start",  shlex.quote(start_time),
        "--end",    shlex.quote(end_time),
    ])

    # Prepend `export VAR=value &&` for each env var we want to pass through.
    prefix_parts = []
    if env_extra:
        for k, v in env_extra.items():
            if v is None or v == "":
                continue
            prefix_parts.append(f"export {k}={shlex.quote(str(v))}")
    prefix = " && ".join(prefix_parts) + " && " if prefix_parts else ""

    cmd = prefix + server_cmd

    log_callback(f"[CloudShell] Initiating Cloud Shell session for task {task_id}...")
    log_callback(f"[CloudShell] URL: {url}")
    log_callback(f"[CloudShell] Trim: {start_time} -> {end_time}")

    success, output = run_command(task_id, cmd, log_callback, timeout=timeout)

    if success:
        log_callback("[CloudShell] Cloud Shell session completed. Connection closed.")
    else:
        # Check if it was a rate-limit failure -- the run_command already
        # emitted a [RATE_LIMIT] marker in that case.
        if "[RATE_LIMIT]" in output:
            log_callback("[CloudShell] Task aborted due to Cloud Shell rate limit. See banner above.")
        else:
            log_callback("[CloudShell] Cloud Shell session ended with errors. Connection closed.")

    return success, output


# ── v2.14.0: Kill the running Cloud Shell SSH session ───────────────────────

def kill_running_session():
    """
    Kill the currently running Cloud Shell SSH subprocess.

    Called by app.py when the user clicks "Stop Task" or "Delete Task"
    while the download step is running. Sends SIGKILL to the entire
    process group (the SSH session + all child processes on Cloud Shell).

    This immediately closes the SSH connection, which stops the yt-dlp
    download + ffmpeg trim + Kaggle upload on the Cloud Shell side.
    """
    global _running_ssh_proc, _running_ssh_pid

    if _running_ssh_proc is None:
        log.info("[CloudShell] No running session to kill")
        return False

    try:
        pid = _running_ssh_pid
        log.info("[CloudShell] Killing SSH session (PID %s)...", pid)

        # Kill the entire process group (preexec_fn=os.setsid created a new group)
        try:
            os.killpg(os.getpgid(pid), signal.SIGKILL)
            log.info("[CloudShell] ✓ Process group killed")
        except ProcessLookupError:
            log.info("[CloudShell] Process already exited")
        except Exception as e:
            # Fall back to killing just the process
            try:
                _running_ssh_proc.kill()
            except Exception:
                pass
            log.info("[CloudShell] Kill fallback: %s", e)

        # Also try proc.terminate() as a backup
        try:
            _running_ssh_proc.terminate()
        except Exception:
            pass

        _running_ssh_proc = None
        _running_ssh_pid = None
        return True

    except Exception as e:
        log.error("[CloudShell] Error killing session: %s", e)
        _running_ssh_proc = None
        _running_ssh_pid = None
        return False