Update tests.py

tests.py

@@ -13,6 +13,7 @@ import openpyxl
 import shutil
 from google import genai
 import pexpect
+
 client = genai.Client(api_key="AIzaSyDtP05TyoIy9j0uPL7_wLEhgQEE75AZQSc")
 
 source_dir = "/app/uploads/temp"
@@ -26,8 +27,6 @@ os.environ["OPENROUTER_API_KEY"] = "sk-or-v1-019ff564f86e6d14b2a78a78be1fb88724e
 mcp = FastMCP("code_sandbox")
 data={}
 result=""
-stdout=""
-stderr=""
 import requests
 import os
 from bs4 import BeautifulSoup  # For parsing HTML
@@ -140,17 +139,29 @@ chat_id = "5075390513"
 from requests_futures.sessions import FuturesSession
 session = FuturesSession()
 
-def run(cmd, timeout_sec):
-    global stdout 
-    global stderr
-    proc = Popen(shlex.split(cmd), stdout=PIPE, stderr=PIPE,cwd="/app/code_interpreter/")
-    timer = Timer(timeout_sec, proc.kill)
-    try:
-        timer.start()
-        stdout, stderr = proc.communicate()
-    finally:
-        timer.cancel()
-
+def run(cmd, timeout_sec,forever_cmd=False):
+    t=time.time()
+    child = pexpect.spawn("bash")
+    output=""
+    command="cd /app/code_interpreter/ && "+cmd
+
+    child.sendline('PROMPT_COMMAND="echo END"')
+    child.readline().decode()
+    child.readline().decode()
+
+    child.sendline(command)
+
+    while (not child.eof() ) and (time.time()-t<90):
+        x=child.readline().decode()
+        output=output+x
+        print(x)
+        if "END" in x :
+            output=output.replace("END","")
+            child.close()
+            break
+        if forever_cmd:
+            break
+    return output
 
 @mcp.tool()
 def analyse_audio(audiopath,query) -> dict:
@@ -221,13 +232,12 @@ def run_code(python_packages:str,filename: str, code: str,start_cmd:str,forever_
         start_cmd:      Command to execute the file (e.g., "python /app/code_interpreter/app.py" 
                         or "bash /app/code_interpreter/app.py").
         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.
+    Notes:
         - 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.
     """
 
     global files_list
-    global stdout
-    global stderr
     package_names = python_packages.strip()
     command="pip install"
     if not package_names:
@@ -242,15 +252,8 @@ def run_code(python_packages:str,filename: str, code: str,start_cmd:str,forever_
     f = open(os.path.join(destination_dir, filename), "w")
     f.write(code)
     f.close()
-    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
+    stdot=run(start_cmd, 300,forever_cmd)
+
     onlyfiles = glob.glob("/app/code_interpreter/*")
     onlyfiles=list(set(onlyfiles)-set(files_list))
     uploaded_filenames=[]
@@ -261,19 +264,14 @@ def run_code(python_packages:str,filename: str, code: str,start_cmd:str,forever_
         except:
             pass
     files_list=onlyfiles
-    return {"info":"task completed","stdout":stdout,"stderr":stderr,"Files_download_link":uploaded_filenames}
+    return {"output":stdot,"Files_download_link":uploaded_filenames}
 
 
 @mcp.tool()
 def run_code_files(start_cmd:str) -> dict:
     """(start_cmd:Example-  python /app/code_interpreter/app.py or bash /app/code_interpreter/app.py).The files must be inside the /app/code_interpreter directory."""
     global files_list
-    global stdout
-    global stderr
-    run(start_cmd, 300)
-    while stderr=="" and stdout=="":
-        pass
-    time.sleep(1.5)
+    stdot=run(start_cmd, 300)
     onlyfiles = glob.glob("/app/code_interpreter/*")
     onlyfiles=list(set(onlyfiles)-set(files_list))
     uploaded_filenames=[]
@@ -284,19 +282,16 @@ def run_code_files(start_cmd:str) -> dict:
         except:
             pass
     files_list=onlyfiles
-    return {"stdout":stdout,"stderr":stderr,"Files_download_link":uploaded_filenames}
+    return {"output":stdot,"Files_download_link":uploaded_filenames}
 
 
 @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.."""
-    process = pexpect.spawn(cmd)
-    strLine=""
-    while not process.eof():
-        strLine = strLine+process.readline().decode()
-    output=strLine
+    global stdout
+    run(cmd, 300)
     transfer_files()
-    return {"output":output}
+    return {"output":stdout}
 
 
 
@@ -309,13 +304,11 @@ def install_python_packages(python_packages:str) -> dict:
     if not package_names:
         return
 
-    run(
+    stdot=run(
         f"{command} --break-system-packages {package_names}", timeout_sec=300
     )
-    while stderr=="" and stdout=="":
-        pass
-    time.sleep(2)
-    return {"stdout":stdout,"stderr":stderr,"info":"Ran package installation command"}
+
+    return {"stdout":stdot,"info":"Ran package installation command"}
 
 @mcp.tool()
 def get_youtube_transcript(videoid:str) -> dict:
@@ -428,4 +421,5 @@ def deepthinking3(query:str,info:str) -> dict:
 
 if __name__ == "__main__":
     # Initialize and run the server
-    mcp.run(transport='stdio')
+    mcp.run(transport='stdio')
+