support passing morph snapshot id

model_tools.py

@@ -518,8 +518,9 @@ def handle_terminal_function_call(function_name: str, function_args: Dict[str, A
         background = function_args.get("background", False)
         idle_threshold = function_args.get("idle_threshold", 5.0)
         timeout = function_args.get("timeout")
+        snapshot_id = function_args.get("snapshot_id")
         # Session management is handled internally - don't pass session_id from model
-        return terminal_tool(command, input_keys, None, background, idle_threshold, timeout)
+        return terminal_tool(command, input_keys, None, background, idle_threshold, timeout, snapshot_id=snapshot_id)
     
     else:
         return json.dumps({"error": f"Unknown terminal function: {function_name}"})

run_agent.py

@@ -51,7 +51,8 @@ class AIAgent:
         disabled_tools: List[str] = None,
         enabled_toolsets: List[str] = None,
         disabled_toolsets: List[str] = None,
-        save_trajectories: bool = False
+        save_trajectories: bool = False,
+        morph_snapshot_id: str | None = None,
     ):
         """
         Initialize the AI Agent.
@@ -67,11 +68,13 @@ class AIAgent:
             enabled_toolsets (List[str]): Only enable tools from these toolsets (optional)
             disabled_toolsets (List[str]): Disable tools from these toolsets (optional)
             save_trajectories (bool): Whether to save conversation trajectories to JSONL files (default: False)
+            morph_snapshot_id (str | None): Morph Cloud snapshot id from which to start terminal tool
         """
         self.model = model
         self.max_iterations = max_iterations
         self.tool_delay = tool_delay
         self.save_trajectories = save_trajectories
+        self.morph_snapshot_id = morph_snapshot_id
         
         # Store tool filtering options
         self.enabled_tools = enabled_tools
@@ -389,6 +392,9 @@ class AIAgent:
                         
                         print(f"  📞 Tool {i}: {function_name}({list(function_args.keys())})")
 
+                        if function_name == "terminal" and self.morph_snapshot_id is not None:
+                            function_args["snapshot_id"] = self.morph_snapshot_id
+                        
                         # Execute the tool
                         function_result = handle_function_call(function_name, function_args)
                         
@@ -480,7 +486,8 @@ def main(
     enabled_toolsets: str = None,
     disabled_toolsets: str = None,
     list_tools: bool = False,
-    save_trajectories: bool = False
+    save_trajectories: bool = False,
+    morph_snapshot_id: str | None = None,
 ):
     """
     Main function for running the agent directly.
@@ -497,6 +504,7 @@ def main(
         disabled_toolsets (str): Comma-separated list of toolsets to disable (e.g., "terminal_tools")
         list_tools (bool): Just list available tools and exit
         save_trajectories (bool): Save conversation trajectories to JSONL files. Defaults to False.
+        morph_snapshot_id (str | None): Morph Cloud snapshot id to start terminal tool from
     """
     print("🤖 AI Agent with Tool Calling")
     print("=" * 50)
@@ -573,7 +581,8 @@ def main(
             disabled_tools=disabled_tools_list,
             enabled_toolsets=enabled_toolsets_list,
             disabled_toolsets=disabled_toolsets_list,
-            save_trajectories=save_trajectories
+            save_trajectories=save_trajectories,
+            morph_snapshot_id=morph_snapshot_id
         )
     except RuntimeError as e:
         print(f"❌ Failed to initialize agent: {e}")

terminal_tool.py

@@ -78,7 +78,8 @@ def terminal_tool(
     session_id: Optional[str] = None,
     background: bool = False,
     idle_threshold: float = 5.0,
-    timeout: Optional[int] = None
+    timeout: Optional[int] = None,
+    snapshot_id: str | None = None,
 ) -> str:
     """
     Execute a command on a Morph VM with optional interactive session support.
@@ -136,7 +137,7 @@ def terminal_tool(
         )
         
         # Execute with lifecycle management
-        result = run_tool_with_lifecycle_management(tool_call)
+        result = run_tool_with_lifecycle_management(tool_call, snapshot_id=snapshot_id)
         
         # Format the result with all possible fields
         # Map hecate's "stdout" to "output" for compatibility