@useatlas/create 0.0.5 → 0.0.7

package/templates/nextjs-standalone/src/lib/tools/python-sandbox.ts

@@ -17,18 +17,16 @@
  */
 
 import type { PythonBackend, PythonResult } from "./python";
+import { PYTHON_SECURITY_AND_SETUP, PYTHON_EXEC_AND_COLLECT } from "./python-wrapper";
+import { sandboxErrorDetail, safeError, MAX_OUTPUT } from "./backends/shared";
 import { randomUUID } from "crypto";
 import { createLogger } from "@atlas/api/lib/logger";
-import { SENSITIVE_PATTERNS } from "@atlas/api/lib/security";
 
 const log = createLogger("python-sandbox");
 
 /** Default Python execution timeout in ms. */
 const DEFAULT_TIMEOUT_MS = 30_000;
 
-/** Maximum bytes to read from stdout/stderr (1 MB). */
-const MAX_OUTPUT = 1024 * 1024;
-
 /** Packages to install in the sandbox. */
 const DATA_SCIENCE_PACKAGES = [
   "pandas",
@@ -40,13 +38,9 @@ const DATA_SCIENCE_PACKAGES = [
 ];
 
 /**
- * Python wrapper script — adapted from the nsjail PYTHON_WRAPPER.
- *
- * Key difference: data is injected via a JSON file (argv[2]) instead of
- * stdin, since Vercel Sandbox's runCommand does not support stdin piping.
- *
- * Reads user code from argv[1], executes in a restricted Python namespace,
- * collects charts + structured output via result marker.
+ * Non-streaming Python wrapper for Vercel Sandbox. Composes shared fragments
+ * (PYTHON_SECURITY_AND_SETUP, PYTHON_EXEC_AND_COLLECT) with file-based
+ * data injection (argv[2]) since runCommand does not support stdin piping.
  */
 const PYTHON_WRAPPER = `
 import sys, json, io, base64, glob, os, ast
@@ -54,53 +48,11 @@ import sys, json, io, base64, glob, os, ast
 _marker = os.environ["ATLAS_RESULT_MARKER"]
 _chart_dir = os.environ.get("ATLAS_CHART_DIR", "/tmp/charts")
 
-# --- Import guard ---
-_BLOCKED_MODULES = {
-    "subprocess", "os", "socket", "shutil", "sys", "ctypes", "importlib",
-    "code", "signal", "multiprocessing", "threading", "pty", "fcntl",
-    "termios", "resource", "posixpath",
-    "http", "urllib", "requests", "httpx", "aiohttp", "webbrowser",
-    "pickle", "tempfile", "pathlib",
-}
-_BLOCKED_BUILTINS = {
-    "compile", "exec", "eval", "__import__", "open", "breakpoint",
-    "getattr", "globals", "locals", "vars", "dir", "delattr", "setattr",
-}
-
-_user_code = open(sys.argv[1]).read()
-try:
-    _tree = ast.parse(_user_code)
-except SyntaxError as e:
-    print(_marker + json.dumps({"success": False, "error": f"SyntaxError: {e.msg} (line {e.lineno})"}))
+def _report_error(msg):
+    print(_marker + json.dumps({"success": False, "error": msg}))
     sys.exit(0)
 
-_blocked = None
-for _node in ast.walk(_tree):
-    if _blocked:
-        break
-    if isinstance(_node, ast.Import):
-        for _alias in _node.names:
-            _mod = _alias.name.split('.')[0]
-            if _mod in _BLOCKED_MODULES:
-                _blocked = f'Blocked import: "{_mod}" is not allowed'
-                break
-    elif isinstance(_node, ast.ImportFrom):
-        if _node.module:
-            _mod = _node.module.split('.')[0]
-            if _mod in _BLOCKED_MODULES:
-                _blocked = f'Blocked import: "{_mod}" is not allowed'
-    elif isinstance(_node, ast.Call):
-        _name = None
-        if isinstance(_node.func, ast.Name):
-            _name = _node.func.id
-        elif isinstance(_node.func, ast.Attribute):
-            _name = _node.func.attr
-        if _name and _name in _BLOCKED_BUILTINS:
-            _blocked = f'Blocked builtin: "{_name}()" is not allowed'
-
-if _blocked:
-    print(_marker + json.dumps({"success": False, "error": _blocked}))
-    sys.exit(0)
+${PYTHON_SECURITY_AND_SETUP}
 
 # --- Data injection (from file, not stdin) ---
 _atlas_data = None
@@ -122,84 +74,9 @@ if _atlas_data:
     except ImportError:
         data = _atlas_data
 
-# Configure matplotlib for headless rendering
-try:
-    import matplotlib
-    matplotlib.use('Agg')
-except ImportError:
-    pass
-
-os.makedirs(_chart_dir, exist_ok=True)
-
-def chart_path(n=0):
-    return os.path.join(_chart_dir, f"chart_{n}.png")
-
-# --- Execute user code ---
-_old_stdout = sys.stdout
-sys.stdout = _captured = io.StringIO()
-
-_user_ns = {"chart_path": chart_path, "data": data, "df": df}
-_atlas_error = None
-try:
-    exec(_user_code, _user_ns)
-except Exception as e:
-    _atlas_error = f"{type(e).__name__}: {e}"
-
-_output = _captured.getvalue()
-sys.stdout = _old_stdout
-
-# --- Collect results ---
-_charts = []
-for f in sorted(glob.glob(os.path.join(_chart_dir, "chart_*.png"))):
-    with open(f, "rb") as fh:
-        _charts.append({"base64": base64.b64encode(fh.read()).decode(), "mimeType": "image/png"})
-
-_result = {"success": _atlas_error is None}
-if _output.strip():
-    _result["output"] = _output.strip()
-if _atlas_error:
-    _result["error"] = _atlas_error
-
-if "_atlas_table" in _user_ns:
-    _result["table"] = _user_ns["_atlas_table"]
-
-if "_atlas_chart" in _user_ns:
-    _ac = _user_ns["_atlas_chart"]
-    if isinstance(_ac, dict):
-        _result["rechartsCharts"] = [_ac]
-    elif isinstance(_ac, list):
-        _result["rechartsCharts"] = _ac
-
-if _charts:
-    _result["charts"] = _charts
-
-print(_marker + json.dumps(_result), file=_old_stdout)
+${PYTHON_EXEC_AND_COLLECT}
 `;
 
-/** Format an error for logging, with extra detail from @vercel/sandbox APIError. */
-function sandboxErrorDetail(err: unknown): string {
-  if (!(err instanceof Error)) return String(err);
-  const detail = err.message;
-  const json = (err as unknown as Record<string, unknown>).json;
-  const text = (err as unknown as Record<string, unknown>).text;
-  if (json) {
-    try {
-      return `${detail} — response: ${JSON.stringify(json)}`;
-    } catch {
-      return `${detail} — response: [unserializable object]`;
-    }
-  }
-  if (typeof text === "string" && text) return `${detail} — body: ${text.slice(0, 500)}`;
-  return detail;
-}
-
-/** Scrub sensitive data from error messages before exposing. */
-function safeError(detail: string): string {
-  return SENSITIVE_PATTERNS.test(detail)
-    ? "sandbox API error (details in server logs)"
-    : detail;
-}
-
 // Sandbox base dir for relative paths
 const SANDBOX_BASE = "/vercel/sandbox";
 

package/templates/nextjs-standalone/src/lib/tools/python-sidecar.ts

@@ -8,8 +8,8 @@
  * Configured via ATLAS_SANDBOX_URL (same as the explore sidecar).
  */
 
-import type { SidecarPythonRequest } from "@atlas/api/lib/sidecar-types";
-import type { PythonResult } from "./python";
+import type { SidecarPythonRequest, SidecarPythonStreamEvent } from "@atlas/api/lib/sidecar-types";
+import type { PythonChart, PythonProgressEvent, PythonResult, RechartsChart } from "./python";
 import { createLogger } from "@atlas/api/lib/logger";
 
 const log = createLogger("python-sidecar");
@@ -36,6 +36,7 @@ export async function executePythonViaSidecar(
   try {
     baseUrl = new URL(sidecarUrl);
   } catch {
+    log.warn({ sidecarUrl }, "Invalid ATLAS_SANDBOX_URL — cannot parse as URL");
     return pythonError(
       `Invalid ATLAS_SANDBOX_URL: "${sidecarUrl}". Expected a valid URL (e.g. http://sandbox-sidecar:8080).`,
     );
@@ -92,6 +93,7 @@ export async function executePythonViaSidecar(
     try {
       errorBody = await response.text();
     } catch {
+      // intentionally ignored: body may already be consumed or unreadable
       errorBody = `HTTP ${response.status}`;
     }
 
@@ -112,7 +114,7 @@ export async function executePythonViaSidecar(
           return parsed;
         }
       } catch {
-        // Not structured — fall through
+        // intentionally ignored: best-effort structured parse; raw error already logged above
       }
     }
 
@@ -148,3 +150,182 @@ export async function executePythonViaSidecar(
 
   return result;
 }
+
+/**
+ * Execute Python via the sidecar's streaming endpoint.
+ *
+ * Reads NDJSON events from POST /exec-python-stream, calling `onProgress`
+ * for each stdout/chart event, and accumulates the final PythonResult.
+ * Falls back to the non-streaming endpoint on connection errors.
+ */
+export async function executePythonViaSidecarStream(
+  sidecarUrl: string,
+  code: string,
+  data: { columns: string[]; rows: unknown[][] } | undefined,
+  onProgress: (event: PythonProgressEvent) => void,
+): Promise<PythonResult> {
+  const authToken = process.env.SIDECAR_AUTH_TOKEN;
+
+  let baseUrl: URL;
+  try {
+    baseUrl = new URL(sidecarUrl);
+  } catch {
+    log.warn({ sidecarUrl }, "Invalid ATLAS_SANDBOX_URL — cannot parse as URL (stream)");
+    return pythonError(
+      `Invalid ATLAS_SANDBOX_URL: "${sidecarUrl}". Expected a valid URL (e.g. http://sandbox-sidecar:8080).`,
+    );
+  }
+
+  const execUrl = new URL("/exec-python-stream", baseUrl).toString();
+  const rawTimeout = parseInt(process.env.ATLAS_PYTHON_TIMEOUT ?? String(DEFAULT_TIMEOUT_MS), 10);
+  const timeout = Number.isNaN(rawTimeout) ? DEFAULT_TIMEOUT_MS : rawTimeout;
+
+  const requestBody: SidecarPythonRequest = { code, timeout };
+  if (data) {
+    requestBody.data = data;
+  }
+
+  let response: Response;
+  try {
+    response = await fetch(execUrl, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        ...(authToken ? { Authorization: `Bearer ${authToken}` } : {}),
+      },
+      body: JSON.stringify(requestBody),
+      signal: AbortSignal.timeout(timeout + HTTP_OVERHEAD_MS),
+    });
+  } catch (err) {
+    const detail = err instanceof Error ? err.message : String(err);
+    const isConnectionError =
+      detail.includes("ECONNREFUSED") ||
+      detail.includes("fetch failed") ||
+      detail.includes("Failed to connect");
+
+    if (isConnectionError) {
+      log.warn({ err: detail }, "Streaming endpoint unreachable, falling back to non-streaming");
+      return executePythonViaSidecar(sidecarUrl, code, data);
+    }
+
+    log.error({ err: detail }, "Unexpected error calling streaming Python endpoint");
+    return pythonError(`Streaming Python execution failed: ${detail}`);
+  }
+
+  if (!response.ok || !response.body) {
+    if (response.status === 404) {
+      // Sidecar doesn't support streaming — expected for older versions
+      log.info("Streaming endpoint not found (404), falling back to non-streaming");
+      return executePythonViaSidecar(sidecarUrl, code, data);
+    }
+    log.error({ status: response.status, hasBody: !!response.body }, "Streaming Python endpoint returned unexpected HTTP error, falling back");
+    return executePythonViaSidecar(sidecarUrl, code, data);
+  }
+
+  // Consume the NDJSON stream
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+
+  // Accumulate output, charts, tables for the final PythonResult
+  const outputParts: string[] = [];
+  const charts: PythonChart[] = [];
+  const rechartsCharts: RechartsChart[] = [];
+  let table: { columns: string[]; rows: unknown[][] } | undefined;
+  let error: string | undefined;
+  let success = true;
+  let receivedTerminal = false;
+  let parseFailures = 0;
+
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+
+      buffer += decoder.decode(value, { stream: true });
+
+      let newlineIdx;
+      while ((newlineIdx = buffer.indexOf("\n")) !== -1) {
+        const line = buffer.slice(0, newlineIdx).trim();
+        buffer = buffer.slice(newlineIdx + 1);
+
+        if (!line) continue;
+
+        let event: SidecarPythonStreamEvent;
+        try {
+          event = JSON.parse(line) as SidecarPythonStreamEvent;
+        } catch (parseErr) {
+          parseFailures++;
+          log.warn(
+            { line: line.slice(0, 200), err: parseErr instanceof Error ? parseErr.message : String(parseErr), parseFailures },
+            "Skipping unparseable NDJSON line from Python stream",
+          );
+          continue;
+        }
+
+        switch (event.type) {
+          case "stdout":
+            outputParts.push(event.data);
+            onProgress({ type: "stdout", content: event.data });
+            break;
+          case "chart":
+            charts.push(event.data);
+            onProgress({ type: "chart", chart: event.data });
+            break;
+          case "recharts":
+            rechartsCharts.push(event.data);
+            onProgress({ type: "recharts", chart: event.data });
+            break;
+          case "table":
+            table = event.data;
+            break;
+          case "done":
+            receivedTerminal = true;
+            break;
+          case "error":
+            receivedTerminal = true;
+            success = false;
+            error = event.data.error;
+            if (event.data.output) outputParts.push(event.data.output);
+            break;
+          default: {
+            // Exhaustiveness guard — new event types added to SidecarPythonStreamEvent
+            // will cause a compile error here until handled.
+            const _exhaustive: never = event;
+            log.warn({ type: (_exhaustive as { type: string }).type }, "Unknown Python stream event type");
+          }
+        }
+      }
+    }
+  } catch (err) {
+    const detail = err instanceof Error ? err.message : String(err);
+    log.error({ err: detail }, "Error reading Python stream");
+    return pythonError(`Stream read failed: ${detail}`);
+  } finally {
+    reader.releaseLock();
+  }
+
+  if (parseFailures > 0) {
+    log.warn({ parseFailures }, "Python stream had unparseable NDJSON lines — result may be incomplete");
+  }
+
+  const output = outputParts.join("").trim() || undefined;
+
+  // If the stream ended without a terminal event, the execution was interrupted
+  if (!receivedTerminal) {
+    log.warn("Python stream ended without done/error event — execution may have been interrupted");
+    return { success: false, error: "Python execution was interrupted (no completion signal received)", ...(output && { output }) };
+  }
+
+  if (!success) {
+    return { success: false, error: error ?? "Unknown error", ...(output && { output }) };
+  }
+
+  return {
+    success: true,
+    ...(output && { output }),
+    ...(table && { table }),
+    ...(charts.length > 0 && { charts }),
+    ...(rechartsCharts.length > 0 && { rechartsCharts }),
+  };
+}

