Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 37 additions & 2 deletions crates/openjd-sessions/src/embedded_files.rs
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,7 @@ fn validate_resolved_filename(resolved: &str) -> Result<(), String> {
}

struct FileRecord {
_symbol: String,
symbol: String,
filename: PathBuf,
file: EmbeddedFile,
}
Expand Down Expand Up @@ -311,14 +311,49 @@ impl EmbeddedFiles {
)
.map_err(|e| SessionError::Runtime(format!("Failed to set {symbol}: {e}")))?;
self.records.push(FileRecord {
_symbol: symbol,
symbol,
filename,
file: file.clone(),
});
}
Ok(())
}

/// Re-register previously allocated file paths into a symbol table.
///
/// This is used when a wrap environment's embedded file paths have already
/// been allocated (on the first wrap-hook invocation) and need to be made
/// visible in subsequent wrap-hook scopes without re-allocating paths or
/// creating files on disk. Contents are NOT written — call
/// `write_file_contents` separately after the scope is fully built.
pub(crate) fn register_file_paths(&self, symtab: &mut SymbolTable) -> Result<(), SessionError> {
let scope_name = match self.scope {
EmbeddedFilesScope::Step => "Task",
EmbeddedFilesScope::Env => "Environment",
};
session_log!(
info,
&self.session_id,
LogContent::FILE_PATH,
"Reusing embedded file paths for {} scope.",
scope_name
);
for record in &self.records {
symtab
.set(
&record.symbol,
ExprValue::new_path(
record.filename.to_string_lossy().to_string(),
PathFormat::host(),
),
)
.map_err(|e| {
SessionError::Runtime(format!("Failed to set {}: {e}", record.symbol))
})?;
}
Ok(())
}

