feat(grpc): serve /v1/responses in PD mode

e2e_test/router/test_pd_responses.py

@@ -0,0 +1,103 @@
+"""Responses API tests for PD (Prefill-Decode) disaggregated gRPC routing.
+
+/v1/responses rides the same mode-parameterized gRPC pipeline as chat
+completions, so every create below exercises prefill/decode pair selection,
+bootstrap injection, and dual dispatch.
+
+Backends:
+- "pd_grpc": gRPC mode (both SGLang and vLLM)
+
+Requirements:
+    - SGLang: sgl_kernel package
+    - vLLM: NIXL or Mooncake KV transfer support
+    - GPUs: num_prefill + num_decode (default: 2 GPUs for 1+1)
+
+Usage:
+    pytest e2e_test/router/test_pd_responses.py -v
+
+    # vLLM
+    E2E_RUNTIME=vllm pytest e2e_test/router/test_pd_responses.py -v
+"""
+
+from __future__ import annotations
+
+import logging
+
+import openai
+import pytest
+import smg_client
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.mark.engine("sglang", "vllm")
+@pytest.mark.gpu(2)
+@pytest.mark.model("meta-llama/Llama-3.1-8B-Instruct")
+@pytest.mark.e2e
+@pytest.mark.gateway(extra_args=["--history-backend", "memory"])
+@pytest.mark.parametrize("setup_backend", ["pd_grpc"], indirect=True)
+@pytest.mark.parametrize("api_client", ["openai", "smg"], indirect=True)
+class TestPDResponsesGrpc:
+    """Responses API tests using PD disaggregation (gRPC mode)."""
+
+    def test_basic_response_creation(self, model, api_client):
+        """Test basic response creation."""
+        resp = api_client.responses.create(model=model, input="What is 2+2?")
+
+        assert resp.id is not None
+        assert resp.error is None
+        assert resp.status == "completed"
+        assert len(resp.output_text) > 0
+        assert resp.usage is not None
+
+    def test_streaming_response(self, model, api_client):
+        """Test streaming response."""
+        resp = api_client.responses.create(
+            model=model, input="Count to 5", stream=True, max_output_tokens=50
+        )
+
+        events = list(resp)
+        created_events = [e for e in events if e.type == "response.created"]
+        assert len(created_events) > 0
+
+        delta_events = [e for e in events if e.type == "response.output_text.delta"]
+        assert len(delta_events) > 0
+
+        completed_events = [e for e in events if e.type == "response.completed"]
+        assert len(completed_events) == 1
+
+    def test_previous_response_id_chaining(self, model, api_client):
+        """Test chaining responses using previous_response_id."""
+        # First response
+        resp1 = api_client.responses.create(
+            model=model, input="My name is Alice and my friend is Bob. Remember it."
+        )
+        assert resp1.error is None
+        assert resp1.status == "completed"
+
+        # Second response referencing first
+        resp2 = api_client.responses.create(
+            model=model, input="What is my name", previous_response_id=resp1.id
+        )
+        assert resp2.error is None
+        assert resp2.status == "completed"
+        assert "Alice" in resp2.output_text
+
+        # Third response referencing second
+        resp3 = api_client.responses.create(
+            model=model,
+            input="What is my friend name?",
+            previous_response_id=resp2.id,
+        )
+        assert resp3.error is None
+        assert resp3.status == "completed"
+        assert "Bob" in resp3.output_text
+
+    def test_store_false_not_retrievable(self, model, api_client):
+        """Test that store=false responses cannot be retrieved."""
+        resp = api_client.responses.create(model=model, input="Hello", store=False)
+        assert resp.id is not None
+        assert resp.status == "completed"
+
+        with pytest.raises((openai.NotFoundError, smg_client.NotFoundError)):
+            api_client.responses.retrieve(response_id=resp.id)

model_gateway/src/routers/factory.rs

@@ -324,7 +324,7 @@ mod grpc_router_type_tests {
 
     /// Build an `AppContext` for a gRPC `RoutingMode` with the components
     /// `GrpcRouter::new` needs (parser factories always; an initialized MCP
-    /// orchestrator for Regular, which is the only mode that reads it).
+    /// orchestrator for Regular/PD, which build the responses contexts).
     async fn grpc_ctx(mode: RoutingMode) -> Arc<AppContext> {
         let config = RouterConfig::builder()
             .mode(mode)

model_gateway/src/routers/grpc/router.rs

@@ -282,8 +282,8 @@ fn transcription_response(format: TranscriptionResponseFormat, text: String) ->
     }
 }
 
-/// `501 NOT_IMPLEMENTED`, returned by Regular-only endpoints when this router is
-/// in PD/EPD mode (matching the `RouterTrait` default).
+/// `501 NOT_IMPLEMENTED`, returned by endpoints this router's mode doesn't
+/// serve (matching the `RouterTrait` default).
 fn not_implemented(message: &'static str) -> Response {
     (StatusCode::NOT_IMPLEMENTED, message).into_response()
 }
