feat(transforms): implement code-mode transform and quickjs worker pool

nitrostack/transforms/__init__.py

@@ -13,9 +13,11 @@
 )
 from nitrostack.transforms.validation import ToolArgumentError, validate_tool_arguments
 from nitrostack.transforms.synthetic import SyntheticTool
+from nitrostack.transforms.code_mode import CodeModeTransform
 
 __all__ = [
     "CatalogTransform", "McpTransform", "McpTransformBase", "TransformContext",
     "TransformError", "TransformPipeline", "TransformRegistry",
     "ToolArgumentError", "validate_tool_arguments", "active_bypass", "bypassing", "SyntheticTool",
+    "CodeModeTransform",
 ]

nitrostack/transforms/code_mode/__init__.py

@@ -0,0 +1,35 @@
+"""Code Mode transform for NitroStack."""
+from __future__ import annotations
+
+from nitrostack.transforms.code_mode.destructive_guard import (
+    DestructiveToolError,
+    assert_tool_allowed,
+    is_destructive_tool,
+)
+from nitrostack.transforms.code_mode.legible_error import format_legible_error
+from nitrostack.transforms.code_mode.sandbox_interface import SandboxProtocol, ToolsApi
+from nitrostack.transforms.code_mode.transform import CodeModeTransform
+from nitrostack.transforms.code_mode.types import (
+    CodeModeLimits,
+    ExecuteScriptInput,
+    GetSchemaInput,
+    SandboxExecutionResult,
+    SearchToolsInput,
+)
+from nitrostack.transforms.code_mode.worker_pool import WorkerPool
+
+__all__ = [
+    "CodeModeTransform",
+    "CodeModeLimits",
+    "SandboxExecutionResult",
+    "SandboxProtocol",
+    "ToolsApi",
+    "WorkerPool",
+    "is_destructive_tool",
+    "assert_tool_allowed",
+    "DestructiveToolError",
+    "format_legible_error",
+    "SearchToolsInput",
+    "GetSchemaInput",
+    "ExecuteScriptInput",
+]

nitrostack/transforms/code_mode/destructive_guard.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+class DestructiveToolError(Exception):
+    """Raised when a destructive tool is called while allowDestructive is false."""
+    pass
+
+
+def _get_annotation(annotations: Any, *keys: str) -> Optional[bool]:
+    if annotations is None:
+        return None
+    for k in keys:
+        if isinstance(annotations, dict):
+            if k in annotations:
+                return annotations[k]
+        elif hasattr(annotations, k):
+            return getattr(annotations, k)
+    return None
+
+
+def is_destructive_tool(tool: Any) -> bool:
+    """A tool is destructive unless readOnlyHint is true or destructiveHint is explicitly false."""
+    annotations = getattr(tool, "annotations", None)
+    if annotations is None and hasattr(tool, "config"):
+        annotations = getattr(tool.config, "annotations", None)
+
+    read_only = _get_annotation(annotations, "readOnlyHint", "read_only_hint", "read_only")
+    if read_only is True:
+        return False
+
+    destructive = _get_annotation(annotations, "destructiveHint", "destructive_hint")
+    if destructive is False:
+        return False
+
+    return True
+
+
+def assert_tool_allowed(tool: Any, allow_destructive: bool) -> None:
+    """Reject tool if destructive and allow_destructive is False."""
+    if not allow_destructive and is_destructive_tool(tool):
+        name = getattr(tool, "name", str(tool))
+        raise DestructiveToolError(
+            f"Tool '{name}' is destructive and cannot be executed in code-mode when allowDestructive is false."
+        )

