harn-vm 0.10.162

//! Explicit persistent-state ownership for isolated VM executions.

use std::cell::RefCell;
use std::marker::PhantomData;
use std::path::{Path, PathBuf};
use std::rc::Rc;

use crate::Vm;

thread_local! {
    static SCOPED_PERSISTENT_STATE_ROOT: RefCell<Option<PersistentStateContext>> = const { RefCell::new(None) };
}

#[derive(Clone, Debug)]
pub(crate) struct PersistentStateContext {
    state_root: PathBuf,
    workspace: Option<PathBuf>,
}

/// One selected workspace and its captured physical store address.
/// Execution and capability roots remain independent of this storage scope.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct SessionStoreScope {
    workspace: PathBuf,
    state_root: PathBuf,
}

impl SessionStoreScope {
    /// Resolve trusted storage configuration once for an existing workspace.
    pub fn resolve(
        cwd: &Path,
        launcher: &crate::security::LauncherEnvironment,
    ) -> std::io::Result<Self> {
        let workspace = cwd.canonicalize()?;
        if !workspace.is_dir() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "session workspace must be a directory",
            ));
        }
        let state_root = crate::runtime_paths::captured_state_root(&workspace, launcher);
        Ok(Self {
            workspace,
            state_root,
        })
    }

    pub fn workspace(&self) -> &Path {
        &self.workspace
    }
    pub fn state_root(&self) -> &Path {
        &self.state_root
    }

    /// Bind a synchronous store open. Use `run` across async suspensions.
    pub fn enter(&self) -> ScopedPersistentStateRoot {
        replace_context(PersistentStateContext {
            state_root: self.state_root.clone(),
            workspace: Some(self.workspace.clone()),
        })
    }

    /// Save and restore this storage scope at every async suspension.
    pub async fn run<F: std::future::Future>(&self, future: F) -> F::Output {
        // Capture when the execution is polled, after outer transport and host
        // scopes are installed, rather than while its future is constructed.
        crate::orchestration::scope_persistent_state_context(
            PersistentStateContext {
                state_root: self.state_root.clone(),
                workspace: Some(self.workspace.clone()),
            },
            future,
        )
        .await
    }
}

/// A caller-owned persistent-state root that bypasses ambient runtime paths.
#[derive(Clone, Copy, Debug)]
pub struct PersistentStateRoot<'a>(&'a Path);

impl<'a> PersistentStateRoot<'a> {
    #[must_use]
    pub fn new(path: &'a Path) -> Self {
        Self(path)
    }

    fn as_path(self) -> &'a Path {
        self.0
    }
}

/// Restores the prior caller-owned persistent-state root on drop.
///
/// The guard is deliberately not `Send`: its scope is tied to the current
/// thread or `LocalSet`, just like the runtime paths it overrides.
#[derive(Debug)]
#[must_use = "retain this guard for the isolated VM execution"]
pub struct ScopedPersistentStateRoot {
    previous: Option<PersistentStateContext>,
    _not_send: PhantomData<Rc<()>>,
}

/// Route every default persistent runtime path through one caller-owned root.
///
/// This is the execution-wide counterpart to
/// [`register_persistent_state_builtins_at_root`]. It covers durable consumers
/// such as the agent session journal that resolve their paths during execution
/// rather than when builtins are registered.
pub fn scope_persistent_state_root(root: PersistentStateRoot<'_>) -> ScopedPersistentStateRoot {
    replace_context(PersistentStateContext {
        state_root: root.as_path().to_path_buf(),
        workspace: None,
    })
}

fn replace_context(context: PersistentStateContext) -> ScopedPersistentStateRoot {
    let previous = SCOPED_PERSISTENT_STATE_ROOT.with(|slot| slot.replace(Some(context)));
    ScopedPersistentStateRoot {
        previous,
        _not_send: PhantomData,
    }
}

impl Drop for ScopedPersistentStateRoot {
    fn drop(&mut self) {
        SCOPED_PERSISTENT_STATE_ROOT.with(|slot| {
            slot.replace(self.previous.take());
        });
    }
}

pub(crate) fn current_persistent_state_root() -> Option<PathBuf> {
    SCOPED_PERSISTENT_STATE_ROOT.with(|slot| {
        slot.borrow()
            .as_ref()
            .map(|context| context.state_root.clone())
    })
}

pub(crate) fn current_session_workspace() -> Option<PathBuf> {
    SCOPED_PERSISTENT_STATE_ROOT.with(|slot| {
        slot.borrow()
            .as_ref()
            .and_then(|context| context.workspace.clone())
    })
}

pub(crate) fn session_matches_current_workspace(meta: &harn_session_store::SessionMeta) -> bool {
    session_matches_workspace(meta, current_session_workspace())
}

fn session_matches_workspace(
    meta: &harn_session_store::SessionMeta,
    workspace: Option<PathBuf>,
) -> bool {
    workspace.is_none_or(|workspace| {
        let expected = workspace.to_string_lossy();
        meta.project_scope.as_deref() == Some(expected.as_ref())
            && meta.cwd.as_deref() == Some(expected.as_ref())
    })
}

pub(crate) fn validate_session_store_workspace(
    state_root: &Path,
    meta: &harn_session_store::SessionMeta,
) -> Result<(), crate::VmError> {
    if session_matches_workspace(meta, session_workspace_for_state(state_root)) {
        Ok(())
    } else {
        Err(crate::VmError::Runtime(format!(
            "session_store: session '{}' belongs to another workspace",
            meta.id
        )))
    }
}

