feat(rl): version table, response stamping, and state-aware routing (M2)

bindings/python/src/lib.rs

@@ -534,6 +534,7 @@ struct Router {
     enable_rl: bool,
     rl_control_timeout_secs: u64,
     rl_fanout_concurrency: usize,
+    rl_version_policy: String,
 }
 
 impl Router {
@@ -832,6 +833,17 @@ impl Router {
                 None
             };
 
+        let rl_version_policy: smg_rl::VersionPolicy =
+            self.rl_version_policy
+                .parse()
+                .map_err(
+                    |e: smg_rl::VersionPolicyError| config::ConfigError::InvalidValue {
+                        field: "rl-version-policy".to_string(),
+                        value: self.rl_version_policy.clone(),
+                        reason: e.reason.to_string(),
+                    },
+                )?;
+
         config::RouterConfig::builder()
             .mode(mode)
             .policy(policy)
@@ -944,6 +956,7 @@ impl Router {
                 enabled: self.enable_rl,
                 control_timeout_secs: self.rl_control_timeout_secs,
                 fanout_concurrency: self.rl_fanout_concurrency,
+                version_policy: rl_version_policy,
             })
             .maybe_client_cert_and_key(
                 self.client_cert_path.as_ref(),
@@ -1117,6 +1130,7 @@ impl Router {
         enable_rl = false,
         rl_control_timeout_secs = 600,
         rl_fanout_concurrency = 32,
+        rl_version_policy = String::from("any"),
     ))]
     #[expect(clippy::too_many_arguments)]
     #[expect(
@@ -1276,6 +1290,7 @@ impl Router {
         enable_rl: bool,
         rl_control_timeout_secs: u64,
         rl_fanout_concurrency: usize,
+        rl_version_policy: String,
     ) -> PyResult<Self> {
         let mut all_urls = worker_urls.clone();
 
@@ -1447,6 +1462,7 @@ impl Router {
             enable_rl,
             rl_control_timeout_secs,
             rl_fanout_concurrency,
+            rl_version_policy,
         })
     }
 

bindings/python/src/smg/rl.py

@@ -11,6 +11,10 @@
 Only HTTP workers can be proxied. A gRPC or ZMQ worker matched by a selector is
 reported in `failed[]` as `unsupported_connection_mode`, which makes `fanout`
 raise `FanoutError` unless `allow_partial=True`.
