windmill/backend/windmill-common/src/worker.rs

#[cfg(unix)]
use anyhow::anyhow;
use axum::http::HeaderMap;
use bytes::Bytes;
use itertools::Itertools;
use regex::Regex;
use reqwest_middleware::ClientWithMiddleware;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serde_json::{json, value::RawValue};
use sqlx::{types::Json, Pool, Postgres};
use std::{
    cmp::Reverse,
    collections::{HashMap, HashSet},
    fs::{self, File},
    io::Write,
    ops::Deref,
    panic::Location,
    path::{Component, Path, PathBuf},
    str::FromStr,
    sync::atomic::{AtomicBool, AtomicI64, AtomicU32},
    time::Duration,
};
#[cfg(windows)]
use sysinfo::System;
use tokio::time::timeout;
use uuid::Uuid;
use windmill_macros::annotations;

use crate::{
    agent_workers::PingJobStatusResponse,
    cache::{unwrap_or_error, RawNode, RawScript},
    error::{self, to_anyhow},
    external_ip::UNKNOWN_IP,
    global_settings::CUSTOM_TAGS_SETTING,
    indexer::TantivyIndexerSettings,
    server::Smtp,
    utils::{merge_nested_raw_values_to_array, merge_raw_values_to_array},
    KillpillSender, DB,
};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct CustomTags {
    pub global: Vec<String>,
    pub specific: HashMap<String, SpecificTagData>,
}

impl CustomTags {
    pub fn from(tags: Vec<String>) -> Self {
        let mut global = vec![];
        let mut specific: HashMap<String, SpecificTagData> = HashMap::new();
        for e in tags {
            if let Some(cap) = CUSTOM_TAG_REGEX.captures(&e) {
                let tag_name = cap.get(1).unwrap().as_str().to_string();
                let workspace_str = cap.get(2).unwrap().as_str();
                let tag_type = SpecificTagType::from_regex_string(workspace_str);
                let workspaces: Vec<WorkspaceMatcher> = workspace_str
                    .split(tag_type.corresponding_separator())
                    .filter(|s| !s.is_empty())
                    .map(WorkspaceMatcher::parse)
                    .collect();
                if workspaces.is_empty() {
                    tracing::warn!("Ignoring tag `{}` with empty exclusion/inclusion list", e);
                    global.push(e);
                    continue;
                }
                specific.insert(tag_name, SpecificTagData { tag_type, workspaces });
            } else {
                global.push(e.to_string());
            }
        }
        Self { global, specific }
    }

    /// `filter_with_workspace` is the workspace's id chain (see [`SpecificTagData::applies_to_workspace`]);
    /// `None` re-emits the authored `tag(ws1+ws2)` strings for the settings editor.
    pub fn to_string_vec(&self, filter_with_workspace: Option<&[String]>) -> Vec<String> {
        let specific = if let Some(chain) = filter_with_workspace {
            self.specific
                .iter()
                .filter(|(_, tag_data)| tag_data.applies_to_workspace(chain))
                .map(|(tag, _)| tag.clone())
                .collect::<Vec<String>>()
        } else {
            self.specific
                .iter()
                .map(|(tag, tag_data)| {
                    let separator = tag_data.tag_type.corresponding_separator();
                    let mut workspaces = tag_data
                        .workspaces
                        .iter()
                        .map(|w| w.to_string())
                        .collect::<Vec<_>>()
                        .join(&*separator.to_string());
                    if tag_data.tag_type == SpecificTagType::AllExcluding {
                        // the AllExcluding tag syntax has a leading separator
                        workspaces.insert(0, separator);
                    }
                    format!("{}({})", tag, workspaces)
                })
                .collect::<Vec<String>>()
        };
        let all_tags = self.global.clone();
        all_tags.into_iter().chain(specific.into_iter()).collect()
    }
}

/// Marker suffixed to a workspace id inside a custom tag's scope (`mytag(prod*)`) to extend the
/// entry to that workspace's forks. `*` cannot appear in a workspace id (the `proper_id` check
/// constraint restricts them to `^\w+(-\w+)*$`), so it can never collide with a real id.
pub const FORK_SCOPE_MARKER: char = '*';

/// One workspace entry in a custom tag's scope. Bare (`prod`) matches that workspace only;
/// with the [`FORK_SCOPE_MARKER`] (`prod*`) it also matches its forks, transitively.
///
/// The marker is opt-in in BOTH scope forms so that no existing tag string changes meaning:
/// `sensitive(^prod)` keeps excluding only `prod` itself, and `sensitive(^prod*)` is how you
/// exclude its forks too.
#[derive(Clone, Serialize, Deserialize, PartialEq)]
pub struct WorkspaceMatcher {
    pub id: String,
    pub include_forks: bool,
}

impl WorkspaceMatcher {
    fn parse(entry: &str) -> Self {
        match entry.strip_suffix(FORK_SCOPE_MARKER) {
            Some(id) => Self { id: id.to_string(), include_forks: true },
            None => Self { id: entry.to_string(), include_forks: false },
        }
    }

    fn matches(&self, workspace_id: &str, fork_ancestors: &[String]) -> bool {
        workspace_id == self.id
            || (self.include_forks && fork_ancestors.iter().any(|a| *a == self.id))
    }
}

impl std::fmt::Display for WorkspaceMatcher {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.id)?;
        if self.include_forks {
            f.write_str(FORK_SCOPE_MARKER.encode_utf8(&mut [0u8; 4]))?;
        }
        Ok(())
    }
}

/// Renders the authored `prod` / `prod*` form rather than the struct fields: `CustomTags` is
/// `{:?}`-dumped into the "tag is not in the allowed CUSTOM_TAGS" error operators see.
impl std::fmt::Debug for WorkspaceMatcher {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Debug::fmt(&self.to_string(), f)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SpecificTagData {
    pub tag_type: SpecificTagType,
    pub workspaces: Vec<WorkspaceMatcher>,
}

impl SpecificTagData {
    /// `chain` is the workspace itself followed by its fork ancestors, nearest-first, as built by
    /// `workspaces::workspace_with_fork_ancestors`. Pass a single-element slice when
    /// [`Self::is_fork_scoped`] is false: the ancestors cannot affect the outcome then.
    pub fn applies_to_workspace(&self, chain: &[String]) -> bool {
        let Some((workspace_id, fork_ancestors)) = chain.split_first() else {
            return false;
        };
        let matched = self
            .workspaces
            .iter()
            .any(|w| w.matches(workspace_id, fork_ancestors));
        match self.tag_type {
            SpecificTagType::AllExcluding => !matched,
            SpecificTagType::NoneExcept => matched,
        }
    }

    /// Whether any entry carries the fork marker, i.e. whether resolving the workspace's fork
    /// lineage can change what [`Self::applies_to_workspace`] returns. Lets hot callers skip the
    /// lineage lookup for the (overwhelmingly common) fork-agnostic tag.
    pub fn is_fork_scoped(&self) -> bool {
        self.workspaces.iter().any(|w| w.include_forks)
    }
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SpecificTagType {
    AllExcluding,
    NoneExcept,
}

impl SpecificTagType {
    pub fn corresponding_separator(&self) -> char {
        match self {
            SpecificTagType::AllExcluding => '^',
            SpecificTagType::NoneExcept => '+',
        }
    }

