fix: use Modal async API for proper streaming

- Use modal.Sandbox.create.aio() and sb.exec.aio()
- Use async for iteration over stdout
- Use process.wait.aio() and sb.terminate.aio()

This ensures proper async streaming instead of mixing sync/async.

Author: nikhilwoodruff

src/policyengine_api/agent_sandbox.py

@@ -41,17 +41,17 @@
 logfire_secret = modal.Secret.from_name("logfire-token")
 
 
-def run_claude_code_in_sandbox(
+async def run_claude_code_in_sandbox_async(
     question: str,
     api_base_url: str = "https://v2.api.policyengine.org",
 ) -> tuple[modal.Sandbox, any]:
     """Create a sandbox running Claude Code with MCP server configured.
 
     Returns the sandbox and process handle for streaming output.
+    Uses Modal's async API for proper streaming support.
     """
     import logfire
 
-    print("[SANDBOX] run_claude_code_in_sandbox starting", flush=True)
     logfire.info(
         "run_claude_code_in_sandbox: starting",
         question=question[:100],
@@ -65,35 +65,20 @@ def run_claude_code_in_sandbox(
     mcp_config_json = json.dumps(mcp_config)
 
     # Get reference to deployed app (required when calling from outside Modal)
-    print("[SANDBOX] looking up Modal app", flush=True)
     logfire.info("run_claude_code_in_sandbox: looking up Modal app")
     sandbox_app = modal.App.lookup("policyengine-sandbox", create_if_missing=True)
-    print("[SANDBOX] Modal app found", flush=True)
     logfire.info("run_claude_code_in_sandbox: Modal app found")
 
-    print("[SANDBOX] creating sandbox", flush=True)
     logfire.info("run_claude_code_in_sandbox: creating sandbox")
-    sb = modal.Sandbox.create(
+    sb = await modal.Sandbox.create.aio(
         app=sandbox_app,
         image=sandbox_image,
         secrets=[anthropic_secret, logfire_secret],
         timeout=600,
         workdir="/tmp",
     )
-    print("[SANDBOX] sandbox created", flush=True)
     logfire.info("run_claude_code_in_sandbox: sandbox created")
 
-    # Run Claude Code with the question
-    # Note: Can't use --dangerously-skip-permissions as root (Modal runs as root)
-    # Use shell wrapper with </dev/null to properly close stdin (prevents hanging)
-    # --max-turns: limit execution to prevent runaway
-    # Use --mcp-config to pass MCP config directly (more reliable than config file)
-    print("[SANDBOX] Starting claude CLI with question", flush=True)
-    logfire.info(
-        "run_claude_code_in_sandbox: starting claude CLI",
-        mcp_url=f"{api_base_url}/mcp",
-    )
-
     # Escape the question and config for shell
     escaped_question = question.replace("'", "'\"'\"'")
     escaped_mcp_config = mcp_config_json.replace("'", "'\"'\"'")
@@ -112,9 +97,8 @@ def run_claude_code_in_sandbox(
         question_len=len(question),
         escaped_question_len=len(escaped_question),
     )
-    # text=True, bufsize=1 enables line-buffered streaming
-    process = sb.exec("sh", "-c", cmd, text=True, bufsize=1)
-    print("[SANDBOX] claude CLI process started", flush=True)
+    # Use async exec for proper streaming
+    process = await sb.exec.aio("sh", "-c", cmd, text=True, bufsize=1)
     logfire.info("run_claude_code_in_sandbox: claude CLI process started, returning")
 
     return sb, process

src/policyengine_api/api/agent.py

@@ -154,58 +154,38 @@ def _parse_claude_stream_event(line: str) -> dict | None:
 
 async def _stream_modal_sandbox(question: str, api_base_url: str):
     """Stream output from Claude Code running in Modal Sandbox."""
-    from concurrent.futures import ThreadPoolExecutor
-
     with logfire.span(
         "agent_stream", question=question[:100], api_base_url=api_base_url
     ):
         sb = None
-        executor = ThreadPoolExecutor(max_workers=2)
         try:
-            from policyengine_api.agent_sandbox import run_claude_code_in_sandbox
+            from policyengine_api.agent_sandbox import run_claude_code_in_sandbox_async
 
             logfire.info("creating_sandbox")
-
-            loop = asyncio.get_event_loop()
-            sb, process = await loop.run_in_executor(
-                executor, run_claude_code_in_sandbox, question, api_base_url
-            )
+            sb, process = await run_claude_code_in_sandbox_async(question, api_base_url)
             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
-
+            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
-            returncode = await loop.run_in_executor(executor, process.wait)
+            # Wait for process using async API
+            returncode = await process.wait.aio()
             logfire.info(
                 "complete",
                 returncode=returncode,
@@ -221,10 +201,9 @@ def read_next_line(stdout_iter):
         finally:
             if sb is not None:
                 try:
-                    await loop.run_in_executor(executor, sb.terminate)
+                    await sb.terminate.aio()
                 except Exception:
                     pass
-            executor.shutdown(wait=False)
 
 
 @router.post("/stream")