Switch to direct observation of costs on a per-operation basis

cel/program.go

@@ -341,7 +341,7 @@ func newProgram(e *Env, a *ast.AST, opts []ProgramOption) (Program, error) {
 			observers = append(observers, interpreter.EvalStateObserver())
 		}
 		if p.evalOpts&OptTrackCost == OptTrackCost {
-			observers = append(observers, interpreter.CostObserver(interpreter.CostTrackerFactory(trackerFactory)))
+			plannerOptions = append(plannerOptions, interpreter.CostObserver(interpreter.CostTrackerFactory(trackerFactory)))
 		}
 		// Enable exhaustive eval over a basic observer since it offers a superset of features.
 		if p.evalOpts&OptExhaustiveEval == OptExhaustiveEval {
@@ -369,19 +369,6 @@ func (p *prog) initInterpretable(a *ast.AST, plannerOptions []interpreter.Planne
 
 // Eval implements the Program interface method.
 func (p *prog) Eval(input any) (out ref.Val, det *EvalDetails, err error) {
-	// Configure error recovery for unexpected panics during evaluation. Note, the use of named
-	// return values makes it possible to modify the error response during the recovery
-	// function.
-	defer func() {
-		if r := recover(); r != nil {
-			switch t := r.(type) {
-			case interpreter.EvalCancelledError:
-				err = t
-			default:
-				err = fmt.Errorf("internal error: %v", r)
-			}
-		}
-	}()
 	// Asynchronous calls cannot be resolved by a single-pass evaluation. Reject before doing any
 	// work (this also covers ContextEval, which delegates here); ConcurrentEval does not call Eval.
 	if p.hasAsync {
@@ -398,19 +385,36 @@ func (p *prog) Eval(input any) (out ref.Val, det *EvalDetails, err error) {
 		}
 		defer frame.Close()
 	}
+	// Configure error recovery and details capture for evaluation.
+	defer func() {
+		if tracker := frame.CostTracker(); tracker != nil {
+			if det == nil {
+				det = &EvalDetails{}
+			}
+			det.costTracker = tracker
+		}
+		if r := recover(); r != nil {
+			switch t := r.(type) {
+			case interpreter.EvalCancelledError:
+				err = t
+			default:
+				err = fmt.Errorf("internal error: %v", r)
+			}
+		}
+	}()
+
 	if p.observable != nil {
 		det = &EvalDetails{}
 		out = p.observable.ObserveExec(frame, func(observed any) {
 			switch o := observed.(type) {
 			case interpreter.EvalState:
 				det.state = o
-			case *interpreter.CostTracker:
-				det.costTracker = o
 			}
 		})
 	} else {
 		out = p.interpretable.Exec(frame)
 	}
+
 	// The output of an internal Eval may have a value (`v`) that is a types.Err. This step
 	// translates the CEL value to a Go error response. This interface does not quite match the
 	// RPC signature which allows for multiple errors to be returned, but should be sufficient.
@@ -511,47 +515,60 @@ func (p *prog) ConcurrentEval(ctx context.Context, input any) <-chan EvalResult
 
 	go func() {
 		defer close(resCh)
+		frame, err := p.newAsyncFrame(ctx, input)
+		if err != nil {
+			resCh <- EvalResult{Err: err}
+			return
+		}
+		defer frame.Close()
+
+		var det *EvalDetails
 		// Ensure concurrent eval handles panic / recovery properly
 		defer func() {
+			if tracker := frame.CostTracker(); tracker != nil {
+				if det == nil {
+					det = &EvalDetails{}
+				}
+				det.costTracker = tracker
+			}
 			if r := recover(); r != nil {
 				switch t := r.(type) {
 				case interpreter.EvalCancelledError:
-					resCh <- EvalResult{Err: t}
+					resCh <- EvalResult{EvalDetails: det, Err: t}
 				default:
-					resCh <- EvalResult{Err: fmt.Errorf("internal error: %v", r)}
+					resCh <- EvalResult{EvalDetails: det, Err: fmt.Errorf("internal error: %v", r)}
 				}
 			}
 		}()
 
-		frame, err := p.newAsyncFrame(ctx, input)
-		if err != nil {
-			resCh <- EvalResult{Err: err}
-			return
-		}
-		defer frame.Close()
-
 		// Completions are signaled to this channel as async calls finish. The asyncCallState
 		// fan-in also selects on ctx.Done(), so the sender will not leak if this loop returns early.
 		completions := make(chan int64, p.resolveCompletionBufferSize())
 		frame.SetCompletions(completions)
 
 		for {
 			var out ref.Val
-			var det *EvalDetails
 
 			if p.observable != nil {
 				det = &EvalDetails{}
 				out = p.observable.ObserveExec(frame, func(observed any) {
 					switch o := observed.(type) {
 					case interpreter.EvalState:
 						det.state = o
-					case *interpreter.CostTracker:
-						det.costTracker = o
 					}
 				})
 			} else {
 				out = p.interpretable.Exec(frame)
 			}
+			// This ensures that cost tracking is present on the result passed through the channel
+			// in the positive outcome case, as opposed to the defer which captures these details
+			// in the evaluation error scenarios.
+			if tracker := frame.CostTracker(); tracker != nil {
+				if det == nil {
+					det = &EvalDetails{}
+				}
+				det.costTracker = tracker
+			}
 
 			// Communicate errors quickly.
 			if types.IsError(out) {

ext/encoders_test.go

@@ -249,7 +249,7 @@ func TestEncodersCosts(t *testing.T) {
 				"x": 100,
 			},
 			estimatedCost: checker.CostEstimate{Min: 2, Max: math.MaxUint64},
-			actualCost:    1,
+			actualCost:    math.MaxUint64,
 			version:       1,
 		},
 	}

interpreter/async.go

@@ -111,7 +111,11 @@ func (fn *evalAsyncFunc) Exec(frame *ExecutionFrame) ref.Val {
 		return unk
 	}
 	result := frame.ComputeResult(fn.ID(), fn.Function(), fn.OverloadID(), fn.impl, argVals)
-	return types.LabelErrNode(fn.id, result)
+	val := types.LabelErrNode(fn.id, result)
+	if costs := frame.CostTracker(); costs != nil {
+		costs.EvalVarArgs(frame, fn.id, fn, argVals, val)
+	}
+	return val
 }
 
 // asyncCallStateTracker manages async call states across re-evaluations of a single program.

interpreter/attributes.go

@@ -1318,6 +1318,11 @@ func applyQualifiers(vars Activation, obj any, qualifiers []Qualifier) (any, boo
 			if err != nil {
 				return nil, false, err
 			}
+			if frame := AsFrame(vars); frame != nil {
+				if costs := frame.CostTracker(); costs != nil {
+					costs.Qualify(qual.ID())
+				}
+			}
 			if !present {
 				// We return optional none here with a presence of 'false' as the layers
 				// above will attempt to call types.OptionalOf() on a present value if any
@@ -1329,6 +1334,11 @@ func applyQualifiers(vars Activation, obj any, qualifiers []Qualifier) (any, boo
 			if err != nil {
 				return nil, false, err
 			}
+			if frame := AsFrame(vars); frame != nil {
+				if costs := frame.CostTracker(); costs != nil {
+					costs.Qualify(qual.ID())
+				}
+			}
 		}
 		obj = qualObj
 	}

interpreter/frame.go

@@ -231,6 +231,22 @@ func (f *ExecutionFrame) CheckInterrupt() bool {
 	return false
 }
 
+// CostTracker returns the active CostTracker for this evaluation pass, or nil.
+func (f *ExecutionFrame) CostTracker() *CostTracker {
+	if f.ctx == nil {
+		return nil
+	}
+	return f.ctx.costs
+}
+
+// SetCostTracker sets the active CostTracker for this evaluation pass.
+func (f *ExecutionFrame) SetCostTracker(tracker *CostTracker) {
+	if f.ctx == nil {
+		f.ctx = evalContextPool.Get().(*evalContext)
+	}
+	f.ctx.costs = tracker
+}
+
 // ComputeResult tracks and computes the result of the given asynchronous function.
 //
 // The first invocation for a given (node id, args) tuple registers the call state and returns an