feat: add flag to disable sandbox replay

[gwarmstrong]

Summary by CodeRabbit

• New Features
  • Added a configuration option to control whether session history is restored after sandbox restarts.
  • Session restoration can now be disabled globally via a new command-line flag in the Python execution tool.
  • When restoration is disabled, previous session history is cleared and the execution environment is notified.

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic bug: the elif doesn't cover the case where session restore is disabled but new_session_created is False

When disable_session_restore=True and new_session_created=False (normal execution without restart), the code falls through without executing anything. The warning message and history clearing should still happen when a restart is detected, but successful code history should never be stored when restore is disabled (which is handled correctly on line 268).

However, the logic structure with elif prevents the code block on lines 243-250 from executing when disable_session_restore=True and the session wasn't recreated. This means normal code execution might be skipped.

Suggested change

	elif session_id is not None and new_session_created and self.disable_session_restore:
	# Session was recreated but restore is disabled — clear stale history and warn the model.
	self.session_histories.pop(session_id_str, None)
	self.session_histories.pop(request_session_id_str, None)
	LOG.warning("Session %s was recreated but restore is disabled; history cleared", session_id)
	output["stderr"] = (
	"RuntimeError: Sandbox state restoration failed after the execution worker restarted. "
	"The interactive session history has been cleared; "
	"please re-run the last code block without relying on prior state.\n"
	) + output.get("stderr", "")
	elif session_id is not None and new_session_created and self.disable_session_restore:
	# Session was recreated but restore is disabled — clear stale history and warn the model.
	self.session_histories.pop(session_id_str, None)
	self.session_histories.pop(request_session_id_str, None)
	LOG.warning("Session %s was recreated but restore is disabled; history cleared", session_id)
	output["stderr"] = (
	"RuntimeError: Sandbox state restoration failed after the execution worker restarted. "
	"The interactive session history has been cleared; "
	"please re-run the last code block without relying on prior state.\n"
	) + output.get("stderr", "")
	# Execute the new code on the fresh session
	exec_request = self._prepare_request(
	generated_code, timeout, language, std_input, max_output_characters, traceback_verbosity
	)
	exec_request["session_id"] = request_session_id_str
	try:
	output = await self._send_request(exec_request, timeout)
	except httpx.TimeoutException:
	output = {"process_status": "timeout", "stdout": "", "stderr": "Client timed out\n"}

[coderabbitai]

📝 Walkthrough

Walkthrough

Introduces a disable_session_restore parameter to the Sandbox class to toggle session history replay on worker restart. The flag can be set via configuration or environment variable. Adds corresponding command-line argument in the Python MCP server to expose this functionality. When disabled, clears stale session histories and injects warning messages to notify about restoration being skipped.

Changes

Cohort / File(s)	Summary
Session Restoration Control nemo_skills/code_execution/sandbox.py	Added disable_session_restore boolean parameter to Sandbox.__init__ (default False, reads from environment variable NEMO_SKILLS_DISABLE_SESSION_RESTORE). Modified session recreation and history append logic to conditionally restore/store history based on this flag. When disabled, clears stale histories and injects warning messages into stderr.
CLI Flag Propagation nemo_skills/mcp/servers/python_tool.py	Added --disable-session-restore command-line argument that sets sandbox_cfg["disable_session_restore"] = True before sandbox initialization, enabling users to disable session restoration at runtime.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: adding a flag to disable sandbox replay/session restoration across the codebase.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch georgea/disable-session-restore

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@nemo_skills/code_execution/sandbox.py`:
- Around line 71-73: The env-var check for initializing
self.disable_session_restore is case-sensitive and treats "False"/"FALSE"
incorrectly; update the expression in the initializer where
disable_session_restore is assigned (reference: disable_session_restore and
self.disable_session_restore) to normalize the environment value with .lower()
(e.g., os.getenv("NEMO_SKILLS_DISABLE_SESSION_RESTORE", "").lower()) before
testing membership against ("", "0", "false") so values like "False" and "FALSE"
are handled correctly.

🧹 Nitpick comments (1)

nemo_skills/code_execution/sandbox.py (1)

252-261: Warning message says "restoration failed" but restore was intentionally disabled.

The stderr message is identical to the one used when restoration genuinely fails (Line 231). When restore is deliberately disabled, a more accurate phrasing (e.g., "restoration is disabled") would avoid confusion in logs/debugging. Consider differentiating the two messages.

Also, unlike the first branch (Line 185), this elif doesn't exclude the process_status == "timeout" case, so on a timeout with a new session the restoration-disabled warning is prepended to the timeout message. This is minor since the net effect (history cleared, user warned) is correct, but could be confusing.

♻️ Suggested message differentiation

-            output["stderr"] = (
-                "RuntimeError: Sandbox state restoration failed after the execution worker restarted. "
-                "The interactive session history has been cleared; "
-                "please re-run the last code block without relying on prior state.\n"
-            ) + output.get("stderr", "")
+            output["stderr"] = (
+                "RuntimeError: Session restoration is disabled. The execution worker restarted "
+                "and the interactive session history has been cleared; "
+                "please re-run the last code block without relying on prior state.\n"
+            ) + output.get("stderr", "")

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Case-insensitive handling missing for the env var.

"False" and "FALSE" are not in ("", "0", "false"), so setting NEMO_SKILLS_DISABLE_SESSION_RESTORE=False would paradoxically enable the flag. Normalize with .lower() before the membership check.

🐛 Proposed fix

-        self.disable_session_restore = disable_session_restore or os.getenv(
-            "NEMO_SKILLS_DISABLE_SESSION_RESTORE", ""
-        ) not in ("", "0", "false")
+        self.disable_session_restore = disable_session_restore or os.getenv(
+            "NEMO_SKILLS_DISABLE_SESSION_RESTORE", ""
+        ).lower() not in ("", "0", "false")

🤖 Prompt for AI Agents

In `@nemo_skills/code_execution/sandbox.py` around lines 71 - 73, The env-var
check for initializing self.disable_session_restore is case-sensitive and treats
"False"/"FALSE" incorrectly; update the expression in the initializer where
disable_session_restore is assigned (reference: disable_session_restore and
self.disable_session_restore) to normalize the environment value with .lower()
(e.g., os.getenv("NEMO_SKILLS_DISABLE_SESSION_RESTORE", "").lower()) before
testing membership against ("", "0", "false") so values like "False" and "FALSE"
are handled correctly.

[Kipok]

I guess this wouldn't harm @gwarmstrong

[coderabbitai]

Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!

[gwarmstrong]

@Kipok fixed

[greptile-apps]

_{No files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}