From 701a7be2919ee1457b810d18eb62420df1de68fc Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Fri, 14 Feb 2025 13:35:37 +0530
Subject: [PATCH] Stop code sandbox on request timeout to allow sandbox process
 restarts

---
 .github/workflows/run_evals.yml      | 2 +-
 src/khoj/processor/tools/run_code.py | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/run_evals.yml b/.github/workflows/run_evals.yml
index 914a0835..6e6f7143 100644
--- a/.github/workflows/run_evals.yml
+++ b/.github/workflows/run_evals.yml
@@ -114,7 +114,7 @@ jobs:
 
           # Start code sandbox
           npm install -g pm2
-          npm run ci --prefix terrarium
+          NODE_ENV=production npm run ci --prefix terrarium
 
           # Wait for server to be ready
           timeout=120
diff --git a/src/khoj/processor/tools/run_code.py b/src/khoj/processor/tools/run_code.py
index ac1463ab..62ffde74 100644
--- a/src/khoj/processor/tools/run_code.py
+++ b/src/khoj/processor/tools/run_code.py
@@ -100,7 +100,13 @@ async def run_code(
             logger.info(f"Executed Code\n----\n{code}\n----\nResult\n----\n{cleaned_result}\n----")
             yield {query: {"code": code, "results": result}}
     except asyncio.TimeoutError as e:
-        raise ValueError(f"Failed to run code for {query} with Timeout error: {e}")
+        # Call the sandbox_url/stop GET API endpoint to stop the code sandbox
+        error = f"Failed to run code for {query} with Timeout error: {e}"
+        try:
+            await aiohttp.ClientSession().get(f"{sandbox_url}/stop", timeout=5)
+        except Exception as e:
+            error += f"\n\nFailed to stop code sandbox with error: {e}"
+        raise ValueError(error)
     except Exception as e:
         raise ValueError(f"Failed to run code for {query} with error: {e}")