harn_vm/stdlib/

agent_state.rs

pub mod backend;

use crate::value::VmDictExt;
use std::collections::BTreeMap;
use std::path::PathBuf;

use backend::{
    BackendScope, BackendWriteOptions, ConflictPolicy, DurableStateBackend, FilesystemBackend,
    WriterIdentity,
};

use crate::stdlib::macros::{harn_builtin, VmBuiltinDef};
use crate::value::{VmError, VmValue};
use crate::vm::Vm;

const HANDLE_TYPE: &str = "state_handle";
const HANDOFF_KEY: &str = "__handoff.json";

pub use backend::{
    BackendWriteOutcome, ConflictRecord, DurableStateBackend as AgentStateBackend,
    FilesystemBackend as AgentStateFilesystemBackend,
};

pub fn register_agent_state_builtins(vm: &mut Vm) {
    for def in MODULE_BUILTINS {
        vm.register_builtin_def(def);
    }
}

pub(crate) const MODULE_BUILTINS: &[&VmBuiltinDef] = &[
    &AGENT_STATE_INIT_IMPL_DEF,
    &AGENT_STATE_RESUME_IMPL_DEF,
    &AGENT_STATE_WRITE_IMPL_DEF,
    &AGENT_STATE_READ_IMPL_DEF,
    &AGENT_STATE_LIST_IMPL_DEF,
    &AGENT_STATE_DELETE_IMPL_DEF,
    &AGENT_STATE_HANDOFF_IMPL_DEF,
];

#[harn_builtin(
    sig = "__agent_state_init(root_or_session: string, options_or_root?: any, options?: dict) -> dict",
    runtime_only = true,
    category = "agent_state"
)]
fn agent_state_init_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
    let backend = FilesystemBackend::new();
    let (scope, writer, conflict_policy) = parse_init_request(args)?;
    backend.ensure_scope(&scope)?;
    Ok(handle_value(&backend, &scope, &writer, conflict_policy))
}

#[harn_builtin(
    sig = "__agent_state_resume(root: string, session_id: string, options?: dict) -> dict",
    runtime_only = true,
    category = "agent_state"
)]
fn agent_state_resume_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
    let backend = FilesystemBackend::new();
    let (scope, writer, conflict_policy) = parse_resume_request(args)?;
    backend.resume_scope(&scope)?;
    Ok(handle_value(&backend, &scope, &writer, conflict_policy))
}

#[harn_builtin(
    sig = "__agent_state_write(handle: dict, key: string, content: string) -> nil",
    runtime_only = true,
    category = "agent_state"
)]
fn agent_state_write_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
    let backend = FilesystemBackend::new();
    let handle = handle_from_args(args, "__agent_state_write")?;
    let key = required_arg_string(args, 1, "__agent_state_write", "key")?;
    let content = required_arg_string(args, 2, "__agent_state_write", "content")?;
    let scope = scope_from_handle(handle)?;
    let options = write_options_from_handle(handle)?;
    let outcome = backend.write(&scope, &key, &content, &options)?;
    enforce_conflict_policy(handle, &outcome)?;
    Ok(VmValue::Nil)
}

#[harn_builtin(
    sig = "__agent_state_read(handle: dict, key: string) -> string?",
    runtime_only = true,
    category = "agent_state"
)]
fn agent_state_read_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
    let backend = FilesystemBackend::new();
    let handle = handle_from_args(args, "__agent_state_read")?;
    let key = required_arg_string(args, 1, "__agent_state_read", "key")?;
    let scope = scope_from_handle(handle)?;
    match backend.read(&scope, &key)? {
        Some(content) => Ok(VmValue::String(arcstr::ArcStr::from(content))),
        None => Ok(VmValue::Nil),
    }
}

#[harn_builtin(
    sig = "__agent_state_list(handle: dict) -> list",
    runtime_only = true,
    category = "agent_state"
)]
fn agent_state_list_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
    let backend = FilesystemBackend::new();
    let handle = handle_from_args(args, "__agent_state_list")?;
    let scope = scope_from_handle(handle)?;
    let items = backend
        .list(&scope)?
        .into_iter()
        .map(|key| VmValue::String(arcstr::ArcStr::from(key)))
        .collect();
    Ok(VmValue::List(std::sync::Arc::new(items)))
}

#[harn_builtin(
    sig = "__agent_state_delete(handle: dict, key: string) -> nil",
    runtime_only = true,
    category = "agent_state"
)]
fn agent_state_delete_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
    let backend = FilesystemBackend::new();
    let handle = handle_from_args(args, "__agent_state_delete")?;
    let key = required_arg_string(args, 1, "__agent_state_delete", "key")?;
    let scope = scope_from_handle(handle)?;
    backend.delete(&scope, &key)?;
    Ok(VmValue::Nil)
}

