diff --git a/src/policyengine_api/api/agent.py b/src/policyengine_api/api/agent.py index 03c060a..e4b23a6 100644 --- a/src/policyengine_api/api/agent.py +++ b/src/policyengine_api/api/agent.py @@ -172,39 +172,27 @@ async def _stream_modal_sandbox(question: str, api_base_url: str): ) logfire.info("sandbox_created") - # Read stdout synchronously in executor, yield lines as we get them - def read_next_line(stdout_iter): - try: - return next(stdout_iter) - except StopIteration: - return None - - stdout_iter = iter(process.stdout) + # Use Modal's async iteration for stdout lines_received = 0 events_sent = 0 - while True: - line = await loop.run_in_executor(executor, read_next_line, stdout_iter) - - if line is None: - # stdout exhausted - logfire.info("stdout_exhausted", total_lines=lines_received) - break - + # Modal StreamReader supports async iteration + async for line in process.stdout: lines_received += 1 logfire.info( "raw_line", line_num=lines_received, - line_len=len(line), - line_preview=line[:300].replace("session", "sess1on"), + line_len=len(line) if line else 0, + line_preview=line[:300].replace("session", "sess1on") + if line + else None, ) - parsed = _parse_claude_stream_event(line) if parsed: events_sent += 1 yield f"data: {json.dumps(parsed)}\n\n" - # Wait for process to finish + # Wait for process returncode = await loop.run_in_executor(executor, process.wait) logfire.info( "complete",