sglang/experimental/sgl-router/tests/proxy/cache_aware_input_ids.rs

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

//! End-to-end at the HTTP layer: the router tokenizes the prompt once at
//! ingress and forwards the ids to the engine as `input_ids` (so the engine
//! skips re-tokenizing the same prompt). Asserts the gating contract through
//! the real chat handler + a MockWorker backend:
//!
//! * A plain text chat request on the engine-equivalent chat-formatter path →
//!   the forwarded body carries `input_ids` AND retains `messages`.
//! * A request carrying `tools` → `input_ids` omitted (the router's encoder
//!   doesn't render tool schemas, so its ids would diverge from the engine).
//! * A request with multimodal (array) content → `input_ids` omitted (a text
//!   tokenizer can't represent image content).

use axum::body::Body;
use axum::http::{Request, StatusCode};
use serde_json::{json, Value};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry;
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::state::kv_events::{BlockSizeOracle, HashTree};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use std::sync::Arc;
use std::time::Duration;
use tower::ServiceExt;

use crate::common::cache_aware_fixture::{config, MODEL};
use crate::common::mock_worker::MockWorker;

fn build_ctx(url: String) -> Arc<AppContext> {
    let cfg = config();
    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap());
    assert!(
        tokenizers.has_chat_formatter(MODEL),
        "deepseek-v4 model id must auto-attach the built-in chat formatter"
    );
    let registry = Arc::new(WorkerRegistry::default());
    let _ = registry.add(WorkerSpec {
        id: WorkerId(url.clone()),
        url,
        mode: WorkerMode::Plain,
        model_ids: vec![ModelId(MODEL.into())],
        bootstrap_port: None,
    });
    // Use the configured tokenizer so the chat path can emit input_ids.
    let policies =
        Arc::new(build_registry(&cfg, Arc::new(HashTree::new()), BlockSizeOracle::new()).unwrap());
    let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
    Arc::new(AppContext::new(cfg, tokenizers, proxy, registry, policies))
}

async fn send(ctx: Arc<AppContext>, body: Value) -> StatusCode {
    let app = build_router(ctx);
    let req = Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json")
        .body(Body::from(serde_json::to_vec(&body).unwrap()))
        .unwrap();
    app.oneshot(req).await.unwrap().status()
}

fn captured(mock: &MockWorker) -> Value {
    let b = mock
        .captured
        .lock()
        .unwrap()
        .last_body
        .clone()
        .expect("worker captured a request body");
    serde_json::from_slice(&b).expect("captured body is valid JSON")
}

#[tokio::test]
async fn plain_chat_forwards_input_ids_and_keeps_messages() {
    let mock = MockWorker::start(vec![]).await;
    let ctx = build_ctx(mock.url.clone());
    let status = send(
        ctx,
        json!({
            "model": MODEL,
            "messages": [{"role": "user", "content": "hello there friend"}],
        }),
    )
    .await;
    assert_eq!(status, StatusCode::OK);

    let body = captured(&mock);
    let ids = body.get("input_ids").and_then(|v| v.as_array());
    assert!(
        ids.is_some_and(|a| !a.is_empty()),
        "engine must receive non-empty input_ids; got {body}"
    );
    assert!(
        body.get("messages").is_some(),
        "messages must be retained alongside input_ids; got {body}"
    );
}

#[tokio::test]
async fn tool_request_omits_input_ids() {
    let mock = MockWorker::start(vec![]).await;
    let ctx = build_ctx(mock.url.clone());
    let status = send(
        ctx,
        json!({
            "model": MODEL,
            "messages": [{"role": "user", "content": "hi"}],
            "tools": [{"type": "function", "function": {"name": "f"}}],
        }),
    )
    .await;
    assert_eq!(status, StatusCode::OK);

    let body = captured(&mock);
    assert!(
        body.get("input_ids").is_none(),
        "tool requests must not forward input_ids; got {body}"
    );
}

#[tokio::test]
async fn thinking_request_omits_input_ids() {
    // `chat_template_kwargs` steers engine-side thinking mode, which the
    // router's encoder renders in the default mode only — forwarding ids would
    // silently run the wrong mode, so the handler must omit them.
    let mock = MockWorker::start(vec![]).await;
    let ctx = build_ctx(mock.url.clone());
    let status = send(
        ctx,
        json!({
            "model": MODEL,
            "messages": [{"role": "user", "content": "hi"}],
            "chat_template_kwargs": {"enable_thinking": true},
        }),
    )
    .await;
    assert_eq!(status, StatusCode::OK);

    let body = captured(&mock);
    assert!(
        body.get("input_ids").is_none(),
        "thinking-mode requests must not forward input_ids; got {body}"
    );
}

#[tokio::test]
async fn multimodal_request_omits_input_ids() {
    let mock = MockWorker::start(vec![]).await;
    let ctx = build_ctx(mock.url.clone());
    let status = send(
        ctx,
        json!({
            "model": MODEL,
            "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": "x"}]}],
        }),
    )
    .await;
    assert_eq!(status, StatusCode::OK);

    let body = captured(&mock);
    assert!(
        body.get("input_ids").is_none(),
        "multimodal requests must not forward input_ids; got {body}"
    );
}

/// Caller-supplied `input_ids` are never re-rendered or replaced: a flat u32
/// array (empty included) drives routing, anything else yields no routing
/// tokens, and the body reaches the engine byte-for-byte for validation.
#[tokio::test]
async fn caller_input_ids_are_used_for_routing_and_preserved() {
    let mock = MockWorker::start(vec![]).await;
    let ctx = build_ctx(mock.url.clone());
    for (ids, expected) in [
        (json!([7, 8]), Some(vec![7, 8])),
        (json!([]), Some(vec![])),
        (json!([7, -1]), None),
        (json!("bad"), None),
    ] {
        let request = json!({
            "model": MODEL,
            "messages": [{"role": "user", "content": "hi"}],
            "input_ids": ids,
        });
        let tokens = sgl_router::policies::request_tokens_for(
            &ctx.tokenizers,
            &ModelId(MODEL.into()),
            &request,
        );
        assert!(!tokens.as_ref().is_some_and(|t| t.rendered_from_chat));
        assert_eq!(tokens.map(|t| t.ids), expected, "input_ids: {ids}");
        assert_eq!(
            send(Arc::clone(&ctx), request.clone()).await,
            StatusCode::OK
        );
        assert_eq!(captured(&mock), request, "body must be forwarded untouched");
    }
    // Bypasses are not rendering failures.
    assert!(!ctx
        .metrics
        .render()
        .contains("sgl_router_ingress_tokenize_errors_total{"));
}

/// `input_ids: null` is the same as absent: the router renders and forwards.
#[tokio::test]
async fn null_input_ids_keep_normal_rendering() {
    let mock = MockWorker::start(vec![]).await;
    let ctx = build_ctx(mock.url.clone());
    let status = send(
        ctx,
        json!({
            "model": MODEL,
            "messages": [{"role": "user", "content": "hello there friend"}],
            "input_ids": null,
        }),
    )
    .await;
    assert_eq!(status, StatusCode::OK);
    let body = captured(&mock);
    assert!(
        body["input_ids"].as_array().is_some_and(|a| !a.is_empty()),
        "null input_ids must not suppress rendering; got {body}"
    );
}