pub fn write_file_contents(
&self,
symtab: &SymbolTable,
Expand Down
178 changes: 170 additions & 8 deletions crates/openjd-sessions/src/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -474,6 +474,12 @@ pub struct Session {
// Environment tracking
environments: HashMap<EnvironmentIdentifier, Environment>,
environments_entered: Vec<EnvironmentIdentifier>,
/// Cached embedded-file records for wrap environments. Paths are allocated
/// once per wrap environment and reused for every subsequent wrap-hook
/// invocation so `Env.File.*` paths stay stable across tasks and unnamed
/// files do not accumulate on disk; contents are re-written per invocation
/// so each invocation starts from the authored content.
wrap_env_file_records: HashMap<EnvironmentIdentifier, EmbeddedFiles>,
// Env var tracking
env_vars: HashMap<String, String>,
process_env: HashMap<String, String>,
Expand Down Expand Up @@ -521,6 +527,7 @@ impl Session {
_files_dir: None,
environments: HashMap::new(),
environments_entered: Vec::new(),
wrap_env_file_records: HashMap::new(),
env_vars: HashMap::new(),
process_env: HashMap::new(),
created_env_vars: HashMap::new(),
Expand Down Expand Up @@ -724,6 +731,7 @@ impl Session {
_files_dir: Some(files_dir),
environments: HashMap::new(),
environments_entered: Vec::new(),
wrap_env_file_records: HashMap::new(),
env_vars: HashMap::new(),
process_env,
created_env_vars: HashMap::new(),
Expand Down Expand Up @@ -1194,6 +1202,19 @@ impl Session {

let lib = self.library.clone();
if let Some((wrap_env, _)) = wrap_action.as_ref() {
// Register wrap env's embedded file paths BEFORE seed so the
// wrap env's let bindings can reference Env.File.*.
let wrap_env_id = self.wrap_env_id_excluding(&identifier).cloned();
if let Some(ref wid) = wrap_env_id {
let files = self
.environments
.get(wid)
.and_then(|e| e.script.as_ref())
.and_then(|s| s.embedded_files.clone());
self.ensure_wrap_env_files(wid, files.as_deref(), &mut action_symtab)
.map_err(|e| self.fail_action_setup(e))?;
}

// The wrapped onEnter resolves against the INNER env's own
// scope (its embedded files and lets) — the same scope
// `runner.enter` would have built had the action run
Expand Down Expand Up @@ -1221,6 +1242,13 @@ impl Session {
"onEnter",
)
.map_err(|e| self.fail_action_setup(e))?;

// Write wrap env file contents AFTER seed (so data resolves
// against the post-lets symbol table).
if let Some(ref wid) = wrap_env_id {
self.write_wrap_env_file_contents(wid, &action_symtab, Some(&lib))
.map_err(|e| self.fail_action_setup(e))?;
}
}

// The effective action is now resolved (wrap hook or the env's
Expand Down Expand Up @@ -1440,6 +1468,19 @@ impl Session {

let lib = self.library.clone();
if let Some((wrap_env, _)) = wrap_action.as_ref() {
// Register wrap env's embedded file paths BEFORE seed so the
// wrap env's let bindings can reference Env.File.*.
let wrap_env_id = self.active_wrap_env_id().cloned();
if let Some(ref wid) = wrap_env_id {
let files = self
.environments
.get(wid)
.and_then(|e| e.script.as_ref())
.and_then(|s| s.embedded_files.clone());
self.ensure_wrap_env_files(wid, files.as_deref(), &mut action_symtab)
.map_err(|e| self.fail_action_setup(e))?;
}

// See the onEnter path: the wrapped onExit resolves against
// the INNER env's own scope.
let inner_symtab = self
Expand All @@ -1464,6 +1505,13 @@ impl Session {
"onExit",
)
.map_err(|e| self.fail_action_setup(e))?;

// Write wrap env file contents AFTER seed (so data resolves
// against the post-lets symbol table).
if let Some(ref wid) = wrap_env_id {
self.write_wrap_env_file_contents(wid, &action_symtab, Some(&lib))
.map_err(|e| self.fail_action_setup(e))?;
}
}

// See the onEnter path: record the effective action's declared
Expand Down Expand Up @@ -1581,6 +1629,13 @@ impl Session {
String::new()
};

// Evict the wrap-env file cache when the wrap environment itself is
// exited. The on-disk files are intentionally NOT deleted — they live
// in the session files directory and are cleaned up with it.
if env_has_any_wrap_hook(&env) {
self.wrap_env_file_records.remove(identifier);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Eviction is unreachable when the wrap environment's own onExit fails: the EnvironmentScriptFailed return a few lines above (and the result? before it) bypasses this line, while self.environments.remove(identifier) / environments_entered.pop() have already run at the top of the function.

The leftover entry is more than a leak. Because identifier is no longer in self.environments, the DuplicateEnvironment check in enter_environment_with_output no longer rejects it, so a caller that retries with the same identifier (worker-agent-supplied ids are stable) will hit ensure_wrap_env_files → cache hit and get the previous environment's FileRecords: symbols named after the old embedded files, pointing at the old paths, while the new environment's files are never allocated or written. A hook referencing one of the new file names then fails with Undefined variable.

Removing the entry right after the environments.remove/pop (i.e. alongside the other tracking teardown, before the exit script runs) makes eviction unconditional and keeps it paired with the state it mirrors.

}

Ok(output)
}

Expand Down Expand Up @@ -1632,15 +1687,37 @@ impl Session {
// step script runs exactly as before — this keeps the non-WRAP_ACTIONS
// path a zero-cost addition.
//
// Scope note: this pass does NOT re-materialize the wrap environment's
// embedded files. Wrap actions that reference `{{Env.File.*}}` will
// see only the names registered when the wrap env was entered, which
// are not persisted across action runs. Inline wrap scripts
// (`command: bash, args: ["-c", "..."]`) work without this. Re-running
// `allocate_file_paths` against the wrap env's embedded_files at task
// dispatch time is the follow-up to enable `Env.File.*` inside wrap
// hooks end-to-end.
// Wrap environment embedded files: paths are allocated once per wrap
// env (on the first hook invocation that triggers a cache miss) and
// reused for every subsequent invocation. File contents are re-written
// per invocation so each hook execution starts from the authored data.
// Registration happens BEFORE seed_wrapped_action_symbols so the wrap
// env's let bindings can reference `{{Env.File.*}}`; writing happens
// AFTER so file data resolves against the post-lets symbol table.
let lib = self.library.clone();

// Hoist wrap-env identity and embedded files BEFORE the closure that
// borrows `self` immutably, so we can call `&mut self` methods for
// the embedded-file cache.
let wrap_env_id_and_files: Option<(
EnvironmentIdentifier,
Option<Vec<openjd_model::job::EmbeddedFile>>,
)> = self.active_wrap_env_id().and_then(|id| {
let env = self.environments.get(id)?;
if env.script.as_ref()?.actions.on_wrap_task_run.is_some() {
let files = env.script.as_ref().and_then(|s| s.embedded_files.clone());
Some((id.clone(), files))
} else {
None
}
});

// Register wrap env file paths BEFORE seed (so lets can reference Env.File.*).
if let Some((ref wrap_id, ref files)) = wrap_env_id_and_files {
self.ensure_wrap_env_files(wrap_id, files.as_deref(), &mut action_symtab)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The wrap env's Env.File.* symbols are registered into action_symtab before build_wrapped_inner_scope is called (line 1738 here, and lines 1214/1224 + 1480/1487 on the enter/exit paths). build_wrapped_inner_scope takes &action_symtab as its base and clones it, so the wrapper's Env.File.* entries end up inside the inner scope too — the scope in which the wrapped action's command/args/timeout and the inner entity's let bindings are resolved.

That inverts the isolation invariant seed_wrapped_action_symbols documents ("a wrapper-defined name can never leak into the wrapped action's resolved values"):

  • For a step (onWrapTaskRun), Env.File.* is not part of a step script's scope at all, yet the wrapped onRun can now resolve {{Env.File.<wrapper-file>}}.
  • For onWrapEnvEnter/onWrapEnvExit, the inner env sees any wrapper file name it does not itself declare (its own names correctly shadow, since it allocates over the same keys afterwards).

Model-level validation rejects most such references, so the blast radius is templates that bypass it (externally supplied / worker-injected environments) — but the resolution result differs from the unwrapped run, which is the property the wrap path is supposed to preserve. Building the inner scope from a snapshot of action_symtab taken before the wrap-env registration would keep both scopes correct while preserving the register→seed→write ordering.

.map_err(|e| self.fail_action_setup(e))?;
}

let wrap_action: Option<openjd_model::job::Action> = self
.active_wrap_env()
.and_then(|wrap_env| {
Expand Down Expand Up @@ -1680,6 +1757,12 @@ impl Session {
.transpose()
.map_err(|e| self.fail_action_setup(e))?;

// Write wrap env file contents AFTER seed (so data resolves against post-lets symtab).
if let Some((ref wrap_id, _)) = wrap_env_id_and_files {
self.write_wrap_env_file_contents(wrap_id, &action_symtab, Some(&lib))
.map_err(|e| self.fail_action_setup(e))?;
}

// Box large locals so they live on the heap instead of inflating
// this async fn's state machine. Without this, the combined future
// (run_task → drive_action → select!) exceeds Windows' default
Expand Down Expand Up @@ -2515,6 +2598,85 @@ impl Session {
}
None
}

/// Like `wrap_env_excluding` but returns the identifier instead of a
/// reference to the environment, avoiding borrow conflicts when `&mut self`
/// is needed after the lookup.
fn wrap_env_id_excluding(&self, self_id: &str) -> Option<&EnvironmentIdentifier> {
for id in self.environments_entered.iter().rev() {
if id == self_id {
continue;
}
if let Some(env) = self.environments.get(id) {
if env_has_any_wrap_hook(env) {
return Some(id);
}
}
}
None
}

/// Return the `EnvironmentIdentifier` of the active wrap environment, if
/// one exists. Like `active_wrap_env` but returns the id rather than a
/// reference, avoiding borrow conflicts when mutation is needed.
fn active_wrap_env_id(&self) -> Option<&EnvironmentIdentifier> {
for id in self.environments_entered.iter().rev() {
if let Some(env) = self.environments.get(id) {
if env_has_any_wrap_hook(env) {
return Some(id);
}
}
}
None
}

/// Ensure the wrap environment's embedded file paths are allocated and
/// registered in `symtab`. On the first call for a given wrap env
/// (cache miss), allocates paths and creates the cache entry. On
/// subsequent calls (cache hit), re-registers the previously allocated
/// paths without allocating new ones.
///
/// No-op when `files` is `None` or empty.
fn ensure_wrap_env_files(
&mut self,
wrap_env_id: &EnvironmentIdentifier,
files: Option<&[openjd_model::job::EmbeddedFile]>,
symtab: &mut SymbolTable,
) -> Result<(), SessionError> {
let files = match files {
Some(f) if !f.is_empty() => f,
_ => return Ok(()),
};
if let Some(cached) = self.wrap_env_file_records.get(wrap_env_id) {
// Cache hit: re-register paths without allocating.
cached.register_file_paths(symtab)?;
} else {
// Cache miss: allocate paths and insert into cache.
let mut ef = EmbeddedFiles::new(
EmbeddedFilesScope::Env,
self.files_directory.clone(),
&self.session_id,
)
.with_user(self.cross_user.user.clone());
ef.allocate_file_paths(files, symtab)?;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This cache is a second, independent allocation of the same wrap environment's embedded files. When the wrap env was entered, EnvironmentScriptRunner::enter built its own EmbeddedFiles and allocated paths for the onEnter script (runner/env_script.rs:236) — that instance is dropped, and the first wrap-hook invocation allocates again here.

For named files both allocations land on <files_dir>/<filename>, so they agree. For unnamed files each allocation calls random_hex_filename(), so {{Env.File.foo}} inside the wrap env's own onEnter and {{Env.File.foo}} inside its onWrapTaskRun resolve to different files. Anything the onEnter script writes into that path is invisible to the hooks — which is a natural way to author a wrap env (stage state in an unnamed embedded file at enter, consume it per task) and would silently read the wrong file. It also leaves an extra unnamed file in the session files dir per wrap env.

Populating the cache from the enter-time allocation (or having enter consult the same cache) would make the symbol mean one file for the whole environment. If the divergence is intentional, it is worth stating explicitly in specs/sessions/embedded-files.md, since the new section reads as though there is a single allocation per wrap environment.

self.wrap_env_file_records.insert(wrap_env_id.clone(), ef);
}
Ok(())
}

/// Write the cached wrap environment's embedded file contents into the
/// symtab's resolved scope. No-op if no cache entry exists for the id.
fn write_wrap_env_file_contents(
&self,
wrap_env_id: &EnvironmentIdentifier,
symtab: &SymbolTable,
lib: Option<&FunctionLibrary>,
) -> Result<(), SessionError> {
if let Some(cached) = self.wrap_env_file_records.get(wrap_env_id) {
cached.write_file_contents(symtab, lib)?;
}
Ok(())
}
}

/// Returns true iff the environment defines any of the three wrap hooks
Expand Down
Loading
Loading