grammar: preserve tagged unions through singleton allOf

common/json-schema-to-grammar.cpp

@@ -964,6 +964,27 @@ class common_schema_converter {
                     properties, required, name,
                     schema.contains("additionalProperties") ? schema["additionalProperties"] : json()));
         }
+        // A singleton union wrapper adds no constraint unless a sibling does.
+        if (schema.contains("allOf") && schema["allOf"].is_array() && schema["allOf"].size() == 1) {
+            bool transparent = true;
+            for (const auto & item : schema.items()) {
+                const auto & key = item.key();
+                if (key != "allOf" && key != "description" && key != "title" && key != "$comment" &&
+                    key != "default" && key != "examples" && key != "definitions" && key != "$defs") {
+                    transparent = false;
+                }
+            }
+            const json * component = &schema["allOf"][0];
+            if (component->is_object() && component->contains("$ref")) {
+                auto ref = _refs.find((*component)["$ref"]);
+                if (ref != _refs.end() && component->size() == 1) {
+                    component = &ref->second;
+                }
+            }
+            if (transparent && (component->contains("oneOf") || component->contains("anyOf"))) {
+                return visit(*component, name);
+            }
+        }
         if ((schema_type.is_null() || schema_type == "object" || schema_type == "string") && schema.contains("allOf")) {
             std::unordered_set<std::string> required;
             std::vector<std::pair<std::string, json>> properties;

common/speculative.cpp

@@ -634,6 +634,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
                 /*.seq_id   =*/ nullptr,
                 /*.logits   =*/ nullptr,
                 /*.label    =*/ nullptr,
+                /*.label_weight =*/ nullptr,
             };
             const int32_t rc = llama_encode(ctx_dft, enc_batch);
             if (rc != 0) {
@@ -1128,6 +1129,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                     /*.seq_id   =*/ nullptr,
                     /*.logits   =*/ nullptr,
                     /*.label    =*/ nullptr,
+                    /*.label_weight =*/ nullptr,
                 };
 
                 int32_t rc = llama_encode(ctx_dft, enc_batch);

tests/CMakeLists.txt

@@ -150,10 +150,6 @@ endif ()
 
 llama_build(test-recurrent-state-rollback.cpp)
 
-# the exact walk trains on the window's whole gradient whatever the chunking: run on a pure-attention AND a
-# hybrid model (-m), on a GPU backend as well as the CPU (a hybrid passed on the CPU while Metal tore its state)
-llama_build(test-walk-exact.cpp)
-
 # GGML_TENSOR_FLAG_GRAD: a leaf whose gradient the backward computes, never optimized (the walk reads the cached prefix K/V gradient)
 llama_build_and_test(test-grad-leaf.cpp)
 
@@ -170,6 +166,9 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     llama_build_and_test(test-llama-grammar.cpp)
     llama_build_and_test(test-batch-alloc.cpp)
     llama_build_and_test(test-walk-plan.cpp)
+    # the exact walk trains on the window's whole gradient whatever the chunking: run on a pure-attention AND a
+    # hybrid model (-m), on a GPU backend as well as the CPU (a hybrid passed on the CPU while Metal tore its state)
+    llama_build(test-walk-exact.cpp)
     llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
     target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server)
     target_link_libraries(test-chat PRIVATE server-context)

tests/test-chat.cpp

@@ -14,6 +14,7 @@
 #include "common.h"
 #include "ggml.h"
 #include "log.h"
+#include "json-schema-to-grammar.h"
 
 #include <algorithm>
 #include <chrono>
@@ -2116,6 +2117,35 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
         // Qwen3.5 (basically same as Nemotron, but keeping separate tests just in case)
         auto tst = peg_tester("models/templates/Qwen3.5-4B.jinja", detailed_debug);
 
