Initial stateful observers prior to evaluation (#1163)

TristonianJones · web-flow · commit 3da6139fe4e1 · 2025-04-14T11:01:44.000-07:00
diff --git a/cel/cel_test.go b/cel/cel_test.go
@@ -17,6 +17,7 @@ package cel
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"reflect"
@@ -1766,6 +1767,95 @@ func TestEstimateCostAndRuntimeCost(t *testing.T) {
 	}
 }
 
+func TestCostLimit(t *testing.T) {
+	cases := []struct {
+		name      string
+		expr      string
+		decls     []EnvOption
+		costLimit uint64
+		in        any
+		err       error
+	}{
+		{
+			name: "greater",
+			expr: `val1 > val2`,
+			decls: []EnvOption{
+				Variable("val1", IntType),
+				Variable("val2", IntType),
+			},
+			in:        map[string]any{"val1": 1, "val2": 2},
+			costLimit: 10,
+		},
+		{
+			name: "greater - error",
+			expr: `val1 > val2`,
+			decls: []EnvOption{
+				Variable("val1", IntType),
+				Variable("val2", IntType),
+			},
+			in:        map[string]any{"val1": 1, "val2": 2},
+			costLimit: 0,
+			err:       errors.New("actual cost limit exceeded"),
+		},
+	}
+
+	for _, tst := range cases {
+		tc := tst
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			envOpts := []EnvOption{
+				CostEstimatorOptions(
+					checker.OverloadCostEstimate(overloads.TimestampToYear, estimateTimestampToYear),
+				),
+			}
+			envOpts = append(envOpts, tc.decls...)
+			env := testEnv(t, envOpts...)
+			ast, iss := env.Compile(tc.expr)
+			if iss.Err() != nil {
+				t.Fatalf("env.Compile(%v) failed: %v", tc.expr, iss.Err())
+			}
+			est, err := env.EstimateCost(ast, testCostEstimator{hints: map[string]uint64{}})
+			if err != nil {
+				t.Fatalf("Env.EstimateCost(ast *Ast, estimator checker.CostEstimator) failed to estimate cost: %s\n", err)
+			}
+
+			checkedAst, iss := env.Check(ast)
+			if iss.Err() != nil {
+				t.Fatalf(`Env.Check(ast *Ast) failed to check expression: %v`, iss.Err())
+			}
+			// Evaluate expression.
+			program, err := env.Program(checkedAst,
+				CostTracking(testRuntimeCostEstimator{}),
+				CostTrackerOptions(
+					interpreter.OverloadCostTracker(overloads.TimestampToYear, trackTimestampToYear),
+				),
+				CostLimit(tc.costLimit),
+			)
+			if err != nil {
+				t.Fatalf(`Env.Program(ast *Ast, opts ...ProgramOption) failed to construct program: %v`, err)
+			}
+			_, details, err := program.Eval(tc.in)
+			if err != nil && tc.err == nil {
+				t.Fatalf(`Program.Eval(vars any) failed to evaluate expression: %v`, err)
+			}
+			actualCost := details.ActualCost()
+			if actualCost == nil {
+				t.Errorf(`EvalDetails.ActualCost() got nil for "%s" cost, wanted %d`, tc.expr, actualCost)
+			}
+			if err == nil {
+				if est.Min > *actualCost || est.Max < *actualCost {
+					t.Errorf("EvalDetails.ActualCost() failed to return a runtime cost %d is the range of estimate cost [%d, %d]", *actualCost,
+						est.Min, est.Max)
+				}
+			} else {
+				if !strings.Contains(err.Error(), tc.err.Error()) {
+					t.Fatalf("program.Eval() got error %v, wanted error containing %v", err, tc.err)
+				}
+			}
+		})
+	}
+}
+
 func TestPartialVars(t *testing.T) {
 	env := testEnv(t,
 		Variable("x", StringType),
diff --git a/interpreter/interpretable.go b/interpreter/interpretable.go
@@ -138,6 +138,9 @@ func (oi *ObservableInterpretable) ObserveEval(vars Activation, observer func(an
 		if err != nil {
 			return types.WrapErr(err)
 		}
+		// Provide an initial reference to the state to ensure state is available
+		// even in cases of interrupting errors generated during evaluation.
+		observer(obs.GetState(vars))
 	}
 	result := oi.Interpretable.Eval(vars)
 	// Get the state which needs to be reported back as having been observed.

Original file line number	Diff line number	Diff line change
`@@ -138,6 +138,9 @@ func (oi *ObservableInterpretable) ObserveEval(vars Activation, observer func(an`
`138`	`138`	`if err != nil {`
`139`	`139`	`return types.WrapErr(err)`
`140`	`140`	`}`
	`141`	`+ // Provide an initial reference to the state to ensure state is available`
	`142`	`+ // even in cases of interrupting errors generated during evaluation.`
	`143`	`+ observer(obs.GetState(vars))`
`141`	`144`	`}`
`142`	`145`	`result := oi.Interpretable.Eval(vars)`
`143`	`146`	`// Get the state which needs to be reported back as having been observed.`