server : allow text-only slot save/restore with mtmd (#25076)

tools/server/server-common.h

@@ -207,6 +207,9 @@ public:
 
     bool empty() const { return tokens.empty(); }
 
+    // true if the sequence actually contains image/audio chunks.
+    bool has_media() const { return !map_idx_to_media.empty(); }
+
     void clear() {
         map_idx_to_media.clear();
         tokens.clear();

tools/server/server-context.cpp

@@ -2011,10 +2011,13 @@ private:
         queue_results.send(std::move(res));
     }
 
-    // if multimodal is enabled, send an error and return false
+    // Gate slot save/restore/erase on slot content (does it hold media),
-    bool check_no_mtmd(const int id_task) {
+    // not model capability: a multimodal model may hold a pure-text slot.
-        if (mctx) {
+    bool check_slot_no_media(const server_slot & slot, const int id_task) {
-            send_error(id_task, "This feature is not supported by multimodal", ERROR_TYPE_NOT_SUPPORTED);
+        if (slot.prompt.tokens.has_media()) {
+            send_error(id_task,
+                "This operation is not supported while the slot holds image/audio tokens (a pure-text prefix is supported)",
+                ERROR_TYPE_NOT_SUPPORTED);
             return false;
         }
         return true;
@@ -2502,16 +2505,15 @@ private:
                 } break;
             case SERVER_TASK_TYPE_SLOT_SAVE:
                 {
-                    if (!check_no_mtmd(task.id)) {
-                        break;
-                    }
-
                     const int id_slot = task.slot_action.id_slot;
                     server_slot * slot = get_slot_by_id(id_slot);
                     if (slot == nullptr) {
                         send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST);
                         break;
                     }
+                    if (!check_slot_no_media(*slot, task.id)) {
+                        break;
+                    }
                     if (slot->is_processing()) {
                         // if requested slot is unavailable, we defer this task for processing later
                         SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id);
@@ -2519,13 +2521,13 @@ private:
                         break;
                     }
 
-                    const size_t token_count = slot->prompt.tokens.size();
                     const int64_t t_start = ggml_time_us();
 
                     std::string filename = task.slot_action.filename;
                     std::string filepath = task.slot_action.filepath;
 
-                    const llama_tokens & tokens = slot->prompt.tokens.get_tokens();
+                    const llama_tokens tokens = slot->prompt.tokens.get_text_tokens();
+                    const size_t token_count = tokens.size();
                     const size_t nwrite = llama_state_seq_save_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), token_count);
 
                     const int64_t t_end = ggml_time_us();
