add a policy for time bound approval

l46kok · l46kok · commit fdd790d455cf · 2026-02-06T13:24:36.000-08:00
diff --git a/tools/src/test/java/dev/cel/tools/ai/AgenticPolicyCompilerTest.java b/tools/src/test/java/dev/cel/tools/ai/AgenticPolicyCompilerTest.java
@@ -21,23 +21,26 @@
 import dev.cel.common.types.SimpleType;
 import dev.cel.common.types.StructTypeReference;
 import dev.cel.expr.ai.Agent;
-import dev.cel.expr.ai.AgentContext; // New Import
+import dev.cel.expr.ai.AgentContext;
 import dev.cel.expr.ai.AgentMessage;
 import dev.cel.expr.ai.Finding;
 import dev.cel.expr.ai.Tool;
 import dev.cel.expr.ai.ToolAnnotations;
 import dev.cel.expr.ai.ToolCall;
-import dev.cel.expr.ai.TrustLevel; // New Import
+import dev.cel.expr.ai.TrustLevel;
 import dev.cel.parser.CelStandardMacro;
 import dev.cel.policy.testing.PolicyTestSuiteHelper;
 import dev.cel.policy.testing.PolicyTestSuiteHelper.PolicyTestSuite;
 import dev.cel.policy.testing.PolicyTestSuiteHelper.PolicyTestSuite.PolicyTestSection;
 import dev.cel.policy.testing.PolicyTestSuiteHelper.PolicyTestSuite.PolicyTestSection.PolicyTestCase;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.CelFunctionBinding;
+import dev.cel.runtime.CelLateFunctionBindings;
 import java.io.IOException;
 import java.net.URL;
+import java.time.Instant;
 import java.util.List;
+import java.util.stream.Collectors;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -61,10 +64,12 @@ public class AgenticPolicyCompilerTest {
 
       .addVar("agent.input", StructTypeReference.create("cel.expr.ai.AgentMessage"))
       .addVar("agent.context", StructTypeReference.create("cel.expr.ai.AgentContext"))
+      .addVar("_test_history", ListType.create(StructTypeReference.create("cel.expr.ai.AgentMessage")))
+      .addVar("now", SimpleType.TIMESTAMP)
+
       .addVar("tool.name", SimpleType.STRING)
       .addVar("tool.annotations", StructTypeReference.create("cel.expr.ai.ToolAnnotations"))
       .addVar("tool.call", StructTypeReference.create("cel.expr.ai.ToolCall"))
-
       .addFunctionDeclarations(
           newFunctionDeclaration(
               "ai.finding",
@@ -100,6 +105,31 @@ public class AgenticPolicyCompilerTest {
                   ListType.create(StructTypeReference.create("cel.expr.ai.Finding")),
                   ListType.create(StructTypeReference.create("cel.expr.ai.Finding"))
               )
+          ),
+          newFunctionDeclaration(
+              "agent.history",
+              newGlobalOverload(
+                  "agent_history",
+                  ListType.create(StructTypeReference.create("cel.expr.ai.AgentMessage"))
+              )
+          ),
+          newFunctionDeclaration(
+              "role",
+              newMemberOverload(
+                  "list_agent_message_role_string",
+                  ListType.create(StructTypeReference.create("cel.expr.ai.AgentMessage")),
+                  ListType.create(StructTypeReference.create("cel.expr.ai.AgentMessage")),
+                  SimpleType.STRING
+              )
+          ),
+          newFunctionDeclaration(
+              "after",
+              newMemberOverload(
+                  "list_agent_message_after_timestamp",
+                  ListType.create(StructTypeReference.create("cel.expr.ai.AgentMessage")),
+                  ListType.create(StructTypeReference.create("cel.expr.ai.AgentMessage")),
+                  SimpleType.TIMESTAMP
+              )
           )
       )
       .addFunctionBindings(
@@ -151,14 +181,40 @@ public class AgenticPolicyCompilerTest {
               (args) -> {
                 List<Finding> actualFindings = (List<Finding>) args[0];
                 List<Finding> expectedFindings = (List<Finding>) args[1];
-
                 return expectedFindings.stream().anyMatch(expected ->
                     actualFindings.stream().anyMatch(actual ->
                         actual.getValue().equals(expected.getValue()) &&
                             actual.getConfidence() >= expected.getConfidence()
                     )
                 );
               }
+          ),
+          CelFunctionBinding.from(
+              "list_agent_message_role_string",
+              ImmutableList.of(List.class, String.class),
+              (args) -> {
+                List<AgentMessage> history = (List<AgentMessage>) args[0];
+                String role = (String) args[1];
+                return history.stream()
+                    .filter(m -> m.getRole().equals(role))
+                    .collect(Collectors.toList());
+              }
+          ),
+          CelFunctionBinding.from(
+              "list_agent_message_after_timestamp",
+              ImmutableList.of(List.class, Instant.class),
+              (args) -> {
+                List<AgentMessage> history = (List<AgentMessage>) args[0];
+                Instant cutoff = (Instant) args[1];
+
+                return history.stream()
+                    .filter(m -> {
+                      com.google.protobuf.Timestamp protoTs = m.getTime();
+                      Instant msgTime = Instant.ofEpochSecond(protoTs.getSeconds(), protoTs.getNanos());
+                      return msgTime.compareTo(cutoff) >= 0;
+                    })
+                    .collect(Collectors.toList());
+              }
           )
       )
       .build();
@@ -188,6 +244,10 @@ private enum AgenticPolicyTestCase {
     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;
@@ -217,7 +277,21 @@ private void runTests(Cel cel, CelAbstractSyntaxTree ast, PolicyTestSuite testSu
             "%s: %s", testSection.getName(), testCase.getName());
         try {
           ImmutableMap<String, Object> inputMap = testCase.toInputMap(cel);
-          Object evalResult = cel.createProgram(ast).eval(inputMap);
+
+          List<AgentMessage> history =
+              inputMap.containsKey("_test_history")
+                  ? (List<AgentMessage>) inputMap.get("_test_history")
+                  : ImmutableList.of();
+
+          CelLateFunctionBindings bindings = CelLateFunctionBindings.from(
+              CelFunctionBinding.from(
+                  "agent_history",
+                  ImmutableList.of(), // No args
+                  (args) -> history
+              )
+          );
+
+          Object evalResult = cel.createProgram(ast).eval(inputMap, bindings);
           Object expectedOutput = cel.createProgram(cel.compile(testCase.getOutput()).getAst()).eval();
           expect.withMessage(testName).that(evalResult).isEqualTo(expectedOutput);
         } catch (CelValidationException e) {
@@ -228,4 +302,4 @@ private void runTests(Cel cel, CelAbstractSyntaxTree ast, PolicyTestSuite testSu
       }
     }
   }
-}
+}
diff --git a/tools/src/test/java/dev/cel/tools/ai/BUILD.bazel b/tools/src/test/java/dev/cel/tools/ai/BUILD.bazel
@@ -22,10 +22,12 @@ java_library(
         "//policy/testing:policy_test_suite_helper",
         "//runtime:evaluation_exception",
         "//runtime:function_binding",
+        "//runtime:late_function_binding",
         "//tools/ai:agentic_policy_compiler",
         "//tools/src/main/java/dev/cel/tools/ai:agent_context_java_proto",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java",
+        "@maven//:com_google_protobuf_protobuf_java_util",
         "@maven//:com_google_testparameterinjector_test_parameter_injector",
         "@maven//:junit_junit",
     ],
