Skip to content

Commit bc489d1

Browse files
committed
```
refactor(utils): change exec_code to use subprocess with timeout protection - Replace threading-based execution with multiprocessing for better isolation - Add proper timeout handling that terminates processes exceeding time limits - Implement graceful process termination and force kill fallback mechanisms - Use multiprocessing Queue for safe inter-process communication - Update function documentation to reflect subprocess execution model - Maintain backward compatibility while improving security and reliability ```
1 parent f85592c commit bc489d1

1 file changed

Lines changed: 54 additions & 30 deletions

File tree

finance_mcp/core/utils/common_utils.py

Lines changed: 54 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -99,46 +99,70 @@ async def execute_task():
9999
await task
100100

101101

102+
import multiprocessing
103+
def _exec_code_in_process(queue: multiprocessing.Queue, code: str) -> None:
104+
"""Helper function to execute code in a subprocess.
105+
106+
This function runs in a separate process and communicates results
107+
back via a multiprocessing Queue.
108+
109+
Args:
110+
queue: Queue to put the execution result into.
111+
code: Python source code to execute.
112+
"""
113+
try:
114+
redirected_output = StringIO()
115+
with contextlib.redirect_stdout(redirected_output):
116+
exec(code) # noqa: S102
117+
queue.put(("success", redirected_output.getvalue()))
118+
except BaseException as e: # noqa: BLE001
119+
queue.put(("error", str(e)))
120+
121+
102122
async def exec_code(code: str, timeout: Optional[float] = 30) -> str:
103123
"""Execute arbitrary Python code and capture its printed output.
104124
105-
The code is executed in the current global context and any text written to
125+
The code is executed in a separate process and any text written to
106126
``stdout`` is captured and returned as a string. If an exception occurs,
107-
its string representation is returned instead.
127+
its string representation is returned instead. If execution exceeds the
128+
timeout, the process is terminated and an error message is returned.
108129
109130
Args:
110131
code: Python source code to execute.
111132
timeout: Maximum time in seconds to wait for code execution. ``None``
112133
disables the timeout and waits indefinitely. Defaults to 30 seconds.
113134
114135
Returns:
115-
Captured ``stdout`` output, or the exception message if execution fails.
116-
117-
Raises:
118-
asyncio.TimeoutError: If code execution exceeds ``timeout``.
136+
Captured ``stdout`` output, the exception message if execution fails,
137+
or a timeout error message if execution exceeds the timeout.
119138
"""
120139

121-
def _exec_code_sync():
122-
"""Synchronous code execution helper."""
123-
try:
124-
redirected_output = StringIO()
125-
with contextlib.redirect_stdout(redirected_output):
126-
exec(code)
127-
return redirected_output.getvalue()
128-
except Exception as e: # noqa: BLE001
129-
return str(e)
130-
except BaseException as e:
131-
return str(e)
132-
133-
try:
134-
# Execute code in a separate thread to avoid blocking the event loop
135-
if timeout:
136-
result = await asyncio.wait_for(
137-
asyncio.to_thread(_exec_code_sync),
138-
timeout=timeout
139-
)
140-
else:
141-
result = await asyncio.to_thread(_exec_code_sync)
142-
return result
143-
except asyncio.TimeoutError:
144-
return f"Code execution timed out after {timeout} seconds"
140+
def _run_in_process() -> str:
141+
"""Run code execution in a subprocess with timeout handling."""
142+
queue: multiprocessing.Queue = multiprocessing.Queue()
143+
process = multiprocessing.Process(
144+
target=_exec_code_in_process,
145+
args=(queue, code),
146+
)
147+
process.start()
148+
process.join(timeout=timeout)
149+
150+
if process.is_alive():
151+
# Process is still running, terminate it
152+
process.terminate()
153+
process.join(timeout=1)
154+
if process.is_alive():
155+
# Force kill if terminate didn't work
156+
process.kill()
157+
process.join(timeout=1)
158+
return f"Code execution timed out after {timeout} seconds"
159+
160+
# Process finished, get the result
161+
if not queue.empty():
162+
status, result = queue.get_nowait()
163+
return result
164+
165+
return "No output"
166+
167+
# Run the blocking process operation in a thread to avoid blocking the event loop
168+
return await asyncio.to_thread(_run_in_process)

0 commit comments

Comments
 (0)