    pub fn from_regex_string(workspaces_str: &str) -> Self {
        if workspaces_str.contains(SpecificTagType::AllExcluding.corresponding_separator()) {
            SpecificTagType::AllExcluding
        } else {
            // Regex match ensures the second branch is correct
            SpecificTagType::NoneExcept
        }
    }
}

pub const DEFAULT_CLOUD_TIMEOUT: u64 = 900;
pub const DEFAULT_SELFHOSTED_TIMEOUT: u64 = 604800; // 7 days
pub const MIN_PERIODIC_SCRIPT_INTERVAL_SECONDS: u64 = 60;
/// Default for [`CONCURRENCY_KEY_MAX_QUEUED`]; also the value the setting loader restores when
/// the setting is cleared or malformed.
pub const CONCURRENCY_KEY_MAX_QUEUED_DEFAULT: u32 = 10_000;
/// Default for [`WORKSPACE_MAX_QUEUED_JOBS`]; also the value the setting loader restores when
/// the setting is cleared or malformed. A workspace spans many keys, so this sits well above
/// the per-key cap.
pub const WORKSPACE_MAX_QUEUED_JOBS_DEFAULT: u32 = 20_000;
/// Default for [`JOB_OOM_SCORE_ADJ`]; also the value used when the env var is out of range or
/// unparseable.
pub const JOB_OOM_SCORE_ADJ_DEFAULT: i32 = 1000;

/// procfs accepts -1000..=1000, but a job must never be *less* killable than the worker that
/// supervises it, so negative adjustments are rejected rather than clamped.
fn parse_job_oom_score_adj(raw: Option<&str>) -> i32 {
    let Some(raw) = raw else {
        return JOB_OOM_SCORE_ADJ_DEFAULT;
    };
    match raw.trim().parse::<i32>() {
        Ok(v) if (0..=1000).contains(&v) => v,
        Ok(v) => {
            tracing::warn!(
                "JOB_OOM_SCORE_ADJ={v} is outside the accepted 0..=1000 range, \
                using {JOB_OOM_SCORE_ADJ_DEFAULT}"
            );
            JOB_OOM_SCORE_ADJ_DEFAULT
        }
        Err(e) => {
            tracing::warn!(
                "Could not parse JOB_OOM_SCORE_ADJ='{raw}': {e}, using {JOB_OOM_SCORE_ADJ_DEFAULT}"
            );
            JOB_OOM_SCORE_ADJ_DEFAULT
        }
    }
}
lazy_static::lazy_static! {
    pub static ref WORKER_GROUP: String = std::env::var("WORKER_GROUP").unwrap_or_else(|_| {
        #[cfg(not(feature = "enterprise"))]
        {
            "default".to_string()
        }

        #[cfg(feature = "enterprise")]
        {
            if let Some(token) = crate::agent_workers::DECODED_AGENT_TOKEN.as_ref() {
                token.worker_group.clone()
            } else {
                "default".to_string()
            }
        }
    });

    pub static ref NO_LOGS: bool = std::env::var("NO_LOGS").ok().is_some_and(|x| x == "1" || x == "true");

    /// Shut the worker process down once it has executed this many jobs, so a supervisor
    /// (docker restart policy, kubernetes, systemd, ...) restarts it on a pristine
    /// environment. Meant for deployments that cannot sandbox jobs with nsjail and rely on
    /// the process/container lifetime to isolate one execution from the next. `0` (or unset)
    /// disables it. Only the jobs the worker's own main loop ran count: ones handed off to a
    /// dedicated worker or a flow runner are executed by another task and never counted.
    /// Workers of one process share that environment, so with `NUM_WORKERS > 1` the first of
    /// them to reach the limit takes the whole process down, cancelling whatever the others
    /// still run in a container or a dedicated worker. Run one worker per process.
    /// A value that does not parse is rejected at startup by [`validate_worker_lifecycle_env`]
    /// rather than read as "disabled" here: a deployment that isolates executions this way
    /// would otherwise keep running with no isolation at all.
    pub static ref EXIT_AFTER_N_JOBS: Option<u64> = std::env::var("EXIT_AFTER_N_JOBS")
        .ok()
        .and_then(|x| x.parse::<u64>().ok())
        .filter(|x| *x > 0);

    /// Replaces the random part of the worker name, which is what makes a restarted process
    /// reclaim its `worker_ping` row rather than register as a new worker. Two worker
    /// processes must never share it: it is only needed when several of them run on one host
    /// under the same worker group, since the name is otherwise derived from the hostname.
    pub static ref WORKER_SUFFIX: Option<String> = std::env::var("WORKER_SUFFIX")
        .ok()
        .filter(|x| !x.is_empty());

    pub static ref NATIVE_MODE: bool = std::env::var("NATIVE_MODE").ok().is_some_and(|x| x == "1" || x == "true");

    pub static ref LIMIT_WINDOWS_TO_1CU: bool = std::env::var("LIMIT_WINDOWS_TO_1CU").ok().is_some_and(|x| x == "1" || x == "true");

    /// `oom_score_adj` applied to job subprocesses. The kernel adds it to the process's memory
    /// use expressed in permille of host RAM, so the job only reliably outranks the worker once
    /// the gap between their two adjustments exceeds the worker's own footprint in permille; the
    /// default maximizes that margin. Userspace OOM daemons (earlyoom, systemd-oomd, nohang) rank
    /// every process on the host by the same score, so at 1000 a tiny job outranks multi-GB
    /// processes and gets killed first. Lowering this trades margin over the worker for a fairer
    /// ranking against everything else on the host.
    pub static ref JOB_OOM_SCORE_ADJ: i32 = parse_job_oom_score_adj(std::env::var("JOB_OOM_SCORE_ADJ").ok().as_deref());

    pub static ref CGROUP_V2_PATH_RE: Regex = Regex::new(r#"(?m)^0::(/.*)$"#).unwrap();
    pub static ref CGROUP_V2_CPU_RE: Regex = Regex::new(r#"(?m)^(\d+) \S+$"#).unwrap();
    pub static ref CGROUP_V1_INACTIVE_FILE_RE: Regex = Regex::new(r#"(?m)^total_inactive_file (\d+)$"#).unwrap();
    pub static ref CGROUP_V2_INACTIVE_FILE_RE: Regex = Regex::new(r#"(?m)^inactive_file (\d+)$"#).unwrap();

    pub static ref DEFAULT_TAGS: Vec<String> = vec![
        "deno".to_string(),
        "python3".to_string(),
        "go".to_string(),
        "bash".to_string(),
        "powershell".to_string(),
        "nativets".to_string(),
        "mysql".to_string(),
        "oracledb".to_string(),
        "bun".to_string(),
        "postgresql".to_string(),
        "bigquery".to_string(),
        "snowflake".to_string(),
        "mssql".to_string(),
        "graphql".to_string(),
        "php".to_string(),
        "rust".to_string(),
        "ansible".to_string(),
        "csharp".to_string(),
        "nu".to_string(),
        "java".to_string(),
        "ruby".to_string(),
        "rlang".to_string(),
        "duckdb".to_string(),
        "dbt".to_string(),
        // for related places search: ADD_NEW_LANG
        "dependency".to_string(),
        "flow".to_string(),
        "other".to_string()
    ];

    pub static ref NATIVE_TAGS: Vec<String> = vec![
        "nativets".to_string(),
        "postgresql".to_string(),
        "mysql".to_string(),
        "graphql".to_string(),
        "snowflake".to_string(),
        "mssql".to_string(),
        "bigquery".to_string(),
        "oracledb".to_string()
        // for related places search: ADD_NEW_LANG
    ];

    pub static ref DEFAULT_TAGS_PER_WORKSPACE: AtomicBool = AtomicBool::new(false);
    pub static ref DEFAULT_TAGS_WORKSPACES: arc_swap::ArcSwap<Option<Vec<String>>> = arc_swap::ArcSwap::from_pointee(None);
    pub static ref FORK_WORKSPACE_TAG_APPEND_FORK_SUFFIX: AtomicBool = AtomicBool::new(false);
    pub static ref PREVIEW_TAGS_OVERRIDE: AtomicBool = AtomicBool::new(false);

    pub static ref MAX_TIMEOUT: u64 = std::env::var("TIMEOUT")
    .ok()
    .and_then(|x| x.parse::<u64>().ok())
    .unwrap_or_else(|| if *CLOUD_HOSTED { DEFAULT_CLOUD_TIMEOUT } else { DEFAULT_SELFHOSTED_TIMEOUT });

    pub static ref SCRIPT_TOKEN_EXPIRY: u64 = std::env::var("SCRIPT_TOKEN_EXPIRY")
        .ok()
        .and_then(|x| x.parse::<u64>().ok())
        .unwrap_or(*MAX_TIMEOUT);

    pub static ref WORKER_CONFIG: arc_swap::ArcSwap<WorkerConfig> = arc_swap::ArcSwap::from_pointee(WorkerConfig {
        worker_tags: Default::default(),
        priority_tags_sorted: Default::default(),
        dedicated_worker: Default::default(),
        dedicated_workers: Default::default(),
        cache_clear: Default::default(),
        init_bash: Default::default(),
        periodic_script_bash: Default::default(),
        periodic_script_interval_seconds: Default::default(),
        additional_python_paths: Default::default(),
        pip_local_dependencies: Default::default(),
        env_vars: Default::default(),
        native_mode: false,
    });

    pub static ref WORKER_PULL_QUERIES: arc_swap::ArcSwap<Vec<String>> = arc_swap::ArcSwap::from_pointee(vec![]);
    pub static ref WORKER_PULL_QUERIES_FAIRNESS: arc_swap::ArcSwap<Vec<String>> = arc_swap::ArcSwap::from_pointee(vec![]);
    pub static ref WORKER_SUSPENDED_PULL_QUERY: arc_swap::ArcSwap<String> = arc_swap::ArcSwap::from_pointee("".to_string());

    // Workspace fairness (cloud-only). When enabled, a workspace whose footprint over the rolling
    // `WORKSPACE_FAIRNESS_DURATION_SECS` window represents >= `WORKSPACE_FAIRNESS_MAX_PERCENT`% of
    // all worker activity gets excluded from the pull query, freeing slots for other workspaces.
    // The list of overloaded workspaces is computed cluster-wide via a single coordinated UPDATE
    // on `background_task_state` so only one process per refresh interval runs the aggregation.
    pub static ref WORKSPACE_FAIRNESS_ENABLED: AtomicBool = AtomicBool::new(false);
    pub static ref WORKSPACE_FAIRNESS_MAX_PERCENT: AtomicU32 = AtomicU32::new(50);
    pub static ref WORKSPACE_FAIRNESS_DURATION_SECS: AtomicU32 = AtomicU32::new(10);
    pub static ref WORKSPACE_FAIRNESS_MIN_TOTAL: AtomicU32 = AtomicU32::new(4);
    pub static ref WORKSPACE_FAIRNESS_OVERLOADED: arc_swap::ArcSwap<Vec<String>> = arc_swap::ArcSwap::from_pointee(vec![]);
    pub static ref WORKSPACE_FAIRNESS_LAST_REFRESH_MICROS: AtomicI64 = AtomicI64::new(0);

    /// Stochastic admission probability for capped workspaces, expressed in
    /// parts per 10_000 (so `420` = 4.2%). The refresh computes this from the
    /// observed worker-second distribution and the configured cap so that
    /// admission converges to the target *worker-second* share — independent
    /// of how the capped vs uncapped workspaces compare on per-job durations.
    /// See `workspace_fairness_ee::refresh_overloaded` for the derivation.
    /// `10_000` (= admit all) is the default until the first refresh
    /// classifies an overloaded set — before that, no workspace is capped so
    /// `should_admit_capped` is moot and "admit all" is the correct no-op.
    pub static ref WORKSPACE_FAIRNESS_ADMISSION_PPM: AtomicU32 = AtomicU32::new(10_000);

    /// Cloud-only ceiling on the number of jobs queued behind a single concurrency key.
    /// A concurrency-limited key drains at most `concurrent_limit` jobs per window, so a
    /// producer pushing faster than that grows an unbounded backlog that no amount of
    /// spare worker capacity can absorb. `0` disables the cap.
    pub static ref CONCURRENCY_KEY_MAX_QUEUED: AtomicU32 =
        AtomicU32::new(CONCURRENCY_KEY_MAX_QUEUED_DEFAULT);

    /// Cloud-only ceiling on the total number of jobs a workspace may have queued at once,
    /// across every concurrency key and script. Guards against a workspace flooding the queue
    /// generally (not just behind one key), including from parallel for-loops. `0` disables it.
    pub static ref WORKSPACE_MAX_QUEUED_JOBS: AtomicU32 =
        AtomicU32::new(WORKSPACE_MAX_QUEUED_JOBS_DEFAULT);


    pub static ref SMTP_CONFIG: arc_swap::ArcSwap<Option<Smtp>> = arc_swap::ArcSwap::from_pointee(None);
    pub static ref INDEXER_CONFIG: arc_swap::ArcSwap<TantivyIndexerSettings> = arc_swap::ArcSwap::from_pointee(TantivyIndexerSettings::default());


    pub static ref CLOUD_HOSTED: bool = std::env::var("CLOUD_HOSTED").is_ok();
    /// Host used to gate cloud-only features that must only ever run on the
    /// production `app.windmill.dev` cluster, not on staging or self-hosted.
    pub static ref CLOUD_PRODUCTION_HOST: &'static str = "app.windmill.dev";

    /// `--no-auth` mode: when set, every API request is treated as
    /// authenticated as the `admin@windmill.dev` superadmin and no login is
    /// ever required. Meant for self-hosted deployments that front Windmill
    /// with their own authenticating gateway. Never honored on the managed
    /// cloud (`CLOUD_HOSTED`), which must always enforce real authentication.
    pub static ref NO_AUTH: bool = !*CLOUD_HOSTED
        && std::env::var("NO_AUTH").ok().is_some_and(|x| x == "1" || x == "true");

    pub static ref CUSTOM_TAGS: Vec<String> = std::env::var("CUSTOM_TAGS")
        .ok()
        .map(|x| x.split(',').map(|x| x.to_string()).collect::<Vec<_>>()).unwrap_or_default();


    pub static ref CUSTOM_TAGS_PER_WORKSPACE: arc_swap::ArcSwap<CustomTags> = arc_swap::ArcSwap::from_pointee(CustomTags::default());

    pub static ref ALL_TAGS: arc_swap::ArcSwap<Vec<String>> = arc_swap::ArcSwap::from_pointee(vec![]);



    //    ^([\w-]+)         # Group 1: tag name
    //    \(                # Literal '('
    //    (                 # Group 2: the full workspace list
    //      (?:[\w-]+\*?\+)*[\w-]+\*?   # NoneExcept pattern: ws1+ws2*
    //      |                          # OR
    //      (?:\^[\w-]+\*?)+           # AllExcluding pattern: ^ws1^ws2*
    //    )
    //    \)$               # Closing ')'
    //
    // The optional `*` after each workspace id is the fork marker, see [`WorkspaceMatcher`].
    static ref CUSTOM_TAG_REGEX: Regex = Regex::new(r"^([\w-]+)\(((?:[\w-]+\*?\+)*[\w-]+\*?|(?:\^[\w-]+\*?)+)\)$").unwrap();

    pub static ref DISABLE_BUNDLING: bool = std::env::var("DISABLE_BUNDLING")
    .ok()
    .and_then(|x| x.parse::<bool>().ok())
    .unwrap_or(false);

    // Features flags:
    pub static ref DISABLE_FLOW_SCRIPT: bool = std::env::var("DISABLE_FLOW_SCRIPT").ok().is_some_and(|x| x == "1" || x == "true");

    pub static ref ROOT_STANDALONE_BUNDLE_DIR: String = format!("{}/.windmill/standalone_bundle", std::env::var("HOME").unwrap_or_else(|_| "/root".to_string()));
}

lazy_static::lazy_static! {
    pub static ref ROOT_CACHE_NOMOUNT_DIR: String = format!("{}/cache_nomount/", *WINDMILL_DIR);
}

/// Refuses to start on an `EXIT_AFTER_N_JOBS` that does not parse. Silently ignoring it
/// would leave a worker meant to recycle its environment running forever without doing so,
/// which is exactly the guarantee the deployment set it for.
pub fn validate_worker_lifecycle_env() -> anyhow::Result<()> {
    match std::env::var("EXIT_AFTER_N_JOBS") {
        Ok(v) if !v.is_empty() && v.parse::<u64>().is_err() => Err(anyhow::anyhow!(
            "EXIT_AFTER_N_JOBS must be a positive integer (or 0 to disable), got '{v}'"
        )),
        _ => Ok(()),
    }
}

/// Whether native mode is forced by the environment (NATIVE_MODE=true env var or WORKER_GROUP=native).
/// This does NOT account for native_mode set in the DB worker group config — for that, read
/// `WORKER_CONFIG.native_mode` which combines all sources.
pub fn is_native_mode_from_env() -> bool {
    *NATIVE_MODE || *WORKER_GROUP == "native"
}

/// True iff this process is configured to act as the production cloud cluster:
/// `CLOUD_HOSTED=true` AND `BASE_URL`'s host matches `CLOUD_PRODUCTION_HOST`.
/// Centralized so the API setter, the runtime pull path, and any future cloud-
/// only feature share one canonical check (rather than re-implementing the
/// scheme/host parser at each call site).
pub fn is_cloud_production_host() -> bool {
    if !*CLOUD_HOSTED {
        return false;
    }
    let base = crate::BASE_URL.load();
    let s = base.as_str();
    if s.is_empty() {
        return false;
    }
    let after_scheme = s
        .strip_prefix("https://")
        .or_else(|| s.strip_prefix("http://"))
        .unwrap_or(s);
    let host = after_scheme
        .split('/')
        .next()
        .unwrap_or("")
        .split(':')
        .next()
        .unwrap_or("");
    host == *CLOUD_PRODUCTION_HOST
}

/// Cached resolved native mode flag, updated when worker config is reloaded.
/// Use this for hot-path checks (e.g. per-job dispatch) to avoid read-locking WORKER_CONFIG.
pub static NATIVE_MODE_RESOLVED: AtomicBool = AtomicBool::new(false);

pub static MIN_VERSION_IS_LATEST: AtomicBool = AtomicBool::new(false);
#[derive(Clone)]
pub struct HttpClient {
    pub client: ClientWithMiddleware,
    pub base_internal_url: String,
}

impl Deref for HttpClient {
    type Target = ClientWithMiddleware;

    fn deref(&self) -> &Self::Target {
        &self.client
    }
}

impl HttpClient {
    pub async fn post<T: Serialize, R: DeserializeOwned>(
        &self,
        url: &str,
        headers: Option<HeaderMap>,
        body: &T,
    ) -> anyhow::Result<R> {
        let base_url = self.base_internal_url.clone();

        let response_builder = self.client.post(format!("{}{}", base_url, url)).json(body);

        let response_builder = match headers {
            Some(headers) => response_builder.headers(headers),
            None => response_builder,
        };

        let response = response_builder
            .send()
            .await
            .map_err(|e| anyhow::anyhow!(e))?;
        let status = response.status();
        if status.is_success() {
            Ok(response.json().await?)
        } else {
            Err(anyhow::anyhow!(format!(
                "HTTP agent request POST {} failed {}",
                url,
                response.status()
            )))
        }
    }

    pub async fn get<R: DeserializeOwned>(&self, url: &str) -> anyhow::Result<R> {
        let base_url = self.base_internal_url.clone();
        let response = self
            .client
            .get(format!("{}{}", base_url, url))
            .send()
            .await
            .map_err(|e| anyhow::anyhow!(e))?;
        let status = response.status();
        if status.is_success() {
            Ok(response.json().await?)
        } else {
            Err(anyhow::anyhow!(format!(
                "HTTP agent request GET {} failed {}",
                url,
                response.status()
            )))
        }
    }

    pub async fn get_bytes(&self, url: &str) -> anyhow::Result<Bytes> {
        let base_url = self.base_internal_url.clone();
        let response = self
            .client
            .get(format!("{}{}", base_url, url))
            .send()
            .await
            .map_err(|e| anyhow::anyhow!(e))?;
        if response.status().is_success() {
            Ok(response.bytes().await?)
        } else {
            Err(anyhow::anyhow!(
                "HTTP agent request GET {} failed {}",
                url,
                response.status()
            ))
        }
    }

    pub async fn put_bytes(&self, url: &str, bytes: Bytes) -> anyhow::Result<()> {
        let base_url = self.base_internal_url.clone();
        let response = self
            .client
            .put(format!("{}{}", base_url, url))
            .body(bytes)
            .send()
            .await
            .map_err(|e| anyhow::anyhow!(e))?;
        if response.status().is_success() {
            Ok(())
        } else {
            Err(anyhow::anyhow!(
                "HTTP agent request PUT {} failed {}",
                url,
                response.status()
            ))
        }
    }
}

#[derive(Clone)]
pub enum Connection {
    Sql(Pool<Postgres>),
    Http(HttpClient),
}

impl std::fmt::Debug for Connection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Connection::Sql(_) => write!(f, "Sql"),
            Connection::Http(_) => write!(f, "Http"),
        }
    }
}

impl Connection {
    pub fn as_sql(&self) -> Option<&Pool<Postgres>> {
        match self {
            Connection::Sql(db) => Some(db),
            Connection::Http(_) => None,
        }
    }
}

impl From<Pool<Postgres>> for Connection {
    fn from(value: Pool<Postgres>) -> Self {
        Connection::Sql(value)
    }
}

impl From<&Pool<Postgres>> for Connection {
    fn from(value: &Pool<Postgres>) -> Self {
        Connection::Sql(value.clone())
    }
}

fn format_pull_query(peek: String) -> String {
    let r = format!(
        "WITH peek AS (
            {}
        ), q AS NOT MATERIALIZED (
            UPDATE v2_job_queue SET
                running = true,
                started_at = coalesce(started_at, now()),
                suspend_until = null,
                worker = $1
            WHERE id = (SELECT id FROM peek)
            RETURNING
                started_at, scheduled_for,
                canceled_by, canceled_reason, worker, cache_ignore_s3_path, runnable_settings_handle
        ), r AS NOT MATERIALIZED (
            UPDATE v2_job_runtime SET
                ping = now()
            WHERE id = (SELECT id FROM peek)
        ), j AS NOT MATERIALIZED (
            SELECT
                id, workspace_id, parent_job, created_by, created_at, runnable_id,
                runnable_path, args, kind, trigger, trigger_kind,
                permissioned_as, permissioned_as_email, script_lang,
                flow_innermost_root_job, root_job, flow_step_id,
                same_worker, pre_run_error, visible_to_owner, tag, concurrent_limit,
                concurrency_time_window_s, timeout, cache_ttl, priority, raw_code, raw_lock,
                raw_flow, script_entrypoint_override, preprocessed
            FROM v2_job
            WHERE id = (SELECT id FROM peek)
        ) SELECT j.id, j.workspace_id, j.parent_job, j.created_by, started_at, scheduled_for,
            j.runnable_id, j.runnable_path, j.args, canceled_by,
            canceled_reason, j.kind, j.trigger, j.trigger_kind, j.permissioned_as,
            flow_status, j.script_lang,
            j.same_worker, j.pre_run_error, j.visible_to_owner,
            j.tag, j.concurrent_limit, j.concurrency_time_window_s, j.flow_innermost_root_job, j.root_job,
            j.timeout, j.flow_step_id, j.cache_ttl, q.cache_ignore_s3_path, q.runnable_settings_handle, j.priority, j.raw_code, j.raw_lock, j.raw_flow,
            j.script_entrypoint_override, j.preprocessed, COALESCE(pj.runnable_path, j.args->>'_FLOW_PATH') as parent_runnable_path,
            COALESCE(p.email, j.permissioned_as_email) as permissioned_as_email, p.username as permissioned_as_username, p.is_admin as permissioned_as_is_admin,
            p.is_operator as permissioned_as_is_operator, p.groups as permissioned_as_groups, p.folders as permissioned_as_folders, p.end_user_email as permissioned_as_end_user_email
        FROM q, j
            LEFT JOIN v2_job_status f USING (id)
            LEFT JOIN job_perms p ON p.job_id = j.id
            LEFT JOIN v2_job pj ON j.parent_job = pj.id
            ",
        peek
    );
    // tracing::debug!("pull query: {}", r);
    r
}

pub fn make_suspended_pull_query(tags: &[String]) -> String {
    format_pull_query(format!(
        "SELECT id
        FROM v2_job_queue
        WHERE suspend_until IS NOT NULL AND (suspend <= 0 OR suspend_until <= now()) AND tag IN ({})
        ORDER BY priority DESC NULLS LAST, created_at
        FOR UPDATE SKIP LOCKED
        LIMIT 1",
        tags.iter().map(|x| format!("'{x}'")).join(", ")
    ))
}
// pub async fn make_suspended
pub async fn store_suspended_pull_query(wc: &WorkerConfig) {
    if wc.worker_tags.len() == 0 {
        tracing::error!("Empty tags in worker tags, skipping");
        return;
    }
    let query = make_suspended_pull_query(&wc.worker_tags);
    WORKER_SUSPENDED_PULL_QUERY.store(std::sync::Arc::new(query));
}

pub fn make_pull_query(tags: &[String]) -> String {
    let query = format_pull_query(format!(
        "SELECT id
        FROM v2_job_queue
        WHERE running = false AND tag IN ({}) AND scheduled_for <= now()
        ORDER BY priority DESC NULLS LAST, scheduled_for
        FOR UPDATE SKIP LOCKED
        LIMIT 1",
        tags.iter().map(|x| format!("'{x}'")).join(", ")
    ));
    query
}

// Variant of `make_pull_query` that additionally excludes jobs whose workspace_id is in the
// overloaded-list bind parameter ($2::text[]). Built as a separate string (rather than reusing
// `make_pull_query` with an always-bound array) so the planner can keep using the same indexes
// when fairness is off — the default `make_pull_query` text stays bit-identical to today's.
//
// `pub(crate)` because only `store_pull_query` consumes it; the resulting query string is what
// crosses crate boundaries via `WORKER_PULL_QUERIES_FAIRNESS`.
pub(crate) fn make_pull_query_fairness(tags: &[String]) -> String {
    let query = format_pull_query(format!(
        "SELECT id
        FROM v2_job_queue
        WHERE running = false AND tag IN ({}) AND scheduled_for <= now()
          AND workspace_id <> ALL($2::text[])
        ORDER BY priority DESC NULLS LAST, scheduled_for
        FOR UPDATE SKIP LOCKED
        LIMIT 1",
        tags.iter().map(|x| format!("'{x}'")).join(", ")
    ));
    query
}

pub async fn store_pull_query(wc: &WorkerConfig) {
    let mut queries = vec![];
    let mut fairness_queries = vec![];
    let fairness_enabled = WORKSPACE_FAIRNESS_ENABLED.load(std::sync::atomic::Ordering::Relaxed);
    for tags in wc.priority_tags_sorted.iter() {
        if tags.tags.len() == 0 {
            tracing::error!("Empty tags in priority tags, skipping");
            continue;
        }
        queries.push(make_pull_query(&tags.tags));
        if fairness_enabled {
            fairness_queries.push(make_pull_query_fairness(&tags.tags));
        }
    }
    WORKER_PULL_QUERIES.store(std::sync::Arc::new(queries));
    WORKER_PULL_QUERIES_FAIRNESS.store(std::sync::Arc::new(fairness_queries));
}

lazy_static::lazy_static! {
    pub static ref WINDMILL_DIR: String = {
        let dir = std::env::var("WINDMILL_DIR")
            .unwrap_or_else(|_| {
                #[cfg(not(windows))]
                { "/tmp/windmill".to_string() }
                #[cfg(windows)]
                {
                    let temp = std::env::temp_dir();
                    let temp_str = temp.to_string_lossy();
                    let normalized = temp_str.trim_end_matches(&['/', '\\'][..]).replace('\\', "/");
                    format!("{}/windmill", normalized)
                }
            });
        if dir.is_empty() {
            panic!("WINDMILL_DIR must not be empty");
        }
        if dir.ends_with('/') || dir.ends_with('\\') {
            panic!("WINDMILL_DIR must not end with a trailing slash, got: {dir}");
        }
        dir
    };
    pub static ref TMP_LOGS_DIR: String = format!("{}/logs", *WINDMILL_DIR);
    pub static ref ROOT_CACHE_DIR: String = format!("{}/cache/", *WINDMILL_DIR);
    pub static ref HUB_CACHE_DIR: String = format!("{}hub", *ROOT_CACHE_DIR);
    pub static ref HUB_RT_CACHE_DIR: String = format!("{}hub_rt", *ROOT_CACHE_DIR);
}

pub fn write_file(dir: &str, path: &str, content: &str) -> error::Result<File> {
    let path = format!("{}/{}", dir, path);
    let mut file = File::create(&path).map_err(|e| {
        tracing::error!("Failed to create file at {path}: {:?}", &e);
        e
    })?;
    file.write_all(content.as_bytes())?;
    file.flush()?;
    Ok(file)
}

pub fn write_file_bytes(path: &str, content: &Bytes) -> error::Result<File> {
    let mut file = File::create(path)?;
    file.write_all(content)?;
    file.flush()?;
    Ok(file)
}
/// from : https://github.com/rust-lang/cargo/blob/fede83ccf973457de319ba6fa0e36ead454d2e20/src/cargo/util/paths.rs#L61
fn normalize_path(path: &Path) -> PathBuf {
    let mut components = path.components().peekable();
    let mut ret = if let Some(c @ Component::Prefix(..)) = components.peek().cloned() {
        components.next();
        PathBuf::from(c.as_os_str())
    } else {
        PathBuf::new()
    };

    for component in components {
        match component {
            Component::Prefix(..) => unreachable!(),
            Component::RootDir => {
                ret.push(component.as_os_str());
            }
            Component::CurDir => {}
            Component::ParentDir => {
                ret.pop();
            }
            Component::Normal(c) => {
                ret.push(c);
            }
        }
    }
    ret
}

pub fn is_allowed_file_location(job_dir: &str, user_defined_path: &str) -> error::Result<PathBuf> {
    let job_dir = Path::new(job_dir);
    let user_path = PathBuf::from(user_defined_path);

    let full_path = job_dir.join(&user_path);

    let normalized_job_dir = normalize_path(job_dir);
    let normalized_full_path = normalize_path(&full_path);

    if !normalized_full_path.starts_with(&normalized_job_dir) {
        return Err(std::io::Error::new(
            std::io::ErrorKind::PermissionDenied,
            "Path is outside the allowed job directory.",
        )
        .into());
    }

    // The lexical check above cannot see symlinks: a symlink planted inside the
    // job dir - e.g. by an earlier Ansible `git_repos` clone whose tracked
    // content includes one - would let a later `git clone` or file write follow
    // it out of the job dir while still passing the textual `starts_with` check.
    // Walk the *normalized* relative path (`..`/`.` already collapsed) so each
    // step matches the real on-disk resolution, and reject any existing component
    // that is a symlink. Walking the raw user path would drift on an in-bounds
    // `..` (e.g. `foo/../link`, which normalizes back inside the job dir) and miss
    // the real symlinked component. Not-yet-existing components are safe: a path
    // that does not exist cannot itself be a symlink.
    let relative = normalized_full_path
        .strip_prefix(&normalized_job_dir)
        .unwrap_or(&normalized_full_path);
    let mut current = normalized_job_dir.clone();
    for component in relative.components() {
        if let Component::Normal(c) = component {
            current.push(c);
            if std::fs::symlink_metadata(&current)
                .map(|m| m.file_type().is_symlink())
                .unwrap_or(false)
            {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::PermissionDenied,
                    "Path traverses a symlink, which is not allowed.",
                )
                .into());
            }
        }
    }