#[harn_builtin(
    sig = "__agent_state_handoff(handle: dict, summary: dict) -> nil",
    runtime_only = true,
    category = "agent_state"
)]
fn agent_state_handoff_impl(args: &[VmValue], _out: &mut String) -> Result<VmValue, VmError> {
    let backend = FilesystemBackend::new();
    let handle = handle_from_args(args, "__agent_state_handoff")?;
    let summary = args.get(1).ok_or_else(|| {
        VmError::Runtime("__agent_state_handoff: `summary` is required".to_string())
    })?;
    let summary_json = crate::llm::vm_value_to_json(summary);
    let serde_json::Value::Object(_) = summary_json else {
        return Err(VmError::Runtime(
            "__agent_state_handoff: `summary` must be a JSON object".to_string(),
        ));
    };
    let typed_handoff =
        crate::orchestration::normalize_handoff_artifact_json(summary_json.clone()).ok();
    let scope = scope_from_handle(handle)?;
    let writer = writer_from_handle(handle);
    let envelope = serde_json::json!({
        "_type": "agent_state_handoff",
        "version": 1,
        "session_id": scope.namespace,
        "root": scope.root.to_string_lossy(),
        "key": HANDOFF_KEY,
        "summary": summary_json,
        "handoff": typed_handoff,
        "writer": writer_json(&writer),
        "written_at": now_epoch_seconds(),
    });
    let content = serde_json::to_string_pretty(&envelope)
        .map_err(|error| VmError::Runtime(format!("agent_state.handoff: encode error: {error}")))?;
    let options = write_options_from_handle(handle)?;
    let outcome = backend.write(&scope, HANDOFF_KEY, &content, &options)?;
    enforce_conflict_policy(handle, &outcome)?;
    Ok(VmValue::Nil)
}

fn handle_from_args<'a>(
    args: &'a [VmValue],
    fn_name: &str,
) -> Result<&'a crate::value::DictMap, VmError> {
    let handle = args
        .first()
        .ok_or_else(|| VmError::Runtime(format!("{fn_name}: `handle` is required")))?;
    let dict = handle.as_dict().ok_or_else(|| {
        VmError::Runtime(format!("{fn_name}: `handle` must be a state_handle dict"))
    })?;
    match dict.get("_type").map(VmValue::display).as_deref() {
        Some(HANDLE_TYPE) => Ok(dict),
        _ => Err(VmError::Runtime(format!(
            "{fn_name}: `handle` must be a state_handle dict"
        ))),
    }
}

fn required_arg_string(
    args: &[VmValue],
    idx: usize,
    fn_name: &str,
    arg_name: &str,
) -> Result<String, VmError> {
    match args.get(idx) {
        Some(VmValue::String(value)) => Ok(value.to_string()),
        Some(value) if !value.display().is_empty() => Ok(value.display()),
        _ => Err(VmError::Runtime(format!(
            "{fn_name}: `{arg_name}` must be a non-empty string"
        ))),
    }
}

fn parse_init_request(
    args: &[VmValue],
) -> Result<(BackendScope, WriterIdentity, ConflictPolicy), VmError> {
    let (root, options, explicit_session_id) =
        match (args.first(), args.get(1), args.get(2)) {
            (Some(VmValue::String(root)), Some(VmValue::Dict(options)), _) => {
                (root.to_string(), Some((**options).clone()), None)
            }
            (Some(VmValue::String(root)), None | Some(VmValue::Nil), _) => {
                (root.to_string(), None, None)
            }
            (Some(VmValue::String(session_id)), Some(VmValue::String(root)), maybe_options) => {
                let options = maybe_options.and_then(VmValue::as_dict).cloned();
                (root.to_string(), options, Some(session_id.to_string()))
            }
            _ => return Err(VmError::Runtime(
                "__agent_state_init: expected `(root, options?)` or `(session_id, root, options?)`"
                    .to_string(),
            )),
        };
    let root = resolve_root(&root);
    let session_id = explicit_session_id
        .or_else(|| option_string(options.as_ref(), "session_id"))
        .unwrap_or_else(default_session_id);
    let writer = writer_identity(options.as_ref(), Some(&session_id));
    let conflict_policy = conflict_policy(options.as_ref())?;
    Ok((
        BackendScope {
            root,
            namespace: session_id,
        },
        writer,
        conflict_policy,
    ))
}

