diff --git a/crates/persisting-gateway/src/gateway/dispatch.rs b/crates/persisting-gateway/src/gateway/dispatch.rs index e8eb9181b..ac7d920ec 100644 --- a/crates/persisting-gateway/src/gateway/dispatch.rs +++ b/crates/persisting-gateway/src/gateway/dispatch.rs @@ -3,7 +3,8 @@ use async_trait::async_trait; use axum::Router; use axum::extract::Request; -use axum::response::Response; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; use bytes::Bytes; use persisting_overlaynet::policy::{DenyReason, NetworkPolicy}; use persisting_overlaynet::server::{OverlayRequestContext, OverlayServerState, OverlaySink}; @@ -64,6 +65,13 @@ impl OverlaySink for GatewayState { peer: std::net::SocketAddr, context: &OverlayRequestContext, ) -> anyhow::Result { + if !self.gateway_enabled { + return Ok(( + StatusCode::BAD_REQUEST, + "relative request requires Gateway mode", + ) + .into_response()); + } llm_capture( self.clone(), request, @@ -75,8 +83,9 @@ impl OverlaySink for GatewayState { } fn accepts(&self, request: &Request) -> bool { - crate::protocol::ProtocolKind::from_path(request.uri().path()) - != crate::protocol::ProtocolKind::Unknown + self.gateway_enabled + && crate::protocol::ProtocolKind::from_path(request.uri().path()) + != crate::protocol::ProtocolKind::Unknown } fn on_denied( diff --git a/crates/persisting-gateway/src/gateway/state.rs b/crates/persisting-gateway/src/gateway/state.rs index e50377af6..7e36d78b8 100644 --- a/crates/persisting-gateway/src/gateway/state.rs +++ b/crates/persisting-gateway/src/gateway/state.rs @@ -22,6 +22,7 @@ use crate::sink::CaptureEventSink; #[derive(Clone)] pub(crate) struct GatewayState { + pub(crate) gateway_enabled: bool, pub(crate) config: Arc, pub(crate) storage: Arc, pub(crate) client: reqwest::Client, @@ -40,6 +41,7 @@ pub(crate) struct GatewayRuntimeControl { pub(crate) interception_metrics: InterceptionMetrics, pub(crate) bandwidth_registry: BandwidthRegistry, pub(crate) attempt_id: Option, + pub(crate) gateway_enabled: bool, } pub async fn serve( @@ -122,6 +124,7 @@ pub async fn serve_with_listeners_and_shutdown( interception_metrics: InterceptionMetrics::default(), bandwidth_registry: BandwidthRegistry::default(), attempt_id: None, + gateway_enabled: true, }, listener, admin_listener, @@ -154,6 +157,7 @@ pub async fn serve_with_runtime_control( interception_metrics: InterceptionMetrics::default(), bandwidth_registry: BandwidthRegistry::default(), attempt_id: None, + gateway_enabled: true, }, ready, shutdown, @@ -268,6 +272,7 @@ async fn serve_with_bound_listeners( interception_metrics: interception_metrics.clone(), bandwidth_registry: runtime_control.bandwidth_registry, attempt_id: runtime_control.attempt_id, + gateway_enabled: runtime_control.gateway_enabled, }; let admin_state = AdminState { diff --git a/crates/persisting-gateway/src/runtime/in_process.rs b/crates/persisting-gateway/src/runtime/in_process.rs index bf02dba01..696063ad9 100644 --- a/crates/persisting-gateway/src/runtime/in_process.rs +++ b/crates/persisting-gateway/src/runtime/in_process.rs @@ -29,6 +29,8 @@ pub struct InProcessRuntime { pub interception_metrics: InterceptionMetrics, pub bandwidth_registry: BandwidthRegistry, pub attempt_id: Option, + /// Disable LLM dispatch for pVisor runs that only need the network proxy. + pub gateway_enabled: bool, } impl Default for InProcessRuntime { @@ -38,6 +40,7 @@ impl Default for InProcessRuntime { interception_metrics: InterceptionMetrics::default(), bandwidth_registry: BandwidthRegistry::default(), attempt_id: None, + gateway_enabled: true, } } } @@ -84,6 +87,7 @@ impl InProcessCapture { interception_metrics: thread_metrics, bandwidth_registry: runtime.bandwidth_registry, attempt_id: runtime.attempt_id, + gateway_enabled: runtime.gateway_enabled, }, None, async { diff --git a/crates/persisting-gateway/tests/network_policy_http.rs b/crates/persisting-gateway/tests/network_policy_http.rs index 22e00e6b8..9df1541d2 100644 --- a/crates/persisting-gateway/tests/network_policy_http.rs +++ b/crates/persisting-gateway/tests/network_policy_http.rs @@ -13,6 +13,7 @@ use persisting_agentctl::{ ControlController, ControlReason, ControlRequest, ControlTransition, PolicyControlController, }; use persisting_gateway::config::ProxyConfig; +use persisting_gateway::runtime::in_process::{InProcessCapture, InProcessRuntime}; use persisting_gateway::sink::SeqOnlySink; use persisting_gateway::{serve_with_runtime_control, serve_with_shutdown_and_ready}; use tokio::sync::oneshot; @@ -967,6 +968,94 @@ upstream = "http://127.0.0.1:9/v1" let _ = mock_stop.send(()); } +#[tokio::test] +async fn e2e_absolute_uri_llm_keeps_gateway_model_routing() { + let (mock_port, captured, mock_stop) = spawn_capturing_llm_http().await; + let config = format!( + r#" +listen = "{{{{LISTEN}}}}" +admin_listen = "{{{{ADMIN}}}}" + +[network] +mode = "allowlist" +allowed_hosts = ["127.0.0.1"] + +[[models]] +name = "client-model" +forward = "upstream-model" + +[[models]] +name = "upstream-model" +upstream = "http://127.0.0.1:{mock_port}/v1" +"# + ); + let (proxy, _storage, stop) = spawn_proxy(&config).await; + let response = reqwest::Client::builder() + .proxy(reqwest::Proxy::all(&proxy).unwrap()) + .timeout(Duration::from_secs(5)) + .build() + .unwrap() + // The URI destination differs from the configured model upstream. + .post("http://127.0.0.1:9/v1/chat/completions") + .header("content-type", "application/json") + .body(r#"{"model":"client-model","messages":[{"role":"user","content":"hi"}]}"#) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.text().await.unwrap(); + assert_eq!( + captured.lock().unwrap().as_ref().unwrap()["model"], + "upstream-model" + ); + let _ = stop.send(()); + let _ = mock_stop.send(()); +} + +#[tokio::test] +async fn e2e_network_only_allowlist_forwards_absolute_llm_uri() { + let (mock_port, mock_stop) = spawn_mock_http().await; + let listen = format!("127.0.0.1:{}", free_port()); + let admin = format!("127.0.0.1:{}", free_port()); + let config = ProxyConfig::from_toml_str(&format!( + r#"listen = "{listen}" +admin_listen = "{admin}" +agent_id = "t" +models = [] +[network] +mode = "allowlist" +allowed_hosts = ["127.0.0.1"] +"# + )) + .unwrap(); + let storage = tempfile::tempdir().unwrap(); + let proxy = InProcessCapture::start_with_runtime( + config, + storage.path().to_path_buf(), + Arc::new(SeqOnlySink::new()), + false, + InProcessRuntime { + gateway_enabled: false, + ..InProcessRuntime::default() + }, + ) + .unwrap(); + let response = reqwest::Client::builder() + .proxy(reqwest::Proxy::all(format!("http://{}", proxy.listen)).unwrap()) + .timeout(Duration::from_secs(5)) + .build() + .unwrap() + .post(format!("http://127.0.0.1:{mock_port}/v1/chat/completions")) + .body(r#"not JSON; OverlayNet must forward it unchanged"#) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert!(response.text().await.unwrap().contains("chatcmpl-test")); + proxy.shutdown().unwrap(); + let _ = mock_stop.send(()); +} + #[tokio::test] async fn e2e_relative_llm_gateway_bypasses_host_allowlist() { let (mock_port, mock_stop) = spawn_mock_http().await; diff --git a/crates/persisting-pvisor/src/runtime/attempt.rs b/crates/persisting-pvisor/src/runtime/attempt.rs index c4ac49d80..2324d6dc0 100644 --- a/crates/persisting-pvisor/src/runtime/attempt.rs +++ b/crates/persisting-pvisor/src/runtime/attempt.rs @@ -352,6 +352,7 @@ pub(crate) fn prepare_attempt( interception_metrics: network_metrics.clone(), bandwidth_registry: bandwidth_registry.clone(), attempt_id: Some(opts.attempt_id.to_owned()), + gateway_enabled: opts.gateway_enabled, }, )?; @@ -1126,6 +1127,17 @@ fn enrich_with_session( { plan.env.insert(key, value); } + if !gateway_enabled { + for key in [ + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "AZURE_OPENAI_ENDPOINT", + "ANTHROPIC_BASE_URL", + "GEMINI_API_BASE", + ] { + plan.env.remove(key); + } + } plan.notes .push(format!("network service: proxy env → http://{listen}")); if uses_krun_executor(spec) {