    Ok(normalized_full_path)
}

pub fn write_file_at_user_defined_location(
    job_dir: &str,
    user_defined_path: &str,
    content: &str,
    mode: Option<u32>,
) -> error::Result<PathBuf> {
    let normalized_full_path = is_allowed_file_location(job_dir, user_defined_path)?;

    let full_path = normalized_full_path.as_path();
    if let Some(parent_dir) = full_path.parent() {
        std::fs::create_dir_all(parent_dir)?;
    }

    let mut file = File::create(full_path)?;

    #[cfg(unix)]
    if let Some(mode) = mode {
        let perm = std::os::unix::fs::PermissionsExt::from_mode(mode);
        file.set_permissions(perm)
            .map_err(|e| anyhow!("Failed to set permissions to {}: {e}", user_defined_path))?;
    }

    #[cfg(windows)]
    if mode.is_some() {
        tracing::error!("Cannot use `mode` to set file permissions on windows workers");
    }

    file.write_all(content.as_bytes())?;
    file.flush()?;
    Ok(normalized_full_path)
}

pub async fn reload_custom_tags_setting(db: &DB) -> error::Result<()> {
    let q =
        crate::global_settings::load_value_from_global_settings(db, CUSTOM_TAGS_SETTING).await?;
    apply_custom_tags_setting(q);
    Ok(())
}

/// The half of [`reload_custom_tags_setting`] after the read, so a batched settings pass can
/// apply a value it already fetched.
pub fn apply_custom_tags_setting(q: Option<serde_json::Value>) {
    let tags = if let Some(q) = q {
        if let Ok(v) = serde_json::from_value::<Vec<String>>(q.clone()) {
            v
        } else {
            tracing::error!(
                "Could not parse custom tags setting as vec of strings, found: {:#?}",
                &q
            );
            vec![]
        }
    } else {
        CUSTOM_TAGS.clone()
    };

    let custom_tags = CustomTags::from(tags);

    tracing::info!(
        "Loaded setting custom_tags, common: {:?}, per-workspace: {:?}",
        custom_tags.global,
        custom_tags.specific,
    );

    CUSTOM_TAGS_PER_WORKSPACE.store(std::sync::Arc::new(custom_tags.clone()));
    ALL_TAGS.store(std::sync::Arc::new(
        [
            custom_tags.global.clone(),
            custom_tags
                .specific
                .keys()
                .map(|x| x.to_string())
                .collect_vec(),
        ]
        .concat(),
    ));
}

#[cfg(not(windows))]
fn parse_file<T: FromStr>(path: &str) -> Option<T> {
    std::process::Command::new("cat")
        .args([path])
        .output()
        .ok()
        .map(|o| {
            String::from_utf8_lossy(&o.stdout)
                .to_string()
                .trim()
                .parse::<T>()
                .ok()
        })
        .flatten()
}

#[annotations("#")]
pub struct RubyAnnotations {
    pub verbose: bool,
}

#[annotations("#")]
pub struct RlangAnnotations {
    pub renv_verbose: bool,
    pub renv_install_verbose: bool,
    pub sandbox: bool,
}

#[annotations("#")]
pub struct PythonAnnotations {
    pub no_cache: bool,
    pub no_postinstall: bool,
    pub py_select_latest: bool,
    pub skip_result_postprocessing: bool,
    pub py310: bool,
    pub py311: bool,
    pub py312: bool,
    pub py313: bool,
    pub sandbox: bool,
}

#[annotations("//")]
pub struct GoAnnotations {
    pub go1_22_compat: bool,
}

#[annotations("//")]
pub struct TypeScriptAnnotations {
    pub npm: bool,
    pub nodejs: bool,
    pub native: bool,
    pub nobundling: bool,
    pub sandbox: bool,
}

#[annotations("--")]
pub struct SqlAnnotations {
    pub return_last_result: bool, // deprecated, use result_collection instead
    pub result_collection: SqlResultCollectionStrategy,
    pub prepare: bool, // Used to prepare datatable queries without executing
    // Emit {columns: [{name, oid, type_name}], rows: [[text|null]]} from the
    // last statement instead of the default `[{col: val}]` JSON shape. Used by
    // `wmill datatable serve` to map Postgres results onto the wire protocol
    // without re-stringifying every JSON value.
    pub raw_output: bool,
}

#[annotations("#")]
pub struct BashAnnotations {
    pub docker: bool,
    pub sandbox: bool,
}

impl BashAnnotations {
    /// If the script declares `# sandbox <image>` (an image ref after the sandbox
    /// annotation), returns that image ref. This selects the daemonless, sandboxed
    /// container runtime: extract the image's rootfs and run it inside the job's
    /// nsjail sandbox.
    ///
    /// A bare `# sandbox` (no image argument) returns `None` and keeps the plain
    /// nsjail-sandboxed-bash behavior (the `sandbox` boolean modifier). `# docker`
    /// is unaffected and keeps the legacy v1 (dind/daemon) path.
    pub fn sandbox_image(code: &str) -> Option<String> {
        for line in code.lines() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            // Mirror the annotation parser: stop at the first non-comment line.
            if !line.starts_with('#') {
                break;
            }
            let mut tokens = line[1..].split_whitespace();
            if tokens.next() == Some("sandbox") {
                // `# sandbox <image>` -> container; bare `# sandbox` -> nsjail bash.
                if let Some(image) = tokens.next() {
                    return Some(image.to_string());
                }
            }
        }
        None
    }

    /// If the script declares `#ssh <resource_path>` (a resource path after the
    /// ssh annotation), returns that path. This reroutes execution to a remote
    /// host over SSH (enterprise feature): the script runs on the host described
    /// by the `ssh_target` resource at `<resource_path>` instead of on the worker.
    /// The `#ssh $<arg_name>` form returns the `$`-prefixed token verbatim; the
    /// executor resolves it from the job argument of that name at run time.
    ///
    /// Mirrors `sandbox_image`: only leading comment lines are scanned, stopping
    /// at the first non-comment line. A bare `#ssh` with no path returns `None`.
    /// Only an exact `#ssh <target>` line triggers the reroute — the target must
    /// be a resource path (`u/...`/`f/...`) or a `$arg` reference (a valid
    /// identifier), with nothing else on the line — so prose comments like
    /// `# ssh into the box`, `# ssh tunnel/proxy setup is below` or
    /// `# ssh $HOST manually first` never match.
    pub fn ssh_target(code: &str) -> Option<String> {
        for line in code.lines() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            if !line.starts_with('#') {
                break;
            }
            let mut tokens = line[1..].split_whitespace();
            if tokens.next() == Some("ssh") {
                if let Some(path) = tokens.next() {
                    let is_path = path.starts_with("u/") || path.starts_with("f/");
                    let is_arg = path.strip_prefix('$').is_some_and(|a| {
                        !a.is_empty()
                            && !a.starts_with(|c: char| c.is_ascii_digit())
                            && a.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
                    });
                    if (is_path || is_arg) && tokens.next().is_none() {
                        return Some(path.to_string());
                    }
                }
            }
        }
        None
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum SqlResultCollectionStrategy {
    LastStatementAllRows,
    LastStatementFirstRow,
    LastStatementAllRowsScalar,
    LastStatementFirstRowScalar,
    AllStatementsAllRows,
    AllStatementsFirstRow,
    AllStatementsAllRowsScalar,
    AllStatementsFirstRowScalar,
    Legacy,
}