@@ -293,10 +293,11 @@ fn not_implemented(message: &'static str) -> Response {
 /// A single `Mode`-parameterized router serving Regular, PrefillDecode, and
 /// EncodePrefillDecode. `mode` selects the disaggregation params baked into
 /// every pipeline and drives the per-mode retry-metric labels, `Debug` output,
-/// and `router_type`. The Regular-only members (`harmony_pipeline`,
-/// `embedding_pipeline`, `classify_pipeline`, `responses_context`,
-/// `harmony_responses_context`) are `Some` only in `Mode::Regular`; PD/EPD leave
-/// them `None` and 501 the corresponding endpoints.
+/// and `router_type`. Optional members are `None` when the mode doesn't serve
+/// them and the corresponding endpoints 501: `embedding_pipeline` and
+/// `classify_pipeline` are Regular-only; `harmony_pipeline`,
+/// `responses_context`, and `harmony_responses_context` exist in every mode
+/// but EPD.
 #[derive(Clone)]
 pub struct GrpcRouter {
     worker_registry: Arc<WorkerRegistry>,
@@ -314,9 +315,9 @@ pub struct GrpcRouter {
 }
 
 impl GrpcRouter {
-    /// Only `Mode::Regular` builds the Harmony, embedding, classify, and
-    /// responses members and requires the MCP orchestrator; PD/EPD leave them
-    /// `None` and 501 those endpoints.
+    /// Regular and PD build the Harmony pipeline and responses contexts and
+    /// require the MCP orchestrator; EPD leaves them `None` and 501s those
+    /// endpoints. Embedding/classify pipelines are Regular-only.
     pub fn new(ctx: &Arc<AppContext>, mode: Mode) -> Result<Self, String> {
         // Get tokenizer registry (no longer requires pre-loaded tokenizer)
         let tokenizer_registry = ctx.tokenizer_registry.clone();
@@ -374,14 +375,16 @@ impl GrpcRouter {
         let completion_pipeline = RequestPipeline::build(Endpoint::Completion, mode, &pair_deps)
             .ok_or_else(|| format!("gRPC router: no completion pipeline for mode {mode:?}"))?;
 
-        // Regular-only pipelines; `None` in PD/EPD (which 501 these endpoints).
+        // `None` when the (endpoint, mode) combo is unsupported; those endpoints 501.
         let harmony_pipeline = RequestPipeline::build(Endpoint::Harmony, mode, &configured_deps);
         let embedding_pipeline = RequestPipeline::build(Endpoint::Embeddings, mode, &pair_deps);
         let classify_pipeline = RequestPipeline::build(Endpoint::Classify, mode, &pair_deps);
 
-        // Responses contexts are Regular-only and are the sole consumer of the MCP
-        // orchestrator; PD/EPD skip both (they don't serve /v1/responses).
-        let (responses_context, harmony_responses_context) = if mode == Mode::Regular {
+        // Responses contexts are the sole consumer of the MCP orchestrator; EPD
+        // builds neither (it doesn't serve /v1/responses).
+        let (responses_context, harmony_responses_context) = if mode == Mode::EncodePrefillDecode {
+            (None, None)
+        } else {
             let mcp_orchestrator = ctx
                 .mcp_orchestrator
                 .get()
@@ -410,11 +413,9 @@ impl GrpcRouter {
                 .as_ref()
                 .map(&create_responses_context)
                 .ok_or_else(|| {
-                    "gRPC router: regular mode must build a harmony pipeline".to_string()
+                    format!("gRPC router: mode {mode:?} must build a harmony pipeline")
                 })?;
             (Some(responses_context), Some(harmony_responses_context))
-        } else {
-            (None, None)
         };
 
         Ok(GrpcRouter {
@@ -482,8 +483,8 @@ impl GrpcRouter {
             return *response;
         }
 
-        // Harmony routing is Regular-only: PD/EPD have no Harmony pipeline, so all
-        // chat requests use the single chat/generate pipeline.
+        // EPD has no Harmony pipeline, so its chat requests all use the single
+        // chat/generate pipeline.
         let is_harmony = self.harmony_pipeline.is_some()
             && HarmonyDetector::is_harmony_model_in_registry(&self.worker_registry, &body.model);
 
@@ -909,8 +910,7 @@ impl RouterTrait for GrpcRouter {
         audio: AudioFile,
         model_id: &str,
     ) -> Response {
-        // Routes through the regular chat pipeline (`execute_chat_for_responses`),
-        // which PD/EPD do not serve.
+        // Qwen3-ASR transcription is Regular-only.
         if self.mode != Mode::Regular {
             return not_implemented("Audio transcriptions not implemented");
         }
@@ -1194,33 +1194,50 @@ mod tests {
 }
 
 #[cfg(test)]
-mod pd_retry_tests {
+mod pd_tests {
     use std::sync::{Arc, OnceLock};
 
     use llm_tokenizer::registry::TokenizerRegistry;
     use reasoning_parser::ParserFactory as ReasoningParserFactory;
     use smg_data_connector::{
         MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
     };
+    use smg_mcp::{McpConfig, McpOrchestrator};
     use tool_parser::ParserFactory as ToolParserFactory;
 
     use super::*;
     use crate::{
         config::{PolicyConfig, RouterConfig, RoutingMode},
         policies::PolicyRegistry,
+        tenant::TenantKey,
         worker::WorkerRegistry,
     };
 
-    /// Minimal `AppContext` for constructing a gRPC PD router. PD/EPD don't read
-    /// the MCP orchestrator, so an empty `OnceLock` suffices.
-    fn pd_ctx() -> Arc<AppContext> {
+    fn pd_routing_mode() -> RoutingMode {
+        RoutingMode::PrefillDecode {
+            prefill_urls: vec![],
+            decode_urls: vec![],
+            prefill_policy: None,
+            decode_policy: None,
+        }
+    }
+
+    fn epd_routing_mode() -> RoutingMode {
+        RoutingMode::EncodePrefillDecode {
+            encode_urls: vec![],
+            prefill_urls: vec![],
+            decode_urls: vec![],
+            encode_policy: None,
+            prefill_policy: None,
+            decode_policy: None,
+        }
+    }
+
+    /// Minimal `AppContext` for constructing a disaggregated gRPC router. PD
+    /// serves /v1/responses, so the MCP orchestrator must be initialized.
+    async fn grpc_ctx(mode: RoutingMode) -> Arc<AppContext> {
         let config = RouterConfig::builder()
-            .mode(RoutingMode::PrefillDecode {
-                prefill_urls: vec![],
-                decode_urls: vec![],
-                prefill_policy: None,
-                decode_policy: None,
-            })
+            .mode(mode)
             .grpc_connection()
             .policy(PolicyConfig::Random)
             .host("127.0.0.1")
@@ -1233,6 +1250,15 @@ mod pd_retry_tests {
             .queue_timeout_secs(60)
             .build_unchecked();
 
+        let mcp_orchestrator = Arc::new(OnceLock::new());
+        mcp_orchestrator
+            .set(Arc::new(
+                McpOrchestrator::new(McpConfig::default())
+                    .await
+                    .expect("mcp orchestrator"),
+            ))
+            .ok();
+
         Arc::new(
             AppContext::builder()
                 .router_config(config.clone())
@@ -1247,17 +1273,24 @@ mod pd_retry_tests {
                 .conversation_item_storage(Arc::new(MemoryConversationItemStorage::new()))
                 .worker_job_queue(Arc::new(OnceLock::new()))
                 .workflow_engines(Arc::new(OnceLock::new()))
-                .mcp_orchestrator(Arc::new(OnceLock::new()))
+                .mcp_orchestrator(mcp_orchestrator)
                 .build()
                 .expect("app context"),
         )
     }
 
+    fn responses_request(model: &str) -> ResponsesRequest {
+        serde_json::from_value(json!({"model": model, "input": "hi"})).expect("responses request")
+    }
+
     /// PD-mode completion must honor a per-model retry override, not the router
     /// default.
-    #[test]
-    fn pd_completion_honors_per_model_retry_override() {
-        let ctx = pd_ctx();
+    // Multi-thread runtime: router construction eagerly builds the Harmony
+    // pipeline, whose one-shot encoding load uses `block_in_place` (which
+    // panics on a current-thread runtime).
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn pd_completion_honors_per_model_retry_override() {
+        let ctx = grpc_ctx(pd_routing_mode()).await;
         let router = GrpcRouter::new(&ctx, Mode::PrefillDecode).expect("pd router");
 
         // Router default differs from the override so the assertion is meaningful.
@@ -1279,4 +1312,39 @@ mod pd_retry_tests {
         let fallback = router.resolve_retry_config("model-without-override");
         assert_eq!(fallback.max_retries, default_retries);
     }
+
+    /// PD serves /v1/responses: an unknown model is rejected per-request (404),
+    /// not gated behind a blanket 501, and cancel reaches storage.
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn pd_router_serves_responses_and_cancel() {
+        let ctx = grpc_ctx(pd_routing_mode()).await;
+        let router = GrpcRouter::new(&ctx, Mode::PrefillDecode).expect("pd router");
+        let tenant_meta = TenantRequestMeta::new(TenantKey::new("test-tenant"));
+
+        let request = responses_request("missing-model");
+        let response = router
+            .route_responses(None, &tenant_meta, &request, "missing-model")
+            .await;
+        assert_eq!(response.status(), StatusCode::NOT_FOUND);
+
+        let cancel = router.cancel_response(None, "resp_missing").await;
+        assert_eq!(cancel.status(), StatusCode::NOT_FOUND);
+    }
+
+    /// EPD still 501s /v1/responses and cancel.
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn epd_router_501s_responses_and_cancel() {
+        let ctx = grpc_ctx(epd_routing_mode()).await;
+        let router = GrpcRouter::new(&ctx, Mode::EncodePrefillDecode).expect("epd router");
+        let tenant_meta = TenantRequestMeta::new(TenantKey::new("test-tenant"));
+
+        let request = responses_request("missing-model");
+        let response = router
+            .route_responses(None, &tenant_meta, &request, "missing-model")
+            .await;
+        assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED);
+
+        let cancel = router.cancel_response(None, "resp_missing").await;
+        assert_eq!(cancel.status(), StatusCode::NOT_IMPLEMENTED);
+    }
 }