From d3db6d8cbed98cd3f2426cd2e6e585372bca11ec Mon Sep 17 00:00:00 2001 From: Mykhailo Chalyi Date: Mon, 1 Jun 2026 23:16:02 -0500 Subject: [PATCH] fix(snapshot): preserve counters on resume --- crates/bashkit/src/interpreter/mod.rs | 6 ++-- crates/bashkit/src/lib.rs | 5 +-- crates/bashkit/src/snapshot.rs | 7 +++- .../tests/integration/snapshot_tests.rs | 33 +++++++++++++++++++ 4 files changed, 45 insertions(+), 6 deletions(-) diff --git a/crates/bashkit/src/interpreter/mod.rs b/crates/bashkit/src/interpreter/mod.rs index e0df17d67..30a8f8ec0 100644 --- a/crates/bashkit/src/interpreter/mod.rs +++ b/crates/bashkit/src/interpreter/mod.rs @@ -1946,10 +1946,10 @@ impl Interpreter { &self.counters } - /// Restore session-level counters from a snapshot. + /// Merge session-level counters from a snapshot without lowering live usage. pub fn restore_session_counters(&mut self, session_commands: u64, session_exec_calls: u64) { - self.counters.session_commands = session_commands; - self.counters.session_exec_calls = session_exec_calls; + self.counters.session_commands = self.counters.session_commands.max(session_commands); + self.counters.session_exec_calls = self.counters.session_exec_calls.max(session_exec_calls); } /// Set an output callback for streaming output during execution. diff --git a/crates/bashkit/src/lib.rs b/crates/bashkit/src/lib.rs index 3d9dfd8a1..f00470774 100644 --- a/crates/bashkit/src/lib.rs +++ b/crates/bashkit/src/lib.rs @@ -1159,10 +1159,11 @@ impl Bash { (c.session_commands, c.session_exec_calls) } - /// Restore session-level counters to resume a session across Bash instances. + /// Merge session-level counters to resume a session across Bash instances. /// /// This is used by external tool hosts to persist cumulative session counters - /// across fresh Bash instances created per tool call. + /// across fresh Bash instances created per tool call. Counters are monotonic: + /// restoring lower values never reduces already-consumed session budget. pub fn restore_session_counters(&mut self, session_commands: u64, session_exec_calls: u64) { self.interpreter .restore_session_counters(session_commands, session_exec_calls); diff --git a/crates/bashkit/src/snapshot.rs b/crates/bashkit/src/snapshot.rs index 1c829af4e..dbbd4265a 100644 --- a/crates/bashkit/src/snapshot.rs +++ b/crates/bashkit/src/snapshot.rs @@ -1,7 +1,7 @@ // Decision: Snapshot format uses serde_json for Phase 1 (debuggable, human-readable). // Phase 2 can add bincode/postcard for compactness. // VFS contents are included by default; SnapshotOptions can opt out for shell-only restores. -// Session counters are serialized for observability; restore paths do not trust them. +// Session counters are serialized and restored monotonically so snapshot/resume cannot reset budgets. //! Snapshot/resume — serialize interpreter state between `exec()` calls. //! @@ -335,6 +335,11 @@ impl crate::Bash { // Shell state cannot fail past validation, and the VFS has already // been restored atomically (or rejected) above. self.interpreter.restore_shell_state(&snap.shell); + // Session counters are part of session accounting. Merge them + // monotonically: authenticated snapshot/resume carries used budget + // forward, while tampered unkeyed bytes cannot lower live counters. + self.interpreter + .restore_session_counters(snap.session_commands, snap.session_exec_calls); Ok(()) } diff --git a/crates/bashkit/tests/integration/snapshot_tests.rs b/crates/bashkit/tests/integration/snapshot_tests.rs index a61610ad4..0649495e2 100644 --- a/crates/bashkit/tests/integration/snapshot_tests.rs +++ b/crates/bashkit/tests/integration/snapshot_tests.rs @@ -567,6 +567,39 @@ async fn snapshot_restore_does_not_reset_session_exec_limit_with_tampered_counte ); } +#[tokio::test] +async fn keyed_snapshot_restore_carries_session_exec_budget_forward() { + let key = b"session-budget-hmac-key"; + let session_limits = SessionLimits::new().max_exec_calls(2); + let mut bash = Bash::builder() + .session_limits(session_limits.clone()) + .build(); + bash.exec("echo first").await.unwrap(); + let bytes = bash.snapshot_to_bytes_keyed(key).unwrap(); + + let mut restored = Bash::builder().session_limits(session_limits).build(); + restored.restore_snapshot_keyed(&bytes, key).unwrap(); + assert_eq!(restored.session_counters().1, 1); + + restored.exec("echo second").await.unwrap(); + let third = restored.exec("echo third").await; + assert!( + third.is_err(), + "authenticated snapshot resume must not grant a fresh exec-call budget" + ); +} + +#[tokio::test] +async fn from_snapshot_keyed_restores_session_counters() { + let key = b"session-counter-hmac-key"; + let mut bash = Bash::new(); + bash.exec("echo first").await.unwrap(); + let bytes = bash.snapshot_to_bytes_keyed(key).unwrap(); + + let restored = Bash::from_snapshot_keyed(&bytes, key).unwrap(); + assert_eq!(restored.session_counters(), bash.session_counters()); +} + #[tokio::test] async fn snapshot_restore_rejects_tampered_shell_state_that_exceeds_memory_limits() { let mut src = Bash::new();