@@ -2543,7 +2545,6 @@ private:
                 } break;
             case SERVER_TASK_TYPE_SLOT_RESTORE:
                 {
-                    if (!check_no_mtmd(task.id)) break;
                     const int id_slot = task.slot_action.id_slot;
                     server_slot * slot = get_slot_by_id(id_slot);
                     if (slot == nullptr) {
@@ -2590,15 +2591,16 @@ private:
                 } break;
             case SERVER_TASK_TYPE_SLOT_ERASE:
                 {
-                    if (!check_no_mtmd(task.id)) {
-                        break;
-                    }
                     const int id_slot = task.slot_action.id_slot;
                     server_slot * slot = get_slot_by_id(id_slot);
                     if (slot == nullptr) {
                         send_error(task, "Invalid slot ID", ERROR_TYPE_INVALID_REQUEST);
                         break;
                     }
+                    // Gate on slot content, consistent with save/restore.
+                    if (!check_slot_no_media(*slot, task.id)) {
+                        break;
+                    }
                     if (slot->is_processing()) {
                         // if requested slot is unavailable, we defer this task for processing later
                         SRV_DBG("requested slot is unavailable, defer task, id_task = %d\n", task.id);

tools/server/tests/unit/test_slot_save.py

@@ -1,5 +1,7 @@
 import pytest
 from utils import *
+import base64
+import requests
 
 server = ServerPreset.tinyllama2()
 
@@ -96,3 +98,127 @@ def test_slot_erase():
     assert res.status_code == 200
     assert match_regex("(Whiskers|Flana)+", res.body["content"])
     assert res.body["timings"]["prompt_n"] == 21  # all tokens are processed
+
+
+#
+# Multimodal server (mmproj loaded) slot save/restore.
+#
+# Regression coverage for issue #21133: slot save/restore/erase must be gated on
+# the slot's CONTENT (does it actually hold image/audio tokens) rather than the
+# model's CAPABILITY (is an mmproj loaded). A pure-text slot on a multimodal
+# server must save/restore/erase normally; a slot that actually holds an image
+# must be rejected with ERROR_TYPE_NOT_SUPPORTED (HTTP 501).
+#
+
+IMG_URL_CAT = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png"
+
+
+def _get_img_base64(url: str) -> str:
+    response = requests.get(url)
+    response.raise_for_status()  # Raise an exception for bad status codes
+    return base64.b64encode(response.content).decode("utf-8")
+
+
+@pytest.fixture
+def mmproj_server():
+    # tinygemma3 is a small multimodal model: the mmproj is provided by the HF
+    # registry API and auto-downloaded on first run.
+    os.environ['LLAMA_MEDIA_MARKER'] = '<__media__>'
+    mm_server = ServerPreset.tinygemma3()
+    mm_server.slot_save_path = "./tmp"
+    mm_server.temperature = 0.0
+    return mm_server
+
+
+def test_slot_save_restore_text_only_on_multimodal(mmproj_server):
+    server = mmproj_server
+    server.start()
+
+    # A pure-text prompt processed on slot 1 of a multimodal server.
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "The quick brown fox jumps over the lazy dog.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    prompt_n = res.body["timings"]["prompt_n"]
+    assert prompt_n > 0  # all tokens are processed
+
+    # Saving a pure-text slot must succeed even though an mmproj is loaded.
+    res = server.make_request("POST", "/slots/1?action=save", data={
+        "filename": "mm_slot1.bin",
+    })
+    assert res.status_code == 200
+    n_saved = res.body["n_saved"]
+    assert n_saved > 0  # the slot KV (prompt + generated tokens) was written
+
+    # Restore the saved state into slot 0; it must round-trip exactly.
+    res = server.make_request("POST", "/slots/0?action=restore", data={
+        "filename": "mm_slot1.bin",
+    })
+    assert res.status_code == 200
+    assert res.body["n_restored"] == n_saved
+
+    # The restored slot is usable for a follow-up completion. We do NOT assert
+    # prefix reuse here: tinygemma3 is a SWA model, which forces full prompt
+    # re-processing after a restore (a model property, not the save/restore gate
+    # under test).
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "The quick brown fox jumps over the lazy dog.",
+        "id_slot": 0,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+
+
+def test_slot_save_rejected_when_slot_holds_image(mmproj_server):
+    server = mmproj_server
+    server.start()
+
+    # Process a prompt that actually contains an image on slot 1.
+    res = server.make_request("POST", "/completions", data={
+        "temperature": 0.0,
+        "top_k": 1,
+        "id_slot": 1,
+        "cache_prompt": True,
+        "prompt": {
+            "prompt_string": "What is this: <__media__>\n",
+            "multimodal_data": [ _get_img_base64(IMG_URL_CAT) ],
+        },
+    })
+    assert res.status_code == 200
+
+    # Saving a slot that holds image tokens must be rejected (HTTP 501,
+    # not_supported_error).
+    res = server.make_request("POST", "/slots/1?action=save", data={
+        "filename": "mm_slot_image.bin",
+    })
+    assert res.status_code != 200
+    assert res.body["error"]["type"] == "not_supported_error"
+
+
+def test_slot_erase_text_only_on_multimodal(mmproj_server):
+    server = mmproj_server
+    server.start()
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "The quick brown fox jumps over the lazy dog.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    prompt_n = res.body["timings"]["prompt_n"]
+    assert prompt_n > 0  # all tokens are processed
+
+    # Erasing a pure-text slot must succeed even though an mmproj is loaded.
+    res = server.make_request("POST", "/slots/1?action=erase")
+    assert res.status_code == 200
+
+    # Re-running the same prompt should process all tokens again.
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "The quick brown fox jumps over the lazy dog.",
+        "id_slot": 1,
+        "cache_prompt": True,
+    })
+    assert res.status_code == 200
+    assert res.body["timings"]["prompt_n"] == prompt_n  # all tokens are processed again