pkg/aflow: add explicit DoWhile.MaxIterations

Add DoWhile.MaxIterations and make it mandatory.
I think it's useful to make workflow implementer to think
explicitly about a reasonable cap on the number of iterations.

4 files changed, 176 insertions, 6 deletions

diff --git a/pkg/aflow/flow/patching/patching.go b/pkg/aflow/flow/patching/patching.go
index ea5780022..a5aba0ccc 100644
--- a/pkg/aflow/flow/patching/patching.go
+++ b/pkg/aflow/flow/patching/patching.go
@@ -78,7 +78,8 @@ func init() {
 						},
 						crash.TestPatch, // -> PatchDiff or TestError
 					),
-					While: "TestError",
+					While:         "TestError",
+					MaxIterations: 10,
 				},
 				&aflow.LLMAgent{
 					Name:        "description-generator",
diff --git a/pkg/aflow/loop.go b/pkg/aflow/loop.go
index e074d4217..18a7cafa8 100644
--- a/pkg/aflow/loop.go
+++ b/pkg/aflow/loop.go
@@ -18,6 +18,9 @@ type DoWhile struct {
 	// Exit condition. It should be a string state variable.
 	// The loop exists when the variable is empty.
 	While string
+	// Max interations for the loop.
+	// Must be specified to avoid unintended effectively infinite loops.
+	MaxIterations int
 
 	loopVars map[string]reflect.Type
 }
@@ -43,8 +46,7 @@ func (dw *DoWhile) loop(ctx *Context) error {
 		}
 		ctx.state[name] = reflect.Zero(typ).Interface()
 	}
-	const maxIters = 100
-	for iter := 0; iter < maxIters; iter++ {
+	for iter := 0; iter < dw.MaxIterations; iter++ {
 		span := &trajectory.Span{
 			Type: trajectory.SpanLoopIteration,
 			Name: fmt.Sprint(iter),
@@ -60,10 +62,14 @@ func (dw *DoWhile) loop(ctx *Context) error {
 			return nil
 		}
 	}
-	return fmt.Errorf("DoWhile loop is going in cycles for %v iterations", maxIters)
+	return fmt.Errorf("DoWhile loop is going in cycles for %v iterations", dw.MaxIterations)
 }
 
 func (dw *DoWhile) verify(ctx *verifyContext) {
+	if max := 1000; dw.MaxIterations <= 0 || dw.MaxIterations >= max {
+		ctx.errorf("DoWhile", "bad MaxIterations value %v, should be within [1, %v]",
+			dw.MaxIterations, max)
+	}
 	// Verification of loops is a bit tricky.
 	// Normally we require each variable to be defined before use, but loops violate
 	// the assumption. An action in a loop body may want to use a variable produced
diff --git a/pkg/aflow/loop_test.go b/pkg/aflow/loop_test.go
index e68191c2b..782a470b1 100644
--- a/pkg/aflow/loop_test.go
+++ b/pkg/aflow/loop_test.go
@@ -47,7 +47,8 @@ func TestDoWhile(t *testing.T) {
 					return testResults{Diff: "diff"}, nil
 				}),
 			),
-			While: "TestError",
+			While:         "TestError",
+			MaxIterations: 10,
 		},
 		nil,
 	)
@@ -63,7 +64,8 @@ func TestDoWhileErrors(t *testing.T) {
 				}) (struct{}, error) {
 					return struct{}{}, nil
 				}),
-				While: "Condition",
+				While:         "Condition",
+				MaxIterations: 10,
 			},
 		))
 
@@ -76,6 +78,7 @@ func TestDoWhileErrors(t *testing.T) {
 				}) (struct{}, error) {
 					return struct{}{}, nil
 				}),
+				MaxIterations: 10,
 			},
 		))
 
@@ -90,7 +93,34 @@ func TestDoWhileErrors(t *testing.T) {
 				Do: NewFuncAction("body", func(ctx *Context, args struct{}) (output, error) {
 					return output{}, nil
 				}),
+				While:         "Output1",
+				MaxIterations: 10,
+			},
+		))
+	testRegistrationError[struct{}, struct{}](t,
+		"flow test: action DoWhile: bad MaxIterations value 0, should be within [1, 1000]",
+		Pipeline(
+			&DoWhile{
+				Do: NewFuncAction("body", func(ctx *Context, args struct{}) (output, error) {
+					return output{}, nil
+				}),
 				While: "Output1",
 			},
 		))
 }
