Skip to content

Commit bdcec05

Browse files
committed
Replace multiprocessing.Queue with Pipe in Code Mode sandbox
- Use multiprocessing.Pipe for parent-child communication instead of Queue, which is lighter-weight for point-to-point message passing - Remove unused queue import
1 parent 1114ba9 commit bdcec05

1 file changed

Lines changed: 10 additions & 12 deletions

File tree

teaagent/code_mode.py

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import ast
44
import multiprocessing
5-
import queue
65
import traceback
76
from dataclasses import dataclass
87
from typing import Any
@@ -96,11 +95,10 @@ def execute_code_mode(
9695
if active_sandbox.memory_bytes <= 0:
9796
raise UnsafeCodeError('Code Mode memory limit must be positive')
9897

99-
context = multiprocessing.get_context('spawn')
100-
result_queue: multiprocessing.Queue[dict[str, Any]] = context.Queue(maxsize=1)
101-
process = context.Process(
98+
parent_conn, child_conn = multiprocessing.Pipe(duplex=False)
99+
process = multiprocessing.Process(
102100
target=_execute_code_mode_child,
103-
args=(code, safe_inputs, active_sandbox, result_queue),
101+
args=(code, safe_inputs, active_sandbox, child_conn),
104102
)
105103
process.start()
106104
process.join(active_sandbox.timeout_seconds)
@@ -109,10 +107,10 @@ def execute_code_mode(
109107
process.join()
110108
raise UnsafeCodeError('Code Mode timed out')
111109

112-
try:
113-
message = result_queue.get_nowait()
114-
except queue.Empty as exc:
115-
raise UnsafeCodeError('Code Mode sandbox exited without a result') from exc
110+
if parent_conn.poll():
111+
message = parent_conn.recv()
112+
else:
113+
raise UnsafeCodeError('Code Mode sandbox exited without a result')
116114

117115
if message['status'] == 'error':
118116
raise UnsafeCodeError(message['error'])
@@ -135,7 +133,7 @@ def _execute_code_mode_child(
135133
code: str,
136134
inputs: dict[str, Any],
137135
sandbox: CodeModeSandbox,
138-
result_queue: multiprocessing.Queue[dict[str, Any]],
136+
result_pipe: multiprocessing.connection.Connection,
139137
) -> None:
140138
try:
141139
_apply_resource_limits(sandbox)
@@ -149,15 +147,15 @@ def _execute_code_mode_child(
149147
}
150148
_validate_plain_data(variables, 'variables')
151149
except Exception as exc: # pragma: no cover - exercised through parent process.
152-
result_queue.put(
150+
result_pipe.send(
153151
{
154152
'status': 'error',
155153
'error': f'{type(exc).__name__}: {exc}',
156154
'traceback': traceback.format_exc(),
157155
}
158156
)
159157
else:
160-
result_queue.put({'status': 'ok', 'variables': variables})
158+
result_pipe.send({'status': 'ok', 'variables': variables})
161159

162160

163161
def _apply_resource_limits(sandbox: CodeModeSandbox) -> None:

0 commit comments

Comments
 (0)