nitrostack/transforms/code_mode/guest_evaluator.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional
+from multiprocessing.connection import Connection
+
+from nitrostack.transforms.code_mode.legible_error import format_legible_error
+from nitrostack.transforms.code_mode.types import CodeModeLimits, SandboxExecutionResult
+
+logger = logging.getLogger(__name__)
+
+JS_PRELUDE = """
+const console = {
+    log: (...args) => { __nitrostack_log('info', __formatArgs(args)); },
+    info: (...args) => { __nitrostack_log('info', __formatArgs(args)); },
+    warn: (...args) => { __nitrostack_log('warn', __formatArgs(args)); },
+    error: (...args) => { __nitrostack_log('error', __formatArgs(args)); }
+};
+
+function __formatArgs(args) {
+    return args.map(a => {
+        if (typeof a === 'object' && a !== null) {
+            try { return JSON.stringify(a); } catch(e) { return String(a); }
+        }
+        return String(a);
+    }).join(' ');
+}
+
+const tools = new Proxy({
+    call(name, args) {
+        const payload = JSON.stringify(args || {});
+        const resStr = __nitrostack_call_tool(String(name), payload);
+        const parsed = JSON.parse(resStr);
+        if (parsed && parsed.__nitrostack_error) {
+            const err = new Error(parsed.__nitrostack_error);
+            err.name = parsed.__nitrostack_error_type || 'Error';
+            throw err;
+        }
+        return parsed.__nitrostack_value;
+    }
+}, {
+    get(target, prop) {
+        if (prop in target) return target[prop];
+        return function(args) {
+            return target.call(prop, args);
+        };
+    }
+});
+
+const call_tool = (name, args) => tools.call(name, args);
+const call = (name, args) => tools.call(name, args);
+"""
+
+
+def get_quickjs_module():
+    try:
+        import quickjs
+        return quickjs
+    except ImportError as exc:
+        raise RuntimeError(
+            "The 'quickjs' package is required for Code Mode. "
+            "Install it with: pip install 'nitrostack[code-mode]'"
+        ) from exc
+
+
+def worker_process_main(conn: Connection) -> None:
+    """Main loop for isolated worker process."""
+    quickjs = get_quickjs_module()
+    conn.send({"type": "ready"})
+
+    while True:
+        try:
+            if not conn.poll(None):
+                break
+            msg = conn.recv()
+        except (EOFError, KeyboardInterrupt):
+            break
+
+        cmd = msg.get("cmd")
+        if cmd == "shutdown":
+            break
+
+        if cmd == "run":
+            code = msg["code"]
+            limits_dict = msg["limits"]
+            memory_limit_mb = limits_dict.get("memory_limit_mb", 100)
+
+            ctx = quickjs.Context()
+            memory_limit_bytes = int(memory_limit_mb * 1024 * 1024)
+            ctx.set_memory_limit(memory_limit_bytes)
+
+            logs: List[str] = []
+            tool_call_count = 0
+
+            def log_bridge(level: str, message: str) -> None:
+                logs.append(f"[{level}] {message}")
+
+            def tool_call_bridge(name: str, args_json: str) -> str:
+                nonlocal tool_call_count
+                tool_call_count += 1
+                try:
+                    conn.send({
+                        "type": "tool_call",
+                        "name": name,
+                        "args_json": args_json,
+                        "call_index": tool_call_count,
+                    })
+                    response = conn.recv()
+                    if response.get("status") == "error":
+                        return json.dumps({
+                            "__nitrostack_error": response.get("error", "Tool execution error"),
+                            "__nitrostack_error_type": response.get("error_type", "Error"),
+                        })
+                    return json.dumps({
+                        "__nitrostack_value": response.get("value")
+                    })
+                except Exception as e:
+                    return json.dumps({
+                        "__nitrostack_error": format_legible_error(e),
+                        "__nitrostack_error_type": type(e).__name__,
+                    })
+
+            ctx.add_callable("__nitrostack_log", log_bridge)
+            ctx.add_callable("__nitrostack_call_tool", tool_call_bridge)
+
+            # Evaluate prelude
+            ctx.eval(JS_PRELUDE)
+
+            start_time = time.perf_counter()
+            success = False
+            value = None
+            error = None
+
+            try:
+                try:
+                    raw_result = ctx.eval(code)
+                except quickjs.JSException as exc:
+                    if "return not in a function" in str(exc):
+                        wrapped = "(function() {\n" + code + "\n})()"
+                        raw_result = ctx.eval(wrapped)
+                    else:
+                        raise
+
+                if hasattr(raw_result, "json"):
+                    try:
+                        value = json.loads(raw_result.json())
+                    except Exception:
+                        value = str(raw_result)
+                else:
+                    value = raw_result
+
+                success = True
+            except quickjs.JSException as exc:
+                err_str = str(exc)
+                if "out of memory" in err_str:
+                    error = f"MemoryLimitExceededError: Script exceeded memory limit of {memory_limit_mb}MB."
+                else:
+                    error = format_legible_error(exc)
+            except Exception as exc:
+                error = format_legible_error(exc)
+            finally:
+                duration_ms = (time.perf_counter() - start_time) * 1000.0
+                memory_used_mb = 0.0
+                try:
+                    mem_info = ctx.memory()
+                    if isinstance(mem_info, dict):
+                        used_bytes = mem_info.get("memory_used_size", 0)
+                    else:
+                        used_bytes = getattr(mem_info, "memory_used_size", 0)
+                    memory_used_mb = used_bytes / (1024 * 1024)
+                except Exception:
+                    pass
+
+                try:
+                    ctx.gc()
+                except Exception:
+                    pass
+
+            conn.send({
+                "type": "run_result",
+                "success": success,
+                "value": value,
+                "error": error,
+                "logs": logs,
+                "tool_call_count": tool_call_count,
+                "duration_ms": duration_ms,
+                "memory_used_mb": memory_used_mb,
+            })

