feat(compiler): support registered host-await builtins for natural function call syntax (#667)

* feat(compiler): support registered host-await builtins

Allow hosts to register function names at compile time so that calls to
those names emit HostAwait instructions directly, enabling natural syntax
like fetch(x) instead of __builtin_host_await(x, "fetch").

- Add host_await_builtins map and register_host_await_builtin() to Compiler
- Validate arg_count == 1 and reject reserved __builtin_host_await name
- Extend determine_call_target() resolution: explicit > registered > user > builtin
- Both explicit and registered paths emit identical HostAwait bytecode
- Add compile_from_policy_with_host_await() entry point in rules.rs
- Extended test harness with HostAwaitBuiltinSpec and args assertion
- 9 YAML test cases: suspend/resume, run-to-completion, multiple names,
  queue, shadowing, object packing, arg_count rejection, reserved name
  rejection, standard builtin override
- Documentation: instruction-set.md, architecture.md

* Update src/languages/rego/compiler/function_calls.rs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Mark Birger <birgerm@yandex.ru>

* fix(compiler): address PR #667 review feedback on host-await registration

- Compiler::register_host_await_builtin now rejects duplicate, empty,
  and whitespace-only names. Previously a duplicate registration would
  silently overwrite the existing entry, which could mask the host's
  own registration mistakes.
- YAML test cases added: empty registration list as no-op, duplicate
  name rejection, empty/whitespace name rejection, out-param (a, out)
  calling syntax with a single-arg registered builtin, and mixed
  __builtin_host_await + registered builtins in the same policy
  consuming from their respective identifier queues.
- Test harness: replace assert_eq! on HostAwait argument mismatch with
  anyhow::Error so mismatches propagate through the case reporter
  instead of panicking and skipping the harness's normal error path.
- YAML comment fix: "Registration panics" -> "Registration fails with
  an error" (registration returns Err, never panics).

Addresses anakrish + Copilot inline review comments on PR #667.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* compiler: split CallTarget::HostAwait into explicit and registered variants

Addresses PR #667 review item #8: at the emit site in
`compile_function_call`, the discrimination between explicit
`__builtin_host_await(arg, id)` and a registered host-awaitable
builtin was being recovered by string-comparing `original_fcn_path`
against `"__builtin_host_await"`. The information was already known
in `determine_call_target` and was being thrown away.

Replace the single `CallTarget::HostAwait` variant with two:

* `ExplicitHostAwait` (unit) — the two-argument call form. The
  identifier register comes from the user's second argument.
* `RegisteredHostAwait { identifier: String }` — the one-argument
  call form for registered builtins. The identifier is the registered
  name and is captured in the variant at recognition time, so the
  emit site never re-derives it from the function path.

This removes the magic-string comparison at the emit site (the source
of truth is now `determine_call_target`) and makes both match sites
in `compile_function_call` exhaustive over the two forms — adding a
third host-await form in the future would force a compile error at
every match site instead of silently falling through.

Arities are now hardcoded in the `expected_args` extraction
(`Some(2)` for explicit, `Some(1)` for registered) rather than
carried in the variant; registered builtins are constrained to
`arg_count == 1` at registration time, so there is no per-call
variability to carry.

Bytecode output is unchanged; the full RVM test suite (97 cases) and
the registered_host_await suite (15 cases) pass without modification.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* docs(compiler): clarify registered host-await intercepts unqualified calls only

PR #667 review (Medium): the docs implied registered host-await names
shadow user functions and builtins unconditionally, but
determine_call_target matches only the bare original_fcn_path. A
package-qualified call such as data.demo.resolve(x) is therefore not
intercepted -- it resolves through the normal path like any other call.

Rather than expand registration to qualified paths (which would let a
registered name leak into every package exposing a same-named rule),
document the unqualified-only behavior and pin it with tests.

- register_host_await_builtin: doc now states only the unqualified call
  form is intercepted; qualified calls resolve normally.
- determine_call_target: inline comment explaining the deliberate
  original_fcn_path-only match.
- docs/rvm/instruction-set.md: describe qualified-call resolution,
  including that builtins have no qualified form.
- tests: cross-package and same-package qualified calls resolve to the
  rule; bare-name shadowing of a standard builtin; Unknown-function
  outcome when no rule exists at the qualified path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(tests): compare host-await argument without re-running process_value

PR #667 review (Low): the suspendable test harness compared the
host-await argument via process_value(argument), but argument is already
a runtime Value. process_value is a YAML-fixture decoder -- it rewrites
"#undefined" to Undefined, {set!: [...]} to a set, and errors on a
runtime Value::Set. Re-running it on the runtime argument could coerce a
legitimate payload into a fixture sentinel (passing for the wrong
reason) or error outright on sets.

Compare the runtime argument directly against the expected value, which
is already decoded once at YAML load time.

Add a regression case (registered_builtin_suspendable_set_argument) that
passes a set payload: it fails under the old double-processing
("unexpected set in value read from json/yaml") and passes with the fix.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(tests): reject `args:` payload expectations in run-to-completion mode

PR #667 review (Low): a run-to-completion host-await response could carry
an `args:` payload expectation, but RTC execution pre-loads responses and
never surfaces the call argument to the harness, so the expectation was
parsed and silently dropped. A case with `args: "WRONG"` passed as long as
the result matched -- asserting a payload that was never checked.

Reject `args:` for run-to-completion fixtures at load time, directing the
author to suspendable mode where arguments are validated. Also only build
the run-to-completion response vector when the case actually runs in RTC
mode, so a suspendable case using the shared host_await_responses field
with `args:` is not wrongly rejected.

Route the fixture-load error through the same want_error handling used for
compilation errors, and add registered_builtin_run_to_completion_rejects_args
which now fails loudly instead of passing silently.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(compiler): reject host-await builtin names with surrounding whitespace

PR #667 review (Low): register_host_await_builtin rejected all-whitespace
names via name.trim().is_empty(), but accepted padded names like " lookup"
or "lookup ". Those were inserted into host_await_builtins, but Rego
function-call paths produce the trimmed identifier, so a padded
registration could never match -- a silent dead registration.

Reject any name that is not already trimmed (name != name.trim()) in
addition to empty names, and update the error message accordingly. Add
test cases for leading and trailing whitespace.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Signed-off-by: Mark Birger <birgerm@yandex.ru>
Co-authored-by: Mark Birger <markbirger@microsoft.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

docs/rvm/architecture.md

@@ -254,7 +254,13 @@ include formatted state snapshots where possible.
 7. **Host await**: In run-to-completion mode, `HostAwait` consumes a response
    from `host_await_responses`. Suspendable mode yields control with a
    `SuspendReason::HostAwait { dest, argument, identifier }` that the host must
-   service.
+   service. The compiler supports two ways to emit `HostAwait`:
+   - **Explicit**: `__builtin_host_await(payload, identifier)` — raw 2-argument
+     form.
+   - **Registered**: `compile_from_policy_with_host_await` accepts a list of
+     `(name, arg_count)` pairs. Calls to registered names are compiled as
+     `HostAwait` with the function name as the identifier literal. Registered
+     names take precedence over user-defined functions and standard builtins.
 8. **Completion**: `Return` wraps the selected register value into
    `InstructionOutcome::Return`, unwinding frames until the entry frame is
    cleared. `RuleReturn` is a specialised variant used by rule execution

docs/rvm/instruction-set.md

@@ -177,6 +177,75 @@ Parameter tables:
 - Suspendable: emits `InstructionOutcome::Suspend` with `SuspendReason::HostAwait`.
   The host must resume with a value that will be written into `dest`.
 
+### Registered host-await builtins
+
+The compiler can be configured with a list of function names that map directly
+to `HostAwait` instructions. This allows policy authors to write natural
+function calls (e.g. `lookup(input.account_id)`) instead of the raw
+`__builtin_host_await(payload, identifier)` builtin.
+
+Registration is done at compile time via `Compiler::compile_from_policy_with_host_await`:
+
+```rust
+let builtins = [("lookup", 1), ("persist", 1)];
+let program = Compiler::compile_from_policy_with_host_await(
+    &compiled_policy, &entry_points, &builtins,
+)?;
+```
+
+Each registered name is a `(name, arg_count)` pair. When the compiler
+encounters a call to a registered name, it emits a `HostAwait` instruction
+with:
+- `arg` = the first argument register
+- `id` = a register loaded with a string literal containing the function name
+
+Both the explicit `__builtin_host_await(arg, id)` call and a registered
+builtin call produce the **same `HostAwait` bytecode instruction**. The only
+difference is how the `id` register is populated: explicit calls take it from
+the second user-supplied argument, while registered calls auto-generate a
+`Load` instruction for the function name string. The VM cannot distinguish
+between the two at runtime.
+
+**Resolution order** in `determine_call_target()`:
+1. `__builtin_host_await` (magic 2-argument form)
+2. Registered host-await builtins (matched by **bare** function name only)
+3. User-defined functions (matched by package-qualified path)
+4. Standard builtins (matched by bare function name)
+
+Registered names shadow both user-defined functions and standard builtins.
+This means `time.parse_duration_ns` can be overridden to route through the
+host instead of the built-in Rust implementation.
+
+**Only unqualified calls are intercepted.** Registration matches a call by
+the name *as written in the policy*. A bare call — `lookup(x)` — is
+intercepted and compiled to a `HostAwait`. A package-qualified call —
+`data.pkg.lookup(x)` — is **not** intercepted; it is resolved normally, as
+if the name were never registered.
+
+```rego
+# "lookup" is registered as a host-await builtin.
+
+package other
+import rego.v1
+lookup(k) := k            # an ordinary rule that happens to share the name
+
+package demo
+import rego.v1
+a := lookup(input.k)             # intercepted     -> HostAwait
+b := data.other.lookup(input.k)  # NOT intercepted -> calls other.lookup
+```
+
+The qualified form is resolved exactly as it would be without registration:
+if a rule exists at that path it is called, otherwise compilation fails with
+`Unknown function`. (A standard builtin like `count` has no qualified form at
+all, so `data.pkg.count(x)` is always an `Unknown function` error, registered
+or not.)
+
+**Argument handling**: The `HostAwait` instruction carries a single `arg`
+register. Registered builtins must use `arg_count: 1`; the compiler rejects
+`arg_count > 1` at registration time. To pass multiple values, use object
+packing: `lookup({"user": x, "resource": y})`.
+
 ---
 
 ## Halt instruction

src/languages/rego/compiler/function_calls.rs

@@ -17,8 +17,20 @@ use crate::lexer::Span;
 use crate::rvm::instructions::{BuiltinCallParams, FunctionCallParams};
 use crate::rvm::Instruction;
 use crate::utils::get_path_string;
-use alloc::{format, string::ToString, vec::Vec};
+use crate::value::Value;
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 
+/// Resolved destination of a Rego function-call expression. Produced by
+/// [`Compiler::determine_call_target`] and consumed by
+/// [`Compiler::compile_function_call`] to choose which instruction to emit.
+/// Carrying the discrimination in the type (rather than re-matching on a
+/// magic name at the emit site) keeps the host-await handling honest under
+/// future refactors — the compiler will refuse to build if a new variant is
+/// added without updating every match site.
 enum CallTarget {
     User {
         rule_index: u16,
@@ -28,9 +40,14 @@ enum CallTarget {
         builtin_index: u16,
         expected_args: Option<usize>,
     },
-    HostAwait {
-        expected_args: Option<usize>,
-    },
+    /// Explicit `__builtin_host_await(arg, id)` call form (2 user args).
+    /// The identifier is supplied by the policy author at runtime via the
+    /// second argument register.
+    ExplicitHostAwait,
+    /// A registered host-awaitable builtin invoked by its registered name
+    /// (1 user arg). The identifier is the registered name itself and is
+    /// baked into the bytecode as a string literal at compile time.
+    RegisteredHostAwait { identifier: String },
 }
 
 impl<'a> Compiler<'a> {
@@ -59,7 +76,11 @@ impl<'a> Compiler<'a> {
         let expected_args = match &call_target {
             CallTarget::User { expected_args, .. } => *expected_args,
             CallTarget::Builtin { expected_args, .. } => *expected_args,
-            CallTarget::HostAwait { expected_args } => *expected_args,
+            // Both host-await variants have a known fixed arity; carrying it
+            // in the variant lets the rest of the compiler depend on the type
+            // rather than re-matching on the magic name `__builtin_host_await`.
+            CallTarget::ExplicitHostAwait => Some(2),
+            CallTarget::RegisteredHostAwait { .. } => Some(1),
         };
 
         if let Some(expected) = expected_args {
@@ -126,7 +147,8 @@ impl<'a> Compiler<'a> {
                 });
                 self.emit_instruction(Instruction::BuiltinCall { params_index }, &span);
             }
-            CallTarget::HostAwait { .. } => {
+            CallTarget::ExplicitHostAwait => {
+                // Explicit __builtin_host_await(arg, id) — 2 arguments
                 if arg_regs.len() != 2 {
                     return Err(CompilerError::General {
                         message: format!(
@@ -136,7 +158,6 @@ impl<'a> Compiler<'a> {
                     }
                     .at(&span));
                 }
-
                 self.emit_instruction(
                     Instruction::HostAwait {
                         dest,
@@ -146,6 +167,37 @@ impl<'a> Compiler<'a> {
                     &span,
                 );
             }
+            CallTarget::RegisteredHostAwait { identifier } => {
+                // Registered host-awaitable builtin — the identifier is the
+                // registered name and is baked into the bytecode as a literal.
+                if arg_regs.len() != 1 {
+                    return Err(CompilerError::General {
+                        message: format!(
+                            "host-awaitable builtin '{}' expects exactly 1 argument, got {}",
+                            identifier,
+                            arg_regs.len()
+                        ),
+                    }
+                    .at(&span));
+                }
+                let id_reg = self.alloc_register();
+                let literal_idx = self.add_literal(Value::String(identifier.into()));
+                self.emit_instruction(
+                    Instruction::Load {
+                        dest: id_reg,
+                        literal_idx,
+                    },
+                    &span,
+                );
+                self.emit_instruction(
+                    Instruction::HostAwait {
+                        dest,
+                        arg: arg_regs[0],
+                        id: id_reg,
+                    },
+                    &span,
+                );
+            }
         }
 
         if let Some((plan, plan_span)) = &out_param_plan {
@@ -187,8 +239,26 @@ impl<'a> Compiler<'a> {
         span: &Span,
     ) -> Result<CallTarget> {
         if original_fcn_path == "__builtin_host_await" {
-            return Ok(CallTarget::HostAwait {
-                expected_args: Some(2),
+            return Ok(CallTarget::ExplicitHostAwait);
+        }
+
+        // Check registered host-awaitable builtins. Registered builtins are
+        // restricted to arg_count == 1 at registration time (see
+        // `Compiler::register_host_await_builtin`), so the variant doesn't
+        // need to carry an arity — it's fixed at 1.
+        //
+        // We deliberately match against `original_fcn_path` only, not
+        // `full_fcn_path`. Registration intercepts the *unqualified* call
+        // form (e.g. `lookup(x)` inside the policy's own package). A
+        // package-qualified call like `data.other.lookup(x)` is left to
+        // resolve through the normal user-defined / builtin path, so a
+        // registered name does not leak into unrelated packages that
+        // happen to expose a rule with the same identifier. This is
+        // documented on `register_host_await_builtin`; the
+        // `registered_host_await.yaml` suite pins the behavior.
+        if self.host_await_builtins.contains_key(original_fcn_path) {
+            return Ok(CallTarget::RegisteredHostAwait {
+                identifier: original_fcn_path.to_string(),
             });
         }
 

src/languages/rego/compiler/mod.rs

@@ -26,7 +26,9 @@ use crate::rvm::program::{Program, RuleType, SpanInfo};
 use crate::CompiledPolicy;
 use crate::Value;
 use alloc::collections::{BTreeMap, BTreeSet};
+use alloc::format;
 use alloc::string::String;
+use alloc::string::ToString as _;
 use alloc::vec;
 use alloc::vec::Vec;
 use indexmap::IndexMap;
@@ -139,6 +141,10 @@ pub struct Compiler<'a> {
     current_call_stack: Vec<u16>,
     entry_points: IndexMap<String, usize>,
     soft_assert_mode: bool,
+    /// Registered host-awaitable builtins: name → expected arg count.
+    /// When the compiler encounters a call to one of these names, it emits a
+    /// `HostAwait` instruction instead of a regular function or builtin call.
+    host_await_builtins: BTreeMap<String, usize>,
 }
 
 impl<'a> Compiler<'a> {
@@ -173,9 +179,75 @@ impl<'a> Compiler<'a> {
             current_call_stack: Vec::new(),
             entry_points: IndexMap::new(),
             soft_assert_mode: false,
+            host_await_builtins: BTreeMap::new(),
         }
     }
 
+    /// Register a function name as a host-awaitable builtin.
+    ///
+    /// When the compiler encounters an **unqualified** call to `name(arg)`
+    /// (i.e. `name(arg)` from inside the policy's own package, not
+    /// `data.pkg.name(arg)` or any other package-qualified form), it will
+    /// emit a `HostAwait` instruction with the argument and `name` as the
+    /// identifier, instead of treating it as a user-defined or standard
+    /// builtin function.
+    ///
+    /// Package-qualified calls (e.g. `data.other.name(arg)`) are **not**
+    /// intercepted by registration. Those resolve through the normal
+    /// user-defined / builtin lookup against their fully-qualified path
+    /// (`data.other.name`).
+    ///
+    /// `arg_count` must be exactly 1. The `HostAwait` instruction carries a
+    /// single argument register; use object packing to pass multiple values
+    /// (e.g. `name({"key1": v1, "key2": v2})`).
+    ///
+    /// Returns `Err` when:
+    /// - `name` is the reserved identifier `__builtin_host_await`,
+    /// - `name` is empty, only whitespace, or has leading/trailing
+    ///   whitespace (whitespace-padded names would never match the
+    ///   trimmed identifier produced by the Rego parser, creating dead
+    ///   registrations),
+    /// - `name` is already registered (duplicate registration is rejected
+    ///   rather than silently overwritten),
+    /// - `arg_count` is not exactly 1.
+    pub fn register_host_await_builtin(&mut self, name: &str, arg_count: usize) -> Result<()> {
+        if name == "__builtin_host_await" {
+            return Err(CompilerError::General {
+                message: "__builtin_host_await is a reserved name and cannot be registered as a host-await builtin"
+                    .to_string(),
+            }
+            .into());
+        }
+        if name.is_empty() || name != name.trim() {
+            return Err(CompilerError::General {
+                message: format!(
+                    "host-await builtin name {name:?} must not be empty or contain leading/trailing whitespace"
+                ),
+            }
+            .into());
+        }
+        if self.host_await_builtins.contains_key(name) {
+            return Err(CompilerError::General {
+                message: format!(
+                    "host-await builtin '{name}' is already registered; \
+                     duplicate registration is not allowed"
+                ),
+            }
+            .into());
+        }
+        if arg_count != 1 {
+            return Err(CompilerError::General {
+                message: format!(
+                    "registered host-await builtin '{name}' must have arg_count == 1, got {arg_count}. \
+                     Use object packing to pass multiple values."
+                ),
+            }
+            .into());
+        }
+        self.host_await_builtins.insert(name.to_string(), arg_count);
+        Ok(())
+    }
+
     pub(super) fn with_soft_assert_mode<F, R>(&mut self, enabled: bool, f: F) -> R
     where
         F: FnOnce(&mut Self) -> R,

src/languages/rego/compiler/rules.rs

@@ -234,8 +234,20 @@ impl<'a> Compiler<'a> {
     pub fn compile_from_policy(
         policy: &CompiledPolicy,
         entry_points: &[&str],
+    ) -> Result<Arc<Program>> {
+        Self::compile_from_policy_with_host_await(policy, entry_points, &[])
+    }
+
+    /// Compile from a CompiledPolicy to RVM Program with registered host-awaitable builtins.
+    pub fn compile_from_policy_with_host_await(
+        policy: &CompiledPolicy,
+        entry_points: &[&str],
+        host_await_builtins: &[(&str, usize)],
     ) -> Result<Arc<Program>> {
         let mut compiler = Compiler::with_policy(policy);
+        for &(name, arg_count) in host_await_builtins {
+            compiler.register_host_await_builtin(name, arg_count)?;
+        }
         compiler.current_rule_path = "".to_string();
         let rules = policy.get_rules();
 

tests/rvm/rego/cases/registered_host_await.yaml

@@ -0,0 +1,610 @@
+cases:
+  - note: registered_builtin_suspendable
+    data: {}
+    input:
+      account_id: "acct-42"
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: get_account
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        allow if {
+          account := get_account({"id": input.account_id})
+          account.status == "active"
+        }
+    query: data.demo.allow
+    host_await_responses_suspendable:
+      - id: "get_account"
+        args:
+          id: "acct-42"
+        value:
+          status: "active"
+          name: "Alice"
+    want_result: true
+
+  - note: registered_builtin_run_to_completion
+    data: {}
+    input:
+      lang: "es"
+    skip_interpreter: true
+    execution_mode: run-to-completion
+    host_await_builtins:
+      - name: translate
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        greeting := msg if {
+          msg := translate(input.lang)
+        }
+    query: data.demo.greeting
+    host_await_responses:
+      - id: "translate"
+        value: "hola"
+    want_result: "hola"
+
+  - note: registered_builtin_run_to_completion_rejects_args
+    data: {}
+    input:
+      lang: "es"
+    skip_interpreter: true
+    execution_mode: run-to-completion
+    # `args:` payload validation is only meaningful in suspendable mode, where
+    # the harness sees each call's argument. In run-to-completion mode the VM
+    # consumes pre-loaded responses internally, so an `args:` expectation can
+    # never be checked. Rather than silently ignore it (which would let a case
+    # "assert" a payload that is never verified), the harness rejects it.
+    host_await_builtins:
+      - name: translate
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        greeting := msg if {
+          msg := translate(input.lang)
+        }
+    query: data.demo.greeting
+    host_await_responses:
+      - id: "translate"
+        args: "es"
+        value: "hola"
+    want_error: "not supported in run-to-completion mode"
+
+  - note: registered_builtin_multiple_names
+    data: {}
+    input:
+      user_id: "user-7"
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: lookup
+        arg_count: 1
+      - name: persist
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        result := {"data": fetched, "stored": saved} if {
+          fetched := lookup(input.user_id)
+          saved := persist({"id": input.user_id, "action": "audit"})
+        }
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "lookup"
+        args: "user-7"
+        value:
+          name: "Charlie"
+      - id: "persist"
+        args:
+          id: "user-7"
+          action: "audit"
+        value: true
+    want_result:
+      data:
+        name: "Charlie"
+      stored: true
+
+  - note: registered_builtin_suspendable_queue
+    data: {}
+    input:
+      items: ["alpha", "beta", "gamma"]
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: enrich
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        results := [r |
+          item := input.items[_]
+          r := enrich(item)
+        ]
+    query: data.demo.results
+    host_await_responses_suspendable:
+      - id: "enrich"
+        args: "alpha"
+        value: "enriched-alpha"
+      - id: "enrich"
+        args: "beta"
+        value: "enriched-beta"
+      - id: "enrich"
+        args: "gamma"
+        value: "enriched-gamma"
+    want_result: ["enriched-alpha", "enriched-beta", "enriched-gamma"]
+
+  - note: registered_builtin_shadows_user_function
+    data: {}
+    input:
+      key: "test-key"
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: resolve
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        # This user-defined function should be shadowed by the registered builtin
+        resolve(x) := {"local": true, "key": x}
+
+        result := resolve(input.key)
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "resolve"
+        args: "test-key"
+        value: "from-host"
+    # The registered builtin takes precedence — result is the host response, not the user function
+    want_result: "from-host"
+
+  - note: registered_builtin_multi_arg_object_packing
+    data: {}
+    input:
+      user: "alice"
+      resource: "/api/data"
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: check_access
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        # Multi-value calls pack arguments into a single object
+        allowed if {
+          result := check_access({"user": input.user, "resource": input.resource})
+          result.granted == true
+        }
+    query: data.demo.allowed
+    host_await_responses_suspendable:
+      - id: "check_access"
+        args:
+          user: "alice"
+          resource: "/api/data"
+        value:
+          granted: true
+          reason: "admin"
+    want_result: true
+
+  - note: registered_builtin_rejects_arg_count_greater_than_one
+    data: {}
+    input:
+      key: "k1"
+      value: "v1"
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: kv_store
+        arg_count: 2
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        result := kv_store(input.key, input.value)
+    query: data.demo.result
+    # Registration fails with an error because arg_count must be 1.
+    # Use object packing instead: kv_store({"key": input.key, "value": input.value})
+    want_error: "arg_count == 1"
+
+  - note: registered_builtin_overrides_standard_builtin
+    data: {}
+    input:
+      duration: "2h30m"
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: time.parse_duration_ns
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+
+        # time.parse_duration_ns is a standard Rego builtin (1 arg, returns nanoseconds).
+        # Registering it as a host-await builtin shadows the standard implementation.
+        duration_ns := time.parse_duration_ns(input.duration)
+    query: data.demo.duration_ns
+    host_await_responses_suspendable:
+      - id: "time.parse_duration_ns"
+        args: "2h30m"
+        value: 9000000000000
+    # Host returns 9000000000000 (custom value) instead of the real parse result.
+    # This proves the registered builtin shadows the standard one.
+    want_result: 9000000000000
+
+  - note: registered_builtin_rejects_reserved_name
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: __builtin_host_await
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := true
+    query: data.demo.result
+    # __builtin_host_await is a reserved name handled by the explicit code path;
+    # registering it as a host-await builtin is rejected at compile time.
+    want_error: "__builtin_host_await is a reserved name"
+
+  - note: registered_builtin_empty_list_is_noop
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Empty registration list: nothing is registered, the policy compiles
+    # normally and no HostAwait machinery fires.
+    host_await_builtins: []
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := 42
+    query: data.demo.result
+    want_result: 42
+
+  - note: registered_builtin_rejects_duplicate_name
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Registering the same name twice is rejected (rather than silently
+    # overwritten) so the host can't accidentally clobber its own registration.
+    host_await_builtins:
+      - name: lookup
+        arg_count: 1
+      - name: lookup
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := lookup("anything")
+    query: data.demo.result
+    want_error: "already registered"
+
+  - note: registered_builtin_rejects_empty_name
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # The empty identifier is meaningless; reject at registration time.
+    host_await_builtins:
+      - name: ""
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := true
+    query: data.demo.result
+    want_error: "must not be empty"
+
+  - note: registered_builtin_rejects_whitespace_name
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Whitespace-only names are equivalent to empty for registration purposes.
+    host_await_builtins:
+      - name: "   "
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := true
+    query: data.demo.result
+    want_error: "leading/trailing whitespace"
+
+  - note: registered_builtin_rejects_leading_whitespace_name
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Whitespace-padded names would never match the trimmed identifier produced
+    # by the Rego parser, creating an unreachable registration. Reject at
+    # registration time so misconfiguration is loud, not silent.
+    host_await_builtins:
+      - name: " lookup"
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := true
+    query: data.demo.result
+    want_error: "leading/trailing whitespace"
+
+  - note: registered_builtin_rejects_trailing_whitespace_name
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    host_await_builtins:
+      - name: "lookup "
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := true
+    query: data.demo.result
+    want_error: "leading/trailing whitespace"
+
+  - note: registered_builtin_out_param_syntax
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Rego's "output-param" calling form: `f(in, out)` desugars to `f(in)`
+    # with the return value unified with `out`. With a registered arg_count=1
+    # builtin, only the first positional argument (input) reaches the host;
+    # the second positional is the output binding, not a second host-await
+    # argument. The host returns the response, which is then unified with the
+    # output binding (here, `out`).
+    host_await_builtins:
+      - name: lookup
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := out if {
+          lookup("ping", out)
+        }
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "lookup"
+        args: "ping"
+        value: "pong"
+    want_result: "pong"
+
+  - note: registered_builtin_mixed_with_explicit_builtin_host_await
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Both invocation forms in the same policy. The explicit
+    # __builtin_host_await call uses the user-supplied identifier ("kv_get"),
+    # while the registered "lookup" name resolves to a HostAwait with the
+    # registered identifier. Both are emitted as HostAwait instructions and
+    # consume from their respective identifier queues.
+    host_await_builtins:
+      - name: lookup
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        registered_value := lookup("alpha")
+        explicit_value   := __builtin_host_await("beta", "kv_get")
+        result := {
+          "registered": registered_value,
+          "explicit":   explicit_value,
+        }
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "lookup"
+        args: "alpha"
+        value: "from_registered"
+      - id: "kv_get"
+        args: "beta"
+        value: "from_explicit"
+    want_result:
+      registered: "from_registered"
+      explicit:   "from_explicit"
+
+  - note: registered_builtin_does_not_intercept_qualified_calls
+    data: {}
+    input:
+      key: "alpha"
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Pins the documented FQN behavior: a registered name (`resolve`) only
+    # intercepts the unqualified call form (`resolve(x)`) inside the
+    # registering package. A package-qualified call (`data.other.resolve(x)`)
+    # resolves through the normal user-defined function path. The other
+    # package's `resolve` is invoked and returns its own value — the host
+    # never sees the call. The test asserts both: the registered intercept
+    # is consumed once (from the unqualified call), and the qualified call
+    # returns the user-rule output without registering as a host-await.
+    host_await_builtins:
+      - name: resolve
+        arg_count: 1
+    modules:
+      - |
+        package other
+        import rego.v1
+        resolve(k) := result if {
+          result := sprintf("user-rule-handled:%s", [k])
+        }
+      - |
+        package demo
+        import rego.v1
+        intercepted := resolve(input.key)
+        bypassed    := data.other.resolve(input.key)
+        result := {
+          "intercepted": intercepted,
+          "bypassed":    bypassed,
+        }
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "resolve"
+        args: "alpha"
+        value: "from_host"
+    want_result:
+      intercepted: "from_host"
+      bypassed: "user-rule-handled:alpha"
+
+  - note: registered_builtin_same_package_qualified_reaches_local_rule
+    data: {}
+    input:
+      key: "alpha"
+    skip_interpreter: true
+    execution_mode: suspendable
+    # When a registered name *also* exists as a rule in the registering
+    # package, the bare call is intercepted (HostAwait) while the
+    # same-package qualified call reaches the local rule. Registration is a
+    # bare-name intercept only; the qualified path resolves as it normally
+    # would. This gives a policy a deliberate escape hatch: register `resolve`
+    # for host interception, yet still reach the local rule via
+    # `data.demo.resolve` when the host should be bypassed.
+    host_await_builtins:
+      - name: resolve
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        resolve(k) := sprintf("local-rule:%s", [k])
+        intercepted := resolve(input.key)
+        bypassed    := data.demo.resolve(input.key)
+        result := {
+          "intercepted": intercepted,
+          "bypassed":    bypassed,
+        }
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "resolve"
+        args: "alpha"
+        value: "from_host"
+    want_result:
+      intercepted: "from_host"
+      bypassed: "local-rule:alpha"
+
+  - note: registered_builtin_shadows_standard_builtin_by_bare_name
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Registering a name that collides with a standard builtin (`count`)
+    # intercepts the bare call form: `count([...])` compiles to a HostAwait
+    # instead of invoking the built-in implementation. The qualified form
+    # `data.demo.count([...])` has no meaning for a builtin (builtins are
+    # callable only by bare name) and would fail to compile with
+    # `Unknown function` whether or not `count` is registered, so it is not
+    # exercised here.
+    host_await_builtins:
+      - name: count
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := count([10, 20, 30])
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "count"
+        args: [10, 20, 30]
+        value: "from_host"
+    want_result: "from_host"
+
+  - note: registered_builtin_qualified_unknown_function_when_no_rule
+    data: {}
+    input:
+      key: "alpha"
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Registration is a bare-name intercept and creates no rule. A qualified
+    # call to the registered name therefore has nothing to resolve to (no rule
+    # exists at `data.demo.resolve`) and fails to compile with
+    # `Unknown function` — the same outcome as without registration. Pins the
+    # documented "otherwise compilation fails with Unknown function" case.
+    host_await_builtins:
+      - name: resolve
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := data.demo.resolve(input.key)
+    query: data.demo.result
+    want_error: "Unknown function"
+
+  - note: registered_builtin_qualified_builtin_name_is_unknown_function
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # A standard builtin has no package-qualified form: `data.demo.count([...])`
+    # fails with `Unknown function` whether or not `count` is registered as a
+    # host-await builtin. Pins the parenthetical in the docs.
+    host_await_builtins:
+      - name: count
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := data.demo.count([10, 20, 30])
+    query: data.demo.result
+    want_error: "Unknown function"
+
+  - note: registered_builtin_suspendable_set_argument
+    data: {}
+    input: {}
+    skip_interpreter: true
+    execution_mode: suspendable
+    # Regression pin for the host-await argument comparison: the policy passes
+    # a *set* as the payload. The runtime argument is already a `Value::Set`,
+    # and the expected `args` decodes (via the `set!` helper) to the same set.
+    # The harness must compare them directly -- re-running `process_value` on the
+    # runtime argument would bail with "unexpected set in value read from
+    # json/yaml" and fail the case for the wrong reason.
+    host_await_builtins:
+      - name: lookup
+        arg_count: 1
+    modules:
+      - |
+        package demo
+        import rego.v1
+        result := lookup({1, 2, 3})
+    query: data.demo.result
+    host_await_responses_suspendable:
+      - id: "lookup"
+        args:
+          set!: [1, 2, 3]
+        value: "found"
+    want_result: "found"

tests/rvm/rego/mod.rs

@@ -41,6 +41,7 @@ struct TestCase {
     pub host_await_responses: Option<Vec<HostAwaitResponseSpec>>,
     pub host_await_responses_run_to_completion: Option<Vec<HostAwaitResponseSpec>>,
     pub host_await_responses_suspendable: Option<Vec<HostAwaitResponseSpec>>,
+    pub host_await_builtins: Option<Vec<HostAwaitBuiltinSpec>>,
 }
 
 fn default_strict() -> bool {
@@ -55,14 +56,24 @@ struct YamlTest {
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
 struct HostAwaitResponseSpec {
     pub id: Value,
+    pub args: Option<Value>,
     pub value: Value,
 }
 
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
+struct HostAwaitBuiltinSpec {
+    pub name: String,
+    pub arg_count: usize,
+}
+
+type HostAwaitResponseMap = BTreeMap<Value, VecDeque<(Option<Value>, Value)>>;
+
 #[derive(Debug, Clone)]
 struct RvmExecutionOptions {
     execution_mode: ExecutionMode,
     host_await_responses_run_to_completion: Option<Vec<(Value, Vec<Value>)>>,
-    host_await_responses_suspendable: Option<BTreeMap<Value, VecDeque<Value>>>,
+    host_await_responses_suspendable: Option<HostAwaitResponseMap>,
+    host_await_builtins: Option<Vec<(String, usize)>>,
 }
 
 impl Default for RvmExecutionOptions {
@@ -71,6 +82,7 @@ impl Default for RvmExecutionOptions {
             execution_mode: ExecutionMode::RunToCompletion,
             host_await_responses_run_to_completion: None,
             host_await_responses_suspendable: None,
+            host_await_builtins: None,
         }
     }
 }
@@ -82,12 +94,13 @@ fn render_program_listing(program: &Program) -> String {
 
 fn build_host_await_response_map(
     responses: &[HostAwaitResponseSpec],
-) -> anyhow::Result<BTreeMap<Value, VecDeque<Value>>> {
-    let mut map: BTreeMap<Value, VecDeque<Value>> = BTreeMap::new();
+) -> anyhow::Result<HostAwaitResponseMap> {
+    let mut map: HostAwaitResponseMap = BTreeMap::new();
     for response in responses {
         let id = process_value(&response.id)?;
+        let expected_args = response.args.as_ref().map(process_value).transpose()?;
         let value = process_value(&response.value)?;
-        map.entry(id).or_default().push_back(value);
+        map.entry(id).or_default().push_back((expected_args, value));
     }
     Ok(map)
 }
@@ -95,10 +108,24 @@ fn build_host_await_response_map(
 fn build_host_await_response_vec(
     responses: &[HostAwaitResponseSpec],
 ) -> anyhow::Result<Vec<(Value, Vec<Value>)>> {
+    // Run-to-completion responses are pre-loaded into the VM, which consumes
+    // them internally without surfacing each call's argument to the harness.
+    // There is therefore no point at which an `args:` expectation could be
+    // checked, so silently dropping it would let a case "assert" a payload
+    // that is never verified. Reject `args:` up front instead, pointing the
+    // author at suspendable mode where argument validation is supported.
+    if let Some(response) = responses.iter().find(|response| response.args.is_some()) {
+        return Err(anyhow::anyhow!(
+            "`args:` payload validation is not supported in run-to-completion mode \
+             (response for id {:?}); drop the `args:` field or move the case to \
+             execution_mode: suspendable",
+            response.id
+        ));
+    }
     let map = build_host_await_response_map(responses)?;
     Ok(map
         .into_iter()
-        .map(|(id, values)| (id, values.into_iter().collect()))
+        .map(|(id, values)| (id, values.into_iter().map(|(_, output)| output).collect()))
         .collect())
 }
 
@@ -111,12 +138,20 @@ fn build_execution_options(case: &TestCase) -> anyhow::Result<RvmExecutionOption
         }
     };
 
-    let rtc_responses = case
-        .host_await_responses_run_to_completion
-        .as_ref()
-        .or(case.host_await_responses.as_ref())
-        .map(|responses| build_host_await_response_vec(responses))
-        .transpose()?;
+    // Only build the run-to-completion response vec when the case actually
+    // runs in RTC mode. A suspendable case may legitimately use the shared
+    // `host_await_responses` field with `args:` expectations (validated via
+    // the suspendable map below); building the RTC vec for it would wrongly
+    // trip the RTC-only `args:` rejection in `build_host_await_response_vec`.
+    let rtc_responses = if execution_mode == ExecutionMode::RunToCompletion {
+        case.host_await_responses_run_to_completion
+            .as_ref()
+            .or(case.host_await_responses.as_ref())
+            .map(|responses| build_host_await_response_vec(responses))
+            .transpose()?
+    } else {
+        None
+    };
 
     let suspendable_responses = case
         .host_await_responses_suspendable
@@ -125,10 +160,18 @@ fn build_execution_options(case: &TestCase) -> anyhow::Result<RvmExecutionOption
         .map(|responses| build_host_await_response_map(responses))
         .transpose()?;
 
+    let ha_builtins = case.host_await_builtins.as_ref().map(|specs| {
+        specs
+            .iter()
+            .map(|s| (s.name.clone(), s.arg_count))
+            .collect()
+    });
+
     Ok(RvmExecutionOptions {
         execution_mode,
         host_await_responses_run_to_completion: rtc_responses,
         host_await_responses_suspendable: suspendable_responses,
+        host_await_builtins: ha_builtins,
     })
 }
 
@@ -227,7 +270,13 @@ fn compile_and_run_rvm_with_all_entry_points(
     listing_out: &mut Option<String>,
     execution_options: &RvmExecutionOptions,
 ) -> anyhow::Result<Vec<Value>> {
-    let program = Compiler::compile_from_policy(compiled_policy, entry_points)?;
+    let ha_builtins = execution_options
+        .host_await_builtins
+        .as_ref()
+        .map(|b| b.iter().map(|(n, a)| (n.as_str(), *a)).collect::<Vec<_>>())
+        .unwrap_or_default();
+    let program =
+        Compiler::compile_from_policy_with_host_await(compiled_policy, entry_points, &ha_builtins)?;
 
     // Basic serialization sanity check keeps regressions visible in CI.
     test_round_trip_serialization(program.as_ref()).map_err(|e| anyhow::anyhow!(e))?;
@@ -269,8 +318,12 @@ fn compile_and_run_rvm_with_all_entry_points(
                         return Err(anyhow::anyhow!("{}", error));
                     }
                     ExecutionState::Suspended { reason, .. } => match reason {
-                        SuspendReason::HostAwait { identifier, .. } => {
-                            let response = suspendable_responses
+                        SuspendReason::HostAwait {
+                            identifier,
+                            argument,
+                            ..
+                        } => {
+                            let (expected_args, response) = suspendable_responses
                                 .get_mut(identifier)
                                 .and_then(|queue| queue.pop_front())
                                 .ok_or_else(|| {
@@ -279,6 +332,24 @@ fn compile_and_run_rvm_with_all_entry_points(
                                         identifier
                                     )
                                 })?;
+                            if let Some(expected) = expected_args {
+                                // `argument` is already a runtime `Value`; the
+                                // expected side has been through `process_value`
+                                // once at YAML decode time (see
+                                // `build_host_await_response_map`). Comparing
+                                // raw runtime values keeps fixture sentinels like
+                                // "#undefined" from coercing a runtime string
+                                // payload into a different shape, which would
+                                // otherwise let tests pass for the wrong reason.
+                                if argument != &expected {
+                                    return Err(anyhow::anyhow!(
+                                        "HostAwait argument mismatch for {:?}: expected {:?}, got {:?}",
+                                        identifier,
+                                        expected,
+                                        argument
+                                    ));
+                                }
+                            }
                             vm.resume(Some(response))?;
                         }
                         other => {
@@ -365,7 +436,34 @@ fn yaml_test_impl(file: &str) -> Result<()> {
             Some(engine.eval_rule(case.query.clone()))
         };
 
-        let execution_options = build_execution_options(&case)?;
+        let execution_options = match build_execution_options(&case) {
+            Ok(options) => options,
+            Err(options_error) => {
+                // A malformed host-await fixture (e.g. an `args:` expectation on
+                // a run-to-completion response, which can never be validated) is
+                // reported here. Mirror the compilation-error handling below: if
+                // the case expects an error, match it; otherwise fail hard.
+                if let (None, Some(expected_error)) = (&case.want_result, &case.want_error) {
+                    let error_str = options_error.to_string();
+                    if error_str.contains(expected_error) {
+                        println!(
+                            "✓ Execution-options error matches expected for case '{}'",
+                            case.note
+                        );
+                        println!("passed");
+                        continue;
+                    }
+                    panic_with_listing!(
+                        &last_listing,
+                        &case.note,
+                        "Execution-options error does not match expected for case '{}':\nExpected: '{expected_error}'\nActual: '{error_str}'",
+                        case.note
+                    );
+                }
+                dump_rvm_listing(&case.note, &last_listing);
+                return Err(options_error);
+            }
+        };
 
         if let Err(compilation_error) = &compilation_result {
             if let (None, Some(expected_error)) = (&case.want_result, &case.want_error) {