impl SqlResultCollectionStrategy {
    pub fn parse(s: &str) -> Self {
        use SqlResultCollectionStrategy::*;
        match s {
            "last_statement_all_rows" => LastStatementAllRows,
            "last_statement_first_row" => LastStatementFirstRow,
            "last_statement_all_rows_scalar" => LastStatementAllRowsScalar,
            "last_statement_first_row_scalar" => LastStatementFirstRowScalar,
            "all_statements_all_rows" => AllStatementsAllRows,
            "all_statements_first_row" => AllStatementsFirstRow,
            "all_statements_all_rows_scalar" => AllStatementsAllRowsScalar,
            "all_statements_first_row_scalar" => AllStatementsFirstRowScalar,
            "legacy" => Legacy,
            _ => SqlResultCollectionStrategy::default(),
        }
    }

    pub fn collect_last_statement_only(&self, query_count: usize) -> bool {
        use SqlResultCollectionStrategy::*;
        match self {
            LastStatementAllRows
            | LastStatementFirstRow
            | LastStatementFirstRowScalar
            | LastStatementAllRowsScalar => true,
            Legacy => query_count == 1,
            _ => false,
        }
    }
    pub fn collect_first_row_only(&self) -> bool {
        use SqlResultCollectionStrategy::*;
        match self {
            LastStatementFirstRow
            | LastStatementFirstRowScalar
            | AllStatementsFirstRow
            | AllStatementsFirstRowScalar => true,
            _ => false,
        }
    }
    pub fn collect_scalar(&self) -> bool {
        use SqlResultCollectionStrategy::*;
        match self {
            LastStatementFirstRowScalar
            | AllStatementsFirstRowScalar
            | LastStatementAllRowsScalar
            | AllStatementsAllRowsScalar => true,
            _ => false,
        }
    }

    // This function transforms the shape (e.g Row[][] -> Row)
    // It is the responsibility of the executor to avoid fetching unnecessary statements/rows
    pub fn collect(
        &self,
        values: Vec<Vec<Box<serde_json::value::RawValue>>>,
    ) -> error::Result<Box<serde_json::value::RawValue>> {
        let null = || serde_json::value::RawValue::from_string("null".to_string()).unwrap();

        let values = if self.collect_last_statement_only(values.len()) {
            values.into_iter().rev().take(1).collect()
        } else {
            values
        };

        let values = if self.collect_first_row_only() {
            values
                .into_iter()
                .map(|rows| rows.into_iter().take(1).collect())
                .collect()
        } else {
            values
        };

        let values = if self.collect_scalar() {
            values
                .into_iter()
                .map(|rows| {
                    rows.into_iter()
                        .map(|row| {
                            // Take the first value in the object
                            let record =
                                match serde_json::from_str(row.get()) {
                                    Ok(serde_json::Value::Object(record)) => record,
                                    Ok(_) => return Err(error::Error::ExecutionErr(
                                        "Could not collect sql scalar value from non-object row"
                                            .to_string(),
                                    )),
                                    Err(e) => {
                                        return Err(error::Error::ExecutionErr(format!(
                                    "Could not collect sql scalar value (failed to parse row): {}",
                                    e
                                )))
                                    }
                                };
                            let Some((_, value)) = record.iter().next() else {
                                return Err(error::Error::ExecutionErr(
                                    "Could not collect sql scalar value from empty row".to_string(),
                                ));
                            };
                            Ok(serde_json::value::RawValue::from_string(
                                serde_json::to_string(value).map_err(to_anyhow)?,
                            )
                            .map_err(to_anyhow)?)
                        })
                        .collect::<error::Result<Vec<_>>>()
                })
                .collect::<error::Result<Vec<_>>>()?
        } else {
            values
        };

        match (
            self.collect_last_statement_only(values.len()),
            self.collect_first_row_only(),
        ) {
            (true, true) => {
                match values
                    .into_iter()
                    .last()
                    .map(|rows| rows.into_iter().next())
                {
                    Some(Some(row)) => Ok(row.clone()),
                    _ => Ok(null()),
                }
            }
            (true, false) => match values.into_iter().last() {
                Some(rows) => Ok(merge_raw_values_to_array(rows.as_slice())),
                None => Ok(null()),
            },
            (false, true) => {
                let values = values
                    .into_iter()
                    .map(|rows| rows.into_iter().next().unwrap_or_else(null))
                    .collect::<Vec<_>>();
                Ok(merge_raw_values_to_array(values.as_slice()))
            }
            (false, false) => Ok(merge_nested_raw_values_to_array(
                values.iter().map(|x| x.iter()),
            )),
        }
    }
}

impl Default for SqlResultCollectionStrategy {
    fn default() -> Self {
        SqlResultCollectionStrategy::LastStatementAllRows
    }
}

/// length = 5
/// value  = "foo"
/// output = "foo  "
///           12345
pub fn pad_string(value: &str, total_length: usize) -> String {
    if value.len() >= total_length {
        value.to_string() // Return the original string if it's already long enough
    } else {
        let padding_needed = total_length - value.len();
        format!("{value}{}", " ".repeat(padding_needed)) // Pad with spaces
    }
}
pub fn copy_dir_recursively(src: &Path, dst: &Path) -> error::Result<()> {
    if !dst.exists() {
        fs::create_dir_all(dst)?;
    }

    tracing::debug!("Copying recursively from {:?} to {:?}", src, dst);

    for entry in fs::read_dir(src)? {
        let entry = entry?;
        let src_path = entry.path();
        let dst_path = dst.join(entry.file_name());

        if src_path.is_dir() && !src_path.is_symlink() {
            copy_dir_recursively(&src_path, &dst_path)?;
        } else {
            fs::copy(&src_path, &dst_path)?;
        }
    }

    tracing::debug!("Finished copying recursively from {:?} to {:?}", src, dst);

    Ok(())
}

/// Write `bytes` to `final_path` atomically: write to a sibling temp file then
/// `rename` into place. A concurrent reader that gates on `metadata(final_path)`
/// therefore only ever observes a fully-written file (`rename` is atomic on a
/// POSIX same-filesystem path). Safe under a thundering herd: concurrent renames
/// to the same path are last-writer-wins and every writer produces identical
/// bytes. This prevents the cold-load race where a parallel for-loop spawns N
/// `//native` sandboxes that each observe a partially-written bundle.
pub fn atomic_write_file_bytes(
    final_path: &str,
    bytes: &[u8],
    executable: bool,
) -> error::Result<()> {
    #[cfg(not(unix))]
    let _ = executable;
    let tmp_path = format!("{}.tmp.{}", final_path, Uuid::new_v4());
    let write = || -> error::Result<()> {
        let mut file = File::create(&tmp_path)?;
        file.write_all(bytes)?;
        #[cfg(unix)]
        if executable {
            use std::os::unix::fs::PermissionsExt;
            file.set_permissions(std::fs::Permissions::from_mode(0o755))?;
        }
        file.flush()?;
        Ok(())
    };
    if let Err(e) = write() {
        let _ = fs::remove_file(&tmp_path);
        return Err(e);
    }
    if let Err(e) = fs::rename(&tmp_path, final_path) {
        let _ = fs::remove_file(&tmp_path);
        // Content-addressed cache: if the destination now exists, a concurrent
        // publisher already wrote byte-identical content via its own
        // tmp+rename, so the cache is correct. This also covers Windows, where
        // `rename` can refuse to replace a destination that an in-flight reader
        // has open even though that destination is already valid.
        if Path::new(final_path).exists() {
            return Ok(());
        }
        return Err(e.into());
    }
    Ok(())
}

/// Copy `origin` to `final_path` atomically via a sibling temp file + `rename`,
/// preserving `std::fs::copy` permission semantics. Same atomicity guarantee as
/// [`atomic_write_file_bytes`].
pub fn atomic_copy_file(origin: &str, final_path: &str) -> error::Result<()> {
    let tmp_path = format!("{}.tmp.{}", final_path, Uuid::new_v4());
    if let Err(e) = fs::copy(origin, &tmp_path) {
        let _ = fs::remove_file(&tmp_path);
        return Err(e.into());
    }
    if let Err(e) = fs::rename(&tmp_path, final_path) {
        let _ = fs::remove_file(&tmp_path);
        // See `atomic_write_file_bytes`: a content-addressed destination that
        // now exists is already correct (peer publish / Windows open-file).
        if Path::new(final_path).exists() {
            return Ok(());
        }
        return Err(e.into());
    }
    Ok(())
}

/// Publish a fully-populated `tmp_dir` to `final_dir` via a single atomic
/// `rename`, so a concurrent reader gating on `metadata(final_dir)` never
/// observes a half-populated directory: `final_dir` is only ever absent or
/// complete (we never extract/copy in place).
///
/// These dir caches are content-addressed (PHP `vendor/{hash}`, Java deps —
/// the path embeds a hash of the inputs), so every concurrent publisher builds
/// byte-identical content. If `rename` fails and `final_dir` already exists, a
/// peer published the identical content (or, on Windows, `rename` refused to
/// replace an existing directory that is nonetheless already valid): treat the
/// existing dir as the correct cache.
///
/// Deliberately simple — no destroy-then-recreate swap. The accepted residual
/// is that a *stale partial* `final_dir` left by a populate hard-killed on a
/// pre-atomic binary is trusted rather than rebuilt. That is a finite,
/// self-draining migration hazard (post-fix code only ever renames a complete
/// `tmp_dir` into place, so it cannot create that state), it is confined to the
/// PHP/Java dependency caches (never the `//native` path this PR targets), and
/// it clears on cache eviction. A swap that rebuilds it introduces
/// concurrent-publisher edge cases that are a worse trade on a critical path.
pub fn atomic_publish_dir(tmp_dir: &str, final_dir: &str) -> error::Result<()> {
    match fs::rename(tmp_dir, final_dir) {
        Ok(()) => Ok(()),
        Err(e) => {
            let _ = fs::remove_dir_all(tmp_dir);
            if Path::new(final_dir).exists() {
                Ok(())
            } else {
                Err(e.into())
            }
        }
    }
}

#[cfg(all(feature = "enterprise", feature = "parquet"))]
pub async fn extract_tar(tar: bytes::Bytes, folder: &str) -> error::Result<()> {
    use std::time::Instant;

    use bytes::Buf;
    use tokio::fs::{self};

    let start: Instant = Instant::now();
    fs::create_dir_all(&folder).await?;

    let mut ar = tar::Archive::new(tar.reader());

    if let Err(e) = ar.unpack(folder) {
        tracing::info!("Failed to untar to {folder}. Error: {:?}", e);
        fs::remove_dir_all(&folder).await?;
        return Err(error::Error::ExecutionErr(format!(
            "Failed to untar tar {folder}"
        )));
    }
    tracing::info!(
        "Finished extracting tar to {folder}. Took {}ms",
        start.elapsed().as_millis(),
    );
    Ok(())
}
#[cfg(all(feature = "enterprise", feature = "parquet"))]
pub fn write_binary_file(main_path: &str, byts: &mut bytes::Bytes) -> error::Result<()> {
    atomic_write_file_bytes(main_path, byts, true)
}

#[cfg(not(windows))]
fn get_cgroupv2_path() -> Option<String> {
    let cgroup_path: String = parse_file("/proc/self/cgroup")?;

    CGROUP_V2_PATH_RE
        .captures(&cgroup_path)
        .map(|x| format!("/sys/fs/cgroup{}", x.get(1).unwrap().as_str()))
}

