pkg/aflow/tool/codeeditor: add actual implementation

pkg/aflow/flow.go

@@ -87,12 +87,7 @@ func registerOne[Inputs, Outputs any](all map[string]*Flow, flow *Flow) error {
 	if all[flow.Name] != nil {
 		return fmt.Errorf("flow %v is already registered", flow.Name)
 	}
-	ctx := &verifyContext{
-		inputs:  true,
-		outputs: true,
-		state:   make(map[string]*varState),
-		models:  make(map[string]bool),
-	}
+	ctx := newVerifyContext()
 	provideOutputs[Inputs](ctx, "flow inputs")
 	flow.Root.verify(ctx)
 	requireInputs[Outputs](ctx, "flow outputs")

pkg/aflow/func_tool.go

@@ -5,7 +5,10 @@ package aflow
 
 import (
 	"fmt"
+	"reflect"
+	"testing"
 
+	"github.com/stretchr/testify/require"
 	"google.golang.org/genai"
 )
 
@@ -77,3 +80,19 @@ func (t *funcTool[State, Args, Results]) verify(ctx *verifyContext) {
 	requireSchema[Results](ctx, t.Name, "Results")
 	requireInputs[State](ctx, t.Name)
 }
+
+func (*funcTool[State, Args, Results]) checkTestTypes(t *testing.T, ctx *verifyContext, state, args, results any) (
+	map[string]any, map[string]any, map[string]any) {
+	require.Equal(t, reflect.TypeFor[State](), reflect.TypeOf(state))
+	require.Equal(t, reflect.TypeFor[Args](), reflect.TypeOf(args))
+	require.Equal(t, reflect.TypeFor[Results](), reflect.TypeOf(results))
+	provideOutputs[State](ctx, "state")
+	return convertToMap(state.(State)), convertToMap(args.(Args)), convertToMap(results.(Results))
+}
+
+func (*funcTool[State, Args, Results]) checkFuzzTypes(t *testing.T, state, args any) (
+	map[string]any, map[string]any) {
+	require.Equal(t, reflect.TypeFor[State](), reflect.TypeOf(state))
+	require.Equal(t, reflect.TypeFor[Args](), reflect.TypeOf(args))
+	return convertToMap(state.(State)), convertToMap(args.(Args))
+}

pkg/aflow/test_tool.go

@@ -0,0 +1,49 @@
+// Copyright 2026 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package aflow
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestTool(t *testing.T, tool Tool, initState, initArgs, wantResults any, wantError string) {
+	type toolTester interface {
+		checkTestTypes(t *testing.T, ctx *verifyContext, state, args, results any) (
+			map[string]any, map[string]any, map[string]any)
+	}
+	vctx := newVerifyContext()
+	state, args, results := tool.(toolTester).checkTestTypes(t, vctx, initState, initArgs, wantResults)
+	tool.verify(vctx)
+	if err := vctx.finalize(); err != nil {
+		t.Fatal(err)
+	}
+	// Just ensure it does not crash.
+	_ = tool.declaration()
+	// We don't init all fields, init more, if necessary.
+	ctx := &Context{
+		state: state,
+	}
+	defer ctx.close()
+	gotResults, err := tool.execute(ctx, args)
+	gotError := ""
+	if err != nil {
+		gotError = err.Error()
+	}
+	require.Equal(t, wantError, gotError)
+	require.Equal(t, results, gotResults)
+}
+
+func FuzzTool(t *testing.T, tool Tool, initState, initArgs any) (map[string]any, error) {
+	type toolFuzzer interface {
+		checkFuzzTypes(t *testing.T, state, args any) (map[string]any, map[string]any)
+	}
+	state, args := tool.(toolFuzzer).checkFuzzTypes(t, initState, initArgs)
+	ctx := &Context{
+		state: state,
+	}
+	defer ctx.close()
+	return tool.execute(ctx, args)
+}

pkg/aflow/tool/codeeditor/codeeditor.go

@@ -4,11 +4,21 @@
 package codeeditor
 
 import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"slices"
+	"strings"
+
 	"github.com/google/syzkaller/pkg/aflow"
+	"github.com/google/syzkaller/pkg/codesearch"
+	"github.com/google/syzkaller/pkg/osutil"
 )
 
 var Tool = aflow.NewFuncTool("codeeditor", codeeditor, `
-The tool does one code edit to form the final patch.
+The tool does one source code edit to form the final patch by replacing full lines
+with new provided lines. If new code is empty, current lines will be deleted.
+Provide full lines of code including new line characters.
 The tool should be called mutiple times to do all required changes one-by-one,
 but avoid changing the same lines multiple times.
 Note: You will not see your edits via the codesearch tool.
@@ -20,18 +30,91 @@ type state struct {
 }
 
 type args struct {
-	SourceFile  string `jsonschema:"Full source file path."`
-	CurrentCode string `jsonschema:"The current code to replace verbatim with new lines, but without line numbers."`
-	NewCode     string `jsonschema:"New code to replace the current code snippet."`
+	SourceFile  string `jsonschema:"Full source file path to edit."`
+	CurrentCode string `jsonschema:"The current code lines to be replaced."`
+	NewCode     string `jsonschema:"New code lines to replace the current code lines."`
 }
 
 func codeeditor(ctx *aflow.Context, state state, args args) (struct{}, error) {
-	// TODO: check that the SourceFile is not escaping.
-	// If SourceFile is incorrect, or CurrentCode is not matched, return aflow.BadCallError
-	// with an explanation. Say that it needs to increase context if CurrentCode is not matched.
-	// Try to do as fuzzy match for CurrentCode as possible (strip line numbers,
-	// ignore white-spaces, etc).
-	// Should we accept a reference line number, or function name to disambiguate in the case
-	// of multiple matches?
-	return struct{}{}, nil
+	if strings.Contains(filepath.Clean(args.SourceFile), "..") {
+		return struct{}{}, aflow.BadCallError("SourceFile %q is outside of the source tree", args.SourceFile)
+	}
+	file := filepath.Join(state.KernelScratchSrc, args.SourceFile)
+	// Filter out not source files too (e.g. .git, etc),
+	// LLM have not seen them and should not be messing with them.
+	if !osutil.IsExist(file) || !codesearch.IsSourceFile(file) {
+		return struct{}{}, aflow.BadCallError("SourceFile %q does not exist", args.SourceFile)
+	}
+	if strings.TrimSpace(args.CurrentCode) == "" {
+		return struct{}{}, aflow.BadCallError("CurrentCode snippet is empty")
+	}
+	fileData, err := os.ReadFile(file)
+	if err != nil {
+		return struct{}{}, err
+	}
+	if len(fileData) == 0 || fileData[len(fileData)-1] != '\n' {
+		// Generally shouldn't happen, but just in case.
+		fileData = append(fileData, '\n')
+	}
+	if args.CurrentCode[len(args.CurrentCode)-1] != '\n' {
+		args.CurrentCode += "\n"
+	}
+	if args.NewCode != "" && args.NewCode[len(args.NewCode)-1] != '\n' {
+		args.NewCode += "\n"
+	}
+	lines := slices.Collect(bytes.Lines(fileData))
+	src := slices.Collect(bytes.Lines([]byte(args.CurrentCode)))
+	dst := slices.Collect(bytes.Lines([]byte(args.NewCode)))
+	// First, try to match as is. If that fails, try a more permissive matching
+	// that ignores whitespaces, empty lines, etc.
+	newLines, matches := replace(lines, src, dst, false)
+	if matches == 0 {
+		newLines, matches = replace(lines, src, dst, true)
+	}
+	if matches == 0 {
+		return struct{}{}, aflow.BadCallError("CurrentCode snippet does not match anything in the source file," +
+			" provide more precise CurrentCode snippet")
+	}
+	if matches > 1 {
+		return struct{}{}, aflow.BadCallError("CurrentCode snippet matched %v places,"+
+			" increase context in CurrentCode to avoid ambiguity", matches)
+	}
+	err = osutil.WriteFile(file, slices.Concat(newLines...))
+	return struct{}{}, err
+}
+
+func replace(lines, src, dst [][]byte, fuzzy bool) (newLines [][]byte, matches int) {
+	for i := 0; i < len(lines); i++ {
+		li, si := i, 0
+		for li < len(lines) && si < len(src) {
+			l, s := lines[li], src[si]
+			if fuzzy {
+				// Ignore whitespaces and empty lines.
+				l, s = bytes.TrimSpace(l), bytes.TrimSpace(s)
+				// Potentially we can remove line numbers from s here if they are present,
+				// or use them to disambiguate in the case of multiple matches.
+				if len(s) == 0 {
+					si++
+					continue
+				}
+				if len(l) == 0 {
+					li++
+					continue
+				}
+			}
+			if !bytes.Equal(l, s) {
+				break
+			}
+			li++
+			si++
+		}
+		if si != len(src) {
+			newLines = append(newLines, lines[i])
+			continue
+		}
+		matches++
+		newLines = append(newLines, dst...)
+		i = li - 1
+	}
+	return
 }