@useatlas/create 0.0.6 → 0.0.7

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

@@ -16,12 +16,16 @@ import { tool } from "ai";
 import { z } from "zod";
 import { createLogger } from "@atlas/api/lib/logger";
 import { withSpan } from "@atlas/api/lib/tracing";
+import { getConfig } from "@atlas/api/lib/config";
+import { useVercelSandbox, useSidecar } from "./backends/detect";
+import { getStreamWriter } from "./python-stream";
 
 const log = createLogger("python");
 
 // --- Import guard (defense-in-depth) ---
 
-const BLOCKED_MODULES = new Set([
+/** Default blocked modules — the baseline when no config overrides are set. */
+export const DEFAULT_BLOCKED_MODULES = new Set([
   "subprocess",
   "os",
   "socket",
@@ -51,6 +55,48 @@ const BLOCKED_MODULES = new Set([
   "pathlib",
 ]);
 
+/**
+ * Critical modules that can never be unblocked, even with `python.allowModules`.
+ * These provide direct OS/process access that no sandbox config should override.
+ */
+export const CRITICAL_MODULES = new Set(["os", "subprocess", "sys", "shutil"]);
+
+/**
+ * Build the effective blocked module set from the default list + config overrides.
+ * Called on each validation rather than cached at module load, so any config
+ * singleton update (e.g. test mocking or future hot-reload) takes effect on
+ * the next call.
+ */
+export function getEffectiveBlockedModules(): Set<string> {
+  const config = getConfig();
+  const pythonConfig = config?.python;
+
+  if (!pythonConfig) return DEFAULT_BLOCKED_MODULES;
+
+  const { blockedModules = [], allowModules = [] } = pythonConfig;
+
+  // Reject attempts to unblock critical modules
+  const criticalViolations = allowModules.filter((m) => CRITICAL_MODULES.has(m));
+  if (criticalViolations.length > 0) {
+    throw new Error(
+      `Cannot unblock critical Python modules: ${criticalViolations.join(", ")}. ` +
+      `These modules (${[...CRITICAL_MODULES].join(", ")}) are blocked regardless of configuration.`,
+    );
+  }
+
+  const allowSet = new Set(allowModules);
+  const effective = new Set<string>();
+
+  for (const mod of DEFAULT_BLOCKED_MODULES) {
+    if (!allowSet.has(mod)) effective.add(mod);
+  }
+  for (const mod of blockedModules) {
+    effective.add(mod);
+  }
+
+  return effective;
+}
+
 const BLOCKED_BUILTINS = new Set([
   "compile",
   "exec",
@@ -71,7 +117,8 @@ const BLOCKED_BUILTINS = new Set([
  * Validate Python code for blocked imports and dangerous builtins.
  *
  * Uses Python's own `ast` module to parse the code, then checks for:
- * - `import X` / `from X import ...` where X is in BLOCKED_MODULES
+ * - `import X` / `from X import ...` where X is in the effective blocked set
+ *   (see {@link getEffectiveBlockedModules})
  * - Calls to blocked builtins (exec, eval, compile, __import__, open, getattr, etc.)
  *
  * This is defense-in-depth — the sidecar container is the security boundary.
@@ -158,8 +205,17 @@ json.dump({"imports": imports, "calls": calls}, sys.stdout)
   }
 
   // Check imports
+  let blockedModules: Set<string>;
+  try {
+    blockedModules = getEffectiveBlockedModules();
+  } catch (err) {
+    const detail = err instanceof Error ? err.message : String(err);
+    log.error({ err: detail }, "Failed to compute effective blocked modules — likely a config error");
+    return { safe: false, reason: detail };
+  }
+
   for (const mod of result.imports ?? []) {
-    if (BLOCKED_MODULES.has(mod)) {
+    if (blockedModules.has(mod)) {
       return { safe: false, reason: `Blocked import: "${mod}" is not allowed` };
     }
   }
@@ -202,14 +258,28 @@ export type PythonResult =
       output?: string;
     };
 
+// --- Progress events ---
+
+/** Events emitted during streaming Python execution. */
+export type PythonProgressEvent =
+  | { type: "stdout"; content: string }
+  | { type: "chart"; chart: PythonChart }
+  | { type: "recharts"; chart: RechartsChart };
+
 // --- Backend interface ---
 
 /**
  * Python execution backend. Implementations handle isolation (sidecar, Vercel sandbox, nsjail).
  * Each backend receives validated code + optional data and returns a structured result.
+ * Backends that support streaming implement the optional `execStream` method.
  */
 export interface PythonBackend {
   exec(code: string, data?: { columns: string[]; rows: unknown[][] }): Promise<PythonResult>;
+  execStream?(
+    code: string,
+    data: { columns: string[]; rows: unknown[][] } | undefined,
+    onProgress: (event: PythonProgressEvent) => void,
+  ): Promise<PythonResult>;
 }
 
 // --- Backend selection ---
@@ -226,16 +296,18 @@ export interface PythonBackend {
  */
 async function getPythonBackend(): Promise<PythonBackend | { error: string }> {
   // 1. Sidecar
-  const sidecarUrl = process.env.ATLAS_SANDBOX_URL;
-  if (sidecarUrl) {
-    const { executePythonViaSidecar } = await import("./python-sidecar");
+  if (useSidecar()) {
+    const sidecarUrl = process.env.ATLAS_SANDBOX_URL!;
+    const { executePythonViaSidecar, executePythonViaSidecarStream } = await import("./python-sidecar");
     return {
       exec: (code, data) => executePythonViaSidecar(sidecarUrl, code, data),
+      execStream: (code, data, onProgress) =>
+        executePythonViaSidecarStream(sidecarUrl, code, data, onProgress),
     };
   }
 
   // 2. Vercel sandbox (Python 3.13 runtime)
-  if (process.env.ATLAS_RUNTIME === "vercel" || process.env.VERCEL) {
+  if (useVercelSandbox()) {
     let createPythonSandboxBackend;
     try {
       ({ createPythonSandboxBackend } = await import("./python-sandbox"));
@@ -250,7 +322,7 @@ async function getPythonBackend(): Promise<PythonBackend | { error: string }> {
   // 3. nsjail explicit (ATLAS_SANDBOX=nsjail) — hard-fail
   if (process.env.ATLAS_SANDBOX === "nsjail") {
     try {
-      const { findNsjailBinary } = await import("./explore-nsjail");
+      const { findNsjailBinary } = await import("./backends/nsjail");
       const nsjailPath = findNsjailBinary();
       if (nsjailPath) {
         const { createPythonNsjailBackend } = await import("./python-nsjail");
@@ -267,7 +339,7 @@ async function getPythonBackend(): Promise<PythonBackend | { error: string }> {
 
   // 4. nsjail auto-detect
   try {
-    const { findNsjailBinary } = await import("./explore-nsjail");
+    const { findNsjailBinary } = await import("./backends/nsjail");
     const nsjailPath = findNsjailBinary();
     if (nsjailPath) {
       const { createPythonNsjailBackend } = await import("./python-nsjail");
@@ -280,7 +352,7 @@ async function getPythonBackend(): Promise<PythonBackend | { error: string }> {
       "code" in err &&
       (err as NodeJS.ErrnoException).code === "MODULE_NOT_FOUND"
     ) {
-      log.debug("explore-nsjail module not available, skipping nsjail Python backend");
+      log.debug("nsjail module not available, skipping nsjail Python backend");
     } else {
       const detail = err instanceof Error ? err.message : String(err);
       log.error({ err: detail }, "Unexpected error initializing nsjail Python backend");
@@ -324,7 +396,7 @@ Blocked: subprocess, os, socket, shutil, sys, ctypes, importlib, exec(), eval(),
       .describe("Optional data payload from a previous SQL query (columns + rows)"),
   }),
 
-  execute: async ({ code, explanation, data }) => {
+  execute: async ({ code, explanation, data }, options) => {
     // 0. Resolve backend
     const backend = await getPythonBackend();
     if ("error" in backend) {
@@ -339,18 +411,46 @@ Blocked: subprocess, os, socket, shutil, sys, ctypes, importlib, exec(), eval(),
       return { success: false, error: validation.reason };
     }
 
-    // 2. Execute via selected backend
+    // 2. Build streaming progress callback if stream writer is available
+    const writer = getStreamWriter();
+    const toolCallId = options?.toolCallId;
+    const canStream = writer && toolCallId && backend.execStream;
+
+    const onProgress = canStream
+      ? (event: PythonProgressEvent) => {
+          try {
+            // Custom data part — the AI SDK's typed data parts require compile-time
+            // registration via UIMessage generics. We bypass with a cast because Atlas
+            // uses dynamic data parts consumed via onData on the client.
+            writer.write({
+              type: "data-python-progress" as const,
+              id: toolCallId,
+              data: event,
+            } as unknown as Parameters<typeof writer.write>[0]);
+          } catch (err) {
+            log.debug(
+              { err: err instanceof Error ? err.message : String(err), toolCallId },
+              "Stream writer closed, Python progress events will not be delivered",
+            );
+          }
+        }
+      : undefined;
+
+    // 3. Execute via selected backend
     const start = performance.now();
     try {
       const result = await withSpan(
         "atlas.python.execute",
-        { "code.length": code.length },
-        () => backend.exec(code, data),
+        { "code.length": code.length, streaming: !!onProgress },
+        () =>
+          onProgress && backend.execStream
+            ? backend.execStream(code, data, onProgress)
+            : backend.exec(code, data),
       );
       const durationMs = Math.round(performance.now() - start);
 
       log.debug(
-        { durationMs, success: result.success, hasCharts: result.success && !!result.charts?.length },
+        { durationMs, success: result.success, hasCharts: result.success && !!result.charts?.length, streaming: !!onProgress },
         "python execution",
       );
 

package/templates/nextjs-standalone/src/lib/tools/registry.ts

@@ -170,16 +170,25 @@ defaultRegistry.register({
 
 defaultRegistry.freeze();
 
+interface BuildRegistryResult {
+  registry: ToolRegistry;
+  warnings: string[];
+}
+
 /**
  * Build a dynamic ToolRegistry with optional action and Python support.
  *
  * Python tool is included when `ATLAS_PYTHON_ENABLED=true`.
  * Action tools are included when `includeActions` is true.
+ *
+ * Returns both the registry and any warnings about tools that failed to load.
+ * Fatal misconfigurations (e.g. Python enabled without sandbox URL) still throw.
  */
 export async function buildRegistry(options?: {
   includeActions?: boolean;
-}): Promise<ToolRegistry> {
+}): Promise<BuildRegistryResult> {
   const registry = new ToolRegistry();
+  const warnings: string[] = [];
 
   registry.register({
     name: "explore",
@@ -238,11 +247,14 @@ export async function buildRegistry(options?: {
         { err: err instanceof Error ? err : new Error(String(err)) },
         "Failed to load action tools — JIRA and email actions will be unavailable",
       );
+      warnings.push(
+        "Action tools (JIRA, email) failed to load and are unavailable for this session. Inform the user and suggest they check server logs or retry later.",
+      );
     }
   }
 
   registry.freeze();
-  return registry;
+  return { registry, warnings };
 }
 
 export { defaultRegistry };