#[cfg(not(windows))]
pub fn get_vcpus() -> Option<i64> {
    if Path::new("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").exists() {
        // cgroup v1
        parse_file("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
            .map(|x: i64| if x > 0 { Some(x) } else { None })
            .flatten()
    } else {
        // cgroup v2
        let cgroup_path = get_cgroupv2_path()?;

        let cpu_max_path = format!("{cgroup_path}/cpu.max");

        parse_file(&cpu_max_path)
            .map(|x: String| {
                CGROUP_V2_CPU_RE
                    .captures(&x)
                    .map(|x| x.get(1).unwrap().as_str().parse::<i64>().ok())
                    .flatten()
            })
            .flatten()
            .map(|x| if x > 0 { Some(x) } else { None })
            .flatten()
    }
}

#[cfg(windows)]
pub fn get_vcpus() -> Option<i64> {
    if *LIMIT_WINDOWS_TO_1CU {
        return Some(100000); // 1 vCPU
    }
    let mut sys = System::new();
    sys.refresh_cpu_all();
    (sys.cpus().len() * 100000).try_into().ok()
}

/// The window `get_vcpus`'s quota is spent over, in the same microseconds. Only
/// their ratio is a number of CPUs, and the window is configurable — 100ms is
/// merely its usual value.
#[cfg(not(windows))]
pub fn get_cpu_period() -> Option<i64> {
    if Path::new("/sys/fs/cgroup/cpu/cpu.cfs_period_us").exists() {
        // cgroup v1
        parse_file("/sys/fs/cgroup/cpu/cpu.cfs_period_us")
    } else {
        // cgroup v2: `cpu.max` is "<quota|max> <period>"
        let cgroup_path = get_cgroupv2_path()?;
        parse_file::<String>(&format!("{cgroup_path}/cpu.max"))?
            .split_whitespace()
            .nth(1)?
            .parse()
            .ok()
    }
    .filter(|period| *period > 0)
}

#[cfg(windows)]
pub fn get_cpu_period() -> Option<i64> {
    Some(100000)
}

/// CPUs the process is allowed to run on, ignoring any bandwidth quota — the count
/// Go's `NumCPU` reports. `available_parallelism` cannot stand in for it: that folds
/// the quota in, so a fraction of a CPU makes it report a single-core machine.
#[cfg(not(windows))]
pub fn get_affinity_cpus() -> Option<usize> {
    // "Cpus_allowed_list:\t0-7,16-23"
    let status = parse_file::<String>("/proc/self/status")?;
    let list = status
        .split("Cpus_allowed_list:")
        .nth(1)?
        .lines()
        .next()?
        .trim();

    let cpus = list
        .split(',')
        .map(|range| {
            let (first, last) = range.split_once('-').unwrap_or((range, range));
            let (first, last) = (first.trim().parse::<usize>(), last.trim().parse::<usize>());
            match (first, last) {
                (Ok(first), Ok(last)) if last >= first => Some(last - first + 1),
                _ => None,
            }
        })
        .sum::<Option<usize>>()?;

    (cpus > 0).then_some(cpus)
}

#[cfg(windows)]
pub fn get_affinity_cpus() -> Option<usize> {
    // The 1CU cap is a policy rather than a bandwidth quota, so it is the whole
    // answer here as it is for `get_vcpus` and `get_memory` — a consumer that reads
    // this as the hardware count would raise the worker back above the cap.
    if *LIMIT_WINDOWS_TO_1CU {
        return Some(1);
    }
    let mut sys = System::new();
    sys.refresh_cpu_all();
    Some(sys.cpus().len()).filter(|cpus| *cpus > 0)
}

#[cfg(not(windows))]
fn get_memory_from_meminfo() -> Option<i64> {
    let memory_info = parse_file::<String>("/proc/meminfo")?;
    if memory_info.contains("MemTotal") {
        let memory_total = memory_info
            .split("MemTotal:")
            .nth(1)?
            .split("kB")
            .next()?
            .trim()
            .parse::<i64>()
            .ok()?;
        return Some(memory_total * 1024);
    }
    None
}

#[cfg(not(windows))]
pub fn get_memory() -> Option<i64> {
    let memory_limit: Option<i64> =
        if Path::new("/sys/fs/cgroup/memory/memory.limit_in_bytes").exists() {
            // cgroup v1
            parse_file("/sys/fs/cgroup/memory/memory.limit_in_bytes")
        } else {
            // cgroup v2
            let cgroup_path = get_cgroupv2_path()?;
            let memory_max_path = format!("{cgroup_path}/memory.max");

            parse_file(&memory_max_path)
        };

    // if memory_max is super high, read machine total memory (meminfo)
    if memory_limit
        .as_ref()
        .is_some_and(|x| x > &(1024 * 1024 * 1024 * 1024 * 1024))
    {
        get_memory_from_meminfo()
    } else {
        memory_limit
    }
}

#[cfg(windows)]
pub fn get_memory() -> Option<i64> {
    if *LIMIT_WINDOWS_TO_1CU {
        return Some(2 * 1024 * 1024 * 1024); // 2 GB
    }
    let mut sys = System::new();
    sys.refresh_memory();
    Some(sys.total_memory() as i64)
}

#[cfg(not(windows))]
pub fn get_worker_memory_usage() -> Option<i64> {
    if Path::new("/sys/fs/cgroup/memory/memory.usage_in_bytes").exists() {
        // cgroup v1
        let total_memory_usage: i64 = parse_file("/sys/fs/cgroup/memory/memory.usage_in_bytes")?;

        let inactive_file = parse_file("/sys/fs/cgroup/memory/memory.stat")
            .map(|x: String| {
                CGROUP_V1_INACTIVE_FILE_RE
                    .captures(&x)
                    .map(|x| x.get(1).unwrap().as_str().parse::<i64>().ok())
                    .flatten()
            })
            .flatten()?;

        Some(total_memory_usage - inactive_file)
    } else {
        // cgroup v2
        let cgroup_path = get_cgroupv2_path()?;
        let memory_current_path = format!("{cgroup_path}/memory.current");

        let total_memory_usage: i64 = parse_file(&memory_current_path)?;

        let memory_stat_path = format!("{cgroup_path}/memory.stat");

        let inactive_file = parse_file(&memory_stat_path)
            .map(|x: String| {
                CGROUP_V2_INACTIVE_FILE_RE
                    .captures(&x)
                    .map(|x| x.get(1).unwrap().as_str().parse::<i64>().ok())
                    .flatten()
            })
            .flatten()?;

        Some(total_memory_usage - inactive_file)
    }
}

#[cfg(windows)]
pub fn get_worker_memory_usage() -> Option<i64> {
    let mut sys = System::new();
    sys.refresh_memory();
    Some(sys.used_memory() as i64)
}

pub fn get_windmill_memory_usage() -> Option<i64> {
    #[cfg(all(not(target_env = "msvc"), feature = "jemalloc"))]
    {
        match tikv_jemalloc_ctl::epoch::advance() {
            Ok(_) => match tikv_jemalloc_ctl::stats::resident::read() {
                Ok(resident) => i64::try_from(resident).ok(),
                Err(e) => {
                    tracing::error!("jemalloc resident memory read failed: {:?}", e);
                    None
                }
            },
            Err(e) => {
                tracing::error!("jemalloc epoch advance failed: {:?}", e);
                None
            }
        }
    }

    #[cfg(any(target_env = "msvc", not(feature = "jemalloc")))]
    {
        None
    }
}

#[derive(Serialize, Deserialize)]
pub enum PingType {
    Initial,
    MainLoop,
    Job,
    InitScript,
}
#[derive(Serialize, Deserialize)]
pub struct Ping {
    pub last_job_executed: Option<Uuid>,
    pub last_job_workspace_id: Option<String>,
    pub worker_instance: Option<String>,
    pub ip: Option<String>,
    pub tags: Option<Vec<String>>,
    pub dw: Option<String>,
    pub dws: Option<Vec<String>>,
    pub version: Option<String>,
    pub vcpus: Option<i64>,
    pub memory: Option<i64>,
    pub memory_usage: Option<i64>,
    pub wm_memory_usage: Option<i64>,
    pub jobs_executed: Option<i32>,
    pub occupancy_rate: Option<f32>,
    pub occupancy_rate_15s: Option<f32>,
    pub occupancy_rate_5m: Option<f32>,
    pub occupancy_rate_30m: Option<f32>,
    pub job_isolation: Option<String>,
    pub native_mode: Option<bool>,
    pub ping_type: PingType,
}
pub async fn update_ping_http(
    insert_ping: Ping,
    worker_name: &str,
    worker_group: &str,
    db: &DB,
) -> anyhow::Result<()> {
    // tracing::info!("update ping: {}", insert_ping.tags.join(","));
    match insert_ping.ping_type {
        PingType::MainLoop => {
            update_worker_ping_main_loop_query(
                worker_name,
                insert_ping.tags.unwrap_or_default().as_slice(),
                insert_ping.vcpus,
                insert_ping.memory,
                insert_ping.jobs_executed,
                insert_ping.occupancy_rate,
                insert_ping.memory_usage,
                insert_ping.wm_memory_usage,
                insert_ping.occupancy_rate_15s,
                insert_ping.occupancy_rate_5m,
                insert_ping.occupancy_rate_30m,
                insert_ping.native_mode.unwrap_or(false),
                insert_ping.ip.as_deref(),
                db,
            )
            .await?
        }
        PingType::Initial => {
            if insert_ping.worker_instance.is_none() || insert_ping.version.is_none() {
                return Err(anyhow::anyhow!("Worker instance and version are required"));
            }

            insert_ping_query(
                &insert_ping.worker_instance.unwrap(),
                &worker_name,
                worker_group,
                // An agent worker sends the sentinel rather than nothing, to stay acceptable to
                // servers that still require an IP here; both mean "not resolved yet".
                insert_ping.ip.as_deref().filter(|ip| *ip != UNKNOWN_IP),
                insert_ping.tags.unwrap_or_default().as_slice(),
                insert_ping.dw,
                insert_ping.dws.as_deref(),
                &insert_ping.version.unwrap(),
                insert_ping.vcpus,
                insert_ping.memory,
                insert_ping.job_isolation,
                insert_ping.native_mode.unwrap_or(false),
                db,
            )
            .await?;
        }
        PingType::Job => {
            update_worker_ping_from_job_query(
                &insert_ping.last_job_executed.unwrap_or_default(),
                &insert_ping.last_job_workspace_id.unwrap_or_default(),
                worker_name,
                insert_ping.memory_usage,
                insert_ping.wm_memory_usage,
                insert_ping.occupancy_rate,
                insert_ping.occupancy_rate_15s,
                insert_ping.occupancy_rate_5m,
                insert_ping.occupancy_rate_30m,
                insert_ping.job_isolation,
                db,
            )
            .await?;
        }
        PingType::InitScript => {
            update_ping_for_failed_init_script_query(
                worker_name,
                insert_ping.last_job_executed.unwrap_or_default(),
                db,
            )
            .await?
        }
    }
    Ok(())
}

#[derive(Debug, Serialize, Deserialize)]
pub struct JobCancelled {
    pub canceled_by: String,
    pub reason: String,
}

pub async fn set_job_cancelled_query(
    job_id: Uuid,
    db: &DB,
    canceled_by: &str,
    reason: &str,
) -> anyhow::Result<()> {
    sqlx::query!(
        "UPDATE v2_job_queue
    SET canceled_by = $1
      , canceled_reason = $2
WHERE id = $3",
        canceled_by,
        reason,
        job_id
    )
    .execute(db)
    .await?;
    Ok(())
}

pub async fn update_ping_for_failed_init_script_query(
    worker_name: &str,
    last_job_id: Uuid,
    db: &DB,
) -> anyhow::Result<()> {
    sqlx::query!(
        "UPDATE worker_ping SET
ping_at = now(),
jobs_executed = 1,
current_job_id = $1,
current_job_workspace_id = 'admins'
WHERE worker = $2",
        last_job_id,
        worker_name
    )
    .execute(db)
    .await?;
    Ok(())
}

pub async fn fetch_flow_node_query(
    db: &DB,
    id: i64,
    loc: &'static Location<'_>,
) -> error::Result<RawNode> {
    let r = sqlx::query!(
        "SELECT \
            code AS \"raw_code: String\", \
            lock AS \"raw_lock: String\", \
            flow AS \"raw_flow: Json<Box<RawValue>>\" \
        FROM flow_node WHERE id = $1 LIMIT 1",
        id,
    )
    .fetch_optional(db)
    .await
    .map_err(Into::into)
    .and_then(unwrap_or_error(loc, "Flow node", id))
    .map(|r| RawNode {
        raw_code: r.raw_code,
        raw_lock: r.raw_lock,
        raw_flow: r.raw_flow.map(|Json(raw_flow)| raw_flow),
    })?;
    Ok(r)
}

pub async fn fetch_raw_script_from_app_query(
    db: &DB,
    id: i64,
    loc: &'static Location<'_>,
) -> error::Result<RawScript> {
    sqlx::query!(
        "SELECT lock, code FROM app_script WHERE id = $1 LIMIT 1",
        id,
    )
    .fetch_optional(db)
    .await
    .map_err(Into::into)
    .and_then(unwrap_or_error(&loc, "Application script", id))
    .map(|r| RawScript { content: r.code, lock: r.lock, meta: None, modules: None })
}

/// Returns the number of jobs the row already accounted for: non-zero when a worker of the
/// same name pinged before, i.e. when this process is a restart of an earlier one (see
/// [`crate::utils::resolve_worker_suffix`]) and its counter is meant to keep climbing.
///
/// Everything else describing the process is overwritten on such a restart — a row left
/// reporting the version or the isolation mode of the process that died would, for
/// `wm_version`, hold the instance-wide `MIN_VERSION` back forever, and one still naming the
/// job that process was killed mid-way through skews the zombie/OOM diagnostics that read it.
/// `started_at` and `jobs_executed` are the only two columns carried over, being the
/// continuity itself — plus `ip` for as long as `ip` is `None`, which means the external IP
/// lookup has not resolved yet and the predecessor's address is still the best guess.
pub async fn insert_ping_query(
    worker_instance: &str,
    worker_name: &str,
    worker_group: &str,
    ip: Option<&str>,
    tags: &[String],
    dw: Option<String>,
    dws: Option<&[String]>,
    version: &str,
    vcpus: Option<i64>,
    memory: Option<i64>,
    job_isolation: Option<String>,
    native_mode: bool,
    db: &DB,
) -> anyhow::Result<i32> {
    // A NULL `ip` means the external IP lookup is still in flight; a later ping fills it in, and
    // meanwhile the value a previous process wrote to a reclaimed row is the best guess we have. A
    // lookup that has failed reports `external_ip::UNRETRIEVABLE_IP`, which does overwrite it. The
    // literal below must stay equal to `external_ip::UNKNOWN_IP`.
    let previous_jobs_executed = sqlx::query_scalar!(
        "INSERT INTO worker_ping (worker_instance, worker, ip, custom_tags, worker_group, dedicated_worker, dedicated_workers, wm_version, vcpus, memory, job_isolation, native_mode) VALUES ($1, $2, COALESCE($3, 'NO IP'), $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT (worker)
        DO UPDATE set ping_at = now(), worker_instance = EXCLUDED.worker_instance, ip = COALESCE($3, worker_ping.ip), custom_tags = EXCLUDED.custom_tags, worker_group = EXCLUDED.worker_group, dedicated_worker = EXCLUDED.dedicated_worker, dedicated_workers = EXCLUDED.dedicated_workers, wm_version = EXCLUDED.wm_version, vcpus = COALESCE(EXCLUDED.vcpus, worker_ping.vcpus), memory = COALESCE(EXCLUDED.memory, worker_ping.memory), job_isolation = EXCLUDED.job_isolation, native_mode = EXCLUDED.native_mode, current_job_id = NULL, current_job_workspace_id = NULL
        RETURNING jobs_executed",
        worker_instance,
        worker_name,
        ip,
        tags,
        worker_group,
        dw,
        dws,
        version,
        vcpus,
        memory,
        job_isolation.as_deref(),
        native_mode,
        )
        .fetch_one(db)
        .await?;
    Ok(previous_jobs_executed)
}

pub async fn update_worker_ping_from_job_query(
    job_id: &Uuid,
    w_id: &str,
    worker_name: &str,
    memory_usage: Option<i64>,
    wm_memory_usage: Option<i64>,
    occupancy_rate: Option<f32>,
    occupancy_rate_15s: Option<f32>,
    occupancy_rate_5m: Option<f32>,
    occupancy_rate_30m: Option<f32>,
    job_isolation: Option<String>,
    db: &DB,
) -> anyhow::Result<()> {
    sqlx::query!(
        "UPDATE worker_ping SET ping_at = now(), current_job_id = $1, current_job_workspace_id = $2, memory_usage = $3, wm_memory_usage = $4,
        occupancy_rate = $6, occupancy_rate_15s = $7, occupancy_rate_5m = $8, occupancy_rate_30m = $9, job_isolation = $10 WHERE worker = $5",
            job_id,
        w_id,
        memory_usage,
        wm_memory_usage,
        worker_name,
        occupancy_rate,
        occupancy_rate_15s,
        occupancy_rate_5m,
        occupancy_rate_30m,
        job_isolation,
    )
    .execute(db)
    .await?;
    Ok(())
}

pub async fn update_job_ping_query(
    job_id: &Uuid,
    db: &DB,
    mem_peak: Option<i32>,
) -> anyhow::Result<PingJobStatusResponse> {
    let ro = sqlx::query!(
        "UPDATE v2_job_runtime r SET
        memory_peak = $1,
        ping = now()
    FROM v2_job_queue q
    WHERE r.id = $2 AND q.id = r.id
    RETURNING canceled_by, canceled_reason",
        mem_peak,
        job_id
    )
    .map(|x| PingJobStatusResponse {
        canceled_by: x.canceled_by,
        canceled_reason: x.canceled_reason,
        already_completed: false,
    })
    .fetch_optional(db)
    .await;

    // TODO: add memory metrics to memory time series

    if let Ok(r) = ro {
        if let Some(i) = r {
            Ok(i)
        } else {
            Ok(PingJobStatusResponse {
                canceled_by: None,
                canceled_reason: None,
                already_completed: true,
            })
        }
    } else {
        Err(to_anyhow(ro.unwrap_err()))
    }
}

pub async fn update_worker_ping_main_loop_query(
    worker_name: &str,
    tags: &[String],
    vcpus: Option<i64>,
    memory: Option<i64>,
    jobs_executed: Option<i32>,
    occupancy_rate: Option<f32>,
    memory_usage: Option<i64>,
    wm_memory_usage: Option<i64>,
    occupancy_rate_15s: Option<f32>,
    occupancy_rate_5m: Option<f32>,
    occupancy_rate_30m: Option<f32>,
    native_mode: bool,
    ip: Option<&str>,
    db: &DB,
) -> anyhow::Result<()> {
    timeout(Duration::from_secs(10), sqlx::query!(
        "UPDATE worker_ping SET ping_at = now(), jobs_executed = $1, custom_tags = $2,
         occupancy_rate = $3, memory_usage = $4, wm_memory_usage = $5, vcpus = COALESCE($7, vcpus),
         memory = COALESCE($8, memory), occupancy_rate_15s = $9, occupancy_rate_5m = $10, occupancy_rate_30m = $11, native_mode = $12, ip = COALESCE($13, ip) WHERE worker = $6",
        jobs_executed,
        tags,
        occupancy_rate,
        memory_usage,
        wm_memory_usage,
        worker_name,
        vcpus,
        memory,
        occupancy_rate_15s,
        occupancy_rate_5m,
        occupancy_rate_30m,
        native_mode,
        ip,
    )
        .execute(db))
    .await??;
    Ok(())
}

// "UPDATE worker_ping SET ping_at = now(), jobs_executed = $1, custom_tags = $2,
// occupancy_rate = $3, memory_usage = $4, wm_memory_usage = $5, vcpus = COALESCE($7, vcpus),
// memory = COALESCE($8, memory), occupancy_rate_15s = $9, occupancy_rate_5m = $10, occupancy_rate_30m = $11 WHERE worker = $6",

const MAX_TAG_LEN: usize = 50;
const HASH_SUFFIX_LEN: usize = 16;

pub fn dedicated_worker_tag(workspace_id: &str, path: &str) -> String {
    let full_tag = format!("{}:{}", workspace_id, path);
    if full_tag.len() <= MAX_TAG_LEN {
        return full_tag;
    }
    let hash = <sha2::Sha256 as sha2::Digest>::digest(full_tag.as_bytes());
    let hex_hash = hex::encode(hash);
    let prefix_len = MAX_TAG_LEN - 1 - HASH_SUFFIX_LEN;
    format!(
        "{}#{}",
        &full_tag[..prefix_len],
        &hex_hash[..HASH_SUFFIX_LEN]
    )
}

/// Configuration for a runner group — a single long-lived subprocess
/// that can execute multiple scripts sharing the same workspace dependency.
/// Auto-detected from script content annotations at worker startup.
#[derive(Clone, PartialEq, Debug)]
pub struct RunnerGroupConfig {
    pub workspace_id: String,
    pub dep_name: String,
    pub language: String,
}

pub async fn load_worker_config(
    db: &DB,
    killpill_tx: KillpillSender,
) -> error::Result<WorkerConfig> {
    tracing::info!("Loading config from WORKER_GROUP: {}", *WORKER_GROUP);
    let mut config: WorkerConfigOpt = sqlx::query_scalar!(
        "SELECT config FROM config WHERE name = $1",
        format!("worker__{}", *WORKER_GROUP)
    )
    .fetch_optional(db)
    .await?
    .flatten()
    .map(|x| serde_json::from_value(x).ok())
    .flatten()
    .unwrap_or_default();

    if config.dedicated_worker.is_none() {
        let dw = std::env::var("DEDICATED_WORKER").ok();
        if dw.is_some() {
            tracing::info!(
                "DEDICATED_WORKER set from env variable: {}",
                dw.as_ref().unwrap()
            );
            config.dedicated_worker = dw;
        }
    } else {
        tracing::info!(
            "DEDICATED_WORKER set from config: {}",
            config.dedicated_worker.as_ref().unwrap()
        );
    }
    let dedicated_worker = config
        .dedicated_worker
        .map(|x| {
            let splitted = x.split(':').to_owned().collect_vec();
            if splitted.len() != 2 {
                killpill_tx.send();
                return Err(anyhow::anyhow!(
                    "Invalid dedicated_worker format. Got {x}, expects <workspace_id>:<path>"
                ));
            } else {
                let workspace = splitted[0];
                let script_path = splitted[1];
                Ok(WorkspacedPath {
                    workspace_id: workspace.to_string(),
                    path: script_path.to_string(),
                })
            }
        })
        .transpose()?;

    // Parse dedicated_workers (multiple dedicated workers)
    let dedicated_workers = config
        .dedicated_workers
        .map(|workers| {
            workers
                .into_iter()
                .map(|x| {
                    let splitted = x.split(':').to_owned().collect_vec();
                    if splitted.len() != 2 {
                        return Err(anyhow::anyhow!(
                            "Invalid dedicated_workers format. Got {x}, expects <workspace_id>:<path>"
                        ));
                    }
                    let workspace = splitted[0];
                    let script_path = splitted[1];
                    Ok(WorkspacedPath {
                        workspace_id: workspace.to_string(),
                        path: script_path.to_string(),
                    })
                })
                .collect::<Result<Vec<_>, _>>()
        })
        .transpose()?;
    if *WORKER_GROUP == "default" && dedicated_worker.is_none() {
        let mut all_tags = config
            .worker_tags
            .unwrap_or_default()
            .into_iter()
            .chain(
                std::env::var("WORKER_TAGS")
                    .ok()
                    .map(|x| x.split(',').map(|x| x.to_string()).collect_vec())
                    .unwrap_or_default(),
            )
            .sorted()
            .collect_vec();
        all_tags.dedup();
        config.worker_tags = Some(all_tags);
    }

    if let Some(force_worker_tags) = std::env::var("FORCE_WORKER_TAGS")
        .ok()
        .filter(|x| !x.is_empty())
        .map(|x| x.split(',').map(|x| x.to_string()).collect::<Vec<String>>())
    {
        tracing::info!(
            "Detected FORCE_WORKER_TAGS, forcing worker tags to: {:#?}",
            force_worker_tags
        );
        config.worker_tags = Some(force_worker_tags);
    }

    // set worker_tags using default if none. If priority tags is set, compute the sorted priority tags as well
    let worker_tags = config
        .worker_tags
        .or_else(|| {
            // Check for multiple dedicated workers
            if let Some(ref dws) = dedicated_workers.as_ref() {
                let mut dedi_tags: Vec<String> = dws
                    .iter()
                    .map(|dw| dedicated_worker_tag(&dw.workspace_id, &dw.path))
                    .collect();
                if std::env::var("ADD_FLOW_TAG").is_ok() {
                    dedi_tags.push("flow".to_string());
                }
                Some(dedi_tags)
            } else if let Some(ref dedicated_worker) = dedicated_worker.as_ref() {
                // Fallback to single dedicated worker for backward compatibility
                let mut dedi_tags = vec![dedicated_worker_tag(
                    &dedicated_worker.workspace_id,
                    &dedicated_worker.path,
                )];
                if std::env::var("ADD_FLOW_TAG").is_ok() {
                    dedi_tags.push("flow".to_string());
                }
                Some(dedi_tags)
            } else {
                std::env::var("WORKER_TAGS")
                    .ok()
                    .map(|x| x.split(',').map(|x| x.to_string()).collect())
            }
        })
        .unwrap_or_else(|| DEFAULT_TAGS.clone());

    let mut priority_tags_sorted: Vec<PriorityTags> = Vec::new();
    let priority_tags_map = config.priority_tags.unwrap_or_else(HashMap::new);
    if priority_tags_map.len() > 0 {
        let mut all_tags_set: HashSet<String> = HashSet::from_iter(worker_tags.clone());

        let mut tags_by_priority: HashMap<u8, Vec<String>> = HashMap::new();

        for (tag, priority) in priority_tags_map.iter() {
            if *priority == 0 {
                // ignore tags with no priority as they will be added at the end from the `all_tags` set
                continue;
            }
            match tags_by_priority.get_mut(priority) {
                Some(tags) => {
                    tags.push(tag.clone());
                }
                None => {
                    let mut t: Vec<String> = Vec::new();
                    t.push(tag.clone());
                    tags_by_priority.insert(*priority, t);
                }
            };
            all_tags_set.remove(tag);
        }
        priority_tags_sorted = tags_by_priority
            .iter()
            .map(|(priority, tags)| PriorityTags { priority: priority.clone(), tags: tags.clone() })
            .collect();
        priority_tags_sorted.push(PriorityTags { priority: 0, tags: Vec::from_iter(all_tags_set) }); // push the tags that were not listed as high priority with a priority = 0
        priority_tags_sorted.sort_by_key(|elt| Reverse(elt.priority)); // sort by priority DESC
    } else {
        // if no priority is used, push all tags with a priority to 0
        priority_tags_sorted.push(PriorityTags { priority: 0, tags: worker_tags.clone() });
    }
    tracing::debug!("Custom tags priority set: {:?}", priority_tags_sorted);

    let env_vars_static = config.env_vars_static.unwrap_or_default().clone();
    let resolved_env_vars: HashMap<String, String> = load_env_vars(
        config
            .env_vars_allowlist
            .unwrap_or_default()
            .into_iter()
            .chain(load_whitelist_env_vars_from_env())
            .chain(env_vars_static.keys().map(|x| x.to_string())),
        &env_vars_static,
    );

    let periodic_script_bash = config
        .periodic_script_bash
        .or_else(|| load_periodic_bash_script_from_env())
        .and_then(|x| if x.is_empty() { None } else { Some(x) });
    let periodic_script_interval_seconds = config
        .periodic_script_interval_seconds
        .or_else(|| load_periodic_bash_script_interval_from_env());

    if periodic_script_bash.is_some() {
        if let Some(interval) = periodic_script_interval_seconds {
            if interval < MIN_PERIODIC_SCRIPT_INTERVAL_SECONDS {
                killpill_tx.send();
                return Err(anyhow::anyhow!(
                    "Periodic script interval must be at least {} seconds, got {} seconds",
                    MIN_PERIODIC_SCRIPT_INTERVAL_SECONDS,
                    interval
                )
                .into());
            }
        } else {
            killpill_tx.send();
            return Err(anyhow::anyhow!(
                "Periodic script interval must be specified when periodic script is configured"
            )
            .into());
        }
    }

    let native_mode = is_native_mode_from_env() || config.native_mode.unwrap_or(false);
    NATIVE_MODE_RESOLVED.store(native_mode, std::sync::atomic::Ordering::Relaxed);

    Ok(WorkerConfig {
        worker_tags,
        priority_tags_sorted,
        dedicated_worker,
        dedicated_workers,
        init_bash: config
            .init_bash
            .or_else(|| load_init_bash_from_env())
            .and_then(|x| if x.is_empty() { None } else { Some(x) }),
        periodic_script_bash,
        periodic_script_interval_seconds,
        cache_clear: config.cache_clear,
        pip_local_dependencies: config
            .pip_local_dependencies
            .or_else(|| load_pip_local_dependencies_from_env()),
        additional_python_paths: config
            .additional_python_paths
            .or_else(|| load_additional_python_paths_from_env()),
        env_vars: resolved_env_vars,
        native_mode,
    })
}

pub fn load_init_bash_from_env() -> Option<String> {
    std::env::var("INIT_SCRIPT")
        .ok()
        .and_then(|x| if x.is_empty() { None } else { Some(x) })
}

pub fn load_periodic_bash_script_from_env() -> Option<String> {
    std::env::var("PERIODIC_BASH_SCRIPT")
        .ok()
        .and_then(|x| if x.is_empty() { None } else { Some(x) })
}

pub fn load_periodic_bash_script_interval_from_env() -> Option<u64> {
    std::env::var("PERIODIC_BASH_SCRIPT_INTERVAL_SECONDS")
        .ok()
        .and_then(|x| x.parse::<u64>().ok())
}

pub fn load_pip_local_dependencies_from_env() -> Option<Vec<String>> {
    std::env::var("PIP_LOCAL_DEPENDENCIES")
        .ok()
        .map(|x| x.split(',').map(|x| x.to_string()).collect_vec())
}

pub fn load_additional_python_paths_from_env() -> Option<Vec<String>> {
    std::env::var("ADDITIONAL_PYTHON_PATHS")
        .ok()
        .map(|x| x.split(':').map(|x| x.to_string()).collect_vec())
}

pub fn load_whitelist_env_vars_from_env() -> std::vec::IntoIter<String> {
    std::env::var("WHITELIST_ENVS")
        .ok()
        .map(|x| x.split(',').map(|x| x.to_string()).collect_vec())
        .unwrap_or_default()
        .into_iter()
}

pub fn load_env_vars(
    iter: impl Iterator<Item = String>,
    env_vars_static: &HashMap<String, String>,
) -> HashMap<String, String> {
    iter.sorted()
        .unique()
        .map(|envvar_name| {
            (
                envvar_name.clone(),
                env_vars_static
                    .get::<String>(&envvar_name)
                    .map(|v| v.to_owned())
                    .unwrap_or_else(|| {
                        std::env::var(envvar_name.clone()).unwrap_or("".to_string())
                    }),
            )
        })
        .collect()
}

pub fn error_to_value(err: &error::Error) -> serde_json::Value {
    match err {
        error::Error::JsonErr(err) => err.clone(),
        _ => json!({"message": err.to_string(), "name": err.name()}),
    }
}

#[derive(Clone, PartialEq, Debug)]
pub struct WorkspacedPath {
    pub workspace_id: String,
    pub path: String,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct WorkerConfigOpt {
    pub worker_tags: Option<Vec<String>>,
    pub priority_tags: Option<HashMap<String, u8>>,
    pub dedicated_worker: Option<String>,
    pub dedicated_workers: Option<Vec<String>>,
    pub init_bash: Option<String>,
    pub periodic_script_bash: Option<String>,
    pub periodic_script_interval_seconds: Option<u64>,
    pub cache_clear: Option<u32>,
    pub additional_python_paths: Option<Vec<String>>,
    pub pip_local_dependencies: Option<Vec<String>>,
    pub env_vars_static: Option<HashMap<String, String>>,
    pub env_vars_allowlist: Option<Vec<String>>,
    pub native_mode: Option<bool>,
}

impl Default for WorkerConfigOpt {
    fn default() -> Self {
        Self {
            worker_tags: Default::default(),
            priority_tags: Default::default(),
            dedicated_worker: Default::default(),
            dedicated_workers: Default::default(),
            init_bash: Default::default(),
            periodic_script_bash: Default::default(),
            periodic_script_interval_seconds: Default::default(),
            cache_clear: Default::default(),
            additional_python_paths: Default::default(),
            pip_local_dependencies: Default::default(),
            env_vars_static: Default::default(),
            env_vars_allowlist: Default::default(),
            native_mode: Default::default(),
        }
    }
}

#[derive(PartialEq, Clone)]
pub struct WorkerConfig {
    pub worker_tags: Vec<String>,
    pub priority_tags_sorted: Vec<PriorityTags>,
    pub dedicated_worker: Option<WorkspacedPath>,
    pub dedicated_workers: Option<Vec<WorkspacedPath>>,
    pub init_bash: Option<String>,
    pub periodic_script_bash: Option<String>,
    pub periodic_script_interval_seconds: Option<u64>,
    pub cache_clear: Option<u32>,
    pub additional_python_paths: Option<Vec<String>>,
    pub pip_local_dependencies: Option<Vec<String>>,
    pub env_vars: HashMap<String, String>,
    pub native_mode: bool,
}

impl std::fmt::Debug for WorkerConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "WorkerConfig {{ worker_tags: {:?}, priority_tags_sorted: {:?}, dedicated_worker: {:?}, dedicated_workers: {:?}, init_bash: {:?}, periodic_script_bash: {:?}, periodic_script_interval_seconds: {:?}, cache_clear: {:?}, additional_python_paths: {:?}, pip_local_dependencies: {:?}, env_vars: {:?}, native_mode: {:?} }}",
        self.worker_tags, self.priority_tags_sorted, self.dedicated_worker, self.dedicated_workers, self.init_bash, self.periodic_script_bash, self.periodic_script_interval_seconds, self.cache_clear, self.additional_python_paths, self.pip_local_dependencies, self.env_vars.iter().map(|(k, v)| format!("{}: {}{} ({} chars)", k, &v[..3.min(v.len())], "***", v.len())).collect::<Vec<String>>().join(", "), self.native_mode)
    }
}