diff --git a/tools/src/test/resources/time_bound_approval.celpolicy b/tools/src/test/resources/time_bound_approval.celpolicy
@@ -0,0 +1,23 @@
+name: "policy.safety.time_bound_approval"
+default: allow
+
+variables:
+  # Define the validity window (30 seconds ago)
+  - approval_cutoff: now - duration('30s')
+
+  # Find approval messages in the valid window
+  - valid_approvals: >
+      agent.history()
+           .after(variables.approval_cutoff)
+           .role('model')
+           .filter(m, has(m.metadata.step) && m.metadata.step == 'approval_granted')
+
+  - has_valid_approval: variables.valid_approvals.size() > 0
+
+rules:
+  - description: "Require approval within the last 30 seconds for sensitive writes"
+    condition: >
+      tool.name == 'database_write' &&
+      !variables.has_valid_approval
+    effect: deny
+    message: "Authorization expired. Please re-approve the database write operation."
diff --git a/tools/src/test/resources/time_bound_approval_tests.yaml b/tools/src/test/resources/time_bound_approval_tests.yaml
@@ -0,0 +1,46 @@
+description: "Time-Bound Approval Policy Tests"
+
+section:
+- name: "Time Window Enforcement"
+  tests:
+  - name: "Approval Expired (Deny)"
+    input:
+      tool.name:
+        value: "database_write"
+      now:
+        expr: timestamp("2024-01-01T12:01:00Z")
+      _test_history:
+        expr: >
+          [
+            AgentMessage{ 
+              role: "model", 
+              time: timestamp("2024-01-01T12:00:00Z"),
+              metadata: { "step": "approval_granted" }
+            }
+          ]
+    output: >
+      {
+        "effect": "deny",
+        "message": "Authorization expired. Please re-approve the database write operation."
+      }
+
+  - name: "Approval Valid (Allow)"
+    input:
+      tool.name:
+        value: "database_write"
+      now:
+        expr: timestamp("2024-01-01T12:00:10Z")
+      _test_history:
+        expr: >
+          [
+            AgentMessage{ 
+              role: "model", 
+              time: timestamp("2024-01-01T12:00:00Z"),
+              metadata: { "step": "approval_granted" }
+            }
+          ]
+    output: >
+      {
+        "effect": "allow",
+        "message": ""
+      }
