From 642b8321ddfc977a1b98d677cba4ddd2605f8de3 Mon Sep 17 00:00:00 2001 From: bkmashiro <53376445+bkmashiro@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:08:34 +0100 Subject: [PATCH 1/4] feat(runtime): add Generic WASM and Python Reactor --- cmd/root.go | 5 +- cmd/shimmy-artifact-check/main.go | 74 ++ go.mod | 1 + go.sum | 2 + internal/execution/dispatcher.go | 42 +- internal/execution/supervisor/config.go | 5 +- internal/execution/supervisor/models.go | 3 + internal/execution/wasm/adapter.go | 177 +++ internal/execution/wasm/agent_python.go | 1040 +++++++++++++++++ .../agent_python_lifecycle_config_test.go | 42 + .../execution/wasm/agent_python_observer.go | 120 ++ .../execution/wasm/agent_python_protocol.go | 484 ++++++++ internal/execution/wasm/agent_python_test.go | 750 ++++++++++++ internal/execution/wasm/artifact_check.go | 176 +++ .../execution/wasm/artifact_check_test.go | 52 + internal/execution/wasm/config.go | 193 +++ internal/execution/wasm/dispatcher.go | 378 ++++++ internal/execution/wasm/dispatcher_test.go | 465 ++++++++ internal/execution/wasm/json_util.go | 17 + internal/execution/wasm/pool.go | 42 + .../wasm/python_preload_config_test.go | 28 + .../execution/wasm/python_reactor_artifact.go | 185 +++ internal/execution/wasm/robustness_test.go | 130 +++ internal/execution/wasm/snapshot.go | 108 ++ internal/execution/wasm/snapshot_test.go | 274 +++++ internal/execution/wasm/supervisor.go | 226 ++++ internal/execution/wasm/testdata/echo.wasm | Bin 0 -> 241 bytes internal/execution/wasm/testdata/echo.wat | 66 ++ internal/execution/wasm/testhelpers_test.go | 46 + 29 files changed, 5125 insertions(+), 6 deletions(-) create mode 100644 cmd/shimmy-artifact-check/main.go create mode 100644 internal/execution/wasm/adapter.go create mode 100644 internal/execution/wasm/agent_python.go create mode 100644 internal/execution/wasm/agent_python_lifecycle_config_test.go create mode 100644 internal/execution/wasm/agent_python_observer.go create mode 100644 internal/execution/wasm/agent_python_protocol.go create mode 100644 internal/execution/wasm/agent_python_test.go create mode 100644 internal/execution/wasm/artifact_check.go create mode 100644 internal/execution/wasm/artifact_check_test.go create mode 100644 internal/execution/wasm/config.go create mode 100644 internal/execution/wasm/dispatcher.go create mode 100644 internal/execution/wasm/dispatcher_test.go create mode 100644 internal/execution/wasm/json_util.go create mode 100644 internal/execution/wasm/pool.go create mode 100644 internal/execution/wasm/python_preload_config_test.go create mode 100644 internal/execution/wasm/python_reactor_artifact.go create mode 100644 internal/execution/wasm/robustness_test.go create mode 100644 internal/execution/wasm/snapshot.go create mode 100644 internal/execution/wasm/snapshot_test.go create mode 100644 internal/execution/wasm/supervisor.go create mode 100644 internal/execution/wasm/testdata/echo.wasm create mode 100644 internal/execution/wasm/testdata/echo.wat create mode 100644 internal/execution/wasm/testhelpers_test.go diff --git a/cmd/root.go b/cmd/root.go index eb6019b..690258f 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -47,7 +47,7 @@ functions on arbitrary, serverless platforms.` &cli.StringFlag{ Name: "interface", Aliases: []string{"i"}, - Usage: "the interface to use for worker process communication. Options: rpc, file.", + Usage: "the interface to use for worker communication. Options: rpc, file, wasm.", Value: "rpc", Category: "function", EnvVars: []string{"FUNCTION_INTERFACE"}, @@ -55,10 +55,9 @@ functions on arbitrary, serverless platforms.` &cli.StringFlag{ Name: "command", Aliases: []string{"c"}, - Usage: "the command to invoke to start the worker process.", + Usage: "the command to invoke to start the worker process, or the WASM module path when --interface=wasm.", Category: "function", EnvVars: []string{"FUNCTION_COMMAND"}, - Required: true, }, &cli.StringFlag{ Name: "cwd", diff --git a/cmd/shimmy-artifact-check/main.go b/cmd/shimmy-artifact-check/main.go new file mode 100644 index 0000000..7c449a0 --- /dev/null +++ b/cmd/shimmy-artifact-check/main.go @@ -0,0 +1,74 @@ +// shimmy-artifact-check validates caller-produced WebAssembly artifacts without +// starting Shimmy's production request path. +package main + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "os" + "os/exec" + + "github.com/lambda-feedback/shimmy/internal/execution/wasm" +) + +func main() { + os.Exit(run(os.Args[1:])) +} + +func run(args []string) int { + flags := flag.NewFlagSet("shimmy-artifact-check", flag.ContinueOnError) + flags.SetOutput(os.Stderr) + profile := flags.String("profile", "generic", "runtime ABI: generic or python-reactor") + module := flags.String("module", "", "path to a prebuilt WebAssembly module") + manifest := flags.String("manifest", "", "Python Reactor manifest path") + buildCommand := flags.String("build-command", "", "explicit producer command to run before validation") + buildDir := flags.String("build-dir", ".", "working directory for --build-command") + jsonOutput := flags.Bool("json", false, "emit a JSON report") + if err := flags.Parse(args); err != nil { + return 2 + } + if flags.NArg() != 0 { + fmt.Fprintf(os.Stderr, "unexpected arguments: %v\n", flags.Args()) + return 2 + } + + if *buildCommand != "" { + command := exec.Command("/bin/sh", "-c", *buildCommand) + command.Dir = *buildDir + command.Stdout = os.Stdout + command.Stderr = os.Stderr + if err := command.Run(); err != nil { + fmt.Fprintf(os.Stderr, "artifact build failed: %v\n", err) + return 1 + } + } + + report, err := wasm.CheckArtifact(context.Background(), wasm.ArtifactCheckOptions{ + Profile: *profile, + ModulePath: *module, + ManifestPath: *manifest, + }) + if err != nil { + fmt.Fprintln(os.Stderr, err) + return 1 + } + if *jsonOutput { + encoder := json.NewEncoder(os.Stdout) + encoder.SetIndent("", " ") + if err := encoder.Encode(report); err != nil { + fmt.Fprintf(os.Stderr, "encode report: %v\n", err) + return 1 + } + return 0 + } + + fmt.Printf("OK %s artifact: %s\n", report.Profile, report.Module) + fmt.Printf("exports: %v\n", report.Exports) + fmt.Printf("imports: %v\n", report.Imports) + for _, warning := range report.Warnings { + fmt.Printf("WARNING: %s\n", warning) + } + return 0 +} diff --git a/go.mod b/go.mod index 10caf84..cee6645 100644 --- a/go.mod +++ b/go.mod @@ -37,6 +37,7 @@ require ( github.com/perimeterx/marshmallow v1.1.5 // indirect github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect github.com/shirou/gopsutil v3.21.4-0.20210419000835-c7a38de76ee5+incompatible // indirect + github.com/tetratelabs/wazero v1.9.0 github.com/tklauser/go-sysconf v0.3.12 // indirect github.com/tklauser/numcpus v0.6.1 // indirect github.com/woodsbury/decimal128 v1.3.0 // indirect diff --git a/go.sum b/go.sum index 014f78c..8aeb9f0 100644 --- a/go.sum +++ b/go.sum @@ -110,6 +110,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/supranational/blst v0.3.11 h1:LyU6FolezeWAhvQk0k6O/d49jqgO52MSDDfYgbeoEm4= github.com/supranational/blst v0.3.11/go.mod h1:jZJtfjgudtNl4en1tzwPIV3KjUnQUvG3/j+w+fVonLw= +github.com/tetratelabs/wazero v1.9.0 h1:IcZ56OuxrtaEz8UYNRHBrUa9bYeX9oVY93KspZZBf/I= +github.com/tetratelabs/wazero v1.9.0/go.mod h1:TSbcXCfFP0L2FGkRPxHphadXPjo1T6W+CseNNY7EkjM= github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= diff --git a/internal/execution/dispatcher.go b/internal/execution/dispatcher.go index 300ca3f..95921e8 100644 --- a/internal/execution/dispatcher.go +++ b/internal/execution/dispatcher.go @@ -2,11 +2,16 @@ package execution import ( "context" + "fmt" + "os" + "sort" + "strings" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/dispatcher" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/execution/wasm" ) type Dispatcher dispatcher.Dispatcher @@ -32,7 +37,8 @@ type Params struct { } func NewDispatcher(params Params) (dispatcher.Dispatcher, error) { - if params.Config.Supervisor.IO.Interface == supervisor.RpcIO { + switch params.Config.Supervisor.IO.Interface { + case supervisor.RpcIO: return dispatcher.NewDedicatedDispatcher( dispatcher.DedicatedDispatcherParams{ Config: dispatcher.DedicatedDispatcherConfig{ @@ -42,7 +48,39 @@ func NewDispatcher(params Params) (dispatcher.Dispatcher, error) { Log: params.Log, }, ) - } else { + + case supervisor.WasmIO: + wasmProfile := strings.ToLower(strings.TrimSpace(os.Getenv("FUNCTION_WASM_PROFILE"))) + if wasmProfile == "" { + wasmProfile = "generic" + } + + cfg := wasm.Config{ + ModulePath: params.Config.Supervisor.StartParams.Cmd, + MaxInstances: params.Config.MaxWorkers, + Timeout: params.Config.Supervisor.SendParams.Timeout, + } + switch wasmProfile { + case "generic": + d := wasm.NewDispatcher(cfg, params.Log) + if err := d.Start(params.Context); err != nil { + return nil, err + } + return d, nil + case "python-reactor": + cfg.PythonScriptPath = os.Getenv("FUNCTION_WASM_PYTHON_SCRIPT") + d := wasm.NewAgentPythonDispatcher(cfg, params.Log) + if err := d.Start(params.Context); err != nil { + return nil, err + } + return d, nil + default: + validProfiles := []string{"generic", "python-reactor"} + sort.Strings(validProfiles) + return nil, fmt.Errorf("unsupported FUNCTION_WASM_PROFILE %q; supported values: %s", wasmProfile, strings.Join(validProfiles, ", ")) + } + + default: return dispatcher.NewPooledDispatcher( dispatcher.PooledDispatcherParams{ Config: dispatcher.PooledDispatcherConfig{ diff --git a/internal/execution/supervisor/config.go b/internal/execution/supervisor/config.go index 520e367..758b0ff 100644 --- a/internal/execution/supervisor/config.go +++ b/internal/execution/supervisor/config.go @@ -24,7 +24,7 @@ type SendConfig struct { // IOInterface describes the interface used to communicate with the worker. type IOConfig struct { // Interface describes the communication between the supervisor - // and the worker. It can be either "rpc" or "file". + // and the worker. It can be "rpc", "file", or "wasm". // // If "rpc", the supervisor will communicate with the worker over // a specified transport. The worker is expected to handle incoming @@ -35,6 +35,9 @@ type IOConfig struct { // containing the message payload and response are passed as args // to the worker process. // + // If "wasm", Shimmy loads a pre-built WASI module from FUNCTION_COMMAND + // or FUNCTION_WASM_MODULE and calls its internal alloc/dispatch adapter ABI. + // // Default is "rpc". Interface IOInterface `conf:"interface"` diff --git a/internal/execution/supervisor/models.go b/internal/execution/supervisor/models.go index e7776db..8f98bcb 100644 --- a/internal/execution/supervisor/models.go +++ b/internal/execution/supervisor/models.go @@ -16,6 +16,9 @@ const ( // FileIO describes communication w/ processes over files FileIO IOInterface = "file" + + // WasmIO describes in-process execution of a pre-built WASI module. + WasmIO IOInterface = "wasm" ) // IOTransport describes the transport mechanism used to communicate with diff --git a/internal/execution/wasm/adapter.go b/internal/execution/wasm/adapter.go new file mode 100644 index 0000000..e9622c7 --- /dev/null +++ b/internal/execution/wasm/adapter.go @@ -0,0 +1,177 @@ +// Package wasm implements a WebAssembly execution backend for shimmy using +// wazero. It exposes a [Dispatcher] that manages a pool of pre-compiled WASM +// module instances and dispatches evaluation requests to them. +// +// # Guest ABI +// +// WASM modules loaded by this backend must export two functions: +// +// alloc(size i32) i32 +// Allocate `size` bytes in guest linear memory and return a pointer to +// the start of the allocation. The host will write the JSON-encoded +// request into this region immediately after the call returns. +// +// dispatch(req_ptr i32, req_len i32) i32 +// Process the JSON request at [req_ptr, req_ptr+req_len). Returns a +// pointer P into guest memory where the response is encoded as: +// bytes [P, P+4) — uint32 little-endian response length L +// bytes [P+4, P+4+L) — L bytes of UTF-8 JSON response +// +// The JSON request envelope has the shape: +// +// {"method": "", "params": {…}} +// +// The JSON response is a plain JSON object (map[string]any) that is returned +// verbatim to the caller. +package wasm + +import ( + "context" + "encoding/binary" + "encoding/json" + "fmt" + "time" + + "github.com/tetratelabs/wazero/api" + "go.uber.org/zap" +) + +// requestEnvelope is the JSON structure written into guest memory for each +// evaluation call. +type requestEnvelope struct { + Method string `json:"method"` + Params map[string]any `json:"params"` +} + +// wasmAdapter performs a single opaque dispatch call against a live wazero api.Module. +// It is stateless and safe to call from one goroutine at a time. +type wasmAdapter struct { + mod api.Module + log *zap.Logger + allocFn api.Function // cached exported "alloc" function (M-4 fix) + dispatchFn api.Function // cached exported "dispatch" function +} + +func newWasmAdapter(mod api.Module, log *zap.Logger) *wasmAdapter { + return &wasmAdapter{ + mod: mod, + log: log.Named("adapter_wasm"), + allocFn: mod.ExportedFunction("alloc"), + dispatchFn: mod.ExportedFunction("dispatch"), + } +} + +// send marshals (method, data) into JSON, writes it into the guest's linear +// memory via alloc, calls dispatch, and reads back the length-prefixed +// response. +func (a *wasmAdapter) send( + ctx context.Context, + method string, + data map[string]any, + timeout time.Duration, +) (map[string]any, error) { + if timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + + // 1. Marshal request envelope. + envelope := requestEnvelope{Method: method, Params: data} + + reqBytes, err := json.Marshal(envelope) + if err != nil { + return nil, fmt.Errorf("wasm: marshal request: %w", err) + } + + reqLen := uint64(len(reqBytes)) + + // 2. Allocate guest memory for the request (cached lookup — M-4 fix). + if a.allocFn == nil { + return nil, fmt.Errorf("wasm: guest module does not export 'alloc'") + } + + allocRes, err := a.allocFn.Call(ctx, reqLen) + if err != nil { + return nil, fmt.Errorf("wasm: alloc(%d): %w", reqLen, err) + } + if len(allocRes) != 1 { + return nil, fmt.Errorf("wasm: alloc returned %d values, expected 1", len(allocRes)) + } + + reqPtr := allocRes[0] + if reqPtr == 0 { + return nil, fmt.Errorf("wasm: alloc returned NULL (out of memory)") + } + + // 3. Write request bytes into guest memory. + mem := a.mod.Memory() + if mem == nil { + return nil, fmt.Errorf("wasm: guest module has no linear memory") + } + + if !mem.Write(uint32(reqPtr), reqBytes) { + return nil, fmt.Errorf( + "wasm: failed to write %d bytes at ptr=%d (memory size=%d)", + len(reqBytes), reqPtr, mem.Size(), + ) + } + + // 4. Call the language- and method-agnostic dispatch ABI. + if a.dispatchFn == nil { + return nil, fmt.Errorf("wasm: guest module does not export 'dispatch'") + } + + a.log.Debug("calling dispatch", + zap.String("method", method), + zap.Uint64("req_ptr", reqPtr), + zap.Uint64("req_len", reqLen), + ) + + dispatchRes, err := a.dispatchFn.Call(ctx, reqPtr, reqLen) + if err != nil { + return nil, fmt.Errorf("wasm: dispatch: %w", err) + } + if len(dispatchRes) != 1 { + return nil, fmt.Errorf("wasm: dispatch returned %d values, expected 1", len(dispatchRes)) + } + + resPtr := uint32(dispatchRes[0]) + + // 5. Read the 4-byte little-endian length prefix. + lenBytes, ok := mem.Read(resPtr, 4) + if !ok { + return nil, fmt.Errorf("wasm: failed to read response length at ptr=%d", resPtr) + } + + resLen := binary.LittleEndian.Uint32(lenBytes) + + // 6. Read the response JSON body. + // Validate bounds before reading to catch corrupt/malicious response pointers. + if uint64(resPtr)+4+uint64(resLen) > uint64(mem.Size()) { + return nil, fmt.Errorf( + "wasm: response out of bounds: resPtr=%d resLen=%d memSize=%d", + resPtr, resLen, mem.Size(), + ) + } + resBytes, ok := mem.Read(resPtr+4, resLen) + if !ok { + return nil, fmt.Errorf( + "wasm: failed to read %d response bytes at ptr=%d", + resLen, resPtr+4, + ) + } + + a.log.Debug("received response", + zap.Uint32("res_ptr", resPtr), + zap.Uint32("res_len", resLen), + ) + + // 7. Unmarshal response. + var result map[string]any + if err := json.Unmarshal(resBytes, &result); err != nil { + return nil, fmt.Errorf("wasm: unmarshal response: %w", err) + } + + return result, nil +} diff --git a/internal/execution/wasm/agent_python.go b/internal/execution/wasm/agent_python.go new file mode 100644 index 0000000..3ea6acf --- /dev/null +++ b/internal/execution/wasm/agent_python.go @@ -0,0 +1,1040 @@ +package wasm + +import ( + "context" + cryptorand "crypto/rand" + "encoding/binary" + "errors" + "fmt" + "math" + "os" + "runtime" + "sync" + "sync/atomic" + "time" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" + "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1" + "go.uber.org/zap" +) + +const ( + agentPythonDefaultMemoryPages = 8192 + agentPythonMaxMemoryPages = 16384 + agentPythonDiagnosticMax = 16 * 1024 +) + +// AgentPythonDispatcher consumes the clean Agent Python Runtime v1 artifact. +// The artifact is compiled once. Module ownership is selected explicitly by +// PythonLifecycle: fresh, never-served single-use candidates, or prepared +// linear-memory snapshot restore. +type AgentPythonDispatcher struct { + cfg Config + log *zap.Logger + + mu sync.Mutex + started bool + closed bool + closedCh chan struct{} + pending sync.WaitGroup + + runtime wazero.Runtime + compiled wazero.CompiledModule + cache wazero.CompilationCache + artifact *AgentPythonArtifact + script string + slots chan struct{} + prepared chan *agentPythonModuleSlot + snapshotSelected string + + refillCtx context.Context + refillCancel context.CancelFunc + refillMu sync.Mutex + refillInFlight int + refills sync.WaitGroup + preparedHits atomic.Uint64 + preparedMisses atomic.Uint64 + preparedRefills atomic.Uint64 + + runCounter atomic.Uint64 + slotCounter atomic.Uint64 +} + +type agentPythonModuleSlot struct { + id uint64 + module api.Module + diagnostic *agentPythonDiagnosticBuffer + strategy SnapshotStrategy + baselineSize uint32 + snapshotSelected string +} + +func (slot *agentPythonModuleSlot) close(ctx context.Context) error { + if slot == nil { + return nil + } + var moduleErr, strategyErr error + if slot.module != nil { + moduleErr = slot.module.Close(ctx) + slot.module = nil + } + if slot.strategy != nil { + strategyErr = slot.strategy.Close() + slot.strategy = nil + } + return errors.Join(moduleErr, strategyErr) +} + +func NewAgentPythonDispatcher(cfg Config, log *zap.Logger) *AgentPythonDispatcher { + if log == nil { + log = zap.NewNop() + } + return &AgentPythonDispatcher{ + cfg: cfg, + log: log.Named("dispatcher_agent_python"), + closedCh: make(chan struct{}), + } +} + +func (d *AgentPythonDispatcher) Start(ctx context.Context) error { + d.mu.Lock() + startupObserver := d.cfg.AgentPythonObserver + var startupEvents []AgentPythonPhaseEvent + if startupObserver != nil { + // Start serializes dispatcher state under d.mu, but external observers must + // never run in that lock domain: they may synchronously inspect or shut down + // the dispatcher. Capture already-timed immutable events and flush them in + // order after releasing the lock. + d.cfg.AgentPythonObserver = func(event AgentPythonPhaseEvent) { + startupEvents = append(startupEvents, event) + } + } + defer func() { + d.cfg.AgentPythonObserver = startupObserver + d.mu.Unlock() + for _, event := range startupEvents { + d.emitAgentPythonPhaseEvent(startupObserver, event) + } + }() + if d.closed { + return errors.New("python-reactor: dispatcher is shut down") + } + if d.started { + return nil + } + + d.cfg.applyEnv() + if d.cfg.Timeout == 0 { + d.cfg.Timeout = 30 * time.Second + } + if d.cfg.MaxMemoryPages == 0 { + d.cfg.MaxMemoryPages = agentPythonDefaultMemoryPages + } + if d.cfg.MaxMemoryPages > agentPythonMaxMemoryPages { + return fmt.Errorf("python-reactor: memory limit %d pages exceeds hard bound %d", d.cfg.MaxMemoryPages, agentPythonMaxMemoryPages) + } + if d.cfg.MaxInstances <= 0 { + d.cfg.MaxInstances = runtime.NumCPU() + if d.cfg.MaxInstances > 4 { + d.cfg.MaxInstances = 4 + } + if d.cfg.MaxInstances < 1 { + d.cfg.MaxInstances = 1 + } + } + if d.cfg.PythonPreloadMode == "" { + d.cfg.PythonPreloadMode = "evaluator" + } + d.cfg.applyAgentPythonDefaults() + if err := d.cfg.validatePythonPreloadMode(); err != nil { + return fmt.Errorf("python-reactor: %w", err) + } + if err := d.cfg.validateAgentPythonLifecycle(); err != nil { + return fmt.Errorf("python-reactor: %w", err) + } + if len(d.cfg.AllowedPaths) != 0 { + return errors.New("agent-python does not expose Host filesystem paths; unset FUNCTION_WASM_ALLOWED_PATHS") + } + if d.cfg.PythonScriptPath == "" { + return errors.New("python-reactor: PythonScriptPath must be set (FUNCTION_WASM_PYTHON_SCRIPT)") + } + scriptBytes, err := os.ReadFile(d.cfg.PythonScriptPath) + if err != nil { + return fmt.Errorf("python-reactor: read script %q: %w", d.cfg.PythonScriptPath, err) + } + if len(scriptBytes) == 0 || len(scriptBytes) > agentPythonPayloadMax { + return fmt.Errorf("python-reactor: trusted script size %d is outside the 1 MiB guest bound", len(scriptBytes)) + } + + phaseStart := time.Now() + artifact, err := verifyAgentPythonArtifact(d.cfg.ModulePath, d.cfg.AgentPythonManifestPath) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseArtifactVerify, Purpose: AgentPythonPurposeStartup, + Started: phaseStart, Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + return err + } + if artifact.ABI == "shimmy-python-runtime/v1" && d.cfg.PythonPreloadMode == "off" { + return errors.New("python-reactor: Shimmy producer ABI requires prepared evaluator preload") + } + + runtimeConfig := wazero.NewRuntimeConfig(). + WithCloseOnContextDone(true). + WithMemoryLimitPages(d.cfg.MaxMemoryPages) + var cache wazero.CompilationCache + if d.cfg.CompileCacheDir != "" { + cache, err = wazero.NewCompilationCacheWithDir(d.cfg.CompileCacheDir) + if err != nil { + return fmt.Errorf("python-reactor: create compilation cache: %w", err) + } + runtimeConfig = runtimeConfig.WithCompilationCache(cache) + } + + phaseStart = time.Now() + wasmRuntime := wazero.NewRuntimeWithConfig(ctx, runtimeConfig) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseRuntimeCreate, Purpose: AgentPythonPurposeStartup, + Started: phaseStart, Outcome: AgentPythonOutcomeOK, + }) + closePartial := func() { + _ = wasmRuntime.Close(context.Background()) + if cache != nil { + _ = cache.Close(context.Background()) + } + } + phaseStart = time.Now() + _, err = wasi_snapshot_preview1.Instantiate(ctx, wasmRuntime) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseWASIImports, Purpose: AgentPythonPurposeStartup, + Started: phaseStart, Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + closePartial() + return fmt.Errorf("python-reactor: instantiate WASI imports: %w", err) + } + phaseStart = time.Now() + _, err = wasmRuntime.NewHostModuleBuilder("agent_runtime_v1"). + NewFunctionBuilder(). + WithFunc(agentPythonDeniedHostCall). + Export("host_call"). + Instantiate(ctx) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseHostImports, Purpose: AgentPythonPurposeStartup, + Started: phaseStart, Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + closePartial() + return fmt.Errorf("python-reactor: instantiate Host imports: %w", err) + } + phaseStart = time.Now() + compiled, err := wasmRuntime.CompileModule(ctx, artifact.WasmBytes) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseCompile, Purpose: AgentPythonPurposeStartup, + Started: phaseStart, Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + closePartial() + return fmt.Errorf("python-reactor: compile guest: %w", err) + } + if err := verifyCompiledPythonReactorArtifact(compiled, artifact); err != nil { + closePartial() + return err + } + + d.runtime = wasmRuntime + d.compiled = compiled + d.cache = cache + d.artifact = artifact + d.script = string(scriptBytes) + d.slots = make(chan struct{}, d.cfg.MaxInstances) + d.refillCtx, d.refillCancel = context.WithCancel(context.Background()) + + switch d.cfg.PythonLifecycle { + case "snapshot": + d.prepared = make(chan *agentPythonModuleSlot, d.cfg.MaxInstances) + for i := 0; i < d.cfg.MaxInstances; i++ { + slot, err := d.newPreparedModuleSlot(ctx, true, AgentPythonPurposeStartup, 0) + if err != nil { + _ = d.closeRuntime(context.Background()) + return err + } + d.snapshotSelected = "memcpy" + d.prepared <- slot + } + case "single-use": + d.prepared = make(chan *agentPythonModuleSlot, d.cfg.PythonPreparedCapacity) + for i := 0; i < d.cfg.PythonPreparedCapacity; i++ { + slot, err := d.newPreparedModuleSlot(ctx, false, AgentPythonPurposeStartup, 0) + if err != nil { + _ = d.closeRuntime(context.Background()) + return err + } + d.prepared <- slot + } + case "fresh": + // Probe the exact artifact and trusted script before reporting readiness. + slot, err := d.newPreparedModuleSlot(ctx, false, AgentPythonPurposeStartup, 0) + if err != nil { + _ = d.closeRuntime(context.Background()) + return err + } + _ = slot.close(context.Background()) + } + + d.started = true + d.log.Info("agent-python dispatcher ready", + zap.String("artifact_sha256", artifact.SHA256), + zap.String("producer_commit", artifact.ProducerCommit), + zap.String("artifact_profile", artifact.Profile), + zap.Int("max_instances", d.cfg.MaxInstances), + zap.Duration("request_timeout", d.cfg.Timeout), + zap.String("lifecycle", d.cfg.PythonLifecycle), + zap.String("snapshot_mode", d.snapshotMode()), + zap.String("reset_mode", d.resetMode()), + ) + return nil +} + +func (d *AgentPythonDispatcher) Send(ctx context.Context, method string, params map[string]any) (map[string]any, error) { + if method == "healthcheck" { + d.mu.Lock() + ready := d.started && !d.closed + profile := "" + if d.artifact != nil { + profile = d.artifact.Profile + } + preparedReady := len(d.prepared) + d.mu.Unlock() + if !ready { + return nil, errors.New("python-reactor: dispatcher is not ready") + } + return map[string]any{ + "command": "healthcheck", + "result": map[string]any{ + "status": "ok", + "profile": profile, + "lifecycle": d.cfg.PythonLifecycle, + "snapshot_mode": d.snapshotMode(), + "snapshot_selected": d.snapshotSelected, + "reset_mode": d.resetMode(), + "prepared_ready": preparedReady, + "prepared_hits": d.preparedHits.Load(), + "prepared_misses": d.preparedMisses.Load(), + "prepared_refills": d.preparedRefills.Load(), + }, + }, nil + } + if !d.tryBeginSend() { + return nil, errors.New("python-reactor: dispatcher is not ready") + } + defer d.pending.Done() + + select { + case d.slots <- struct{}{}: + defer func() { <-d.slots }() + case <-d.closedCh: + return nil, errors.New("python-reactor: dispatcher is shut down") + case <-ctx.Done(): + return nil, fmt.Errorf("python-reactor: acquire execution slot: %w", ctx.Err()) + } + + requestID := d.runCounter.Add(1) + var request []byte + var err error + if d.artifact.ABI == "shimmy-python-runtime/v1" { + request, err = buildShimmyPythonRunRequest(method, params) + } else { + runID := fmt.Sprintf("shimmy-%s-%d", d.artifact.SHA256[:12], requestID) + scriptInRequest := "" + if d.cfg.PythonPreloadMode == "off" { + scriptInRequest = d.script + } + request, err = buildAgentPythonRunRequest(runID, method, params, scriptInRequest) + } + if err != nil { + return nil, err + } + + runContext, cancel := context.WithTimeout(ctx, d.cfg.Timeout) + defer cancel() + + var slot *agentPythonModuleSlot + checkoutStart := time.Now() + switch d.cfg.PythonLifecycle { + case "snapshot": + slot, err = acquireAgentPythonSnapshotSlot( + runContext, + d.prepared, + d.closedCh, + d.snapshotRefillInFlight, + func(createContext context.Context) (*agentPythonModuleSlot, error) { + return d.newPreparedModuleSlot(createContext, true, AgentPythonPurposeReplacement, requestID) + }, + ) + if err != nil { + return nil, err + } + case "single-use": + select { + case slot = <-d.prepared: + d.preparedHits.Add(1) + default: + d.preparedMisses.Add(1) + } + d.scheduleSingleUseRefill(requestID) + if slot == nil { + slot, err = d.newPreparedModuleSlot(runContext, false, AgentPythonPurposeFresh, requestID) + if err != nil { + return nil, err + } + } + case "fresh": + slot, err = d.newPreparedModuleSlot(runContext, false, AgentPythonPurposeFresh, requestID) + if err != nil { + return nil, err + } + } + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseCheckout, Purpose: AgentPythonPurposeRequest, + RequestID: requestID, SlotID: slot.id, Started: checkoutStart, + MemoryBytes: uint64(slot.module.Memory().Size()), SnapshotSelected: slot.snapshotSelected, + Outcome: AgentPythonOutcomeOK, + }) + if d.cfg.PythonLifecycle != "snapshot" { + defer func() { + phaseStart := time.Now() + closeErr := slot.close(context.Background()) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseClose, Purpose: AgentPythonPurposeRequest, + RequestID: requestID, SlotID: slot.id, Started: phaseStart, + Outcome: agentPythonPhaseOutcome(closeErr), Err: closeErr, + }) + }() + } + + phaseStart := time.Now() + payload, callErr := callAgentPythonExecute(runContext, slot.module, d.artifact.ExecuteExport, request) + if callErr != nil && runContext.Err() != nil { + callErr = errors.Join(callErr, runContext.Err()) + } + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseExecute, Purpose: AgentPythonPurposeRequest, + RequestID: requestID, SlotID: slot.id, Started: phaseStart, + MemoryBytes: uint64(slot.module.Memory().Size()), SnapshotSelected: slot.snapshotSelected, + Outcome: agentPythonPhaseOutcome(callErr), Err: callErr, + }) + + if d.cfg.PythonLifecycle == "snapshot" { + var restoreErr error + if callErr == nil { + phaseStart = time.Now() + restoreErr = restoreAgentPythonSnapshot(slot) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseRestore, Purpose: AgentPythonPurposeRequest, + RequestID: requestID, SlotID: slot.id, Started: phaseStart, + MemoryBytes: uint64(slot.module.Memory().Size()), SnapshotSelected: slot.snapshotSelected, + Outcome: agentPythonPhaseOutcome(restoreErr), Err: restoreErr, + }) + } + if callErr != nil || restoreErr != nil { + diagnostic := slot.diagnostic.String() + d.discardSnapshotSlotAsync(slot, requestID) + d.scheduleSnapshotRefill(requestID) + return nil, withAgentPythonDiagnostic(errors.Join(callErr, restoreErr), diagnostic) + } + slot.diagnostic.Reset() + d.prepared <- slot + } + if callErr != nil { + return nil, withAgentPythonDiagnostic(callErr, slot.diagnostic.String()) + } + phaseStart = time.Now() + var result map[string]any + if d.artifact.ABI == "shimmy-python-runtime/v1" { + result, err = decodeShimmyPythonResponse(payload) + } else { + result, err = decodeAgentPythonResponse(payload) + } + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseDecode, Purpose: AgentPythonPurposeRequest, + RequestID: requestID, SlotID: slot.id, Started: phaseStart, + SnapshotSelected: slot.snapshotSelected, + Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + return nil, err + } + return map[string]any{"command": method, "result": result}, nil +} + +func acquireAgentPythonSnapshotSlot( + ctx context.Context, + prepared <-chan *agentPythonModuleSlot, + closed <-chan struct{}, + refillInFlight func() bool, + create func(context.Context) (*agentPythonModuleSlot, error), +) (*agentPythonModuleSlot, error) { + select { + case slot := <-prepared: + if slot != nil { + return slot, nil + } + case <-closed: + return nil, errors.New("python-reactor: dispatcher is shut down") + case <-ctx.Done(): + return nil, fmt.Errorf("python-reactor: acquire prepared module: %w", ctx.Err()) + default: + } + if refillInFlight != nil && refillInFlight() { + select { + case slot := <-prepared: + if slot != nil { + return slot, nil + } + case <-closed: + return nil, errors.New("python-reactor: dispatcher is shut down") + case <-ctx.Done(): + return nil, fmt.Errorf("python-reactor: wait for snapshot refill: %w", ctx.Err()) + } + } + + slot, err := create(ctx) + if err != nil { + return nil, fmt.Errorf("python-reactor: replenish missing prepared snapshot slot: %w", err) + } + if slot == nil { + return nil, errors.New("python-reactor: replenish missing prepared snapshot slot returned nil") + } + return slot, nil +} + +func (d *AgentPythonDispatcher) resetMode() string { + switch d.cfg.PythonLifecycle { + case "snapshot": + return "linear-memory-" + d.snapshotSelected + case "single-use": + return "single-use-prepared" + default: + return "fresh-instance" + } +} + +func (d *AgentPythonDispatcher) snapshotMode() string { + if d.cfg.PythonLifecycle == "snapshot" { + return "memcpy" + } + return "" +} + +func (d *AgentPythonDispatcher) tryBeginSend() bool { + d.mu.Lock() + defer d.mu.Unlock() + if !d.started || d.closed { + return false + } + d.pending.Add(1) + return true +} + +func (d *AgentPythonDispatcher) newInitializedModule( + ctx context.Context, + prepare bool, + purpose AgentPythonPurpose, + requestID uint64, + slotID uint64, +) (api.Module, *agentPythonDiagnosticBuffer, error) { + diagnostic := &agentPythonDiagnosticBuffer{} + phaseStart := time.Now() + module, err := d.runtime.InstantiateModule( + ctx, + d.compiled, + wazero.NewModuleConfig().WithName("").WithRandSource(cryptorand.Reader).WithStderr(diagnostic), + ) + memoryBytes := uint64(0) + if module != nil && module.Memory() != nil { + memoryBytes = uint64(module.Memory().Size()) + } + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseInstantiate, Purpose: purpose, RequestID: requestID, SlotID: slotID, + Started: phaseStart, MemoryBytes: memoryBytes, Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + return nil, diagnostic, fmt.Errorf("python-reactor: instantiate guest: %w", err) + } + failed := true + defer func() { + if failed { + _ = module.Close(context.Background()) + } + }() + phaseStart = time.Now() + err = callAgentPythonNoArgs(ctx, module, "_initialize") + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseInitialize, Purpose: purpose, RequestID: requestID, SlotID: slotID, + Started: phaseStart, MemoryBytes: uint64(module.Memory().Size()), Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + return nil, diagnostic, err + } + phaseStart = time.Now() + if d.artifact.ABI == "shimmy-python-runtime/v1" { + err = callAgentPythonNoArgsValue(ctx, module, "shimmy_python_runtime_identity", 0x53505231) + if err == nil { + err = callAgentPythonNoArgsValue(ctx, module, d.artifact.InitExport, 0) + } + } else { + err = callAgentPythonStatus(ctx, module, d.artifact.InitExport, []byte("{}")) + } + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseRuntimeInit, Purpose: purpose, RequestID: requestID, SlotID: slotID, + Started: phaseStart, MemoryBytes: uint64(module.Memory().Size()), Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + return nil, diagnostic, err + } + if prepare { + phaseStart = time.Now() + err = callAgentPythonStatus(ctx, module, d.artifact.PrepareExport, []byte(d.script)) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseRuntimePrepare, Purpose: purpose, RequestID: requestID, SlotID: slotID, + Started: phaseStart, MemoryBytes: uint64(module.Memory().Size()), Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + return nil, diagnostic, err + } + } + diagnostic.Reset() + failed = false + return module, diagnostic, nil +} + +func reserveAgentPythonSnapshotHeadroom(ctx context.Context, module api.Module, bytes uint64) (retErr error) { + if bytes == 0 { + return nil + } + if bytes > math.MaxUint32 { + return fmt.Errorf("python-reactor: snapshot headroom %d exceeds wasm32 allocation limit", bytes) + } + allocate := module.ExportedFunction("alloc") + deallocate := module.ExportedFunction("dealloc") + if allocate == nil || deallocate == nil { + return errors.New("python-reactor: snapshot headroom requires alloc and dealloc exports") + } + + const chunkBytes = uint64(1024 * 1024) + pointers := make([]uint64, 0, (bytes+chunkBytes-1)/chunkBytes) + defer func() { + for i := len(pointers) - 1; i >= 0; i-- { + if _, err := deallocate.Call(context.Background(), pointers[i]); err != nil { + retErr = errors.Join(retErr, fmt.Errorf("python-reactor: release snapshot headroom: %w", err)) + } + } + }() + + for remaining := bytes; remaining > 0; { + chunk := chunkBytes + if remaining < chunk { + chunk = remaining + } + result, err := allocate.Call(ctx, chunk) + if err != nil { + return fmt.Errorf("python-reactor: reserve %d snapshot headroom bytes: %w", bytes, err) + } + if len(result) != 1 || result[0] == 0 { + return fmt.Errorf("python-reactor: reserve %d snapshot headroom bytes: guest allocator returned no pointer", bytes) + } + pointers = append(pointers, result[0]) + remaining -= chunk + } + return nil +} + +func (d *AgentPythonDispatcher) newPreparedModuleSlot( + ctx context.Context, + takeSnapshot bool, + purpose AgentPythonPurpose, + requestID uint64, +) (*agentPythonModuleSlot, error) { + slotID := d.slotCounter.Add(1) + module, diagnostic, err := d.newInitializedModule( + ctx, + d.cfg.PythonPreloadMode != "off", + purpose, + requestID, + slotID, + ) + if err != nil { + return nil, withAgentPythonDiagnostic(err, diagnostic.String()) + } + slot := &agentPythonModuleSlot{ + id: slotID, + module: module, + diagnostic: diagnostic, + } + if !takeSnapshot { + return slot, nil + } + phaseStart := time.Now() + err = reserveAgentPythonSnapshotHeadroom(ctx, module, d.cfg.PythonSnapshotHeadroomBytes) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseHeadroom, Purpose: purpose, RequestID: requestID, SlotID: slotID, + Started: phaseStart, MemoryBytes: uint64(module.Memory().Size()), Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + _ = slot.close(context.Background()) + return nil, err + } + phaseStart = time.Now() + slot.strategy = NewFullMemcpyStrategy() + slot.snapshotSelected = "memcpy" + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseStrategySelect, Purpose: purpose, RequestID: requestID, SlotID: slotID, + Started: phaseStart, MemoryBytes: uint64(module.Memory().Size()), SnapshotSelected: slot.snapshotSelected, + Outcome: AgentPythonOutcomeOK, + }) + phaseStart = time.Now() + err = slot.strategy.Take(module.Memory()) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseSnapshotTake, Purpose: purpose, RequestID: requestID, SlotID: slotID, + Started: phaseStart, MemoryBytes: uint64(module.Memory().Size()), SnapshotSelected: slot.snapshotSelected, + Outcome: agentPythonPhaseOutcome(err), Err: err, + }) + if err != nil { + _ = slot.close(context.Background()) + return nil, fmt.Errorf("python-reactor: take prepared snapshot: %w", err) + } + slot.baselineSize = module.Memory().Size() + return slot, nil +} + +func restoreAgentPythonSnapshot(slot *agentPythonModuleSlot) error { + if slot == nil || slot.module == nil || slot.strategy == nil { + return errors.New("python-reactor: prepared snapshot slot is incomplete") + } + memory := slot.module.Memory() + if memory == nil { + return errors.New("python-reactor: prepared snapshot slot has no memory") + } + if memory.Size() != slot.baselineSize { + return fmt.Errorf("python-reactor: memory size drift: got %d bytes, baseline %d", memory.Size(), slot.baselineSize) + } + return slot.strategy.Restore(memory) +} + +func (d *AgentPythonDispatcher) discardSnapshotSlotAsync(slot *agentPythonModuleSlot, requestID uint64) { + // Send already holds one pending count, so this Add cannot race Shutdown's + // Wait. Closing a context-cancelled wazero module or its snapshot strategy + // can block and must not extend the request deadline. + d.pending.Add(1) + go func() { + defer d.pending.Done() + phaseStart := time.Now() + closeErr := slot.close(context.Background()) + d.observeAgentPythonPhase(AgentPythonPhaseObservation{ + Phase: AgentPythonPhaseClose, Purpose: AgentPythonPurposeRequest, + RequestID: requestID, SlotID: slot.id, Started: phaseStart, + Outcome: agentPythonPhaseOutcome(closeErr), Err: closeErr, + }) + }() +} + +func (d *AgentPythonDispatcher) snapshotRefillInFlight() bool { + d.refillMu.Lock() + defer d.refillMu.Unlock() + return d.refillInFlight > 0 +} + +func (d *AgentPythonDispatcher) scheduleSnapshotRefill(requestID uint64) { + if d.refillCtx == nil || d.prepared == nil { + return + } + d.refillMu.Lock() + if len(d.prepared)+d.refillInFlight >= cap(d.prepared) { + d.refillMu.Unlock() + return + } + d.refillInFlight++ + d.refills.Add(1) + refillCtx := d.refillCtx + d.refillMu.Unlock() + + go func() { + defer d.refills.Done() + defer func() { + d.refillMu.Lock() + d.refillInFlight-- + d.refillMu.Unlock() + }() + + timeout := d.cfg.Timeout + if timeout < 30*time.Second { + timeout = 30 * time.Second + } + ctx, cancel := context.WithTimeout(refillCtx, timeout) + defer cancel() + slot, err := d.newPreparedModuleSlot(ctx, true, AgentPythonPurposeReplacement, requestID) + if err != nil { + if refillCtx.Err() == nil { + d.log.Warn("agent-python snapshot refill failed", zap.Error(err)) + } + return + } + select { + case d.prepared <- slot: + d.preparedRefills.Add(1) + case <-refillCtx.Done(): + _ = slot.close(context.Background()) + } + }() +} + +func (d *AgentPythonDispatcher) scheduleSingleUseRefill(requestID uint64) { + if d.refillCtx == nil || d.prepared == nil { + return + } + d.refillMu.Lock() + if len(d.prepared)+d.refillInFlight >= cap(d.prepared) { + d.refillMu.Unlock() + return + } + d.refillInFlight++ + d.refills.Add(1) + refillCtx := d.refillCtx + d.refillMu.Unlock() + + go func() { + defer d.refills.Done() + defer func() { + d.refillMu.Lock() + d.refillInFlight-- + d.refillMu.Unlock() + }() + + timeout := d.cfg.Timeout + if timeout < 30*time.Second { + timeout = 30 * time.Second + } + ctx, cancel := context.WithTimeout(refillCtx, timeout) + defer cancel() + slot, err := d.newPreparedModuleSlot(ctx, false, AgentPythonPurposeRefill, requestID) + if err != nil { + if refillCtx.Err() == nil { + d.log.Warn("agent-python single-use refill failed", zap.Error(err)) + } + return + } + select { + case d.prepared <- slot: + d.preparedRefills.Add(1) + case <-refillCtx.Done(): + _ = slot.close(context.Background()) + } + }() +} + +func (d *AgentPythonDispatcher) Shutdown(ctx context.Context) error { + d.mu.Lock() + if d.closed { + d.mu.Unlock() + return nil + } + d.closed = true + close(d.closedCh) + if d.refillCancel != nil { + d.refillCancel() + } + d.mu.Unlock() + + d.pending.Wait() + d.refills.Wait() + d.mu.Lock() + defer d.mu.Unlock() + return d.closeRuntime(ctx) +} + +func (d *AgentPythonDispatcher) closeRuntime(ctx context.Context) error { + if d.refillCancel != nil { + d.refillCancel() + d.refillCancel = nil + } + d.refillCtx = nil + var slotErr error + if d.prepared != nil { + for { + select { + case slot := <-d.prepared: + slotErr = errors.Join(slotErr, slot.close(ctx)) + default: + d.prepared = nil + goto preparedClosed + } + } + } + +preparedClosed: + var compiledErr, runtimeErr, cacheErr error + if d.compiled != nil { + compiledErr = d.compiled.Close(ctx) + d.compiled = nil + } + if d.runtime != nil { + runtimeErr = d.runtime.Close(ctx) + d.runtime = nil + } + if d.cache != nil { + cacheErr = d.cache.Close(ctx) + d.cache = nil + } + d.started = false + return errors.Join(slotErr, compiledErr, runtimeErr, cacheErr) +} + +func agentPythonDeniedHostCall(context.Context, api.Module, uint32, uint32, uint32, uint32) int32 { + return -1 +} + +func callAgentPythonNoArgs(ctx context.Context, module api.Module, name string) error { + function := module.ExportedFunction(name) + if function == nil { + return fmt.Errorf("python-reactor: required export %q is missing", name) + } + if _, err := function.Call(ctx); err != nil { + return fmt.Errorf("python-reactor: call %s: %w", name, err) + } + return nil +} + +func callAgentPythonNoArgsValue(ctx context.Context, module api.Module, name string, expected uint32) error { + function := module.ExportedFunction(name) + if function == nil { + return fmt.Errorf("python-reactor: required export %q is missing", name) + } + results, err := function.Call(ctx) + if err != nil { + return fmt.Errorf("python-reactor: call %s: %w", name, err) + } + if len(results) != 1 || uint32(results[0]) != expected { + return fmt.Errorf("python-reactor: %s returned identity/status %v; want %d", name, results, expected) + } + return nil +} + +func callAgentPythonStatus(ctx context.Context, module api.Module, name string, data []byte) error { + results, release, err := callAgentPythonWithBytes(ctx, module, name, data) + if release != nil { + defer release() + } + if err != nil { + return err + } + if len(results) != 1 || uint32(results[0]) != 0 { + return fmt.Errorf("python-reactor: %s returned non-zero status", name) + } + return nil +} + +func callAgentPythonExecute(ctx context.Context, module api.Module, name string, request []byte) ([]byte, error) { + if name == "" { + name = "execute" + } + results, release, err := callAgentPythonWithBytes(ctx, module, name, request) + if release != nil { + defer release() + } + if err != nil { + return nil, err + } + if len(results) != 1 { + return nil, errors.New("python-reactor: execute returned an unexpected result count") + } + return readAgentPythonResponse(module.Memory(), uint32(results[0])) +} + +func callAgentPythonWithBytes(ctx context.Context, module api.Module, name string, data []byte) ([]uint64, func(), error) { + if len(data) == 0 || len(data) > agentPythonPayloadMax || len(data) > math.MaxUint32 { + return nil, nil, fmt.Errorf("python-reactor: %s input size %d is outside the guest bound", name, len(data)) + } + allocate := module.ExportedFunction("alloc") + deallocate := module.ExportedFunction("dealloc") + function := module.ExportedFunction(name) + if allocate == nil || deallocate == nil || function == nil { + return nil, nil, fmt.Errorf("python-reactor: required allocation or %s export is missing", name) + } + allocated, err := allocate.Call(ctx, uint64(uint32(len(data)))) + if err != nil || len(allocated) != 1 || allocated[0] == 0 { + return nil, nil, fmt.Errorf("python-reactor: guest allocation failed: %w", err) + } + pointer := uint32(allocated[0]) + var once sync.Once + release := func() { + once.Do(func() { + releaseContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, _ = deallocate.Call(releaseContext, uint64(pointer)) + }) + } + if !module.Memory().Write(pointer, data) { + release() + return nil, nil, errors.New("python-reactor: guest input write is out of bounds") + } + results, err := function.Call(ctx, uint64(pointer), uint64(uint32(len(data)))) + if err != nil { + // A failed guest call is followed by module disposal in every caller. + // Calling guest dealloc here can itself consume another deadline and + // delay the timeout response without reclaiming reusable memory. + return nil, nil, fmt.Errorf("python-reactor: call %s: %w", name, err) + } + return results, release, nil +} + +func readAgentPythonResponse(memory api.Memory, pointer uint32) ([]byte, error) { + if memory == nil { + return nil, errors.New("python-reactor: guest module has no linear memory") + } + header, ok := memory.Read(pointer, 4) + if !ok { + return nil, errors.New("python-reactor: response length prefix is out of bounds") + } + length := binary.LittleEndian.Uint32(header) + if length > agentPythonPayloadMax { + return nil, fmt.Errorf("python-reactor: response payload length %d exceeds limit %d", length, agentPythonPayloadMax) + } + if uint64(pointer)+4+uint64(length) > uint64(memory.Size()) { + return nil, errors.New("python-reactor: response frame is out of bounds") + } + payload, ok := memory.Read(pointer+4, length) + if !ok { + return nil, errors.New("python-reactor: response payload is out of bounds") + } + return append([]byte(nil), payload...), nil +} + +type agentPythonDiagnosticBuffer struct { + data []byte +} + +func (buffer *agentPythonDiagnosticBuffer) Write(data []byte) (int, error) { + length := len(data) + if length >= agentPythonDiagnosticMax { + buffer.data = append(buffer.data[:0], data[length-agentPythonDiagnosticMax:]...) + return length, nil + } + if overflow := len(buffer.data) + length - agentPythonDiagnosticMax; overflow > 0 { + copy(buffer.data, buffer.data[overflow:]) + buffer.data = buffer.data[:len(buffer.data)-overflow] + } + buffer.data = append(buffer.data, data...) + return length, nil +} + +func (buffer *agentPythonDiagnosticBuffer) String() string { return string(buffer.data) } +func (buffer *agentPythonDiagnosticBuffer) Reset() { buffer.data = buffer.data[:0] } + +func withAgentPythonDiagnostic(base error, diagnostic string) error { + if diagnostic == "" { + return base + } + return fmt.Errorf("%w; guest stderr: %s", base, diagnostic) +} diff --git a/internal/execution/wasm/agent_python_lifecycle_config_test.go b/internal/execution/wasm/agent_python_lifecycle_config_test.go new file mode 100644 index 0000000..6a719f7 --- /dev/null +++ b/internal/execution/wasm/agent_python_lifecycle_config_test.go @@ -0,0 +1,42 @@ +package wasm + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAgentPythonLifecycleDefaultsToSnapshotMemcpy(t *testing.T) { + cfg := Config{} + cfg.applyAgentPythonDefaults() + + assert.Equal(t, "snapshot", cfg.PythonLifecycle) + + assert.Equal(t, 1, cfg.PythonPreparedCapacity) + assert.Equal(t, uint64(8*1024*1024), cfg.PythonSnapshotHeadroomBytes) + require.NoError(t, cfg.validateAgentPythonLifecycle()) +} + +func TestAgentPythonLifecycleReadsExplicitSingleUseCapacity(t *testing.T) { + t.Setenv("FUNCTION_WASM_PYTHON_LIFECYCLE", "single-use") + t.Setenv("FUNCTION_WASM_PYTHON_PREPARED_CAPACITY", "2") + cfg := Config{} + cfg.applyEnv() + cfg.applyAgentPythonDefaults() + + assert.Equal(t, "single-use", cfg.PythonLifecycle) + assert.Equal(t, 2, cfg.PythonPreparedCapacity) + require.NoError(t, cfg.validateAgentPythonLifecycle()) +} + +func TestAgentPythonLifecycleRejectsUnknownAndOversizedCapacity(t *testing.T) { + for _, cfg := range []Config{ + {PythonLifecycle: "reuse-maybe"}, + {PythonLifecycle: "single-use", PythonPreparedCapacity: 5}, + {PythonLifecycle: "snapshot", MaxInstances: 5}, + } { + cfg.applyAgentPythonDefaults() + require.Error(t, cfg.validateAgentPythonLifecycle()) + } +} diff --git a/internal/execution/wasm/agent_python_observer.go b/internal/execution/wasm/agent_python_observer.go new file mode 100644 index 0000000..8ee3478 --- /dev/null +++ b/internal/execution/wasm/agent_python_observer.go @@ -0,0 +1,120 @@ +package wasm + +import "time" + +// AgentPythonPhase identifies one measured Agent Python lifecycle boundary. +type AgentPythonPhase string + +const ( + AgentPythonPhaseArtifactVerify AgentPythonPhase = "artifact-verify" + AgentPythonPhaseRuntimeCreate AgentPythonPhase = "runtime-create" + AgentPythonPhaseWASIImports AgentPythonPhase = "wasi-imports" + AgentPythonPhaseHostImports AgentPythonPhase = "host-imports" + AgentPythonPhaseCompile AgentPythonPhase = "compile" + AgentPythonPhaseInstantiate AgentPythonPhase = "instantiate" + AgentPythonPhaseInitialize AgentPythonPhase = "initialize" + AgentPythonPhaseRuntimeInit AgentPythonPhase = "runtime-init" + AgentPythonPhaseRuntimePrepare AgentPythonPhase = "runtime-prepare" + AgentPythonPhaseHeadroom AgentPythonPhase = "headroom" + AgentPythonPhaseStrategySelect AgentPythonPhase = "strategy-select" + AgentPythonPhaseSnapshotTake AgentPythonPhase = "snapshot-take" + AgentPythonPhaseCheckout AgentPythonPhase = "checkout" + AgentPythonPhaseExecute AgentPythonPhase = "execute" + AgentPythonPhaseDecode AgentPythonPhase = "decode" + AgentPythonPhaseRestore AgentPythonPhase = "restore" + AgentPythonPhaseClose AgentPythonPhase = "close" +) + +// AgentPythonPurpose explains why a slot or phase was created. +type AgentPythonPurpose string + +const ( + AgentPythonPurposeStartup AgentPythonPurpose = "startup" + AgentPythonPurposeRequest AgentPythonPurpose = "request" + AgentPythonPurposeFresh AgentPythonPurpose = "fresh" + AgentPythonPurposeRefill AgentPythonPurpose = "refill" + AgentPythonPurposeReplacement AgentPythonPurpose = "replacement" +) + +// AgentPythonOutcome is the terminal state of one observed phase. +type AgentPythonOutcome string + +const ( + AgentPythonOutcomeOK AgentPythonOutcome = "ok" + AgentPythonOutcomeError AgentPythonOutcome = "error" +) + +func agentPythonPhaseOutcome(err error) AgentPythonOutcome { + if err != nil { + return AgentPythonOutcomeError + } + return AgentPythonOutcomeOK +} + +// AgentPythonPhaseEvent is immutable phase evidence delivered after timing stops. +// Observer callbacks may be concurrent during single-use refill. +type AgentPythonPhaseEvent struct { + Phase AgentPythonPhase `json:"phase"` + Purpose AgentPythonPurpose `json:"purpose,omitempty"` + Lifecycle string `json:"lifecycle,omitempty"` + SnapshotRequested string `json:"snapshot_requested,omitempty"` + SnapshotSelected string `json:"snapshot_selected,omitempty"` + RequestID uint64 `json:"request_id,omitempty"` + SlotID uint64 `json:"slot_id,omitempty"` + Duration time.Duration `json:"duration_ns"` + MemoryBytes uint64 `json:"memory_bytes,omitempty"` + Outcome AgentPythonOutcome `json:"outcome"` + Error string `json:"error,omitempty"` +} + +// AgentPythonPhaseObservation is the internal input used to finish a phase. +type AgentPythonPhaseObservation struct { + Phase AgentPythonPhase + Purpose AgentPythonPurpose + RequestID uint64 + SlotID uint64 + Started time.Time + MemoryBytes uint64 + SnapshotSelected string + Outcome AgentPythonOutcome + Err error +} + +func (d *AgentPythonDispatcher) observeAgentPythonPhase(observation AgentPythonPhaseObservation) { + observer := d.cfg.AgentPythonObserver + if observer == nil { + return + } + duration := time.Duration(0) + if !observation.Started.IsZero() { + duration = time.Since(observation.Started) + } + event := AgentPythonPhaseEvent{ + Phase: observation.Phase, + Purpose: observation.Purpose, + Lifecycle: d.cfg.PythonLifecycle, + SnapshotRequested: d.snapshotMode(), + SnapshotSelected: observation.SnapshotSelected, + RequestID: observation.RequestID, + SlotID: observation.SlotID, + Duration: duration, + MemoryBytes: observation.MemoryBytes, + Outcome: observation.Outcome, + } + if observation.Err != nil { + event.Error = observation.Err.Error() + } + d.emitAgentPythonPhaseEvent(observer, event) +} + +func (d *AgentPythonDispatcher) emitAgentPythonPhaseEvent(observer func(AgentPythonPhaseEvent), event AgentPythonPhaseEvent) { + if observer == nil { + return + } + defer func() { + if recovered := recover(); recovered != nil { + d.log.Warn("agent-python observer panicked") + } + }() + observer(event) +} diff --git a/internal/execution/wasm/agent_python_protocol.go b/internal/execution/wasm/agent_python_protocol.go new file mode 100644 index 0000000..1700802 --- /dev/null +++ b/internal/execution/wasm/agent_python_protocol.go @@ -0,0 +1,484 @@ +package wasm + +// This file carries the consumer copy of the neutral Agent Python Runtime v1 +// request/response and artifact contract. The source contract was pinned from +// bkmashiro/agent-python-runtime guest commit +// 9a571176bb58c2d6a41312d01ad789abdd6b82e6 with repository-owner approval. + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "regexp" + "sort" +) + +const agentPythonPayloadMax = 1024 * 1024 + +const agentPythonPreparedCall = `_shimmy_dispatch = globals().get("dispatch") +if not callable(_shimmy_dispatch): + raise RuntimeError("python reactor artifact must define callable dispatch(method, payload)") +result = _shimmy_dispatch(inputs["method"], inputs["params"]) +` + +const agentPythonUnpreparedCall = `exec(compile(inputs["script"], "", "exec"), globals(), globals()) +` + agentPythonPreparedCall + +type AgentPythonArtifact struct { + WasmBytes []byte + ABI string + Profile string + PythonModules []string + ProducerCommit string + SHA256 string + ManifestPath string + InitExport string + PrepareExport string + ExecuteExport string + DeclaredExports []string + DeclaredImports []pythonReactorImport +} + +type agentPythonManifest struct { + SchemaVersion int `json:"schema_version"` + ABIVersion string `json:"abi_version"` + ArtifactProfile string `json:"artifact_profile"` + Target string `json:"target"` + Artifact struct { + Filename string `json:"filename"` + Size int64 `json:"size"` + SHA256 string `json:"sha256"` + } `json:"artifact"` + Build struct { + RepositoryCommit string `json:"repository_commit"` + SourceDateEpoch string `json:"source_date_epoch"` + CompilerTarget string `json:"compiler_target"` + ExecutionModel string `json:"execution_model"` + } `json:"build"` + Wasm struct { + Exports []string `json:"exports"` + Imports []pythonReactorImport `json:"imports"` + } `json:"wasm"` +} + +type shimmyPythonManifestEntry struct { + Name string `json:"name"` + Module string `json:"module"` + Kind string `json:"kind"` +} + +type shimmyPythonManifest struct { + Schema string `json:"schema"` + ArtifactContract string `json:"artifact_contract"` + Profile string `json:"profile"` + Target string `json:"target"` + ExecutionModel string `json:"execution_model"` + PythonModules []string `json:"python_modules"` + IdentityU32 uint32 `json:"identity_u32"` + Producer struct { + Project string `json:"project"` + Repository string `json:"repository"` + Commit string `json:"commit"` + Dirty bool `json:"dirty"` + } `json:"producer"` + SourceDateEpoch int64 `json:"source_date_epoch"` + Artifact struct { + Name string `json:"name"` + Size int64 `json:"size"` + SHA256 string `json:"sha256"` + } `json:"artifact"` + Wasm struct { + Exports []shimmyPythonManifestEntry `json:"exports"` + Imports []shimmyPythonManifestEntry `json:"imports"` + } `json:"wasm"` +} + +var agentPythonCommitPattern = regexp.MustCompile(`^[0-9a-f]{40}$`) + +func verifyAgentPythonArtifact(modulePath, manifestPath string) (*AgentPythonArtifact, error) { + if modulePath == "" { + return nil, errors.New("python-reactor: ModulePath must be set (FUNCTION_WASM_MODULE)") + } + if manifestPath == "" { + manifestPath = filepath.Join(filepath.Dir(modulePath), "manifest.json") + } + manifestBytes, err := os.ReadFile(manifestPath) + if err != nil { + return nil, fmt.Errorf("python-reactor: read manifest %q: %w", manifestPath, err) + } + var format struct { + Schema string `json:"schema"` + } + if err := json.Unmarshal(manifestBytes, &format); err != nil { + return nil, fmt.Errorf("python-reactor: parse manifest: %w", err) + } + if format.Schema != "" { + return verifyShimmyPythonArtifact(modulePath, manifestPath, manifestBytes) + } + var manifest agentPythonManifest + if err := json.Unmarshal(manifestBytes, &manifest); err != nil { + return nil, fmt.Errorf("python-reactor: parse manifest: %w", err) + } + if manifest.SchemaVersion != 2 || manifest.ABIVersion != "v1" { + return nil, fmt.Errorf("python-reactor: unsupported manifest schema/ABI %d/%q", manifest.SchemaVersion, manifest.ABIVersion) + } + if manifest.Target != "wasm32-wasip1" || manifest.Build.CompilerTarget != "wasm32-wasip1" || manifest.Build.ExecutionModel != "reactor" { + return nil, errors.New("python-reactor: manifest target must be a wasm32-wasip1 reactor") + } + if manifest.ArtifactProfile != "base" && manifest.ArtifactProfile != "numpy-core" { + return nil, fmt.Errorf("python-reactor: unsupported artifact profile %q", manifest.ArtifactProfile) + } + if !agentPythonCommitPattern.MatchString(manifest.Build.RepositoryCommit) { + return nil, errors.New("python-reactor: manifest producer commit must be 40 lowercase hex characters") + } + if manifest.Build.SourceDateEpoch == "" { + return nil, errors.New("python-reactor: manifest SOURCE_DATE_EPOCH is missing") + } + if filepath.Base(manifest.Artifact.Filename) != manifest.Artifact.Filename || manifest.Artifact.Filename != filepath.Base(modulePath) { + return nil, fmt.Errorf("python-reactor: manifest artifact filename %q does not bind module %q", manifest.Artifact.Filename, filepath.Base(modulePath)) + } + + wasmBytes, err := os.ReadFile(modulePath) + if err != nil { + return nil, fmt.Errorf("python-reactor: read artifact %q: %w", modulePath, err) + } + if len(wasmBytes) < 8 || !bytes.Equal(wasmBytes[:8], []byte("\x00asm\x01\x00\x00\x00")) { + return nil, errors.New("python-reactor: artifact is not a WebAssembly core module") + } + if int64(len(wasmBytes)) != manifest.Artifact.Size { + return nil, fmt.Errorf("python-reactor: artifact size %d does not match manifest %d", len(wasmBytes), manifest.Artifact.Size) + } + digest := sha256.Sum256(wasmBytes) + digestHex := hex.EncodeToString(digest[:]) + if digestHex != manifest.Artifact.SHA256 { + return nil, fmt.Errorf("python-reactor: artifact SHA-256 %s does not match manifest %s", digestHex, manifest.Artifact.SHA256) + } + + exports := make(map[string]struct{}, len(manifest.Wasm.Exports)) + for _, name := range manifest.Wasm.Exports { + if _, duplicate := exports[name]; duplicate { + return nil, fmt.Errorf("python-reactor: manifest repeats export %q", name) + } + exports[name] = struct{}{} + } + requiredExports := []string{"memory", "_initialize", "runtime_init", "runtime_prepare", "alloc", "dealloc", "execute"} + var missing []string + for _, name := range requiredExports { + if _, ok := exports[name]; !ok { + missing = append(missing, name) + } + } + if len(missing) > 0 { + sort.Strings(missing) + return nil, fmt.Errorf("python-reactor: manifest is missing required exports: %v", missing) + } + + hostCallCount := 0 + imports := make(map[pythonReactorImport]struct{}, len(manifest.Wasm.Imports)) + for _, imported := range manifest.Wasm.Imports { + if _, duplicate := imports[imported]; duplicate { + return nil, fmt.Errorf("python-reactor: manifest repeats import %q.%q", imported.Module, imported.Name) + } + imports[imported] = struct{}{} + if imported.Module == "wasi_snapshot_preview1" { + continue + } + if imported.Module == "agent_runtime_v1" && imported.Name == "host_call" { + hostCallCount++ + continue + } + return nil, fmt.Errorf("python-reactor: unexpected custom import %q.%q", imported.Module, imported.Name) + } + if hostCallCount != 1 { + return nil, fmt.Errorf("python-reactor: expected exactly one agent_runtime_v1.host_call import, got %d", hostCallCount) + } + + return &AgentPythonArtifact{ + WasmBytes: wasmBytes, + ABI: "agent-python-runtime/v1", + Profile: manifest.ArtifactProfile, + ProducerCommit: manifest.Build.RepositoryCommit, + SHA256: digestHex, + ManifestPath: manifestPath, + InitExport: "runtime_init", + PrepareExport: "runtime_prepare", + ExecuteExport: "execute", + DeclaredExports: append([]string(nil), manifest.Wasm.Exports...), + DeclaredImports: append([]pythonReactorImport(nil), manifest.Wasm.Imports...), + }, nil +} + +type agentPythonRunRequest struct { + RunID string `json:"run_id"` + Code string `json:"code"` + Inputs map[string]any `json:"inputs"` +} + +func verifyShimmyPythonArtifact(modulePath, manifestPath string, manifestBytes []byte) (*AgentPythonArtifact, error) { + var manifest shimmyPythonManifest + if err := json.Unmarshal(manifestBytes, &manifest); err != nil { + return nil, fmt.Errorf("python-reactor: parse Shimmy producer manifest: %w", err) + } + if manifest.Schema != "shimmy-python-runtime-artifact/v1" || manifest.ArtifactContract != "shimmy-python-runtime/v1" { + return nil, errors.New("python-reactor: unsupported Shimmy producer artifact contract") + } + if manifest.Target != "wasm32-wasip1" || manifest.ExecutionModel != "reactor" || manifest.IdentityU32 != 0x53505231 { + return nil, errors.New("python-reactor: Shimmy producer target, execution model, or identity mismatch") + } + if manifest.Producer.Project != "shimmy" || manifest.Producer.Dirty || !agentPythonCommitPattern.MatchString(manifest.Producer.Commit) { + return nil, errors.New("python-reactor: Shimmy producer identity is invalid or dirty") + } + if manifest.SourceDateEpoch <= 0 { + return nil, errors.New("python-reactor: Shimmy producer SOURCE_DATE_EPOCH is invalid") + } + expectedModules := map[string][]string{ + "base": {}, "numpy-core": {"numpy"}, "sympy": {"mpmath", "sympy"}, + } + modules, ok := expectedModules[manifest.Profile] + if !ok || !equalAgentPythonStrings(manifest.PythonModules, modules) { + return nil, fmt.Errorf("python-reactor: manifest python_modules do not match profile %q", manifest.Profile) + } + if filepath.Base(manifest.Artifact.Name) != manifest.Artifact.Name || manifest.Artifact.Name != filepath.Base(modulePath) { + return nil, fmt.Errorf("python-reactor: manifest artifact name %q does not bind module %q", manifest.Artifact.Name, filepath.Base(modulePath)) + } + wasmBytes, err := os.ReadFile(modulePath) + if err != nil { + return nil, fmt.Errorf("python-reactor: read artifact %q: %w", modulePath, err) + } + if len(wasmBytes) < 8 || !bytes.Equal(wasmBytes[:8], []byte("\x00asm\x01\x00\x00\x00")) { + return nil, errors.New("python-reactor: artifact is not a WebAssembly v1 module") + } + if manifest.Artifact.Size != int64(len(wasmBytes)) { + return nil, fmt.Errorf("python-reactor: artifact size mismatch: manifest=%d actual=%d", manifest.Artifact.Size, len(wasmBytes)) + } + digest := sha256.Sum256(wasmBytes) + digestHex := hex.EncodeToString(digest[:]) + if manifest.Artifact.SHA256 != digestHex { + return nil, errors.New("python-reactor: artifact SHA-256 does not match manifest") + } + + exports := make([]string, 0, len(manifest.Wasm.Exports)) + seenExports := make(map[string]struct{}, len(manifest.Wasm.Exports)) + for _, entry := range manifest.Wasm.Exports { + if entry.Name == "" || entry.Kind == "" { + return nil, errors.New("python-reactor: malformed Shimmy producer export declaration") + } + if _, duplicate := seenExports[entry.Name]; duplicate { + return nil, fmt.Errorf("python-reactor: duplicate manifest export %q", entry.Name) + } + seenExports[entry.Name] = struct{}{} + exports = append(exports, entry.Name) + } + imports := make([]pythonReactorImport, 0, len(manifest.Wasm.Imports)) + seenImports := make(map[pythonReactorImport]struct{}, len(manifest.Wasm.Imports)) + for _, entry := range manifest.Wasm.Imports { + declared := pythonReactorImport{Module: entry.Module, Name: entry.Name} + if declared.Module != "wasi_snapshot_preview1" || declared.Name == "" || entry.Kind == "" { + return nil, fmt.Errorf("python-reactor: unexpected Shimmy producer import %q.%q", declared.Module, declared.Name) + } + if _, duplicate := seenImports[declared]; duplicate { + return nil, fmt.Errorf("python-reactor: duplicate manifest import %q.%q", declared.Module, declared.Name) + } + seenImports[declared] = struct{}{} + imports = append(imports, declared) + } + return &AgentPythonArtifact{ + WasmBytes: wasmBytes, ABI: "shimmy-python-runtime/v1", Profile: manifest.Profile, + PythonModules: append([]string(nil), manifest.PythonModules...), + ProducerCommit: manifest.Producer.Commit, SHA256: digestHex, ManifestPath: manifestPath, + InitExport: "shimmy_python_init", PrepareExport: "shimmy_python_prepare", ExecuteExport: "evaluate", + DeclaredExports: exports, DeclaredImports: imports, + }, nil +} + +func equalAgentPythonStrings(left, right []string) bool { + if len(left) != len(right) { + return false + } + for index := range left { + if left[index] != right[index] { + return false + } + } + return true +} + +func buildAgentPythonRunRequest(runID, method string, params map[string]any, script string) ([]byte, error) { + if runID == "" { + return nil, errors.New("python-reactor: run ID is required") + } + if method == "" { + method = "eval" + } + if params == nil { + params = map[string]any{} + } + inputs := map[string]any{"method": method, "params": params} + code := agentPythonPreparedCall + if script != "" { + inputs["script"] = script + code = agentPythonUnpreparedCall + } + payload, err := json.Marshal(agentPythonRunRequest{RunID: runID, Code: code, Inputs: inputs}) + if err != nil { + return nil, fmt.Errorf("python-reactor: encode run request: %w", err) + } + if len(payload) > agentPythonPayloadMax { + return nil, fmt.Errorf("python-reactor: run request exceeds %d-byte guest bound", agentPythonPayloadMax) + } + return payload, nil +} + +func buildShimmyPythonRunRequest(method string, params map[string]any) ([]byte, error) { + if method == "" { + method = "eval" + } + if params == nil { + params = map[string]any{} + } + payload, err := json.Marshal(map[string]any{"method": method, "params": params}) + if err != nil { + return nil, fmt.Errorf("python-reactor: encode Shimmy producer request: %w", err) + } + if len(payload) > agentPythonPayloadMax { + return nil, fmt.Errorf("python-reactor: run request exceeds %d-byte guest bound", agentPythonPayloadMax) + } + return payload, nil +} + +type shimmyPythonRunResponse struct { + Status string `json:"status"` + Result json.RawMessage `json:"result"` + Error *struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` +} + +func decodeShimmyPythonResponse(payload []byte) (map[string]any, error) { + decoder := json.NewDecoder(bytes.NewReader(payload)) + decoder.DisallowUnknownFields() + var response shimmyPythonRunResponse + if err := decoder.Decode(&response); err != nil { + return nil, fmt.Errorf("python-reactor: decode Shimmy producer response: %w", err) + } + if err := ensureAgentPythonJSONEOF(decoder); err != nil { + return nil, err + } + switch response.Status { + case "ok": + if response.Error != nil || len(response.Result) == 0 || bytes.Equal(response.Result, []byte("null")) { + return nil, errors.New("python-reactor: successful Shimmy producer response is malformed") + } + var result map[string]any + if err := json.Unmarshal(response.Result, &result); err != nil || result == nil { + return nil, errors.New("python-reactor: evaluator result must be a JSON object") + } + return result, nil + case "error": + if response.Error == nil || response.Error.Type == "" || response.Error.Message == "" || len(response.Result) != 0 { + return nil, errors.New("python-reactor: failed Shimmy producer response is malformed") + } + return nil, &PythonReactorExecutionError{ + Code: "guest_error", Message: response.Error.Message, ErrorType: response.Error.Type, + } + default: + return nil, fmt.Errorf("python-reactor: unsupported response status %q", response.Status) + } +} + +type agentPythonRunResponse struct { + Status string `json:"status"` + Result json.RawMessage `json:"result"` + Receipts []json.RawMessage `json:"receipts"` + Metrics *struct { + GuestTimeMS *float64 `json:"guest_time_ms,omitempty"` + CapabilityCalls uint32 `json:"capability_calls"` + ResultBytes uint32 `json:"result_bytes"` + } `json:"metrics"` + Error *struct { + Code string `json:"code"` + Message string `json:"message"` + ErrorType *string `json:"error_type,omitempty"` + Traceback *string `json:"traceback,omitempty"` + } `json:"error"` +} + +// PythonReactorExecutionError preserves a structured error returned by the +// evaluator-owned dispatcher. The sandbox does not reinterpret it as a normal +// result or map it to a different business method. +type PythonReactorExecutionError struct { + Code string + Message string + ErrorType string + Traceback string +} + +func (e *PythonReactorExecutionError) Error() string { + if e == nil { + return "python-reactor: execution failed" + } + if e.Code == "" { + return "python-reactor: " + e.Message + } + return fmt.Sprintf("python-reactor: %s: %s", e.Code, e.Message) +} + +func decodeAgentPythonResponse(payload []byte) (map[string]any, error) { + decoder := json.NewDecoder(bytes.NewReader(payload)) + decoder.DisallowUnknownFields() + var response agentPythonRunResponse + if err := decoder.Decode(&response); err != nil { + return nil, fmt.Errorf("python-reactor: decode response: %w", err) + } + if err := ensureAgentPythonJSONEOF(decoder); err != nil { + return nil, err + } + if response.Metrics == nil || (response.Metrics.GuestTimeMS != nil && *response.Metrics.GuestTimeMS < 0) { + return nil, errors.New("python-reactor: response metrics are invalid") + } + switch response.Status { + case "ok": + if response.Error != nil || len(response.Result) == 0 || bytes.Equal(response.Result, []byte("null")) { + return nil, errors.New("python-reactor: successful response has invalid result/error fields") + } + var result map[string]any + if err := json.Unmarshal(response.Result, &result); err != nil || result == nil { + return nil, errors.New("python-reactor: evaluator result must be a JSON object") + } + return result, nil + case "error": + if response.Error == nil || response.Error.Code == "" || response.Error.Message == "" || !bytes.Equal(response.Result, []byte("null")) { + return nil, errors.New("python-reactor: failed response has invalid result/error fields") + } + executionErr := &PythonReactorExecutionError{ + Code: response.Error.Code, + Message: response.Error.Message, + } + if response.Error.ErrorType != nil { + executionErr.ErrorType = *response.Error.ErrorType + } + if response.Error.Traceback != nil { + executionErr.Traceback = *response.Error.Traceback + } + return nil, executionErr + default: + return nil, fmt.Errorf("python-reactor: unsupported response status %q", response.Status) + } +} + +func ensureAgentPythonJSONEOF(decoder *json.Decoder) error { + var trailing any + if err := decoder.Decode(&trailing); errors.Is(err, io.EOF) { + return nil + } else if err != nil { + return fmt.Errorf("python-reactor: decode trailing response JSON: %w", err) + } + return errors.New("python-reactor: response contains trailing JSON") +} diff --git a/internal/execution/wasm/agent_python_test.go b/internal/execution/wasm/agent_python_test.go new file mode 100644 index 0000000..8b4c1da --- /dev/null +++ b/internal/execution/wasm/agent_python_test.go @@ -0,0 +1,750 @@ +package wasm + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "os" + "os/exec" + "path/filepath" + "runtime" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" + "go.uber.org/zap" +) + +func writeAgentPythonManifestFixture(t *testing.T, customModule, customName string) (string, string) { + t.Helper() + dir := t.TempDir() + wasmPath := filepath.Join(dir, "agent-python-runtime.wasm") + wasmBytes := []byte("\x00asm\x01\x00\x00\x00fixture") + require.NoError(t, os.WriteFile(wasmPath, wasmBytes, 0o644)) + digest := sha256.Sum256(wasmBytes) + manifest := map[string]any{ + "schema_version": 2, + "abi_version": "v1", + "artifact_profile": "base", + "target": "wasm32-wasip1", + "artifact": map[string]any{ + "filename": filepath.Base(wasmPath), + "size": len(wasmBytes), + "sha256": hex.EncodeToString(digest[:]), + }, + "build": map[string]any{ + "repository_commit": "a3b7c9d1e5f80123456789abcdef0123456789ab", + "source_date_epoch": "1784781655", + "compiler_target": "wasm32-wasip1", + "execution_model": "reactor", + }, + "wasm": map[string]any{ + "exports": []string{ + "memory", "runtime_init", "runtime_prepare", "alloc", "dealloc", "execute", "_initialize", + }, + "imports": []map[string]string{ + {"module": customModule, "name": customName}, + {"module": "wasi_snapshot_preview1", "name": "random_get"}, + }, + }, + } + encoded, err := json.MarshalIndent(manifest, "", " ") + require.NoError(t, err) + manifestPath := filepath.Join(dir, "manifest.json") + require.NoError(t, os.WriteFile(manifestPath, append(encoded, '\n'), 0o644)) + return wasmPath, manifestPath +} + +func TestVerifyAgentPythonArtifactAcceptsPinnedV1Contract(t *testing.T) { + wasmPath, manifestPath := writeAgentPythonManifestFixture(t, "agent_runtime_v1", "host_call") + + artifact, err := verifyAgentPythonArtifact(wasmPath, manifestPath) + + require.NoError(t, err) + assert.Equal(t, "base", artifact.Profile) + assert.Equal(t, "a3b7c9d1e5f80123456789abcdef0123456789ab", artifact.ProducerCommit) + assert.Len(t, artifact.WasmBytes, 15) +} + +func TestVerifyAgentPythonArtifactRejectsUnexpectedCustomImport(t *testing.T) { + wasmPath, manifestPath := writeAgentPythonManifestFixture(t, "legacy_env", "stub") + + _, err := verifyAgentPythonArtifact(wasmPath, manifestPath) + + require.Error(t, err) + assert.Contains(t, err.Error(), `unexpected custom import "legacy_env"."stub"`) +} + +func TestVerifyAgentPythonArtifactRejectsDigestDrift(t *testing.T) { + wasmPath, manifestPath := writeAgentPythonManifestFixture(t, "agent_runtime_v1", "host_call") + require.NoError(t, os.WriteFile(wasmPath, []byte("\x00asm\x01\x00\x00\x00changed"), 0o644)) + + _, err := verifyAgentPythonArtifact(wasmPath, manifestPath) + + require.Error(t, err) + assert.Contains(t, err.Error(), "artifact SHA-256") +} + +func writeShimmyPythonManifestFixture(t *testing.T, profile string, modules []string) (string, string) { + t.Helper() + dir := t.TempDir() + wasmPath := filepath.Join(dir, "shimmy-python-runtime-"+profile+".wasm") + wasmBytes := []byte("\x00asm\x01\x00\x00\x00producer") + require.NoError(t, os.WriteFile(wasmPath, wasmBytes, 0o644)) + digest := sha256.Sum256(wasmBytes) + manifest := map[string]any{ + "schema": "shimmy-python-runtime-artifact/v1", + "artifact_contract": "shimmy-python-runtime/v1", + "profile": profile, "target": "wasm32-wasip1", "execution_model": "reactor", + "python_modules": modules, "identity_u32": 1397772849, + "producer": map[string]any{ + "project": "shimmy", "repository": "lambda-feedback/shimmy", + "commit": "a3b7c9d1e5f80123456789abcdef0123456789ab", "dirty": false, + }, + "source_date_epoch": 1784781655, + "artifact": map[string]any{ + "name": filepath.Base(wasmPath), "size": len(wasmBytes), "sha256": hex.EncodeToString(digest[:]), + }, + "wasm": map[string]any{ + "exports": []map[string]string{ + {"name": "memory", "kind": "memory"}, {"name": "_initialize", "kind": "function"}, + {"name": "shimmy_python_runtime_identity", "kind": "function"}, + {"name": "shimmy_python_init", "kind": "function"}, + {"name": "shimmy_python_prepare", "kind": "function"}, + {"name": "alloc", "kind": "function"}, {"name": "dealloc", "kind": "function"}, + {"name": "evaluate", "kind": "function"}, + }, + "imports": []map[string]string{{"module": "wasi_snapshot_preview1", "name": "fd_write", "kind": "function"}}, + }, + } + encoded, err := json.MarshalIndent(manifest, "", " ") + require.NoError(t, err) + manifestPath := filepath.Join(dir, "manifest.json") + require.NoError(t, os.WriteFile(manifestPath, append(encoded, '\n'), 0o644)) + return wasmPath, manifestPath +} + +func TestVerifyAgentPythonArtifactAcceptsShimmyProducerContract(t *testing.T) { + wasmPath, manifestPath := writeShimmyPythonManifestFixture(t, "sympy", []string{"mpmath", "sympy"}) + artifact, err := verifyAgentPythonArtifact(wasmPath, manifestPath) + require.NoError(t, err) + assert.Equal(t, "shimmy-python-runtime/v1", artifact.ABI) + assert.Equal(t, []string{"mpmath", "sympy"}, artifact.PythonModules) + assert.Equal(t, "shimmy_python_init", artifact.InitExport) + assert.Equal(t, "shimmy_python_prepare", artifact.PrepareExport) + assert.Equal(t, "evaluate", artifact.ExecuteExport) +} + +func TestVerifyAgentPythonArtifactRejectsFalseProfileModules(t *testing.T) { + wasmPath, manifestPath := writeShimmyPythonManifestFixture(t, "base", []string{"sympy"}) + _, err := verifyAgentPythonArtifact(wasmPath, manifestPath) + require.Error(t, err) + assert.Contains(t, err.Error(), "python_modules") +} + +func validPythonReactorModuleShape() pythonReactorModuleShape { + i32 := api.ValueTypeI32 + return pythonReactorModuleShape{ + Exports: map[string]pythonReactorFunctionSignature{ + "_initialize": {}, + "runtime_init": {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + "runtime_prepare": {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + "alloc": {Params: []api.ValueType{i32}, Results: []api.ValueType{i32}}, + "dealloc": {Params: []api.ValueType{i32}}, + "execute": {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + }, + ExportedMemories: map[string]struct{}{"memory": {}}, + Imports: map[pythonReactorImport]struct{}{ + {Module: "agent_runtime_v1", Name: "host_call"}: {}, + {Module: "wasi_snapshot_preview1", Name: "fd_write"}: {}, + }, + } +} + +func validPythonReactorArtifactContract() *AgentPythonArtifact { + return &AgentPythonArtifact{ + DeclaredExports: []string{"memory", "_initialize", "runtime_init", "runtime_prepare", "alloc", "dealloc", "execute"}, + DeclaredImports: []pythonReactorImport{ + {Module: "agent_runtime_v1", Name: "host_call"}, + {Module: "wasi_snapshot_preview1", Name: "fd_write"}, + }, + } +} + +func TestVerifyPythonReactorModuleShapeAcceptsShimmyProducerABI(t *testing.T) { + i32 := api.ValueTypeI32 + shape := pythonReactorModuleShape{ + Exports: map[string]pythonReactorFunctionSignature{ + "_initialize": {}, + "shimmy_python_runtime_identity": {Results: []api.ValueType{i32}}, + "shimmy_python_init": {Results: []api.ValueType{i32}}, + "shimmy_python_prepare": {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + "alloc": {Params: []api.ValueType{i32}, Results: []api.ValueType{i32}}, + "dealloc": {Params: []api.ValueType{i32}}, + "evaluate": {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + }, + ExportedMemories: map[string]struct{}{"memory": {}}, + Imports: map[pythonReactorImport]struct{}{ + {Module: "wasi_snapshot_preview1", Name: "fd_write"}: {}, + }, + } + artifact := &AgentPythonArtifact{ + ABI: "shimmy-python-runtime/v1", InitExport: "shimmy_python_init", + PrepareExport: "shimmy_python_prepare", ExecuteExport: "evaluate", + DeclaredExports: []string{"memory", "_initialize", "shimmy_python_runtime_identity", "shimmy_python_init", "shimmy_python_prepare", "alloc", "dealloc", "evaluate"}, + DeclaredImports: []pythonReactorImport{{Module: "wasi_snapshot_preview1", Name: "fd_write"}}, + } + require.NoError(t, verifyPythonReactorModuleShape(shape, artifact)) +} + +func TestVerifyPythonReactorModuleShapeAcceptsExactContract(t *testing.T) { + err := verifyPythonReactorModuleShape(validPythonReactorModuleShape(), validPythonReactorArtifactContract()) + require.NoError(t, err) +} + +func TestVerifyPythonReactorModuleShapeRejectsUndeclaredActualImport(t *testing.T) { + shape := validPythonReactorModuleShape() + shape.Imports[pythonReactorImport{Module: "wasi_snapshot_preview1", Name: "sock_send"}] = struct{}{} + + err := verifyPythonReactorModuleShape(shape, validPythonReactorArtifactContract()) + + require.Error(t, err) + assert.Contains(t, err.Error(), `actual import "wasi_snapshot_preview1"."sock_send" is not declared by manifest`) +} + +func TestVerifyPythonReactorModuleShapeRejectsWrongDispatchABISignature(t *testing.T) { + shape := validPythonReactorModuleShape() + shape.Exports["execute"] = pythonReactorFunctionSignature{ + Params: []api.ValueType{api.ValueTypeI64}, + Results: []api.ValueType{api.ValueTypeI32}, + } + + err := verifyPythonReactorModuleShape(shape, validPythonReactorArtifactContract()) + + require.Error(t, err) + assert.Contains(t, err.Error(), `export "execute" has ABI`) +} + +func TestBuildAgentPythonRunRequestPreservesArbitraryMethodAndOpaqueParams(t *testing.T) { + params := map[string]any{ + "messages": []any{map[string]any{"role": "USER", "content": "hello"}}, + "future_field": map[string]any{"nested": true}, + } + + request, err := buildAgentPythonRunRequest("shimmy-run-1", "future/chat.v2", params, "") + + require.NoError(t, err) + var envelope struct { + RunID string `json:"run_id"` + Code string `json:"code"` + Inputs map[string]any `json:"inputs"` + } + require.NoError(t, json.Unmarshal(request, &envelope)) + assert.Equal(t, "shimmy-run-1", envelope.RunID) + assert.Equal(t, agentPythonPreparedCall, envelope.Code) + assert.Equal(t, "future/chat.v2", envelope.Inputs["method"]) + assert.Equal(t, params["messages"], envelope.Inputs["params"].(map[string]any)["messages"]) + assert.Equal(t, true, envelope.Inputs["params"].(map[string]any)["future_field"].(map[string]any)["nested"]) + assert.Contains(t, envelope.Code, `dispatch(inputs["method"], inputs["params"])`) + assert.NotContains(t, envelope.Code, "evaluation_function") + assert.NotContains(t, envelope.Code, "preview_function") + assert.NotContains(t, envelope.Code, "shimmy-run-1") +} + +func TestShimmyProducerRequestAndResponseContract(t *testing.T) { + request, err := buildShimmyPythonRunRequest("preview", map[string]any{"response": "x", "params": map[string]any{}}) + require.NoError(t, err) + assert.JSONEq(t, `{"method":"preview","params":{"response":"x","params":{}}}`, string(request)) + + result, err := decodeShimmyPythonResponse([]byte(`{"status":"ok","result":{"preview":{"sympy":"x"}}}`)) + require.NoError(t, err) + assert.Equal(t, "x", result["preview"].(map[string]any)["sympy"]) +} + +func TestShimmyProducerResponsePreservesTypedError(t *testing.T) { + _, err := decodeShimmyPythonResponse([]byte(`{"status":"error","error":{"type":"ImportError","message":"No module named scipy"}}`)) + var executionErr *PythonReactorExecutionError + require.ErrorAs(t, err, &executionErr) + assert.Equal(t, "ImportError", executionErr.ErrorType) + assert.Equal(t, "No module named scipy", executionErr.Message) +} + +func TestBuildAgentPythonRunRequestSupportsExplicitPreloadOff(t *testing.T) { + request, err := buildAgentPythonRunRequest( + "shimmy-run-2", + "eval", + map[string]any{"response": "1", "answer": "1"}, + "def dispatch(method, payload): return {'method': method, 'payload': payload}", + ) + + require.NoError(t, err) + var envelope map[string]any + require.NoError(t, json.Unmarshal(request, &envelope)) + inputs := envelope["inputs"].(map[string]any) + assert.Contains(t, envelope["code"], `inputs["script"]`) + assert.Contains(t, inputs["script"], "def dispatch(method, payload)") + assert.NotContains(t, envelope["code"], "evaluation_function") + assert.NotContains(t, envelope["code"], "preview_function") +} + +func TestDecodeAgentPythonResponsePreservesSuccessResult(t *testing.T) { + payload := []byte(`{"status":"ok","result":{"opaque":{"value":true}},"receipts":[],"metrics":{"capability_calls":0,"result_bytes":25},"error":null}`) + + result, err := decodeAgentPythonResponse(payload) + + require.NoError(t, err) + assert.Equal(t, map[string]any{"value": true}, result["opaque"]) +} + +func TestDecodeAgentPythonResponseReturnsTypedExecutionError(t *testing.T) { + payload := []byte(`{"status":"error","result":null,"receipts":[],"metrics":{"capability_calls":0,"result_bytes":0},"error":{"code":"unsupported_method","message":"method is not registered","error_type":"UnsupportedMethod","traceback":"trace"}}`) + + result, err := decodeAgentPythonResponse(payload) + + require.Nil(t, result) + var executionErr *PythonReactorExecutionError + require.ErrorAs(t, err, &executionErr) + assert.Equal(t, "unsupported_method", executionErr.Code) + assert.Equal(t, "method is not registered", executionErr.Message) + assert.Equal(t, "UnsupportedMethod", executionErr.ErrorType) + assert.Equal(t, "trace", executionErr.Traceback) +} + +func TestAgentPythonRejectsHostFilesystemPaths(t *testing.T) { + t.Setenv("FUNCTION_WASM_ALLOWED_PATHS", "/tmp") + dispatcher := NewAgentPythonDispatcher(Config{}, zap.NewNop()) + err := dispatcher.Start(context.Background()) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not expose Host filesystem paths") +} + +func TestAgentPythonDispatcherRealNumPyArtifactCompatibility(t *testing.T) { + wasmPath := os.Getenv("AGENT_PYTHON_RUNTIME_WASM") + manifestPath := os.Getenv("AGENT_PYTHON_RUNTIME_MANIFEST") + if wasmPath == "" || manifestPath == "" { + t.Skip("AGENT_PYTHON_RUNTIME_WASM and AGENT_PYTHON_RUNTIME_MANIFEST are required") + } + + scriptPath := filepath.Join(t.TempDir(), "eval.py") + script := ` +import numpy as np +_counter = 0 + +def dispatch(method, payload): + if method == "preview": + return {"preview": f"response={payload.get('response')}"} + if method != "eval": + raise LookupError("unsupported method: " + method) + response = payload.get("response") + answer = payload.get("answer") + global _counter + _counter += 1 + if response == "explode": + raise ValueError("expected explosion") + if response == "host_call": + from agent_runtime.tools import fetch_many + return fetch_many([{"request_id": "r1", "target": "fixture", "path": "/ok"}]) + if response == "float128": + one = np.longdouble("1") + wide = np.longdouble("1.0000000000000000000000000000000002") + return { + "longdouble_itemsize": int(np.dtype(np.longdouble).itemsize), + "longdouble_nmant": int(np.finfo(np.longdouble).nmant), + "double_nmant": int(np.finfo(np.double).nmant), + "preserves_extra_precision": bool(wide > one), + "narrows_to_double_one": bool(float(wide) == 1.0), + "epsilon_is_narrower": bool(np.finfo(np.longdouble).eps < np.finfo(np.double).eps), + "counter": _counter, + } + return {"is_correct": response == answer, "counter": _counter} +` + require.NoError(t, os.WriteFile(scriptPath, []byte(script), 0o644)) + + dispatcher := NewAgentPythonDispatcher(Config{ + ModulePath: wasmPath, + AgentPythonManifestPath: manifestPath, + PythonScriptPath: scriptPath, + PythonLifecycle: "snapshot", + MaxInstances: 1, + MaxMemoryPages: 8192, + Timeout: 120 * time.Second, + }, zap.NewNop()) + startContext, startCancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer startCancel() + require.NoError(t, dispatcher.Start(startContext)) + t.Cleanup(func() { + shutdownContext, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _ = dispatcher.Shutdown(shutdownContext) + }) + + health, err := dispatcher.Send(context.Background(), "healthcheck", nil) + require.NoError(t, err) + assert.Equal(t, "snapshot", health["result"].(map[string]any)["lifecycle"]) + assert.Equal(t, "memcpy", health["result"].(map[string]any)["snapshot_mode"]) + assert.Equal(t, "linear-memory-memcpy", health["result"].(map[string]any)["reset_mode"]) + + first, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "42", "answer": "42"}) + require.NoError(t, err) + assert.Equal(t, true, first["result"].(map[string]any)["is_correct"]) + assert.Equal(t, float64(1), first["result"].(map[string]any)["counter"]) + + second, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "42", "answer": "42"}) + require.NoError(t, err) + assert.Equal(t, float64(1), second["result"].(map[string]any)["counter"], "snapshot restore must not retain globals") + + preview, err := dispatcher.Send(context.Background(), "preview", map[string]any{"response": "3.14", "answer": "3.14"}) + require.NoError(t, err) + assert.Equal(t, "response=3.14", preview["result"].(map[string]any)["preview"]) + + failure, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "explode", "answer": "x"}) + require.Nil(t, failure) + var failureErr *PythonReactorExecutionError + require.ErrorAs(t, err, &failureErr) + assert.Equal(t, "ValueError", failureErr.ErrorType) + assert.Equal(t, "expected explosion", failureErr.Message) + + denied, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "host_call", "answer": "x"}) + require.Nil(t, denied) + var deniedErr *PythonReactorExecutionError + require.ErrorAs(t, err, &deniedErr) + assert.Equal(t, "RuntimeError", deniedErr.ErrorType) + assert.Contains(t, deniedErr.Message, "Host capability bridge rejected") + + binary128, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "float128", "answer": "x"}) + require.NoError(t, err) + value := binary128["result"].(map[string]any) + assert.Equal(t, float64(16), value["longdouble_itemsize"]) + assert.GreaterOrEqual(t, value["longdouble_nmant"].(float64), float64(112)) + assert.Equal(t, true, value["preserves_extra_precision"]) + assert.Equal(t, true, value["narrows_to_double_one"]) + assert.Equal(t, true, value["epsilon_is_narrower"]) +} + +func TestAcquireAgentPythonSnapshotSlotReplenishesMissingSlot(t *testing.T) { + prepared := make(chan *agentPythonModuleSlot, 1) + closed := make(chan struct{}) + want := &agentPythonModuleSlot{snapshotSelected: "memcpy"} + calls := 0 + + got, err := acquireAgentPythonSnapshotSlot( + context.Background(), + prepared, + closed, + nil, + func(context.Context) (*agentPythonModuleSlot, error) { + calls++ + return want, nil + }, + ) + + require.NoError(t, err) + assert.Same(t, want, got) + assert.Equal(t, 1, calls) +} + +func TestAcquireAgentPythonSnapshotSlotReturnsReplenishFailure(t *testing.T) { + prepared := make(chan *agentPythonModuleSlot, 1) + closed := make(chan struct{}) + wantErr := errors.New("replacement unavailable") + + _, err := acquireAgentPythonSnapshotSlot( + context.Background(), + prepared, + closed, + nil, + func(context.Context) (*agentPythonModuleSlot, error) { + return nil, wantErr + }, + ) + + require.ErrorIs(t, err, wantErr) +} + +func TestAcquireAgentPythonSnapshotSlotWaitsForInFlightRefill(t *testing.T) { + prepared := make(chan *agentPythonModuleSlot, 1) + closed := make(chan struct{}) + want := &agentPythonModuleSlot{snapshotSelected: "memcpy"} + createCalls := 0 + go func() { + time.Sleep(10 * time.Millisecond) + prepared <- want + }() + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + got, err := acquireAgentPythonSnapshotSlot( + ctx, + prepared, + closed, + func() bool { return true }, + func(context.Context) (*agentPythonModuleSlot, error) { + createCalls++ + return nil, errors.New("must not construct a duplicate slot") + }, + ) + + require.NoError(t, err) + assert.Same(t, want, got) + assert.Zero(t, createCalls) +} + +func TestRestoreAgentPythonSnapshotRejectsMemoryGrowth(t *testing.T) { + ctx := context.Background() + rt, compiled := compileEchoModule(t, ctx, echoWasmBytes(t)) + t.Cleanup(func() { require.NoError(t, rt.Close(ctx)) }) + module, err := rt.InstantiateModule(ctx, compiled, wazero.NewModuleConfig()) + require.NoError(t, err) + + strategy := NewFullMemcpyStrategy() + require.NoError(t, strategy.Take(module.Memory())) + slot := &agentPythonModuleSlot{ + module: module, + strategy: strategy, + baselineSize: module.Memory().Size(), + } + _, grew := module.Memory().Grow(1) + require.True(t, grew) + + err = restoreAgentPythonSnapshot(slot) + require.Error(t, err) + assert.Contains(t, err.Error(), "memory size drift") +} + +func TestAgentPythonDispatcherProducerTimeoutReturnsBeforeSnapshotRefill(t *testing.T) { + wasmPath := os.Getenv("AGENT_PYTHON_RUNTIME_WASM") + manifestPath := os.Getenv("AGENT_PYTHON_RUNTIME_MANIFEST") + evaluatorPath := os.Getenv("SAFE_EVAL_PYTHON_SCRIPT") + if wasmPath == "" || manifestPath == "" || evaluatorPath == "" { + t.Skip("AGENT_PYTHON_RUNTIME_WASM, AGENT_PYTHON_RUNTIME_MANIFEST, and SAFE_EVAL_PYTHON_SCRIPT are required") + } + + dispatcher := NewAgentPythonDispatcher(Config{ + ModulePath: wasmPath, + AgentPythonManifestPath: manifestPath, + PythonScriptPath: evaluatorPath, + PythonLifecycle: "snapshot", + MaxInstances: 1, + MaxMemoryPages: 8192, + Timeout: 2 * time.Second, + }, zap.NewNop()) + startContext, startCancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer startCancel() + require.NoError(t, dispatcher.Start(startContext)) + t.Cleanup(func() { + shutdownContext, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + _ = dispatcher.Shutdown(shutdownContext) + }) + + started := time.Now() + _, err := dispatcher.Send(context.Background(), "eval", map[string]any{ + "response": "while True:\n pass", "answer": "", "params": map[string]any{"mode": "demo"}, + }) + elapsed := time.Since(started) + t.Logf("timeout request returned in %s", elapsed) + require.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, elapsed, 10*time.Second, "request must not wait for close or snapshot refill") + + require.Eventually(t, func() bool { + health, healthErr := dispatcher.Send(context.Background(), "healthcheck", nil) + if healthErr != nil { + return false + } + return health["result"].(map[string]any)["prepared_ready"] == 1 + }, time.Minute, 100*time.Millisecond) + + recovered, err := dispatcher.Send(context.Background(), "eval", map[string]any{ + "response": "print(7 * 6)", "answer": "", "params": map[string]any{"mode": "demo"}, + }) + require.NoError(t, err) + assert.Equal(t, "42\n", recovered["result"].(map[string]any)["stdout"]) +} + +func TestAgentPythonDispatcherSingleUsePreparedRefillsNeverServedCandidates(t *testing.T) { + wasmPath := os.Getenv("AGENT_PYTHON_RUNTIME_WASM") + manifestPath := os.Getenv("AGENT_PYTHON_RUNTIME_MANIFEST") + if wasmPath == "" || manifestPath == "" { + t.Skip("AGENT_PYTHON_RUNTIME_WASM and AGENT_PYTHON_RUNTIME_MANIFEST are required") + } + + scriptPath := filepath.Join(t.TempDir(), "single-use.py") + script := ` +_counter = 0 + +def dispatch(method, payload): + if method != "eval": + raise LookupError("unsupported method: " + method) + global _counter + _counter += 1 + return {"counter": _counter, "is_correct": payload.get("response") == payload.get("answer")} +` + require.NoError(t, os.WriteFile(scriptPath, []byte(script), 0o644)) + + dispatcher := NewAgentPythonDispatcher(Config{ + ModulePath: wasmPath, + AgentPythonManifestPath: manifestPath, + PythonScriptPath: scriptPath, + PythonLifecycle: "single-use", + PythonPreparedCapacity: 1, + MaxInstances: 1, + MaxMemoryPages: 8192, + Timeout: 120 * time.Second, + }, zap.NewNop()) + startContext, startCancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer startCancel() + require.NoError(t, dispatcher.Start(startContext)) + t.Cleanup(func() { + shutdownContext, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _ = dispatcher.Shutdown(shutdownContext) + }) + + health, err := dispatcher.Send(context.Background(), "healthcheck", nil) + require.NoError(t, err) + assert.Equal(t, "single-use", health["result"].(map[string]any)["lifecycle"]) + assert.Equal(t, 1, health["result"].(map[string]any)["prepared_ready"]) + assert.Equal(t, "single-use-prepared", health["result"].(map[string]any)["reset_mode"]) + + first, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "42", "answer": "42"}) + require.NoError(t, err) + assert.Equal(t, float64(1), first["result"].(map[string]any)["counter"]) + + // The hit starts a slow background refill. An immediate next request must not + // wait for it; it initializes one fresh single-use fallback synchronously. + second, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "42", "answer": "42"}) + require.NoError(t, err) + assert.Equal(t, float64(1), second["result"].(map[string]any)["counter"]) + + require.Eventually(t, func() bool { + health, err = dispatcher.Send(context.Background(), "healthcheck", nil) + if err != nil { + return false + } + state := health["result"].(map[string]any) + return state["prepared_ready"] == 1 && state["prepared_refills"] == uint64(1) + }, 2*time.Minute, 100*time.Millisecond) + + health, err = dispatcher.Send(context.Background(), "healthcheck", nil) + require.NoError(t, err) + state := health["result"].(map[string]any) + assert.Equal(t, uint64(1), state["prepared_hits"]) + assert.Equal(t, uint64(1), state["prepared_misses"]) + + third, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "42", "answer": "42"}) + require.NoError(t, err) + assert.Equal(t, float64(1), third["result"].(map[string]any)["counter"]) + + health, err = dispatcher.Send(context.Background(), "healthcheck", nil) + require.NoError(t, err) + assert.Equal(t, uint64(2), health["result"].(map[string]any)["prepared_hits"]) +} + +func TestAgentPythonDispatcherTimeoutDoesNotPoisonRuntime(t *testing.T) { + wasmPath := os.Getenv("AGENT_PYTHON_RUNTIME_WASM") + manifestPath := os.Getenv("AGENT_PYTHON_RUNTIME_MANIFEST") + if wasmPath == "" || manifestPath == "" { + t.Skip("AGENT_PYTHON_RUNTIME_WASM and AGENT_PYTHON_RUNTIME_MANIFEST are required") + } + + scriptPath := filepath.Join(t.TempDir(), "timeout.py") + script := ` +def dispatch(method, payload): + if method != "eval": + raise LookupError("unsupported method: " + method) + response = payload.get("response") + if response == "loop": + while True: + pass + return {"is_correct": response == payload.get("answer")} +` + require.NoError(t, os.WriteFile(scriptPath, []byte(script), 0o644)) + + dispatcher := NewAgentPythonDispatcher(Config{ + ModulePath: wasmPath, + AgentPythonManifestPath: manifestPath, + PythonScriptPath: scriptPath, + MaxInstances: 1, + MaxMemoryPages: 8192, + Timeout: 12 * time.Second, + }, zap.NewNop()) + startContext, startCancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer startCancel() + require.NoError(t, dispatcher.Start(startContext)) + t.Cleanup(func() { + shutdownContext, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _ = dispatcher.Shutdown(shutdownContext) + }) + + _, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "loop", "answer": "x"}) + require.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + + after, err := dispatcher.Send(context.Background(), "eval", map[string]any{"response": "42", "answer": "42"}) + require.NoError(t, err) + assert.Equal(t, true, after["result"].(map[string]any)["is_correct"]) +} + +func TestAgentPythonDispatcherRealLambdaFeedbackBundle(t *testing.T) { + wasmPath := os.Getenv("AGENT_PYTHON_RUNTIME_WASM") + manifestPath := os.Getenv("AGENT_PYTHON_RUNTIME_MANIFEST") + if wasmPath == "" || manifestPath == "" { + t.Skip("set AGENT_PYTHON_RUNTIME_WASM and AGENT_PYTHON_RUNTIME_MANIFEST") + } + + _, currentFile, _, ok := runtime.Caller(0) + require.True(t, ok) + repoRoot := filepath.Clean(filepath.Join(filepath.Dir(currentFile), "..", "..", "..")) + bundlePath := filepath.Join(t.TempDir(), "boilerplate.bundle.py") + command := exec.Command("python3", + filepath.Join(repoRoot, "tools", "lf-bundle-python", "lf_bundle_python.py"), + "--root", filepath.Join(repoRoot, "examples", "lambda-feedback-fixtures", "boilerplate-python"), + "--adapter-root", filepath.Join(repoRoot, "examples", "lambda-feedback-adapter"), + "--eval-entrypoint", "evaluation_function.evaluation:evaluation_function", + "--preview-entrypoint", "evaluation_function.preview:preview_function", + "--out", bundlePath, + ) + command.Env = append(os.Environ(), "PYTHONDONTWRITEBYTECODE=1") + output, err := command.CombinedOutput() + require.NoError(t, err, string(output)) + + dispatcher := NewAgentPythonDispatcher(Config{ + ModulePath: wasmPath, + AgentPythonManifestPath: manifestPath, + PythonScriptPath: bundlePath, + MaxMemoryPages: 8192, + MaxInstances: 1, + Timeout: 2 * time.Minute, + }, zap.NewNop()) + startContext, startCancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer startCancel() + require.NoError(t, dispatcher.Start(startContext)) + t.Cleanup(func() { + shutdownContext, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + _ = dispatcher.Shutdown(shutdownContext) + }) + + evalResult, err := dispatcher.Send(context.Background(), "eval", map[string]any{ + "response": "same", + "answer": "same", + "params": map[string]any{}, + }) + require.NoError(t, err) + assert.Equal(t, true, evalResult["result"].(map[string]any)["is_correct"]) + + previewResult, err := dispatcher.Send(context.Background(), "preview", map[string]any{ + "response": "x+y", + "params": map[string]any{}, + }) + require.NoError(t, err) + preview := previewResult["result"].(map[string]any)["preview"].(map[string]any) + assert.Equal(t, "x+y", preview["sympy"]) +} diff --git a/internal/execution/wasm/artifact_check.go b/internal/execution/wasm/artifact_check.go new file mode 100644 index 0000000..5d0d4a9 --- /dev/null +++ b/internal/execution/wasm/artifact_check.go @@ -0,0 +1,176 @@ +package wasm + +import ( + "context" + "fmt" + "os" + "sort" + "strings" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" +) + +// ArtifactCheckOptions selects an explicitly declared runtime ABI. The checker +// never infers a source language or evaluator framework. +type ArtifactCheckOptions struct { + Profile string + ModulePath string + ManifestPath string +} + +// ArtifactCheckReport describes objective module facts plus advisory warnings. +// Warnings do not make an artifact invalid. +type ArtifactCheckReport struct { + Profile string `json:"profile"` + Module string `json:"module"` + Exports []string `json:"exports"` + Imports []string `json:"imports"` + Warnings []string `json:"warnings,omitempty"` +} + +// CheckArtifact compiles and validates a caller-produced module against the +// explicitly selected Shimmy runtime ABI. +func CheckArtifact(ctx context.Context, options ArtifactCheckOptions) (*ArtifactCheckReport, error) { + profile := strings.ToLower(strings.TrimSpace(options.Profile)) + if profile != "generic" && profile != "python-reactor" { + return nil, fmt.Errorf("artifact checker: unsupported profile %q; use generic or python-reactor", options.Profile) + } + if options.ModulePath == "" { + return nil, fmt.Errorf("artifact checker: module path is required") + } + + var ( + moduleBytes []byte + artifact *AgentPythonArtifact + err error + ) + if profile == "python-reactor" { + artifact, err = verifyAgentPythonArtifact(options.ModulePath, options.ManifestPath) + if err != nil { + return nil, err + } + moduleBytes = artifact.WasmBytes + } else { + moduleBytes, err = os.ReadFile(options.ModulePath) + if err != nil { + return nil, fmt.Errorf("artifact checker: read module %q: %w", options.ModulePath, err) + } + } + + runtime := wazero.NewRuntime(ctx) + defer runtime.Close(ctx) //nolint:errcheck -- validation result has precedence + compiled, err := runtime.CompileModule(ctx, moduleBytes) + if err != nil { + return nil, fmt.Errorf("artifact checker: compile module: %w", err) + } + + if profile == "python-reactor" { + if err := verifyCompiledPythonReactorArtifact(compiled, artifact); err != nil { + return nil, err + } + } else if err := verifyGenericWasmArtifact(compiled); err != nil { + return nil, err + } + + report := reportCompiledArtifact(profile, options.ModulePath, compiled) + if profile == "generic" { + report.Warnings = genericWasmWarnings(compiled) + } + return report, nil +} + +func verifyGenericWasmArtifact(compiled wazero.CompiledModule) error { + if compiled == nil { + return fmt.Errorf("generic wasm: compiled module is nil") + } + + exports := compiled.ExportedFunctions() + required := map[string]pythonReactorFunctionSignature{ + "alloc": { + Params: []api.ValueType{api.ValueTypeI32}, + Results: []api.ValueType{api.ValueTypeI32}, + }, + "dispatch": { + Params: []api.ValueType{api.ValueTypeI32, api.ValueTypeI32}, + Results: []api.ValueType{api.ValueTypeI32}, + }, + } + for name, expected := range required { + definition, ok := exports[name] + if !ok { + return fmt.Errorf("generic wasm: required export %q is missing", name) + } + actual := pythonReactorFunctionSignature{Params: definition.ParamTypes(), Results: definition.ResultTypes()} + if !samePythonReactorSignature(actual, expected) { + return fmt.Errorf("generic wasm: export %q has ABI %s; expected %s", name, formatPythonReactorSignature(actual), formatPythonReactorSignature(expected)) + } + } + if _, ok := compiled.ExportedMemories()["memory"]; !ok { + return fmt.Errorf("generic wasm: required exported memory %q is missing", "memory") + } + + for _, definition := range compiled.ImportedFunctions() { + module, name, imported := definition.Import() + if imported && module != "wasi_snapshot_preview1" { + return fmt.Errorf("generic wasm: unsupported custom import %q.%q", module, name) + } + } + for _, definition := range compiled.ImportedMemories() { + module, name, imported := definition.Import() + if imported && module != "wasi_snapshot_preview1" { + return fmt.Errorf("generic wasm: unsupported custom memory import %q.%q", module, name) + } + } + return nil +} + +func reportCompiledArtifact(profile, modulePath string, compiled wazero.CompiledModule) *ArtifactCheckReport { + exports := make([]string, 0, len(compiled.ExportedFunctions())+len(compiled.ExportedMemories())) + for name := range compiled.ExportedFunctions() { + exports = append(exports, name) + } + for name := range compiled.ExportedMemories() { + exports = append(exports, name) + } + imports := make([]string, 0, len(compiled.ImportedFunctions())+len(compiled.ImportedMemories())) + for _, definition := range compiled.ImportedFunctions() { + module, name, imported := definition.Import() + if imported { + imports = append(imports, module+"."+name) + } + } + for _, definition := range compiled.ImportedMemories() { + module, name, imported := definition.Import() + if imported { + imports = append(imports, module+"."+name) + } + } + sort.Strings(exports) + sort.Strings(imports) + return &ArtifactCheckReport{Profile: profile, Module: modulePath, Exports: exports, Imports: imports} +} + +func genericWasmWarnings(compiled wazero.CompiledModule) []string { + var hasFilesystem, hasNetwork bool + for _, definition := range compiled.ImportedFunctions() { + module, name, imported := definition.Import() + if !imported || module != "wasi_snapshot_preview1" { + continue + } + if strings.HasPrefix(name, "sock_") { + hasNetwork = true + } + if strings.HasPrefix(name, "path_") || strings.HasPrefix(name, "fd_") { + hasFilesystem = true + } + } + warnings := make([]string, 0, 2) + if hasFilesystem { + warnings = append(warnings, "module imports WASI filesystem operations; behavior depends on explicitly allowed sandbox paths and may differ from native execution") + } + if hasNetwork { + warnings = append(warnings, "module imports WASI socket operations; network behavior may be unavailable or differ from native execution") + } + return warnings +} diff --git a/internal/execution/wasm/artifact_check_test.go b/internal/execution/wasm/artifact_check_test.go new file mode 100644 index 0000000..c483d2c --- /dev/null +++ b/internal/execution/wasm/artifact_check_test.go @@ -0,0 +1,52 @@ +package wasm + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCheckArtifactAcceptsGenericDispatchABI(t *testing.T) { + report, err := CheckArtifact(context.Background(), ArtifactCheckOptions{ + Profile: "generic", + ModulePath: echoModulePath(t), + }) + + require.NoError(t, err) + assert.Equal(t, "generic", report.Profile) + assert.Contains(t, report.Exports, "dispatch") + assert.NotContains(t, report.Exports, "evaluate") +} + +func TestCheckArtifactRejectsLegacyBusinessNamedExport(t *testing.T) { + data, err := os.ReadFile(echoModulePath(t)) + require.NoError(t, err) + require.Contains(t, string(data), "dispatch") + data = []byte(replaceEqualLength(string(data), "dispatch", "evaluate")) + modulePath := filepath.Join(t.TempDir(), "legacy.wasm") + require.NoError(t, os.WriteFile(modulePath, data, 0o644)) + + _, err = CheckArtifact(context.Background(), ArtifactCheckOptions{ + Profile: "generic", + ModulePath: modulePath, + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), `required export "dispatch" is missing`) +} + +func replaceEqualLength(value, old, replacement string) string { + if len(old) != len(replacement) { + panic("replacement must preserve binary length") + } + for index := 0; index+len(old) <= len(value); index++ { + if value[index:index+len(old)] == old { + return value[:index] + replacement + value[index+len(old):] + } + } + return value +} diff --git a/internal/execution/wasm/config.go b/internal/execution/wasm/config.go new file mode 100644 index 0000000..983d512 --- /dev/null +++ b/internal/execution/wasm/config.go @@ -0,0 +1,193 @@ +package wasm + +import ( + "fmt" + "os" + "strconv" + "strings" + "time" +) + +// Config holds the configuration for the WASM execution backend. +// +// Configuration is read from environment variables via koanf (the same +// mechanism used by the rest of shimmy). The "conf" struct tags map to +// the koanf key names derived from the FUNCTION_* env-var prefix. +type Config struct { + // ModulePath is the path to the .wasm file to load. + // Populated from FUNCTION_COMMAND (the command field re-used as the + // .wasm file path when FUNCTION_INTERFACE=wasm). + ModulePath string `conf:"cmd"` + + // AgentPythonManifestPath binds the clean Python reactor artifact to its + // producer manifest. When empty, the agent-python dispatcher reads + // manifest.json next to ModulePath. FUNCTION_WASM_MANIFEST overrides it. + AgentPythonManifestPath string `conf:"wasm_manifest"` + + // MaxInstances is the maximum number of concurrently active module + // instances. When the pool is exhausted requests block until a slot is + // available. Defaults to runtime.NumCPU() when <= 0. + // Populated from FUNCTION_MAX_PROCS / max_workers. + MaxInstances int `conf:"max_workers"` + + // Timeout is the per-request deadline passed to the WASM call. + // Populated from FUNCTION_WORKER_SEND_TIMEOUT / send.timeout. + Timeout time.Duration `conf:"timeout"` + + // --- Sandbox limits --- + + // MaxMemoryPages limits WASM linear memory (1 page = 64KB). + // Default: 256 pages = 16MB. 0 means use module's own max. + MaxMemoryPages uint32 `conf:"wasm_max_memory_pages"` + + // AllowedPaths is a list of host paths the module may read (read-only). + // Empty means no filesystem access at all. + AllowedPaths []string `conf:"wasm_allowed_paths"` + + // AllowedEnv is a list of env var names the module may read. + // Empty means no env vars exposed. + AllowedEnv []string `conf:"wasm_allowed_env"` + + // PythonScriptPath is the host path to the trusted Python evaluation script. + // Used by Python Reactor and the independent resident Python compatibility path. + // Python Reactor scripts must define dispatch(method, payload). + PythonScriptPath string `conf:"wasm_python_script"` + + // PythonPreloadMode controls whether Agent Python passes the trusted evaluator + // through runtime_prepare. "evaluator" is the default; "off" executes the + // trusted script in each fresh request namespace. + PythonPreloadMode string `conf:"wasm_python_preload"` + + // PythonLifecycle selects whether Agent Python modules are initialized for + // every request, consumed once from a prepared pool, or restored to their + // prepared linear-memory snapshot and reused. + PythonLifecycle string `conf:"wasm_python_lifecycle"` + + // PythonPreparedCapacity bounds never-served candidates retained by the + // single-use lifecycle. The current numpy-core artifact retains 128 MiB of + // Guest linear memory per candidate, so this surface is deliberately small. + PythonPreparedCapacity int `conf:"wasm_python_prepared_capacity"` + + // PythonSnapshotHeadroomBytes reserves allocator capacity before Take so + // normal requests do not immediately grow memory beyond a restorable baseline. + PythonSnapshotHeadroomBytes uint64 `conf:"wasm_python_snapshot_headroom_bytes"` + + // CompileCacheDir, if non-empty, enables wazero's on-disk compilation cache. + // Set via FUNCTION_WASM_COMPILE_CACHE env var. Shared across all runners and + // processes that point at the same directory, making cold starts much faster + // after the first compile. + CompileCacheDir string `conf:"wasm_compile_cache"` + + // AgentPythonObserver receives optional phase evidence. Callbacks may be + // concurrent during refill and must return promptly. It is never populated + // from operator configuration. + AgentPythonObserver func(AgentPythonPhaseEvent) `conf:"-"` +} + +// applyDefaults fills in zero-value fields with sensible defaults. +func (c *Config) applyDefaults() { + if c.Timeout == 0 { + c.Timeout = 30 * time.Second + } + if c.MaxMemoryPages == 0 { + c.MaxMemoryPages = 256 // 16 MB + } + if c.PythonPreloadMode == "" { + c.PythonPreloadMode = "evaluator" + } +} + +func (c *Config) validatePythonPreloadMode() error { + switch c.PythonPreloadMode { + case "evaluator", "off": + return nil + default: + return fmt.Errorf("python preload mode %q is invalid; use \"evaluator\" or \"off\"", c.PythonPreloadMode) + } +} + +func (c *Config) applyAgentPythonDefaults() { + if c.PythonLifecycle == "" { + c.PythonLifecycle = "snapshot" + } + if c.PythonPreparedCapacity == 0 { + c.PythonPreparedCapacity = 1 + } + if c.PythonSnapshotHeadroomBytes == 0 { + c.PythonSnapshotHeadroomBytes = 8 * 1024 * 1024 + } + +} + +func (c *Config) validateAgentPythonLifecycle() error { + switch c.PythonLifecycle { + case "fresh", "single-use", "snapshot": + default: + return fmt.Errorf("agent Python lifecycle %q is invalid; use \"fresh\", \"single-use\", or \"snapshot\"", c.PythonLifecycle) + } + if c.PythonPreparedCapacity < 1 || c.PythonPreparedCapacity > 4 { + return fmt.Errorf("agent Python prepared capacity %d is outside the supported range 1..4", c.PythonPreparedCapacity) + } + if c.MaxInstances > 4 { + return fmt.Errorf("agent Python max instances %d exceeds the supported limit 4", c.MaxInstances) + } + return nil +} + +// applyEnv reads sandbox fields from FUNCTION_WASM_* environment variables. +// This allows operators to configure sandbox limits without threading them +// through the full koanf config chain. +func (c *Config) applyEnv() { + // FUNCTION_WASM_MODULE overrides FUNCTION_COMMAND as the .wasm file path. + if v := os.Getenv("FUNCTION_WASM_MODULE"); v != "" { + c.ModulePath = v + } + if v := os.Getenv("FUNCTION_WASM_MANIFEST"); v != "" { + c.AgentPythonManifestPath = v + } + if v := os.Getenv("FUNCTION_WASM_MAX_MEMORY_PAGES"); v != "" { + if n, err := strconv.ParseUint(v, 10, 32); err == nil { + c.MaxMemoryPages = uint32(n) + } + } + if v := os.Getenv("FUNCTION_WASM_ALLOWED_PATHS"); v != "" { + c.AllowedPaths = splitNonEmpty(v, ",") + } + if v := os.Getenv("FUNCTION_WASM_ALLOWED_ENV"); v != "" { + c.AllowedEnv = splitNonEmpty(v, ",") + } + + if v := os.Getenv("FUNCTION_WASM_PYTHON_SCRIPT"); v != "" { + c.PythonScriptPath = v + } + if v := os.Getenv("FUNCTION_WASM_PYTHON_PRELOAD"); v != "" { + c.PythonPreloadMode = v + } + if v := os.Getenv("FUNCTION_WASM_PYTHON_LIFECYCLE"); v != "" { + c.PythonLifecycle = strings.TrimSpace(v) + } + if v := os.Getenv("FUNCTION_WASM_PYTHON_PREPARED_CAPACITY"); v != "" { + if n, err := strconv.Atoi(v); err == nil { + c.PythonPreparedCapacity = n + } + } + if v := os.Getenv("FUNCTION_WASM_PYTHON_SNAPSHOT_HEADROOM_BYTES"); v != "" { + if n, err := strconv.ParseUint(v, 10, 64); err == nil { + c.PythonSnapshotHeadroomBytes = n + } + } + + if v := os.Getenv("FUNCTION_WASM_COMPILE_CACHE"); v != "" { + c.CompileCacheDir = v + } +} + +func splitNonEmpty(s, sep string) []string { + var out []string + for _, p := range strings.Split(s, sep) { + if t := strings.TrimSpace(p); t != "" { + out = append(out, t) + } + } + return out +} diff --git a/internal/execution/wasm/dispatcher.go b/internal/execution/wasm/dispatcher.go new file mode 100644 index 0000000..0cd85e9 --- /dev/null +++ b/internal/execution/wasm/dispatcher.go @@ -0,0 +1,378 @@ +package wasm + +import ( + "context" + "errors" + "fmt" + "os" + "runtime" + "sync" + "time" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1" + "go.uber.org/zap" + + "github.com/lambda-feedback/shimmy/internal/execution/dispatcher" +) + +// ErrDispatcherClosed is returned by Send after the dispatcher has begun (or +// completed) Shutdown. Callers should treat it as a terminal error. +var ErrDispatcherClosed = fmt.Errorf("wasm: dispatcher is shut down") + +// Dispatcher implements [dispatcher.Dispatcher] for the WASM execution +// backend. It compiles the .wasm module once at startup, then maintains a pool +// of pre-initialised [wasmSupervisor] instances (one compiled module, N module +// instances). Requests are dispatched by acquiring a supervisor from the pool, +// calling its Send, and returning it to the pool. +type Dispatcher struct { + cfg Config + rt wazero.Runtime + compiled wazero.CompiledModule + modCfg wazero.ModuleConfig + pool chan *wasmSupervisor + log *zap.Logger + + // mu protects closed and serialises the closed/push transitions so that a + // replacement supervisor cannot land in the pool after Shutdown has begun + // draining it. + mu sync.Mutex + closed bool + // closedCh is closed atomically with closed=true (under mu) by Shutdown. + // Send selects on it to (a) unblock a pool acquire that is racing Shutdown + // and (b) avoid waiting on an empty pool that Shutdown is about to drain. + closedCh chan struct{} + // pending tracks BOTH in-flight Sends (Add in tryBeginSend, Done via Send's + // defer) AND background goroutines spawned during a Send (replacement + // spawns, discard-shutdowns). Shutdown waits on it before draining the + // pool / closing the runtime. + // + // Invariant: every pending.Add is either (a) made under d.mu after + // observing !closed, or (b) made by code that is itself holding a pending + // count (e.g. discardAsync called from inside Send). This keeps Add from + // racing Shutdown's Wait — if closed is already set, branch (a) skips the + // Add and falls back to a synchronous close; in branch (b) Shutdown is + // guaranteed to still be blocked at Wait on the caller's count. + pending sync.WaitGroup +} + +var _ dispatcher.Dispatcher = (*Dispatcher)(nil) + +// NewDispatcher creates a new WASM dispatcher. Compilation and pool +// initialisation happen in Start. +func NewDispatcher(cfg Config, log *zap.Logger) *Dispatcher { + return &Dispatcher{ + cfg: cfg, + log: log.Named("dispatcher_wasm"), + closedCh: make(chan struct{}), + } +} + +// tryBeginSend atomically checks the closed flag and increments pending. It +// returns false if Shutdown has begun (caller must abort with +// ErrDispatcherClosed); on true the caller MUST call pending.Done exactly +// once when finished. Holding a pending count across the entire Send keeps +// Shutdown's Wait blocked while the Send is mid-flight, which is what lets +// discardAsync inside Send safely Add to pending without racing Wait. +func (d *Dispatcher) tryBeginSend() bool { + d.mu.Lock() + defer d.mu.Unlock() + if d.closed { + return false + } + d.pending.Add(1) + return true +} + +// Start reads and compiles the .wasm file, sets up WASI host functions, and +// pre-warms the supervisor pool. +func (d *Dispatcher) Start(ctx context.Context) error { + // Pick up sandbox overrides from FUNCTION_WASM_* env vars (including + // FUNCTION_WASM_MODULE as an alternative to FUNCTION_COMMAND), then apply + // sensible defaults for any fields still at their zero values. + d.cfg.applyEnv() + d.cfg.applyDefaults() + + if d.cfg.ModulePath == "" { + return fmt.Errorf("wasm: ModulePath must be set (FUNCTION_COMMAND or FUNCTION_WASM_MODULE)") + } + + maxInstances := d.cfg.MaxInstances + if maxInstances <= 0 { + maxInstances = runtime.NumCPU() + } + + d.log.Info("starting wasm dispatcher", + zap.String("module", d.cfg.ModulePath), + zap.Int("max_instances", maxInstances), + zap.Uint32("max_memory_pages", d.cfg.MaxMemoryPages), + zap.Duration("timeout", d.cfg.Timeout), + ) + + // Read the .wasm bytes from disk. + wasmBytes, err := os.ReadFile(d.cfg.ModulePath) + if err != nil { + return fmt.Errorf("wasm: read module file %q: %w", d.cfg.ModulePath, err) + } + + // Build the runtime config with memory limit and context-done interruption. + rtCfg := wazero.NewRuntimeConfig(). + // WithCloseOnContextDone causes wazero to interrupt a running WASM module + // when the call context is cancelled or times out, preventing goroutine leaks. + WithCloseOnContextDone(true) + if d.cfg.MaxMemoryPages > 0 { + rtCfg = rtCfg.WithMemoryLimitPages(d.cfg.MaxMemoryPages) + } + + // Wire in on-disk compilation cache when configured. + if d.cfg.CompileCacheDir != "" { + cache, err := wazero.NewCompilationCacheWithDir(d.cfg.CompileCacheDir) + if err != nil { + d.log.Warn("failed to create wazero compilation cache, continuing without cache", + zap.String("dir", d.cfg.CompileCacheDir), + zap.Error(err)) + } else { + rtCfg = rtCfg.WithCompilationCache(cache) + d.log.Info("wazero compilation cache enabled", zap.String("dir", d.cfg.CompileCacheDir)) + } + } + + // Create a single wazero runtime shared by all instances. + rt := wazero.NewRuntimeWithConfig(ctx, rtCfg) + d.rt = rt + + // Instantiate WASI host functions. Most evaluation functions will need at + // least minimal WASI support (e.g. for memory allocation helpers compiled + // from C/Rust/TinyGo). + if _, err := wasi_snapshot_preview1.Instantiate(ctx, rt); err != nil { + _ = rt.Close(ctx) + return fmt.Errorf("wasm: instantiate wasi: %w", err) + } + + // Compile the module once; all instances share the compiled code. + compiled, err := rt.CompileModule(ctx, wasmBytes) + if err != nil { + _ = rt.Close(ctx) + return fmt.Errorf("wasm: compile module: %w", err) + } + d.compiled = compiled + + // Build a locked-down ModuleConfig: no filesystem, no env vars, no + // stdin/stdout/stderr, no args. Only allow nanosleep and wall/mono clocks + // which the Go runtime needs. + modCfg := wazero.NewModuleConfig(). + WithName(""). + WithSysNanosleep(). + WithSysWalltime(). + WithSysNanotime() + + // Filesystem: mount allowed paths read-only; no access by default. + fsCfg := wazero.NewFSConfig() + for _, p := range d.cfg.AllowedPaths { + fsCfg = fsCfg.WithReadOnlyDirMount(p, p) + } + modCfg = modCfg.WithFSConfig(fsCfg) + + // Env vars: expose only explicitly whitelisted variables. + for _, key := range d.cfg.AllowedEnv { + if val, ok := os.LookupEnv(key); ok { + modCfg = modCfg.WithEnv(key, val) + } + } + d.modCfg = modCfg + + // Build the pool. + d.pool = make(chan *wasmSupervisor, maxInstances) + + for i := 0; i < maxInstances; i++ { + sv := newWasmSupervisor(rt, compiled, modCfg, d.cfg.Timeout, d.log) + + if err := sv.Start(ctx); err != nil { + // Clean up already-started supervisors. + drainPool(ctx, d.pool, d.log) + _ = rt.Close(ctx) + d.rt = nil + return fmt.Errorf("wasm: start instance %d: %w", i, err) + } + + d.pool <- sv + } + + d.log.Info("wasm dispatcher ready", zap.Int("instances", maxInstances)) + + return nil +} + +// Send acquires a supervisor from the pool, dispatches the request, and +// returns the supervisor to the pool. +func (d *Dispatcher) Send( + ctx context.Context, + method string, + data map[string]any, +) (map[string]any, error) { + if !d.tryBeginSend() { + return nil, ErrDispatcherClosed + } + defer d.pending.Done() + + // Acquire a supervisor, honouring the caller's context AND the shutdown + // signal so we never block forever on a drained pool. + var sv *wasmSupervisor + select { + case sv = <-d.pool: + case <-d.closedCh: + return nil, ErrDispatcherClosed + case <-ctx.Done(): + return nil, fmt.Errorf("wasm: acquire instance: %w", ctx.Err()) + } + + result, err := sv.Send(ctx, method, data) + + // Return the supervisor to the pool only if it is healthy. + // If the snapshot restore failed inside Send, sv.healthy is false and the + // supervisor's state is undefined — discard it and spawn a replacement so + // pool capacity is eventually restored. + if sv.IsHealthy() { + d.returnOrDiscard(sv) + } else { + d.log.Warn("wasm supervisor unhealthy after request — dropping from pool, spawning replacement") + d.discardAsync(sv) + d.spawnReplacementAsync() + } + + if err != nil { + return nil, fmt.Errorf("wasm: send: %w", err) + } + + return result, nil +} + +// returnOrDiscard puts a healthy supervisor back in the pool unless Shutdown +// has begun, in which case the supervisor is closed asynchronously so it does +// not leak past a drained pool. +// +// Must be called from a goroutine that already holds a pending count (i.e. +// from inside Send) so that the Add issued by discardAsync is guaranteed to +// happen before Shutdown's pending.Wait can return. +func (d *Dispatcher) returnOrDiscard(sv *wasmSupervisor) { + d.mu.Lock() + if d.closed { + d.mu.Unlock() + d.discardAsync(sv) + return + } + // Push under the lock so it interleaves correctly with Shutdown's + // closed=true → drainPool sequence: either we push before closed is set + // (drainPool sees the supervisor) or we discard via the branch above. + d.pool <- sv + d.mu.Unlock() +} + +// discardAsync closes a discarded supervisor in the background and tracks it +// via the pending WaitGroup so Shutdown can wait for the close to complete +// before tearing down the runtime. +// +// Must be called from a goroutine that already holds a pending count +// (Send, via tryBeginSend). That invariant keeps Shutdown.pending.Wait +// blocked across this Add, eliminating the Add-after-Wait race. +func (d *Dispatcher) discardAsync(sv *wasmSupervisor) { + d.pending.Add(1) + go func() { + defer d.pending.Done() + _ = sv.Shutdown(context.Background()) + }() +} + +// spawnReplacementAsync kicks off spawnOne in a background goroutine, but only +// if the dispatcher is still open. If Shutdown has begun, no replacement is +// scheduled. Tracked via the pending WaitGroup. +func (d *Dispatcher) spawnReplacementAsync() { + d.mu.Lock() + if d.closed { + d.mu.Unlock() + return + } + d.pending.Add(1) + d.mu.Unlock() + go d.spawnOne() +} + +// spawnOne initialises a fresh wasmSupervisor and adds it to the pool. +// Called in a goroutine when an unhealthy supervisor is discarded so that +// pool capacity is eventually restored. Failures are logged but not fatal. +// +// If Shutdown begins while Start is running, the freshly initialised +// supervisor is closed immediately rather than inserted into the drained pool. +func (d *Dispatcher) spawnOne() { + defer d.pending.Done() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + + d.log.Info("wasm: initialising replacement supervisor") + sv := newWasmSupervisor(d.rt, d.compiled, d.modCfg, d.cfg.Timeout, d.log) + if err := sv.Start(ctx); err != nil { + d.log.Error("wasm: replacement supervisor init failed", zap.Error(err)) + return + } + + d.mu.Lock() + if d.closed { + d.mu.Unlock() + d.log.Info("wasm: replacement supervisor born during shutdown — closing immediately") + _ = sv.Shutdown(context.Background()) + return + } + d.pool <- sv + d.mu.Unlock() + d.log.Info("wasm: replacement supervisor ready") +} + +// Shutdown closes all module instances and the wazero runtime. Idempotent. +func (d *Dispatcher) Shutdown(ctx context.Context) error { + d.mu.Lock() + if d.closed { + d.mu.Unlock() + return nil + } + d.closed = true + // Close the channel under mu so the closed=true / close(closedCh) pair is + // atomic with respect to tryBeginSend: any Send that observes !closed has + // also pending.Add'd before Shutdown can reach pending.Wait. + close(d.closedCh) + d.mu.Unlock() + + d.log.Debug("shutting down wasm dispatcher") + + // Wait for in-flight Sends AND any background goroutines (replacement + // spawns / discard shutdowns) to finish so that no late-created supervisor + // lands in the pool after the drain below, no module is mid-Close while we + // close the runtime, and no Send is running against the wazero runtime + // when we tear it down. + d.pending.Wait() + + // Non-blocking drain: after pending.Wait, no spawn or returnOrDiscard + // goroutine will push to the pool, so we just close everything currently + // buffered. (drainPool's blocking-for-cap-items semantics would deadlock + // here when spawnOne took the closed-shortcut and never pushed.) + for { + select { + case sv := <-d.pool: + if err := sv.Shutdown(ctx); err != nil { + d.log.Warn("error shutting down pooled supervisor", zap.Error(err)) + } + default: + goto drained + } + } +drained: + + var closeErr error + if d.rt != nil { + if err := d.rt.Close(ctx); err != nil { + closeErr = errors.Join(closeErr, fmt.Errorf("wasm: close runtime: %w", err)) + } + d.rt = nil + } + return closeErr +} diff --git a/internal/execution/wasm/dispatcher_test.go b/internal/execution/wasm/dispatcher_test.go new file mode 100644 index 0000000..c0b5939 --- /dev/null +++ b/internal/execution/wasm/dispatcher_test.go @@ -0,0 +1,465 @@ +package wasm + +import ( + "context" + "errors" + "path/filepath" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tetratelabs/wazero" + "go.uber.org/zap" +) + +// echoModulePath returns the absolute path to the pre-compiled echo.wasm test +// fixture. The fixture is a minimal guest module that always returns +// {"ok":true} regardless of the request, which lets us test the host-side Go +// code (alloc call, memory write, dispatch call, length-prefix parsing, JSON +// unmarshal) without implementing a full language runtime in WAT. +func echoModulePath(t *testing.T) string { + t.Helper() + // __file__ is not available in Go, but runtime.Caller gives us the source + // file path so we can derive testdata/ relative to the test file. + _, filename, _, ok := runtime.Caller(0) + require.True(t, ok, "runtime.Caller failed") + return filepath.Join(filepath.Dir(filename), "testdata", "echo.wasm") +} + +// newTestLogger returns a no-op zap logger suitable for unit tests. +func newTestLogger(t *testing.T) *zap.Logger { + t.Helper() + log, err := zap.NewDevelopment() + require.NoError(t, err) + return log +} + +// newEchoDispatcher creates a Dispatcher backed by the echo fixture and starts +// it. The caller is responsible for calling Shutdown. +func newEchoDispatcher(t *testing.T, maxInstances int) *Dispatcher { + t.Helper() + cfg := Config{ + ModulePath: echoModulePath(t), + MaxInstances: maxInstances, + Timeout: 5 * time.Second, + } + d := NewDispatcher(cfg, newTestLogger(t)) + require.NoError(t, d.Start(context.Background()), "dispatcher start") + return d +} + +// TestDispatcher_StartStop verifies that a Dispatcher can be started and shut +// down cleanly without any interaction in between. +func TestDispatcher_StartStop(t *testing.T) { + d := newEchoDispatcher(t, 1) + err := d.Shutdown(context.Background()) + assert.NoError(t, err) +} + +// TestDispatcher_StartStop_MultipleInstances verifies start/stop with the +// default pool size (NumCPU). +func TestDispatcher_StartStop_MultipleInstances(t *testing.T) { + d := newEchoDispatcher(t, runtime.NumCPU()) + err := d.Shutdown(context.Background()) + assert.NoError(t, err) +} + +// TestDispatcher_Send_BasicResponse sends a single request and checks that the +// echo module returns {"ok":true}. +func TestDispatcher_Send_BasicResponse(t *testing.T) { + d := newEchoDispatcher(t, 1) + t.Cleanup(func() { _ = d.Shutdown(context.Background()) }) + + result, err := d.Send(context.Background(), "test", map[string]any{"hello": "world"}) + require.NoError(t, err) + require.NotNil(t, result) + + ok, exists := result["ok"] + assert.True(t, exists, "response should contain 'ok' key") + assert.Equal(t, true, ok, "response 'ok' should be true") +} + +// TestDispatcher_Send_EmptyParams verifies that Send works with nil params. +func TestDispatcher_Send_EmptyParams(t *testing.T) { + d := newEchoDispatcher(t, 1) + t.Cleanup(func() { _ = d.Shutdown(context.Background()) }) + + result, err := d.Send(context.Background(), "noop", nil) + require.NoError(t, err) + require.NotNil(t, result) + assert.Equal(t, true, result["ok"]) +} + +// TestDispatcher_Send_Concurrent sends 10 concurrent requests using a pool of +// 3 instances and verifies that all succeed. +func TestDispatcher_Send_Concurrent(t *testing.T) { + const ( + numWorkers = 10 + numRequests = 20 + poolSize = 3 + ) + + d := newEchoDispatcher(t, poolSize) + t.Cleanup(func() { _ = d.Shutdown(context.Background()) }) + + type result struct { + res map[string]any + err error + } + + results := make([]result, numRequests) + var wg sync.WaitGroup + wg.Add(numRequests) + + sem := make(chan struct{}, numWorkers) + for i := range numRequests { + sem <- struct{}{} + go func(i int) { + defer wg.Done() + defer func() { <-sem }() + res, err := d.Send(context.Background(), "eval", map[string]any{"i": i}) + results[i] = result{res, err} + }(i) + } + + wg.Wait() + + for i, r := range results { + require.NoError(t, r.err, "request %d failed", i) + require.NotNil(t, r.res, "request %d returned nil result", i) + assert.Equal(t, true, r.res["ok"], "request %d: unexpected result", i) + } +} + +// TestDispatcher_Send_AfterShutdown checks that Send after Shutdown returns +// ErrDispatcherClosed immediately, rather than blocking on the drained pool +// until the caller's context expires. +func TestDispatcher_Send_AfterShutdown(t *testing.T) { + d := newEchoDispatcher(t, 1) + require.NoError(t, d.Shutdown(context.Background())) + + _, err := d.Send(context.Background(), "test", nil) + assert.ErrorIs(t, err, ErrDispatcherClosed, "Send after Shutdown must return ErrDispatcherClosed") +} + +// TestDispatcher_Shutdown_Idempotent verifies that calling Shutdown twice does +// not return an error or double-close the runtime. +func TestDispatcher_Shutdown_Idempotent(t *testing.T) { + d := newEchoDispatcher(t, 1) + require.NoError(t, d.Shutdown(context.Background())) + require.NoError(t, d.Shutdown(context.Background()), "second Shutdown must be a no-op") +} + +// TestDispatcher_ReplacementDuringShutdown exercises the race where Send has +// just discarded an unhealthy supervisor and scheduled a replacement spawn +// while Shutdown begins. The replacement spawn must NOT insert a supervisor +// into a drained pool, and Shutdown must wait for the spawn goroutine to +// finish before closing the runtime (otherwise the late supervisor would +// reference a torn-down wazero.Runtime). +func TestDispatcher_ReplacementDuringShutdown(t *testing.T) { + d := newEchoDispatcher(t, 1) + + // Consume the only supervisor in the pool to mimic an in-flight Send. + sv := <-d.pool + + // Simulate Send's unhealthy-path bookkeeping: schedule the discard close + // of the bad supervisor and the spawn of a replacement. + d.discardAsync(sv) + d.spawnReplacementAsync() + + // Shutdown races with the spawn. It must wait for pending background work + // (via d.pending.Wait) before draining the pool and closing the runtime. + require.NoError(t, d.Shutdown(context.Background())) + + // After Shutdown the pool must be empty: any replacement that finished + // initialising during the race window was closed by spawnOne's + // closed-guard rather than inserted. + assert.Equal(t, 0, len(d.pool), "drained pool must be empty after Shutdown") + + // Send after Shutdown returns ErrDispatcherClosed promptly. + _, err := d.Send(context.Background(), "test", nil) + assert.ErrorIs(t, err, ErrDispatcherClosed) +} + +// TestDispatcher_Shutdown_WaitsForInFlightSends drives the original race the +// lifecycle patch is meant to fix: many concurrent Sends are issued while +// Shutdown runs partway through. Without the in-flight tracking, Shutdown +// could close the wazero runtime out from under a live Send (use-after-close), +// or returnOrDiscard/discardAsync could call pending.Add after Shutdown's +// pending.Wait already returned. Both surfaces are caught by -race or by an +// outright panic. +// +// Acceptance: every Send either succeeds or returns ErrDispatcherClosed, never +// any other error; Shutdown returns nil; no panic. +func TestDispatcher_Shutdown_WaitsForInFlightSends(t *testing.T) { + d := newEchoDispatcher(t, runtime.NumCPU()) + + const numWorkers = 128 + var ( + wg sync.WaitGroup + successes atomic.Int64 + closedExits atomic.Int64 + unexpected atomic.Int64 + ) + wg.Add(numWorkers) + + start := make(chan struct{}) + for i := 0; i < numWorkers; i++ { + go func() { + defer wg.Done() + <-start + for j := 0; j < 5; j++ { + _, err := d.Send(context.Background(), "eval", map[string]any{"j": j}) + switch { + case err == nil: + successes.Add(1) + case errors.Is(err, ErrDispatcherClosed): + closedExits.Add(1) + return // dispatcher is gone; stop hammering + default: + unexpected.Add(1) + t.Errorf("unexpected error: %v", err) + return + } + } + }() + } + + close(start) + // Give some Sends a chance to begin. + time.Sleep(5 * time.Millisecond) + + require.NoError(t, d.Shutdown(context.Background())) + wg.Wait() + + assert.Zero(t, unexpected.Load(), "no Send may return a non-closed error") + // Post-shutdown Send must return ErrDispatcherClosed promptly (not block). + postCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := d.Send(postCtx, "eval", nil) + assert.ErrorIs(t, err, ErrDispatcherClosed) + t.Logf("successes=%d closed_exits=%d", successes.Load(), closedExits.Load()) +} + +// TestDispatcher_Shutdown_UnblocksBlockedSend covers the second race called +// out in the patch: a Send that passed tryBeginSend but finds the pool empty +// (all supervisors are in-use or have been drained by a racing Shutdown). +// Without selecting on closedCh, the Send would block on the empty pool until +// the caller's context expired. With the patch it must return +// ErrDispatcherClosed as soon as Shutdown begins. +func TestDispatcher_Shutdown_UnblocksBlockedSend(t *testing.T) { + d := newEchoDispatcher(t, 1) + // Empty the pool so a Send is forced to block on acquire. + sv := <-d.pool + + type sendResult struct { + err error + } + res := make(chan sendResult, 1) + go func() { + _, err := d.Send(context.Background(), "eval", nil) + res <- sendResult{err: err} + }() + + // Let Send reach the empty-pool select. + time.Sleep(50 * time.Millisecond) + + // Put sv back so the dispatcher's drain has something to clean up + // (otherwise Shutdown sees an empty pool, which is also fine). + d.pool <- sv + + require.NoError(t, d.Shutdown(context.Background())) + + select { + case r := <-res: + // Either the Send got the supervisor before Shutdown drained it + // (succeeded), or Shutdown's closedCh fired first. + if r.err != nil { + assert.ErrorIs(t, r.err, ErrDispatcherClosed) + } + case <-time.After(2 * time.Second): + t.Fatal("Send did not return after Shutdown — closedCh select missing") + } +} + +// TestDispatcher_SpawnReplacementAsync_NoopAfterShutdown asserts that calling +// spawnReplacementAsync on a closed dispatcher is a no-op: it must not +// increment pending and must not launch a goroutine that touches the closed +// runtime. +func TestDispatcher_SpawnReplacementAsync_NoopAfterShutdown(t *testing.T) { + d := newEchoDispatcher(t, 1) + require.NoError(t, d.Shutdown(context.Background())) + + // Should return immediately without scheduling work. + d.spawnReplacementAsync() + + // Wait briefly with a deadline — pending.Wait would block forever if the + // no-op guard regressed and a goroutine were leaked with a stale runtime. + done := make(chan struct{}) + go func() { + d.pending.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("pending.Wait did not return — spawn goroutine leaked after Shutdown") + } +} + +// TestDispatcher_MissingModule checks that Start fails when ModulePath does +// not point to a valid file. +func TestDispatcher_MissingModule(t *testing.T) { + cfg := Config{ + ModulePath: "/nonexistent/path/module.wasm", + MaxInstances: 1, + } + d := NewDispatcher(cfg, newTestLogger(t)) + err := d.Start(context.Background()) + assert.Error(t, err, "Start with missing module should fail") +} + +// TestSupervisor_MemoryRestored sends two sequential requests through the same +// supervisor and verifies that both succeed with the same response. This +// exercises the snapshot/restore cycle: after the first dispatch the bump +// allocator's heap_top is advanced, but restoreSnapshot rewinds memory so the +// second call starts from the exact same state. +func TestSupervisor_MemoryRestored(t *testing.T) { + // Use a pool of exactly 1 so both sends use the same supervisor instance. + d := newEchoDispatcher(t, 1) + t.Cleanup(func() { _ = d.Shutdown(context.Background()) }) + + ctx := context.Background() + + r1, err := d.Send(ctx, "first", map[string]any{"seq": 1}) + require.NoError(t, err) + require.NotNil(t, r1) + + r2, err := d.Send(ctx, "second", map[string]any{"seq": 2}) + require.NoError(t, err) + require.NotNil(t, r2) + + // Both responses must be identical {"ok":true}. + assert.Equal(t, r1, r2, "responses must be equal, proving memory was restored between calls") + assert.Equal(t, true, r1["ok"]) + assert.Equal(t, true, r2["ok"]) +} + +// TestSupervisor_MemoryRestored_ManyTimes exercises many sequential calls +// through a single-instance pool to ensure the snapshot/restore cycle is +// stable over repeated invocations. +func TestSupervisor_MemoryRestored_ManyTimes(t *testing.T) { + d := newEchoDispatcher(t, 1) + t.Cleanup(func() { _ = d.Shutdown(context.Background()) }) + + ctx := context.Background() + const iters = 50 + + for i := range iters { + res, err := d.Send(ctx, "loop", map[string]any{"i": i}) + require.NoError(t, err, "iteration %d", i) + assert.Equal(t, true, res["ok"], "iteration %d", i) + } +} + +// TestSupervisor_Start_Idempotent verifies that calling Start twice on the +// same supervisor does not error (the second call is a no-op). +func TestSupervisor_Start_Idempotent(t *testing.T) { + ctx := context.Background() + log := newTestLogger(t) + + wasmBytes := echoWasmBytes(t) + + rt, compiled := compileEchoModule(t, ctx, wasmBytes) + t.Cleanup(func() { _ = rt.Close(ctx) }) + + sv := newWasmSupervisor(rt, compiled, wazero.NewModuleConfig().WithName(""), 5*time.Second, log) + require.NoError(t, sv.Start(ctx)) + require.NoError(t, sv.Start(ctx), "second Start must be a no-op") + require.NoError(t, sv.Shutdown(ctx)) +} + +// TestSupervisor_Send_NotStarted checks that Send before Start returns an +// error. +func TestSupervisor_Send_NotStarted(t *testing.T) { + ctx := context.Background() + log := newTestLogger(t) + + wasmBytes := echoWasmBytes(t) + rt, compiled := compileEchoModule(t, ctx, wasmBytes) + t.Cleanup(func() { _ = rt.Close(ctx) }) + + sv := newWasmSupervisor(rt, compiled, wazero.NewModuleConfig().WithName(""), 5*time.Second, log) + // Do NOT call sv.Start. + + _, err := sv.Send(ctx, "test", nil) + assert.Error(t, err, "Send without Start should return an error") +} + +// TestSupervisor_Send_MemoryGrowDetected is the regression test for +// memory.grow snapshot isolation: if the guest expands linear memory during a +// request, the supervisor must (a) detect the growth, (b) zero the grown tail +// so the next request cannot read leaked guest data, (c) surface +// ErrMemoryGrew, and (d) mark itself unhealthy so the dispatcher discards it +// instead of returning it to the pool. +// +// The echo fixture itself never grows memory, so we simulate a request that +// did by growing the module's memory from host code (between Start and Send) +// and writing a recognisable poison pattern into the new pages. After Send +// runs, restoreSnapshot observes mem.Size() > snapshotSize and must trip the +// defensive path. +func TestSupervisor_Send_MemoryGrowDetected(t *testing.T) { + ctx := context.Background() + log := newTestLogger(t) + + wasmBytes := echoWasmBytes(t) + rt, compiled := compileEchoModule(t, ctx, wasmBytes) + t.Cleanup(func() { _ = rt.Close(ctx) }) + + sv := newWasmSupervisor(rt, compiled, wazero.NewModuleConfig().WithName(""), 5*time.Second, log) + require.NoError(t, sv.Start(ctx)) + t.Cleanup(func() { _ = sv.Shutdown(ctx) }) + + require.True(t, sv.IsHealthy(), "supervisor should be healthy after Start") + + // Capture the snapshot size, then grow memory by 1 page (64 KiB) and + // poison the new pages. This simulates a guest that called memory.grow + // during execution and wrote sensitive data into the new pages. + mem := sv.mod.Memory() + require.NotNil(t, mem) + origSize := mem.Size() + require.Equal(t, origSize, sv.snapshotSize, "snapshotSize must be recorded at Take time") + + prevPages, ok := mem.Grow(1) + require.True(t, ok, "memory.Grow must succeed (echo fixture has no max)") + require.Equal(t, origSize/(64*1024), prevPages) + + grownSize := mem.Size() + require.Greater(t, grownSize, origSize, "memory must have grown") + + poison := make([]byte, grownSize-origSize) + for i := range poison { + poison[i] = 0xAB + } + require.True(t, mem.Write(origSize, poison), "poison tail") + + // Issue a request. The echo guest doesn't itself grow memory, but Send's + // post-call restoreSnapshot will observe the host-injected growth and + // trip the defensive path. + _, err := sv.Send(ctx, "test", map[string]any{"hello": "world"}) + require.Error(t, err, "Send must return the restore error") + assert.ErrorIs(t, err, ErrMemoryGrew, "error must wrap ErrMemoryGrew") + + assert.False(t, sv.IsHealthy(), "supervisor must be marked unhealthy after grow detected") + + // The grown tail must have been zeroed so no leftover guest data remains + // in the (now-unhealthy but still-instantiated) module. + tail, readOK := mem.Read(origSize, grownSize-origSize) + require.True(t, readOK) + expected := make([]byte, grownSize-origSize) + assert.Equal(t, expected, []byte(tail), "tail must be zero-filled, not contain poison bytes") +} diff --git a/internal/execution/wasm/json_util.go b/internal/execution/wasm/json_util.go new file mode 100644 index 0000000..5eba909 --- /dev/null +++ b/internal/execution/wasm/json_util.go @@ -0,0 +1,17 @@ +package wasm + +import ( + "encoding/json" + "fmt" + "strings" +) + +// parseJSONResponse unmarshals a JSON object from the given string. +func parseJSONResponse(s string) (map[string]any, error) { + s = strings.TrimSpace(s) + var result map[string]any + if err := json.Unmarshal([]byte(s), &result); err != nil { + return nil, fmt.Errorf("unmarshal JSON: %w (raw: %.200s)", err, s) + } + return result, nil +} diff --git a/internal/execution/wasm/pool.go b/internal/execution/wasm/pool.go new file mode 100644 index 0000000..08cbeff --- /dev/null +++ b/internal/execution/wasm/pool.go @@ -0,0 +1,42 @@ +package wasm + +import ( + "context" + + "go.uber.org/zap" +) + +// poolItem is the interface satisfied by any item that can be shut down when +// draining a pool (for example wasmSupervisor or ResidentPythonRunner). +type poolItem interface { + Shutdown(ctx context.Context) error +} + +// drainPool receives up to cap(pool) items from the channel and calls +// Shutdown on each. If the context is cancelled before all items are drained, +// it logs a warning and returns early, avoiding the deadlock that occurs when +// an unhealthy item was discarded and its replacement goroutine hasn't +// finished yet. +func drainPool[T poolItem](ctx context.Context, pool chan T, log *zap.Logger) error { + if pool == nil { + return nil + } + + var firstErr error + for i := 0; i < cap(pool); i++ { + select { + case item := <-pool: + if err := item.Shutdown(ctx); err != nil { + log.Error("error shutting down pool item", zap.Error(err)) + if firstErr == nil { + firstErr = err + } + } + case <-ctx.Done(): + log.Warn("drainPool: context cancelled, some items may not be shut down", + zap.Int("remaining", cap(pool)-i)) + return ctx.Err() + } + } + return firstErr +} diff --git a/internal/execution/wasm/python_preload_config_test.go b/internal/execution/wasm/python_preload_config_test.go new file mode 100644 index 0000000..d3d94fa --- /dev/null +++ b/internal/execution/wasm/python_preload_config_test.go @@ -0,0 +1,28 @@ +package wasm + +import "testing" + +func TestPythonPreloadModeDefaultsToEvaluator(t *testing.T) { + var cfg Config + cfg.applyDefaults() + if cfg.PythonPreloadMode != "evaluator" { + t.Fatalf("default preload mode = %q, want evaluator", cfg.PythonPreloadMode) + } +} + +func TestPythonPreloadModeCanBeDisabled(t *testing.T) { + t.Setenv("FUNCTION_WASM_PYTHON_PRELOAD", "off") + var cfg Config + cfg.applyEnv() + cfg.applyDefaults() + if cfg.PythonPreloadMode != "off" { + t.Fatalf("preload mode = %q, want off", cfg.PythonPreloadMode) + } +} + +func TestPythonPreloadModeRejectsUnknownValue(t *testing.T) { + cfg := Config{PythonPreloadMode: "typo"} + if err := cfg.validatePythonPreloadMode(); err == nil { + t.Fatal("unknown preload mode must fail closed") + } +} diff --git a/internal/execution/wasm/python_reactor_artifact.go b/internal/execution/wasm/python_reactor_artifact.go new file mode 100644 index 0000000..e382490 --- /dev/null +++ b/internal/execution/wasm/python_reactor_artifact.go @@ -0,0 +1,185 @@ +package wasm + +import ( + "fmt" + "slices" + "sort" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" +) + +type pythonReactorImport struct { + Module string `json:"module"` + Name string `json:"name"` +} + +type pythonReactorFunctionSignature struct { + Params []api.ValueType + Results []api.ValueType +} + +type pythonReactorModuleShape struct { + Exports map[string]pythonReactorFunctionSignature + ExportedMemories map[string]struct{} + Imports map[pythonReactorImport]struct{} +} + +func inspectPythonReactorCompiledModule(compiled wazero.CompiledModule) (pythonReactorModuleShape, error) { + if compiled == nil { + return pythonReactorModuleShape{}, fmt.Errorf("python-reactor: compiled module is nil") + } + shape := pythonReactorModuleShape{ + Exports: make(map[string]pythonReactorFunctionSignature), + ExportedMemories: make(map[string]struct{}), + Imports: make(map[pythonReactorImport]struct{}), + } + for name, definition := range compiled.ExportedFunctions() { + shape.Exports[name] = pythonReactorFunctionSignature{ + Params: append([]api.ValueType(nil), definition.ParamTypes()...), + Results: append([]api.ValueType(nil), definition.ResultTypes()...), + } + } + for name := range compiled.ExportedMemories() { + shape.ExportedMemories[name] = struct{}{} + } + for _, definition := range compiled.ImportedFunctions() { + module, name, imported := definition.Import() + if !imported { + return pythonReactorModuleShape{}, fmt.Errorf("python-reactor: imported function has no import identity") + } + shape.Imports[pythonReactorImport{Module: module, Name: name}] = struct{}{} + } + for _, definition := range compiled.ImportedMemories() { + module, name, imported := definition.Import() + if !imported { + return pythonReactorModuleShape{}, fmt.Errorf("python-reactor: imported memory has no import identity") + } + shape.Imports[pythonReactorImport{Module: module, Name: name}] = struct{}{} + } + return shape, nil +} + +func verifyCompiledPythonReactorArtifact(compiled wazero.CompiledModule, artifact *AgentPythonArtifact) error { + shape, err := inspectPythonReactorCompiledModule(compiled) + if err != nil { + return err + } + return verifyPythonReactorModuleShape(shape, artifact) +} + +func verifyPythonReactorModuleShape(shape pythonReactorModuleShape, artifact *AgentPythonArtifact) error { + if artifact == nil { + return fmt.Errorf("python-reactor: artifact contract is nil") + } + + initExport := artifact.InitExport + prepareExport := artifact.PrepareExport + executeExport := artifact.ExecuteExport + if initExport == "" { + initExport, prepareExport, executeExport = "runtime_init", "runtime_prepare", "execute" + } + i32 := api.ValueTypeI32 + required := map[string]pythonReactorFunctionSignature{ + "_initialize": {}, + initExport: {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + prepareExport: {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + "alloc": {Params: []api.ValueType{i32}, Results: []api.ValueType{i32}}, + "dealloc": {Params: []api.ValueType{i32}}, + executeExport: {Params: []api.ValueType{i32, i32}, Results: []api.ValueType{i32}}, + } + if artifact.ABI == "shimmy-python-runtime/v1" { + required[initExport] = pythonReactorFunctionSignature{Results: []api.ValueType{i32}} + required["shimmy_python_runtime_identity"] = pythonReactorFunctionSignature{Results: []api.ValueType{i32}} + } + for name, expected := range required { + actual, ok := shape.Exports[name] + if !ok { + return fmt.Errorf("python-reactor: actual module is missing required export %q", name) + } + if !samePythonReactorSignature(actual, expected) { + return fmt.Errorf("python-reactor: export %q has ABI params=%s results=%s; want params=%s results=%s", name, formatWasmValueTypes(actual.Params), formatWasmValueTypes(actual.Results), formatWasmValueTypes(expected.Params), formatWasmValueTypes(expected.Results)) + } + } + if _, ok := shape.ExportedMemories["memory"]; !ok { + return fmt.Errorf("python-reactor: actual module is missing required exported memory %q", "memory") + } + + for _, name := range artifact.DeclaredExports { + if _, ok := shape.Exports[name]; ok { + continue + } + if _, ok := shape.ExportedMemories[name]; ok { + continue + } + return fmt.Errorf("python-reactor: manifest export %q is absent from actual module", name) + } + + declaredImports := make(map[pythonReactorImport]struct{}, len(artifact.DeclaredImports)) + for _, imported := range artifact.DeclaredImports { + declaredImports[imported] = struct{}{} + } + var undeclared []pythonReactorImport + for imported := range shape.Imports { + if _, ok := declaredImports[imported]; !ok { + undeclared = append(undeclared, imported) + } + } + sort.Slice(undeclared, func(i, j int) bool { + if undeclared[i].Module == undeclared[j].Module { + return undeclared[i].Name < undeclared[j].Name + } + return undeclared[i].Module < undeclared[j].Module + }) + if len(undeclared) > 0 { + return fmt.Errorf("python-reactor: actual import %q.%q is not declared by manifest", undeclared[0].Module, undeclared[0].Name) + } + var absent []pythonReactorImport + for imported := range declaredImports { + if _, ok := shape.Imports[imported]; !ok { + absent = append(absent, imported) + } + } + sort.Slice(absent, func(i, j int) bool { + if absent[i].Module == absent[j].Module { + return absent[i].Name < absent[j].Name + } + return absent[i].Module < absent[j].Module + }) + if len(absent) > 0 { + return fmt.Errorf("python-reactor: manifest import %q.%q is absent from actual module", absent[0].Module, absent[0].Name) + } + return nil +} + +func samePythonReactorSignature(actual, expected pythonReactorFunctionSignature) bool { + return slices.Equal(actual.Params, expected.Params) && slices.Equal(actual.Results, expected.Results) +} + +func formatPythonReactorSignature(signature pythonReactorFunctionSignature) string { + return fmt.Sprintf("params=%s results=%s", formatWasmValueTypes(signature.Params), formatWasmValueTypes(signature.Results)) +} + +func formatWasmValueTypes(types []api.ValueType) string { + if len(types) == 0 { + return "[]" + } + names := make([]string, len(types)) + for i, valueType := range types { + switch valueType { + case api.ValueTypeI32: + names[i] = "i32" + case api.ValueTypeI64: + names[i] = "i64" + case api.ValueTypeF32: + names[i] = "f32" + case api.ValueTypeF64: + names[i] = "f64" + case api.ValueTypeExternref: + names[i] = "externref" + default: + names[i] = fmt.Sprintf("0x%x", valueType) + } + } + return fmt.Sprintf("%v", names) +} diff --git a/internal/execution/wasm/robustness_test.go b/internal/execution/wasm/robustness_test.go new file mode 100644 index 0000000..cb858bb --- /dev/null +++ b/internal/execution/wasm/robustness_test.go @@ -0,0 +1,130 @@ +package wasm + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func wasmULEB(v uint32) []byte { + var buf []byte + for { + b := byte(v & 0x7f) + v >>= 7 + if v != 0 { + b |= 0x80 + } + buf = append(buf, b) + if v == 0 { + break + } + } + return buf +} + +func wasmSection(id byte, payload []byte) []byte { + out := []byte{id} + out = append(out, wasmULEB(uint32(len(payload)))...) + out = append(out, payload...) + return out +} + +func wasmName(s string) []byte { + out := wasmULEB(uint32(len(s))) + out = append(out, []byte(s)...) + return out +} + +func malformedABIWasm(allocReturnsValue, dispatchReturnsValue bool) []byte { + module := []byte{0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00} + + // Types: alloc(i32) [-> i32], dispatch(i32, i32) [-> i32]. + var types []byte + types = append(types, 0x02) + types = append(types, 0x60, 0x01, 0x7f) + if allocReturnsValue { + types = append(types, 0x01, 0x7f) + } else { + types = append(types, 0x00) + } + types = append(types, 0x60, 0x02, 0x7f, 0x7f) + if dispatchReturnsValue { + types = append(types, 0x01, 0x7f) + } else { + types = append(types, 0x00) + } + module = append(module, wasmSection(1, types)...) + + // Two functions: alloc uses type 0; dispatch uses type 1. + module = append(module, wasmSection(3, []byte{0x02, 0x00, 0x01})...) + + // One memory page. + module = append(module, wasmSection(5, []byte{0x01, 0x00, 0x01})...) + + // Export memory, alloc, dispatch. + var exports []byte + exports = append(exports, 0x03) + exports = append(exports, wasmName("memory")...) + exports = append(exports, 0x02, 0x00) + exports = append(exports, wasmName("alloc")...) + exports = append(exports, 0x00, 0x00) + exports = append(exports, wasmName("dispatch")...) + exports = append(exports, 0x00, 0x01) + module = append(module, wasmSection(7, exports)...) + + // Code bodies. + var code []byte + code = append(code, 0x02) + allocBody := []byte{0x00} + if allocReturnsValue { + allocBody = append(allocBody, 0x41, 0x08) // i32.const 8 + } + allocBody = append(allocBody, 0x0b) // end + code = append(code, wasmULEB(uint32(len(allocBody)))...) + code = append(code, allocBody...) + + dispatchBody := []byte{0x00} + if dispatchReturnsValue { + dispatchBody = append(dispatchBody, 0x41, 0x08) // i32.const 8 + } + dispatchBody = append(dispatchBody, 0x0b) // end + code = append(code, wasmULEB(uint32(len(dispatchBody)))...) + code = append(code, dispatchBody...) + module = append(module, wasmSection(10, code)...) + + return module +} + +func writeTempWasm(t *testing.T, bytes []byte) string { + t.Helper() + path := filepath.Join(t.TempDir(), "eval.wasm") + require.NoError(t, os.WriteFile(path, bytes, 0o644)) + return path +} + +func TestDispatcher_Send_ReturnsErrorForAllocWithoutReturnValue(t *testing.T) { + path := writeTempWasm(t, malformedABIWasm(false, true)) + d := NewDispatcher(Config{ModulePath: path, MaxInstances: 1, Timeout: time.Second}, newTestLogger(t)) + require.NoError(t, d.Start(context.Background())) + t.Cleanup(func() { _ = d.Shutdown(context.Background()) }) + + _, err := d.Send(context.Background(), "eval", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "alloc returned 0 values") +} + +func TestDispatcher_Send_ReturnsErrorForDispatchWithoutReturnValue(t *testing.T) { + path := writeTempWasm(t, malformedABIWasm(true, false)) + d := NewDispatcher(Config{ModulePath: path, MaxInstances: 1, Timeout: time.Second}, newTestLogger(t)) + require.NoError(t, d.Start(context.Background())) + t.Cleanup(func() { _ = d.Shutdown(context.Background()) }) + + _, err := d.Send(context.Background(), "eval", nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "dispatch returned 0 values") +} diff --git a/internal/execution/wasm/snapshot.go b/internal/execution/wasm/snapshot.go new file mode 100644 index 0000000..67a1a82 --- /dev/null +++ b/internal/execution/wasm/snapshot.go @@ -0,0 +1,108 @@ +package wasm + +import ( + "errors" + "fmt" + + "github.com/tetratelabs/wazero/api" +) + +// ErrSnapshotMemoryDrifted means the guest changed linear-memory size after +// the post-initialisation snapshot. WASM memory cannot shrink, so restoring +// only the captured prefix would leave request state in the grown tail. +var ErrSnapshotMemoryDrifted = errors.New("snapshot: wasm linear memory size drifted") + +// SnapshotStrategy abstracts the full-memory snapshot used by both generic +// WASM and Python Reactor execution. FullMemcpyStrategy is the only +// implementation. +// +// Contract (I-4 fix — document ordering and concurrency expectations): +// - Take must be called at least once before Restore. +// - Take may be called multiple times; each call overwrites the previous +// snapshot. +// - Calling Restore without a prior Take is a no-op (returns nil) but +// logically meaningless. +// - Implementations are NOT safe for concurrent calls to Take / Restore. +// The caller (wasmSupervisor) must serialise access. +type SnapshotStrategy interface { + // Take captures the current state of the WASM linear memory. + // It is called once after module initialisation. + Take(mem api.Memory) error + + // Restore writes the captured snapshot back into WASM linear memory. + // It is called after every request so the next request sees a clean state. + Restore(mem api.Memory) error + + // Close releases the owned snapshot buffer. + Close() error +} + +// --------------------------------------------------------------------------- +// FullMemcpyStrategy +// --------------------------------------------------------------------------- + +// FullMemcpyStrategy is the always-available baseline: it copies the entire +// linear memory into a []byte on Take and writes it all back on Restore. +// Cost is O(total memory size) regardless of how many pages were actually +// written during the request. +type FullMemcpyStrategy struct { + snapshot []byte + size uint32 +} + +// NewFullMemcpyStrategy returns a ready-to-use FullMemcpyStrategy. +func NewFullMemcpyStrategy() *FullMemcpyStrategy { + return &FullMemcpyStrategy{} +} + +// Take implements SnapshotStrategy. +func (f *FullMemcpyStrategy) Take(mem api.Memory) error { + if mem == nil { + f.snapshot = nil + f.size = 0 + return nil + } + + size := mem.Size() + if size == 0 { + f.snapshot = nil + f.size = 0 + return nil + } + + buf, ok := mem.Read(0, size) + if !ok { + return fmt.Errorf("snapshot: could not read %d bytes of linear memory", size) + } + + // Make an owned copy — mem.Read may return a slice backed by the wazero + // memory buffer which could be modified by subsequent guest execution. + f.snapshot = make([]byte, len(buf)) + copy(f.snapshot, buf) + f.size = size + + return nil +} + +// Restore implements SnapshotStrategy. +func (f *FullMemcpyStrategy) Restore(mem api.Memory) error { + if f.snapshot == nil || mem == nil { + return nil + } + if mem.Size() != f.size { + return fmt.Errorf("%w: captured=%d current=%d", ErrSnapshotMemoryDrifted, f.size, mem.Size()) + } + + if !mem.Write(0, f.snapshot) { + return fmt.Errorf("snapshot: failed to restore %d bytes", len(f.snapshot)) + } + + return nil +} + +// Close implements SnapshotStrategy. FullMemcpyStrategy holds no OS resources. +func (f *FullMemcpyStrategy) Close() error { + f.snapshot = nil + f.size = 0 + return nil +} diff --git a/internal/execution/wasm/snapshot_test.go b/internal/execution/wasm/snapshot_test.go new file mode 100644 index 0000000..ea19993 --- /dev/null +++ b/internal/execution/wasm/snapshot_test.go @@ -0,0 +1,274 @@ +//go:build !plan9 + +package wasm + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" +) + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// leb128Encode encodes a uint32 as an unsigned LEB128 byte slice. +func leb128Encode(v uint32) []byte { + var buf []byte + for { + b := byte(v & 0x7f) + v >>= 7 + if v != 0 { + b |= 0x80 + } + buf = append(buf, b) + if v == 0 { + break + } + } + return buf +} + +// buildTestMemoryModule constructs a minimal WASM binary that declares exactly +// `pages` pages (64 KiB each) of linear memory. wazero's Module.Memory() +// returns the first memory regardless of whether it is exported, so no export +// section is needed. +// +// Binary layout (WASM spec §5): +// +// \0asm (magic) + version (1) + memory section +// +// This mirrors buildMinimalMemoryModule from snapshot_bench_test.go but +// accepts *testing.T so it can be used in unit tests. +func buildTestMemoryModule(t testing.TB, pages int) []byte { + t.Helper() + + // Memory section payload: count=1, limits type=0x00 (min only), min=pages + pagesLEB := leb128Encode(uint32(pages)) + memPayload := append([]byte{0x01, 0x00}, pagesLEB...) + + // Section: id=5 (memory), size=len(payload), payload + memSec := append([]byte{0x05}, append(leb128Encode(uint32(len(memPayload))), memPayload...)...) + + // Full module: magic + version + memory section + module := []byte{0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00} + module = append(module, memSec...) + return module +} + +// newTestWazeroMemory instantiates a minimal WASM module with the given number +// of 64 KiB pages and returns its api.Memory. The runtime and module are +// closed via t.Cleanup. +func newTestWazeroMemory(t testing.TB, pages int) api.Memory { + t.Helper() + ctx := context.Background() + + wasmBin := buildTestMemoryModule(t, pages) + + rt := wazero.NewRuntime(ctx) + t.Cleanup(func() { _ = rt.Close(ctx) }) + + compiled, err := rt.CompileModule(ctx, wasmBin) + require.NoError(t, err, "compile minimal module") + t.Cleanup(func() { _ = compiled.Close(ctx) }) + + mod, err := rt.InstantiateModule(ctx, compiled, wazero.NewModuleConfig().WithName("")) + require.NoError(t, err, "instantiate minimal module") + t.Cleanup(func() { _ = mod.Close(ctx) }) + + mem := mod.Memory() + require.NotNil(t, mem, "module must have linear memory") + return mem +} + +// --------------------------------------------------------------------------- +// TestFullMemcpyStrategy_TakeRestoreRoundtrip +// --------------------------------------------------------------------------- + +// TestFullMemcpyStrategy_TakeRestoreRoundtrip verifies the core contract: +// after Take, mutating the memory and calling Restore brings it back to the +// snapshotted state. +func TestFullMemcpyStrategy_TakeRestoreRoundtrip(t *testing.T) { + mem := newTestWazeroMemory(t, 1) // 1 page = 64 KiB + + // Fill memory with a known pattern. + size := mem.Size() + pattern := make([]byte, size) + for i := range pattern { + pattern[i] = byte(i % 251) + } + require.True(t, mem.Write(0, pattern), "write initial pattern") + + s := NewFullMemcpyStrategy() + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + // Take snapshot. + require.NoError(t, s.Take(mem)) + + // Overwrite memory with zeros (simulated guest write). + zeros := make([]byte, size) + require.True(t, mem.Write(0, zeros), "overwrite with zeros") + + after, ok := mem.Read(0, size) + require.True(t, ok) + require.Equal(t, zeros, []byte(after), "sanity: memory should be all-zeros now") + + // Restore and verify memory matches original pattern. + require.NoError(t, s.Restore(mem)) + + restored, ok := mem.Read(0, size) + require.True(t, ok) + assert.Equal(t, pattern, []byte(restored), "Restore must return memory to snapshotted state") +} + +func TestFullMemcpyStrategy_RejectsMemoryGrowth(t *testing.T) { + mem := newTestWazeroMemory(t, 1) + s := NewFullMemcpyStrategy() + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + require.NoError(t, s.Take(mem)) + previousPages, ok := mem.Grow(1) + require.True(t, ok) + require.Equal(t, uint32(1), previousPages) + + err := s.Restore(mem) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrSnapshotMemoryDrifted)) +} + +// --------------------------------------------------------------------------- +// TestFullMemcpyStrategy_TakeNilMemory +// --------------------------------------------------------------------------- + +// TestFullMemcpyStrategy_TakeNilMemory checks that Take(nil) is safe and +// results in a nil snapshot (no panic, no error). +func TestFullMemcpyStrategy_TakeNilMemory(t *testing.T) { + s := NewFullMemcpyStrategy() + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + require.NoError(t, s.Take(nil)) + assert.Nil(t, s.snapshot, "snapshot should be nil after Take(nil)") + + // A subsequent Restore(nil) must also be a no-op. + require.NoError(t, s.Restore(nil)) +} + +// --------------------------------------------------------------------------- +// TestFullMemcpyStrategy_RestoreBeforeTake +// --------------------------------------------------------------------------- + +// TestFullMemcpyStrategy_RestoreBeforeTake verifies that calling Restore on a +// zero-value / never-initialised strategy is a no-op that does not modify +// memory or return an error. +func TestFullMemcpyStrategy_RestoreBeforeTake(t *testing.T) { + s := NewFullMemcpyStrategy() + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + mem := newTestWazeroMemory(t, 1) + size := mem.Size() + + // Fill with recognisable data. + data := make([]byte, size) + for i := range data { + data[i] = byte(i % 97) + } + require.True(t, mem.Write(0, data), "write initial data") + + // Snapshot the state so we can compare after Restore. + before, ok := mem.Read(0, size) + require.True(t, ok) + beforeCopy := make([]byte, len(before)) + copy(beforeCopy, before) + + // Restore before any Take — must be a no-op (snapshot is nil). + require.NoError(t, s.Restore(mem)) + + after, ok := mem.Read(0, size) + require.True(t, ok) + assert.Equal(t, beforeCopy, []byte(after), "Restore before Take must leave memory unchanged") +} + +// --------------------------------------------------------------------------- +// TestFullMemcpyStrategy_EmptyMemory +// --------------------------------------------------------------------------- + +// TestFullMemcpyStrategy_EmptyMemory checks that a zero-size case in snapshot +// logic produces a nil snapshot (size==0 branch). We test this by calling +// Take with nil (which mirrors the zero-size code path in the implementation: +// both nil and zero-size result in snapshot=nil). +func TestFullMemcpyStrategy_EmptyMemory(t *testing.T) { + s := NewFullMemcpyStrategy() + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + // Take(nil) exercises the "mem == nil" branch which sets snapshot=nil. + require.NoError(t, s.Take(nil)) + assert.Nil(t, s.snapshot, "snapshot must be nil when memory is nil") + + // Restore(nil) must be a no-op. + require.NoError(t, s.Restore(nil)) +} + +// --------------------------------------------------------------------------- +// TestFullMemcpyStrategy_CloseIdempotent +// --------------------------------------------------------------------------- + +// TestFullMemcpyStrategy_CloseIdempotent verifies that Close can be called +// multiple times without panicking or returning an error. +func TestFullMemcpyStrategy_CloseIdempotent(t *testing.T) { + s := NewFullMemcpyStrategy() + + mem := newTestWazeroMemory(t, 1) + require.NoError(t, s.Take(mem)) + assert.NotNil(t, s.snapshot, "snapshot should be set after Take") + + // First Close should succeed and clear the snapshot. + require.NoError(t, s.Close()) + assert.Nil(t, s.snapshot, "snapshot should be nil after first Close") + + // Second Close must also be safe. + require.NoError(t, s.Close()) +} + +// --------------------------------------------------------------------------- +// TestFullMemcpyStrategy_SnapshotIsOwnedCopy +// --------------------------------------------------------------------------- + +// TestFullMemcpyStrategy_SnapshotIsOwnedCopy confirms that the snapshot is an +// independent copy of the memory buffer, not an alias into wazero's backing +// store. If Take stored a slice backed by the same underlying array, a +// subsequent guest write would silently corrupt the snapshot. +func TestFullMemcpyStrategy_SnapshotIsOwnedCopy(t *testing.T) { + mem := newTestWazeroMemory(t, 1) + size := mem.Size() + + // Write distinct pattern. + pattern := make([]byte, size) + for i := range pattern { + pattern[i] = byte(i % 199) + } + require.True(t, mem.Write(0, pattern), "write pattern") + + s := NewFullMemcpyStrategy() + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + require.NoError(t, s.Take(mem)) + + // Overwrite memory entirely with 0xFF. + corrupt := make([]byte, size) + for i := range corrupt { + corrupt[i] = 0xFF + } + require.True(t, mem.Write(0, corrupt)) + + // Restore: snapshot must be independent of the wazero buffer. + require.NoError(t, s.Restore(mem)) + + restored, ok := mem.Read(0, size) + require.True(t, ok) + assert.Equal(t, pattern, []byte(restored), "snapshot must be independent copy of original data") +} diff --git a/internal/execution/wasm/supervisor.go b/internal/execution/wasm/supervisor.go new file mode 100644 index 0000000..3162337 --- /dev/null +++ b/internal/execution/wasm/supervisor.go @@ -0,0 +1,226 @@ +package wasm + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" + "go.uber.org/zap" +) + +// ErrMemoryGrew indicates that the guest expanded linear memory during a +// request beyond the size captured at snapshot time. wazero (and the WASM +// spec) does not allow shrinking linear memory, so the original snapshotted +// state cannot be fully reproduced and the supervisor must be discarded. +var ErrMemoryGrew = errors.New("wasm: linear memory grew beyond snapshotted size") + +// wasmSupervisor manages a single instantiated WASM module. After the module +// is initialised its linear memory is snapshotted; the snapshot is restored +// after every Send so that the next request sees a clean initial state. This +// gives cheap warm-start semantics without re-compiling the module. +type wasmSupervisor struct { + mu sync.Mutex + + runtime wazero.Runtime + compiled wazero.CompiledModule + modCfg wazero.ModuleConfig + + mod api.Module + adapter *wasmAdapter + + // strategy owns the full linear-memory copy restored after each request. + strategy SnapshotStrategy + + // healthy is true when the supervisor is in a known-good state and can be + // safely returned to the pool. It is set to false when restoreSnapshot fails, + // indicating the WASM module's memory state is undefined. + healthy bool + + // snapshotSize is the linear-memory size (in bytes) captured at Take time. + // restoreSnapshot compares this against the post-request memory size to + // detect memory.grow during execution — wazero cannot shrink memory, so + // any growth invalidates the snapshot and must mark the supervisor unhealthy. + snapshotSize uint32 + + timeout time.Duration + log *zap.Logger +} + +func newWasmSupervisor( + rt wazero.Runtime, + compiled wazero.CompiledModule, + modCfg wazero.ModuleConfig, + timeout time.Duration, + log *zap.Logger, +) *wasmSupervisor { + return &wasmSupervisor{ + runtime: rt, + compiled: compiled, + modCfg: modCfg, + timeout: timeout, + log: log.Named("supervisor_wasm"), + } +} + +// Start instantiates the compiled module, runs any WASI start function, then +// snapshots linear memory. +func (s *wasmSupervisor) Start(ctx context.Context) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.mod != nil { + return nil + } + + s.log.Debug("instantiating wasm module") + + // Apply start functions on top of the provided (sandboxed) module config. + instCfg := s.modCfg.WithStartFunctions("_initialize", "_start") + + mod, err := s.runtime.InstantiateModule(ctx, s.compiled, instCfg) + if err != nil { + releaseErr := s.closeResources(ctx) + return errors.Join(fmt.Errorf("wasm: instantiate module: %w", err), releaseErr) + } + + s.mod = mod + s.adapter = newWasmAdapter(mod, s.log) + s.healthy = true + + s.strategy = NewFullMemcpyStrategy() + + // Snapshot linear memory so we can restore it before each request. + if err := s.takeSnapshot(); err != nil { + releaseErr := s.closeResources(ctx) + return errors.Join(fmt.Errorf("wasm: snapshot memory: %w", err), releaseErr) + } + + memSize := uint32(0) + if m := s.mod.Memory(); m != nil { + memSize = m.Size() + } + s.log.Debug("wasm module ready", + zap.Uint32("snapshot_bytes", memSize), + zap.String("strategy", fmt.Sprintf("%T", s.strategy)), + ) + + return nil +} + +// Send calls the guest's dispatch function, then restores linear memory from +// the snapshot so the next request starts from a clean state. +func (s *wasmSupervisor) Send( + ctx context.Context, + method string, + data map[string]any, +) (map[string]any, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.mod == nil || s.adapter == nil { + return nil, fmt.Errorf("wasm: supervisor not started") + } + + result, err := s.adapter.send(ctx, method, data, s.timeout) + + // Restore memory snapshot to keep state clean for the next request. + // If restore fails, mark the supervisor unhealthy so the dispatcher + // discards it rather than returning it to the pool with undefined state. + if restoreErr := s.restoreSnapshot(); restoreErr != nil { + s.log.Error("failed to restore memory snapshot — marking supervisor unhealthy", zap.Error(restoreErr)) + s.healthy = false + if err == nil { + err = fmt.Errorf("wasm: restore snapshot: %w", restoreErr) + } + } + + return result, err +} + +// IsHealthy reports whether the supervisor is in a known-good state. +// Safe to call without holding s.mu (acquires the lock internally). (I-3 fix) +func (s *wasmSupervisor) IsHealthy() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.healthy +} + +// Shutdown closes the module instance and releases resources. +func (s *wasmSupervisor) Shutdown(ctx context.Context) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.mod == nil && s.strategy == nil { + return nil + } + + s.log.Debug("shutting down wasm module instance") + return s.closeResources(ctx) +} + +// closeResources must run only after guest execution has stopped and while +// s.mu is held. +func (s *wasmSupervisor) closeResources(ctx context.Context) error { + var moduleErr, strategyErr error + if s.mod != nil { + moduleErr = s.mod.Close(ctx) + s.mod = nil + s.adapter = nil + } + if s.strategy != nil { + strategyErr = s.strategy.Close() + s.strategy = nil + } + return errors.Join(moduleErr, strategyErr) +} + +// takeSnapshot captures the guest's linear memory via the active strategy and +// records the memory size so restoreSnapshot can detect post-snapshot growth. +// Must be called with s.mu held. +func (s *wasmSupervisor) takeSnapshot() error { + mem := s.mod.Memory() + if mem == nil { + s.snapshotSize = 0 + return nil + } + if err := s.strategy.Take(mem); err != nil { + return err + } + s.snapshotSize = mem.Size() + return nil +} + +// restoreSnapshot restores the guest's linear memory from the last snapshot +// via the active strategy. If the guest grew memory during the request +// (memory.grow), it zero-fills the tail beyond the snapshotted size to prevent +// leaking guest data into the next request and returns ErrMemoryGrew so the +// caller (Send) marks the supervisor unhealthy and discards it. Must be called +// with s.mu held. +func (s *wasmSupervisor) restoreSnapshot() error { + if s.mod == nil { + return nil + } + mem := s.mod.Memory() + if mem == nil { + return nil + } + if cur := mem.Size(); cur > s.snapshotSize { + tail := cur - s.snapshotSize + zeros := make([]byte, tail) + if !mem.Write(s.snapshotSize, zeros) { + return fmt.Errorf("wasm: memory grew by %d bytes; zero-fill failed: %w", tail, ErrMemoryGrew) + } + // The instance is discarded after this error, so restoring the captured + // prefix has no value. Returning before strategy.Restore also avoids + // asking pointer/size-sensitive strategies to touch a drifted backing. + return fmt.Errorf("wasm: memory grew by %d bytes (tail zero-filled): %w", tail, ErrMemoryGrew) + } + if err := s.strategy.Restore(mem); err != nil { + return err + } + return nil +} diff --git a/internal/execution/wasm/testdata/echo.wasm b/internal/execution/wasm/testdata/echo.wasm new file mode 100644 index 0000000000000000000000000000000000000000..17ec87da1eb19faacfc0f6d02fcb819a4a3b5bc7 GIT binary patch literal 241 zcmX}nF%Q8|6b0aO?`x^PVl+w0bwr#)k387M^iC~((u7`U RpT6lJRcm8`hYbQg`T>CeB{l#6 literal 0 HcmV?d00001 diff --git a/internal/execution/wasm/testdata/echo.wat b/internal/execution/wasm/testdata/echo.wat new file mode 100644 index 0000000..b7fb29d --- /dev/null +++ b/internal/execution/wasm/testdata/echo.wat @@ -0,0 +1,66 @@ +;; echo.wat — minimal guest ABI fixture for wasm package tests. +;; +;; Implements: +;; alloc(size i32) i32 — bump allocator; heap pointer stored at mem[0..3] +;; dispatch(req_ptr i32, req_len i32) i32 +;; — ignores input; always returns fixed response {"ok":true} +;; as a length-prefixed blob: [4-byte LE uint32 len][JSON bytes] +;; +;; The compiled binary (echo.wasm) was generated from this source. +;; {"ok":true} is 11 bytes: 7b 22 6f 6b 22 3a 74 72 75 65 7d +;; +;; Design note: the heap pointer is stored IN linear memory (offset 0, 4 bytes) +;; rather than in a WASM global. This means the snapshot/restore mechanism +;; (which copies linear memory) correctly resets the allocator state between +;; requests. If a global were used, snapshot/restore would not reset it and +;; the heap pointer would keep advancing across requests. +(module + (memory (export "memory") 1) + + ;; mem[0..3]: heap pointer (i32, LE), initialized to 4 + ;; (offset 0..3 reserved for the pointer itself, so allocations start at 4) + (data (i32.const 0) "\04\00\00\00") + + ;; alloc(size i32) i32 + (func (export "alloc") (param $size i32) (result i32) + (local $ptr i32) + ;; ptr = i32.load(mem[0]) + (local.set $ptr (i32.load (i32.const 0))) + ;; mem[0] = ptr + size + (i32.store (i32.const 0) (i32.add (local.get $ptr) (local.get $size))) + (local.get $ptr) + ) + + ;; dispatch(req_ptr i32, req_len i32) i32 + ;; Returns pointer P where: + ;; mem[P .. P+4) = little-endian uint32 length (11) + ;; mem[P+4 .. P+15) = {"ok":true} + (func (export "dispatch") (param $req_ptr i32) (param $req_len i32) (result i32) + (local $resp_ptr i32) + ;; resp_ptr = i32.load(mem[0]) + (local.set $resp_ptr (i32.load (i32.const 0))) + ;; mem[0] = resp_ptr + 15 (4 bytes length prefix + 11 bytes JSON) + (i32.store (i32.const 0) (i32.add (local.get $resp_ptr) (i32.const 15))) + + ;; Write little-endian length prefix: 11, 0, 0, 0 + (i32.store8 offset=0 (local.get $resp_ptr) (i32.const 11)) + (i32.store8 offset=1 (local.get $resp_ptr) (i32.const 0)) + (i32.store8 offset=2 (local.get $resp_ptr) (i32.const 0)) + (i32.store8 offset=3 (local.get $resp_ptr) (i32.const 0)) + + ;; Write {"ok":true} + (i32.store8 offset=4 (local.get $resp_ptr) (i32.const 0x7b)) ;; { + (i32.store8 offset=5 (local.get $resp_ptr) (i32.const 0x22)) ;; " + (i32.store8 offset=6 (local.get $resp_ptr) (i32.const 0x6f)) ;; o + (i32.store8 offset=7 (local.get $resp_ptr) (i32.const 0x6b)) ;; k + (i32.store8 offset=8 (local.get $resp_ptr) (i32.const 0x22)) ;; " + (i32.store8 offset=9 (local.get $resp_ptr) (i32.const 0x3a)) ;; : + (i32.store8 offset=10 (local.get $resp_ptr) (i32.const 0x74)) ;; t + (i32.store8 offset=11 (local.get $resp_ptr) (i32.const 0x72)) ;; r + (i32.store8 offset=12 (local.get $resp_ptr) (i32.const 0x75)) ;; u + (i32.store8 offset=13 (local.get $resp_ptr) (i32.const 0x65)) ;; e + (i32.store8 offset=14 (local.get $resp_ptr) (i32.const 0x7d)) ;; } + + (local.get $resp_ptr) + ) +) diff --git a/internal/execution/wasm/testhelpers_test.go b/internal/execution/wasm/testhelpers_test.go new file mode 100644 index 0000000..dbee194 --- /dev/null +++ b/internal/execution/wasm/testhelpers_test.go @@ -0,0 +1,46 @@ +package wasm + +import ( + "context" + "os" + "testing" + + "github.com/stretchr/testify/require" + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1" +) + +// echoWasmBytes reads the pre-compiled echo fixture from testdata/echo.wasm. +// The fixture is a minimal WASM module that: +// - exports a bump-allocator alloc(size i32) i32 +// - exports dispatch(req_ptr i32, req_len i32) i32 that always returns the +// fixed JSON {"ok":true} as a 4-byte LE length-prefixed blob +// +// The WAT source is kept alongside the binary at testdata/echo.wat for +// reference. The binary was generated using a pure-Go WASM assembler so that +// the test suite requires no external toolchain. +func echoWasmBytes(t *testing.T) []byte { + t.Helper() + path := echoModulePath(t) + b, err := os.ReadFile(path) + require.NoError(t, err, "read echo.wasm fixture") + return b +} + +// compileEchoModule creates a wazero runtime, wires up WASI host functions, +// and compiles the echo WASM bytes into a CompiledModule. The runtime must be +// closed by the caller. +func compileEchoModule(t *testing.T, ctx context.Context, wasmBytes []byte) (wazero.Runtime, wazero.CompiledModule) { + t.Helper() + + rt := wazero.NewRuntime(ctx) + _, err := wasi_snapshot_preview1.Instantiate(ctx, rt) + require.NoError(t, err, "instantiate WASI") + + compiled, err := rt.CompileModule(ctx, wasmBytes) + require.NoError(t, err, "compile echo module") + + t.Cleanup(func() { _ = compiled.Close(ctx) }) + + return rt, compiled +} From 152860770d022f5aee51c799443e81b1f514b0b3 Mon Sep 17 00:00:00 2001 From: bkmashiro <53376445+bkmashiro@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:08:42 +0100 Subject: [PATCH 2/4] test(wasm): add Linux Python Reactor HTTP E2E --- docs/execution-paths.md | 118 +++++++++++++++++++++++ scripts/e2e-python-reactor.sh | 133 ++++++++++++++++++++++++++ tests/e2e/python-reactor/evaluator.py | 45 +++++++++ 3 files changed, 296 insertions(+) create mode 100644 docs/execution-paths.md create mode 100755 scripts/e2e-python-reactor.sh create mode 100644 tests/e2e/python-reactor/evaluator.py diff --git a/docs/execution-paths.md b/docs/execution-paths.md new file mode 100644 index 0000000..b21a46d --- /dev/null +++ b/docs/execution-paths.md @@ -0,0 +1,118 @@ +# WebAssembly execution paths + +Shimmy keeps the existing `rpc` and `file` process interfaces and adds two +explicit, opt-in WebAssembly paths. Selection is configuration-driven; Shimmy +does not inspect source files or silently retry a request under another backend. + +## Generic WebAssembly + +```bash +FUNCTION_INTERFACE=wasm +FUNCTION_WASM_PROFILE=generic +FUNCTION_WASM_MODULE=/opt/evaluator/evaluator.wasm +``` + +The module runs in-process under wazero and exports `memory`, `alloc`, and +`dispatch`. Shimmy copies each request into guest linear memory, copies the +response out, and restores the prepared memory before reusing the instance. + +Memory reset uses one portable implementation: a full copy of linear memory. +There is no snapshot-strategy selector in this path. If a request grows linear +memory, the instance is discarded because WebAssembly memory cannot shrink back +to the captured size. + +The generic path has no host filesystem access unless paths are explicitly +allowed with `FUNCTION_WASM_ALLOWED_PATHS`. Environment variables are similarly +allowlisted with `FUNCTION_WASM_ALLOWED_ENV`. + +## Python Reactor + +```bash +FUNCTION_INTERFACE=wasm +FUNCTION_WASM_PROFILE=python-reactor +FUNCTION_WASM_MODULE=/opt/runtime/python-reactor.wasm +FUNCTION_WASM_MANIFEST=/opt/runtime/manifest.json +FUNCTION_WASM_PYTHON_SCRIPT=/opt/evaluator/evaluator.py +FUNCTION_WASM_PYTHON_LIFECYCLE=snapshot +``` + +The prepared trusted script owns `dispatch(method, payload)`. Shimmy verifies +the following before serving requests: + +- artifact SHA-256 against the manifest; +- Reactor ABI name and version; +- required imports, exports, and function signatures; and +- manifest-declared `python_modules` against the artifact capability section. + +This proves that the selected artifact and manifest are internally consistent. +Artifact authenticity, trusted Producer commit policy, signatures, and release +provenance remain deployment-system responsibilities. + +### Lifecycle choices + +| Value | Behavior | +|---|---| +| `snapshot` | Prepare once per slot and restore the full linear-memory copy after each successful request. Failed or timed-out slots are discarded and replenished asynchronously. | +| `single-use` | Prepare candidates ahead of time, serve each candidate once, then replace it. | +| `fresh` | Instantiate and prepare a new module for every request. | + +`snapshot` is the default and the only lifecycle that restores memory. Its reset +implementation is always full-memory copy; there is no snapshot-strategy +configuration. `single-use` and `fresh` are lifecycle alternatives, not hidden +fallbacks. Shimmy never changes lifecycle after a request fails. + +Python Reactor does not expose host paths. Leave +`FUNCTION_WASM_ALLOWED_PATHS` unset. Runtime modules are selected by the +manifest-validated artifact profile, for example `base`, `numpy-core`, or +`sympy`. + +### Linux HTTP verification + +Run the HTTP startup and request-flow check against a real Producer artifact and +its exact manifest: + +```bash +SHIMMY_PYTHON_REACTOR_WASM=/opt/runtime/python-reactor.wasm \ +SHIMMY_PYTHON_REACTOR_MANIFEST=/opt/runtime/manifest.json \ + scripts/e2e-python-reactor.sh +``` + +The check starts Shimmy, sends two `eval` requests and one `preview` request, +and verifies prepared-state restoration between requests. + +## Safe Python evaluator example + +[`examples/safe-eval-python`](../examples/safe-eval-python/README.md) is a +backend-level Python Reactor example for student Python in `demo`, `io_test`, +`unit_test`, and `preview` modes. It uses: + +- wazero's WebAssembly capability boundary; +- request deadlines; +- artifact, ABI, and manifest-capability validation; +- full-copy memory reset and failed-slot replacement; and +- evaluator limits for code, input, tests, and output. + +It does not depend on nsjail, privileged Lambda configuration, Node, Docker, or +runtime package installation. AST checks provide early feedback and defense in +depth; they are not a containment boundary. + +```bash +SHIMMY_PYTHON_REACTOR_WASM=/path/to/base.wasm \ +SHIMMY_PYTHON_REACTOR_MANIFEST=/path/to/base.manifest.json \ + scripts/e2e-safe-eval-python.sh +``` + +For the guided base, NumPy, and SymPy examples, follow the +[quick start](../examples/safe-eval-python/README.md#start-here-first-successful-evaluation). + +## Security boundary + +WebAssembly isolation, request deadlines, state reset, and evaluator-level +limits do not form a complete operating-system sandbox. Deployment policy still +owns process memory, aggregate concurrency, authentication, request-size limits, +logging, artifact provenance, and network exposure. + +AWS Lambda cannot grant the namespaces or capabilities required to use nsjail +as a security boundary. On supported Linux hosts or containers, an external OS +sandbox may be added as a separate deployment layer; Shimmy does not claim that +boundary for Lambda. diff --git a/scripts/e2e-python-reactor.sh b/scripts/e2e-python-reactor.sh new file mode 100755 index 0000000..305a41a --- /dev/null +++ b/scripts/e2e-python-reactor.sh @@ -0,0 +1,133 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +WASM="${SHIMMY_PYTHON_REACTOR_WASM:?set SHIMMY_PYTHON_REACTOR_WASM to a Producer artifact}" +MANIFEST="${SHIMMY_PYTHON_REACTOR_MANIFEST:?set SHIMMY_PYTHON_REACTOR_MANIFEST to its manifest.json}" +EVALUATOR="${SHIMMY_E2E_EVALUATOR:-${ROOT}/tests/e2e/python-reactor/evaluator.py}" +HOST="127.0.0.1" +TMP="$(mktemp -d "${TMPDIR:-/tmp}/shimmy-python-reactor-e2e.XXXXXX")" +PORT="${SHIMMY_E2E_PORT:-}" +SERVER_PID="" +PREBUILT_BIN="${SHIMMY_E2E_BIN:-}" +PREBUILT_CHECK="${SHIMMY_E2E_ARTIFACT_CHECK_BIN:-}" + +cleanup() { + if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then + kill "${SERVER_PID}" 2>/dev/null || true + wait "${SERVER_PID}" 2>/dev/null || true + fi + rm -rf "${TMP}" +} +trap cleanup EXIT + +for cmd in curl python3; do + command -v "${cmd}" >/dev/null 2>&1 || { echo "missing required command: ${cmd}" >&2; exit 1; } +done +[[ "$(uname -s)" == "Linux" ]] || { echo "Python Reactor E2E requires Linux" >&2; exit 1; } +[[ -r "${WASM}" && -r "${MANIFEST}" && -r "${EVALUATOR}" ]] || { echo "artifact, manifest, and evaluator must be readable" >&2; exit 1; } + +if [[ -z "${PORT}" ]]; then + PORT="$(python3 - <<'PY' +import socket +with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + print(sock.getsockname()[1]) +PY +)" +fi +BASE_URL="http://${HOST}:${PORT}" +BIN="${PREBUILT_BIN:-${TMP}/shimmy}" +CHECK="${PREBUILT_CHECK:-${TMP}/shimmy-artifact-check}" +LOG="${TMP}/server.log" + +if [[ -n "${PREBUILT_BIN}" || -n "${PREBUILT_CHECK}" ]]; then + [[ -n "${PREBUILT_BIN}" && -n "${PREBUILT_CHECK}" ]] || { + echo "set both SHIMMY_E2E_BIN and SHIMMY_E2E_ARTIFACT_CHECK_BIN" >&2 + exit 1 + } + [[ -x "${BIN}" && -x "${CHECK}" ]] || { echo "prebuilt Linux binaries must be executable" >&2; exit 1; } +else + command -v go >/dev/null 2>&1 || { echo "missing required command: go" >&2; exit 1; } + ( + cd "${ROOT}" + go build -trimpath -buildvcs=true -o "${BIN}" . + go build -trimpath -buildvcs=true -o "${CHECK}" ./cmd/shimmy-artifact-check + ) +fi + +"${CHECK}" -profile python-reactor -module "${WASM}" -manifest "${MANIFEST}" -json >"${TMP}/artifact-check.json" + +( + cd "${ROOT}" + exec env \ + LOG_LEVEL=error \ + FUNCTION_INTERFACE=wasm \ + FUNCTION_WASM_PROFILE=python-reactor \ + FUNCTION_WASM_MODULE="${WASM}" \ + FUNCTION_WASM_MANIFEST="${MANIFEST}" \ + FUNCTION_WASM_PYTHON_SCRIPT="${EVALUATOR}" \ + FUNCTION_WASM_PYTHON_LIFECYCLE=snapshot \ + FUNCTION_MAX_PROCS=1 \ + FUNCTION_WORKER_SEND_TIMEOUT=30s \ + "${BIN}" serve --host "${HOST}" --port "${PORT}" +) >"${LOG}" 2>&1 & +SERVER_PID="$!" + +ready=false +for _ in $(seq 1 150); do + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then + echo "Shimmy exited during startup" >&2 + sed -n '1,200p' "${LOG}" >&2 + exit 1 + fi + if curl -fsS "${BASE_URL}/health" >/dev/null 2>&1; then + ready=true + break + fi + sleep 0.2 +done +[[ "${ready}" == true ]] || { echo "Shimmy did not become ready" >&2; sed -n '1,200p' "${LOG}" >&2; exit 1; } + +request() { + local command="$1" + local body="$2" + curl -fsS -X POST "${BASE_URL}/" \ + -H 'Content-Type: application/json' \ + -H "Command: ${command}" \ + --data "${body}" +} + +EVAL_OK="$(request eval '{"response":"42","answer":"42","params":{"tolerance":0}}')" +EVAL_BAD="$(request eval '{"response":"41","answer":"42","params":{"tolerance":0}}')" +PREVIEW="$(request preview '{"response":"41","params":{}}')" + +EVAL_OK="${EVAL_OK}" EVAL_BAD="${EVAL_BAD}" PREVIEW="${PREVIEW}" python3 - <<'PY' +import json +import os + +def result(name): + body = json.loads(os.environ[name]) + if "error" in body: + raise SystemExit(f"{name} returned an error: {body['error']}") + return body["result"] + +ok = result("EVAL_OK") +bad = result("EVAL_BAD") +preview = result("PREVIEW") +checks = [ + (ok.get("is_correct") is True, "correct eval result"), + (bad.get("is_correct") is False, "incorrect eval result"), + (preview.get("preview") == "submitted: 41", "preview result"), + (ok.get("invocation_count") == 1, "first request starts from prepared state"), + (bad.get("invocation_count") == 1, "second request is reset"), + (preview.get("invocation_count") == 1, "preview request is reset"), +] +failed = [label for passed, label in checks if not passed] +if failed: + raise SystemExit("failed checks: " + ", ".join(failed)) +print(json.dumps({"eval_correct": ok, "eval_incorrect": bad, "preview": preview}, sort_keys=True)) +PY + +printf 'PASS: Linux Python Reactor HTTP E2E\n' +printf 'configuration: FUNCTION_INTERFACE=wasm FUNCTION_WASM_PROFILE=python-reactor FUNCTION_WASM_PYTHON_LIFECYCLE=snapshot\n' diff --git a/tests/e2e/python-reactor/evaluator.py b/tests/e2e/python-reactor/evaluator.py new file mode 100644 index 0000000..f0e04ff --- /dev/null +++ b/tests/e2e/python-reactor/evaluator.py @@ -0,0 +1,45 @@ +"""Linux E2E fixture shaped like a Lambda Feedback evaluator. + +The evaluator owns its public functions and the thin dispatch adapter. Shimmy +only passes the command and validated request payload. +""" + +_invocation_count = 0 + + +def evaluation_function(response, answer, params): + global _invocation_count + _invocation_count += 1 + tolerance = float((params or {}).get("tolerance", 0.0)) + actual = float(response) + expected = float(answer) + is_correct = abs(actual - expected) <= tolerance + return { + "is_correct": is_correct, + "feedback": "correct" if is_correct else "incorrect", + "invocation_count": _invocation_count, + } + + +def preview_function(response, params): + global _invocation_count + _invocation_count += 1 + return { + "preview": f"submitted: {response}", + "invocation_count": _invocation_count, + } + + +def dispatch(method, payload): + if method == "eval": + return evaluation_function( + payload.get("response"), + payload.get("answer"), + payload.get("params", {}), + ) + if method == "preview": + return preview_function( + payload.get("response"), + payload.get("params", {}), + ) + raise LookupError("unsupported method: " + str(method)) From ea13dd430cc8b273e77d16a5715497973c86b496 Mon Sep 17 00:00:00 2001 From: bkmashiro <53376445+bkmashiro@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:08:52 +0100 Subject: [PATCH 3/4] feat(examples): add safe Python Reactor evaluator --- README.md | 23 +- examples/safe-eval-python/README.md | 268 +++++++++++++ examples/safe-eval-python/requests/demo.json | 7 + .../requests/io-tests-fail.json | 17 + .../requests/io-tests-pass.json | 18 + .../safe-eval-python/requests/numpy-core.json | 8 + .../requests/preview-blocked.json | 4 + examples/safe-eval-python/requests/sympy.json | 8 + .../safe-eval-python/requests/unit-tests.json | 8 + examples/safe-eval-python/safe_eval.py | 360 ++++++++++++++++++ examples/safe-eval-python/safe_eval_test.py | 158 ++++++++ examples/safe-eval-python/serve.sh | 87 +++++ examples/safe-eval-python/try.sh | 93 +++++ scripts/e2e-safe-eval-python.sh | 171 +++++++++ 14 files changed, 1227 insertions(+), 3 deletions(-) create mode 100644 examples/safe-eval-python/README.md create mode 100644 examples/safe-eval-python/requests/demo.json create mode 100644 examples/safe-eval-python/requests/io-tests-fail.json create mode 100644 examples/safe-eval-python/requests/io-tests-pass.json create mode 100644 examples/safe-eval-python/requests/numpy-core.json create mode 100644 examples/safe-eval-python/requests/preview-blocked.json create mode 100644 examples/safe-eval-python/requests/sympy.json create mode 100644 examples/safe-eval-python/requests/unit-tests.json create mode 100644 examples/safe-eval-python/safe_eval.py create mode 100644 examples/safe-eval-python/safe_eval_test.py create mode 100755 examples/safe-eval-python/serve.sh create mode 100755 examples/safe-eval-python/try.sh create mode 100755 scripts/e2e-safe-eval-python.sh diff --git a/README.md b/README.md index e8963d4..055bc9e 100644 --- a/README.md +++ b/README.md @@ -202,6 +202,15 @@ The shim keeps the evaluation function running as a persistent process and commu | `tcp` | Raw TCP connection. | | `ws` | WebSocket connection. Experimental — custom dialer configuration is not yet supported. | +Generic WASM and Python Reactor are explicit opt-in execution paths. See +[Execution paths](docs/execution-paths.md) for their environment contracts, +lifecycle behavior, and compatibility boundaries. + +To try Python Reactor without assembling requests by hand, follow the +[`safe-eval-python` three-command quick start](examples/safe-eval-python/README.md#start-here-first-successful-evaluation). +It includes runnable base, NumPy, and SymPy fixtures plus both passing and +failing student-code examples. + The shim injects the following environment variables into the evaluation function process so it can identify the transport it should listen on: | Variable | Value | @@ -246,16 +255,24 @@ For example, a Wolfram Language evaluation function in `evaluation.wl` would be wolframscript -file evaluation.wl /tmp/shimmy/abc/request-data-123 /tmp/shimmy/abc/response-data-456 ``` -### Sandboxed Execution (Linux only, experimental) +### Sandboxed Execution (Linux host/container only, experimental) -Shimmy can wrap each worker process in an [nsjail](https://github.com/google/nsjail) sandbox to safely execute arbitrary, untrusted code. The sandbox provides: +On supported Linux hosts, Shimmy can wrap each worker process in an [nsjail](https://github.com/google/nsjail) sandbox to execute untrusted code with an additional OS boundary. The sandbox provides: - **Filesystem confinement** — the worker can only access explicitly bind-mounted paths - **Resource limits** — CPU time, memory, and file descriptor caps - **Network isolation** — optional; disables all outbound connections - **Unprivileged UID** — worker runs as `nobody` (uid 65534) inside the jail -Sandboxing requires Linux and the `nsjail` binary. The Docker image built from the project's `Dockerfile` includes nsjail at `/usr/sbin/nsjail`. On the host, install it with `sudo apt install nsjail` (Ubuntu 22.04+) or build from source. +Sandboxing requires Linux, the `nsjail` binary, and permission to create the +required namespaces/capabilities. The Docker image built from the project's +`Dockerfile` includes nsjail at `/usr/sbin/nsjail`. On the host, install it with +`sudo apt install nsjail` (Ubuntu 22.04+) or build from source. + +> **AWS Lambda:** Lambda does not grant the namespace/capability controls needed +> to enable this nsjail path. Shipping the binary in a Lambda container image +> does not make it an available security boundary. Use the in-process WASM +> execution profiles for Lambda-compatible isolation. Enable sandboxing with `--sandbox` and configure it with the flags below: diff --git a/examples/safe-eval-python/README.md b/examples/safe-eval-python/README.md new file mode 100644 index 0000000..d0699c0 --- /dev/null +++ b/examples/safe-eval-python/README.md @@ -0,0 +1,268 @@ +# `safe-eval-python` Reactor example + +This example provides a small `demo` / `io_test` / `unit_test` evaluator for +student Python. It is selected at the **Shimmy backend boundary** and runs inside +the Python Reactor WASM profile; it does not adapt the Linux +`evaluatePython` implementation and does not start CPython or Node subprocesses. + +```text +Shimmy HTTP + → wazero + → verified Python Reactor artifact + → safe_eval.py + → student code +``` + +## Start here: first successful evaluation + +You need a Shimmy Python Reactor artifact and its matching manifest. Producer +artifacts are immutable CI outputs rather than Git blobs: obtain both files from +the same trusted Producer build, verify the bundle's published checksums and +provenance, and keep them together. + +From the repository root, start the evaluator: + +```bash +examples/safe-eval-python/serve.sh \ + /path/to/shimmy-python-runtime-base.wasm \ + /path/to/manifest.json +``` + +The launcher validates the artifact/manifest contract before starting Shimmy. +It uses `go run` by default, so contributors do not need a preinstalled Shimmy +binary. It requires Bash, Python 3, and curl; Go is only required when the two +prebuilt Shimmy binaries are not supplied. In another terminal, run the guided examples: + +```bash +examples/safe-eval-python/try.sh base +``` + +This sends real HTTP requests for: + +1. captured demo output; +2. passing public and hidden input/output tests; +3. a failing test and its student-facing feedback; +4. evaluator-defined unit tests; and +5. preview rejection of a blocked host capability. + +Every response is printed and checked. `try.sh` waits up to 90 seconds for the +listener, so it can be started while the Reactor is still preparing. The command +exits non-zero if the running system does not match the documented contract. + +For a richer artifact, use the same flow and name its profile when trying it: + +```bash +examples/safe-eval-python/serve.sh /path/to/numpy-core.wasm /path/to/manifest.json +examples/safe-eval-python/try.sh numpy-core + +examples/safe-eval-python/serve.sh /path/to/sympy.wasm /path/to/manifest.json +examples/safe-eval-python/try.sh sympy +``` + +The onboarding launcher uses a 5-second worker deadline for `base` and +`numpy-core`, and 30 seconds for SymPy's heavier first import. Override it with +`SHIMMY_SAFE_EVAL_TIMEOUT`. These are demonstration defaults, not production +SLOs: measure the chosen profile on the deployment platform and set the shortest +deadline that supports legitimate exercises. + +The request bodies are ordinary JSON files under [`requests/`](requests/). +Copy one and change `response`, `mode`, and `tests` to prototype a real exercise; +no client SDK is required. + +## What to hand to another team + +The smallest useful handoff bundle is: + +```text +shimmy-safe-eval/ +├── shimmy +├── shimmy-artifact-check +├── runtime.wasm +├── manifest.json +├── SHA256SUMS +└── examples/safe-eval-python/ + ├── safe_eval.py + ├── serve.sh + ├── try.sh + └── requests/ +``` + +Set `SHIMMY_BIN` and `SHIMMY_ARTIFACT_CHECK_BIN` to the two shipped binaries; +then `serve.sh` needs no Go toolchain. The deployment owner must still verify +the bundle's signature/provenance and apply platform memory, concurrency, and +request-deadline policy. Do not give users a loose WASM file without its exact +manifest and provenance receipt. + +### Keep the roles separate + +| Role | Starts from | Usually changes | Must not control | +|---|---|---|---| +| Platform owner | `serve.sh`, artifact, manifest | deployment paths, signatures, memory/concurrency/deadlines | per-request capability expansion | +| Evaluator author | `safe_eval.py` and its tests | trusted grading modes and fixed limits | artifact provenance or host mounts | +| Exercise author | a file in `requests/` | student starter code, public/hidden tests, expected output | trusted evaluator source or runtime limits | +| Student/client | HTTP `response` field | submitted Python | tests, manifest, filesystem/network policy | + +For a first workshop, the platform owner starts one `base` instance and runs +`try.sh` once. Exercise authors then copy `io-tests-pass.json` or +`unit-tests.json`; they should not need to understand WASI or modify deployment +environment variables. Move to `numpy-core` or `sympy` only when an exercise +actually requires those packages. + +`serve.sh` is a contributor/onboarding launcher. Production should use the same +validated inputs with a pinned Shimmy binary and platform-managed process, +logging, authentication, resource limits, and artifact provenance policy. + +## Why this path + +AWS Lambda cannot grant the namespaces or capabilities required to make nsjail +a usable runtime boundary. A Python engine that exposes host-process or +JavaScript bridges is likewise not the capability boundary used by this example. + +The Reactor path works within Lambda's normal process constraints: + +- wazero provides the host boundary; +- `FUNCTION_WASM_ALLOWED_PATHS` must remain empty, so no host directory is + mounted into WASI; +- the selected artifact and manifest are verified before execution; +- request deadlines close a non-terminating WASM module; +- snapshot lifecycle restores prepared memory after every request; +- a failed or timed-out snapshot slot is closed and replaced; +- code, input, test count, and output have explicit evaluator limits. + +The AST checks in `safe_eval.py` provide early feedback and reduce accidental +misuse. They are not claimed as the sandbox. The WASM capability boundary, +request deadline, memory limit, and state reset are the security controls. + +## Backend selection + +Use a signed Producer artifact and its matching manifest: + +```bash +export FUNCTION_INTERFACE=wasm +export FUNCTION_WASM_PROFILE=python-reactor +export FUNCTION_WASM_MODULE=/opt/shimmy/runtime/shimmy-python-runtime-base.wasm +export FUNCTION_WASM_MANIFEST=/opt/shimmy/runtime/shimmy-python-runtime-base.manifest.json +export FUNCTION_WASM_PYTHON_SCRIPT="$PWD/examples/safe-eval-python/safe_eval.py" +export FUNCTION_WASM_PYTHON_LIFECYCLE=snapshot +export FUNCTION_WASM_ALLOWED_PATHS= +export FUNCTION_MAX_PROCS=1 +export FUNCTION_WORKER_SEND_TIMEOUT=5s + +shimmy serve --host 127.0.0.1 --port 8080 +``` + +The evaluator is below the 1 MiB trusted-script limit and only uses the standard +library. Package availability comes from the manifest-validated artifact profile, never +from runtime pip or network installation: + +| Student-code requirement | Backend/artifact | +|---|---| +| Standard library | Python Reactor `base` | +| NumPy | Python Reactor `numpy-core` | +| SymPy + mpmath | Python Reactor `sympy` | +| Existing trusted evaluator requiring SciPy or subprocesses | Existing RPC/container backend (outside this example) | + +Do not silently fall back between these paths. Switching the module, manifest, +and runner is deployment configuration. + +Runtime manifest validation establishes digest, ABI, imports/exports, and +capability consistency. Artifact authenticity, trusted Producer commit policy, +and digital-signature verification remain deployment-system responsibilities. + +## Request examples + +Shimmy's `eval` schema requires a non-null `answer`; use an empty string when a +mode does not need one. The `preview` schema does not accept `answer`. + +### Demo + +```bash +curl -sS -X POST http://127.0.0.1:8080/ \ + -H 'Content-Type: application/json' -H 'Command: eval' \ + --data '{"response":"print(6 * 7)","answer":"","params":{"mode":"demo"}}' +``` + +Demo returns captured stdout and `is_correct: false`; it displays execution +output but does not claim a pass condition. + +### Input/output tests + +```bash +curl -sS -X POST http://127.0.0.1:8080/ \ + -H 'Content-Type: application/json' -H 'Command: eval' \ + --data '{ + "response":"n = int(input())\nprint(n * n)", + "answer":"", + "params":{"mode":"io_test","tests":[ + {"input":"5\n","expected_output":"25\n"}, + {"input":"3\n","expected_output":"9\n","hidden":true} + ]} + }' +``` + +Each test receives a fresh Python namespace. Hidden test details omit actual and +expected output. An `inject` object can initialize variables before execution: + +```json +{"inject":{"n":5},"expected_output":"25\n"} +``` + +### Unit tests + +The example intentionally implements a small contract: zero-argument functions +whose names begin with `test_`; a failed `assert` fails that test. + +```bash +curl -sS -X POST http://127.0.0.1:8080/ \ + -H 'Content-Type: application/json' -H 'Command: eval' \ + --data '{ + "response":"def square(n):\n return n * n", + "answer":"", + "params":{"mode":"unit_test","test_code":"def test_square():\n assert square(5) == 25"} + }' +``` + +### Preview + +```bash +curl -sS -X POST http://127.0.0.1:8080/ \ + -H 'Content-Type: application/json' -H 'Command: preview' \ + --data '{"response":"import socket","params":{}}' +``` + +## Built-in evaluator limits + +The limits are constants in the trusted script and cannot be raised by request +parameters: + +| Limit | Value | +|---|---:| +| Student code | 64 KiB | +| Captured stdout/stderr retained in memory per stream/execution | 64 KiB, enforced while writing | +| Input per test | 64 KiB | +| Tests per request | 32 | + +The deployment additionally controls the Reactor memory-page limit and Shimmy +request deadline. Keep the HTTP/worker deadline short enough to bound infinite +loops and long enough for the selected profile's normal work. + +## Verification + +Host-side behavior tests: + +```bash +python3 -m unittest examples/safe-eval-python/safe_eval_test.py -v +``` + +Full Linux path with a real Producer artifact: + +```bash +SHIMMY_PYTHON_REACTOR_WASM=/path/to/base.wasm \ +SHIMMY_PYTHON_REACTOR_MANIFEST=/path/to/base.manifest.json \ +scripts/e2e-safe-eval-python.sh +``` + +The E2E covers all three modes, preview rejection, timeout of an infinite loop, +and successful recovery through a replacement snapshot slot. No Docker, +privileged Lambda configuration, runtime package installation, or nsjail is +required. diff --git a/examples/safe-eval-python/requests/demo.json b/examples/safe-eval-python/requests/demo.json new file mode 100644 index 0000000..7857f5e --- /dev/null +++ b/examples/safe-eval-python/requests/demo.json @@ -0,0 +1,7 @@ +{ + "response": "print(6 * 7)", + "answer": "", + "params": { + "mode": "demo" + } +} diff --git a/examples/safe-eval-python/requests/io-tests-fail.json b/examples/safe-eval-python/requests/io-tests-fail.json new file mode 100644 index 0000000..caac6d8 --- /dev/null +++ b/examples/safe-eval-python/requests/io-tests-fail.json @@ -0,0 +1,17 @@ +{ + "response": "print(input().strip().upper())", + "answer": "", + "params": { + "mode": "io_test", + "tests": [ + { + "input": "hello\n", + "expected_output": "HELLO\n" + }, + { + "input": "shimmy\n", + "expected_output": "NOT-SHIMMY\n" + } + ] + } +} diff --git a/examples/safe-eval-python/requests/io-tests-pass.json b/examples/safe-eval-python/requests/io-tests-pass.json new file mode 100644 index 0000000..eaec88f --- /dev/null +++ b/examples/safe-eval-python/requests/io-tests-pass.json @@ -0,0 +1,18 @@ +{ + "response": "n = int(input())\nprint(n * n)", + "answer": "", + "params": { + "mode": "io_test", + "tests": [ + { + "input": "5\n", + "expected_output": "25\n" + }, + { + "input": "3\n", + "expected_output": "9\n", + "hidden": true + } + ] + } +} diff --git a/examples/safe-eval-python/requests/numpy-core.json b/examples/safe-eval-python/requests/numpy-core.json new file mode 100644 index 0000000..90b8d8a --- /dev/null +++ b/examples/safe-eval-python/requests/numpy-core.json @@ -0,0 +1,8 @@ +{ + "response": "import numpy as np\n\ndef vector_norm(values):\n return float(np.linalg.norm(np.array(values)))", + "answer": "", + "params": { + "mode": "unit_test", + "test_code": "def test_vector_norm():\n assert abs(vector_norm([3, 4]) - 5.0) < 1e-9" + } +} diff --git a/examples/safe-eval-python/requests/preview-blocked.json b/examples/safe-eval-python/requests/preview-blocked.json new file mode 100644 index 0000000..d4f46c3 --- /dev/null +++ b/examples/safe-eval-python/requests/preview-blocked.json @@ -0,0 +1,4 @@ +{ + "response": "import socket\nsocket.create_connection(('example.com', 80))", + "params": {} +} diff --git a/examples/safe-eval-python/requests/sympy.json b/examples/safe-eval-python/requests/sympy.json new file mode 100644 index 0000000..7a0a703 --- /dev/null +++ b/examples/safe-eval-python/requests/sympy.json @@ -0,0 +1,8 @@ +{ + "response": "import sympy as sp\n\ndef derivative_at_two():\n x = sp.symbols('x')\n return sp.diff(x ** 3, x).subs(x, 2)", + "answer": "", + "params": { + "mode": "unit_test", + "test_code": "def test_derivative():\n assert derivative_at_two() == 12" + } +} diff --git a/examples/safe-eval-python/requests/unit-tests.json b/examples/safe-eval-python/requests/unit-tests.json new file mode 100644 index 0000000..7991bfa --- /dev/null +++ b/examples/safe-eval-python/requests/unit-tests.json @@ -0,0 +1,8 @@ +{ + "response": "def square(n):\n return n * n", + "answer": "", + "params": { + "mode": "unit_test", + "test_code": "def test_positive():\n assert square(5) == 25\n\ndef test_negative():\n assert square(-4) == 16" + } +} diff --git a/examples/safe-eval-python/safe_eval.py b/examples/safe-eval-python/safe_eval.py new file mode 100644 index 0000000..ff79258 --- /dev/null +++ b/examples/safe-eval-python/safe_eval.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import ast +import builtins +import contextlib +import io +import json +import traceback +from typing import Any + +DEFAULT_LIMITS = { + "max_output_bytes": 64 * 1024, + "max_code_bytes": 64 * 1024, + "max_tests": 32, + "max_input_bytes": 64 * 1024, +} + +_BLOCKED_MODULES = { + "builtins", + "ctypes", + "http", + "importlib", + "js", + "micropip", + "multiprocessing", + "os", + "pathlib", + "pickle", + "pyodide", + "requests", + "shutil", + "socket", + "subprocess", + "sys", + "threading", + "urllib", +} +_BLOCKED_CALLS = {"compile", "eval", "exec", "open", "__import__"} + + +class ValidationError(ValueError): + pass + + +class _SafetyVisitor(ast.NodeVisitor): + def __init__(self) -> None: + self.violations: list[str] = [] + + def visit_Import(self, node: ast.Import) -> None: + for alias in node.names: + root = alias.name.split(".", 1)[0] + if root in _BLOCKED_MODULES: + self.violations.append(f"import of '{root}' is not allowed") + self.generic_visit(node) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + if node.module: + root = node.module.split(".", 1)[0] + if root in _BLOCKED_MODULES: + self.violations.append(f"import of '{root}' is not allowed") + self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> None: + if isinstance(node.func, ast.Name) and node.func.id in _BLOCKED_CALLS: + self.violations.append(f"use of '{node.func.id}()' is not allowed") + self.generic_visit(node) + + def visit_Attribute(self, node: ast.Attribute) -> None: + if node.attr.startswith("__") and node.attr.endswith("__"): + self.violations.append(f"dunder attribute '{node.attr}' is not allowed") + self.generic_visit(node) + + +def validate_source(source: str) -> list[str]: + try: + tree = ast.parse(source) + except SyntaxError as exc: + line = exc.lineno or 0 + return [f"SyntaxError: {exc.msg} (line {line})"] + visitor = _SafetyVisitor() + visitor.visit(tree) + return visitor.violations + + +def _bounded_text(value: Any, name: str, limit: int) -> str: + text = "" if value is None else str(value) + if len(text.encode("utf-8")) > limit: + raise ValidationError(f"{name} exceeds {limit} bytes") + return text + + +class _BoundedTextWriter(io.TextIOBase): + """Text sink that never retains more than limit UTF-8 bytes.""" + + def __init__(self, limit: int) -> None: + super().__init__() + self._limit = max(0, limit) + self._buffer = bytearray() + self.truncated = False + + @property + def retained_bytes(self) -> int: + return len(self._buffer) + + def writable(self) -> bool: + return True + + def write(self, value: str) -> int: + if not isinstance(value, str): + raise TypeError("write() argument must be str") + if not value: + return 0 + remaining = self._limit - len(self._buffer) + if remaining <= 0: + self.truncated = True + return len(value) + + offset = 0 + while offset < len(value) and remaining > 0: + chunk = value[offset : offset + min(4096, remaining)] + encoded = chunk.encode("utf-8") + available = remaining + self._buffer.extend(encoded[:available]) + offset += len(chunk) + remaining = self._limit - len(self._buffer) + if len(encoded) > available: + self.truncated = True + break + if offset < len(value): + self.truncated = True + return len(value) + + def getvalue(self) -> str: + if not self.truncated: + return bytes(self._buffer).decode("utf-8", errors="ignore") + return _render_truncated(bytes(self._buffer), self._limit) + + +def _render_truncated(encoded: bytes, limit: int) -> str: + suffix = b"\n[output truncated]" + if limit <= len(suffix): + return suffix[:limit].decode("utf-8", errors="ignore") + budget = limit - len(suffix) + return encoded[:budget].decode("utf-8", errors="ignore") + suffix.decode() + + +def _bounded_traceback(limit: int) -> tuple[str, bool]: + writer = _BoundedTextWriter(limit) + traceback.print_exc(file=writer) + return writer.getvalue(), writer.truncated + + +class _TextBudget: + """Shared UTF-8 budget for dynamic strings in one structured result.""" + + def __init__(self, limit: int) -> None: + self.remaining = max(0, limit) + self.truncated = False + + def take(self, value: str) -> str: + writer = _BoundedTextWriter(self.remaining) + writer.write(value) + rendered = writer.getvalue() + self.remaining -= len(rendered.encode("utf-8")) + self.truncated = self.truncated or writer.truncated + return rendered + + +def _execute(source: str, stdin: str, inject: dict[str, Any], output_limit: int) -> dict[str, Any]: + violations = validate_source(source) + if violations: + return {"ok": False, "kind": "validation", "error": "\n".join(violations), "stdout": "", "stderr": ""} + + input_lines = iter(stdin.splitlines()) + + def safe_input(prompt: str = "") -> str: + if prompt: + print(prompt, end="") + try: + return next(input_lines) + except StopIteration as exc: + raise EOFError("input exhausted") from exc + + namespace: dict[str, Any] = {"__name__": "__student__", **inject} + stdout = _BoundedTextWriter(output_limit) + stderr = _BoundedTextWriter(output_limit) + original_input = builtins.input + try: + builtins.input = safe_input + with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + exec(compile(source, "", "exec"), namespace, namespace) + return { + "ok": True, + "namespace": namespace, + "stdout": stdout.getvalue(), + "stderr": stderr.getvalue(), + "truncated": stdout.truncated or stderr.truncated, + } + except BaseException: + error, traceback_truncated = _bounded_traceback(output_limit) + return { + "ok": False, + "kind": "runtime", + "error": error, + "stdout": stdout.getvalue(), + "stderr": "", + "truncated": stdout.truncated or stderr.truncated or traceback_truncated, + } + finally: + builtins.input = original_input + + +def _demo(code: str, limits: dict[str, int]) -> dict[str, Any]: + run = _execute(code, "", {}, limits["max_output_bytes"]) + if not run["ok"]: + return {"is_correct": False, "feedback": run["error"], "stdout": run["stdout"], "kind": run["kind"]} + return { + "is_correct": False, + "feedback": "Program completed.", + "stdout": run["stdout"], + "stderr": run["stderr"], + "truncated": run["truncated"], + } + + +def _io_test(code: str, tests: list[Any], limits: dict[str, int]) -> dict[str, Any]: + if len(tests) > limits["max_tests"]: + raise ValidationError(f"tests exceeds {limits['max_tests']} entries") + details = [] + passed = 0 + result_budget = _TextBudget(limits["max_output_bytes"]) + for index, raw_test in enumerate(tests, 1): + if not isinstance(raw_test, dict): + raise ValidationError(f"test {index} must be an object") + stdin = _bounded_text(raw_test.get("input", ""), f"test {index} input", limits["max_input_bytes"]) + expected = _bounded_text(raw_test.get("expected_output", ""), f"test {index} expected_output", limits["max_output_bytes"]) + inject = raw_test.get("inject", {}) + if not isinstance(inject, dict): + raise ValidationError(f"test {index} inject must be an object") + run = _execute(code, stdin, inject, limits["max_output_bytes"]) + actual = run["stdout"].rstrip() + correct = bool(run["ok"] and actual == expected.rstrip()) + passed += int(correct) + hidden = bool(raw_test.get("hidden", False)) + detail: dict[str, Any] = {"index": index, "passed": correct, "hidden": hidden} + if not hidden: + detail.update({"actual": result_budget.take(actual), "expected": result_budget.take(expected.rstrip())}) + if not run["ok"]: + detail["error"] = result_budget.take(run["error"] if not hidden else "hidden test failed") + details.append(detail) + total = len(tests) + return { + "is_correct": total > 0 and passed == total, + "feedback": f"{passed}/{total} tests passed.", + "passed": passed, + "total": total, + "tests": details, + "truncated": result_budget.truncated, + } + + +def _unit_test(code: str, test_code: str, limits: dict[str, int]) -> dict[str, Any]: + student = _execute(code, "", {}, limits["max_output_bytes"]) + if not student["ok"]: + return {"is_correct": False, "feedback": student["error"], "kind": student["kind"]} + violations = validate_source(test_code) + if violations: + return {"is_correct": False, "feedback": "\n".join(violations), "kind": "validation"} + + namespace = student["namespace"] + test_stdout = _BoundedTextWriter(limits["max_output_bytes"]) + test_stderr = _BoundedTextWriter(limits["max_output_bytes"]) + try: + with contextlib.redirect_stdout(test_stdout), contextlib.redirect_stderr(test_stderr): + exec(compile(test_code, "", "exec"), namespace, namespace) + except BaseException: + error, _ = _bounded_traceback(limits["max_output_bytes"]) + return {"is_correct": False, "feedback": error, "kind": "test_setup"} + + tests = sorted((name, value) for name, value in namespace.items() if name.startswith("test_") and callable(value)) + if len(tests) > limits["max_tests"]: + raise ValidationError(f"unit tests exceeds {limits['max_tests']} entries") + details = [] + result_budget = _TextBudget(limits["max_output_bytes"]) + for name, test in tests: + try: + with contextlib.redirect_stdout(test_stdout), contextlib.redirect_stderr(test_stderr): + test() + details.append({"name": result_budget.take(name), "passed": True}) + except BaseException: + error, _ = _bounded_traceback(limits["max_output_bytes"]) + details.append({ + "name": result_budget.take(name), + "passed": False, + "error": result_budget.take(error), + }) + passed = sum(int(item["passed"]) for item in details) + return { + "is_correct": bool(details) and passed == len(details), + "feedback": f"{passed}/{len(details)} unit tests passed.", + "passed": passed, + "total": len(details), + "tests": details, + "truncated": result_budget.truncated or test_stdout.truncated or test_stderr.truncated, + } + + +def _dispatch(method: str, payload: dict[str, Any], limits: dict[str, int]) -> dict[str, Any]: + params = payload.get("params") or {} + try: + code = _bounded_text(payload.get("response", ""), "response", limits["max_code_bytes"]) + if method == "preview": + violations = validate_source(code) + result = { + "is_correct": None, + "preview": "Valid Python syntax." if not violations else "\n".join(violations), + "valid": not violations, + } + else: + mode = params.get("mode", "demo") + if mode == "demo": + result = _demo(code, limits) + elif mode == "io_test": + result = _io_test(code, params.get("tests", []), limits) + elif mode == "unit_test": + test_code = _bounded_text(params.get("test_code", payload.get("answer", "")), "test_code", limits["max_code_bytes"]) + result = _unit_test(code, test_code, limits) + else: + raise ValidationError("mode must be demo, io_test, or unit_test") + except ValidationError as exc: + result = {"is_correct": False, "feedback": str(exc), "kind": "validation"} + return result + + +def evaluation_function(response: Any, answer: Any, params: dict[str, Any]) -> dict[str, Any]: + return _dispatch( + "eval", + {"response": response, "answer": answer, "params": dict(params or {})}, + DEFAULT_LIMITS, + ) + + +def preview_function(response: Any, params: dict[str, Any]) -> dict[str, Any]: + return _dispatch( + "preview", + {"response": response, "params": dict(params or {})}, + DEFAULT_LIMITS, + ) + + +def invoke(request_json: str, limits_json: str) -> str: + """Host-CPython test adapter; the Reactor calls the functions above directly.""" + request = json.loads(request_json) + limits = json.loads(limits_json) + result = _dispatch( + request.get("method", "eval"), + request.get("payload") or {}, + limits, + ) + return json.dumps(result, ensure_ascii=False, separators=(",", ":")) diff --git a/examples/safe-eval-python/safe_eval_test.py b/examples/safe-eval-python/safe_eval_test.py new file mode 100644 index 0000000..969c9b4 --- /dev/null +++ b/examples/safe-eval-python/safe_eval_test.py @@ -0,0 +1,158 @@ +import importlib.util +import json +import pathlib +import unittest + +MODULE_PATH = pathlib.Path(__file__).with_name("safe_eval.py") +SPEC = importlib.util.spec_from_file_location("safe_eval", MODULE_PATH) +assert SPEC is not None +SAFE_EVAL = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +SPEC.loader.exec_module(SAFE_EVAL) + +LIMITS = { + "max_code_bytes": 65536, + "max_input_bytes": 65536, + "max_output_bytes": 65536, + "max_tests": 32, +} + + +def invoke(response, params=None, answer=None, method="eval", limits=None): + request = { + "method": method, + "payload": {"response": response, "answer": answer, "params": params or {}}, + } + return json.loads(SAFE_EVAL.invoke(json.dumps(request), json.dumps(limits or LIMITS))) + + +class SafeEvalTest(unittest.TestCase): + def test_reactor_entrypoints_are_directly_callable(self): + evaluated = SAFE_EVAL.evaluation_function( + "print(6 * 7)", "", {"mode": "demo"} + ) + previewed = SAFE_EVAL.preview_function("import socket", {}) + self.assertEqual(evaluated["stdout"], "42\n") + self.assertFalse(previewed["valid"]) + + def test_trusted_script_fits_reactor_payload_bound(self): + self.assertLess(MODULE_PATH.stat().st_size, 1024 * 1024) + + def test_demo_captures_stdout(self): + result = invoke("print(6 * 7)", {"mode": "demo"}) + self.assertFalse(result["is_correct"]) + self.assertEqual(result["stdout"], "42\n") + + def test_io_tests_use_fresh_namespaces_and_hide_hidden_values(self): + result = invoke( + "value = int(input())\nprint(value * value)", + { + "mode": "io_test", + "tests": [ + {"input": "5\n", "expected_output": "25\n"}, + {"input": "3\n", "expected_output": "8\n", "hidden": True}, + ], + }, + ) + self.assertFalse(result["is_correct"]) + self.assertEqual(result["feedback"], "1/2 tests passed.") + self.assertNotIn("actual", result["tests"][1]) + self.assertNotIn("expected", result["tests"][1]) + + def test_injected_io_test(self): + result = invoke( + "print(n + 1)", + {"mode": "io_test", "tests": [{"inject": {"n": 4}, "expected_output": "5\n"}]}, + ) + self.assertTrue(result["is_correct"]) + + def test_unit_test_discovers_plain_test_functions(self): + result = invoke( + "def square(value):\n return value * value", + {"mode": "unit_test", "test_code": "def test_square():\n assert square(5) == 25"}, + ) + self.assertTrue(result["is_correct"]) + self.assertEqual(result["feedback"], "1/1 unit tests passed.") + + def test_preview_reports_blocked_host_capabilities(self): + result = invoke("import js\njs.process.exit(0)", method="preview") + self.assertFalse(result["valid"]) + self.assertIn("import of 'js' is not allowed", result["preview"]) + + def test_runtime_rejects_blocked_host_capabilities(self): + result = invoke("import subprocess\nsubprocess.run(['id'])") + self.assertFalse(result["is_correct"]) + self.assertEqual(result["kind"], "validation") + + def test_code_limit_fails_closed(self): + limits = {**LIMITS, "max_code_bytes": 8} + result = invoke("print('too long')", limits=limits) + self.assertFalse(result["is_correct"]) + self.assertEqual(result["kind"], "validation") + self.assertIn("exceeds 8 bytes", result["feedback"]) + + def test_output_is_truncated(self): + limits = {**LIMITS, "max_output_bytes": 32} + result = invoke("print('x' * (1024 * 1024))", limits=limits) + self.assertTrue(result["truncated"]) + self.assertLessEqual(len(result["stdout"].encode()), 32) + + def test_output_writer_retains_at_most_the_byte_limit(self): + writer = SAFE_EVAL._BoundedTextWriter(31) + writer.write("λ" * (1024 * 1024)) + self.assertEqual(writer.retained_bytes, 31) + self.assertTrue(writer.truncated) + self.assertLessEqual(len(writer.getvalue().encode()), 31) + + def test_unit_test_output_is_bounded_while_running(self): + limits = {**LIMITS, "max_output_bytes": 32} + result = invoke( + "def square(value):\n return value * value", + { + "mode": "unit_test", + "test_code": ( + "print('x' * (1024 * 1024))\n" + "def test_square():\n" + " print('y' * (1024 * 1024))\n" + " assert square(5) == 25" + ), + }, + limits=limits, + ) + self.assertTrue(result["is_correct"]) + def test_io_test_detail_strings_share_one_output_budget(self): + limits = {**LIMITS, "max_output_bytes": 32} + result = invoke( + "print('x' * 32)", + {"mode": "io_test", "tests": [ + {"expected_output": "x" * 32}, + {"expected_output": "x" * 32}, + ]}, + limits=limits, + ) + retained = sum( + len(detail.get(key, "").encode()) + for detail in result["tests"] + for key in ("actual", "expected", "error") + ) + self.assertLessEqual(retained, 32) + self.assertTrue(result["truncated"]) + + def test_unit_test_errors_share_one_output_budget(self): + limits = {**LIMITS, "max_output_bytes": 32} + test_code = "\n".join( + f"def test_{index}():\n raise AssertionError('x' * 1000)" + for index in range(32) + ) + result = invoke("pass", {"mode": "unit_test", "test_code": test_code}, limits=limits) + retained = sum( + len(detail.get(key, "").encode()) + for detail in result["tests"] + for key in ("name", "error") + ) + self.assertLessEqual(retained, 32) + self.assertTrue(result["truncated"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/safe-eval-python/serve.sh b/examples/safe-eval-python/serve.sh new file mode 100755 index 0000000..e152849 --- /dev/null +++ b/examples/safe-eval-python/serve.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +WASM="${1:-${SHIMMY_PYTHON_REACTOR_WASM:-}}" +MANIFEST="${2:-${SHIMMY_PYTHON_REACTOR_MANIFEST:-}}" +PORT="${SHIMMY_SAFE_EVAL_PORT:-8080}" +HOST="${SHIMMY_SAFE_EVAL_HOST:-127.0.0.1}" +EVALUATOR="${SHIMMY_SAFE_EVAL_EVALUATOR:-${ROOT}/examples/safe-eval-python/safe_eval.py}" + +usage() { + cat >&2 <<'EOF' +Usage: + examples/safe-eval-python/serve.sh /path/to/runtime.wasm /path/to/manifest.json + +Or set: + SHIMMY_PYTHON_REACTOR_WASM=/path/to/runtime.wasm + SHIMMY_PYTHON_REACTOR_MANIFEST=/path/to/manifest.json + +Optional: + SHIMMY_SAFE_EVAL_HOST=127.0.0.1 + SHIMMY_SAFE_EVAL_PORT=8080 + SHIMMY_SAFE_EVAL_TIMEOUT=10s # optional; profile-aware default otherwise + SHIMMY_BIN=/path/to/shimmy + SHIMMY_ARTIFACT_CHECK_BIN=/path/to/shimmy-artifact-check +EOF + exit 2 +} + +[[ -n "${WASM}" && -n "${MANIFEST}" ]] || usage +[[ -r "${WASM}" ]] || { echo "runtime artifact is not readable: ${WASM}" >&2; exit 2; } +[[ -r "${MANIFEST}" ]] || { echo "runtime manifest is not readable: ${MANIFEST}" >&2; exit 2; } +[[ -r "${EVALUATOR}" ]] || { echo "evaluator is not readable: ${EVALUATOR}" >&2; exit 2; } +command -v python3 >/dev/null 2>&1 || { echo "Python 3 is required by the quick-start launcher" >&2; exit 2; } + +if [[ -n "${SHIMMY_ARTIFACT_CHECK_BIN:-}" ]]; then + [[ -x "${SHIMMY_ARTIFACT_CHECK_BIN}" ]] || { echo "artifact checker is not executable: ${SHIMMY_ARTIFACT_CHECK_BIN}" >&2; exit 2; } + CHECK_CMD=("${SHIMMY_ARTIFACT_CHECK_BIN}") +else + command -v go >/dev/null 2>&1 || { echo "Go is required unless SHIMMY_ARTIFACT_CHECK_BIN is set" >&2; exit 2; } + CHECK_CMD=(go run ./cmd/shimmy-artifact-check) +fi + +if [[ -n "${SHIMMY_BIN:-}" ]]; then + [[ -x "${SHIMMY_BIN}" ]] || { echo "Shimmy binary is not executable: ${SHIMMY_BIN}" >&2; exit 2; } + SHIMMY_CMD=("${SHIMMY_BIN}") +else + command -v go >/dev/null 2>&1 || { echo "Go is required unless SHIMMY_BIN is set" >&2; exit 2; } + SHIMMY_CMD=(go run .) +fi + +cd "${ROOT}" +echo "Validating Python Reactor artifact and manifest..." +"${CHECK_CMD[@]}" -profile python-reactor -module "${WASM}" -manifest "${MANIFEST}" +ARTIFACT_PROFILE="$(python3 - "${MANIFEST}" <<'PY' +import json +import pathlib +import sys + +print(json.loads(pathlib.Path(sys.argv[1]).read_text()).get("profile", "")) +PY +)" +case "${ARTIFACT_PROFILE}" in + base|numpy-core) DEFAULT_TIMEOUT=5s ;; + sympy) DEFAULT_TIMEOUT=30s ;; + *) echo "unsupported Python Reactor profile in manifest: ${ARTIFACT_PROFILE:-}" >&2; exit 2 ;; +esac +WORKER_TIMEOUT="${SHIMMY_SAFE_EVAL_TIMEOUT:-${DEFAULT_TIMEOUT}}" + +echo +echo "safe-eval-python (${ARTIFACT_PROFILE}) is starting at http://${HOST}:${PORT}" +echo "onboarding worker deadline: ${WORKER_TIMEOUT}" +echo "In another terminal run:" +echo " examples/safe-eval-python/try.sh ${ARTIFACT_PROFILE} http://${HOST}:${PORT}" +echo + +exec env \ + FUNCTION_INTERFACE=wasm \ + FUNCTION_WASM_PROFILE=python-reactor \ + FUNCTION_WASM_MODULE="${WASM}" \ + FUNCTION_WASM_MANIFEST="${MANIFEST}" \ + FUNCTION_WASM_PYTHON_SCRIPT="${EVALUATOR}" \ + FUNCTION_WASM_PYTHON_LIFECYCLE=snapshot \ + FUNCTION_WASM_ALLOWED_PATHS= \ + FUNCTION_MAX_PROCS=1 \ + FUNCTION_WORKER_SEND_TIMEOUT="${WORKER_TIMEOUT}" \ + "${SHIMMY_CMD[@]}" --worker-send-timeout "${WORKER_TIMEOUT}" serve --host "${HOST}" --port "${PORT}" diff --git a/examples/safe-eval-python/try.sh b/examples/safe-eval-python/try.sh new file mode 100755 index 0000000..cac26dc --- /dev/null +++ b/examples/safe-eval-python/try.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +REQUESTS="${ROOT}/examples/safe-eval-python/requests" +PROFILE="${1:-base}" +BASE_URL="${2:-http://127.0.0.1:8080}" +CURL_TIMEOUT=15 +TMP="$(mktemp -d)" +trap 'rm -rf "${TMP}"' EXIT + +case "${PROFILE}" in + base|numpy-core) ;; + sympy) CURL_TIMEOUT=45 ;; + *) echo "profile must be base, numpy-core, or sympy" >&2; exit 2 ;; +esac + +echo "Waiting for ${BASE_URL} ..." +python3 - "${BASE_URL}" <<'PY' +import socket +import sys +import time +import urllib.parse + +url = urllib.parse.urlsplit(sys.argv[1]) +if url.scheme != "http" or not url.hostname: + raise SystemExit("quick start URL must be an http:// URL with a host") +port = url.port or 80 +deadline = time.monotonic() + 90 +while True: + try: + with socket.create_connection((url.hostname, port), timeout=0.5): + break + except OSError as error: + if time.monotonic() >= deadline: + raise SystemExit(f"server did not become ready within 90 seconds: {error}") + time.sleep(0.5) +PY + +post() { + local label="$1" command="$2" request="$3" assertion="$4" + local response="${TMP}/${assertion}.json" + echo + echo "== ${label} ==" + echo "request: ${request#"${ROOT}"/}" + curl --fail-with-body --max-time "${CURL_TIMEOUT}" -sS \ + -X POST "${BASE_URL}/" \ + -H 'Content-Type: application/json' \ + -H "Command: ${command}" \ + --data-binary "@${request}" \ + -o "${response}" + python3 - "${response}" "${assertion}" <<'PY' +import json +import pathlib +import sys + +path = pathlib.Path(sys.argv[1]) +assertion = sys.argv[2] +value = json.loads(path.read_text()) +result = value.get("result", value) + +checks = { + "demo": lambda r: r.get("stdout") == "42\n" and r.get("is_correct") is False, + "io-pass": lambda r: r.get("is_correct") is True and r.get("passed") == 2, + "io-fail": lambda r: r.get("is_correct") is False and r.get("passed") == 1, + "unit": lambda r: r.get("is_correct") is True and r.get("passed") == 2, + "preview": lambda r: r.get("valid") is False and "socket" in r.get("preview", ""), + "profile": lambda r: r.get("is_correct") is True and r.get("passed") == 1, +} +if assertion not in checks or not checks[assertion](result): + print(json.dumps(value, indent=2, sort_keys=True)) + raise SystemExit(f"unexpected response for {assertion}") +print(json.dumps(value, indent=2, sort_keys=True)) +PY +} + +post "demo" eval "${REQUESTS}/demo.json" demo +post "passing I/O tests (including one hidden test)" eval "${REQUESTS}/io-tests-pass.json" io-pass +post "failing I/O test feedback" eval "${REQUESTS}/io-tests-fail.json" io-fail +post "unit tests" eval "${REQUESTS}/unit-tests.json" unit +post "preview rejects a blocked host capability" preview "${REQUESTS}/preview-blocked.json" preview + +case "${PROFILE}" in + numpy-core) + post "NumPy profile" eval "${REQUESTS}/numpy-core.json" profile + ;; + sympy) + post "SymPy profile" eval "${REQUESTS}/sympy.json" profile + ;; +esac + +echo +echo "Quick start completed for profile: ${PROFILE}" diff --git a/scripts/e2e-safe-eval-python.sh b/scripts/e2e-safe-eval-python.sh new file mode 100755 index 0000000..44812cd --- /dev/null +++ b/scripts/e2e-safe-eval-python.sh @@ -0,0 +1,171 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +WASM="${SHIMMY_PYTHON_REACTOR_WASM:?set SHIMMY_PYTHON_REACTOR_WASM to a Producer base artifact}" +MANIFEST="${SHIMMY_PYTHON_REACTOR_MANIFEST:?set SHIMMY_PYTHON_REACTOR_MANIFEST to its manifest.json}" +EVALUATOR="${SHIMMY_SAFE_EVAL_EVALUATOR:-${ROOT}/examples/safe-eval-python/safe_eval.py}" +HOST=127.0.0.1 +TMP="$(mktemp -d "${TMPDIR:-/tmp}/shimmy-safe-eval-python-e2e.XXXXXX")" +PORT="${SHIMMY_E2E_PORT:-}" +BIN="${SHIMMY_E2E_BIN:-${TMP}/shimmy}" +CHECK="${SHIMMY_E2E_ARTIFACT_CHECK_BIN:-${TMP}/shimmy-artifact-check}" +SERVER_PID="" + +cleanup() { + if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then + kill "${SERVER_PID}" 2>/dev/null || true + wait "${SERVER_PID}" 2>/dev/null || true + fi + if [[ -n "${SHIMMY_E2E_SERVER_LOG:-}" && -f "${LOG:-}" ]]; then + cp "${LOG}" "${SHIMMY_E2E_SERVER_LOG}" + fi + rm -rf "${TMP}" +} +trap cleanup EXIT + +for cmd in curl python3; do + command -v "${cmd}" >/dev/null 2>&1 || { echo "missing required command: ${cmd}" >&2; exit 1; } +done +[[ "$(uname -s)" == "Linux" ]] || { echo "safeEvalPython Reactor E2E requires Linux" >&2; exit 1; } +[[ -r "${WASM}" && -r "${MANIFEST}" && -r "${EVALUATOR}" ]] || { echo "artifact, manifest, and evaluator must be readable" >&2; exit 1; } + +if [[ -z "${PORT}" ]]; then + PORT="$(python3 - <<'PY' +import socket +with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + print(sock.getsockname()[1]) +PY +)" +fi +if [[ -z "${SHIMMY_E2E_BIN:-}" || -z "${SHIMMY_E2E_ARTIFACT_CHECK_BIN:-}" ]]; then + [[ -z "${SHIMMY_E2E_BIN:-}" && -z "${SHIMMY_E2E_ARTIFACT_CHECK_BIN:-}" ]] || { + echo "set both SHIMMY_E2E_BIN and SHIMMY_E2E_ARTIFACT_CHECK_BIN" >&2 + exit 1 + } + command -v go >/dev/null 2>&1 || { echo "missing required command: go" >&2; exit 1; } + ( + cd "${ROOT}" + go build -trimpath -buildvcs=true -o "${BIN}" . + go build -trimpath -buildvcs=true -o "${CHECK}" ./cmd/shimmy-artifact-check + ) +fi +[[ -x "${BIN}" && -x "${CHECK}" ]] || { echo "Shimmy binaries must be executable" >&2; exit 1; } + +"${CHECK}" -profile python-reactor -module "${WASM}" -manifest "${MANIFEST}" -json >"${TMP}/artifact-check.json" +LOG="${TMP}/server.log" +( + cd "${ROOT}" + exec env \ + LOG_LEVEL="${SHIMMY_E2E_LOG_LEVEL:-error}" \ + FUNCTION_INTERFACE=wasm \ + FUNCTION_WASM_PROFILE=python-reactor \ + FUNCTION_WASM_MODULE="${WASM}" \ + FUNCTION_WASM_MANIFEST="${MANIFEST}" \ + FUNCTION_WASM_PYTHON_SCRIPT="${EVALUATOR}" \ + FUNCTION_WASM_PYTHON_LIFECYCLE=snapshot \ + FUNCTION_WASM_ALLOWED_PATHS= \ + FUNCTION_MAX_PROCS=1 \ + FUNCTION_WORKER_SEND_TIMEOUT=2s \ + "${BIN}" --worker-send-timeout 2s serve --host "${HOST}" --port "${PORT}" +) >"${LOG}" 2>&1 & +SERVER_PID="$!" +BASE_URL="http://${HOST}:${PORT}" +ready=false +for _ in $(seq 1 300); do + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then + echo "Shimmy exited during Reactor startup" >&2 + python3 - "${LOG}" <<'PY' >&2 +import pathlib, sys +print(pathlib.Path(sys.argv[1]).read_text(errors="replace")) +PY + exit 1 + fi + if curl -fsS "${BASE_URL}/health" >/dev/null 2>&1; then ready=true; break; fi + sleep 0.2 +done +[[ "${ready}" == true ]] || { echo "Shimmy did not become ready" >&2; exit 1; } + +request() { + local output + if ! output="$(curl --fail-with-body -sS -X POST "${BASE_URL}/" \ + -H 'Content-Type: application/json' -H "Command: $1" --data "$2")"; then + printf '%s\n' "${output}" >&2 + python3 - "${LOG}" <<'PY' >&2 +import pathlib, sys +print(pathlib.Path(sys.argv[1]).read_text(errors="replace")) +PY + return 1 + fi + printf '%s' "${output}" +} + +DEMO="$(request eval '{"response":"print(6 * 7)","answer":"","params":{"mode":"demo"}}')" +IO="$(request eval '{"response":"try:\n n\nexcept NameError:\n n = int(input())\nprint(n * n)","answer":"","params":{"mode":"io_test","tests":[{"input":"5\n","expected_output":"25\n"},{"inject":{"n":3},"expected_output":"9\n","hidden":true}]}}')" +UNIT="$(request eval '{"response":"def square(n):\n return n * n","answer":"","params":{"mode":"unit_test","test_code":"def test_square():\n assert square(5) == 25"}}')" +PREVIEW="$(request preview '{"response":"import socket","params":{}}')" +BLOCKED="$(request eval '{"response":"import socket","answer":"","params":{"mode":"demo"}}')" + +python3 - "${DEMO}" "${IO}" "${UNIT}" "${PREVIEW}" "${BLOCKED}" <<'PY' +import json, sys + +def unwrap(raw): + value = json.loads(raw) + return value.get("result", value) + +demo, io_result, unit, preview, blocked = map(unwrap, sys.argv[1:]) +assert demo["stdout"] == "42\n" and demo["is_correct"] is False +assert io_result["is_correct"] is True and io_result["passed"] == 2 +assert io_result["tests"][1]["hidden"] is True and io_result["tests"][1]["passed"] is True +assert "actual" not in io_result["tests"][1] and "expected" not in io_result["tests"][1] +assert unit["is_correct"] is True and unit["tests"] == [{"name": "test_square", "passed": True}] +assert preview["valid"] is False and "socket" in preview["preview"] +assert blocked["is_correct"] is False and blocked["kind"] == "validation" +print(json.dumps({"demo": demo, "io_test": io_result, "unit_test": unit, "preview": preview}, sort_keys=True)) +PY + +TIMEOUT_BODY="${TMP}/timeout.json" +TIMEOUT_META="$(curl --max-time 10 -sS -o "${TIMEOUT_BODY}" -w '%{http_code} %{time_total}' -X POST "${BASE_URL}/" \ + -H 'Content-Type: application/json' -H 'Command: eval' \ + --data '{"response":"while True:\n pass","answer":"","params":{"mode":"demo"}}')" +read -r TIMEOUT_STATUS TIMEOUT_SECONDS <<<"${TIMEOUT_META}" +case "${TIMEOUT_STATUS}" in + 5??) ;; + *) echo "expected timeout 5xx, got ${TIMEOUT_STATUS}" >&2; exit 1 ;; +esac +echo "timeout_http_status=${TIMEOUT_STATUS}" +echo "timeout_seconds=${TIMEOUT_SECONDS}" +python3 - "${TIMEOUT_BODY}" <<'PY' +import json, pathlib, sys +body = json.loads(pathlib.Path(sys.argv[1]).read_text()) +text = json.dumps(body).lower() +assert any(term in text for term in ("deadline", "timeout", "closed")), body +PY + +RECOVERY_BODY="${TMP}/recovery.json" +RECOVERY_READY=0 +RECOVERY_ATTEMPTS=0 +for _ in $(seq 1 12); do + RECOVERY_ATTEMPTS=$((RECOVERY_ATTEMPTS + 1)) + RECOVERY_STATUS="$(curl --max-time 10 -sS -o "${RECOVERY_BODY}" -w '%{http_code}' -X POST "${BASE_URL}/" \ + -H 'Content-Type: application/json' -H 'Command: eval' \ + --data '{"response":"print(7 * 6)","answer":"","params":{"mode":"demo"}}' || true)" + if [[ "${RECOVERY_STATUS}" == 200 ]]; then + RECOVERY_READY=1 + break + fi + sleep 1 +done +test "${RECOVERY_READY}" = 1 +python3 - "${RECOVERY_BODY}" <<'PY' +import json, pathlib, sys +value = json.loads(pathlib.Path(sys.argv[1]).read_text()) +value = value.get("result", value) +assert value["stdout"] == "42\n" +PY + +echo "timeout_recovery_attempts=${RECOVERY_ATTEMPTS}" + +echo "timeout_recovery=PASS" +echo "safe_eval_python_reactor_e2e=PASS" From 4f1c38265a6c29e940c5161a66adec218fc4b6c4 Mon Sep 17 00:00:00 2001 From: bkmashiro <53376445+bkmashiro@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:51:59 +0100 Subject: [PATCH 4/4] fix(runtime): address WASM review feedback --- cmd/root.go | 19 ++++++++++- cmd/root_test.go | 37 +++++++++++++++++++++ examples/safe-eval-python/safe_eval.py | 2 +- examples/safe-eval-python/safe_eval_test.py | 11 ++++++ internal/execution/wasm/adapter.go | 11 ++++++ internal/execution/wasm/adapter_test.go | 12 +++++++ internal/execution/wasm/dispatcher_test.go | 2 +- internal/execution/wasm/supervisor.go | 14 ++++++-- 8 files changed, 102 insertions(+), 6 deletions(-) create mode 100644 cmd/root_test.go create mode 100644 internal/execution/wasm/adapter_test.go diff --git a/cmd/root.go b/cmd/root.go index 690258f..af32996 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "strings" "time" "github.com/urfave/cli/v2" @@ -359,6 +360,22 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { if err != nil { return config.Config{}, err } + if err := validateRootConfig(cfg, os.Getenv("FUNCTION_WASM_MODULE")); err != nil { + return config.Config{}, err + } + + return cfg, nil +} - return cfg, err +func validateRootConfig(cfg config.Config, wasmModule string) error { + if strings.TrimSpace(cfg.Runtime.Supervisor.StartParams.Cmd) != "" { + return nil + } + if cfg.Runtime.Supervisor.IO.Interface == "wasm" { + if strings.TrimSpace(wasmModule) != "" { + return nil + } + return fmt.Errorf("wasm interface requires --command or FUNCTION_WASM_MODULE") + } + return fmt.Errorf("%s interface requires --command", cfg.Runtime.Supervisor.IO.Interface) } diff --git a/cmd/root_test.go b/cmd/root_test.go new file mode 100644 index 0000000..71314ed --- /dev/null +++ b/cmd/root_test.go @@ -0,0 +1,37 @@ +package cmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/execution/supervisor" +) + +func TestValidateRootConfigRequiresCommandForProcessInterfaces(t *testing.T) { + for _, iface := range []supervisor.IOInterface{supervisor.RpcIO, supervisor.FileIO} { + t.Run(string(iface), func(t *testing.T) { + var cfg config.Config + cfg.Runtime.Supervisor.IO.Interface = iface + err := validateRootConfig(cfg, "") + require.Error(t, err) + assert.Contains(t, err.Error(), "--command") + }) + } +} + +func TestValidateRootConfigAcceptsWasmModuleOverride(t *testing.T) { + var cfg config.Config + cfg.Runtime.Supervisor.IO.Interface = supervisor.WasmIO + require.NoError(t, validateRootConfig(cfg, "/opt/evaluator.wasm")) +} + +func TestValidateRootConfigRequiresWasmModulePath(t *testing.T) { + var cfg config.Config + cfg.Runtime.Supervisor.IO.Interface = supervisor.WasmIO + err := validateRootConfig(cfg, "") + require.Error(t, err) + assert.Contains(t, err.Error(), "FUNCTION_WASM_MODULE") +} diff --git a/examples/safe-eval-python/safe_eval.py b/examples/safe-eval-python/safe_eval.py index ff79258..ef53fbd 100644 --- a/examples/safe-eval-python/safe_eval.py +++ b/examples/safe-eval-python/safe_eval.py @@ -203,7 +203,7 @@ def safe_input(prompt: str = "") -> str: "kind": "runtime", "error": error, "stdout": stdout.getvalue(), - "stderr": "", + "stderr": stderr.getvalue(), "truncated": stdout.truncated or stderr.truncated or traceback_truncated, } finally: diff --git a/examples/safe-eval-python/safe_eval_test.py b/examples/safe-eval-python/safe_eval_test.py index 969c9b4..74206db 100644 --- a/examples/safe-eval-python/safe_eval_test.py +++ b/examples/safe-eval-python/safe_eval_test.py @@ -1,6 +1,7 @@ import importlib.util import json import pathlib +import sys import unittest MODULE_PATH = pathlib.Path(__file__).with_name("safe_eval.py") @@ -43,6 +44,16 @@ def test_demo_captures_stdout(self): self.assertFalse(result["is_correct"]) self.assertEqual(result["stdout"], "42\n") + def test_execute_preserves_bounded_stderr_on_runtime_error(self): + result = SAFE_EVAL._execute( + "print('diagnostic', file=sys.stderr)\nraise RuntimeError('boom')", + "", + {"sys": sys}, + 32, + ) + self.assertFalse(result["ok"]) + self.assertEqual(result["stderr"], "diagnostic\n") + def test_io_tests_use_fresh_namespaces_and_hide_hidden_values(self): result = invoke( "value = int(input())\nprint(value * value)", diff --git a/internal/execution/wasm/adapter.go b/internal/execution/wasm/adapter.go index e9622c7..f85ff3f 100644 --- a/internal/execution/wasm/adapter.go +++ b/internal/execution/wasm/adapter.go @@ -36,6 +36,14 @@ import ( "go.uber.org/zap" ) +func validateWasm32RequestLength(length uint64) error { + const maxWasm32ByteLength = uint64(1<<32 - 1) + if length > maxWasm32ByteLength { + return fmt.Errorf("wasm: request length %d exceeds wasm32 address space", length) + } + return nil +} + // requestEnvelope is the JSON structure written into guest memory for each // evaluation call. type requestEnvelope struct { @@ -85,6 +93,9 @@ func (a *wasmAdapter) send( } reqLen := uint64(len(reqBytes)) + if err := validateWasm32RequestLength(reqLen); err != nil { + return nil, err + } // 2. Allocate guest memory for the request (cached lookup — M-4 fix). if a.allocFn == nil { diff --git a/internal/execution/wasm/adapter_test.go b/internal/execution/wasm/adapter_test.go new file mode 100644 index 0000000..5ced932 --- /dev/null +++ b/internal/execution/wasm/adapter_test.go @@ -0,0 +1,12 @@ +package wasm + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestValidateWasm32RequestLength(t *testing.T) { + require.NoError(t, validateWasm32RequestLength(1<<32-1)) + require.Error(t, validateWasm32RequestLength(1<<32)) +} diff --git a/internal/execution/wasm/dispatcher_test.go b/internal/execution/wasm/dispatcher_test.go index c0b5939..4953567 100644 --- a/internal/execution/wasm/dispatcher_test.go +++ b/internal/execution/wasm/dispatcher_test.go @@ -434,7 +434,7 @@ func TestSupervisor_Send_MemoryGrowDetected(t *testing.T) { origSize := mem.Size() require.Equal(t, origSize, sv.snapshotSize, "snapshotSize must be recorded at Take time") - prevPages, ok := mem.Grow(1) + prevPages, ok := mem.Grow(3) require.True(t, ok, "memory.Grow must succeed (echo fixture has no max)") require.Equal(t, origSize/(64*1024), prevPages) diff --git a/internal/execution/wasm/supervisor.go b/internal/execution/wasm/supervisor.go index 3162337..225a234 100644 --- a/internal/execution/wasm/supervisor.go +++ b/internal/execution/wasm/supervisor.go @@ -210,9 +210,17 @@ func (s *wasmSupervisor) restoreSnapshot() error { } if cur := mem.Size(); cur > s.snapshotSize { tail := cur - s.snapshotSize - zeros := make([]byte, tail) - if !mem.Write(s.snapshotSize, zeros) { - return fmt.Errorf("wasm: memory grew by %d bytes; zero-fill failed: %w", tail, ErrMemoryGrew) + var zeros [64 * 1024]byte + for offset := s.snapshotSize; offset < cur; { + remaining := cur - offset + chunkSize := uint32(len(zeros)) + if remaining < chunkSize { + chunkSize = remaining + } + if !mem.Write(offset, zeros[:chunkSize]) { + return fmt.Errorf("wasm: memory grew by %d bytes; zero-fill failed: %w", tail, ErrMemoryGrew) + } + offset += chunkSize } // The instance is discarded after this error, so restoring the captured // prefix has no value. Returning before strategy.Restore also avoids