Repositories / agent-snapshot.git

agent-snapshot.git

Clone (read-only): git clone http://git.guha-anderson.com/git/agent-snapshot.git

Branch

Allow direct traced command arguments

Author
Arjun Guha <a.guha@northeastern.edu>
Date
2026-05-03 17:42:47 -0400
Commit
073aabc1fbb07071ff839cb89e854e8934fdffc6
README.md
index 0eab6fc..b306e2f 100644
--- a/README.md
+++ b/README.md
@@ -187,15 +187,18 @@ the first run.
 Run a command under Agent Snapshot:
 
 ```bash
-dune exec -- agent-snapshot --output snapshot-dir -- command arg1 arg2
+dune exec -- agent-snapshot --snapshot-dir snapshot-dir command arg1 arg2
 ```
 
 For example:
 
 ```bash
-dune exec -- agent-snapshot --output snapshot-python -- /usr/bin/python3 script.py
+dune exec -- agent-snapshot --snapshot-dir snapshot-python /usr/bin/python3 script.py
 ```
 
+`--output` remains available as an alias for `--snapshot-dir`.
+
+<!-- Restore is temporarily disabled while the bare command form is the default.
 Restore captured final-state blobs and tombstones in place:
 
 ```bash
@@ -205,6 +208,7 @@ dune exec -- agent-snapshot restore snapshot-dir
 Restore only applies files that have blobs and tombstones. Clean Git-tracked
 files and reconstructable system files are represented in the manifest but are
 not rewritten by restore.
+-->
 
 ## Configuration
 
src/ocaml/agent_snapshot.ml
index aedc751..1f6c70b 100644
--- a/src/ocaml/agent_snapshot.ml
+++ b/src/ocaml/agent_snapshot.ml
@@ -777,7 +777,7 @@ let output_arg =
   Arg.(
     value
     & opt (some string) None
-    & info [ "o"; "output" ] ~docv:"SNAPDIR"
+    & info [ "o"; "output"; "snapshot-dir" ] ~docv:"SNAPDIR"
         ~doc:
           "Write the snapshot under this directory. When omitted, a new timestamped directory is created next to the \
            ignore configuration file (the same directory as ignore.json).")
@@ -799,13 +799,44 @@ let snapshot_term : unit Term.t =
   Term.(
     const (fun output cmd ->
       if cmd = [] then
-        failwith "usage: agent-snapshot [--output SNAPDIR] -- command args...";
+        failwith "usage: agent-snapshot [--snapshot-dir SNAPDIR] command args...";
       ignore (run_snapshot output cmd : int))
     $ output_arg
     $ command_arg)
 
+(* Restore is temporarily disabled while the default command form treats the
+   first non-option argument as the traced program name.
 let restore_term : unit Term.t =
   Term.(const (fun dir -> restore_snapshot dir) $ restore_snapdir)
+*)
+
+let option_takes_value (arg : string) : bool =
+  match arg with
+  | "-o" | "--output" | "--snapshot-dir" -> true
+  | _ -> false
+
+let is_help_option (arg : string) : bool =
+  arg = "-h" || arg = "--help" || String.starts_with ~prefix:"--help=" arg
+
+let normalize_snapshot_argv (argv : string array) : string array =
+  let len = Array.length argv in
+  let rec loop i =
+    if i >= len then None
+    else
+      let arg = argv.(i) in
+      if arg = "--" || is_help_option arg then None
+      else if option_takes_value arg then loop (i + 2)
+      else if String.starts_with ~prefix:"--output=" arg || String.starts_with ~prefix:"--snapshot-dir=" arg then loop (i + 1)
+      else if String.starts_with ~prefix:"-" arg then loop (i + 1)
+      else Some i
+  in
+  match loop 1 with
+  | None -> argv
+  | Some command_index ->
+      Array.init (len + 1) (fun i ->
+          if i < command_index then argv.(i)
+          else if i = command_index then "--"
+          else argv.(i - 1))
 
 let cmd_main : unit Cmd.t =
   let doc = "Filesystem snapshotter for traced commands" in
@@ -814,23 +845,25 @@ let cmd_main : unit Cmd.t =
       `S Manpage.s_description;
       `P "Traces a command with ptrace and records a manifest plus file blobs.";
       `S Manpage.s_examples;
-      `P "$(tool) $(b,--output) /tmp/snap $(b,--) make all";
+      `P "$(tool) $(b,--snapshot-dir) /tmp/snap make all";
       `P "$(tool) $(b,--output) /tmp/snap make all";
-      `P "$(tool) $(b,--) make all";
-      `P "$(tool) $(b,restore) /tmp/snap";
+      `P "$(tool) make all";
     ]
   in
   let main_info = Cmd.info "agent-snapshot" ~doc ~man in
