fix: run modal sandbox in thread pool to avoid blocking event loop (#25)

* fix: run modal sandbox in thread pool to avoid blocking event loop

The _stream_modal_sandbox function was running synchronous Modal SDK
calls inside an async generator, blocking the entire FastAPI app.

Also renamed DEMO_USE_MODAL → AGENT_USE_MODAL in env files.

* fix: move imports to top of file in analysis.py

Author: nikhilwoodruff

.env.example

@@ -27,5 +27,5 @@ MODAL_TOKEN_SECRET=as-...
 
 # Demo agent
 ANTHROPIC_API_KEY=sk-ant-...
-DEMO_USE_MODAL=false
+AGENT_USE_MODAL=false
 POLICYENGINE_API_URL=http://localhost:8000

src/policyengine_api/agent_sandbox.py

@@ -82,7 +82,9 @@ def run_claude_code_in_sandbox(
     return sb, process
 
 
-@app.function(image=sandbox_image, secrets=[anthropic_secret, logfire_secret], timeout=600)
+@app.function(
+    image=sandbox_image, secrets=[anthropic_secret, logfire_secret], timeout=600
+)
 def run_policy_analysis(
     question: str, api_base_url: str = "https://v2.api.policyengine.org"
 ) -> dict:
@@ -98,11 +100,15 @@ def run_policy_analysis(
 
     logfire.configure(service_name="policyengine-agent-sandbox")
 
-    with logfire.span("run_policy_analysis", question=question[:100], api_base_url=api_base_url):
+    with logfire.span(
+        "run_policy_analysis", question=question[:100], api_base_url=api_base_url
+    ):
         # Write MCP config
         os.makedirs("/root/.claude", exist_ok=True)
         mcp_config = {
-            "mcpServers": {"policyengine": {"type": "sse", "url": f"{api_base_url}/mcp"}}
+            "mcpServers": {
+                "policyengine": {"type": "sse", "url": f"{api_base_url}/mcp"}
+            }
         }
         with open("/root/.claude/mcp_servers.json", "w") as f:
             json.dump(mcp_config, f)

src/policyengine_api/api/agent.py

@@ -86,38 +86,91 @@ async def _stream_claude_code(question: str, api_base_url: str):
 
 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
+
     import logfire
 
     sb = None
+    executor = ThreadPoolExecutor(max_workers=1)
     try:
         from policyengine_api.agent_sandbox import run_claude_code_in_sandbox
 
-        logfire.info("Creating Modal sandbox", question=question[:100], api_base_url=api_base_url)
-        sb, process = run_claude_code_in_sandbox(question, api_base_url)
+        logfire.info(
+            "Creating Modal sandbox", question=question[:100], api_base_url=api_base_url
+        )
+
+        # Run blocking Modal SDK calls in thread pool to avoid blocking event loop
+        loop = asyncio.get_event_loop()
+        sb, process = await loop.run_in_executor(
+            executor, run_claude_code_in_sandbox, question, api_base_url
+        )
         logfire.info("Modal sandbox created, streaming output")
 
-        # Stream stdout line by line
-        for line in process.stdout:
-            yield f"data: {json.dumps({'type': 'output', 'content': line})}\n\n"
+        # Poll for lines with timeout to allow other async tasks
+        import queue
+        import threading
 
-        process.wait()
+        line_queue = queue.Queue()
 
-        if process.returncode != 0:
-            stderr = process.stderr.read()
-            logfire.error("Claude Code failed in sandbox", returncode=process.returncode, stderr=stderr[:500])
-            yield f"data: {json.dumps({'type': 'error', 'content': stderr})}\n\n"
+        def stream_reader():
+            try:
+                for line in process.stdout:
+                    line_queue.put(("line", line))
+                process.wait()
+                if process.returncode != 0:
+                    stderr = process.stderr.read()
+                    line_queue.put(("error", (process.returncode, stderr)))
+                else:
+                    line_queue.put(("done", process.returncode))
+            except Exception as e:
+                line_queue.put(("exception", str(e)))
+
+        reader_thread = threading.Thread(target=stream_reader, daemon=True)
+        reader_thread.start()
+
+        while True:
+            try:
+                # Non-blocking check with short timeout
+                item = await loop.run_in_executor(
+                    executor, lambda: line_queue.get(timeout=0.1)
+                )
+                event_type, data = item
+
+                if event_type == "line":
+                    yield f"data: {json.dumps({'type': 'output', 'content': data})}\n\n"
+                elif event_type == "error":
+                    returncode, stderr = data
+                    logfire.error(
+                        "Claude Code failed in sandbox",
+                        returncode=returncode,
+                        stderr=stderr[:500],
+                    )
+                    yield f"data: {json.dumps({'type': 'error', 'content': stderr})}\n\n"
+                    yield f"data: {json.dumps({'type': 'done', 'returncode': returncode})}\n\n"
+                    break
+                elif event_type == "done":
+                    yield f"data: {json.dumps({'type': 'done', 'returncode': data})}\n\n"
+                    break
+                elif event_type == "exception":
+                    raise Exception(data)
+            except Exception as e:
+                if "Empty" in type(e).__name__:
+                    # Queue timeout, continue polling
+                    await asyncio.sleep(0)
+                    continue
+                raise
 
-        yield f"data: {json.dumps({'type': 'done', 'returncode': process.returncode})}\n\n"
     except Exception as e:
         logfire.exception("Modal sandbox failed", error=str(e))
         yield f"data: {json.dumps({'type': 'error', 'content': f'Sandbox error: {str(e)}'})}\n\n"
         yield f"data: {json.dumps({'type': 'done', 'returncode': 1})}\n\n"
     finally:
         if sb is not None:
             try:
-                sb.terminate()
+                await loop.run_in_executor(executor, sb.terminate)
             except Exception:
                 pass
+        executor.shutdown(wait=False)
 
 
 @router.post("/stream")

src/policyengine_api/api/analysis.py

@@ -24,16 +24,6 @@
 from pydantic import BaseModel, Field
 from sqlmodel import Session, select
 
-
-def _safe_float(value: float | None) -> float | None:
-    """Convert NaN/inf to None for JSON serialization."""
-    if value is None:
-        return None
-    if math.isnan(value) or math.isinf(value):
-        return None
-    return value
-
-
 from policyengine_api.models import (
     Dataset,
     DecileImpact,
@@ -49,6 +39,15 @@ def _safe_float(value: float | None) -> float | None:
 )
 from policyengine_api.services.database import get_session
 
+
+def _safe_float(value: float | None) -> float | None:
+    """Convert NaN/inf to None for JSON serialization."""
+    if value is None:
+        return None
+    if math.isnan(value) or math.isinf(value):
+        return None
+    return value
+
 # Namespace for deterministic UUIDs
 SIMULATION_NAMESPACE = UUID("a1b2c3d4-e5f6-7890-abcd-ef1234567890")
 REPORT_NAMESPACE = UUID("b2c3d4e5-f6a7-8901-bcde-f12345678901")