diff --git a/Makefile b/Makefile index 80fbd3a..2dbcc76 100644 --- a/Makefile +++ b/Makefile @@ -137,7 +137,7 @@ image-minimal: sudo bash images/templates/minimal/build.sh image-python: - sudo bash images/templates/python311/build.sh + sudo bash images/templates/python312/build.sh image-node: sudo bash images/templates/node20/build.sh diff --git a/images/wrenn-init.sh b/images/wrenn-init.sh index 266e516..d83be1c 100644 --- a/images/wrenn-init.sh +++ b/images/wrenn-init.sh @@ -20,6 +20,14 @@ echo "+cpu +memory +io" > /sys/fs/cgroup/cgroup.subtree_control 2>/dev/null || t # Set hostname hostname sandbox +# Configure networking if the kernel ip= boot arg did not already set it up. +if ! ip addr show eth0 2>/dev/null | grep -q "169.254.0.21"; then + ip link set lo up 2>/dev/null || true + ip link set eth0 up 2>/dev/null || true + ip addr add 169.254.0.21/30 dev eth0 2>/dev/null || true + ip route add default via 169.254.0.22 2>/dev/null || true +fi + # Configure DNS resolver. echo "nameserver 8.8.8.8" > /etc/resolv.conf echo "nameserver 8.8.4.4" >> /etc/resolv.conf diff --git a/internal/recipe/context.go b/internal/recipe/context.go index db4c39c..51e44c1 100644 --- a/internal/recipe/context.go +++ b/internal/recipe/context.go @@ -1,6 +1,9 @@ package recipe -import "strings" +import ( + "regexp" + "strings" +) // ExecContext holds mutable state that persists across recipe steps. // It is initialized empty and updated by ENV and WORKDIR steps. @@ -9,6 +12,12 @@ type ExecContext struct { EnvVars map[string]string } +// This regex matches: +// 1. $$ (escaped dollar) +// 2. ${VAR} or ${} (braced variable, possibly empty) +// 3. $VAR (bare variable) +var envRegex = regexp.MustCompile(`\$\$|\$\{([a-zA-Z0-9_]*)\}|\$([a-zA-Z0-9_]+)`) + // WrappedCommand returns the full shell command for a RUN step with context // applied. The result is passed as the argument to /bin/sh -c. // @@ -56,6 +65,32 @@ func (c *ExecContext) shellPrefix() string { return sb.String() } +// expandEnv replaces $var and ${var} placeholders in the string s with their +// corresponding values from the vars map. +// It supports escaping with $$, which is replaced by a single $. +// If a variable is not found in the vars map, it is replaced with an empty +// string. +func expandEnv(s string, vars map[string]string) string { + return envRegex.ReplaceAllStringFunc(s, func(match string) string { + if match == "$$" { + return "$" + } + + var name string + if len(match) > 1 && match[1] == '{' { + name = match[2 : len(match)-1] + } else { + name = match[1:] + } + + if v, ok := vars[name]; ok { + return v + } + + return "" + }) +} + // shellescape wraps s in single quotes, escaping any embedded single quotes. // This is POSIX-safe for paths, env values, and shell commands. func shellescape(s string) string { diff --git a/internal/recipe/context_test.go b/internal/recipe/context_test.go index b00dfce..3074353 100644 --- a/internal/recipe/context_test.go +++ b/internal/recipe/context_test.go @@ -4,10 +4,11 @@ import "testing" func TestExecContext_WrappedCommand(t *testing.T) { tests := []struct { - name string - ctx ExecContext - cmd string - want string + name string + ctx ExecContext + cmd string + want string + wantOneOf []string }{ { name: "no context", @@ -45,12 +46,35 @@ func TestExecContext_WrappedCommand(t *testing.T) { cmd: "echo $MSG", want: "MSG='it'\\''s fine' /bin/sh -c 'echo $MSG'", }, + { + name: "env expansion with pre-expanded PATH", + ctx: ExecContext{ + EnvVars: map[string]string{"PATH": "/usr/bin", "FOO": "/opt/venv/bin:/usr/bin"}, + }, + cmd: "make build", + // Map iteration order is non-deterministic; accept either ordering. + wantOneOf: []string{ + "FOO='/opt/venv/bin:/usr/bin' PATH='/usr/bin' /bin/sh -c 'make build'", + "PATH='/usr/bin' FOO='/opt/venv/bin:/usr/bin' /bin/sh -c 'make build'", + }, + }, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { got := tc.ctx.WrappedCommand(tc.cmd) - if got != tc.want { + if len(tc.wantOneOf) > 0 { + matched := false + for _, w := range tc.wantOneOf { + if got == w { + matched = true + break + } + } + if !matched { + t.Errorf("WrappedCommand(%q)\n got %q\n want one of %q", tc.cmd, got, tc.wantOneOf) + } + } else if got != tc.want { t.Errorf("WrappedCommand(%q)\n got %q\n want %q", tc.cmd, got, tc.want) } }) @@ -94,6 +118,109 @@ func TestExecContext_StartCommand(t *testing.T) { } } +func TestExpandEnv(t *testing.T) { + tests := []struct { + s string + vars map[string]string + want string + }{ + { + s: "hello", + vars: nil, + want: "hello", + }, + { + s: "$PATH", + vars: map[string]string{"PATH": "/usr/bin"}, + want: "/usr/bin", + }, + { + s: "${PATH}", + vars: map[string]string{"PATH": "/usr/bin"}, + want: "/usr/bin", + }, + { + s: "/opt/venv/bin:$PATH", + vars: map[string]string{"PATH": "/usr/bin"}, + want: "/opt/venv/bin:/usr/bin", + }, + { + s: "${HOME}/code", + vars: map[string]string{"HOME": "/root"}, + want: "/root/code", + }, + { + s: "hello $USER", + vars: map[string]string{"USER": "admin"}, + want: "hello admin", + }, + { + s: "$UNSET", + vars: map[string]string{"PATH": "/usr/bin"}, + want: "", + }, + { + s: "${UNSET}", + vars: map[string]string{"PATH": "/usr/bin"}, + want: "", + }, + { + s: "$$", + vars: map[string]string{"PATH": "/usr/bin"}, + want: "$", + }, + { + s: "price is $$100", + vars: nil, + want: "price is $100", + }, + { + s: "$FOO:$BAR", + vars: map[string]string{"FOO": "a", "BAR": "b"}, + want: "a:b", + }, + { + s: "${FOO}_${BAR}", + vars: map[string]string{"FOO": "hello", "BAR": "world"}, + want: "hello_world", + }, + { + s: "no vars here", + vars: nil, + want: "no vars here", + }, + { + s: "$", + vars: nil, + want: "$", + }, + { + s: "${", + vars: nil, + want: "${", + }, + { + s: "${}", + vars: nil, + want: "", + }, + { + s: "$VAR1$VAR2", + vars: map[string]string{"VAR1": "a", "VAR2": "b"}, + want: "ab", + }, + } + + for _, tc := range tests { + t.Run(tc.s, func(t *testing.T) { + got := expandEnv(tc.s, tc.vars) + if got != tc.want { + t.Errorf("expandEnv(%q, %v)\n got %q\n want %q", tc.s, tc.vars, got, tc.want) + } + }) + } +} + func TestShellescape(t *testing.T) { tests := []struct { input string diff --git a/internal/recipe/executor.go b/internal/recipe/executor.go index 3df45dc..aa4b305 100644 --- a/internal/recipe/executor.go +++ b/internal/recipe/executor.go @@ -68,7 +68,7 @@ func Execute( if bctx.EnvVars == nil { bctx.EnvVars = make(map[string]string) } - bctx.EnvVars[st.Key] = st.Value + bctx.EnvVars[st.Key] = expandEnv(st.Value, bctx.EnvVars) entries = append(entries, BuildLogEntry{Step: step, Phase: phase, Cmd: st.Raw, Ok: true}) case KindWORKDIR: diff --git a/internal/recipe/healthcheck.go b/internal/recipe/healthcheck.go new file mode 100644 index 0000000..6e48f73 --- /dev/null +++ b/internal/recipe/healthcheck.go @@ -0,0 +1,94 @@ +package recipe + +import ( + "fmt" + "strconv" + "strings" + "time" +) + +// HealthcheckConfig holds the parsed configuration for a build healthcheck. +// A healthcheck is a shell command that is executed repeatedly inside the +// sandbox until it succeeds or the retry/timeout budget is exhausted. +// +// Retries of 0 means unlimited retries (bounded only by the overall deadline) +type HealthcheckConfig struct { + Cmd string + Interval time.Duration + Timeout time.Duration + StartPeriod time.Duration + Retries int // 0 = unlimited +} + +// ParseHealthcheck parses a healthcheck string with optional flag prefix into +// a HealthcheckConfig. The syntax is: +// +// [--interval=] [--timeout=] [--start-period=] +// [--retries=] +// +// Flags must use the form --flag=value. The first token that does not start +// with "--" and everything after it is treated as the command. Defaults: +// interval=3s, timeout=10s, start-period=0, retries=0 (unlimited) +func ParseHealthcheck(s string) (HealthcheckConfig, error) { + s = strings.TrimSpace(s) + if s == "" { + return HealthcheckConfig{}, fmt.Errorf("empty healthcheck") + } + + hc := HealthcheckConfig{ + Interval: 3 * time.Second, + Timeout: 10 * time.Second, + } + + tokens := strings.Fields(s) + cmdIndex := -1 + + for i, token := range tokens { + if !strings.HasPrefix(token, "--") { + cmdIndex = i + break + } + + parts := strings.SplitN(token, "=", 2) + if len(parts) != 2 { + return HealthcheckConfig{}, fmt.Errorf("malformed flag (missing '='): %q", token) + } + + key, val := parts[0], parts[1] + switch key { + case "--interval": + d, err := time.ParseDuration(val) + if err != nil { + return HealthcheckConfig{}, fmt.Errorf("parse interval: %w", err) + } + hc.Interval = d + case "--timeout": + d, err := time.ParseDuration(val) + if err != nil { + return HealthcheckConfig{}, fmt.Errorf("parse timeout: %w", err) + } + hc.Timeout = d + case "--start-period": + d, err := time.ParseDuration(val) + if err != nil { + return HealthcheckConfig{}, fmt.Errorf("parse start period: %w", err) + } + hc.StartPeriod = d + case "--retries": + r, err := strconv.Atoi(val) + if err != nil { + return HealthcheckConfig{}, fmt.Errorf("parse retries: %w", err) + } + hc.Retries = r + default: + return HealthcheckConfig{}, fmt.Errorf("unknown healthcheck flag: %q", token) + } + } + + if cmdIndex == -1 { + return HealthcheckConfig{}, fmt.Errorf("healthcheck has no command") + } + + hc.Cmd = strings.Join(tokens[cmdIndex:], " ") + return hc, nil +} diff --git a/internal/recipe/healthcheck_test.go b/internal/recipe/healthcheck_test.go new file mode 100644 index 0000000..528e109 --- /dev/null +++ b/internal/recipe/healthcheck_test.go @@ -0,0 +1,126 @@ +package recipe + +import ( + "testing" + "time" +) + +func TestParseHealthcheck(t *testing.T) { + tests := []struct { + name string + input string + want HealthcheckConfig + wantErr bool + }{ + { + name: "plain command", + input: "curl -f http://localhost:8080", + want: HealthcheckConfig{ + Cmd: "curl -f http://localhost:8080", + Interval: 3 * time.Second, + Timeout: 10 * time.Second, + }, + wantErr: false, + }, + { + name: "all flags", + input: "--interval=5s --timeout=2s --start-period=15s --retries=3 ping -c 1 8.8.8.8", + want: HealthcheckConfig{ + Cmd: "ping -c 1 8.8.8.8", + Interval: 5 * time.Second, + Timeout: 2 * time.Second, + StartPeriod: 15 * time.Second, + Retries: 3, + }, + wantErr: false, + }, + { + name: "partial flags", + input: "--timeout=5s my-custom-check --verbose", + want: HealthcheckConfig{ + Cmd: "my-custom-check --verbose", + Interval: 3 * time.Second, + Timeout: 5 * time.Second, + }, + wantErr: false, + }, + { + name: "retries only", + input: "--retries=5 test.sh", + want: HealthcheckConfig{ + Cmd: "test.sh", + Interval: 3 * time.Second, + Timeout: 10 * time.Second, + Retries: 5, + }, + wantErr: false, + }, + { + name: "empty string", + input: "", + wantErr: true, + }, + { + name: "whitespace only", + input: " \t \n ", + wantErr: true, + }, + { + name: "flags but no command", + input: "--interval=5s --retries=2", + wantErr: true, + }, + { + name: "unknown flag", + input: "--magic=true my-check", + wantErr: true, + }, + { + name: "invalid duration", + input: "--interval=5smiles check.sh", + wantErr: true, + }, + { + name: "invalid retries", + input: "--retries=five check.sh", + wantErr: true, + }, + { + name: "command with dashes", + input: "--interval=2s command-with-dash --flag=value", + want: HealthcheckConfig{ + Cmd: "command-with-dash --flag=value", + Interval: 2 * time.Second, + Timeout: 10 * time.Second, + }, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := ParseHealthcheck(tt.input) + if (err != nil) != tt.wantErr { + t.Errorf("ParseHealthcheck() error = %v, wantErr %v", err, tt.wantErr) + return + } + if !tt.wantErr { + if got.Cmd != tt.want.Cmd { + t.Errorf("Cmd got = %v, want %v", got.Cmd, tt.want.Cmd) + } + if got.Interval != tt.want.Interval { + t.Errorf("Interval got = %v, want %v", got.Interval, tt.want.Interval) + } + if got.Timeout != tt.want.Timeout { + t.Errorf("Timeout got = %v, want %v", got.Timeout, tt.want.Timeout) + } + if got.StartPeriod != tt.want.StartPeriod { + t.Errorf("StartPeriod got = %v, want %v", got.StartPeriod, tt.want.StartPeriod) + } + if got.Retries != tt.want.Retries { + t.Errorf("Retries got = %v, want %v", got.Retries, tt.want.Retries) + } + } + }) + } +} diff --git a/internal/service/build.go b/internal/service/build.go index 1108044..175f74e 100644 --- a/internal/service/build.go +++ b/internal/service/build.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "log/slog" + "strings" "sync" "time" @@ -23,8 +24,6 @@ import ( const ( buildQueueKey = "wrenn:build_queue" buildCommandTimeout = 30 * time.Second - healthcheckInterval = 1 * time.Second - healthcheckTimeout = 60 * time.Second ) // preBuildCmds run before the user recipe to prepare the build environment. @@ -321,11 +320,18 @@ func (s *BuildService) executeBuild(ctx context.Context, buildIDStr string) { panic(fmt.Sprintf("invalid post-build recipe: %v", err)) } - // Execute build phases: pre-build → user recipe → post-build. - // bctx carries working directory and env vars across all phases. var logs []recipe.BuildLogEntry step := 0 - bctx := &recipe.ExecContext{} + + envVars, err := s.fetchSandboxEnv(buildCtx, agent, sandboxIDStr) + if err != nil { + log.Warn("failed to fetch sandbox env, using defaults", "error", err) + envVars = map[string]string{ + "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin", + "HOME": "/root", + } + } + bctx := &recipe.ExecContext{EnvVars: envVars} runPhase := func(phase string, steps []recipe.Step, defaultTimeout time.Duration) bool { newEntries, nextStep, ok := recipe.Execute(buildCtx, phase, steps, sandboxIDStr, step, defaultTimeout, bctx, agent.Exec) @@ -365,8 +371,14 @@ func (s *BuildService) executeBuild(ctx context.Context, buildIDStr string) { // Healthcheck or direct snapshot. var sizeBytes int64 if build.Healthcheck != "" { - log.Info("running healthcheck", "cmd", build.Healthcheck) - if err := s.waitForHealthcheck(buildCtx, agent, sandboxIDStr, build.Healthcheck); err != nil { + hc, err := recipe.ParseHealthcheck(build.Healthcheck) + if err != nil { + s.destroySandbox(buildCtx, agent, sandboxIDStr) + s.failBuild(buildCtx, buildID, fmt.Sprintf("invalid healthcheck: %v", err)) + return + } + log.Info("running healthcheck", "cmd", hc.Cmd, "interval", hc.Interval, "timeout", hc.Timeout, "start_period", hc.StartPeriod, "retries", hc.Retries) + if err := s.waitForHealthcheck(buildCtx, agent, sandboxIDStr, hc); err != nil { s.destroySandbox(buildCtx, agent, sandboxIDStr) if buildCtx.Err() != nil { return @@ -445,36 +457,64 @@ func (s *BuildService) executeBuild(ctx context.Context, buildIDStr string) { log.Info("template build completed successfully", "name", build.Name) } -func (s *BuildService) waitForHealthcheck(ctx context.Context, agent buildAgentClient, sandboxIDStr, cmd string) error { - deadline := time.NewTimer(healthcheckTimeout) - defer deadline.Stop() - ticker := time.NewTicker(healthcheckInterval) +// waitForHealthcheck repeatedly executes the healthcheck command inside the +// sandbox according to the config's interval, timeout, start-period, and +// retries. +// During the start period, failures are not counted toward the retry budget. +// Returns nil on the first successful check, or an error if retries are +// exhausted, the deadline passes, or the context is cancelled. +func (s *BuildService) waitForHealthcheck(ctx context.Context, agent buildAgentClient, sandboxIDStr string, hc recipe.HealthcheckConfig) error { + ticker := time.NewTicker(hc.Interval) defer ticker.Stop() + // When retries > 0, set a deadline based on the retry budget. + // When retries == 0 (unlimited), rely solely on the parent context deadline. + var deadlineCh <-chan time.Time + if hc.Retries > 0 { + deadline := time.NewTimer(hc.StartPeriod + time.Duration(hc.Retries+1)*hc.Interval) + defer deadline.Stop() + deadlineCh = deadline.C + } + + startedAt := time.Now() + failCount := 0 + for { select { case <-ctx.Done(): return ctx.Err() - case <-deadline.C: - return fmt.Errorf("healthcheck timed out after %s", healthcheckTimeout) + case <-deadlineCh: + return fmt.Errorf("healthcheck timed out: exceeded %d attempts over %s", failCount, time.Since(startedAt)) case <-ticker.C: - execCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + execCtx, cancel := context.WithTimeout(ctx, hc.Timeout) resp, err := agent.Exec(execCtx, connect.NewRequest(&pb.ExecRequest{ SandboxId: sandboxIDStr, Cmd: "/bin/sh", - Args: []string{"-c", cmd}, - TimeoutSec: 10, + Args: []string{"-c", hc.Cmd}, + TimeoutSec: int32(hc.Timeout.Seconds()), })) cancel() if err != nil { slog.Debug("healthcheck exec error (retrying)", "error", err) + if time.Since(startedAt) >= hc.StartPeriod { + failCount++ + if hc.Retries > 0 && failCount >= hc.Retries { + return fmt.Errorf("healthcheck failed after %d retries: exec error: %w", failCount, err) + } + } continue } if resp.Msg.ExitCode == 0 { return nil } slog.Debug("healthcheck failed (retrying)", "exit_code", resp.Msg.ExitCode) + if time.Since(startedAt) >= hc.StartPeriod { + failCount++ + if hc.Retries > 0 && failCount >= hc.Retries { + return fmt.Errorf("healthcheck failed after %d retries: exit code %d", failCount, resp.Msg.ExitCode) + } + } } } } @@ -517,3 +557,49 @@ func (s *BuildService) destroySandbox(_ context.Context, agent buildAgentClient, slog.Warn("failed to destroy build sandbox", "sandbox_id", sandboxIDStr, "error", err) } } + +// fetchSandboxEnv executes the 'env' command inside the specified sandbox via +// the build agent and returns environment variables +func (s *BuildService) fetchSandboxEnv(ctx context.Context, + agent buildAgentClient, sandboxIDStr string) (map[string]string, error) { + resp, err := agent.Exec(ctx, connect.NewRequest(&pb.ExecRequest{ + SandboxId: sandboxIDStr, + Cmd: "/bin/sh", + Args: []string{"-c", "env"}, + TimeoutSec: 10, + })) + if err != nil { + return nil, fmt.Errorf("fetch env: %w", err) + } + + if resp.Msg.ExitCode != 0 { + return nil, fmt.Errorf("fetch env: command exited with code %d", + resp.Msg.ExitCode) + } + + return parseSandboxEnv(string(resp.Msg.Stdout)), nil +} + +// parseSandboxEnv converts the raw newline-separated output of an 'env' +// command into a map. +// It skips empty lines and malformed entries, and correctly handles values +// containing '='. +func parseSandboxEnv(raw string) map[string]string { + envVars := make(map[string]string) + + for line := range strings.SplitSeq(raw, "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + + parts := strings.SplitN(line, "=", 2) + if len(parts) != 2 { + continue + } + + envVars[parts[0]] = parts[1] + } + + return envVars +} diff --git a/recipes/python-interpreter-v0-beta.healthcheck b/recipes/python-interpreter-v0-beta.healthcheck new file mode 100644 index 0000000..ca2555c --- /dev/null +++ b/recipes/python-interpreter-v0-beta.healthcheck @@ -0,0 +1 @@ +--interval=5s --timeout=3s --start-period=3s --retries=3 python3 -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8888/api/status', timeout=3)" diff --git a/recipes/python-interpreter-v0-beta.recipefile b/recipes/python-interpreter-v0-beta.recipefile new file mode 100644 index 0000000..e83f5da --- /dev/null +++ b/recipes/python-interpreter-v0-beta.recipefile @@ -0,0 +1,7 @@ +RUN apt-get install -y --no-install-recommends python3 python3-pip python3-venv +RUN python3 -m venv /opt/venv +ENV PATH=/opt/venv/bin:$PATH + +RUN --timeout=5m pip install --no-cache-dir jupyter-server ipykernel + +START jupyter server --ServerApp.ip=0.0.0.0 --ServerApp.port=8888 --ServerApp.token='' --ServerApp.allow_origin='*' --ServerApp.disable_check_xsrf=True --no-browser --allow-root diff --git a/recipes/test-jupyter-kernel.py b/recipes/test-jupyter-kernel.py new file mode 100755 index 0000000..3f0d315 --- /dev/null +++ b/recipes/test-jupyter-kernel.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +import argparse +import json +import sys +import urllib.request +import uuid + +try: + import websocket +except ImportError: + print("websocket-client is required: pip install websocket-client") + sys.exit(1) + + +def create_kernel(base_url: str, token: str) -> str: + url = f"{base_url}/api/kernels" + headers = {} + if token: + headers["X-API-Key"] = token + + req = urllib.request.Request(url, method="POST", data=b"", headers=headers) + resp = urllib.request.urlopen(req) + data = json.loads(resp.read()) + kernel_id = data["id"] + print(f"Created kernel: {kernel_id}") + return kernel_id + + +def execute_code(ws: websocket.WebSocket, code: str) -> dict: + msg_id = str(uuid.uuid4()) + session_id = str(uuid.uuid4()) + msg = { + "header": { + "msg_type": "execute_request", + "msg_id": msg_id, + "username": "", + "session": session_id, + "version": "5.3", + "date": "", + }, + "parent_header": {}, + "metadata": {}, + "content": { + "code": code, + "silent": False, + "store_history": True, + "user_expressions": {}, + }, + "buffers": [], + "channel": "shell", + } + ws.send(json.dumps(msg)) + + result = {"stdout": "", "stderr": "", "output": None, "error": None} + + while True: + resp = json.loads(ws.recv()) + + # Filter out messages from other executions by matching msg_id + parent_id = resp.get("parent_header", {}).get("msg_id") + if parent_id != msg_id: + continue + + msg_type = resp.get("msg_type", "") + + if msg_type == "stream": + result["stdout"] += resp["content"]["text"] + elif msg_type == "error": + result["error"] = "\n".join(resp["content"].get("traceback", [])) + elif msg_type == "execute_result": + result["output"] = resp["content"]["data"] + elif msg_type == "status": + if resp["content"]["execution_state"] == "idle": + break + + return result + + +def main(): + parser = argparse.ArgumentParser( + description="Test Jupyter kernel state management in a sandbox" + ) + parser.add_argument( + "sandbox_id", + help="Sandbox ID (e.g. cl-8nxizn9ygtczplsnn9jve38be)", + ) + parser.add_argument( + "--domain", + default="localhost:8080", + help="Proxy domain (default: localhost:8080)", + ) + parser.add_argument( + "--port", + default="8888", + help="Jupyter port inside the sandbox (default: 8888)", + ) + parser.add_argument( + "--key", + default="", + help="Wrenn API Token", + ) + args = parser.parse_args() + + base_url = f"http://{args.port}-{args.sandbox_id}.{args.domain}" + ws_base = base_url.replace("http", "ws", 1) + + print(f"Testing Jupyter kernel at {base_url}") + print() + + kernel_id = create_kernel(base_url, args.key) + + ws_url = f"{ws_base}/api/kernels/{kernel_id}/channels" + + # Pass auth headers to the WebSocket if a token was provided + ws_headers = {} + if args.key: + ws_headers["X-API-Key"] = args.key + + ws = websocket.create_connection(ws_url, header=ws_headers) + print("Connected to kernel WebSocket") + print() + + tests = [ + ("variable assignment", "x = 42", None), + ("read variable", "x * 2", "84"), + ("import", "import math", None), + ("use import", "math.sqrt(144)", "12.0"), + ("function definition", "def greet(name): return f'hello {name}'", None), + # Fixed: Jupyter 'execute_result' strings include the literal single quotes + ("call function", "greet('sandbox')", "'hello sandbox'"), + ("list mutation", "items = [1, 2, 3]; items.append(4); items", "[1, 2, 3, 4]"), + ] + + passed = 0 + failed = 0 + + for name, code, expected in tests: + print(f" {name}: {code}") + result = execute_code(ws, code) + + if result["error"]: + print(f" ERROR: {result['error']}") + failed += 1 + continue + + output = result["stdout"].strip() + if not output and result["output"]: + if "text/plain" in result["output"]: + output = result["output"]["text/plain"].strip() + + if expected is not None: + if output == expected: + print(f" PASS (got: {output})") + passed += 1 + else: + print(f" FAIL (expected: {expected}, got: {output})") + failed += 1 + else: + print(" OK") + passed += 1 + + ws.close() + print() + print(f"Results: {passed} passed, {failed} failed") + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main()