Allow AsyncProgram to accept an async resolver for an attribute

runtime/src/main/java/dev/cel/runtime/UnknownContext.java

@@ -55,7 +55,7 @@ private UnknownContext(
       ImmutableList<CelAttributePattern> unresolvedAttributes,
       ImmutableMap<CelAttribute, Object> resolvedAttributes) {
     this.unresolvedAttributes = unresolvedAttributes;
-    variableResolver = resolver;
+    this.variableResolver = resolver;
     this.resolvedAttributes = resolvedAttributes;
   }
 
@@ -76,6 +76,17 @@ public static UnknownContext create(
         createExprVariableResolver(resolver), ImmutableList.copyOf(attributes), ImmutableMap.of());
   }
 
+  /** Extends an existing {@code UnknownContext} by adding more attribute patterns to it. */
+  public UnknownContext extend(Collection<CelAttributePattern> attributePatterns) {
+    return new UnknownContext(
+        this.variableResolver(),
+        ImmutableList.<CelAttributePattern>builder()
+            .addAll(this.unresolvedAttributes)
+            .addAll(attributePatterns)
+            .build(),
+        this.resolvedAttributes);
+  }
+
   /** Adapts a CelVariableResolver to the legacy impl equivalent GlobalResolver. */
   private static GlobalResolver createExprVariableResolver(CelVariableResolver resolver) {
     return (String name) -> resolver.find(name).orElse(null);

runtime/src/main/java/dev/cel/runtime/async/AsyncProgramImpl.java

@@ -14,12 +14,14 @@
 
 package dev.cel.runtime.async;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 import static com.google.common.util.concurrent.Futures.transformAsync;
 import static com.google.common.util.concurrent.Futures.whenAllSucceed;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.Futures;
@@ -55,6 +57,7 @@
 final class AsyncProgramImpl implements CelAsyncRuntime.AsyncProgram {
   // Safety limit for resolution rounds.
   private final int maxEvaluateIterations;
+  private final UnknownContext startingUnknownContext;
   private final Program program;
   private final ListeningExecutorService executor;
   private final ImmutableMap<CelAttributePattern, CelUnknownAttributeValueResolver> resolvers;
@@ -63,15 +66,26 @@ final class AsyncProgramImpl implements CelAsyncRuntime.AsyncProgram {
       Program program,
       ListeningExecutorService executor,
       ImmutableMap<CelAttributePattern, CelUnknownAttributeValueResolver> resolvers,
-      int maxEvaluateIterations) {
+      int maxEvaluateIterations,
+      UnknownContext startingUnknownContext) {
     this.program = program;
     this.executor = executor;
     this.resolvers = resolvers;
     this.maxEvaluateIterations = maxEvaluateIterations;
+    // The following is populated from CelAsyncRuntime. The impl is immutable, thus safe to reuse as
+    // a starting context.
+    this.startingUnknownContext = startingUnknownContext;
   }
 
-  private Optional<CelUnknownAttributeValueResolver> lookupResolver(CelAttribute attribute) {
+  private Optional<CelUnknownAttributeValueResolver> lookupResolver(
+      Iterable<CelResolvableAttributePattern> resolvableAttributePatterns, CelAttribute attribute) {
     // TODO: may need to handle multiple resolvers for partial case.
+    for (CelResolvableAttributePattern entry : resolvableAttributePatterns) {
+      if (entry.attributePattern().isPartialMatch(attribute)) {
+        return Optional.of(entry.resolver());
+      }
+    }
+
     for (Map.Entry<CelAttributePattern, CelUnknownAttributeValueResolver> entry :
         resolvers.entrySet()) {
       if (entry.getKey().isPartialMatch(attribute)) {
@@ -102,10 +116,14 @@ private ListenableFuture<ImmutableMap<CelAttribute, Object>> allAsMapOnSuccess(
   }
 
   private ListenableFuture<Object> resolveAndReevaluate(
-      CelUnknownSet unknowns, UnknownContext ctx, int iteration) {
+      CelUnknownSet unknowns,
+      UnknownContext ctx,
+      Iterable<CelResolvableAttributePattern> resolvableAttributePatterns,
+      int iteration) {
     Map<CelAttribute, ListenableFuture<Object>> futureMap = new LinkedHashMap<>();
     for (CelAttribute attr : unknowns.attributes()) {
-      Optional<CelUnknownAttributeValueResolver> maybeResolver = lookupResolver(attr);
+      Optional<CelUnknownAttributeValueResolver> maybeResolver =
+          lookupResolver(resolvableAttributePatterns, attr);
 
       maybeResolver.ifPresent((resolver) -> futureMap.put(attr, resolver.resolve(executor, attr)));
     }
@@ -120,13 +138,21 @@ private ListenableFuture<Object> resolveAndReevaluate(
     // need to be configurable in the future.
     return transformAsync(
         allAsMapOnSuccess(futureMap),
-        (result) -> evalPass(ctx.withResolvedAttributes(result), unknowns, iteration),
+        (result) ->
+            evalPass(
+                ctx.withResolvedAttributes(result),
+                resolvableAttributePatterns,
+                unknowns,
+                iteration),
         executor);
   }
 
   private ListenableFuture<Object> evalPass(
-      UnknownContext ctx, CelUnknownSet lastSet, int iteration) {
-    Object result = null;
+      UnknownContext ctx,
+      Iterable<CelResolvableAttributePattern> resolvableAttributePatterns,
+      CelUnknownSet lastSet,
+      int iteration) {
+    Object result;
     try {
       result = program.advanceEvaluation(ctx);
     } catch (CelEvaluationException e) {
@@ -143,14 +169,29 @@ private ListenableFuture<Object> evalPass(
         return immediateFailedFuture(
             new CelEvaluationException("Max Evaluation iterations exceeded: " + iteration));
       }
-      return resolveAndReevaluate((CelUnknownSet) result, ctx, iteration);
+      return resolveAndReevaluate(
+          (CelUnknownSet) result, startingUnknownContext, resolvableAttributePatterns, iteration);
     }
 
     return immediateFuture(result);
   }
 
   @Override
-  public ListenableFuture<Object> evaluateToCompletion(UnknownContext ctx) {
-    return evalPass(ctx, CelUnknownSet.create(ImmutableSet.of()), 0);
+  public ListenableFuture<Object> evaluateToCompletion(
+      CelResolvableAttributePattern... resolvableAttributes) {
+    return evaluateToCompletion(ImmutableList.copyOf(resolvableAttributes));
+  }
+
+  @Override
+  public ListenableFuture<Object> evaluateToCompletion(
+      Iterable<CelResolvableAttributePattern> resolvableAttributePatterns) {
+    UnknownContext newAsyncContext =
+        startingUnknownContext.extend(
+            ImmutableList.copyOf(resolvableAttributePatterns).stream()
+                .map(CelResolvableAttributePattern::attributePattern)
+                .collect(toImmutableList()));
+
+    return evalPass(
+        newAsyncContext, resolvableAttributePatterns, CelUnknownSet.create(ImmutableSet.of()), 0);
   }
 }

runtime/src/main/java/dev/cel/runtime/async/BUILD.bazel

@@ -13,6 +13,7 @@ package(
 
 ASYNC_RUNTIME_SOURCES = [
     "AsyncProgramImpl.java",
+    "CelResolvableAttributePattern.java",
     "CelAsyncRuntimeImpl.java",
     "CelAsyncRuntime.java",
     "CelAsyncRuntimeBuilder.java",

runtime/src/main/java/dev/cel/runtime/async/CelAsyncRuntime.java

@@ -36,17 +36,13 @@
 @ThreadSafe
 public interface CelAsyncRuntime {
 
-  /**
-   * Initialize a new async context for iterative evaluation.
-   *
-   * <p>This maintains the state related to tracking which parts of the environment are unknown or
-   * have been resolved.
-   */
-  UnknownContext newAsyncContext();
-
   /** AsyncProgram wraps a CEL Program with a driver to resolve unknowns as they are encountered. */
   interface AsyncProgram {
-    ListenableFuture<Object> evaluateToCompletion(UnknownContext ctx);
+    ListenableFuture<Object> evaluateToCompletion(
+        CelResolvableAttributePattern... resolvableAttributes);
+
+    ListenableFuture<Object> evaluateToCompletion(
+        Iterable<CelResolvableAttributePattern> resolvableAttributes);
   }
 
   /**

runtime/src/main/java/dev/cel/runtime/async/CelAsyncRuntimeBuilder.java

@@ -1,5 +1,4 @@
 // Copyright 2022 Google LLC
-//
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
@@ -17,23 +16,36 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import dev.cel.runtime.CelAttributePattern;
 import dev.cel.runtime.CelRuntime;
+import dev.cel.runtime.async.CelAsyncRuntime.AsyncProgram;
 import java.util.concurrent.ExecutorService;
 
 /** Builder interface for {@link CelAsyncRuntime}. */
 public interface CelAsyncRuntimeBuilder {
-  public static final int DEFAULT_MAX_EVALUATE_ITERATIONS = 10;
+  int DEFAULT_MAX_EVALUATE_ITERATIONS = 10;
 
   /** Set the CEL runtime for running incremental evaluation. */
   @CanIgnoreReturnValue
-  public CelAsyncRuntimeBuilder setRuntime(CelRuntime runtime);
+  CelAsyncRuntimeBuilder setRuntime(CelRuntime runtime);
 
-  /** Add attributes that are declared as Unknown, without any resolver. */
+  /**
+   * Add attributes that are declared as Unknown, without any resolver.
+   *
+   * @deprecated Use {@link AsyncProgram#evaluateToCompletion(CelResolvableAttributePattern...)}
+   *     instead to propagate the unknown attributes along with the resolvers into the program.
+   */
   @CanIgnoreReturnValue
-  public CelAsyncRuntimeBuilder addUnknownAttributePatterns(CelAttributePattern... attributes);
+  @Deprecated
+  CelAsyncRuntimeBuilder addUnknownAttributePatterns(CelAttributePattern... attributes);
 
-  /** Marks an attribute pattern as unknown and associates a resolver with it. */
+  /**
+   * Marks an attribute pattern as unknown and associates a resolver with it.
+   *
+   * @deprecated Use {@link AsyncProgram#evaluateToCompletion(CelResolvableAttributePattern...)}
+   *     instead to propagate the unknown attributes along with the resolvers into the program.
+   */
   @CanIgnoreReturnValue
-  public CelAsyncRuntimeBuilder addResolvableAttributePattern(
+  @Deprecated
+  CelAsyncRuntimeBuilder addResolvableAttributePattern(
       CelAttributePattern attribute, CelUnknownAttributeValueResolver resolver);
 
   /**
@@ -42,10 +54,10 @@ public CelAsyncRuntimeBuilder addResolvableAttributePattern(
    * <p>This is a safety mechanism for expressions that chain dependent unknowns (e.g. via the
    * conditional operator or nested function calls).
    *
-   * <p>Implementations should default to {@value DEFAULT_MAX_EVALUATION_ITERATIONS}.
+   * <p>Implementations should default to {@value DEFAULT_MAX_EVALUATE_ITERATIONS}.
    */
   @CanIgnoreReturnValue
-  public CelAsyncRuntimeBuilder setMaxEvaluateIterations(int n);
+  CelAsyncRuntimeBuilder setMaxEvaluateIterations(int n);
 
   /**
    * Sets the variable resolver for simple CelVariable names (e.g. 'x' or 'com.google.x').
@@ -67,7 +79,7 @@ public CelAsyncRuntimeBuilder addResolvableAttributePattern(
    * resolvers.
    */
   @CanIgnoreReturnValue
-  public CelAsyncRuntimeBuilder setExecutorService(ExecutorService executorService);
+  CelAsyncRuntimeBuilder setExecutorService(ExecutorService executorService);
 
-  public CelAsyncRuntime build();
+  CelAsyncRuntime build();
 }

runtime/src/main/java/dev/cel/runtime/async/CelAsyncRuntimeImpl.java

@@ -55,8 +55,7 @@ private CelAsyncRuntimeImpl(
     this.maxEvaluateIterations = maxEvaluateIterations;
   }
 
-  @Override
-  public UnknownContext newAsyncContext() {
+  private UnknownContext newAsyncContext() {
     return UnknownContext.create(variableResolver, unknownAttributePatterns);
   }
 
@@ -66,7 +65,8 @@ public AsyncProgram createProgram(CelAbstractSyntaxTree ast) throws CelEvaluatio
         runtime.createProgram(ast),
         executorService,
         unknownAttributeResolvers,
-        maxEvaluateIterations);
+        maxEvaluateIterations,
+        newAsyncContext());
   }
 
   static Builder newBuilder() {

runtime/src/main/java/dev/cel/runtime/async/CelResolvableAttributePattern.java

@@ -0,0 +1,36 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.runtime.async;
+
+import com.google.auto.value.AutoValue;
+import dev.cel.runtime.CelAttributePattern;
+
+/**
+ * CelResolvableAttributePattern wraps {@link CelAttributePattern} to represent a CEL attribute
+ * whose value is initially unknown and needs to be resolved. It couples the attribute pattern with
+ * a {@link CelUnknownAttributeValueResolver} that can fetch the actual value for the attribute when
+ * it becomes available.
+ */
+@AutoValue
+public abstract class CelResolvableAttributePattern {
+  public abstract CelAttributePattern attributePattern();
+
+  public abstract CelUnknownAttributeValueResolver resolver();
+
+  public static CelResolvableAttributePattern of(
+      CelAttributePattern attribute, CelUnknownAttributeValueResolver resolver) {
+    return new AutoValue_CelResolvableAttributePattern(attribute, resolver);
+  }
+}