Fix multi-turn chat and tool-call defects across model families

cactus-engine/cactus_engine.h

@@ -65,6 +65,15 @@ CACTUS_FFI_EXPORT int cactus_tokenize(
     size_t* out_token_len
 );
 
+CACTUS_FFI_EXPORT int cactus_render_prompt(
+    cactus_model_t model,
+    const char* messages_json,
+    const char* options_json,
+    const char* tools_json,
+    char* prompt_buffer,
+    size_t buffer_size
+);
+
 CACTUS_FFI_EXPORT int cactus_score_window(
     cactus_model_t model,
     const uint32_t* tokens,

cactus-engine/src/complete.cpp

@@ -387,6 +387,7 @@ struct PreparedPrompt {
     std::vector<std::string> audio_paths;
     std::vector<ChatMessage> messages;
     std::vector<ToolFunction> tools;
+    std::string rendered;
     std::vector<uint32_t> tokens;
     size_t context_token_count = 0;
     std::vector<std::vector<CactusModelHandle::ProcessedImage>> images;
@@ -577,23 +578,23 @@ PreparedPrompt prepare_prompt(
     } else if (tokenizer->is_lfm2_family()) {
         formatted_tools = chat_tools::serialize_tools_lfm2(prompt.tools);
     } else {
-        formatted_tools = gemma::format_tools(prompt.tools, true);
+        formatted_tools = gemma::format_tools(prompt.tools, !Config::is_gemma3_family(prompt.model_type));
     }
 
     if (apply_tool_constraints) {
         setup_tool_constraints(handle, prompt.tools, prompt.options.force_tools, prompt.options.temperature);
     }
 
-    std::string full_prompt = tokenizer->format_chat_prompt(
+    prompt.rendered = tokenizer->format_chat_prompt(
         prompt.messages,
         add_generation_prompt,
         formatted_tools,
         prompt.options.enable_thinking_if_supported
     );
-    if (full_prompt.find("ERROR:") == 0) {
-        throw std::runtime_error(full_prompt.substr(6));
+    if (prompt.rendered.find("ERROR:") == 0) {
+        throw std::runtime_error(prompt.rendered.substr(6));
     }
-    prompt.tokens = tokenizer->encode(full_prompt);
+    prompt.tokens = tokenizer->encode(prompt.rendered);
     prompt.context_token_count = prompt.tokens.size();
     prompt.images = images_from_message(prompt.messages);
     return prompt;
@@ -787,7 +788,7 @@ int cactus_complete(
         bool has_images = prompt.has_images();
         bool has_audio = prompt.has_audio();
         const bool cloud_disabled = env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF");
-        const bool cloud_eligible = !cloud_disabled &&
+        const bool cloud_eligible = !cloud_disabled && !handle->cloud_handoff_disabled &&
             prompt.options.auto_handoff && (!has_images || prompt.options.handoff_with_images);
         handle->model->reset_handoff_probe_rollout();
         const bool defer_local_stream_until_probe = cloud_eligible && handle->model->has_handoff_probe();
@@ -809,6 +810,14 @@ int cactus_complete(
             return request;
         };
 
+        auto disable_handoff_on_auth_failure = [&](const std::string& error) {
+            if (error.rfind("http_401", 0) == 0 || error.rfind("http_403", 0) == 0) {
+                handle->cloud_handoff_disabled = true;
+                CACTUS_LOG_WARN("cloud_handoff", "Cloud auth failed (" << error
+                    << "); disabling cloud handoff for this session");
+            }
+        };
+
         auto return_cloud_completion = [&](const CloudCompletionResult& cloud_result,
                                            double ttft_ms,
                                            double total_ms,
@@ -859,6 +868,7 @@ int cactus_complete(
             }
             std::string cloud_error = cloud_result.error.empty() ? "cloud completion failed" : cloud_result.error;
             CACTUS_LOG_WARN("cloud_handoff", "Cloud completion failed before local generation: " << cloud_error);
+            disable_handoff_on_auth_failure(cloud_error);
             handle_error_response(("cloud handoff failed before local generation: " + cloud_error).c_str(),
                                   response_buffer, buffer_size);
             return -1;
@@ -923,7 +933,8 @@ int cactus_complete(
         entropy.add(first_token_entropy);
 
         if (!matches_stop_sequence(generated_tokens, stop_token_sequences)) {
-            if (!defer_local_stream_until_probe
+            if (cloud_eligible
+                && !defer_local_stream_until_probe
                 && !pre_generation_cloud_attempted
                 && confidence < prompt.options.confidence_threshold) {
                 CACTUS_LOG_WARN("cloud_handoff", "Cloud handoff triggered before local streaming; waiting up to "
@@ -941,6 +952,7 @@ int cactus_complete(
                 }
                 cloud_error = cloud_result.error.empty() ? "cloud completion failed" : cloud_result.error;
                 CACTUS_LOG_WARN("cloud_handoff", "Cloud completion failed before local streaming, falling back to local output: " << cloud_error);
+                disable_handoff_on_auth_failure(cloud_error);
             }
 
             if (callback && !defer_local_stream_until_probe) {
@@ -1057,6 +1069,7 @@ int cactus_complete(
             cloud_error = cloud_result.error.empty() ? "cloud completion failed" : cloud_result.error;
             CACTUS_LOG_WARN("cloud_handoff", "Cloud completion failed after probe handoff, falling back to local output: "
                 << cloud_error);
+            disable_handoff_on_auth_failure(cloud_error);
         }
 
         if (callback && defer_local_stream_until_probe && !primary_response.empty()) {
@@ -1258,6 +1271,27 @@ int cactus_tokenize(
     }
 }
 
+int cactus_render_prompt(
+    cactus_model_t model,
+    const char* messages_json,
+    const char* options_json,
+    const char* tools_json,
+    char* prompt_buffer,
+    size_t buffer_size
+) {
+    if (!model || !messages_json || !prompt_buffer || buffer_size == 0) return -1;
+
+    try {
+        auto* handle = static_cast<CactusModelHandle*>(model);
+        auto prompt = prepare_prompt(handle, messages_json, options_json, tools_json, false, true);
+        if (prompt.rendered.size() >= buffer_size) return -2;
+        std::strcpy(prompt_buffer, prompt.rendered.c_str());
+        return static_cast<int>(prompt.rendered.size());
+    } catch (...) {
+        return -1;
+    }
+}
+
 int cactus_score_window(
     cactus_model_t model,
     const uint32_t* tokens,

cactus-engine/src/constraints.cpp

@@ -264,9 +264,12 @@ void ToolCallConstrainer::tokenize_grammar_elements() {
     escape_tokens_.clear();
     gemma_response_start_tokens_.clear();
 
-    add_tokens_for_string(call_start_tag_, gemma_call_start_tokens_);
-    add_tokens_for_string(call_end_tag_, gemma_call_end_tokens_);
-    add_tokens_for_string("<|tool_response>", gemma_response_start_tokens_);
+    call_start_sequence_ = tokenizer_->encode(call_start_tag_);
+    call_end_sequence_ = tokenizer_->encode(call_end_tag_);
+    std::vector<uint32_t> response_sequence = tokenizer_->encode(response_start_tag_);
+    if (call_start_sequence_.size() == 1) gemma_call_start_tokens_.insert(call_start_sequence_[0]);
+    if (call_end_sequence_.size() == 1) gemma_call_end_tokens_.insert(call_end_sequence_[0]);
+    if (response_sequence.size() == 1) gemma_response_start_tokens_.insert(response_sequence[0]);
     add_tokens_for_string("call:", gemma_call_prefix_tokens_);
 
     add_tokens_for_string("{", open_brace_tokens_);
@@ -289,15 +292,24 @@ void ToolCallConstrainer::init(Config::ModelType model_type,
     }
     tokenizer_ = tokenizer;
     generated_text_.clear();
+    current_bias_.clear();
     brace_depth_ = 0;
     in_argument_string_ = false;
+    forced_tag_progress_ = 0;
     const bool model_supported = is_needle() || Config::is_gemma_family(model_type_);
     active_ = model_supported && !function_names_.empty() && tokenizer != nullptr;
 
     reset_needle_constraints();
     state_ = is_needle() ? State::NEEDLE_START : State::GEMMA_START;
-    call_start_tag_ = "<|tool_call>";
-    call_end_tag_ = "<tool_call|>";
+    if (Config::is_gemma3_family(model_type_)) {
+        call_start_tag_ = "<start_function_call>";
+        call_end_tag_ = "<end_function_call>";
+        response_start_tag_ = "<start_function_response>";
+    } else {
+        call_start_tag_ = "<|tool_call>";
+        call_end_tag_ = "<tool_call|>";
+        response_start_tag_ = "<|tool_response>";
+    }
 
     if (!active_) {
         return;
@@ -310,7 +322,15 @@ void ToolCallConstrainer::init(Config::ModelType model_type,
     compute_bias();
 }
 
-void ToolCallConstrainer::update(uint32_t /*token_id*/, const std::string& decoded_text) {
+void ToolCallConstrainer::advance_forced_tag(const std::vector<uint32_t>& sequence, uint32_t token_id) {
+    if (forced_tag_progress_ < sequence.size() && token_id == sequence[forced_tag_progress_]) {
+        forced_tag_progress_++;
+    } else {
+        forced_tag_progress_ = (!sequence.empty() && token_id == sequence[0]) ? 1 : 0;
+    }
+}
+
+void ToolCallConstrainer::update(uint32_t token_id, const std::string& decoded_text) {
     if (!active_) return;
 
     generated_text_ += decoded_text;
@@ -324,9 +344,11 @@ void ToolCallConstrainer::update(uint32_t /*token_id*/, const std::string& decod
 
     switch (state_) {
         case State::GEMMA_START:
+            advance_forced_tag(call_start_sequence_, token_id);
             if (generated_text_.find(call_start_tag_) != std::string::npos) {
                 state_ = State::GEMMA_EXPECT_CALL;
                 generated_text_.clear();
+                forced_tag_progress_ = 0;
             }
             break;
 
@@ -369,6 +391,7 @@ void ToolCallConstrainer::update(uint32_t /*token_id*/, const std::string& decod
                         if (brace_depth_ == 0) {
                             state_ = State::GEMMA_EXPECT_END;
                             in_argument_string_ = false;
+                            forced_tag_progress_ = 0;
                             break;
                         }
                     }
@@ -377,6 +400,7 @@ void ToolCallConstrainer::update(uint32_t /*token_id*/, const std::string& decod
             break;
 
         case State::GEMMA_EXPECT_END:
+            advance_forced_tag(call_end_sequence_, token_id);
             if (generated_text_.find(call_end_tag_) != std::string::npos) {
                 state_ = State::DONE;
                 generated_text_.clear();
@@ -451,8 +475,8 @@ void ToolCallConstrainer::compute_bias() {
 
     switch (state_) {
         case State::GEMMA_START:
-            for (uint32_t t : gemma_call_start_tokens_) {
-                current_bias_[t] = FORCE_BIAS;
+            if (forced_tag_progress_ < call_start_sequence_.size()) {
+                current_bias_[call_start_sequence_[forced_tag_progress_]] = FORCE_BIAS;
             }
             for (uint32_t t : open_brace_tokens_) {
                 current_bias_[t] = BLOCK_BIAS;
@@ -522,8 +546,8 @@ void ToolCallConstrainer::compute_bias() {
             break;
 
         case State::GEMMA_EXPECT_END:
-            for (uint32_t t : gemma_call_end_tokens_) {
-                current_bias_[t] = FORCE_BIAS;
+            if (forced_tag_progress_ < call_end_sequence_.size()) {
+                current_bias_[call_end_sequence_[forced_tag_progress_]] = FORCE_BIAS;
             }
             for (uint32_t t : open_brace_tokens_) {
                 current_bias_[t] = BLOCK_BIAS;
@@ -543,6 +567,7 @@ void ToolCallConstrainer::reset() {
     current_bias_.clear();
     brace_depth_ = 0;
     in_argument_string_ = false;
+    forced_tag_progress_ = 0;
     reset_needle_constraints();
 
     state_ = is_needle() ? State::NEEDLE_START : State::GEMMA_START;
@@ -558,7 +583,6 @@ void Model::set_tool_constraints(const std::vector<ToolConstraintSpec>& tools) {
 }
 
 void Model::clear_tool_constraints() {
-    tool_constrainer_.reset();
     tool_constrainer_.init(config_.model_type, {}, tokenizer_.get());
 }
 

cactus-engine/src/engine.h

@@ -240,6 +240,10 @@ struct Config {
         return t == ModelType::GEMMA || t == ModelType::GEMMA3N || t == ModelType::GEMMA4;
     }
 
+    static bool is_gemma3_family(ModelType t) {
+        return t == ModelType::GEMMA || t == ModelType::GEMMA3N;
+    }
+
     bool from_json(const std::string& json_path);
     std::string to_json() const;
     // Disable rolling unless 0 < target < trigger (when trigger > 0).
@@ -337,6 +341,7 @@ class Tokenizer {
     virtual bool has_chat_template() const { return has_chat_template_; }
     bool is_qwen_family() const { return model_type_ == ModelType::QWEN; }
     bool is_lfm2_family() const { return model_type_ == ModelType::LFM2; }
+    bool has_function_call_tokens() const { return encode("<start_function_call>").size() == 1; }
     std::string get_default_stop_sequence() const;
 
     virtual bool load_vocabulary_with_config(const std::string& vocab_file, const std::string& merges_file, const std::string& config_file) = 0;
@@ -535,6 +540,10 @@ class ToolCallConstrainer {
 
     std::string call_start_tag_;
     std::string call_end_tag_;
+    std::string response_start_tag_;
+    std::vector<uint32_t> call_start_sequence_;
+    std::vector<uint32_t> call_end_sequence_;
+    size_t forced_tag_progress_ = 0;
 
     std::unordered_set<uint32_t> all_func_name_tokens_;
     std::unordered_map<std::string, std::vector<uint32_t>> func_name_sequences_;
@@ -557,6 +566,7 @@ class ToolCallConstrainer {
     void tokenize_grammar_elements();
     void add_tokens_for_string(const std::string& str, std::unordered_set<uint32_t>& token_set);
     void add_tokens_for_prefix_string(const std::string& prefix, std::unordered_set<uint32_t>& token_set);
+    void advance_forced_tag(const std::vector<uint32_t>& sequence, uint32_t token_id);
     void tokenize_function_names(bool quote_names);
     void init_common_tokens();