fn parse_resume_request(
    args: &[VmValue],
) -> Result<(BackendScope, WriterIdentity, ConflictPolicy), VmError> {
    let root = required_arg_string(args, 0, "__agent_state_resume", "root")?;
    let session_id = required_arg_string(args, 1, "__agent_state_resume", "session_id")?;
    let options = args.get(2).and_then(VmValue::as_dict).cloned();
    let writer = writer_identity(options.as_ref(), Some(&session_id));
    let conflict_policy = conflict_policy(options.as_ref())?;
    Ok((
        BackendScope {
            root: resolve_root(&root),
            namespace: session_id,
        },
        writer,
        conflict_policy,
    ))
}

fn resolve_root(root: &str) -> PathBuf {
    crate::stdlib::process::resolve_source_relative_path(root)
}

fn conflict_policy(options: Option<&crate::value::DictMap>) -> Result<ConflictPolicy, VmError> {
    let Some(options) = options else {
        return Ok(ConflictPolicy::Ignore);
    };
    let raw = options
        .get("conflict_policy")
        .or_else(|| options.get("two_writer"))
        .map(VmValue::display)
        .unwrap_or_else(|| "ignore".to_string());
    ConflictPolicy::parse(&raw)
}

fn option_string(options: Option<&crate::value::DictMap>, key: &str) -> Option<String> {
    options
        .and_then(|options| options.get(key))
        .map(VmValue::display)
        .filter(|value| !value.trim().is_empty())
}

fn writer_identity(
    options: Option<&crate::value::DictMap>,
    session_id: Option<&str>,
) -> WriterIdentity {
    let mutation = crate::orchestration::current_mutation_session();
    let current_session = crate::agent_sessions::current_session_id();
    let session_id = session_id
        .map(|value| value.to_string())
        .or_else(|| current_session.clone())
        .or_else(|| mutation.as_ref().map(|session| session.session_id.clone()))
        .filter(|value| !value.is_empty());

    let worker_id = option_string(options, "worker_id").or_else(|| {
        mutation
            .as_ref()
            .and_then(|session| session.worker_id.clone())
    });
    let stage_id = option_string(options, "stage_id")
        .or_else(|| worker_id.clone())
        .or_else(|| mutation.as_ref().and_then(|session| session.run_id.clone()))
        .or_else(|| {
            mutation
                .as_ref()
                .and_then(|session| session.execution_kind.clone())
        });
    let writer_id = option_string(options, "writer_id")
        .or_else(|| worker_id.clone())
        .or_else(|| stage_id.clone())
        .or_else(|| session_id.clone());

    WriterIdentity {
        writer_id,
        stage_id,
        session_id,
        worker_id,
    }
}

fn default_session_id() -> String {
    crate::agent_sessions::current_session_id()
        .or_else(|| {
            crate::orchestration::current_mutation_session()
                .map(|session| session.session_id)
                .filter(|value| !value.is_empty())
        })
        .unwrap_or_else(|| uuid::Uuid::now_v7().to_string())
}

fn handle_value(
    backend: &impl DurableStateBackend,
    scope: &BackendScope,
    writer: &WriterIdentity,
    conflict_policy: ConflictPolicy,
) -> VmValue {
    let mut handle = BTreeMap::new();
    handle.put_str("_type", HANDLE_TYPE);
    handle.put_str("backend", backend.backend_name());
    handle.put_str("root", scope.root.to_string_lossy());
    handle.put_str("session_id", scope.namespace.clone());
    handle.put_str("handoff_key", HANDOFF_KEY);
    handle.put_str("conflict_policy", conflict_policy.as_str());
    handle.insert("writer".to_string(), writer_vm_value(writer));
    VmValue::dict(handle)
}

fn writer_vm_value(writer: &WriterIdentity) -> VmValue {
    let mut value = BTreeMap::new();
    value.insert(
        "writer_id".to_string(),
        writer
            .writer_id
            .as_ref()
            .map(|item| VmValue::String(arcstr::ArcStr::from(item.clone())))
            .unwrap_or(VmValue::Nil),
    );
    value.insert(
        "stage_id".to_string(),
        writer
            .stage_id
            .as_ref()
            .map(|item| VmValue::String(arcstr::ArcStr::from(item.clone())))
            .unwrap_or(VmValue::Nil),
    );
    value.insert(
        "session_id".to_string(),
        writer
            .session_id
            .as_ref()
            .map(|item| VmValue::String(arcstr::ArcStr::from(item.clone())))
            .unwrap_or(VmValue::Nil),
    );
    value.insert(
        "worker_id".to_string(),
        writer
            .worker_id
            .as_ref()
            .map(|item| VmValue::String(arcstr::ArcStr::from(item.clone())))
            .unwrap_or(VmValue::Nil),
    );
    VmValue::dict(value)
}