#[derive(PartialEq, Debug, Clone)]
pub struct PriorityTags {
    pub priority: u8,
    pub tags: Vec<String>,
}

pub fn to_raw_value<T: Serialize>(result: &T) -> Box<RawValue> {
    serde_json::value::to_raw_value(result)
        .unwrap_or_else(|_| RawValue::from_string("{}".to_string()).unwrap())
}

pub fn to_raw_value_owned(result: serde_json::Value) -> Box<RawValue> {
    serde_json::value::to_raw_value(&result)
        .unwrap_or_else(|_| RawValue::from_string("{}".to_string()).unwrap())
}

pub fn split_python_requirements<T: AsRef<str>>(requirements: T) -> Vec<String> {
    requirements
        .as_ref()
        .lines()
        .filter(|x| !x.trim_start().starts_with("--") && !x.trim().is_empty())
        .map(String::from)
        .collect()
}

#[derive(Eq, PartialEq, Clone, Copy, Default, Debug)]
#[repr(u32)]
pub enum PyVAlias {
    Py310 = 10,
    Py311,
    #[default]
    Py312,
    Py313,
}

impl Into<pep440_rs::Version> for PyVAlias {
    fn into(self) -> pep440_rs::Version {
        pep440_rs::Version::new([self.major() as u64, self as u64])
    }
}