pub(crate) fn session_workspace_for_state(state_root: &Path) -> Option<PathBuf> {
    SCOPED_PERSISTENT_STATE_ROOT.with(|slot| {
        slot.borrow()
            .as_ref()
            .filter(|context| context.state_root == state_root)
            .and_then(|context| context.workspace.clone())
    })
}

pub(crate) fn canonical_store_root(workspace: &Path) -> PathBuf {
    SCOPED_PERSISTENT_STATE_ROOT.with(|slot| {
        slot.borrow()
            .as_ref()
            .filter(|context| context.workspace.as_deref() == Some(workspace))
            .map(|context| context.state_root.clone())
            .unwrap_or_else(|| workspace.join(".harn"))
    })
}

pub(crate) fn swap_persistent_state_context(
    replacement: Option<PersistentStateContext>,
) -> Option<PersistentStateContext> {
    SCOPED_PERSISTENT_STATE_ROOT.with(|slot| slot.replace(replacement))
}

/// Register store, metadata, and checkpoint builtins at an exact state root.
///
/// Unlike the individual runtime registrars, this function does not consult
/// `HARN_STATE_DIR`. Embedders use it when concurrent executions require
/// hermetic state without mutating process-global environment variables.
pub fn register_persistent_state_builtins_at_root(
    vm: &mut Vm,
    base_dir: &Path,
    state_root: PersistentStateRoot<'_>,
    pipeline_name: &str,
) {
    let state_root = state_root.as_path();
    crate::store::register_store_builtins_at_state_root(vm, state_root);
    crate::metadata::register_metadata_builtins_at_state_root(vm, base_dir, state_root);
    crate::checkpoint::register_checkpoint_builtins_at_state_root(vm, state_root, pipeline_name);
}

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

    #[test]
    fn scoped_persistent_state_root_restores_nested_owner() {
        let outer = Path::new("/isolated/outer/.harn");
        let inner = Path::new("/isolated/inner/.harn");
        assert_eq!(current_persistent_state_root(), None);
        let outer_guard = scope_persistent_state_root(PersistentStateRoot::new(outer));
        assert_eq!(current_persistent_state_root().as_deref(), Some(outer));
        {
            let _inner_guard = scope_persistent_state_root(PersistentStateRoot::new(inner));
            assert_eq!(current_persistent_state_root().as_deref(), Some(inner));
        }
        assert_eq!(current_persistent_state_root().as_deref(), Some(outer));
        drop(outer_guard);
        assert_eq!(current_persistent_state_root(), None);
    }

    #[tokio::test(flavor = "current_thread")]
    async fn scoped_store_row_validation_preserves_explicit_root_and_refuses_foreign_rows() {
        use harn_session_store::{CreateSession, SessionStore, SqliteSessionStore};

        let dir = tempfile::tempdir().unwrap();
        let launcher = crate::security::LauncherEnvironment::from_snapshot(Default::default());
        let scope = SessionStoreScope::resolve(dir.path(), &launcher).unwrap();
        let store = SqliteSessionStore::open_in_memory().unwrap();
        let workspace = scope.workspace().to_string_lossy().into_owned();
        let mut meta = store
            .create(CreateSession {
                cwd: Some(workspace.clone()),
                project_scope: Some(workspace),
                ..CreateSession::default()
            })
            .await
            .unwrap();
        let _guard = scope.enter();
        assert!(validate_session_store_workspace(scope.state_root(), &meta).is_ok());
        meta.cwd = Some("foreign workspace".to_string());
        assert!(validate_session_store_workspace(scope.state_root(), &meta).is_err());
        meta.cwd = None;
        meta.project_scope = None;
        assert!(validate_session_store_workspace(scope.state_root(), &meta).is_err());
        assert!(
            validate_session_store_workspace(&dir.path().join("explicit-other-state"), &meta)
                .is_ok()
        );
    }

    #[tokio::test(flavor = "current_thread")]
    async fn interleaved_session_store_scopes_restore_and_follow_child_tasks() {
        let left_dir = tempfile::tempdir().expect("left workspace");
        let right_dir = tempfile::tempdir().expect("right workspace");
        let launcher = crate::security::LauncherEnvironment::from_snapshot(Default::default());
        let left = SessionStoreScope::resolve(left_dir.path(), &launcher).unwrap();
        let right = SessionStoreScope::resolve(right_dir.path(), &launcher).unwrap();
        let check = |scope: SessionStoreScope| async move {
            scope
                .run(async {
                    for _ in 0..3 {
                        assert_eq!(
                            current_persistent_state_root().as_deref(),
                            Some(scope.state_root())
                        );
                        assert_eq!(
                            current_session_workspace().as_deref(),
                            Some(scope.workspace())
                        );
                        tokio::task::yield_now().await;
                    }
                    let expected = scope.state_root().to_path_buf();
                    crate::orchestration::run_blocking_with_ambient(move || {
                        assert_eq!(
                            current_persistent_state_root().as_deref(),
                            Some(expected.as_path())
                        );
                    })
                    .await
                    .expect("captured child host task");
                    crate::orchestration::scope_ambient(
                        crate::orchestration::AmbientExecutionScope::capture_inherited(),
                        async {
                            tokio::task::yield_now().await;
                            assert_eq!(
                                current_persistent_state_root().as_deref(),
                                Some(scope.state_root())
                            );
                            assert_eq!(
                                current_session_workspace().as_deref(),
                                Some(scope.workspace())
                            );
                        },
                    )
                    .await;
                })
                .await;
        };
        tokio::join!(check(left), check(right));
        assert_eq!(current_persistent_state_root(), None);
        assert_eq!(current_session_workspace(), None);
    }
}