Custom spec algorithm can handle server args (#28162)

test/registered/unit/spec/test_spec_registry.py

@@ -1,8 +1,10 @@
 """Unit tests for the speculative algorithm plugin registry."""
 
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
+from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_registry import (
     _REGISTRY,
@@ -207,6 +209,39 @@ class TestValidatorHook(_RegistryIsolated):
         validator.assert_not_called()
 
 
+class TestServerArgsHook(_RegistryIsolated):
+    def test_handle_speculative_decoding_invokes_custom_handle_server_args(self):
+        class CustomHandleServerArgs(CustomSpecAlgo):
+            def handle_server_args(self, server_args):
+                server_args.custom_spec_handle_seen = self.name
+                server_args.speculative_num_draft_tokens = 7
+
+        @SpeculativeAlgorithm.register(
+            "MY_HANDLE_ARGS", supports_overlap=True, spec_class=CustomHandleServerArgs
+        )
+        def _factory(server_args):
+            return MagicMock
+
+        server_args = SimpleNamespace(
+            speculative_draft_model_path=None,
+            speculative_draft_model_revision=None,
+            speculative_moe_runner_backend=None,
+            moe_runner_backend="auto",
+            speculative_algorithm="my_handle_args",
+            decrypted_draft_config_file=None,
+            trust_remote_code=False,
+            speculative_draft_window_size=None,
+            speculative_skip_dp_mlp_sync=False,
+            speculative_adaptive=False,
+        )
+
+        handle_speculative_decoding(server_args)
+
+        self.assertEqual(server_args.speculative_algorithm, "MY_HANDLE_ARGS")
+        self.assertEqual(server_args.custom_spec_handle_seen, "MY_HANDLE_ARGS")
+        self.assertEqual(server_args.speculative_num_draft_tokens, 7)
+
+
 class TestSubclassOverride(_RegistryIsolated):
     """Plugins can subclass CustomSpecAlgo to override is_*() / create_worker."""