+
+func TestDoWhileMaxIters(t *testing.T) {
+	type actionResults struct {
+		Error string
+	}
+	testFlow[struct{}, struct{}](t, nil, "DoWhile loop is going in cycles for 3 iterations",
+		&DoWhile{
+			Do: NewFuncAction("nop", func(ctx *Context, args struct{}) (actionResults, error) {
+				return actionResults{"failed"}, nil
+			}),
+			While:         "Error",
+			MaxIterations: 3,
+		},
+		nil,
+	)
+}
diff --git a/pkg/aflow/testdata/TestDoWhileMaxIters.trajectory.json b/pkg/aflow/testdata/TestDoWhileMaxIters.trajectory.json
new file mode 100644
index 000000000..35a2c0f7a
--- /dev/null
+++ b/pkg/aflow/testdata/TestDoWhileMaxIters.trajectory.json
@@ -0,0 +1,133 @@
+[
+	{
+		"Seq": 0,
+		"Nesting": 0,
+		"Type": "flow",
+		"Name": "test",
+		"Started": "0001-01-01T00:00:01Z"
+	},
+	{
+		"Seq": 1,
+		"Nesting": 1,
+		"Type": "loop",
+		"Name": "",
+		"Started": "0001-01-01T00:00:02Z"
+	},
+	{
+		"Seq": 2,
+		"Nesting": 2,
+		"Type": "iteration",
+		"Name": "0",
+		"Started": "0001-01-01T00:00:03Z"
+	},
+	{
+		"Seq": 3,
+		"Nesting": 3,
+		"Type": "action",
+		"Name": "nop",
+		"Started": "0001-01-01T00:00:04Z"
+	},
+	{
+		"Seq": 3,
+		"Nesting": 3,
+		"Type": "action",
+		"Name": "nop",
+		"Started": "0001-01-01T00:00:04Z",
+		"Finished": "0001-01-01T00:00:05Z",
+		"Results": {
+			"Error": "failed"
+		}
+	},
+	{
+		"Seq": 2,
+		"Nesting": 2,
+		"Type": "iteration",
+		"Name": "0",
+		"Started": "0001-01-01T00:00:03Z",
+		"Finished": "0001-01-01T00:00:06Z"
+	},
+	{
+		"Seq": 4,
+		"Nesting": 2,
+		"Type": "iteration",
+		"Name": "1",
+		"Started": "0001-01-01T00:00:07Z"
+	},
+	{
+		"Seq": 5,
+		"Nesting": 3,
+		"Type": "action",
+		"Name": "nop",
+		"Started": "0001-01-01T00:00:08Z"
+	},
+	{
+		"Seq": 5,
+		"Nesting": 3,
+		"Type": "action",
+		"Name": "nop",
+		"Started": "0001-01-01T00:00:08Z",
+		"Finished": "0001-01-01T00:00:09Z",
+		"Results": {
+			"Error": "failed"
+		}
+	},
+	{
+		"Seq": 4,
+		"Nesting": 2,
+		"Type": "iteration",
+		"Name": "1",
+		"Started": "0001-01-01T00:00:07Z",
+		"Finished": "0001-01-01T00:00:10Z"
+	},
+	{
+		"Seq": 6,
+		"Nesting": 2,
+		"Type": "iteration",
+		"Name": "2",
+		"Started": "0001-01-01T00:00:11Z"
+	},
+	{
+		"Seq": 7,
+		"Nesting": 3,
+		"Type": "action",
+		"Name": "nop",
+		"Started": "0001-01-01T00:00:12Z"
+	},
+	{
+		"Seq": 7,
+		"Nesting": 3,
+		"Type": "action",
+		"Name": "nop",
+		"Started": "0001-01-01T00:00:12Z",
+		"Finished": "0001-01-01T00:00:13Z",
+		"Results": {
+			"Error": "failed"
+		}
+	},
+	{
+		"Seq": 6,
+		"Nesting": 2,
+		"Type": "iteration",
+		"Name": "2",
+		"Started": "0001-01-01T00:00:11Z",
+		"Finished": "0001-01-01T00:00:14Z"
+	},
+	{
+		"Seq": 1,
+		"Nesting": 1,
+		"Type": "loop",
+		"Name": "",
+		"Started": "0001-01-01T00:00:02Z",
+		"Finished": "0001-01-01T00:00:15Z",
+		"Error": "DoWhile loop is going in cycles for 3 iterations"
+	},
+	{
+		"Seq": 0,
+		"Nesting": 0,
+		"Type": "flow",
+		"Name": "test",
+		"Started": "0001-01-01T00:00:01Z",
+		"Finished": "0001-01-01T00:00:16Z",
+		"Error": "DoWhile loop is going in cycles for 3 iterations"
+	}
+]
\ No newline at end of file