+
+`set_fleet_version` (like `set_version`) records a version SMG did not observe
+itself; it is a local table write, never an engine call. Every response SMG
+serves for a versioned engine carries `x-smg-weight-version`.
 """
 
 from __future__ import annotations
@@ -54,12 +58,15 @@ class Worker:
     weight_version: str | None
     labels: dict[str, str] = field(default_factory=dict)
     capabilities: dict[str, Any] = field(default_factory=dict)
+    control: str = "active"
+    version_source: str | None = None
 
     @classmethod
     def from_json(cls, d: dict[str, Any]) -> Worker:
         kwargs = {k: d.get(k) for k in cls.__dataclass_fields__}
         kwargs["labels"] = d.get("labels") or {}
         kwargs["capabilities"] = d.get("capabilities") or {}
+        kwargs["control"] = d.get("control") or "active"
         return cls(**kwargs)  # type: ignore[arg-type]
 
 
@@ -177,6 +184,39 @@ def worker(self, worker_id: str) -> Worker:
         _raise_for(status, payload)
         return Worker.from_json(payload)
 
+    def set_version(self, worker_id: str, weight_version: str) -> Worker:
+        """Record the weight version one engine holds (SMG-local; no engine call)."""
+        status, payload = self._request(
+            "POST", f"/v1/rl/workers/{worker_id}/version", {"weight_version": str(weight_version)}
+        )
+        _raise_for(status, payload)
+        return Worker.from_json(payload)
+
+    def set_fleet_version(self, weight_version: str, *, selector: str) -> list[Worker]:
+        status, payload = self._request(
+            "POST",
+            "/v1/rl/version",
+            {"weight_version": str(weight_version)},
+            {"selector": selector},
+        )
+        _raise_for(status, payload)
+        return [Worker.from_json(w) for w in payload.get("workers", [])]
+
+    def set_state(self, worker_id: str, control: str) -> Worker:
+        """Mark one engine active, paused, or asleep for routing (SMG-local)."""
+        status, payload = self._request(
+            "POST", f"/v1/rl/workers/{worker_id}/state", {"control": control}
+        )
+        _raise_for(status, payload)
+        return Worker.from_json(payload)
+
+    def set_fleet_state(self, control: str, *, selector: str) -> list[Worker]:
+        status, payload = self._request(
+            "POST", "/v1/rl/state", {"control": control}, {"selector": selector}
+        )
+        _raise_for(status, payload)
+        return [Worker.from_json(w) for w in payload.get("workers", [])]
+
     def call(
         self,
         worker_id: str,

bindings/python/src/smg/router_args.py

@@ -257,6 +257,7 @@ class RouterArgs:
     enable_rl: bool = False  # Mount the RL control plane under /v1/rl
     rl_control_timeout_secs: int = 600  # Timeout for one proxied engine control call
     rl_fanout_concurrency: int = 32  # Max concurrent engine calls in one fan-out
+    rl_version_policy: str = "any"  # any | latest-only | min-version:<v> | max-staleness:<k>
 
     @staticmethod
     def add_cli_args(
@@ -813,6 +814,16 @@ def add_cli_args(
             default=RouterArgs.rl_fanout_concurrency,
             help="Maximum concurrent engine calls in one fan-out (default: 32)",
         )
+        rl_group.add_argument(
+            f"--{prefix}rl-version-policy",
+            type=str,
+            default=RouterArgs.rl_version_policy,
+            help=(
+                "Which engines a request may be routed to, judged against the model's newest "
+                "known weight version: any | latest-only | min-version:<v> | max-staleness:<k> "
+                "(default: any). Per-request override: the x-smg-version-policy header."
+            ),
+        )
 
         # PD/EPD-specific arguments
         pd_group.add_argument(

bindings/python/tests/test_arg_parser.py

@@ -50,6 +50,7 @@ def test_default_values(self):
         assert args.enable_rl is False
         assert args.rl_control_timeout_secs == 600
         assert args.rl_fanout_concurrency == 32
+        assert args.rl_version_policy == "any"
 
     def test_rl_flags_parse(self):
         """RL control-plane flags land on the dataclass."""
@@ -63,12 +64,15 @@ def test_rl_flags_parse(self):
                     "1200",
                     "--rl-fanout-concurrency",
                     "8",
+                    "--rl-version-policy",
+                    "max-staleness:2",
                 ]
             )
         )
         assert args.enable_rl is True
         assert args.rl_control_timeout_secs == 1200
         assert args.rl_fanout_concurrency == 8
+        assert args.rl_version_policy == "max-staleness:2"
 
     def test_parse_selector_valid(self):
         """Test parsing valid selector arguments."""
@@ -1468,6 +1472,7 @@ class TestRouterArgsFieldOrder:
         "enable_rl",
         "rl_control_timeout_secs",
         "rl_fanout_concurrency",
+        "rl_version_policy",
     ]
 
     def test_complete_field_sequence_is_frozen(self):
@@ -1506,6 +1511,7 @@ def test_new_fields_appended_after_positional_reserve(self):
             "enable_rl",
             "rl_control_timeout_secs",
             "rl_fanout_concurrency",
+            "rl_version_policy",
         ):
             assert names.index(appended) > marker, (
                 f"{appended} must be appended after worker_startup_delay to "

bindings/python/tests/test_rl_client.py

@@ -26,6 +26,8 @@
     "role": None,
     "health": "ready",
     "weight_version": "7",
+    "version_source": "api",
+    "control": "active",
     "labels": {"tp_size": "1"},
     "capabilities": {
         "source": "static",
@@ -93,6 +95,32 @@ def test_workers_and_worker(stub):
     assert _Stub.seen[0]["auth"] == "Bearer k"
 
 
+def test_set_version_and_state_post_bodies(stub):
+    _Stub.stub_responses["/v1/rl/workers/w1/version"] = (200, WORKER)
+    rl = RL(stub)
+    w = rl.set_version("w1", "42")
+    assert _Stub.seen[-1]["body"] == {"weight_version": "42"}
+    assert w.control == "active" and w.version_source == "api"
+
+
+def test_fleet_forms_send_the_selector(stub):
+    _Stub.stub_responses["/v1/rl/state"] = (
+        200,
+        {"protocol_version": 1, "workers": [WORKER], "total": 1},
+    )
+    rl = RL(stub)
+    workers = rl.set_fleet_state("paused", selector="engine=sglang")
+    assert _Stub.seen[-1]["query"] == {"selector": ["engine=sglang"]}
+    assert len(workers) == 1 and workers[0].id == "w1"
+
+
+def test_worker_from_json_defaults_control():
+    d = dict(WORKER)
+    del d["control"]
+    w = Worker.from_json(d)
+    assert w.control == "active"
+
+
 def test_call_sends_method_params_body(stub):
     _Stub.stub_responses["/v1/rl/workers/w1/engine/pause"] = (
         200,

clients/openapi-gen/src/main.rs

@@ -639,6 +639,8 @@ fn main() -> anyhow::Result<()> {
     let rl_entry_name = collect_schema(schema_for!(RlWorkerEntry), &mut schemas)?;
     let rl_outcome_name = collect_schema(schema_for!(RlCallOutcome), &mut schemas)?;
     let rl_fanout_name = collect_schema(schema_for!(RlFanoutResponse), &mut schemas)?;
+    let rl_set_version_name = collect_schema(schema_for!(RlSetVersionRequest), &mut schemas)?;
+    let rl_set_control_name = collect_schema(schema_for!(RlSetControlRequest), &mut schemas)?;
     let json_content = |name: &str| -> BTreeMap<String, MediaType> {
         let mut content = BTreeMap::new();
         content.insert(
@@ -748,6 +750,71 @@ fn main() -> anyhow::Result<()> {
         },
     );
 
+    let json_body = |name: &str| RequestBody {
+        required: true,
+        content: json_content(name),
+    };
+    let selector_param = || Parameter {
+        name: "selector".to_string(),
+        location: "query".to_string(),
+        required: true,
+        schema: ParameterSchema {
+            schema_type: "string".to_string(),
+        },
+    };
+    paths.insert(
+        "/v1/rl/workers/{worker_id}/version".to_string(),
+        PathItem {
+            post: Some(operation(
+                "setRlWorkerVersion",
+                "Record the weight version one engine holds (no engine call)",
+                Some(vec![path_param("worker_id")]),
+                Some(json_body(&rl_set_version_name)),
+                json_response(&rl_entry_name, "The updated worker"),
+            )),
+            ..PathItem::default()
+        },
+    );
+    paths.insert(
+        "/v1/rl/workers/{worker_id}/state".to_string(),
+        PathItem {
+            post: Some(operation(
+                "setRlWorkerState",
+                "Record whether one engine is active, paused, or asleep (no engine call)",
+                Some(vec![path_param("worker_id")]),
+                Some(json_body(&rl_set_control_name)),
+                json_response(&rl_entry_name, "The updated worker"),
+            )),
+            ..PathItem::default()
+        },
+    );
+    paths.insert(
+        "/v1/rl/version".to_string(),
+        PathItem {
+            post: Some(operation(
+                "setRlFleetVersion",
+                "Record the weight version every worker matching `selector` holds",
+                Some(vec![selector_param()]),
+                Some(json_body(&rl_set_version_name)),
+                json_response(&rl_workers_name, "The updated workers"),
+            )),
+            ..PathItem::default()
+        },
+    );
+    paths.insert(
+        "/v1/rl/state".to_string(),
+        PathItem {
+            post: Some(operation(
+                "setRlFleetState",
+                "Record the control state of every worker matching `selector`",
+                Some(vec![selector_param()]),
+                Some(json_body(&rl_set_control_name)),
+                json_response(&rl_workers_name, "The updated workers"),
+            )),
+            ..PathItem::default()
+        },
+    );
+
     // ---- Generate (SGLang native) ----
     use openai_protocol::generate::*;
     let req_name = collect_schema(schema_for!(GenerateRequest), &mut schemas)?;

crates/external_router/src/header_utils.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 use axum::{
     body::Body,
@@ -138,12 +138,33 @@ pub fn preserve_response_headers(reqwest_headers: &HeaderMap) -> HeaderMap {
 /// [`preserve_response_headers`] so the gateway's value wins over anything an
 /// upstream set under the same name. Skipped if the URL isn't representable as
 /// a header value.
+///
+/// This is the header half of [`stamp_routed_worker`], which is what callers
+/// with a whole response in hand should use: it sets this header and the
+/// [`RoutedWorker`] extension together, and layers keyed on the extension
+/// (the RL version stamp among them) see nothing when only the header is set.
 pub fn insert_routed_worker_id(headers: &mut HeaderMap, worker_url: &str) {
     if let Ok(value) = HeaderValue::from_str(worker_url) {
         headers.insert(HEADER_ROUTED_WORKER_ID.clone(), value);
     }
 }
 
+/// Response extension naming the worker that served the request. Set only by
+/// the gateway, never derived from an upstream header, so a layer keyed on it
+/// cannot be fooled through [`preserve_response_headers`] (the same reason
+/// the error code travels as an extension).
+#[derive(Clone, Debug)]
+pub struct RoutedWorker(pub Arc<str>);
+
+/// Stamp `x-smg-routed-worker-id` and the [`RoutedWorker`] extension on a
+/// response. Call after the headers are in place so the gateway's value wins.
+pub fn stamp_routed_worker<B>(response: &mut http::Response<B>, worker_url: &str) {
+    insert_routed_worker_id(response.headers_mut(), worker_url);
+    response
+        .extensions_mut()
+        .insert(RoutedWorker(Arc::from(worker_url)));
+}
+
 /// Whether `name` is SMG's own error-code response header.
 ///
 /// Backend responses must not mint this marker: downstream gateways use it to
@@ -409,6 +430,23 @@ pub fn should_forward_request_header(name: &str) -> bool {
 mod tests {
     use super::*;
 
+    #[test]
+    fn stamp_routed_worker_sets_header_and_extension() {
+        let mut response = http::Response::new(());
+        response
+            .headers_mut()
+            .insert("x-smg-routed-worker-id", "http://forged".parse().unwrap());
+        stamp_routed_worker(&mut response, "http://w1:1@2");
+        assert_eq!(
+            response.headers()["x-smg-routed-worker-id"],
+            "http://w1:1@2"
+        );
+        assert_eq!(
+            &*response.extensions().get::<RoutedWorker>().unwrap().0,
+            "http://w1:1@2"
+        );
+    }
+
     #[test]
     fn preserve_response_headers_drops_smg_owned_error_code() {
         let mut upstream = HeaderMap::new();