runner.go

// Package spine is a small, fresh-instance Wasm CPython runner.
package spine

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
)

// Runner owns compiled code, tools and an optional clean image, never a live user Guest.
// Close must happen after all Run calls return. Tool implementations must honor context.
type Runner struct {
	runtime    wazero.Runtime
	code       wazero.CompiledModule
	tools      map[string]Tool
	earlyReads map[string]bool
	image      []byte // Immutable initialized linear memory; nil means fresh CPython init.
}

type Output struct {
	Value       json.RawMessage
	Stdout      string
	Transformed string // Optional actual Guest AST rendering for learning, not execution input.
}

// New compiles the actual artifact once; every Run instantiates private memory.
func New(ctx context.Context, wasm []byte, tools map[string]Tool, earlyReads ...string) (*Runner, error) {
	r := &Runner{
		runtime:    wazero.NewRuntimeWithConfig(ctx, wazero.NewRuntimeConfig().WithCloseOnContextDone(true)),
		tools:      make(map[string]Tool, len(tools)),
		earlyReads: make(map[string]bool),
	}
	for name, fn := range tools {
		r.tools[name] = fn
	}
	for _, name := range earlyReads {
		if r.tools[name] == nil {
			r.Close(ctx)
			return nil, fmt.Errorf("unknown early-read tool: %s", name)
		}
		r.earlyReads[name] = true
	}
	if _, err := wasi_snapshot_preview1.Instantiate(ctx, r.runtime); err != nil {
		r.Close(ctx)
		return nil, err
	}
	if _, err := r.runtime.NewHostModuleBuilder("spine").NewFunctionBuilder().WithFunc(r.hostCall).Export("call").
		NewFunctionBuilder().WithFunc(r.hostPrepare).Export("prepare").
		NewFunctionBuilder().WithFunc(r.hostResolve).Export("resolve").Instantiate(ctx); err != nil {
		r.Close(ctx)
		return nil, err
	}
	var err error
	r.code, err = r.runtime.CompileModule(ctx, wasm)
	if err != nil {
		r.Close(ctx)
		return nil, err
	}
	for _, name := range []string{"_initialize", "init", "alloc", "release", "execute"} {
		if r.code.ExportedFunctions()[name] == nil {
			r.Close(ctx)
			return nil, fmt.Errorf("missing Guest export: %s", name)
		}
	}
	return r, nil
}

func (r *Runner) Close(ctx context.Context) error {
	r.image = nil
	return r.runtime.Close(ctx)
}

// Run owns the instance, I/O and all Guest allocations until it returns.
func (r *Runner) Run(ctx context.Context, source string, inputs any) (Output, error) {
	return r.run(ctx, source, inputs, false, nil)
}

// RunPLM enables the built-in Guest AST pass, only for explicitly allowed snapshot reads.
func (r *Runner) RunPLM(ctx context.Context, source string, inputs any) (Output, error) {
	return r.run(ctx, source, inputs, true, nil)
}

func (r *Runner) run(ctx context.Context, source string, inputs any, plm bool, chunks <-chan string) (Output, error) {
	state := newRun(ctx, r, plm)
	defer state.close()
	ctx = context.WithValue(state.ctx, runKey{}, state)
	request, err := json.Marshal(struct {
		Source string `json:"source"`
		Inputs any    `json:"inputs"`
		PLM    bool   `json:"plm"`
	}{source, inputs, plm})
	if err != nil {
		return Output{}, err
	}
	if len(request) > maxMessage {
		return Output{}, errors.New("request exceeds 1 MiB")
	}
	stdout, stderr := &boundedText{}, &boundedText{}
	// No host directories, environment, stdin, network or process capabilities.
	m, err := r.newGuest(ctx, stdout, stderr)
	if err != nil {
		return Output{}, err
	}
	defer m.Close(context.Background())
	fail := func(err error) (Output, error) {
		return Output{Stdout: stdout.String()}, fmt.Errorf("%w%s", err, stderr.String())
	}
	if chunks != nil {
		if err := receiveSource(ctx, m, request, chunks); err != nil {
			return fail(err)
		}
		// Input and source already belong to this Guest. Do not send a second copy.
		request = []byte(`{"prefix":true}`)
	}
	packed, err := callWithBytes(ctx, m, "execute", request)
	if err != nil {
		return fail(err)
	}
	// execute returns (length << 32) | pointer. Go copies before release/Close.
	p, n := uint32(packed[0]), uint32(packed[0]>>32)
	if p == 0 {
		return fail(errors.New("Guest execution bridge failed: "))
	}
	defer m.ExportedFunction("release").Call(ctx, uint64(p))
	if n > maxMessage {
		return fail(errors.New("result exceeds 1 MiB"))
	}
	data, ok := m.Memory().Read(p, n)
	if !ok {
		return fail(errors.New("Guest result outside linear memory"))
	}
	var response struct {
		Value       json.RawMessage `json:"value"`
		Error       string          `json:"error"`
		Transformed string          `json:"transformed"`
	}
	if err := json.Unmarshal(data, &response); err != nil {
		return fail(err)
	}
	out := Output{Value: response.Value, Stdout: stdout.String(), Transformed: response.Transformed}
	if response.Error != "" {
		return out, errors.New(response.Error)
	}
	return out, nil
}

// The 1 MiB message/output limit is a demo policy, not a Wasm limit.
const maxMessage = 1 << 20

type boundedText struct{ strings.Builder }

func (b *boundedText) Write(p []byte) (int, error) {
	if b.Len()+len(p) > maxMessage {
		return 0, errors.New("stdout/stderr exceeds 1 MiB")
	}
	return b.Builder.Write(p)
}

// writeResponse is the only Host → Guest memory write for tool results.
func writeResponse(m api.Module, ptr, capacity uint32, data []byte) uint32 {
	if len(data) > maxMessage || uint64(len(data)) > uint64(capacity) || !m.Memory().Write(ptr, data) {
		return ^uint32(0)
	}
	return uint32(len(data))
}