fix(patcher): ctx into nested custom funcs and Env

expr.go

@@ -195,9 +195,14 @@ 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,
+			Env:       &c.Env,
+			NtCache:   &c.NtCache,
+		})
+	}
 }
 
 // Timezone sets default timezone for date() and now() builtin functions.

patcher/with_context.go

@@ -4,11 +4,16 @@ import (
 	"reflect"
 
 	"github.com/expr-lang/expr/ast"
+	"github.com/expr-lang/expr/checker/nature"
+	"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.
+	Env       *nature.Nature      // Optional: used to look up method types when callee type is unknown.
+	NtCache   *nature.Cache       // Optional: cache for nature lookups.
 }
 
 // Visit adds WithContext.Name argument to all functions calls with a context.Context argument.
@@ -19,6 +24,24 @@ func (w WithContext) Visit(node *ast.Node) {
 		if fn == nil {
 			return
 		}
+		// If callee type is interface{} (unknown), look up the function type from
+		// the Functions table or Env. This handles cases where the checker returns early
+		// without visiting nested call arguments (e.g., Date2() in Now2().After(Date2()))
+		// because the outer call's type is unknown due to missing context arguments.
+		if fn.Kind() == reflect.Interface {
+			if ident, ok := call.Callee.(*ast.IdentifierNode); ok {
+				if w.Functions != nil {
+					if f, ok := w.Functions[ident.Value]; ok {
+						fn = f.Type()
+					}
+				}
+				if fn.Kind() == reflect.Interface && w.Env != nil {
+					if m, ok := w.Env.MethodByName(w.NtCache, ident.Value); ok {
+						fn = m.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,72 @@ 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")
+}
+
+// envWithMethods tests that Env methods with context.Context work correctly
+// when nested in method chains (similar to TestIssue823 but with Env methods).
+type envWithMethods struct {
+	Ctx context.Context `expr:"ctx"`
+}
+
+func (e *envWithMethods) Now2(ctx context.Context) time.Time {
+	return time.Now()
+}
+
+func (e *envWithMethods) Date2(ctx context.Context) time.Time {
+	return time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
+}
+
+func TestIssue823_EnvMethods(t *testing.T) {
+	p, err := expr.Compile(
+		"Now2().After(Date2())",
+		expr.Env(&envWithMethods{}),
+		expr.WithContext("ctx"),
+	)
+	require.NoError(t, err)
+
+	r, err := expr.Run(p, &envWithMethods{Ctx: context.Background()})
+	require.NoError(t, err)
+	require.True(t, r.(bool))
 }