package/templates/nextjs-standalone/src/lib/tools/python-stream.ts

@@ -0,0 +1,33 @@
+/**
+ * Request-scoped store for UI message stream writers.
+ *
+ * During streaming Python execution, the Python tool needs to write progress
+ * events (stdout chunks, chart renders) to the same UI message stream that
+ * carries the agent's text output. This module provides a simple Map-based
+ * store keyed by request ID, set up by the chat route and read by the tool.
+ *
+ * Uses getRequestContext() (AsyncLocalStorage) to locate the writer for
+ * the current request without threading it through function arguments.
+ */
+
+import type { UIMessageStreamWriter } from "ai";
+import { getRequestContext } from "@atlas/api/lib/logger";
+
+const writers = new Map<string, UIMessageStreamWriter>();
+
+/** Register a stream writer for the current request. Called by the chat route. */
+export function setStreamWriter(requestId: string, writer: UIMessageStreamWriter): void {
+  writers.set(requestId, writer);
+}
+
+/** Remove the stream writer when the stream completes. */
+export function clearStreamWriter(requestId: string): void {
+  writers.delete(requestId);
+}
+
+/** Get the stream writer for the current request. Returns undefined if not in a streaming context. */
+export function getStreamWriter(): UIMessageStreamWriter | undefined {
+  const ctx = getRequestContext();
+  if (!ctx?.requestId) return undefined;
+  return writers.get(ctx.requestId);
+}

