fix(patcher): inject ctx into nested custom funcs

WithContext patcher now looks up function types from the Functions
table when the callee type is interface{}. This fixes context
injection for custom functions nested as arguments inside method
calls with unknown callee types (e.g., now2().After(date2())).

Also improved the regression test to actually verify context is passed
to both functions, which would have caught this bug.

Signed-off-by: Ville Vesilehto <ville@vesilehto.fi>

Author: thevilledev

expr.go

@@ -195,9 +195,12 @@ func EnableBuiltin(name string) Option {
 
 // WithContext passes context to all functions calls with a context.Context argument.
 func WithContext(name string) Option {
-	return Patch(patcher.WithContext{
-		Name: name,
-	})
+	return func(c *conf.Config) {
+		c.Visitors = append(c.Visitors, patcher.WithContext{
+			Name:      name,
+			Functions: c.Functions,
+		})
+	}
 }
 
 // Timezone sets default timezone for date() and now() builtin functions.

patcher/with_context.go

@@ -4,11 +4,13 @@ import (
 	"reflect"
 
 	"github.com/expr-lang/expr/ast"
+	"github.com/expr-lang/expr/conf"
 )
 
 // WithContext adds WithContext.Name argument to all functions calls with a context.Context argument.
 type WithContext struct {
-	Name string
+	Name      string
+	Functions conf.FunctionsTable // Optional: used to look up function types when callee type is unknown.
 }
 
 // Visit adds WithContext.Name argument to all functions calls with a context.Context argument.
@@ -19,6 +21,16 @@ func (w WithContext) Visit(node *ast.Node) {
 		if fn == nil {
 			return
 		}
+		// If callee type is interface{} (unknown), try to look up the function type
+		// from the Functions table. This handles cases where the call is nested
+		// inside another call with an unknown callee type.
+		if fn.Kind() == reflect.Interface && w.Functions != nil {
+			if ident, ok := call.Callee.(*ast.IdentifierNode); ok {
+				if f, ok := w.Functions[ident.Value]; ok {
+					fn = f.Type()
+				}
+			}
+		}
 		if fn.Kind() != reflect.Func {
 			return
 		}

test/issues/823/issue_test.go

@@ -2,7 +2,6 @@ package issue_test
 
 import (
 	"context"
-	"fmt"
 	"testing"
 	"time"
 
@@ -14,26 +13,45 @@ type env struct {
 	Ctx context.Context `expr:"ctx"`
 }
 
+// TestIssue823 verifies that WithContext injects context into nested custom
+// function calls. The bug was that date2() nested as an argument to After()
+// didn't receive the context because its callee type was unknown.
 func TestIssue823(t *testing.T) {
+	now2Called := false
+	date2Called := false
+
 	p, err := expr.Compile(
 		"now2().After(date2())",
 		expr.Env(env{}),
 		expr.WithContext("ctx"),
 		expr.Function(
 			"now2",
-			func(params ...any) (any, error) { return time.Now(), nil },
+			func(params ...any) (any, error) {
+				require.Len(t, params, 1, "now2 should receive context")
+				_, ok := params[0].(context.Context)
+				require.True(t, ok, "now2 first param should be context.Context")
+				now2Called = true
+				return time.Now(), nil
+			},
 			new(func(context.Context) time.Time),
 		),
 		expr.Function(
 			"date2",
-			func(params ...any) (any, error) { return time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC), nil },
+			func(params ...any) (any, error) {
+				require.Len(t, params, 1, "date2 should receive context")
+				_, ok := params[0].(context.Context)
+				require.True(t, ok, "date2 first param should be context.Context")
+				date2Called = true
+				return time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC), nil
+			},
 			new(func(context.Context) time.Time),
 		),
 	)
-	fmt.Printf("Compile result err: %v\n", err)
 	require.NoError(t, err)
 
 	r, err := expr.Run(p, &env{Ctx: context.Background()})
 	require.NoError(t, err)
 	require.True(t, r.(bool))
+	require.True(t, now2Called, "now2 should have been called")
+	require.True(t, date2Called, "date2 should have been called")
 }