fn scope_from_handle(handle: &crate::value::DictMap) -> Result<BackendScope, VmError> {
    let root = handle
        .get("root")
        .map(VmValue::display)
        .filter(|value| !value.is_empty())
        .ok_or_else(|| VmError::Runtime("state_handle is missing `root`".to_string()))?;
    let session_id = handle
        .get("session_id")
        .map(VmValue::display)
        .filter(|value| !value.is_empty())
        .ok_or_else(|| VmError::Runtime("state_handle is missing `session_id`".to_string()))?;
    Ok(BackendScope {
        root: PathBuf::from(root),
        namespace: session_id,
    })
}

fn writer_from_handle(handle: &crate::value::DictMap) -> WriterIdentity {
    let writer = handle.get("writer").and_then(VmValue::as_dict);
    WriterIdentity {
        writer_id: option_string(writer, "writer_id"),
        stage_id: option_string(writer, "stage_id"),
        session_id: option_string(writer, "session_id"),
        worker_id: option_string(writer, "worker_id"),
    }
}

fn write_options_from_handle(
    handle: &crate::value::DictMap,
) -> Result<BackendWriteOptions, VmError> {
    let policy = handle
        .get("conflict_policy")
        .map(VmValue::display)
        .unwrap_or_else(|| "ignore".to_string());
    Ok(BackendWriteOptions {
        writer: writer_from_handle(handle),
        conflict_policy: ConflictPolicy::parse(&policy)?,
    })
}

fn enforce_conflict_policy(
    handle: &crate::value::DictMap,
    outcome: &backend::BackendWriteOutcome,
) -> Result<(), VmError> {
    let Some(conflict) = &outcome.conflict else {
        return Ok(());
    };
    let options = write_options_from_handle(handle)?;
    let message = format!(
        "agent_state.write: key '{}' was previously written by '{}' and is now being written by '{}'",
        conflict.key,
        conflict.previous.display_name(),
        conflict.current.display_name()
    );
    match options.conflict_policy {
        ConflictPolicy::Ignore => Ok(()),
        ConflictPolicy::Warn => {
            let mut metadata = BTreeMap::new();
            metadata.insert("key".to_string(), serde_json::json!(conflict.key));
            metadata.insert(
                "previous_writer".to_string(),
                writer_json(&conflict.previous),
            );
            metadata.insert("current_writer".to_string(), writer_json(&conflict.current));
            crate::events::log_warn_meta("agent_state.write", &message, metadata);
            Ok(())
        }
        ConflictPolicy::Error => Err(VmError::Runtime(message)),
    }
}

fn writer_json(writer: &WriterIdentity) -> serde_json::Value {
    serde_json::json!({
        "writer_id": writer.writer_id,
        "stage_id": writer.stage_id,
        "session_id": writer.session_id,
        "worker_id": writer.worker_id,
    })
}

