Prefix：从源码片段到原位领取

// RunPrefix receives append-only UTF-8 source fragments. Closing chunks commits
// the source for one execution in the same Guest; ctx cancels an unfinished stream.
func (r *Runner) RunPrefix(ctx context.Context, chunks <-chan string, inputs any) (Output, error) {
	if chunks == nil {
		return Output{}, errors.New("nil source stream")
	}
	for _, name := range []string{"prefix_begin", "prefix_feed"} {
		if r.code.ExportedFunctions()[name] == nil {
			return Output{}, fmt.Errorf("missing prefix Guest export: %s", name)
		}
	}
	return r.run(ctx, "", inputs, true, chunks)
}

func (r *Runner) Close(ctx context.Context) error {
	r.image = nil
	err := r.runtime.Close(ctx)
	if r.cow != nil {
		err = errors.Join(err, r.cow.close())
		r.cow = nil
	}
	return err
}

// 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) {

func receiveSource(ctx context.Context, m api.Module, initial []byte, chunks <-chan string) error {
	status, err := callWithBytes(ctx, m, "prefix_begin", initial)
	if err != nil {
		return err
	}
	if status[0] != 0 {
		return errors.New("prefix initialization failed: ")
	}
	received := len(initial)
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case chunk, ok := <-chunks:
			if !ok {
				return ctx.Err()
			}
			received += len(chunk)
			if received > maxMessage {
				return errors.New("prefix input exceeds 1 MiB")
			}
			if chunk == "" {
				continue
			}
			status, err = callWithBytes(ctx, m, "prefix_feed", []byte(chunk))
			if err != nil {
				return err
			}
			if status[0] != 0 {
				return errors.New("prefix feed failed: ")
			}
		}
	}

        return 0


def resolve_call(handle, name, **args):
    return decode(resolve(handle, json.dumps({"tool": name, "args": args})))


def prefix_begin(request):

class Prefix:
    def __init__(self, inputs, manifest, prepare):
        self.inputs = inputs
        self.manifest = manifest
        self.early_names = {spec["name"] for spec in manifest if spec["allow_early_read"]}
        self.prepare = prepare
        self.source = ""
        self.seen = 0

        self.claimed = 0

    def feed(self, chunk):
        self.source += chunk
        # A trailing partial line is not a complete received statement.
        complete = self.source[:self.source.rfind("\n") + 1]
        try:
            tree = ast.parse(complete, filename="<spine>")

                self.ready.append((request, handle))

    def candidate(self, statement):
        if not (isinstance(statement, ast.Assign) and len(statement.targets) == 1
                and isinstance(statement.targets[0], ast.Name)
                and statement.targets[0].id != "inputs"
                and statement.targets[0].id not in self.early_names):
            return False
        call = statement.value
        return (isinstance(call, ast.Call) and isinstance(call.func, ast.Name)
                and call.func.id in self.early_names and not call.args
                and all(k.arg is not None and self.literal_or_input(k.value) for k in call.keywords))

    @staticmethod
    def literal_or_input(node):
        if isinstance(node, ast.Constant):
            return isinstance(node.value, (str, int, float, bool, type(None)))
        return (isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name)
                and node.value.id == "inputs" and isinstance(node.slice, ast.Constant)
                and isinstance(node.slice.value, (str, int)))

    def argument(self, node):
        if isinstance(node, ast.Constant):

        except SyntaxError:
            return  # Incomplete/invalid source is diagnosed at final compile.
        for index, statement in enumerate(tree.body):
            if not self.candidate(statement):
                break  # Never pass a branch, dependency, mutation, import, or other code.
            if index < self.seen:
                continue
            self.seen = index + 1
            call = statement.value
            try:
                args = {k.arg: self.argument(k.value) for k in call.keywords}
                request = json.dumps({"tool": call.func.id, "args": args})
            except Exception:
                continue  # Argument errors are still raised at the actual Python call.
            handle = self.prepare(request)

func (s *runState) prepare(request callRequest) uint32 {
	// 64 pending reads is a demo resource policy, not a Wasm limit.
	spec, ok := s.runner.manifest[request.Tool]
	if !s.enabled || !ok || !spec.AllowEarlyRead || len(s.futures) >= 64 {
		return 0 // Explicitly not prepared; original call will execute normally.
	}
	ctx, cancel := context.WithCancel(s.ctx)
	f := &future{request: request, done: make(chan struct{}), cancel: cancel}
	s.next++
	id := s.next
	s.futures[id] = f
	s.workers.Add(1)
	go func() {
		defer s.workers.Done()
		defer cancel()
		defer close(f.done)
		value, err := spec.Call(ctx, request.Args)
		f.response = encodeResponse(value, err)
	}()

	request, err := json.Marshal(struct {
		Source   string          `json:"source"`
		Inputs   any             `json:"inputs"`
		PLM      bool            `json:"plm"`
		Manifest []guestToolSpec `json:"manifest"`
	}{source, inputs, plm, r.guestManifest})
	if err != nil {
		return Output{}, err

    request = json.loads(request)
    _prefix = Prefix(request["inputs"], request["manifest"], prepare)


def prefix_feed(chunk):
    _prefix.feed(chunk)


def execute(request):
    transformed = ""
    try:
        request = json.loads(request)
        if request.get("prefix"):
            request = {
                "source": _prefix.source,
                "inputs": _prefix.inputs,
                "manifest": _prefix.manifest,
                "plm": True,

def invoke_tool(name, args):
    request = json.dumps({"tool": name, "args": args})
    handle = _prefix.claim(request) if _prefix is not None else None
    return decode(call(request) if handle is None else resolve(handle, request))


def make_tool(name):
    def proxy(**args):
        return invoke_tool(name, args)
    proxy.__name__ = name
    return proxy


def prepare_call(thunk):
    try:
        name, args = thunk()
        request = json.dumps({"tool": name, "args": args})
        handle = _prefix.claim(request) if _prefix is not None else None
        return prepare(request) if handle is None else handle

        return self.inputs[node.slice.value]

    def claim(self, request):
        # Leading calls execute in source order even if the full program is not optimized.
        # Match the actual request too; never search ahead or reuse a different read.
        if self.claimed < len(self.ready) and self.ready[self.claimed][0] == request:
            handle = self.ready[self.claimed][1]
            self.claimed += 1

func (s *runState) resolve(id uint32, request callRequest) []byte {
	if id == 0 {
		return s.runner.invoke(s.ctx, request)
	}
	f := s.futures[id]
	if f == nil {
		return encodeResponse(nil, errors.New("unknown or consumed future"))
	}
	delete(s.futures, id) // Single-use handle, not a static source-site slot.
	defer f.cancel()
	if f.request.Tool != request.Tool || !bytes.Equal(f.request.Args, request.Args) {
		f.cancel()
		// Only registered snapshot reads were started early. Never reuse mismatched data.
		return s.runner.invoke(s.ctx, request)
	}
	select {
	case <-f.done:
		return f.response // Includes tool errors: do not retry a failed operation.
	case <-s.ctx.Done():
		return encodeResponse(nil, s.ctx.Err())
	}

func newRun(ctx context.Context, r *Runner, enabled bool) *runState {
	ctx, cancel := context.WithCancel(ctx)
	return &runState{runner: r, ctx: ctx, cancel: cancel, enabled: enabled, futures: make(map[uint32]*future)}
}

func (s *runState) close() {
	s.cancel()
	s.workers.Wait() // Cooperative tools finish before Run returns, including discarded reads.
}