+  (*
   let restore_info =
     Cmd.info "restore" ~docs:Manpage.s_commands
       ~doc:"Restore blobbed files and tombstones from a snapshot directory."
   in
   let restore_cmd = Cmd.v restore_info restore_term in
   Cmd.group main_info ~default:snapshot_term [ restore_cmd ]
+  *)
+  Cmd.v main_info snapshot_term
 
 let main () : unit =
   try
-    let rc = Cmd.eval ~catch:false cmd_main in
+    let rc = Cmd.eval ~catch:false ~argv:(normalize_snapshot_argv Sys.argv) cmd_main in
     Ocaml_git.shutdown ();
     exit rc
   with exn ->
tests/test_agent_snapshot.py
index b224fa8..37becb8 100644
--- a/tests/test_agent_snapshot.py
+++ b/tests/test_agent_snapshot.py
@@ -112,9 +112,9 @@ class Snapshot:
 
 def capture(tmp_path: Path, *command: str) -> Snapshot:
     # Every test gets a fresh bundle directory so stale blobs cannot mask capture
-    # bugs. The command is passed after -- to exercise the intended CLI parsing.
+    # bugs.
     out = tmp_path / "snapshot"
-    run([str(BIN), "--output", str(out), "--", *command])
+    run([str(BIN), "--snapshot-dir", str(out), *command])
     return Snapshot(out)
 
 
@@ -122,7 +122,7 @@ def test_missing_ignore_config_creates_defaults(tmp_path, ignore_config):
     ignore_config.unlink()
     assert not ignore_config.exists()
     out = tmp_path / "snapshot"
-    run([str(BIN), "--output", str(out), "--", PYTHON, "test_programs/read_clean.py"])
+    run([str(BIN), "--snapshot-dir", str(out), PYTHON, "test_programs/read_clean.py"])
 
     assert ignore_config.exists()
     data = json.loads(ignore_config.read_text())
@@ -185,7 +185,7 @@ def test_ignore_config_expands_xdg_config_home_fallback(tmp_path, ignore_config)
     env["HOME"] = str(home)
     out = tmp_path / "snapshot"
     subprocess.run(
-        [str(BIN), "--output", str(out), "--", PYTHON, "test_programs/read_ignored_paths.py"],
+        [str(BIN), "--snapshot-dir", str(out), PYTHON, "test_programs/read_ignored_paths.py"],
         cwd=ROOT,
         text=True,
         check=True,
@@ -303,25 +303,39 @@ def test_fork_usr_and_directory_traversal(tmp_path):
     assert "directory" in directory["operations"]
 
 
-def test_restore_applies_final_state(tmp_path):
-    # Restore is intentionally tested from a damaged filesystem state rather than
-    # immediately after capture. That proves the bundle contains enough payload
-    # to recreate final captured files and enough tombstone information to remove
-    # files that should not exist after the traced command.
-    (WORKTREE / "dirty.txt").write_text("changed before capture\n")
-    (WORKTREE / "untracked_runtime.txt").write_text("untracked input\n")
-    (WORKTREE / "deleted_by_program.txt").write_text("delete me\n")
-    snap = capture(tmp_path, PYTHON, "test_programs/dirty_untracked_write.py")
+def test_restore_word_is_treated_as_traced_command(tmp_path, monkeypatch):
+    restore_bin = tmp_path / "bin" / "restore"
+    restore_bin.parent.mkdir()
+    restore_bin.write_text("#!/bin/sh\nprintf 'restore program ran\\n'\n")
+    restore_bin.chmod(0o755)
+    monkeypatch.setenv("PATH", f"{restore_bin.parent}:{os.environ['PATH']}")
+    out = tmp_path / "snapshot"
+
+    completed = subprocess.run(
+        [str(BIN), "--snapshot-dir", str(out), "restore"],
+        cwd=ROOT,
+        text=True,
+        check=True,
+        capture_output=True,
+    )
+
+    assert completed.stdout == "restore program ran\n"
+    assert (out / "manifest.json").exists()
 
-    shutil.rmtree(WORKTREE)
-    WORKTREE.mkdir()
-    (WORKTREE / "created_by_program.txt").write_text("wrong\n")
-    (WORKTREE / "deleted_by_program.txt").write_text("should disappear\n")
 
-    run([str(BIN), "restore", str(snap.path)])
+def test_traced_command_options_do_not_need_separator(tmp_path):
+    out = tmp_path / "snapshot"
+
+    completed = subprocess.run(
+        [str(BIN), "--snapshot-dir", str(out), PYTHON, "-c", "print('direct command')"],
+        cwd=ROOT,
+        text=True,
+        check=True,
+        capture_output=True,
+    )
 
-    assert (WORKTREE / "created_by_program.txt").read_text() == "created final\n"
-    assert not (WORKTREE / "deleted_by_program.txt").exists()
+    assert completed.stdout == "direct command\n"
+    assert Snapshot(out).manifest["command"] == [PYTHON, "-c", "print('direct command')"]
 
 
 def test_rename_records_source_tombstone_and_destination_content(tmp_path):