nitrostack/transforms/code_mode/legible_error.py

@@ -0,0 +1,43 @@
+"""Format legible error messages for models without leaking host paths or stack frames."""
+from __future__ import annotations
+
+import re
+
+_HOST_PATH_RE = re.compile(
+    r"(?:/[a-zA-Z0-9_\-\.]+)+|(?:[a-zA-Z]:\\[a-zA-Z0-9_\-\.\\]+)"
+)
+
+# Match Python traceback line patterns
+_PY_TRACEBACK_LINE_RE = re.compile(
+    r"^\s*File\s+\".*?\",\s+line\s+\d+.*$",
+    re.MULTILINE
+)
+
+
+def format_legible_error(exc: BaseException | str) -> str:
+    """Format an error message so it is readable (message and location) and leaks no host paths or stack frames."""
+    raw = str(exc)
+
+    # Clean up input markers
+    raw = raw.replace("<eval>", "<script>").replace("<input>", "<script>")
+
+    # Remove python traceback lines
+    lines = []
+    for line in raw.splitlines():
+        if _PY_TRACEBACK_LINE_RE.match(line):
+            continue
+        if "Traceback (most recent call last):" in line:
+            continue
+        lines.append(line)
+
+    cleaned = "\n".join(lines).strip()
+
+    # Scrub any host filesystem paths that may remain
+    def _path_replacer(match: re.Match) -> str:
+        matched = match.group(0)
+        if "<script>" in matched:
+            return matched
+        return "<sanitized-path>"
+
+    cleaned = _HOST_PATH_RE.sub(_path_replacer, cleaned)
+    return cleaned or "ScriptExecutionError: Unknown error occurred."

nitrostack/transforms/code_mode/sandbox_interface.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+from nitrostack.transforms.code_mode.types import CodeModeLimits, SandboxExecutionResult
+
+
+@runtime_checkable
+class ToolsApi(Protocol):
+    async def call(self, name: str, arguments: dict[str, Any]) -> Any:
+        ...
+
+
+@runtime_checkable
+class SandboxProtocol(Protocol):
+    async def run(
+        self,
+        code: str,
+        tools_api: Any,
+        limits: CodeModeLimits,
+    ) -> SandboxExecutionResult:
+        ...