+        // Captured tagged edit schema: allOf -> local reference -> oneOf.
+        const common_chat_tool tagged_edit_tool{
+            "edit_file", "Edit an existing file", R"SCHEMA({"definitions":{"EditMode":{"description":"How to edit a file (four modes).","oneOf":[{"description":"Replace content between line numbers (1-indexed, inclusive).","properties":{"end_line":{"format":"uint32","minimum":0.0,"type":"integer"},"new_content":{"type":"string"},"start_line":{"format":"uint32","minimum":1.0,"type":"integer"},"type":{"enum":["line_range"],"type":"string"}},"required":["end_line","new_content","start_line","type"],"type":"object"},{"description":"Find text and replace it.","properties":{"all":{"default":false,"type":"boolean"},"replace":{"type":"string"},"search":{"type":"string"},"type":{"enum":["search_replace"],"type":"string"}},"required":["replace","search","type"],"type":"object"},{"description":"Insert content at a specific line (pushes existing lines down).","properties":{"content":{"type":"string"},"line":{"format":"uint32","minimum":1.0,"type":"integer"},"type":{"enum":["insert_at"],"type":"string"}},"required":["content","line","type"],"type":"object"},{"description":"Append content to end of file.","properties":{"content":{"type":"string"},"type":{"enum":["append"],"type":"string"}},"required":["content","type"],"type":"object"}]}},"properties":{"description":{"description":"Optional note recorded in change history.","type":["string","null"]},"edit_mode":{"allOf":[{"$ref":"#/definitions/EditMode"}],"description":"Operation to apply."},"file_path":{"description":"Path relative to the workspace root.","type":"string"}},"required":["edit_mode","file_path"],"type":"object"})SCHEMA"
+        };
+        const auto tagged_schema = common_json::parse(tagged_edit_tool.parameters);
+        const auto tagged_grammar = json_schema_to_grammar(tagged_schema);
+        const std::string edit_args = R"({"edit_mode":{"replace":"new","search":"old","type":"search_replace"},"file_path":"file.rs"})";
+        assert_equals(true, match_string(edit_args, build_grammar(tagged_grammar).get()));
+        assert_equals(true, match_string(R"({"edit_mode":{"content":"new","type":"append"},"file_path":"file.rs"})",
+                            build_grammar(tagged_grammar).get()));
+        assert_equals(false, match_string(R"({"edit_mode":{},"file_path":"file.rs"})", build_grammar(tagged_grammar).get()));
+        assert_equals(false, match_string(R"({"edit_mode":{"type":"unknown"},"file_path":"file.rs"})",
+                             build_grammar(tagged_grammar).get()));
+        auto closed = tagged_schema;
+        closed["properties"]["edit_mode"]["additionalProperties"] = false;
+        assert_equals(false, match_string(edit_args, build_grammar(json_schema_to_grammar(closed)).get()));
+        auto constrained = tagged_schema;
+        constrained["properties"]["edit_mode"]["properties"] = {{"guard", {{"const", "kept"}}}};
+        constrained["properties"]["edit_mode"]["required"] = {"guard"};
+        assert_equals(false, match_string(edit_args, build_grammar(json_schema_to_grammar(constrained)).get()));
+        tst.test("<tool_call>\n<function=edit_file>\n"
+                 "<parameter=edit_mode>\n{\"replace\":\"new\",\"search\":\"old\",\"type\":\"search_replace\"}\n</parameter>\n"
+                 "<parameter=file_path>\nfile.rs\n</parameter>\n</function>\n</tool_call>")
+            .enable_thinking(false)
+            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
+            .tools({tagged_edit_tool})
+            .expect_tool_calls({{"edit_file", edit_args, {}}})
+            .run();
+
         tst.test("I'm\nthinking\n</think>\n\nHello, world!\nWhat's up?")
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .enable_thinking(true)

tests/test-walk-exact.cpp

@@ -829,7 +829,7 @@ int main(int argc, char ** argv) {
             GGML_ASSERT(llama_state_seq_set_data(ctx, generation_prefix.data(), generation_prefix.size(), 0) == generation_prefix.size());
             llama_batch retained = llama_batch_init(WINDOW - r0, 0, 1);
             common_batch_clear(retained);
-            for (int32_t i = 0; i < WINDOW - r0; ++i) {
+            for (uint32_t i = 0; i < WINDOW - r0; ++i) {
                 common_batch_add(retained, tokens[r0 + i], r0 + i, {0}, true);
                 retained.label[i] = labels[i];
             }