Fix time bound approval and trust cascading test cases

l46kok · l46kok · commit c0ed7c5868f9 · 2026-03-17T14:00:04.000-07:00
diff --git a/tools/src/main/java/dev/cel/tools/ai/AgentMessageSet.java b/tools/src/main/java/dev/cel/tools/ai/AgentMessageSet.java
@@ -114,10 +114,10 @@ abstract static class Builder {
    * filtered history.
    */
   AgentContext filteredContext() {
-    if (!context().hasExtension(AgentContextExtensions.agentContextMessageHistory)) {
+    List<AgentMessage> msgs = context().getExtension(AgentContextExtensions.agentContextMessageHistory);
+    if (msgs.isEmpty()) {
       return context();
     }
-    List<AgentMessage> msgs = context().getExtension(AgentContextExtensions.agentContextMessageHistory);
     List<AgentMessage> filteredMsgs = new ArrayList<>();
 
     for (AgentMessage msg : msgs) {
@@ -162,10 +162,6 @@ AgentContext filteredContext() {
         filteredParts.add(part);
       }
 
-      if (filteredParts.isEmpty()) {
-        continue;
-      }
-
       filteredMsgs.add(msg.toBuilder().clearParts().addAllParts(filteredParts).build());
     }
 
@@ -198,19 +194,19 @@ AgentMessageSet filterResultType(String resultType) {
    * Returns a new {@link AgentMessageSet} filtered to include messages before the
    * given timestamp.
    */
-  AgentMessageSet filterBefore(Timestamp timestamp) {
+  AgentMessageSet filterBefore(Instant timestamp) {
     return toBuilder()
-        .setBefore(Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()))
+        .setBefore(timestamp)
         .build();
   }
 
   /**
    * Returns a new {@link AgentMessageSet} filtered to include messages after the
    * given timestamp.
    */
-  AgentMessageSet filterAfter(Timestamp timestamp) {
+  AgentMessageSet filterAfter(Instant timestamp) {
     return toBuilder()
-        .setAfter(Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos()))
+        .setAfter(timestamp)
         .build();
   }
 
diff --git a/tools/src/main/java/dev/cel/tools/ai/AgenticPolicyEnvironment.java b/tools/src/main/java/dev/cel/tools/ai/AgenticPolicyEnvironment.java
@@ -4,6 +4,7 @@
 
 import com.google.common.base.Ascii;
 import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.io.Resources;
 import dev.cel.bundle.Cel;
@@ -27,6 +28,7 @@
 import dev.cel.runtime.CelFunctionBinding;
 import java.io.IOException;
 import java.net.URL;
+import java.time.Instant;
 import java.util.List;
 import java.util.Optional;
 
@@ -40,6 +42,7 @@ final class AgenticPolicyEnvironment {
   @SuppressWarnings("Immutable")
   static Cel newInstance(AgentClassifier classifier) {
     AgenticPolicyClassifiers classifiers = new AgenticPolicyClassifiers(classifier);
+
     CelBuilder builder = CelFactory.standardCelBuilder()
         .setContainer(CelContainer.ofName("cel.expr.ai"))
         .addFileTypes(Agent.getDescriptor().getFile())
@@ -109,11 +112,50 @@ static Cel newInstance(AgentClassifier classifier) {
             AgentMessage.class,
             String.class,
             (msg, toolName) -> AgentMessageSet.of(msg).filterToolCall(toolName)),
+        CelFunctionBinding.from(
+            "AgentMessageSet_messages",
+            Object.class,
+            (set) -> {
+              AgentMessageSet messageSet = (AgentMessageSet) set;
+              List<AgentMessage> result = messageSet.filteredContext()
+                  .getExtension(AgentContextExtensions.agentContextMessageHistory);
+              return ImmutableList.copyOf(result);
+            }),
+        CelFunctionBinding.from(
+            "list(Finding)_hasAll_list(Finding)",
+            List.class,
+            List.class,
+            (source, required) -> hasAllFindings(Optional.of((List<Finding>) source), (List<Finding>) required)),
         CelFunctionBinding.from(
             "AgentMessage_role_string",
             AgentMessage.class,
-            String.class,
-            (msg, role) -> AgentMessageSet.of(msg).filterRole(role)));
+            Object.class,
+            (msg, role) -> AgentMessageSet.of(msg).filterRole(String.valueOf(role))),
+        CelFunctionBinding.from(
+            "AgentMessageSet_role_string",
+            AgentMessageSet.class,
+            Object.class,
+            (set, role) -> set.filterRole(String.valueOf(role))),
+        CelFunctionBinding.from(
+            "AgentMessageSet_before_timestamp",
+            AgentMessageSet.class,
+            Instant.class,
+            AgentMessageSet::filterBefore),
+        CelFunctionBinding.from(
+            "AgentMessage_before_timestamp",
+            AgentMessage.class,
+            Instant.class,
+            (msg, ts) -> AgentMessageSet.of(msg).filterBefore(ts)),
+        CelFunctionBinding.from(
+            "AgentMessageSet_after_timestamp",
+            AgentMessageSet.class,
+            Instant.class,
+            AgentMessageSet::filterAfter),
+        CelFunctionBinding.from(
+            "AgentMessage_after_timestamp",
+            AgentMessage.class,
+            Instant.class,
+            (msg, ts) -> AgentMessageSet.of(msg).filterAfter(ts)));
 
     Cel celEnv = builder.build();
     celEnv = extendFromConfig(celEnv, "environment/agent_env.yaml");
@@ -131,10 +173,6 @@ private static boolean hasAllFindings(Optional<List<Finding>> sourceOpt, List<Fi
         act.getConfidence() >= req.getConfidence()));
   }
 
-  static Cel newInstance() {
-    return newInstance(AgentClassifier.DEFAULT);
-  }
-
   private static Cel extendFromConfig(Cel cel, String yamlConfigPath) {
     String yamlEnv;
     try {
diff --git a/tools/src/main/java/dev/cel/tools/ai/BUILD.bazel b/tools/src/main/java/dev/cel/tools/ai/BUILD.bazel
@@ -53,12 +53,14 @@ java_library(
         "//common:container",
         "//common:options",
         "//common/types",
+        "//common/types:message_type_provider",
         "//common/types:type_providers",
         "//parser:macro",
         "//runtime:function_binding",
         "//:auto_value",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java",
+        "@maven//:com_google_protobuf_protobuf_java_util",
     ],
 )
 
diff --git a/tools/src/main/resources/environment/common_env.yaml b/tools/src/main/resources/environment/common_env.yaml
@@ -690,6 +690,22 @@ functions:
       params:
       - type_name: cel.expr.ai.AgentMessage.Part
 
+- name: "messages"
+  description: |
+    Returns the ordered list of AgentMessages in the message set.
+  overloads:
+  - id: "AgentMessageSet_messages"
+    examples:
+    - |
+      // Returns the ordered list of messages in the message set.
+      agent.history.messages()
+    target:
+      type_name: cel.expr.ai.AgentMessageSet
+    return:
+      type_name: list
+      params:
+      - type_name: dyn
+
 - name: "spec"
   description: |
     Returns the specification for the tool.
diff --git a/tools/src/test/java/dev/cel/tools/ai/AgenticPolicyCompilerTest.java b/tools/src/test/java/dev/cel/tools/ai/AgenticPolicyCompilerTest.java
@@ -4,7 +4,6 @@
 
 import com.google.common.base.Ascii;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.io.Resources;
 import com.google.common.truth.Expect;
 import com.google.testing.junit.testparameterinjector.TestParameter;
@@ -28,6 +27,7 @@
 import java.io.IOException;
 import java.net.URL;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Optional;
 import org.junit.Rule;
@@ -58,15 +58,13 @@ private enum AgenticPolicyTestCase {
         "require_user_confirmation_for_tool_tests.yaml"),
     OPEN_WORLD_TOOL_REPLAY(
         "open_world_tool_replay.celpolicy",
-        "open_world_tool_replay_tests.yaml");
-    // TRUST_CASCADING(
-    // "trust_cascading.celpolicy",
-    // "trust_cascading_tests.yaml"
-    // ),
-    // TIME_BOUND_APPROVAL(
-    // "time_bound_approval.celpolicy",
-    // "time_bound_approval_tests.yaml"
-    // );
+        "open_world_tool_replay_tests.yaml"),
+    TRUST_CASCADING(
+        "trust_cascading.celpolicy",
+        "trust_cascading_tests.yaml"),
+    TIME_BOUND_APPROVAL(
+        "time_bound_approval.celpolicy",
+        "time_bound_approval_tests.yaml");
 
     private final String policyFilePath;
     private final String policyTestCaseFilePath;
@@ -94,17 +92,24 @@ private void runTests(Cel cel, CelAbstractSyntaxTree ast, PolicyTestSuite testSu
         String testName = String.format(
             "%s: %s", testSection.getName(), testCase.getName());
         try {
-          ImmutableMap<String, Object> inputMap = testCase.toInputMap(cel);
+          HashMap<String, Object> inputMap = new HashMap<>(testCase.toInputMap(cel));
 
           List<AgentMessage> history = inputMap.containsKey("_test_history")
               ? (List<AgentMessage>) inputMap.get("_test_history")
               : ImmutableList.of();
 
+          AgentContext context = AgentContext.newBuilder()
+              .setExtension(AgentContextExtensions.agentContextMessageHistory, history)
+              .build();
+          AgentMessageSet messageSet = AgentMessageSet.of(context);
+
+          inputMap.put("agent.history", messageSet);
+
           @SuppressWarnings("Immutable")
           CelLateFunctionBindings bindings = CelLateFunctionBindings.from(
               CelFunctionBinding.from(
                   "agent_history",
-                  ImmutableList.of(), // No args
+                  ImmutableList.of(),
                   (args) -> history));
 
           Object evalResult = cel.createProgram(ast).eval(inputMap, bindings);
diff --git a/tools/src/test/resources/policy/require_user_confirmation_for_tool.celpolicy b/tools/src/test/resources/policy/require_user_confirmation_for_tool.celpolicy
@@ -3,7 +3,7 @@ default: deny
 
 variables:
   - high_confidence_pii: >
-      tool.call.sensitivityFindings('pii').orValue([]).exists(f, f.confidence >= 0.8)
+      tool.call.sensitivityFindings('pii').hasAll([ai.finding("pii_score", 0.8)])
 
 rules:
   - description: "Confirm tool calls if high-confidence PII is detected"
diff --git a/tools/src/test/resources/policy/time_bound_approval.celpolicy b/tools/src/test/resources/policy/time_bound_approval.celpolicy
@@ -5,11 +5,11 @@ variables:
   # Define the validity window (30 seconds ago)
   - approval_cutoff: now - duration('30s')
 
-  # Find approval messages in the valid window
   - valid_approvals: >
-      agent.history()
+      agent.history
            .after(variables.approval_cutoff)
            .role('model')
+           .messages()
            .filter(m, has(m.metadata.step) && m.metadata.step == 'approval_granted')
 
   - has_valid_approval: variables.valid_approvals.size() > 0
diff --git a/tools/src/test/resources/policy/trust_cascading.celpolicy b/tools/src/test/resources/policy/trust_cascading.celpolicy
@@ -4,11 +4,11 @@ default: allow
 variables:
   # Critical security threats
   - is_compromised: >
-      agent.context.trust.findings.contains([ai.finding("compromised_session", 0.9)])
+      agent.context.trust.findings.hasAll([ai.finding("compromised_session", 0.9)])
 
   # Compliance and/or hygiene issues with the source
   - is_unverified: >
-      agent.context.trust.findings.contains([ai.finding("unverified_source", 0.8)])
+      agent.context.trust.findings.hasAll([ai.finding("unverified_source", 0.8)])
 
 rules:
   - description: "Block sessions with high-confidence compromise indicators"
