diff --git a/rust/nix-flake/src/lib.rs b/rust/nix-flake/src/lib.rs
index fc7ae72..77ccca7 100644
--- a/rust/nix-flake/src/lib.rs
+++ b/rust/nix-flake/src/lib.rs
@@ -1,4 +1,4 @@
-use std::ptr::NonNull;
+use std::{ffi::CString, ptr::NonNull};
 
 use anyhow::{Context as _, Result};
 use nix_c_raw as raw;
@@ -145,6 +145,45 @@ impl FlakeLockFlags {
         let s = unsafe { context::check_call!(raw::flake_lock_flags_new(&mut ctx, settings.ptr)) }?;
         Ok(FlakeLockFlags { ptr: s })
     }
+    pub fn set_mode_write_as_needed(&mut self) -> Result<()> {
+        let mut ctx = Context::new();
+        unsafe {
+            context::check_call!(raw::flake_lock_flags_set_mode_write_as_needed(
+                &mut ctx, self.ptr
+            ))
+        }?;
+        Ok(())
+    }
+    pub fn set_mode_check(&mut self) -> Result<()> {
+        let mut ctx = Context::new();
+        unsafe { context::check_call!(raw::flake_lock_flags_set_mode_check(&mut ctx, self.ptr)) }?;
+        Ok(())
+    }
+    pub fn set_mode_virtual(&mut self) -> Result<()> {
+        let mut ctx = Context::new();
+        unsafe {
+            context::check_call!(raw::flake_lock_flags_set_mode_virtual(&mut ctx, self.ptr))
+        }?;
+        Ok(())
+    }
+    pub fn add_input_override(
+        &mut self,
+        override_path: &str,
+        override_ref: &FlakeReference,
+    ) -> Result<()> {
+        let mut ctx = Context::new();
+        unsafe {
+            context::check_call!(raw::flake_lock_flags_add_input_override(
+                &mut ctx,
+                self.ptr,
+                CString::new(override_path)
+                    .context("Failed to create CString for override_path")?
+                    .as_ptr(),
+                override_ref.ptr.as_ptr()
+            ))
+        }?;
+        Ok(())
+    }
 }
 
 pub struct LockedFlake {
@@ -294,4 +333,250 @@ mod tests {
         drop(tmp_dir);
         drop(gc_registration);
     }
+
+    #[test]
+    fn flake_lock_load_flake_with_flags() {
+        init();
+        let gc_registration = gc_register_my_thread();
+        let store = Store::open(None, []).unwrap();
+        let flake_settings = FlakeSettings::new().unwrap();
+        let mut eval_state = EvalStateBuilder::new(store)
+            .unwrap()
+            .flakes(&flake_settings)
+            .unwrap()
+            .build()
+            .unwrap();
+
+        let tmp_dir = tempfile::tempdir().unwrap();
+
+        let flake_dir_a = tmp_dir.path().join("a");
+        let flake_dir_b = tmp_dir.path().join("b");
+        let flake_dir_c = tmp_dir.path().join("c");
+
+        std::fs::create_dir_all(&flake_dir_a).unwrap();
+        std::fs::create_dir_all(&flake_dir_b).unwrap();
+        std::fs::create_dir_all(&flake_dir_c).unwrap();
+
+        let flake_dir_a_str = flake_dir_a.to_str().unwrap();
+        let flake_dir_c_str = flake_dir_c.to_str().unwrap();
+
+        // a
+        std::fs::write(
+            &tmp_dir.path().join("a/flake.nix"),
+            r#"
+            {
+                inputs.b.url = "@flake_dir_b@";
+                outputs = { b, ... }: {
+                    hello = b.hello;
+                };
+            }
+            "#
+            .replace("@flake_dir_b@", flake_dir_b.to_str().unwrap()),
+        )
+        .unwrap();
+
+        // b
+        std::fs::write(
+            &tmp_dir.path().join("b/flake.nix"),
+            r#"
+            {
+                outputs = { ... }: {
+                    hello = "BOB";
+                };
+            }
+            "#,
+        )
+        .unwrap();
+
+        // c
+        std::fs::write(
+            &tmp_dir.path().join("c/flake.nix"),
+            r#"
+            {
+                outputs = { ... }: {
+                    hello = "Claire";
+                };
+            }
+            "#,
+        )
+        .unwrap();
+
+        let mut flake_lock_flags = FlakeLockFlags::new(&flake_settings).unwrap();
+
+        let mut flake_reference_parse_flags =
+            FlakeReferenceParseFlags::new(&flake_settings).unwrap();
+
+        flake_reference_parse_flags
+            .set_base_directory(tmp_dir.path().to_str().unwrap())
+            .unwrap();
+
+        let (flake_ref_a, fragment) = FlakeReference::parse_with_fragment(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &flake_reference_parse_flags,
+            &format!("path:{}", &flake_dir_a_str),
+        )
+        .unwrap();
+
+        assert_eq!(fragment, "");
+
+        // Step 1: Do not update (check), fails
+
+        flake_lock_flags.set_mode_check().unwrap();
+
+        let locked_flake = LockedFlake::lock(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &eval_state,
+            &flake_lock_flags,
+            &flake_ref_a,
+        );
+        // Has not been locked and would need to write a lock file.
+        assert!(locked_flake.is_err());
+        let saved_err = match locked_flake {
+            Ok(_) => panic!("Expected error, but got Ok"),
+            Err(e) => e,
+        };
+
+        // Step 2: Update but do not write, succeeds
+        flake_lock_flags.set_mode_virtual().unwrap();
+
+        let locked_flake = LockedFlake::lock(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &eval_state,
+            &flake_lock_flags,
+            &flake_ref_a,
+        )
+        .unwrap();
+
+        let outputs = locked_flake
+            .outputs(&flake_settings, &mut eval_state)
+            .unwrap();
+
+        let hello = eval_state.require_attrs_select(&outputs, &"hello").unwrap();
+        let hello = eval_state.require_string(&hello).unwrap();
+
+        assert_eq!(hello, "BOB");
+
+        // Step 3: The lock was not written, so Step 1 would fail again
+
+        flake_lock_flags.set_mode_check().unwrap();
+
+        let locked_flake = LockedFlake::lock(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &eval_state,
+            &flake_lock_flags,
+            &flake_ref_a,
+        );
+        // Has not been locked and would need to write a lock file.
+        assert!(locked_flake.is_err());
+        match locked_flake {
+            Ok(_) => panic!("Expected error, but got Ok"),
+            Err(e) => {
+                assert_eq!(e.to_string(), saved_err.to_string());
+            }
+        };
+
+        // Step 4: Update and write, succeeds
+
+        flake_lock_flags.set_mode_write_as_needed().unwrap();
+
+        let locked_flake = LockedFlake::lock(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &eval_state,
+            &flake_lock_flags,
+            &flake_ref_a,
+        )
+        .unwrap();
+
+        let outputs = locked_flake
+            .outputs(&flake_settings, &mut eval_state)
+            .unwrap();
+        let hello = eval_state.require_attrs_select(&outputs, &"hello").unwrap();
+        let hello = eval_state.require_string(&hello).unwrap();
+        assert_eq!(hello, "BOB");
+
+        // Step 5: Lock was written, so Step 1 succeeds
+
+        flake_lock_flags.set_mode_check().unwrap();
+
+        let locked_flake = LockedFlake::lock(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &eval_state,
+            &flake_lock_flags,
+            &flake_ref_a,
+        )
+        .unwrap();
+
+        let outputs = locked_flake
+            .outputs(&flake_settings, &mut eval_state)
+            .unwrap();
+        let hello = eval_state.require_attrs_select(&outputs, &"hello").unwrap();
+        let hello = eval_state.require_string(&hello).unwrap();
+        assert_eq!(hello, "BOB");
+
+        // Step 6: Lock with override, do not write
+
+        // This shouldn't matter; write_as_needed will be overridden
+        flake_lock_flags.set_mode_write_as_needed().unwrap();
+
+        let (flake_ref_c, fragment) = FlakeReference::parse_with_fragment(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &flake_reference_parse_flags,
+            &format!("path:{}", &flake_dir_c_str),
+        )
+        .unwrap();
+        assert_eq!(fragment, "");
+
+        flake_lock_flags
+            .add_input_override("b", &flake_ref_c)
+            .unwrap();
+
+        let locked_flake = LockedFlake::lock(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &eval_state,
+            &flake_lock_flags,
+            &flake_ref_a,
+        )
+        .unwrap();
+
+        let outputs = locked_flake
+            .outputs(&flake_settings, &mut eval_state)
+            .unwrap();
+        let hello = eval_state.require_attrs_select(&outputs, &"hello").unwrap();
+        let hello = eval_state.require_string(&hello).unwrap();
+        assert_eq!(hello, "Claire");
+
+        // Can't delete overrides, so trash it
+        let mut flake_lock_flags = FlakeLockFlags::new(&flake_settings).unwrap();
+
+        // Step 7: Override was not written; lock still points to b
+
+        flake_lock_flags.set_mode_check().unwrap();
+
+        let locked_flake = LockedFlake::lock(
+            &FetchersSettings::new().unwrap(),
+            &flake_settings,
+            &eval_state,
+            &flake_lock_flags,
+            &flake_ref_a,
+        )
+        .unwrap();
+
+        let outputs = locked_flake
+            .outputs(&flake_settings, &mut eval_state)
+            .unwrap();
+        let hello = eval_state.require_attrs_select(&outputs, &"hello").unwrap();
+        let hello = eval_state.require_string(&hello).unwrap();
+        assert_eq!(hello, "BOB");
+
+        drop(tmp_dir);
+        drop(gc_registration);
+    }
 }