feat: add WM_ROOT_WORKSPACE, the closest dev or prod workspace of a job (#10776)

* feat: add WM_ROOT_WORKSPACE, the closest dev or prod workspace of a job

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HBMJwogo6jJ1P55uvB3YpF

* fix: do not cache a failed root-workspace lookup, and sweep on fork create

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HBMJwogo6jJ1P55uvB3YpF

* fix: shorten the agent-worker root-workspace TTL and pin the sweep wiring

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HBMJwogo6jJ1P55uvB3YpF

* chore: update ee-repo-ref to a2fa58e5301d3865dd06ad73519e20ba7a5af0f0

This commit updates the EE repository reference after PR #736 was merged in windmill-ee-private.

Previous ee-repo-ref: 07a9d26a79a403ae27c48abd508a6699f2c87c49

New ee-repo-ref: a2fa58e5301d3865dd06ad73519e20ba7a5af0f0

Automated by sync-ee-ref workflow.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>

backend/.sqlx/query-a09eb290f2fc7c2de0e4b009e7ba7046de30385983ea39bf203aa8dc20b7aca1.json

@@ -0,0 +1,22 @@
+{
+  "db_name": "PostgreSQL",
+  "query": "\n            WITH RECURSIVE chain AS (\n                SELECT id, parent_workspace_id, is_dev_workspace, 0 AS depth\n                FROM workspace WHERE id = $1\n                UNION ALL\n                SELECT w.id, w.parent_workspace_id, w.is_dev_workspace, chain.depth + 1\n                FROM workspace w\n                JOIN chain ON w.id = chain.parent_workspace_id\n                WHERE chain.depth < 20\n            )\n            SELECT id AS \"id!\" FROM chain\n            WHERE parent_workspace_id IS NULL OR is_dev_workspace\n            ORDER BY depth LIMIT 1\n        ",
+  "describe": {
+    "columns": [
+      {
+        "ordinal": 0,
+        "name": "id!",
+        "type_info": "Varchar"
+      }
+    ],
+    "parameters": {
+      "Left": [
+        "Text"
+      ]
+    },
+    "nullable": [
+      null
+    ]
+  },
+  "hash": "a09eb290f2fc7c2de0e4b009e7ba7046de30385983ea39bf203aa8dc20b7aca1"
+}

backend/ee-repo-ref.txt

@@ -1 +1 @@
-d30af67d38954f9012f7bad08da23e347344b4c6
+a2fa58e5301d3865dd06ad73519e20ba7a5af0f0

backend/src/main.rs

@@ -1736,7 +1736,9 @@ async fn process_notify_event(
             windmill_common::variables::WORKSPACE_CRYPT_CACHE.remove(payload);
         }
         c if c == windmill_queue::tags::FORK_LINEAGE_CHANGE_CHANNEL => {
-            tracing::info!("Fork lineage change detected ({payload}), dropping tag workspace cache");
+            tracing::info!(
+                "Fork lineage change detected ({payload}), dropping lineage-derived caches"
+            );
             windmill_queue::tags::apply_fork_lineage_change(payload);
         }
         "notify_workspace_premium_change" => {

backend/windmill-api-workspaces/src/workspaces.rs

@@ -7445,8 +7445,11 @@ async fn create_workspace_fork(
     tx.commit().await?;
 
     // A pre-creation lookup could have cached an EMPTY ancestor chain for this id, which
-    // would bypass ducklake fork isolation for the TTL.
+    // would bypass ducklake fork isolation for the TTL. The same lookup could have cached the id
+    // as its own root workspace, which would make the fork's first jobs report themselves as their
+    // own environment instead of the parent.
     windmill_common::workspaces::invalidate_fork_ancestor_chain_cache(&forked_id);
+    windmill_queue::tags::invalidate_fork_parent_cache(&forked_id);
 
     if locked_prod {
         windmill_common::workspaces::invalidate_protection_rules_cache(&parent_workspace_id);

backend/windmill-common/src/variables.rs

@@ -433,6 +433,7 @@ pub async fn get_reserved_variables(
     };
 
     let custom_envs = get_cached_workspace_envs(conn, w_id).await;
+    let root_workspace = crate::workspaces::root_workspace_id(conn, w_id).await;
 
     let joined_schedule_path = schedule_path
         .clone()
@@ -462,6 +463,15 @@ pub async fn get_reserved_variables(
         description: "Workspace id of the current script".to_string(),
         is_custom: false,
     },
+    ContextualVariable {
+        name: "WM_ROOT_WORKSPACE".to_string(),
+        value: root_workspace,
+        description: "Workspace id of the nearest dev or prod workspace at or above the current \
+                      one. Equal to WM_WORKSPACE unless the script runs in a fork, in which case \
+                      it is the closest dev or prod workspace the fork descends from - not \
+                      necessarily its direct parent, since a fork of a fork skips past it".to_string(),
+        is_custom: false,
+    },
     ContextualVariable {
         name: "WM_TOKEN".to_string(),
         value: token.to_string(),

backend/windmill-common/src/workspaces.rs

@@ -1494,6 +1494,124 @@ pub async fn workspace_with_fork_ancestors(db: &crate::DB, w_id: &str) -> Result
     Ok(chain)
 }
 
+lazy_static::lazy_static! {
+    /// workspace id -> (root workspace id, expiry ts). Read once per job start, so correctness
+    /// rests on the invalidation rather than on the TTL: every mutation that can change the answer
+    /// sweeps the ids it touches through `windmill_queue::tags::invalidate_fork_parent_cache` and
+    /// broadcasts on `FORK_LINEAGE_CHANGE_CHANNEL`. A process that receives no broadcast — an agent
+    /// worker polls no notify events — has only the TTL, and takes the shorter one.
+    static ref ROOT_WORKSPACE_CACHE: Cache<String, (String, i64)> = Cache::new(5000);
+}
+
+const ROOT_WORKSPACE_CACHE_TTL_S: i64 = 300;
+/// An agent worker consumes no `notify_event`, so no sweep ever reaches its cache and the TTL is
+/// the whole invalidation story there. Hold its entries for the same 60s the other lineage caches
+/// (`FORK_ANCESTOR_CHAIN_CACHE`, `BILLING_WORKSPACE_CACHE`) accept as their staleness bound,
+/// rather than the long TTL that only a broadcast-fed process has earned.
+const ROOT_WORKSPACE_AGENT_CACHE_TTL_S: i64 = 60;
+/// An id the walk finds nothing for is cached far more briefly than a resolved one: it becomes
+/// resolvable the moment its workspace row lands, and creating a workspace is not a lineage change,
+/// so no sweep would drop the entry.
+const ROOT_WORKSPACE_UNRESOLVED_CACHE_TTL_S: i64 = 30;
+
+/// Drop the cached root workspace of one id. Called for every id whose lineage-derived caches are
+/// swept, so it needs no call site of its own — see
+/// `windmill_queue::tags::invalidate_fork_parent_cache`.
+pub fn invalidate_root_workspace_cache(w_id: &str) {
+    ROOT_WORKSPACE_CACHE.remove(w_id);
+}
+
+/// Drop every cached root workspace: the answer depends on the whole ancestor chain, so a mutation
+/// that reshapes the tree moves an unbounded set of descendants.
+pub fn clear_root_workspace_cache() {
+    ROOT_WORKSPACE_CACHE.clear();
+}
+
+/// Nearest ancestor-or-self of `w_id` that is an environment of its own: a root ("prod") workspace
+/// or a dev workspace. Equal to `w_id` for either of those, and to the standing workspace a
+/// throwaway fork was forked from otherwise. Exposed to jobs as `WM_ROOT_WORKSPACE`.
+///
+/// Falls back to `w_id` when the chain cannot be resolved (unknown id, broken or cyclic chain,
+/// failed lookup).
+///
+/// Not the same question as `get_billing_workspace_id`, which walks all the way to the parentless
+/// root: a fork under a dev workspace bills to prod but belongs to the dev environment.
+///
+/// Unauthenticated helper: reads workspace hierarchy for any `w_id`, so callers must already be
+/// authorized for that workspace (or run in trusted server-side code).
+pub async fn root_workspace_id(conn: &crate::worker::Connection, w_id: &str) -> String {
+    let now = chrono::Utc::now().timestamp();
+    let cached = ROOT_WORKSPACE_CACHE.get(w_id);
+    if let Some((root, expiry)) = &cached {
+        if *expiry > now {
+            return root.clone();
+        }
+    }
+
+    let (resolved, fresh_ttl) = match conn {
+        crate::worker::Connection::Sql(db) => (
+            lookup_root_workspace_id(db, w_id).await,
+            ROOT_WORKSPACE_CACHE_TTL_S,
+        ),
+        crate::worker::Connection::Http(client) => (
+            client
+                .get::<Option<String>>(&format!("/api/w/{w_id}/agent_workers/root_workspace"))
+                .await
+                .map_err(Error::from),
+            ROOT_WORKSPACE_AGENT_CACHE_TTL_S,
+        ),
+    };
+
+    let (root, ttl) = match resolved {
+        Ok(Some(root)) => (root, fresh_ttl),
+        Ok(None) => (w_id.to_string(), ROOT_WORKSPACE_UNRESOLVED_CACHE_TTL_S),
+        // A failed lookup is NOT cached, for the reason `lookup_tag_workspace` gives: the fallback
+        // is indistinguishable from a legitimate answer, so pinning one failure would make every
+        // job in a fork report the fork as its own environment until the entry expired. The expired
+        // entry is still the last answer this process actually resolved, so prefer it to that
+        // fallback — an agent talking to a server too old to serve the route would otherwise
+        // demote every fork to itself for the whole rolling upgrade.
+        Err(e) => {
+            tracing::warn!("failed to resolve root workspace of {w_id}: {e:#}");
+            return cached
+                .map(|(root, _)| root)
+                .unwrap_or_else(|| w_id.to_string());
+        }
+    };
+    ROOT_WORKSPACE_CACHE.insert(w_id.to_string(), (root.clone(), now + ttl));
+    root
+}
+
+/// Uncached lookup behind [`root_workspace_id`]. `None` when the chain resolves to nothing: an
+/// unknown id, or a (malformed) cycle that saturates the depth bound — the same cycle-safety
+/// backstop convention as [`fork_ancestor_chain`].
+///
+/// Unauthenticated helper: reads workspace hierarchy for any `w_id`, so callers must already be
+/// authorized for that workspace (or run in trusted server-side code). Answering for an arbitrary
+/// id discloses that the workspace exists and which environment it belongs to.
+pub async fn lookup_root_workspace_id(db: &crate::DB, w_id: &str) -> Result<Option<String>> {
+    sqlx::query_scalar!(
+        r#"
+            WITH RECURSIVE chain AS (
+                SELECT id, parent_workspace_id, is_dev_workspace, 0 AS depth
+                FROM workspace WHERE id = $1
+                UNION ALL
+                SELECT w.id, w.parent_workspace_id, w.is_dev_workspace, chain.depth + 1
+                FROM workspace w
+                JOIN chain ON w.id = chain.parent_workspace_id
+                WHERE chain.depth < 20
+            )
+            SELECT id AS "id!" FROM chain
+            WHERE parent_workspace_id IS NULL OR is_dev_workspace
+            ORDER BY depth LIMIT 1
+        "#,
+        w_id
+    )
+    .fetch_optional(db)
+    .await
+    .map_err(|e| Error::internal_err(format!("resolving root workspace of {w_id}: {e:#}")))
+}
+
 /// Resolve which live descendant workspace (and its inherited repo entry) a git
 /// branch pushed to `parent_repo_path` — a git-sync repo on `parent_w_id`
 /// tracking `expected_base` — deploys to via parent-managed fork sync, or `None`

backend/windmill-common/tests/root_workspace.rs

@@ -0,0 +1,60 @@
+//! Regression guard for the workspace a job reports as `WM_ROOT_WORKSPACE`.
+//!
+//! Run with:
+//!   cargo test -p windmill-common --test root_workspace
+
+use sqlx::{Pool, Postgres};
+use windmill_common::worker::Connection;
+use windmill_common::workspaces::{invalidate_root_workspace_cache, root_workspace_id};
+
+async fn insert_ws(db: &Pool<Postgres>, id: &str, parent: Option<&str>, is_dev: bool) {
+    sqlx::query(
+        "INSERT INTO workspace (id, name, owner, parent_workspace_id, is_dev_workspace)
+         VALUES ($1, $1, 'test-user', $2, $3)",
+    )
+    .bind(id)
+    .bind(parent)
+    .bind(is_dev)
+    .execute(db)
+    .await
+    .expect("insert workspace");
+    // The resolver caches per id in a process-global cache shared across tests.
+    invalidate_root_workspace_cache(id);
+}
+
+#[sqlx::test(migrations = "../migrations", fixtures("base"))]
+async fn root_workspace_is_the_nearest_dev_or_prod_ancestor(db: Pool<Postgres>) {
+    let conn = Connection::Sql(db.clone());
+    insert_ws(&db, "rwt-prod", None, false).await;
+    insert_ws(&db, "rwt-dev", Some("rwt-prod"), true).await;
+    insert_ws(&db, "wm-fork-rwt", Some("rwt-prod"), false).await;
+    insert_ws(&db, "wm-fork-underdev", Some("rwt-dev"), false).await;
+    insert_ws(&db, "wm-fork-nested", Some("wm-fork-underdev"), false).await;
+    // A generated-id workspace re-designated as a dev workspace, under its own root since only one
+    // dev workspace is allowed per parent.
+    insert_ws(&db, "rwt-prod2", None, false).await;
+    insert_ws(&db, "wm-fork-asdev", Some("rwt-prod2"), true).await;
+    insert_ws(&db, "wm-fork-asdev-fork", Some("wm-fork-asdev"), false).await;
+
+    assert_eq!(root_workspace_id(&conn, "rwt-prod").await, "rwt-prod");
+    assert_eq!(root_workspace_id(&conn, "rwt-dev").await, "rwt-dev");
+    assert_eq!(root_workspace_id(&conn, "wm-fork-rwt").await, "rwt-prod");
+    assert_eq!(
+        root_workspace_id(&conn, "wm-fork-underdev").await,
+        "rwt-dev"
+    );
+    assert_eq!(root_workspace_id(&conn, "wm-fork-nested").await, "rwt-dev");
+    // A dev workspace ends the walk whatever its id. Tag resolution deliberately walks past a
+    // generated-id dev workspace because nothing provisions workers for such an id, but the
+    // environment a fork belongs to is still that workspace.
+    assert_eq!(
+        root_workspace_id(&conn, "wm-fork-asdev").await,
+        "wm-fork-asdev"
+    );
+    assert_eq!(
+        root_workspace_id(&conn, "wm-fork-asdev-fork").await,
+        "wm-fork-asdev"
+    );
+    // An id with no lineage to resolve is its own environment.
+    assert_eq!(root_workspace_id(&conn, "rwt-missing").await, "rwt-missing");
+}

backend/windmill-queue/src/tags.rs

@@ -31,20 +31,26 @@ pub const FORK_LINEAGE_CHANGE_CHANNEL: &str = "notify_fork_lineage_change";
 /// set of descendants, which is not worth enumerating; a change to what a single id denotes is.
 const CLEAR_ALL: &str = "*";
 
-/// Drop one cached tag workspace in THIS process only. Resolution walks ancestors, so a mutation
+/// Drop the lineage-derived caches of one workspace id in THIS process only: the tag workspace
+/// here, and the root workspace a job reads as `WM_ROOT_WORKSPACE`. Both answer a walk up the
+/// ancestor chain, so every mutation invalidating one invalidates the other, and both are swept
+/// together rather than through parallel call sites. Resolution walks ancestors, so a mutation
 /// also invalidates every descendant; sweep them here as well. Replicas need a broadcast to match:
 /// [`notify_fork_lineage_reset`] when a subtree moved (attach, detach, archive, rename, a delete
 /// that orphans), or [`notify_fork_lineage_change`] when a single id changed what it denotes.
 pub fn invalidate_fork_parent_cache(workspace_id: &str) {
     FORK_PARENT_CACHE.remove(workspace_id);
+    windmill_common::workspaces::invalidate_root_workspace_cache(workspace_id);
 }
 
-/// Apply a broadcast lineage change to this process's cache.
+/// Apply a broadcast lineage change to this process's caches.
 pub fn apply_fork_lineage_change(payload: &str) {
     if payload == CLEAR_ALL {
         FORK_PARENT_CACHE.clear();
+        windmill_common::workspaces::clear_root_workspace_cache();
     } else {
         FORK_PARENT_CACHE.remove(payload);
+        windmill_common::workspaces::invalidate_root_workspace_cache(payload);
     }
 }
 

backend/windmill-queue/tests/tag_workspace_test.rs

@@ -8,6 +8,8 @@
 //!   cargo test -p windmill-queue --test tag_workspace_test
 
 use sqlx::{Pool, Postgres};
+use windmill_common::worker::Connection;
+use windmill_common::workspaces::root_workspace_id;
 use windmill_queue::tags::{
     apply_fork_lineage_change, invalidate_fork_parent_cache, tag_workspace_id,
 };
@@ -82,3 +84,40 @@ async fn reclaimed_fork_id_follows_its_new_parent(db: Pool<Postgres>) {
     apply_fork_lineage_change("wm-fork-rc");
     assert_eq!(tag_workspace_id("wm-fork-rc", &db).await, "rc-second");
 }
+
+/// `WM_ROOT_WORKSPACE` answers the same walk up the ancestor chain as the tag workspace, so its
+/// cache rides these two sweeps instead of carrying call sites of its own. That is only true while
+/// they actually drop it: dropping either line below leaves a job reporting an environment its
+/// workspace left minutes earlier, and no test in `windmill-common` can catch it — that crate
+/// cannot depend on this one.
+#[sqlx::test(migrations = "../migrations", fixtures("base"))]
+async fn the_sweeps_here_also_drop_the_root_workspace(db: Pool<Postgres>) {
+    let conn = Connection::Sql(db.clone());
+    insert_ws(&db, "rwx-prod", None, false).await;
+    insert_ws(&db, "rwx-mid", Some("rwx-prod"), false).await;
+    insert_ws(&db, "wm-fork-rwx", Some("rwx-mid"), false).await;
+
+    // Warm both caches the way pushing and then starting a job would.
+    assert_eq!(root_workspace_id(&conn, "wm-fork-rwx").await, "rwx-prod");
+
+    // Promoting a mid-chain workspace to a dev workspace is the mutation that moves the answer for
+    // its whole subtree. Done in SQL because the handler that does it lives in the API crate.
+    sqlx::query("UPDATE workspace SET is_dev_workspace = true WHERE id = 'rwx-mid'")
+        .execute(&db)
+        .await
+        .expect("promote to dev workspace");
+
+    invalidate_fork_parent_cache("wm-fork-rwx");
+    assert_eq!(root_workspace_id(&conn, "wm-fork-rwx").await, "rwx-mid");
+
+    // And again through the broadcast leg, which is the only one a replica ever sees. Only the
+    // single-id payload is exercised: `"*"` clears the process-global caches, which would yank the
+    // warmed entries out from under the other tests running in this same process.
+    sqlx::query("UPDATE workspace SET is_dev_workspace = false WHERE id = 'rwx-mid'")
+        .execute(&db)
+        .await
+        .expect("demote dev workspace");
+
+    apply_fork_lineage_change("wm-fork-rwx");
+    assert_eq!(root_workspace_id(&conn, "wm-fork-rwx").await, "rwx-prod");
+}

frontend/src/lib/deno_fetch.d.ts.txt

@@ -424,6 +424,7 @@ interface Process {
     env: {
       WM_TOKEN: string;
       WM_WORKSPACE: string;
+      WM_ROOT_WORKSPACE: string;
       WM_EMAIL: string;
       WM_USERNAME: string;
       WM_BASE_URL: string;

frontend/src/lib/process.d.ts.txt

@@ -160,6 +160,7 @@ declare module 'process' {
 				TZ?: string
 				WM_TOKEN: string
 				WM_WORKSPACE: string
+				WM_ROOT_WORKSPACE: string
 				WM_EMAIL: string
 				WM_USERNAME: string
 				WM_BASE_URL: string