Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions crates/bashkit/src/interpreter/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
5 changes: 3 additions & 2 deletions crates/bashkit/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
7 changes: 6 additions & 1 deletion crates/bashkit/src/snapshot.rs
Original file line number Diff line number Diff line change
@@ -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.
//!
Expand Down Expand Up @@ -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(())
}

Expand Down
33 changes: 33 additions & 0 deletions crates/bashkit/tests/integration/snapshot_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Loading