Update tests.py

tests.py

@@ -12,7 +12,7 @@ import json
 import openpyxl
 import shutil
 from google import genai
-
+import pexpect
 client = genai.Client(api_key="AIzaSyDtP05TyoIy9j0uPL7_wLEhgQEE75AZQSc")
 
 source_dir = "/app/uploads/temp"
@@ -210,10 +210,9 @@ def create_code_files(filename: str, code: str) -> dict:
     return {"info":"task completed. The referenced code files were created successfully. "}
     
 @mcp.tool()
-def run_code(python_packages:str,filename: str, code: str,start_cmd:str) -> dict:
+def run_code(python_packages:str,filename: str, code: str,start_cmd:str,forever_cmd:bool) -> dict:
     """  
     Execute code in a controlled environment with package installation and file handling.
-
     Args:
         python_packages[Output an empty string ,if using any other language.]: Space-separated list of packages to install (e.g., "numpy matplotlib").
                         Preinstalled packages: gradio, XlsxWriter, openpyxl.
@@ -221,8 +220,7 @@ def run_code(python_packages:str,filename: str, code: str,start_cmd:str) -> dict
         code:           Full code to write to the file.
         start_cmd:      Command to execute the file (e.g., "python /app/code_interpreter/app.py" 
                         or "bash /app/code_interpreter/app.py").
-
-    Notes:
+        forever_cmd:    If True, the command will run indefinitely.Set to True, when runnig a website/server. If False, the command will time out after 300 second and the result will be returned.
         - All user-uploaded files are in /app/code_interpreter/.
         - After execution, embed a download link (or display images/gifs/videos directly in markdown format) in your response.
     """
@@ -244,9 +242,15 @@ def run_code(python_packages:str,filename: str, code: str,start_cmd:str) -> dict
     f = open(os.path.join(destination_dir, filename), "w")
     f.write(code)
     f.close()
-    run(start_cmd, 300)
-    while stderr=="" and stdout=="":
-        pass
+    if forever_cmd:
+        child=pexpect.spawn(command)
+        time.sleep(5)
+        stdout=str(child.readline().decode())
+        stderr=""
+    else:
+        run(start_cmd, 300)
+        while stderr=="" and stdout=="":
+            pass
     onlyfiles = glob.glob("/app/code_interpreter/*")
     onlyfiles=list(set(onlyfiles)-set(files_list))
     uploaded_filenames=[]
@@ -286,15 +290,13 @@ def run_code_files(start_cmd:str) -> dict:
 @mcp.tool()
 def run_shell_command(cmd:str) -> dict:
     """(cmd:Example- mkdir test.By default , the command is run inside the /app/code_interpreter/ directory.).Remember, the code_interpreter is running on **alpine linux** , so write commands accordingly.Eg-sudo does not work and is not required.."""
-    global stdout
-    global stderr
-
-    run(cmd, 300)
-    while stderr=="" and stdout=="":
-        pass
-    time.sleep(1.5)
+    process = pexpect.spawn(cmd)
+    strLine=""
+    while not process.eof():
+        strLine = strLine+process.readline().decode()
+    output=strLine
     transfer_files()
-    return {"stdout":stdout,"stderr":stderr}
+    return {"output":output}
 
 
 
@@ -426,14 +428,4 @@ def deepthinking3(query:str,info:str) -> dict:
 
 if __name__ == "__main__":
     # Initialize and run the server
-    mcp.run(transport='stdio')
-
-
-# @mcp.tool()
-# def run_website(start_cmd:str,port=8501) -> dict:
-#     """(start_cmd:streamlit run app.py).Always specify sandbox id.Specify port (int) if different from 8501."""
-#     output=sbx.commands.run(start_cmd,sandbox_id)
-#     url = sbx.get_host(port)
-#     info={"info":f"Your Application is live [here](https://{url})"}
-
-#     return info
+    mcp.run(transport='stdio')