fix: compile render issue (#404)

Author: yesudeep

java.MODULE.bazel

@@ -39,12 +39,12 @@ maven.install(
         "com.fasterxml.jackson.core:jackson-databind:2.20.1",
         "com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.20.1",
     ],
-    lock_file = "//:maven_install.json",
     known_contributing_modules = [
         "bazel_worker_java",
         "dotprompt",
         "protobuf",
     ],
+    lock_file = "//:maven_install.json",
     repositories = [
         "https://repo1.maven.org/maven2",
     ],

java/com/google/dotprompt/Dotprompt.java

@@ -700,12 +700,16 @@ private RenderedPrompt renderWithTemplate(
       mergedData.putAll(data);
     }
 
-    // Context handling
+    // Context handling: Add all context entries as @-prefixed variables
+    // Each key in context becomes accessible as @key in templates
+    // e.g., context: {state: {...}, auth: {...}} creates @state and @auth
     Map<String, Object> contextData = new HashMap<>();
     if (data != null && data.containsKey("context")) {
+      @SuppressWarnings("unchecked")
       Map<String, Object> ctx = (Map<String, Object>) data.get("context");
-      if (ctx.containsKey("state")) {
-        contextData.put("state", ctx.get("state"));
+      if (ctx != null) {
+        // Add all context entries, not just state
+        contextData.putAll(ctx);
       }
     }
 

java/com/google/dotprompt/SpecTest.java

@@ -142,9 +142,23 @@ private void runTestGroup(Map<String, Object> group) throws Exception {
       Map<String, Object> options = (Map<String, Object>) test.get("options");
       Map<String, Object> expect = (Map<String, Object>) test.get("expect");
 
-      Map<String, Object> renderData = data;
+      // Build render data: merge input values into top-level, but also preserve context
+      // for @-prefixed variable access and messages for history helper
+      Map<String, Object> renderData = new HashMap<>();
       if (data != null && data.containsKey("input")) {
-        renderData = (Map<String, Object>) data.get("input");
+        @SuppressWarnings("unchecked")
+        Map<String, Object> inputData = (Map<String, Object>) data.get("input");
+        if (inputData != null) {
+          renderData.putAll(inputData);
+        }
+      }
+      // Preserve context for @-prefixed variables (e.g., @auth, @state, @user)
+      if (data != null && data.containsKey("context")) {
+        renderData.put("context", data.get("context"));
+      }
+      // Preserve messages for {{history}} helper
+      if (data != null && data.containsKey("messages")) {
+        renderData.put("messages", data.get("messages"));
       }
 
       try {

python/dotpromptz/tests/dotpromptz/dotprompt_test.py

@@ -39,7 +39,12 @@
 import pytest
 
 from dotpromptz.dotprompt import Dotprompt, _identify_partials
-from dotpromptz.typing import ModelConfigT, ParsedPrompt, PromptMetadata, ToolDefinition
+from dotpromptz.typing import (
+    ModelConfigT,
+    ParsedPrompt,
+    PromptMetadata,
+    ToolDefinition,
+)
 from handlebarrz import HelperFn, HelperOptions
 
 
@@ -144,6 +149,29 @@ def test_define_tool(mock_handlebars: Mock) -> None:
     assert result == dotprompt
 
 
+class TestCompileRender(IsolatedAsyncioTestCase):
+    async def test_compile_render_mock(self) -> None:
+        """Test that handlebarrz compile produces a working render function.
+
+        The handlebarrz.compile() method returns a function that takes:
+        - context: A dict with the template variables
+        - options: RuntimeOptions with a 'data' key for @ variables
+        """
+        dotprompt = Dotprompt()
+        template_source = 'hello {{name}} ({{@state.name}}, {{@auth.email}})'
+
+        render_fn = dotprompt._handlebars.compile(template_source)
+
+        # Call with positional args as the compiled function expects
+        # context dict for regular variables, RuntimeOptions for @ variables
+        result = render_fn(
+            {'name': 'foo'},  # context - accessed as {{name}}
+            {'data': {'state': {'name': 'bar'}, 'auth': {'email': 'a@b.c'}}},  # RuntimeOptions - accessed as {{@...}}
+        )
+
+        assert result == 'hello foo (bar, a@b.c)'
+
+
 @patch('dotpromptz.dotprompt.parse_document')
 def test_parse(mock_parse_document: Mock, mock_handlebars: Mock) -> None:
     """Test parsing a prompt."""

rs/dotprompt/src/dotprompt.rs

@@ -341,41 +341,37 @@ impl Dotprompt {
             },
         );
 
-        // Add @state from context.state if available
+        // Add all context variables as @-prefixed variables
+        // Each key in context becomes accessible as @key in templates
+        // e.g., context: {state: {...}, auth: {...}} creates @state and @auth
+        let mut template_to_render = parsed.template.clone();
         if let (serde_json::Value::Object(map), Some(context)) =
             (&mut render_context, &data.context)
         {
-            // context is HashMap<String, Value>, get "state" key
-            // Add state as __state (workaround for Handlebars @ prefix)
-            if let Some(state) = context.get("state") {
-                if let Some(state_obj) = state.as_object() {
-                    for (k, v) in state_obj {
-                        // Add each state field as __state.field
-                        let at_state = map
-                            .entry("__state".to_string())
-                            .or_insert(serde_json::Value::Object(serde_json::Map::new()));
-                        if let serde_json::Value::Object(at_state_map) = at_state {
-                            at_state_map.insert(k.clone(), v.clone());
-                        }
-                    }
-                } else {
-                    // If state is not an object, just insert it directly
-                    map.insert("__state".to_string(), state.clone());
-                }
+            for (key, value) in context {
+                // Use __ctx_{key} as workaround since Handlebars treats @ as private data prefix
+                let internal_key = format!("__ctx_{key}");
+                map.insert(internal_key.clone(), value.clone());
+
+                // Replace all occurrences of @{key} with __ctx_{key} in template
+                template_to_render = template_to_render
+                    .replace(&format!("{{{{@{key}."), &format!("{{{{__{key}."))
+                    .replace(&format!("{{{{ @{key}."), &format!("{{{{ __{key}."));
+
+                // Also handle shortcut for just the key (e.g., {{@isAdmin}})
+                template_to_render = template_to_render
+                    .replace(&format!("{{{{@{key}}}}}"), &format!("{{{{__{key}}}}}"))
+                    .replace(&format!("{{{{ @{key} }}}}"), &format!("{{{{ __{key} }}}}"));
+
+                // Insert with __ prefix for template access
+                map.insert(format!("__{key}"), value.clone());
             }
         }
 
-        // Preprocess template to replace @state with __state for Handlebars compatibility
-        // Handlebars treats @ as special prefix for private data, so we use __state as workaround
-        let preprocessed_template = parsed
-            .template
-            .replace("{{@state.", "{{__state.")
-            .replace("{{ @state.", "{{ __state.");
-
         // Render template
         let rendered_string = self
             .handlebars
-            .render_template(&preprocessed_template, &render_context)
+            .render_template(&template_to_render, &render_context)
             .map_err(|e| DotpromptError::RenderError(e.to_string()))?;
 
         // Convert to messages (passing data for history)

spec/metadata.yaml

@@ -59,6 +59,55 @@
           - role: user
             content: [{ text: "Current count is 100\nStatus is pending\n" }]
 
+# Tests accessing arbitrary context variables beyond @state, such as
+# @auth and @user. Any key in the context object should be accessible
+# as an @ variable in the template.
+- name: metadata_context_variables
+  template: |
+    Hello {{name}} ({{@auth.email}}, {{@user.role}})
+  tests:
+    - desc: accesses arbitrary context variables
+      data:
+        input:
+          name: "Alice"
+        context:
+          auth:
+            email: "alice@example.com"
+          user:
+            role: "admin"
+      expect:
+        messages:
+          - role: user
+            content: [{ text: "Hello Alice (alice@example.com, admin)\n" }]
+
+    - desc: handles missing context variables gracefully
+      data:
+        input:
+          name: "Bob"
+        context:
+          auth:
+            email: "bob@example.com"
+      expect:
+        messages:
+          - role: user
+            content: [{ text: "Hello Bob (bob@example.com, )\n" }]
+
+    - desc: handles deeply nested context variables
+      data:
+        input:
+          name: "Carol"
+        context:
+          auth:
+            email: "carol@example.com"
+            permissions:
+              canEdit: true
+          user:
+            role: "editor"
+      expect:
+        messages:
+          - role: user
+            content: [{ text: "Hello Carol (carol@example.com, editor)\n" }]
+
 # Tests that raw frontmatter is preserved alongside parsed frontmatter,
 # allowing access to both structured and unstructured metadata.
 - name: raw