package/templates/nextjs-standalone/src/lib/tools/python-wrapper.ts

@@ -0,0 +1,129 @@
+/**
+ * Shared Python wrapper fragments for data analysis sandboxes.
+ *
+ * Both python-nsjail.ts and python-sandbox.ts compose their full wrapper
+ * scripts from these shared fragments, eliminating code duplication.
+ * The sidecar (packages/sandbox-sidecar) maintains its own copy since
+ * it runs as a standalone container.
+ */
+
+/**
+ * Shared security guard + environment setup.
+ *
+ * Expects the caller to have already imported: sys, json, os, ast
+ *
+ * Expects the caller to have already defined:
+ * - _chart_dir: str — path to chart output directory
+ * - _report_error(msg: str) — emit error and exit the process
+ *
+ * Enforces the AST-based import guard (exits via _report_error on violation),
+ * configures a headless matplotlib backend, ensures _chart_dir exists,
+ * and makes available: _user_code, chart_path().
+ */
+export const PYTHON_SECURITY_AND_SETUP = `# --- Import guard (AST-based enforcement) ---
+_BLOCKED_MODULES = {
+    "subprocess", "os", "socket", "shutil", "sys", "ctypes", "importlib",
+    "code", "signal", "multiprocessing", "threading", "pty", "fcntl",
+    "termios", "resource", "posixpath",
+    "http", "urllib", "requests", "httpx", "aiohttp", "webbrowser",
+    "pickle", "tempfile", "pathlib",
+}
+_BLOCKED_BUILTINS = {
+    "compile", "exec", "eval", "__import__", "open", "breakpoint",
+    "getattr", "globals", "locals", "vars", "dir", "delattr", "setattr",
+}
+
+_user_code = open(sys.argv[1]).read()
+try:
+    _tree = ast.parse(_user_code)
+except SyntaxError as e:
+    _report_error(f"SyntaxError: {e.msg} (line {e.lineno})")
+
+_blocked = None
+for _node in ast.walk(_tree):
+    if _blocked:
+        break
+    if isinstance(_node, ast.Import):
+        for _alias in _node.names:
+            _mod = _alias.name.split('.')[0]
+            if _mod in _BLOCKED_MODULES:
+                _blocked = f'Blocked import: "{_mod}" is not allowed'
+                break
+    elif isinstance(_node, ast.ImportFrom):
+        if _node.module:
+            _mod = _node.module.split('.')[0]
+            if _mod in _BLOCKED_MODULES:
+                _blocked = f'Blocked import: "{_mod}" is not allowed'
+    elif isinstance(_node, ast.Call):
+        _name = None
+        if isinstance(_node.func, ast.Name):
+            _name = _node.func.id
+        elif isinstance(_node.func, ast.Attribute):
+            _name = _node.func.attr
+        if _name and _name in _BLOCKED_BUILTINS:
+            _blocked = f'Blocked builtin: "{_name}()" is not allowed'
+
+if _blocked:
+    _report_error(_blocked)
+
+# Configure matplotlib for headless rendering
+try:
+    import matplotlib
+    matplotlib.use('Agg')
+except ImportError:
+    pass
+
+os.makedirs(_chart_dir, exist_ok=True)  # no-op if pre-created by host (e.g. nsjail bind-mount)
+
+def chart_path(n=0):
+    return os.path.join(_chart_dir, f"chart_{n}.png")`;
+
+/**
+ * Non-streaming execution: stdout capture, exec, result collection.
+ *
+ * Expects: sys, json, io, base64, glob, os imported;
+ * _marker, _chart_dir, _user_code, data, df, chart_path defined.
+ *
+ * Captures stdout, executes user code in an isolated namespace, collects
+ * charts and structured results, and emits a single JSON result via _marker.
+ */
+export const PYTHON_EXEC_AND_COLLECT = `# --- Execute user code in isolated namespace ---
+_old_stdout = sys.stdout
+sys.stdout = _captured = io.StringIO()
+
+_user_ns = {"chart_path": chart_path, "data": data, "df": df}
+_atlas_error = None
+try:
+    exec(_user_code, _user_ns)
+except Exception as e:
+    _atlas_error = f"{type(e).__name__}: {e}"
+
+_output = _captured.getvalue()
+sys.stdout = _old_stdout
+
+# --- Collect results ---
+_charts = []
+for f in sorted(glob.glob(os.path.join(_chart_dir, "chart_*.png"))):
+    with open(f, "rb") as fh:
+        _charts.append({"base64": base64.b64encode(fh.read()).decode(), "mimeType": "image/png"})
+
+_result = {"success": _atlas_error is None}
+if _output.strip():
+    _result["output"] = _output.strip()
+if _atlas_error:
+    _result["error"] = _atlas_error
+
+if "_atlas_table" in _user_ns:
+    _result["table"] = _user_ns["_atlas_table"]
+
+if "_atlas_chart" in _user_ns:
+    _ac = _user_ns["_atlas_chart"]
+    if isinstance(_ac, dict):
+        _result["rechartsCharts"] = [_ac]
+    elif isinstance(_ac, list):
+        _result["rechartsCharts"] = _ac
+
+if _charts:
+    _result["charts"] = _charts
+
+print(_marker + json.dumps(_result), file=_old_stdout)`;