1+ import asyncio
12import base64
23from typing import AsyncIterator , Dict , Optional , Union
34
1617 SandboxProcessStdinParams as SandboxProcessStdinParamsDict ,
1718)
1819from ...sandboxes .shared import _normalize_exec_params
20+ from ...sandboxes .process_output import (
21+ DEFAULT_MAX_PROCESS_OUTPUT_BYTES ,
22+ ProcessOutput ,
23+ validate_output_limit ,
24+ )
1925from .sandbox_transport import RuntimeTransport
2026
2127DEFAULT_PROCESS_KILL_WAIT_SECONDS = 5.0
@@ -25,6 +31,10 @@ class SandboxProcessHandle:
2531 def __init__ (self , transport : RuntimeTransport , summary : SandboxProcessSummary ):
2632 self ._transport = transport
2733 self ._summary = summary
34+ self ._output = None
35+ self ._collector = None
36+ self ._events = None
37+ self ._changed = asyncio .Event ()
2838
2939 @property
3040 def id (self ) -> str :
@@ -50,6 +60,15 @@ async def wait(
5060 timeout_ms : Optional [int ] = None ,
5161 timeout_sec : Optional [int ] = None ,
5262 ) -> SandboxProcessResult :
63+ if self ._collector is not None :
64+ timeout = None
65+ if timeout_sec is not None and timeout_sec > 0 :
66+ timeout = timeout_sec
67+ elif timeout_ms is not None and timeout_ms > 0 :
68+ timeout = timeout_ms / 1000
69+ if not self ._collector .done ():
70+ await asyncio .wait_for (asyncio .shield (self ._collector ), timeout )
71+ return self ._collected_result ()
5372 payload = await self ._transport .request_json (
5473 f"/sandbox/processes/{ self .id } /wait" ,
5574 method = "POST" ,
@@ -60,6 +79,10 @@ async def wait(
6079 headers = {"content-type" : "application/json" },
6180 )
6281 result = SandboxProcessResult (** payload ["result" ])
82+ if result .output_truncated :
83+ raise ProcessOutput (self .id , 0 ).failure (
84+ "Retained process output is incomplete; collect output from process start"
85+ )
6386 self ._summary = SandboxProcessSummary (
6487 id = result .id ,
6588 status = result .status ,
@@ -133,6 +156,21 @@ async def write_stdin(
133156 )
134157
135158 async def stream (self , from_seq : Optional [int ] = None ) -> AsyncIterator [object ]:
159+ if self ._output is not None :
160+ index = 0
161+ while True :
162+ self ._changed .clear ()
163+ while index < len (self ._output .events ):
164+ event = self ._output .events [index ]
165+ index += 1
166+ if from_seq is None or event .seq >= from_seq :
167+ yield event
168+ if self ._collector .done ():
169+ yield SandboxProcessExitEvent (
170+ type = "exit" , result = self ._collected_result ()
171+ )
172+ return
173+ await self ._changed .wait ()
136174 params = {"from_seq" : from_seq } if from_seq and from_seq > 0 else None
137175 async for event in self ._transport .stream_sse (
138176 f"/sandbox/processes/{ self .id } /stream" ,
@@ -153,6 +191,70 @@ async def stream(self, from_seq: Optional[int] = None) -> AsyncIterator[object]:
153191 result = SandboxProcessResult (** data ),
154192 )
155193
194+ def _collected_result (self ) -> SandboxProcessResult :
195+ if self ._output .error is not None :
196+ raise self ._output .error
197+ if self ._output .result is None :
198+ raise self ._output .failure (
199+ "Command stream ended before its completion event"
200+ )
201+ result = self ._output .result
202+ self ._summary = self ._summary .model_copy (
203+ update = {
204+ "status" : result .status ,
205+ "exit_code" : result .exit_code ,
206+ "completed_at" : result .completed_at ,
207+ }
208+ )
209+ return result
210+
211+ async def _collect (self ) -> None :
212+ try :
213+ async for event in self ._events :
214+ self ._output .consume (event )
215+ self ._changed .set ()
216+ if self ._output .result is not None :
217+ return
218+ self ._output .error = self ._output .failure (
219+ "Command stream ended before its completion event"
220+ )
221+ except asyncio .CancelledError :
222+ if self ._output .result is None :
223+ self ._output .error = self ._output .failure (
224+ "Command output collection disconnected"
225+ )
226+ except Exception as error :
227+ self ._output .error = (
228+ self ._output .failure (str (error ))
229+ if not hasattr (error , "code" )
230+ else error
231+ )
232+ finally :
233+ try :
234+ await self ._events .aclose ()
235+ except Exception as error :
236+ if self ._output .result is None and self ._output .error is None :
237+ self ._output .error = self ._output .failure (str (error ))
238+ finally :
239+ self ._changed .set ()
240+
241+ async def disconnect (self ) -> None :
242+ """Stop collecting output; the detached command continues running."""
243+ if self ._collector is not None and not self ._collector .done ():
244+ if self ._output .result is None and self ._output .error is None :
245+ self ._output .error = self ._output .failure (
246+ "Command output collection disconnected"
247+ )
248+ self ._collector .cancel ()
249+ try :
250+ await self ._collector
251+ except asyncio .CancelledError :
252+ pass
253+ finally :
254+ # Cancellation may happen before the collector gets its first turn.
255+ await self ._events .aclose ()
256+ self ._changed .set ()
257+
156258 async def result (self ) -> SandboxProcessResult :
157259 return await self .wait ()
158260
@@ -170,22 +272,21 @@ async def exec(
170272 timeout_ms : Optional [int ] = None ,
171273 timeout_sec : Optional [int ] = None ,
172274 run_as : Optional [str ] = None ,
275+ max_output_bytes : int = DEFAULT_MAX_PROCESS_OUTPUT_BYTES ,
173276 ) -> SandboxProcessResult :
174- params = _normalize_exec_params (
277+ handle = await self . start (
175278 input ,
176279 cwd = cwd ,
177280 env = env ,
178281 timeout_ms = timeout_ms ,
179282 timeout_sec = timeout_sec ,
180283 run_as = run_as ,
284+ max_output_bytes = max_output_bytes ,
181285 )
182- payload = await self ._transport .request_json (
183- "/sandbox/exec" ,
184- method = "POST" ,
185- json_body = dump_request (params , SandboxExecParams ),
186- headers = {"content-type" : "application/json" },
187- )
188- return SandboxProcessResult (** payload ["result" ])
286+ try :
287+ return await handle .wait ()
288+ finally :
289+ await handle .disconnect ()
189290
190291 async def start (
191292 self ,
@@ -196,7 +297,9 @@ async def start(
196297 timeout_ms : Optional [int ] = None ,
197298 timeout_sec : Optional [int ] = None ,
198299 run_as : Optional [str ] = None ,
300+ max_output_bytes : int = DEFAULT_MAX_PROCESS_OUTPUT_BYTES ,
199301 ) -> SandboxProcessHandle :
302+ validate_output_limit (max_output_bytes )
200303 params = _normalize_exec_params (
201304 input ,
202305 cwd = cwd ,
@@ -205,16 +308,25 @@ async def start(
205308 timeout_sec = timeout_sec ,
206309 run_as = run_as ,
207310 )
208- payload = await self ._transport .request_json (
311+ events = self ._transport .stream_sse (
209312 "/sandbox/processes" ,
210313 method = "POST" ,
211314 json_body = dump_request (params , SandboxExecParams ),
212- headers = {"content-type" : "application/json" },
213- )
214- return SandboxProcessHandle (
215- self ._transport ,
216- SandboxProcessSummary (** payload ["process" ]),
217315 )
316+ try :
317+ started = await events .__anext__ ()
318+ if started ["event" ] != "started" :
319+ raise RuntimeError ("Expected process start event" )
320+ handle = SandboxProcessHandle (
321+ self ._transport , SandboxProcessSummary (** started ["data" ])
322+ )
323+ except BaseException :
324+ await events .aclose ()
325+ raise
326+ handle ._events = events
327+ handle ._output = ProcessOutput (handle .id , max_output_bytes )
328+ handle ._collector = asyncio .create_task (handle ._collect ())
329+ return handle
218330
219331 async def get (self , process_id : str ) -> SandboxProcessHandle :
220332 payload = await self ._transport .request_json (f"/sandbox/processes/{ process_id } " )
0 commit comments