fn now_epoch_seconds() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

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

    use std::path::PathBuf;
    use std::process::Command;

    use crate::agent_sessions;

    #[test]
    fn default_session_id_prefers_active_agent_session() {
        agent_sessions::push_current_session("session_test".to_string());
        assert_eq!(default_session_id(), "session_test");
        agent_sessions::pop_current_session();
    }

    #[test]
    fn writer_identity_defaults_to_current_session() {
        agent_sessions::push_current_session("session_writer".to_string());
        let writer = writer_identity(None, None);
        assert_eq!(writer.writer_id.as_deref(), Some("session_writer"));
        assert_eq!(writer.session_id.as_deref(), Some("session_writer"));
        agent_sessions::pop_current_session();
    }

    #[test]
    fn filesystem_round_trip_delete_and_list() {
        let dir = tempfile::tempdir().unwrap();
        let backend = FilesystemBackend::new();
        let scope = BackendScope {
            root: dir.path().to_path_buf(),
            namespace: "session-a".to_string(),
        };
        backend.ensure_scope(&scope).unwrap();
        let options = BackendWriteOptions {
            writer: WriterIdentity {
                writer_id: Some("writer-a".to_string()),
                stage_id: None,
                session_id: Some("session-a".to_string()),
                worker_id: None,
            },
            conflict_policy: ConflictPolicy::Ignore,
        };
        backend.write(&scope, "plan.md", "plan", &options).unwrap();
        backend
            .write(&scope, "evidence/a.json", "{\"ok\":true}", &options)
            .unwrap();

        assert_eq!(
            backend.read(&scope, "plan.md").unwrap().as_deref(),
            Some("plan")
        );
        assert_eq!(
            backend.list(&scope).unwrap(),
            vec!["evidence/a.json".to_string(), "plan.md".to_string()]
        );

        backend.delete(&scope, "plan.md").unwrap();
        assert_eq!(backend.read(&scope, "plan.md").unwrap(), None);
        assert_eq!(
            backend.list(&scope).unwrap(),
            vec!["evidence/a.json".to_string()]
        );
    }

    #[test]
    fn filesystem_rejects_parent_escape_keys() {
        let dir = tempfile::tempdir().unwrap();
        let backend = FilesystemBackend::new();
        let scope = BackendScope {
            root: dir.path().to_path_buf(),
            namespace: "session-b".to_string(),
        };
        backend.ensure_scope(&scope).unwrap();
        let error = backend
            .write(&scope, "../oops", "bad", &BackendWriteOptions::default())
            .unwrap_err();
        assert!(error.to_string().contains("must not escape"));
    }

    #[test]
    fn filesystem_detects_two_writer_conflicts() {
        let dir = tempfile::tempdir().unwrap();
        let backend = FilesystemBackend::new();
        let scope = BackendScope {
            root: dir.path().to_path_buf(),
            namespace: "session-c".to_string(),
        };
        backend.ensure_scope(&scope).unwrap();
        let first = BackendWriteOptions {
            writer: WriterIdentity {
                writer_id: Some("writer-a".to_string()),
                stage_id: Some("stage-a".to_string()),
                session_id: Some("session-c".to_string()),
                worker_id: None,
            },
            conflict_policy: ConflictPolicy::Ignore,
        };
        let second = BackendWriteOptions {
            writer: WriterIdentity {
                writer_id: Some("writer-b".to_string()),
                stage_id: Some("stage-b".to_string()),
                session_id: Some("session-c".to_string()),
                worker_id: None,
            },
            conflict_policy: ConflictPolicy::Error,
        };

        assert!(backend
            .write(&scope, "ledger.md", "one", &first)
            .unwrap()
            .conflict
            .is_none());
        let error = backend
            .write(&scope, "ledger.md", "two", &second)
            .unwrap_err();
        assert!(error.to_string().contains("writer-a"));
        assert!(error.to_string().contains("writer-b"));
        assert_eq!(
            backend.read(&scope, "ledger.md").unwrap().as_deref(),
            Some("one")
        );
    }

    #[test]
    fn crash_helper_aborts_after_temp_write() {
        let helper = std::env::var("HARN_AGENT_STATE_CRASH_HELPER").ok();
        let Some(target) = helper else {
            return;
        };
        let root = std::env::var("HARN_AGENT_STATE_CRASH_ROOT").unwrap();
        let scope = BackendScope {
            root: PathBuf::from(root),
            namespace: "session-crash".to_string(),
        };
        let backend = FilesystemBackend::new();
        backend.ensure_scope(&scope).unwrap();
        let options = BackendWriteOptions {
            writer: WriterIdentity {
                writer_id: Some("writer-crash".to_string()),
                stage_id: None,
                session_id: Some("session-crash".to_string()),
                worker_id: None,
            },
            conflict_policy: ConflictPolicy::Ignore,
        };
        if target == "abort" {
            let _ = backend.write(&scope, "plan.md", "after", &options);
        }
    }

    #[test]
    fn atomic_write_survives_abort_without_partial_content() {
        let exe = std::env::current_exe().unwrap();
        let root = tempfile::tempdir().unwrap();
        let session_dir = root.path().join("session-crash");
        std::fs::create_dir_all(&session_dir).unwrap();
        let target_file = session_dir.join("plan.md");
        std::fs::write(&target_file, "before").unwrap();

        let status = Command::new(exe)
            .arg("--exact")
            .arg("stdlib::agent_state::tests::crash_helper_aborts_after_temp_write")
            .arg("--nocapture")
            .env("HARN_AGENT_STATE_CRASH_HELPER", "abort")
            .env(
                "HARN_AGENT_STATE_CRASH_ROOT",
                root.path().to_string_lossy().into_owned(),
            )
            .env("HARN_AGENT_STATE_ABORT_AFTER_TMP_WRITE", "1")
            .status()
            .unwrap();

        assert!(!status.success());
        assert_eq!(std::fs::read_to_string(&target_file).unwrap(), "before");
    }
}