Merge pull request 'Changes for a python code interpreter' (#12) from feat/python-code-interpreter into dev

Reviewed-on: wrenn/sandbox#12

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

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

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 {

internal/recipe/context_test.go

@@ -8,6 +8,7 @@ func TestExecContext_WrappedCommand(t *testing.T) {
 		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

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:

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=<duration>] [--timeout=<duration>] [--start-period=<duration>]
+// [--retries=<n>] <command>
+//
+// 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
+}

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)
+				}
+			}
+		})
+	}
+}

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
+}

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)"

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

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()