sglang/experimental/sgl-router/tests/component/workers/manager.rs

// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0

use axum::{routing::get, Json, Router};
use serde_json::{json, Value};
use sgl_router::discovery::{DiscoveryEvent, ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::workers::{manager, WireProtocol, WorkerRegistry};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::TcpListener;
use tokio::sync::{mpsc, oneshot, Barrier};

/// Spin up a tiny fake worker that returns `body` on both introspection
/// endpoints. Returns the worker base URL and a shutdown channel.
///
/// One body for both because a real engine reports `served_model_name` on
/// each.
async fn spawn_fake_worker(body: Value) -> (String, oneshot::Sender<()>) {
    spawn_worker_serving(body, true).await
}

/// A worker that answers `/server_info` only — an SGLang predating
/// `served_model_name` on `/model_info`, and the shape the kind e2e fleet
/// still has (`tests/e2e/k8s_integration/fake_worker.py` defines no
/// `/model_info`). Registration must resolve the model name from the
/// `/server_info` fallback for this worker.
async fn spawn_server_info_only_worker(body: Value) -> (String, oneshot::Sender<()>) {
    spawn_worker_serving(body, false).await
}

async fn spawn_worker_serving(body: Value, with_model_info: bool) -> (String, oneshot::Sender<()>) {
    let body = Arc::new(body);
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let serve_body = move || {
        let body = body.clone();
        async move { Json((*body).clone()) }
    };
    let mut app = Router::new().route("/server_info", get(serve_body.clone()));
    if with_model_info {
        app = app.route("/model_info", get(serve_body));
    }
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    (format!("http://127.0.0.1:{port}"), tx)
}

fn spec_for(id: &str, url: &str, mode: WorkerMode) -> WorkerSpec {
    // model_ids are intentionally empty: the manager resolves them via
    // /server_info introspection.  Pre-populating here would lie about
    // what discovery backends actually emit.
    WorkerSpec {
        id: WorkerId(id.into()),
        url: url.into(),
        mode,
        model_ids: Vec::new(),
        bootstrap_port: None,
    }
}

async fn wait_until(condition: impl Fn() -> bool, description: &str) {
    tokio::time::timeout(Duration::from_secs(2), async {
        loop {
            if condition() {
                return;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .unwrap_or_else(|_| panic!("timed out waiting for {description}"));
}

#[tokio::test]
async fn manager_processes_added_then_removed() {
    let (url_a, _s_a) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
    let (url_b, _s_b) = spawn_fake_worker(json!({"served_model_name": "m"})).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url_a,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();
    tx.send(DiscoveryEvent::Added(spec_for(
        "w2",
        &url_b,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();

    wait_until(
        || registry.workers_for(&ModelId("m".into())).len() == 2,
        "both workers to register",
    )
    .await;
    assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 2);

    tx.send(DiscoveryEvent::Removed {
        id: WorkerId("w1".into()),
    })
    .await
    .unwrap();
    wait_until(
        || registry.workers_for(&ModelId("m".into())).len() == 1,
        "removed worker to leave the registry",
    )
    .await;
    assert_eq!(registry.workers_for(&ModelId("m".into())).len(), 1);

    drop(tx);
    h.await.unwrap();
}

#[tokio::test]
async fn manager_handles_mode_changed() {
    let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url,
        WorkerMode::Prefill,
    )))
    .await
    .unwrap();
    wait_until(
        || {
            registry
                .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
                .len()
                == 1
        },
        "prefill worker to register",
    )
    .await;
    assert_eq!(
        registry
            .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
            .len(),
        1
    );

    tx.send(DiscoveryEvent::ModeChanged {
        id: WorkerId("w1".into()),
        mode: WorkerMode::Decode,
    })
    .await
    .unwrap();
    wait_until(
        || {
            registry
                .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
                .is_empty()
                && registry
                    .workers_for_mode(&ModelId("m".into()), WorkerMode::Decode)
                    .len()
                    == 1
        },
        "worker mode to change to decode",
    )
    .await;
    assert_eq!(
        registry
            .workers_for_mode(&ModelId("m".into()), WorkerMode::Prefill)
            .len(),
        0
    );
    assert_eq!(
        registry
            .workers_for_mode(&ModelId("m".into()), WorkerMode::Decode)
            .len(),
        1
    );

    drop(tx);
    h.await.unwrap();
}

/// A worker that reports `enable_http2: true` is registered with
/// [`WireProtocol::H2c`], so the proxy forwards to it over cleartext h2c.
#[tokio::test]
async fn manager_resolves_h2c_protocol_from_introspection() {
    let (url, _s) =
        spawn_fake_worker(json!({"served_model_name": "m", "enable_http2": true})).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();

    let resolved = wait_for_protocol(&registry, "w1", WireProtocol::H2c).await;
    assert!(
        resolved,
        "worker reporting enable_http2 must resolve to h2c, got {:?}",
        registry.get(&WorkerId("w1".into())).map(|w| w.protocol()),
    );

    drop(tx);
    h.await.unwrap();
}

/// A worker that omits `enable_http2` (older SGLang) keeps the safe HTTP/1.1
/// default on its registry entry.
#[tokio::test]
async fn manager_defaults_http1_when_enable_http2_absent() {
    let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();

    // Non-empty `model_ids` is what proves introspection ran: `spawn_fake_worker`
    // answers from one body, so a resolved model id means `/server_info`
    // answered too and its silence on `enable_http2` is the engine's, not the
    // fixture's. The positive direction — that `enable_http2` is actually read
    // — is pinned by `manager_resolves_protocol_per_worker_no_fleet_lock`
    // below.
    let registered = wait_for(Duration::from_secs(2), || {
        registry
            .get(&WorkerId("w1".into()))
            .is_some_and(|w| !w.model_ids.is_empty())
    })
    .await;
    assert!(registered, "worker registered");
    let w = registry.get(&WorkerId("w1".into())).unwrap();
    assert_eq!(w.protocol(), WireProtocol::Http1);

    drop(tx);
    h.await.unwrap();
}

/// A worker with no `/model_info` at all still registers, resolving its model
/// name from the `/server_info` fallback and its protocol from the same
/// response.
///
/// Every other test here uses a fixture that answers both endpoints, so without
/// this one the fallback chain in `introspect::fetch` has no component-level
/// coverage — while the kind e2e fleet runs exactly this shape.
#[tokio::test]
async fn manager_registers_worker_that_serves_server_info_only() {
    let (url, _s) =
        spawn_server_info_only_worker(json!({"served_model_name": "m", "enable_http2": true}))
            .await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w-no-model-info",
        &url,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();

    assert!(
        wait_for_protocol(&registry, "w-no-model-info", WireProtocol::H2c).await,
        "a worker without /model_info must still resolve its model name and \
         protocol from /server_info; got {:?}",
        registry
            .get(&WorkerId("w-no-model-info".into()))
            .map(|w| (w.model_ids.clone(), w.protocol())),
    );

    drop(tx);
    h.await.unwrap();
}

/// The load-bearing regression for the production h2c bug: per-worker protocol
/// means one worker that resolved HTTP/1.1 first does NOT lock the rest of the
/// fleet off h2c. A mixed fleet (one h2c-capable worker registered AFTER a
/// plain HTTP/1.1 worker) ends with each worker on its own protocol — the old
/// single-client first-write-wins design forced both to HTTP/1.1.
#[tokio::test]
async fn manager_resolves_protocol_per_worker_no_fleet_lock() {
    // w-http1 registers first and reports no enable_http2 (resolves Http1);
    // w-h2c registers second and reports enable_http2: true (must still get
    // h2c — it is not dragged down by w-http1's earlier Http1 resolution).
    let (url_http1, _s1) = spawn_fake_worker(json!({"served_model_name": "m"})).await;
    let (url_h2c, _s2) =
        spawn_fake_worker(json!({"served_model_name": "m", "enable_http2": true})).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w-http1",
        &url_http1,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();
    // Let the HTTP/1.1 worker resolve first so it is the one that would have
    // "won" the old fleet-wide client.
    assert!(
        wait_for_protocol(&registry, "w-http1", WireProtocol::Http1).await,
        "first worker should resolve Http1",
    );

    tx.send(DiscoveryEvent::Added(spec_for(
        "w-h2c",
        &url_h2c,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();

    assert!(
        wait_for_protocol(&registry, "w-h2c", WireProtocol::H2c).await,
        "an h2c-capable worker registered after an Http1 worker must still resolve h2c \
         (per-worker protocol, no fleet-wide lock); got {:?}",
        registry
            .get(&WorkerId("w-h2c".into()))
            .map(|w| w.protocol()),
    );
    // The earlier worker is untouched.
    let first = registry.get(&WorkerId("w-http1".into())).unwrap();
    assert_eq!(
        first.protocol(),
        WireProtocol::Http1,
        "the Http1 worker must stay Http1",
    );

    drop(tx);
    h.await.unwrap();
}

/// Block until `worker_id`'s registry entry reports `expected`, or 2 s elapse.
/// Returns whether it converged.
///
/// The `model_ids` check is half the predicate on purpose. `H2c` proves itself
/// — it is reachable only through a successful `/server_info` read — but
/// `Http1` is also the default for a worker built without a resolved protocol,
/// so waiting on that value alone would be satisfied by an implementation that
/// never reads one. Requiring a resolved model id as well means the worker at
/// least completed introspection against a fixture that answers both
/// endpoints.
async fn wait_for_protocol(
    registry: &Arc<WorkerRegistry>,
    worker_id: &str,
    expected: WireProtocol,
) -> bool {
    let id = WorkerId(worker_id.into());
    wait_for(Duration::from_secs(2), || {
        registry
            .get(&id)
            .is_some_and(|w| !w.model_ids.is_empty() && w.protocol() == expected)
    })
    .await
}

/// Poll `cond` every 20 ms until it returns true or `budget` elapses.
async fn wait_for(budget: Duration, mut cond: impl FnMut() -> bool) -> bool {
    let deadline = Instant::now() + budget;
    while Instant::now() < deadline {
        if cond() {
            return true;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    cond()
}

#[tokio::test]
async fn mode_changed_preserves_active_requests_and_breaker() {
    let (url, _s) = spawn_fake_worker(json!({"served_model_name": "m"})).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url,
        WorkerMode::Prefill,
    )))
    .await
    .unwrap();
    tokio::time::sleep(Duration::from_millis(200)).await;

    // Grab a handle, bump active_requests, and open the breaker.
    let w = registry.get(&WorkerId("w1".into())).unwrap();
    w.active_requests.fetch_add(5, Ordering::Relaxed);
    // Default threshold is 3 — record 10 failures to guarantee Open state.
    for _ in 0..10 {
        w.breaker.record_failure();
    }
    let breaker_open_before = !w.breaker.allow();
    assert!(
        breaker_open_before,
        "breaker should be open after 10 failures"
    );

    // Flip mode via ModeChanged.
    tx.send(DiscoveryEvent::ModeChanged {
        id: WorkerId("w1".into()),
        mode: WorkerMode::Decode,
    })
    .await
    .unwrap();
    tokio::time::sleep(Duration::from_millis(50)).await;

    // Re-fetch the Worker handle from the registry.
    let w_after = registry.get(&WorkerId("w1".into())).unwrap();

    assert_eq!(
        w_after.mode(),
        WorkerMode::Decode,
        "mode should have flipped to Decode"
    );
    assert_eq!(
        w_after.active_requests.load(Ordering::Relaxed),
        5,
        "active_requests should be preserved across mode change"
    );
    assert!(
        !w_after.breaker.allow(),
        "breaker open state should be preserved across mode change"
    );

    // Critical: the Arc identity must be the same — mutation in place.
    assert!(
        Arc::ptr_eq(&w, &w_after),
        "Worker handle should be the SAME Arc, not a fresh replacement"
    );

    drop(tx);
    h.await.unwrap();
}

/// An out-of-order `ModeChanged` for a worker the registry does not know
/// about (e.g. a buggy discovery backend reordered `Removed` and
/// `ModeChanged`) must not panic, must not silently log INFO claiming the
/// mode flip happened, and must leave the registry untouched.
#[tokio::test]
async fn manager_handles_orphan_mode_changed_without_panic() {
    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::ModeChanged {
        id: WorkerId("ghost".into()),
        mode: WorkerMode::Decode,
    })
    .await
    .unwrap();
    tokio::time::sleep(Duration::from_millis(50)).await;

    assert!(
        registry.get(&WorkerId("ghost".into())).is_none(),
        "an orphan ModeChanged must not create a phantom worker",
    );
    assert_eq!(
        registry.workers_for(&ModelId("m".into())).len(),
        0,
        "registry must be empty after an orphan event",
    );

    drop(tx);
    h.await.unwrap();
}

/// A `Removed` for an unknown id is a no-op — registry stays empty, manager
/// keeps running.
#[tokio::test]
async fn manager_handles_orphan_removed_without_panic() {
    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Removed {
        id: WorkerId("ghost".into()),
    })
    .await
    .unwrap();
    tokio::time::sleep(Duration::from_millis(50)).await;
    assert!(registry.is_empty());

    drop(tx);
    h.await.unwrap();
}

/// Duplicate `Added` for the same id is an upsert — the registry ends up
/// with exactly one worker. The model resolved by /server_info wins on
/// re-add (a different worker may advertise a different served model).
#[tokio::test]
async fn manager_handles_duplicate_added_as_upsert() {
    let (url_first, _s_first) = spawn_fake_worker(json!({"served_model_name": "m1"})).await;
    let (url_second, _s_second) = spawn_fake_worker(json!({"served_model_name": "m1"})).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url_first,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();
    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url_second,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();
    wait_until(
        || {
            registry
                .get(&WorkerId("w1".into()))
                .is_some_and(|worker| worker.url == url_second)
        },
        "replacement Added event to update the worker",
    )
    .await;

    assert_eq!(
        registry.workers_for(&ModelId("m1".into())).len(),
        1,
        "w1 still serves m1 after the second Added",
    );

    drop(tx);
    h.await.unwrap();
}

/// Spawn a fake worker whose `/server_info` returns `body` only after
/// sleeping for `delay`. Returns the worker URL and a shutdown channel.
async fn spawn_slow_worker(body: Value, delay: Duration) -> (String, oneshot::Sender<()>) {
    let body = Arc::new(body);
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let app = Router::new().route(
        "/server_info",
        get(move || {
            let body = body.clone();
            async move {
                tokio::time::sleep(delay).await;
                Json((*body).clone())
            }
        }),
    );
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    (format!("http://127.0.0.1:{port}"), tx)
}

async fn spawn_gated_worker(body: Value, gate: Arc<Barrier>) -> (String, oneshot::Sender<()>) {
    let body = Arc::new(body);
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let app = Router::new().route(
        "/server_info",
        get(move || {
            let body = body.clone();
            let gate = Arc::clone(&gate);
            async move {
                gate.wait().await;
                Json((*body).clone())
            }
        }),
    );
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    (format!("http://127.0.0.1:{port}"), tx)
}

/// Spawn a fake worker that counts each `GET /server_info` hit in the
/// returned `AtomicUsize`.  Used to assert the manager makes exactly
/// one round-trip per worker.
async fn spawn_counting_worker(body: Value) -> (String, Arc<AtomicUsize>, oneshot::Sender<()>) {
    let body = Arc::new(body);
    let counter = Arc::new(AtomicUsize::new(0));
    let counter_clone = counter.clone();
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let app = Router::new().route(
        "/server_info",
        get(move || {
            let body = body.clone();
            let counter = counter_clone.clone();
            async move {
                counter.fetch_add(1, Ordering::SeqCst);
                Json((*body).clone())
            }
        }),
    );
    let (tx, rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = rx.await;
            })
            .await;
    });
    (format!("http://127.0.0.1:{port}"), counter, tx)
}

/// Registration must start every `/server_info` fetch before any response is
/// released. A sequential manager stalls at the first worker's barrier.
#[tokio::test]
async fn added_events_run_in_parallel() {
    let n = 5;
    let gate = Arc::new(Barrier::new(n + 1));
    let mut workers = Vec::new();
    for _ in 0..n {
        workers
            .push(spawn_gated_worker(json!({"served_model_name": "m"}), Arc::clone(&gate)).await);
    }

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    for (i, (url, _s)) in workers.iter().enumerate() {
        tx.send(DiscoveryEvent::Added(spec_for(
            &format!("w{i}"),
            url,
            WorkerMode::Plain,
        )))
        .await
        .unwrap();
    }
    assert!(
        tokio::time::timeout(Duration::from_secs(2), gate.wait())
            .await
            .is_ok(),
        "manager did not start all {n} /server_info requests concurrently"
    );
    let registered = tokio::time::timeout(Duration::from_secs(2), async {
        loop {
            if registry.workers_for(&ModelId("m".into())).len() == n {
                return;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await;
    assert!(registered.is_ok(), "manager failed to register {n} workers");

    drop(tx);
    h.await.unwrap();
}

/// A `Removed` issued while the matching `Added` is still mid-fetch
/// must await the in-flight registration handle before removing.
/// Without that ordering the removal runs first (registry has nothing
/// to remove), then the Added's deferred registry write leaks the
/// worker.
#[tokio::test]
async fn removed_awaits_pending_added() {
    let (url, _s) = spawn_slow_worker(
        json!({"served_model_name": "m"}),
        Duration::from_millis(300),
    )
    .await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run(rx, registry.clone()));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w-slow",
        &url,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();
    tx.send(DiscoveryEvent::Removed {
        id: WorkerId("w-slow".into()),
    })
    .await
    .unwrap();

    // Wait long enough for the Added's /server_info to complete (300ms),
    // then assert the worker is gone. If Removed ran before Added's
    // registry write, the post-fetch write would leak the entry.
    tokio::time::sleep(Duration::from_millis(600)).await;
    assert!(
        registry.get(&WorkerId("w-slow".into())).is_none(),
        "Removed must await the in-flight Added; otherwise the deferred \
         registry write leaks the worker"
    );

    drop(tx);
    h.await.unwrap();
}

/// The manager must make exactly ONE `/server_info` request per worker.
/// Before this fix the worker manager fetched `served_model_name` and
/// `KvEventIndex::add_worker` fetched the `kv_events` block
/// independently — 2N round-trips for N workers.
#[tokio::test]
async fn manager_emits_single_server_info_fetch_per_worker() {
    use sgl_router::state::kv_events::KvEventIndex;

    let body = json!({
        "served_model_name": "m",
        "kv_events": {
            "publisher": "zmq",
            "endpoint_host": "127.0.0.1",
            "endpoint_port_base": 60100,
            "topic": "",
            "block_size": 64,
            "dp_size": 1,
        }
    });
    let (url, counter, _s) = spawn_counting_worker(body).await;

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let kv_index = KvEventIndex::new();
    let h = tokio::spawn(manager::run_with_config(
        rx,
        registry.clone(),
        None,
        Some(kv_index.clone()),
        None,
    ));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w1",
        &url,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();

    // Wait for both the registry and kv-events index to reflect the worker.
    let ready = tokio::time::timeout(Duration::from_secs(2), async {
        loop {
            if registry.get(&WorkerId("w1".into())).is_some() && kv_index.known_worker_count() == 1
            {
                return true;
            }
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
    })
    .await;
    assert!(
        ready.is_ok(),
        "manager did not finish onboarding the worker"
    );

    let hits = counter.load(Ordering::SeqCst);
    assert_eq!(
        hits, 1,
        "manager must fetch /server_info exactly once per worker (got {hits})"
    );

    drop(tx);
    h.await.unwrap();
    kv_index.shutdown().await;
}

/// Public-API integration for the reconcile loop: a worker that registers
/// with empty `model_ids` (its `/server_info` was failing when discovery
/// first reported it — the "EndpointSlice ready before the engine can
/// answer /server_info" race) must be recovered by the manager's
/// reconcile loop once `/server_info` starts answering. Exercised through
/// the public `run_with_introspector_and_reconcile` entrypoint with no
/// new discovery event after the initial `Added`.
#[tokio::test]
async fn reconcile_recovers_worker_with_unresolved_model_ids() {
    use axum::http::StatusCode;
    use axum::response::IntoResponse;
    use sgl_router::workers::WorkerIntrospector;
    use std::sync::atomic::AtomicBool;

    // Switchable fake engine: 503 on /server_info until `ready` flips
    // true, then serves a body advertising model "m".
    let ready = Arc::new(AtomicBool::new(false));
    let ready_handler = ready.clone();
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let app = Router::new().route(
        "/server_info",
        get(move || {
            let ready = ready_handler.clone();
            async move {
                if ready.load(Ordering::SeqCst) {
                    Json(json!({"served_model_name": "m"})).into_response()
                } else {
                    StatusCode::SERVICE_UNAVAILABLE.into_response()
                }
            }
        }),
    );
    let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
    tokio::spawn(async move {
        let _ = axum::serve(listener, app)
            .with_graceful_shutdown(async move {
                let _ = shutdown_rx.await;
            })
            .await;
    });
    let url = format!("http://127.0.0.1:{port}");

    let (tx, rx) = mpsc::channel(16);
    let registry = Arc::new(WorkerRegistry::default());
    let h = tokio::spawn(manager::run_with_introspector_and_reconcile(
        rx,
        registry.clone(),
        None,
        None,
        None,
        Arc::new(WorkerIntrospector::new(Duration::from_millis(300))),
        Duration::from_millis(150),
    ));

    tx.send(DiscoveryEvent::Added(spec_for(
        "w-slow",
        &url,
        WorkerMode::Plain,
    )))
    .await
    .unwrap();

    // Phase 1: the worker registers but stays out of the model pool while
    // /server_info keeps failing.
    let stuck = tokio::time::timeout(Duration::from_secs(2), async {
        loop {
            if let Some(w) = registry.get(&WorkerId("w-slow".into())) {
                if w.model_ids.is_empty() && registry.workers_for(&ModelId("m".into())).is_empty() {
                    return true;
                }
            }
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
    })
    .await;
    assert!(
        stuck.is_ok(),
        "worker should register with empty model_ids (invisible to routing) while /server_info fails",
    );

    // Engine finishes coming up.
    ready.store(true, Ordering::SeqCst);

    // Phase 2: the reconcile loop re-introspects and the worker joins the
    // model pool — no new discovery event was sent.
    let recovered = tokio::time::timeout(Duration::from_secs(3), async {
        loop {
            if registry.workers_for(&ModelId("m".into())).len() == 1 {
                return true;
            }
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
    })
    .await;
    assert!(
        recovered.is_ok(),
        "reconcile loop must add the worker to the model pool once /server_info recovers",
    );

    drop(tx);
    h.await.unwrap();
    let _ = shutdown_tx.send(());
}