impl Into<u32> for PyVAlias {
    fn into(self) -> u32 {
        self.major() * 100 + self as u32
    }
}

impl PyVAlias {
    pub fn all<T: From<PyVAlias>>() -> Vec<T> {
        use PyVAlias::*;
        vec![Py310.into(), Py311.into(), Py312.into(), Py313.into()]
    }
    // Get MAJOR part of alias. (semver: MAJOR.MINOR.PATCH)
    fn major(&self) -> u32 {
        use PyVAlias::*;
        match self {
            Py310 | Py311 | Py312 | Py313 => 3,
            // Py400 | Py401 => 4
        }
    }

    /// Converts numeric format to alias
    /// Example:
    /// 310u32 (in) -> PyVAlias::Py310 (out)
    pub fn try_from_v1<T: ToString>(numeric: T) -> Option<Self> {
        use PyVAlias::*;
        match numeric.to_string().as_str() {
            "310" => Some(Py310),
            "311" => Some(Py311),
            "312" => Some(Py312),
            "313" => Some(Py313),
            _ => None,
        }
    }
}

/// Parse lockfile for assigned python version.
/// If not found returns None
pub fn try_parse_locked_python_version_from_requirements<S: AsRef<str>>(
    requirements_lines: &[S],
) -> Option<pep440_rs::Version> {
    let parse_version = |s: &str| -> Option<pep440_rs::Version> {
        // Possible inputs:
        // V2:
        // # py: 3.11.0 or #py:3.11.0 or #py: 3.11.0
        //
        // V1:
        // # py311 or #py311
        let version_unparsed = s
            .to_owned()
            // Remove whitespaces. That leaves us with:
            // V2: #py:3.11.0
            // V1: #py311
            //
            // Remove #
            // V2: py:3.11.0
            // V1: py311
            //
            // Remove :
            // V2: py3.11.0
            // V1: py311
            .replace([' ', '#', ':'], "")
            // Remove "py"
            // V2: 3.11.0
            // V1: 311
            .replace("py", "");

        // We will support reading V1 syntax, but it will be overwritten next deploy
        PyVAlias::try_from_v1(&version_unparsed)
            .map(PyVAlias::into)
            .or(pep440_rs::Version::from_str(&version_unparsed)
                .ok()
                .map(pep440_rs::Version::into))
    };

    requirements_lines
        .iter()
        .find(|s| {
            let s = s.as_ref();
            s.starts_with("#py") || s.starts_with("# py")
        })
        .map(S::as_ref)
        .and_then(parse_version)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    fn matcher(id: &str) -> WorkspaceMatcher {
        WorkspaceMatcher { id: id.to_string(), include_forks: false }
    }

    fn fork_matcher(id: &str) -> WorkspaceMatcher {
        WorkspaceMatcher { id: id.to_string(), include_forks: true }
    }

    /// A workspace id chain: the workspace itself, then its fork ancestors nearest-first.
    fn chain(ids: &[&str]) -> Vec<String> {
        ids.iter().map(|s| s.to_string()).collect()
    }

    #[test]
    fn test_parse_job_oom_score_adj() {
        assert_eq!(parse_job_oom_score_adj(Some("300")), 300);
        assert_eq!(parse_job_oom_score_adj(Some(" 0\n")), 0);
        assert_eq!(parse_job_oom_score_adj(None), JOB_OOM_SCORE_ADJ_DEFAULT);
        // Out of range and unparseable both fall back rather than weaken the worker's protection.
        assert_eq!(
            parse_job_oom_score_adj(Some("-500")),
            JOB_OOM_SCORE_ADJ_DEFAULT
        );
        assert_eq!(
            parse_job_oom_score_adj(Some("1001")),
            JOB_OOM_SCORE_ADJ_DEFAULT
        );
        assert_eq!(
            parse_job_oom_score_adj(Some("high")),
            JOB_OOM_SCORE_ADJ_DEFAULT
        );
    }

    #[test]
    fn test_bash_sandbox_image_annotation() {
        // `# sandbox <image>` selects the container runtime and returns the image.
        assert_eq!(
            BashAnnotations::sandbox_image("# sandbox alpine:latest\necho hi"),
            Some("alpine:latest".to_string())
        );
        // Extra whitespace and a leading non-spaced `#` still work.
        assert_eq!(
            BashAnnotations::sandbox_image("#sandbox   python:3.12-slim\n"),
            Some("python:3.12-slim".to_string())
        );
        // A bare `# sandbox` (no image) keeps the nsjail-bash modifier -> None.
        assert_eq!(BashAnnotations::sandbox_image("# sandbox\necho hi"), None);
        // `sandbox` must be its own token, not a prefix.
        assert_eq!(BashAnnotations::sandbox_image("# sandboxed foo"), None);
        // Stops at the first non-comment line (image declared too late is ignored).
        assert_eq!(
            BashAnnotations::sandbox_image("echo hi\n# sandbox alpine"),
            None
        );
        // `# docker` is a different annotation -> not a sandbox image.
        assert_eq!(
            BashAnnotations::sandbox_image("# docker alpine\necho hi"),
            None
        );
    }

    #[test]
    fn test_bash_ssh_target_annotation() {
        // `#ssh <path>` returns the resource path (reroutes to remote execution).
        assert_eq!(
            BashAnnotations::ssh_target("#ssh f/infra/jump_node\nset -e\ndf -h"),
            Some("f/infra/jump_node".to_string())
        );
        // `# ssh <path>` with a space after `#` also works.
        assert_eq!(
            BashAnnotations::ssh_target("# ssh u/me/box\n"),
            Some("u/me/box".to_string())
        );
        // `#ssh $arg` (dynamic target from a job argument) is returned verbatim.
        assert_eq!(
            BashAnnotations::ssh_target("#ssh $jump_host\necho hi"),
            Some("$jump_host".to_string())
        );
        // A bare `#ssh` with no path -> None.
        assert_eq!(BashAnnotations::ssh_target("#ssh\necho hi"), None);
        // `ssh` must be its own token, not a prefix.
        assert_eq!(BashAnnotations::ssh_target("# sshd restart"), None);
        // Prose comments mentioning ssh must not trigger the reroute: the line
        // must be exactly `#ssh <target>` with a `u/`/`f/` path or `$identifier`.
        assert_eq!(
            BashAnnotations::ssh_target("# ssh into the box and restart nginx\necho hi"),
            None
        );
        assert_eq!(
            BashAnnotations::ssh_target("# ssh tunnel/proxy setup is below\necho hi"),
            None
        );
        assert_eq!(
            BashAnnotations::ssh_target("# ssh $HOST manually first\necho hi"),
            None
        );
        // Trailing tokens after a real-looking target also disqualify the line.
        assert_eq!(
            BashAnnotations::ssh_target("#ssh f/infra/box then reboot\necho hi"),
            None
        );
        // ...but a real directive on a later comment line is still found.
        assert_eq!(
            BashAnnotations::ssh_target("# ssh into the box\n#ssh f/infra/box\necho hi"),
            Some("f/infra/box".to_string())
        );
        // Stops at the first non-comment line (declared too late is ignored).
        assert_eq!(
            BashAnnotations::ssh_target("echo hi\n#ssh f/infra/box"),
            None
        );
        // No annotation -> None (normal local bash).
        assert_eq!(BashAnnotations::ssh_target("echo hello"), None);
    }

    #[test]
    fn test_mixed_tags() {
        let input = vec![
            "global".to_string(),
            "feat(ws1+ws2)".to_string(),
            "hotfix(^ws3^ws4)".to_string(),
        ];
        let result = CustomTags::from(input);

        assert_eq!(result.global, vec!["global"]);

        let mut expected = HashMap::new();
        expected.insert(
            "feat".to_string(),
            SpecificTagData {
                tag_type: SpecificTagType::NoneExcept,
                workspaces: vec![matcher("ws1"), matcher("ws2")],
            },
        );
        expected.insert(
            "hotfix".to_string(),
            SpecificTagData {
                tag_type: SpecificTagType::AllExcluding,
                workspaces: vec![matcher("ws3"), matcher("ws4")],
            },
        );

        assert_eq!(result.specific, expected);
    }

    #[test]
    fn test_invalid_specific_tag_format() {
        let input = vec!["invalid(custom+format".to_string()];
        let result = CustomTags::from(input);

        // Regex does not match, so it's treated as a global tag.
        assert_eq!(result.global, vec!["invalid(custom+format"]);
        assert!(result.specific.is_empty());
    }

    #[test]
    fn test_empty_input() {
        let input: Vec<String> = vec![];
        let result = CustomTags::from(input);
        assert!(result.global.is_empty());
        assert!(result.specific.is_empty());
    }

    #[test]
    fn test_custom_tags_from_parses_global_tags_correctly() {
        let input = vec!["frontend".to_string(), "urgent".to_string()];
        let tags = CustomTags::from(input.clone());
        assert_eq!(tags.global, input);
        assert!(tags.specific.is_empty());
    }

    #[test]
    fn test_custom_tags_from_parses_specific_tags_none_except() {
        let input = vec!["urgent(ws1+ws2)".to_string()];
        let tags = CustomTags::from(input.clone());

        assert!(tags.global.is_empty());
        assert_eq!(tags.specific.len(), 1);

        let data = tags.specific.get("urgent").unwrap();
        assert_eq!(data.tag_type, SpecificTagType::NoneExcept);
        assert_eq!(data.workspaces, vec![matcher("ws1"), matcher("ws2")]);
    }

    #[test]
    fn test_custom_tags_from_parses_specific_tags_all_excluding() {
        let input = vec!["legacy(^ws1^ws2)".to_string()];
        let tags = CustomTags::from(input.clone());

        assert!(tags.global.is_empty());
        assert_eq!(tags.specific.len(), 1);

        let data = tags.specific.get("legacy").unwrap();
        assert_eq!(data.tag_type, SpecificTagType::AllExcluding);
        assert_eq!(data.workspaces, vec![matcher("ws1"), matcher("ws2")]);
    }

    #[test]
    fn test_custom_tags_from_ignores_empty_workspace_list_as_global() {
        let input = vec!["foo()".to_string(), "bar(^)".to_string()];
        let tags = CustomTags::from(input.clone());

        assert_eq!(tags.global, input);
        assert!(tags.specific.is_empty());
    }

    #[test]
    fn test_custom_tags_from_filters_specific_tags_by_workspace_none_except() {
        let input = vec!["urgent(ws1+ws2)".to_string()];
        let tags = CustomTags::from(input);

        let output = tags.to_string_vec(Some(&chain(&["ws1"])));
        assert_eq!(output, vec!["urgent"]);

        let output_none = tags.to_string_vec(Some(&chain(&["ws3"])));
        assert!(output_none.is_empty());
    }

    #[test]
    fn test_custom_tags_from_filters_specific_tags_by_workspace_all_excluding() {
        let input = vec!["legacy(^ws1^ws2)".to_string()];
        let tags = CustomTags::from(input);

        let output = tags.to_string_vec(Some(&chain(&["ws3"])));
        assert_eq!(output, vec!["legacy"]);

        let output_excluded = tags.to_string_vec(Some(&chain(&["ws1"])));
        assert!(output_excluded.is_empty());
    }

    #[test]
    fn test_custom_tags_from_reconstructs_all_tags_when_no_filter() {
        let input = vec![
            "foo".to_string(),
            "urgent(ws1+ws2)".to_string(),
            "legacy(^ws1^ws2)".to_string(),
        ];
        let tags = CustomTags::from(input);

        let mut result = tags.to_string_vec(None);
        result.sort();
        assert_eq!(result, vec!["foo", "legacy(^ws1^ws2)", "urgent(ws1+ws2)"]);
    }

    #[test]
    fn test_fork_marker_parses_and_round_trips() {
        let tags = CustomTags::from(vec![
            "urgent(prod*+ws2)".to_string(),
            "legacy(^prod*)".to_string(),
        ]);

        let urgent = tags.specific.get("urgent").unwrap();
        assert_eq!(
            urgent.workspaces,
            vec![fork_matcher("prod"), matcher("ws2")]
        );
        assert!(urgent.is_fork_scoped());

        let legacy = tags.specific.get("legacy").unwrap();
        assert_eq!(legacy.workspaces, vec![fork_matcher("prod")]);

        // The settings editor re-emits what it parsed; dropping `*` here would silently widen
        // an excluding tag / narrow an including one on every save.
        let mut result = tags.to_string_vec(None);
        result.sort();
        assert_eq!(result, vec!["legacy(^prod*)", "urgent(prod*+ws2)"]);
    }

    #[test]
    fn test_fork_marker_extends_none_except_to_forks_only_when_present() {
        let fork = chain(&["wm-fork-x", "prod"]);
        let nested = chain(&["wm-fork-y", "wm-fork-x", "prod"]);

        let marked = CustomTags::from(vec!["urgent(prod*)".to_string()]);
        let marked = marked.specific.get("urgent").unwrap();
        assert!(marked.applies_to_workspace(&chain(&["prod"])));
        assert!(marked.applies_to_workspace(&fork));
        assert!(marked.applies_to_workspace(&nested));
        assert!(!marked.applies_to_workspace(&chain(&["wm-fork-z", "other"])));

        // Without the marker a fork must NOT inherit the parent's tag.
        let unmarked = CustomTags::from(vec!["urgent(prod)".to_string()]);
        let unmarked = unmarked.specific.get("urgent").unwrap();
        assert!(unmarked.applies_to_workspace(&chain(&["prod"])));
        assert!(!unmarked.applies_to_workspace(&fork));
        // Gates the ancestor lookup, so a wrong answer here silently disables the marker.
        assert!(!unmarked.is_fork_scoped());
    }

    #[test]
    fn test_fork_marker_extends_all_excluding_to_forks_only_when_present() {
        let fork = chain(&["wm-fork-x", "prod"]);

        // Pre-existing exclusions keep their exact meaning: only `prod` itself is excluded.
        let unmarked = CustomTags::from(vec!["legacy(^prod)".to_string()]);
        let unmarked = unmarked.specific.get("legacy").unwrap();
        assert!(!unmarked.applies_to_workspace(&chain(&["prod"])));
        assert!(unmarked.applies_to_workspace(&fork));

        let marked = CustomTags::from(vec!["legacy(^prod*)".to_string()]);
        let marked = marked.specific.get("legacy").unwrap();
        assert!(!marked.applies_to_workspace(&chain(&["prod"])));
        assert!(!marked.applies_to_workspace(&fork));
        assert!(marked.applies_to_workspace(&chain(&["other"])));
    }

    #[test]
    fn test_dedicated_worker_tag_short() {
        let tag = dedicated_worker_tag("demo", "u/alice/script");
        assert_eq!(tag, "demo:u/alice/script");
        assert!(tag.len() <= MAX_TAG_LEN);
    }

    #[test]
    fn test_dedicated_worker_tag_exactly_50() {
        // 50 chars exactly should not be hashed
        let workspace = "ws";
        let path = "a".repeat(50 - workspace.len() - 1); // -1 for ':'
        let tag = dedicated_worker_tag(workspace, &path);
        assert_eq!(tag.len(), 50);
        assert!(!tag.contains('#'));
    }

    #[test]
    fn test_dedicated_worker_tag_long_is_hashed() {
        let tag = dedicated_worker_tag(
            "my_workspace",
            "u/engineering/team/automation/critical_workflow_script_v2",
        );
        assert_eq!(tag.len(), MAX_TAG_LEN);
        assert_eq!(tag, "my_workspace:u/engineering/team/a#5bc26db79926d4f0");
    }

    #[test]
    fn test_dedicated_worker_tag_deterministic() {
        let a = dedicated_worker_tag(
            "ws",
            "some/very/long/path/that/exceeds/the/fifty/char/limit/easily",
        );
        assert_eq!(a, "ws:some/very/long/path/that/excee#bbb038d4268a0b41");
        let b = dedicated_worker_tag(
            "ws",
            "some/very/long/path/that/exceeds/the/fifty/char/limit/easily",
        );
        assert_eq!(a, b);
    }

    #[test]
    fn test_dedicated_worker_tag_different_paths_differ() {
        let a = dedicated_worker_tag(
            "ws",
            "some/very/long/path/that/exceeds/the/fifty/char/limit/easily_a",
        );
        let b = dedicated_worker_tag(
            "ws",
            "some/very/long/path/that/exceeds/the/fifty/char/limit/easily_b",
        );
        assert_ne!(a, b);
    }

    #[test]
    fn test_python_sandbox_annotation() {
        let content = "# sandbox\ndef main():\n    pass";
        let annotations = PythonAnnotations::parse(content);
        assert!(annotations.sandbox);
    }

    #[test]
    fn test_sql_raw_output_annotation() {
        let with_flag = SqlAnnotations::parse("-- raw_output\nSELECT 1");
        assert!(with_flag.raw_output);
        assert!(!with_flag.prepare);

        let without_flag = SqlAnnotations::parse("SELECT 1");
        assert!(!without_flag.raw_output);

        let combined = SqlAnnotations::parse("-- prepare\n-- raw_output\nSELECT 1");
        assert!(combined.raw_output);
        assert!(combined.prepare);
    }

    #[test]
    fn test_python_sandbox_annotation_with_other_annotations() {
        let content = "# no_cache\n# sandbox\ndef main():\n    pass";
        let annotations = PythonAnnotations::parse(content);
        assert!(annotations.sandbox);
        assert!(annotations.no_cache);
    }

    #[test]
    fn test_python_no_sandbox_annotation() {
        let content = "# no_cache\ndef main():\n    pass";
        let annotations = PythonAnnotations::parse(content);
        assert!(!annotations.sandbox);
    }

    #[test]
    fn test_typescript_sandbox_annotation() {
        let content = "// sandbox\nexport function main() {}";
        let annotations = TypeScriptAnnotations::parse(content);
        assert!(annotations.sandbox);
    }

    #[test]
    fn test_typescript_sandbox_annotation_with_other_annotations() {
        let content = "// npm\n// sandbox\nexport function main() {}";
        let annotations = TypeScriptAnnotations::parse(content);
        assert!(annotations.sandbox);
        assert!(annotations.npm);
    }

    #[test]
    fn test_typescript_no_sandbox_annotation() {
        let content = "// npm\nexport function main() {}";
        let annotations = TypeScriptAnnotations::parse(content);
        assert!(!annotations.sandbox);
    }

    #[test]
    fn test_python_sandbox_no_space() {
        let content = "#sandbox\ndef main():\n    pass";
        let annotations = PythonAnnotations::parse(content);
        assert!(annotations.sandbox);
    }

    #[test]
    fn test_typescript_sandbox_no_space() {
        let content = "//sandbox\nexport function main() {}";
        let annotations = TypeScriptAnnotations::parse(content);
        assert!(annotations.sandbox);
    }

    // Regression: a parallel for-loop cold-loading a //native bundle spawns
    // many sandboxes that each gate on `metadata(bundle).is_ok()` then read it.
    // The pre-fix non-atomic `File::create + write_all` made the path visible
    // while still empty/truncated, so concurrent readers saw a stub bundle
    // (`module.main is not a function`) or a truncated one (`Unexpected end of
    // input`). With atomic temp+rename publish, a visible path is always a
    // complete file. This test fails against the old non-atomic implementation.
    #[test]
    fn test_atomic_write_file_bytes_concurrent_cold_load() {
        use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
        use std::sync::{Arc, Barrier};

        let dir =
            std::env::temp_dir().join(format!("wm_atomic_write_test_{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(&dir).unwrap();
        let final_path = dir.join("bundle.js");
        let final_path_str = final_path.to_str().unwrap().to_string();

        // Wide payload so a hypothetical non-atomic writer has a large
        // partial-read window for the reader to catch.
        let payload = vec![b'x'; 4 * 1024 * 1024];
        let expected_len = payload.len();

        for _ in 0..12 {
            let _ = std::fs::remove_file(&final_path);

            let stop = Arc::new(AtomicBool::new(false));
            let partial_reads = Arc::new(AtomicUsize::new(0));
            let barrier = Arc::new(Barrier::new(2));

            let reader = {
                let stop = stop.clone();
                let partial_reads = partial_reads.clone();
                let barrier = barrier.clone();
                let path = final_path_str.clone();
                std::thread::spawn(move || {
                    barrier.wait();
                    while !stop.load(Ordering::Relaxed) {
                        if std::fs::metadata(&path).is_ok() {
                            if let Ok(content) = std::fs::read(&path) {
                                if content.len() != expected_len {
                                    partial_reads.fetch_add(1, Ordering::Relaxed);
                                }
                            }
                        }
                        std::thread::yield_now();
                    }
                })
            };

            barrier.wait();
            atomic_write_file_bytes(&final_path_str, &payload, false).unwrap();
            stop.store(true, Ordering::Relaxed);
            reader.join().unwrap();

            assert_eq!(
                partial_reads.load(Ordering::Relaxed),
                0,
                "a concurrent reader observed a partially-written bundle"
            );
            assert_eq!(std::fs::read(&final_path).unwrap().len(), expected_len);

            let leftovers: Vec<_> = std::fs::read_dir(&dir)
                .unwrap()
                .filter_map(|e| e.ok())
                .filter(|e| e.file_name().to_string_lossy().contains(".tmp."))
                .map(|e| e.path())
                .collect();
            assert!(leftovers.is_empty(), "temp files leaked: {:?}", leftovers);
        }

        let _ = std::fs::remove_dir_all(&dir);
    }

    // Under a thundering herd of cold-loads, N populators each build their own
    // temp dir and publish to the same final dir. Every publish must succeed
    // (loser-of-the-race discards its byte-identical copy), the final dir must
    // be complete, and no temp dirs may leak.
    #[test]
    fn test_atomic_publish_dir_thundering_herd() {
        use std::sync::{Arc, Barrier};

        let base =
            std::env::temp_dir().join(format!("wm_atomic_dir_test_{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(&base).unwrap();
        let final_dir = base.join("cache_dir");
        let final_dir_str = final_dir.to_str().unwrap().to_string();

        let n = 8;
        let barrier = Arc::new(Barrier::new(n));
        let mut handles = vec![];
        for _ in 0..n {
            let barrier = barrier.clone();
            let final_dir_str = final_dir_str.clone();
            let base = base.clone();
            handles.push(std::thread::spawn(move || {
                let tmp = base.join(format!("cache_dir.tmp.{}", uuid::Uuid::new_v4()));
                std::fs::create_dir_all(&tmp).unwrap();
                std::fs::write(tmp.join("main.js"), b"export function main() {}").unwrap();
                std::fs::write(tmp.join("meta.txt"), b"v1").unwrap();
                barrier.wait();
                atomic_publish_dir(tmp.to_str().unwrap(), &final_dir_str).unwrap();
            }));
        }
        for h in handles {
            h.join().unwrap();
        }

        assert!(final_dir.join("main.js").is_file());
        assert!(final_dir.join("meta.txt").is_file());
        let leftovers: Vec<_> = std::fs::read_dir(&base)
            .unwrap()
            .filter_map(|e| e.ok())
            .filter(|e| {
                let n = e.file_name();
                let n = n.to_string_lossy();
                n.contains(".tmp.") || n.contains(".bak.")
            })
            .map(|e| e.path())
            .collect();
        assert!(
            leftovers.is_empty(),
            "temp/bak dirs leaked: {:?}",
            leftovers
        );

        let _ = std::fs::remove_dir_all(&base);
    }

    // Contract guard: when `final_dir` ALREADY exists (a complete prior/peer
    // publish — content-addressed, so identical bytes), concurrent publishers
    // must all return Ok via the exists-fallback and must never corrupt or
    // partially-overwrite the existing dir. Covers the preexisting+concurrent
    // case; documents the accepted simple-form behavior (an existing dir is
    // trusted, not rebuilt).
    #[test]
    fn test_atomic_publish_dir_existing_is_trusted_not_corrupted() {
        use std::sync::{Arc, Barrier};

        let base =
            std::env::temp_dir().join(format!("wm_atomic_exist_test_{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(&base).unwrap();
        let final_dir = base.join("cache_dir");
        let final_dir_str = final_dir.to_str().unwrap().to_string();

        // A complete dir already published at the final path.
        std::fs::create_dir_all(&final_dir).unwrap();
        std::fs::write(final_dir.join("main.js"), b"export function main() {}").unwrap();
        std::fs::write(final_dir.join("meta.txt"), b"v1").unwrap();

        let n = 8;
        let barrier = Arc::new(Barrier::new(n));
        let mut handles = vec![];
        for _ in 0..n {
            let barrier = barrier.clone();
            let final_dir_str = final_dir_str.clone();
            let base = base.clone();
            handles.push(std::thread::spawn(move || {
                let tmp = base.join(format!("cache_dir.tmp.{}", uuid::Uuid::new_v4()));
                std::fs::create_dir_all(&tmp).unwrap();
                std::fs::write(tmp.join("main.js"), b"export function main() {}").unwrap();
                std::fs::write(tmp.join("meta.txt"), b"v1").unwrap();
                barrier.wait();
                // Every publisher must succeed (exists-fallback), none error.
                atomic_publish_dir(tmp.to_str().unwrap(), &final_dir_str).unwrap();
            }));
        }
        for h in handles {
            h.join().unwrap();
        }

        // Existing dir intact and complete — never partially overwritten.
        assert_eq!(
            std::fs::read(final_dir.join("main.js")).unwrap(),
            b"export function main() {}"
        );
        assert_eq!(std::fs::read(final_dir.join("meta.txt")).unwrap(), b"v1");
        let leftovers: Vec<_> = std::fs::read_dir(&base)
            .unwrap()
            .filter_map(|e| e.ok())
            .filter(|e| {
                let n = e.file_name();
                let n = n.to_string_lossy();
                n.contains(".tmp.") || n.contains(".bak.")
            })
            .map(|e| e.path())
            .collect();
        assert!(
            leftovers.is_empty(),
            "temp/bak dirs leaked: {:?}",
            leftovers
        );

        let _ = std::fs::remove_dir_all(&base);
    }

    #[test]
    fn test_is_allowed_file_location_allows_plain_relative() {
        let base = std::env::temp_dir().join(format!("wm_allowed_loc_ok_{}", uuid::Uuid::new_v4()));
        let job_dir = base.join("job");
        std::fs::create_dir_all(&job_dir).unwrap();
        let job_dir_str = job_dir.to_str().unwrap();

        let out = is_allowed_file_location(job_dir_str, "repo/sub/playbook.yml").unwrap();
        assert_eq!(out, normalize_path(&job_dir.join("repo/sub/playbook.yml")));

        let _ = std::fs::remove_dir_all(&base);
    }

    #[test]
    fn test_is_allowed_file_location_rejects_parent_and_absolute() {
        let base =
            std::env::temp_dir().join(format!("wm_allowed_loc_esc_{}", uuid::Uuid::new_v4()));
        let job_dir = base.join("job");
        std::fs::create_dir_all(&job_dir).unwrap();
        let job_dir_str = job_dir.to_str().unwrap();

        assert!(is_allowed_file_location(job_dir_str, "../escape").is_err());
        assert!(is_allowed_file_location(job_dir_str, "a/../../escape").is_err());
        assert!(is_allowed_file_location(job_dir_str, "/etc/passwd").is_err());

        let _ = std::fs::remove_dir_all(&base);
    }

    // Regression for GHSA-v934-cvpf-6fjw: a symlink planted inside the job dir
    // (e.g. by an earlier `git_repos` clone) must not let a later target traverse
    // it out of the job dir, even though the lexical path stays "inside".
    #[cfg(unix)]
    #[test]
    fn test_is_allowed_file_location_rejects_symlink_traversal() {
        let base =
            std::env::temp_dir().join(format!("wm_allowed_loc_symlink_{}", uuid::Uuid::new_v4()));
        let job_dir = base.join("job");
        std::fs::create_dir_all(&job_dir).unwrap();
        // Stand-in for the shared cache dir living outside the job dir.
        let outside = base.join("outside");
        std::fs::create_dir_all(&outside).unwrap();
        let job_dir_str = job_dir.to_str().unwrap();

        // Plant `job/repo` -> `../outside`, as a malicious first clone would.
        let planted = job_dir.join("repo");
        std::os::unix::fs::symlink(&outside, &planted).unwrap();

        // Both the symlink itself and any path traversing it are rejected.
        assert!(is_allowed_file_location(job_dir_str, "repo").is_err());
        assert!(is_allowed_file_location(job_dir_str, "repo/payload").is_err());
        assert!(is_allowed_file_location(job_dir_str, "repo/sub/payload").is_err());

        // An in-bounds `..` must not bypass the check: `foo/../repo/payload`
        // normalizes back to `repo/payload` and still traverses the symlink.
        assert!(is_allowed_file_location(job_dir_str, "foo/../repo/payload").is_err());
        std::fs::create_dir(job_dir.join("real")).unwrap();
        assert!(is_allowed_file_location(job_dir_str, "real/../repo/payload").is_err());

        // A dangling symlink (target does not exist yet) is still caught:
        // `symlink_metadata` does not follow the link.
        let dangling = job_dir.join("dangling");
        std::os::unix::fs::symlink(base.join("nonexistent"), &dangling).unwrap();
        assert!(is_allowed_file_location(job_dir_str, "dangling/payload").is_err());

        let _ = std::fs::remove_dir_all(&base);
    }
}