Track in-flight async call IDs in AccumulatedUnknowns

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

@@ -16,9 +16,8 @@
 
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import dev.cel.common.annotations.Internal;
-import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import org.jspecify.annotations.Nullable;
@@ -35,6 +34,7 @@ public final class AccumulatedUnknowns {
   private static final int MAX_UNKNOWN_ATTRIBUTE_SIZE = 500_000;
   private final Set<Long> exprIds;
   private final Set<CelAttribute> attributes;
+  private final Set<Long> callIds;
 
   Set<Long> exprIds() {
     return exprIds;
@@ -44,6 +44,17 @@ Set<CelAttribute> attributes() {
     return attributes;
   }
 
+  /**
+   * Returns the in-flight asynchronous call IDs this unknown is waiting on.
+   *
+   * <p>The returned set is an unmodifiable <em>view</em> over this mutable accumulator, not a
+   * snapshot: a subsequent {@link #merge} on this instance is visible through it. Callers that
+   * retain the set beyond the current evaluation step must copy it.
+   */
+  public Set<Long> callIds() {
+    return Collections.unmodifiableSet(callIds);
+  }
+
   /**
    * Evaluates if the right hand side is an accumulated unknown, and if so, merges it into the
    * accumulator.
@@ -62,20 +73,29 @@ public AccumulatedUnknowns merge(AccumulatedUnknowns arg) {
     enforceMaxAttributeSize(this.attributes, arg.attributes);
     this.exprIds.addAll(arg.exprIds);
     this.attributes.addAll(arg.attributes);
+    this.callIds.addAll(arg.callIds);
     return this;
   }
 
-  static AccumulatedUnknowns create(Long... ids) {
-    return create(Arrays.asList(ids));
-  }
-
-  static AccumulatedUnknowns create(Collection<Long> ids) {
-    return create(ids, new ArrayList<>());
+  static AccumulatedUnknowns create(long exprId) {
+    return new AccumulatedUnknowns(
+        Collections.singletonList(exprId), Collections.emptyList(), Collections.emptyList());
   }
 
   public static AccumulatedUnknowns create(
       Collection<Long> exprIds, Collection<CelAttribute> attributes) {
-    return new AccumulatedUnknowns(new HashSet<>(exprIds), new HashSet<>(attributes));
+    return new AccumulatedUnknowns(exprIds, attributes, Collections.emptyList());
+  }
+
+  /**
+   * Creates an accumulated unknown for a pending asynchronous call, recording {@code exprId} so the
+   * unknown retains its origin when adapted into a {@link CelUnknownSet}.
+   */
+  public static AccumulatedUnknowns createForAsyncCall(long exprId, long callId) {
+    return new AccumulatedUnknowns(
+        Collections.singletonList(exprId),
+        Collections.emptyList(),
+        Collections.singletonList(callId));
   }
 
   private static void enforceMaxAttributeSize(
@@ -88,8 +108,10 @@ private static void enforceMaxAttributeSize(
     }
   }
 
-  private AccumulatedUnknowns(Set<Long> exprIds, Set<CelAttribute> attributes) {
-    this.exprIds = exprIds;
-    this.attributes = attributes;
+  private AccumulatedUnknowns(
+      Collection<Long> exprIds, Collection<CelAttribute> attributes, Collection<Long> callIds) {
+    this.exprIds = new HashSet<>(exprIds);
+    this.attributes = new HashSet<>(attributes);
+    this.callIds = new HashSet<>(callIds);
   }
 }

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

@@ -16,6 +16,7 @@
 
 import dev.cel.common.annotations.Internal;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Optional;
 
 /**
@@ -100,7 +101,7 @@ Optional<Object> maybeUnknowns() {
     }
 
     if (!exprIds.isEmpty()) {
-      return Optional.of(AccumulatedUnknowns.create(exprIds));
+      return Optional.of(AccumulatedUnknowns.create(exprIds, Collections.emptyList()));
     }
 
     return Optional.empty();

runtime/src/test/java/dev/cel/runtime/AccumulatedUnknownsTest.java

@@ -0,0 +1,131 @@
+// Copyright 2026 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;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
+
+import com.google.common.collect.ImmutableList;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import java.util.Set;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(TestParameterInjector.class)
+public final class AccumulatedUnknownsTest {
+
+  @Test
+  public void createForAsyncCall_success() {
+    AccumulatedUnknowns unknowns = AccumulatedUnknowns.createForAsyncCall(7L, 42L);
+
+    assertThat(unknowns.callIds()).containsExactly(42L);
+    assertThat(unknowns.exprIds()).containsExactly(7L);
+    assertThat(unknowns.attributes()).isEmpty();
+  }
+
+  @Test
+  public void createForAsyncCall_adaptedToCelUnknownSet_retainsExprId() {
+    AccumulatedUnknowns unknowns = AccumulatedUnknowns.createForAsyncCall(7L, 42L);
+
+    Object adapted = InterpreterUtil.maybeAdaptToCelUnknownSet(unknowns);
+
+    assertThat(adapted).isEqualTo(CelUnknownSet.create(7L));
+  }
+
+  @Test
+  public void createForAsyncCall_distinctExprIds_adaptToDistinctUnknownSets() {
+    Object first =
+        InterpreterUtil.maybeAdaptToCelUnknownSet(AccumulatedUnknowns.createForAsyncCall(7L, 42L));
+    Object second =
+        InterpreterUtil.maybeAdaptToCelUnknownSet(AccumulatedUnknowns.createForAsyncCall(8L, 43L));
+
+    assertThat(first).isNotEqualTo(second);
+  }
+
+  @Test
+  public void callIds_returnsUnmodifiableSet() {
+    AccumulatedUnknowns unknowns = AccumulatedUnknowns.createForAsyncCall(7L, 42L);
+    Set<Long> callIds = unknowns.callIds();
+
+    assertThrows(UnsupportedOperationException.class, () -> callIds.add(99L));
+  }
+
+  @Test
+  public void merge_mergesCallIdsAndExprIdsAndAttributes() {
+    AccumulatedUnknowns u1 =
+        AccumulatedUnknowns.create(ImmutableList.of(1L), ImmutableList.of(CelAttribute.EMPTY));
+    u1.merge(AccumulatedUnknowns.createForAsyncCall(1L, 100L));
+    AccumulatedUnknowns u2 = AccumulatedUnknowns.create(ImmutableList.of(2L), ImmutableList.of());
+    u2.merge(AccumulatedUnknowns.createForAsyncCall(2L, 200L));
+
+    AccumulatedUnknowns merged = u1.merge(u2);
+
+    assertThat(merged).isSameInstanceAs(u1);
+    assertThat(merged.exprIds()).containsExactly(1L, 2L);
+    assertThat(merged.attributes()).containsExactly(CelAttribute.EMPTY);
+    assertThat(merged.callIds()).containsExactly(100L, 200L);
+  }
+
+  @Test
+  public void maybeMerge_withNullAccumulator_returnsNewUnknowns() {
+    AccumulatedUnknowns u = AccumulatedUnknowns.createForAsyncCall(1L, 1L);
+
+    AccumulatedUnknowns result = AccumulatedUnknowns.maybeMerge(null, u);
+
+    assertThat(result).isSameInstanceAs(u);
+  }
+
+  @Test
+  public void maybeMerge_withExistingAccumulator_mergesBoth() {
+    AccumulatedUnknowns u1 = AccumulatedUnknowns.createForAsyncCall(1L, 1L);
+    AccumulatedUnknowns u2 = AccumulatedUnknowns.createForAsyncCall(2L, 2L);
+
+    AccumulatedUnknowns result = AccumulatedUnknowns.maybeMerge(u1, u2);
+
+    assertThat(result).isSameInstanceAs(u1);
+    assertThat(result.callIds()).containsExactly(1L, 2L);
+    assertThat(result.exprIds()).containsExactly(1L, 2L);
+  }
+
+  @Test
+  public void maybeMerge_withNonUnknownObject_returnsOriginalAccumulator() {
+    AccumulatedUnknowns u = AccumulatedUnknowns.createForAsyncCall(1L, 1L);
+
+    AccumulatedUnknowns result = AccumulatedUnknowns.maybeMerge(u, "not an unknown");
+
+    assertThat(result).isSameInstanceAs(u);
+    assertThat(result.callIds()).containsExactly(1L);
+  }
+
+  @Test
+  public void create_singleExprId_populatesExprIdsOnly() {
+    AccumulatedUnknowns u = AccumulatedUnknowns.create(10L);
+
+    assertThat(u.exprIds()).containsExactly(10L);
+    assertThat(u.attributes()).isEmpty();
+    assertThat(u.callIds()).isEmpty();
+  }
+
+  @Test
+  public void create_exprIdsAndAttributes_leavesCallIdsEmpty() {
+    AccumulatedUnknowns u =
+        AccumulatedUnknowns.create(
+            ImmutableList.of(10L, 20L), ImmutableList.of(CelAttribute.EMPTY));
+
+    assertThat(u.exprIds()).containsExactly(10L, 20L);
+    assertThat(u.attributes()).containsExactly(CelAttribute.EMPTY);
+    assertThat(u.callIds()).isEmpty();
+  }
+}

runtime/src/test/java/dev/cel/runtime/BUILD.bazel

@@ -65,6 +65,7 @@ java_library(
         "//parser:macro",
         "//parser:unparser",
         "//runtime",
+        "//runtime:accumulated_unknowns",
         "//runtime:activation",
         "//runtime:dispatcher",
         "//runtime:evaluation_exception_builder",
@@ -73,6 +74,7 @@ java_library(
         "//runtime:function_resolver",
         "//runtime:interpretable",
         "//runtime:interpreter",
+        "//runtime:interpreter_util",
         "//runtime:late_function_binding",
         "//runtime:lite_runtime",
         "//runtime:lite_runtime_factory",