diff --git a/crates/tui/src/client.rs b/crates/tui/src/client.rs index 5051141363..da9dd7d6dc 100644 --- a/crates/tui/src/client.rs +++ b/crates/tui/src/client.rs @@ -1353,27 +1353,27 @@ impl DeepSeekClient { )); } - let http_client = Self::build_http_client_with_auth_mode( - &api_key, - &http_headers, - api_provider, - &base_url, - wire_format, - auth_disabled, - false, - )?; + let build_client = |force_http1| -> Result { + let builder = Self::build_http_builder_with_auth_mode( + &api_key, + &http_headers, + api_provider, + &base_url, + wire_format, + auth_disabled, + force_http1, + )?; + let builder = match config.model_http_policy.as_ref() { + Some(policy) => policy.configure(builder, &base_url)?, + None => builder, + }; + builder.build().map_err(Into::into) + }; + let http_client = build_client(false)?; // Always keep an HTTP/1.1 twin for automatic stream-header fallback // when H2 stalls. When CODEWHALE_FORCE_HTTP1 is set, both clients are // HTTP/1.1 and the fallback is a no-op retry path. - let http1_client = Self::build_http_client_with_auth_mode( - &api_key, - &http_headers, - api_provider, - &base_url, - wire_format, - auth_disabled, - true, - )?; + let http1_client = build_client(true)?; Ok(Self { http_client, @@ -1657,7 +1657,7 @@ impl DeepSeekClient { api_provider: ApiProvider, base_url: &str, ) -> Result { - Self::build_http_client_with_auth_mode( + Self::build_http_builder_with_auth_mode( api_key, extra_headers, api_provider, @@ -1665,10 +1665,12 @@ impl DeepSeekClient { provider_default_wire_format(api_provider), false, false, - ) + )? + .build() + .map_err(Into::into) } - fn build_http_client_with_auth_mode( + fn build_http_builder_with_auth_mode( api_key: &str, extra_headers: &HashMap, api_provider: ApiProvider, @@ -1676,7 +1678,7 @@ impl DeepSeekClient { wire_format: WireFormat, auth_disabled: bool, force_http1: bool, - ) -> Result { + ) -> Result { let headers = build_default_headers( api_key, extra_headers, @@ -1705,7 +1707,7 @@ impl DeepSeekClient { { builder = add_extra_root_certs(builder, &cert_path); } - builder.build().map_err(Into::into) + Ok(builder) } /// HTTP/1.1 client for automatic stream-header fallback. @@ -4435,6 +4437,9 @@ mod provider_native_search; mod responses; mod role_placement; mod stream_entry; +#[cfg(test)] +#[path = "client/transport_policy_tests.rs"] +mod transport_policy_tests; mod wire; // Retain the crate-visible accounting helpers at the existing client seam. diff --git a/crates/tui/src/client/transport_policy_tests.rs b/crates/tui/src/client/transport_policy_tests.rs new file mode 100644 index 0000000000..05fa79c564 --- /dev/null +++ b/crates/tui/src/client/transport_policy_tests.rs @@ -0,0 +1,144 @@ +//! Real inference-client builders, including the fallback and standalone vision. +use super::*; +use crate::config::{ModelHttpPolicy, VisionModelConfig}; +use crate::tools::spec::{ToolContext, ToolSpec}; +use crate::vision::tools::ImageAnalyzeTool; +use std::net::SocketAddr; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, method}, +}; + +#[derive(Debug)] +struct FixturePolicy { + address: SocketAddr, + calls: AtomicUsize, +} + +impl ModelHttpPolicy for FixturePolicy { + fn configure( + &self, + builder: reqwest::ClientBuilder, + base_url: &str, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + anyhow::ensure!( + reqwest::Url::parse(base_url)?.host_str() == Some("localhost"), + "unprepared route" + ); + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + "x-host-policy", + reqwest::header::HeaderValue::from_static("applied"), + ); + Ok(builder + .default_headers(headers) + .no_proxy() + .redirect(reqwest::redirect::Policy::none()) + .resolve_to_addrs("localhost", &[self.address])) + } +} + +#[tokio::test] +async fn forkguard_host_http_policy_reaches_primary_fallback_and_vision() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(header("x-host-policy", "applied")) + .respond_with(ResponseTemplate::new(200).set_body_json( + serde_json::json!({"choices":[{"message":{"content":"owned fixture"}}]}), + )) + .expect(3) + .mount(&server) + .await; + let address: SocketAddr = server.uri().trim_start_matches("http://").parse().unwrap(); + let base = format!("http://localhost:{}/v1", address.port()); + let policy = Arc::new(FixturePolicy { + address, + calls: AtomicUsize::new(0), + }); + let config = Config { + provider: Some("deepseek".into()), + api_key: Some("owned-fixture-key".into()), + base_url: Some(base.clone()), + default_text_model: Some("deepseek-v4-pro".into()), + model_http_policy: Some(policy.clone()), + vision_model: Some(VisionModelConfig { + model: "owned-vision-model".into(), + api_key: None, + base_url: Some(base.clone()), + model_http_policy: None, + }), + ..Config::default() + }; + let client = DeepSeekClient::new(&config).unwrap(); + assert_eq!(policy.calls.load(Ordering::SeqCst), 2); + for transport in [&client.http_client, client.http1_fallback_client()] { + transport + .post(format!("{base}/chat/completions")) + .body("owned fixture") + .send() + .await + .unwrap() + .error_for_status() + .unwrap(); + } + let vision = ImageAnalyzeTool::new_with_route_client( + config.clone().vision_model_config().unwrap(), + Some(client.clone()), + ); + assert_eq!(policy.calls.load(Ordering::SeqCst), 3); + let workspace = tempfile::tempdir().unwrap(); + std::fs::write(workspace.path().join("sample.png"), b"owned image fixture").unwrap(); + let result = vision + .execute( + serde_json::json!({"image_path":"sample.png"}), + &ToolContext::new(workspace.path().to_path_buf()), + ) + .await + .unwrap(); + assert!(result.success); + assert_eq!(server.received_requests().await.unwrap().len(), 3); +} + +#[tokio::test] +async fn forkguard_host_http_policy_denial_fails_primary_and_vision_closed() { + let untrusted: Config = + serde_json::from_value(serde_json::json!({"model_http_policy": {"disable": true}})) + .unwrap(); + assert!(untrusted.model_http_policy.is_none()); + let policy = Arc::new(FixturePolicy { + address: "127.0.0.1:9".parse().unwrap(), + calls: AtomicUsize::new(0), + }); + let config = Config { + provider: Some("deepseek".into()), + api_key: Some("owned-fixture-key".into()), + base_url: Some("https://unprepared.invalid/v1".into()), + model_http_policy: Some(policy.clone()), + ..Config::default() + }; + assert!(DeepSeekClient::new(&config).is_err()); + assert_eq!(policy.calls.load(Ordering::SeqCst), 1); + let tool = ImageAnalyzeTool::new_with_route_client( + VisionModelConfig { + model: "owned-fixture".into(), + api_key: None, + base_url: config.base_url, + model_http_policy: Some(policy.clone()), + }, + None, + ); + assert!( + tool.execute( + serde_json::json!({"image_path":"never-read.png"}), + &ToolContext::new(std::env::temp_dir()) + ) + .await + .is_err() + ); + assert_eq!(policy.calls.load(Ordering::SeqCst), 2); +} diff --git a/crates/tui/src/config.rs b/crates/tui/src/config.rs index d78aea3fc6..50ce5dbd41 100644 --- a/crates/tui/src/config.rs +++ b/crates/tui/src/config.rs @@ -2961,6 +2961,16 @@ impl TranscriptConfig { } } +/// Host-owned constraints applied to every inference HTTP client, including +/// the HTTP/1 fallback and standalone vision route. Never loaded from TOML. +pub trait ModelHttpPolicy: std::fmt::Debug + Send + Sync { + fn configure( + &self, + builder: reqwest::ClientBuilder, + base_url: &str, + ) -> Result; +} + /// Resolved CLI configuration, including defaults and environment overrides. #[derive(Debug, Clone, Default, Deserialize)] pub struct Config { @@ -2986,6 +2996,8 @@ pub struct Config { /// a static Idempotency-Key conflicts with this opt-in and is rejected. #[serde(skip)] pub request_idempotency_header: Option, + #[serde(skip)] + pub model_http_policy: Option>, /// Optional user-facing tab/window title shown as `[title] …` in front of /// the terminal window title (the `Codewhale` / `reasoning…` / `done.` /// states). This is the default for every session in this config scope; @@ -3649,6 +3661,9 @@ pub struct VisionModelConfig { /// Base URL for the vision model API. Defaults to OpenAI. #[serde(default)] pub base_url: Option, + /// In-memory host policy, retained when a vision route is cloned. + #[serde(skip)] + pub model_http_policy: Option>, } /// `[runtime_api]` table — knobs for the local HTTP/SSE daemon. @@ -7344,6 +7359,9 @@ impl Config { #[must_use] pub fn vision_model_config(&self) -> Option { let mut config = self.vision_model.clone()?; + if config.model_http_policy.is_none() { + config.model_http_policy = self.model_http_policy.clone(); + } if config.api_key.is_none() { config.api_key = self.api_key.clone(); } @@ -10545,6 +10563,7 @@ fn merge_config(base: Config, override_cfg: Config) -> Config { request_idempotency_header: override_cfg .request_idempotency_header .or(base.request_idempotency_header), + model_http_policy: override_cfg.model_http_policy.or(base.model_http_policy), default_text_model: override_cfg.default_text_model.or(base.default_text_model), auth_mode: override_cfg.auth_mode.or(base.auth_mode), reasoning_effort: override_cfg.reasoning_effort.or(base.reasoning_effort), diff --git a/crates/tui/src/vision/tools.rs b/crates/tui/src/vision/tools.rs index 1653e437c6..535d16e0d6 100644 --- a/crates/tui/src/vision/tools.rs +++ b/crates/tui/src/vision/tools.rs @@ -17,7 +17,7 @@ use crate::tools::spec::{ pub struct ImageAnalyzeTool { config: VisionModelConfig, - client: reqwest::Client, + client: Result, route_client: Option, } @@ -56,15 +56,27 @@ impl ImageAnalyzeTool { config: VisionModelConfig, route_client: Option, ) -> Self { - let client = crate::tls::reqwest_client_builder() - // Bound only the connect handshake. A client- or request-level - // `read_timeout` would start counting at `send()` and never - // reset before the response headers, quietly re-introducing a - // total deadline on the upload + long non-streaming generation; - // the total bound lives in VISION_REQUEST_ENVELOPE instead. - .connect_timeout(Duration::from_secs(10)) - .build() - .expect("Failed to build HTTP client"); + let client = (|| -> anyhow::Result { + let builder = crate::tls::reqwest_client_builder() + // Bound only the connect handshake. A client- or request-level + // `read_timeout` would start counting at `send()` and never + // reset before the response headers, quietly re-introducing a + // total deadline on the upload + long non-streaming generation; + // the total bound lives in VISION_REQUEST_ENVELOPE instead. + .connect_timeout(Duration::from_secs(10)); + let builder = match config.model_http_policy.as_ref() { + Some(policy) => policy.configure( + builder, + config + .base_url + .as_deref() + .unwrap_or("https://api.openai.com/v1"), + )?, + None => builder, + }; + Ok(builder.build()?) + })() + .map_err(|_| ()); Self { config, client, @@ -292,6 +304,10 @@ impl ToolSpec for ImageAnalyzeTool { } async fn execute(&self, input: Value, context: &ToolContext) -> Result { + let client = self + .client + .as_ref() + .map_err(|_| ToolError::execution_failed("Vision HTTP transport unavailable"))?; let image_path = required_str(&input, "image_path")?; let prompt = input .get("prompt") @@ -338,7 +354,7 @@ impl ToolSpec for ImageAnalyzeTool { || { let client = idempotency_route .map(|route| route.http_client.clone()) - .unwrap_or_else(|| self.client.clone()); + .unwrap_or_else(|| client.clone()); let url = url.clone(); let api_key = api_key.clone(); let payload = payload.clone(); @@ -448,6 +464,7 @@ mod tests { model: "test-vision-model".to_string(), api_key: Some("test-key".to_string()), base_url: Some("https://example.invalid/v1".to_string()), + model_http_policy: None, } } @@ -568,6 +585,7 @@ mod tests { model, api_key: None, base_url: Some(base_url), + model_http_policy: None, }, Some(client), ); @@ -776,6 +794,7 @@ mod tests { model: "test-vision-model".to_string(), api_key: Some("test-key".to_string()), base_url: Some(base_url), + model_http_policy: None, }) } @@ -1021,6 +1040,7 @@ mod tests { model: "test-vision-model".into(), api_key: Some(key.into()), base_url: Some(server.uri()), + model_http_policy: None, }, Some(route), ); @@ -1108,6 +1128,7 @@ mod tests { model: "test-vision-model".into(), api_key: Some("owned-test-key".into()), base_url: Some(base.clone()), + model_http_policy: None, }, Some(route), ); @@ -1172,6 +1193,7 @@ mod tests { model: "test-vision-model".into(), api_key: Some("owned-test-key".into()), base_url: Some(format!("{}/v1", server.uri())), + model_http_policy: None, }, Some(route.clone()), ); @@ -1218,6 +1240,7 @@ mod tests { model: "test-vision-model".into(), api_key: Some("ordinary-vision-key".into()), base_url: Some(format!("{}/v1", server.uri())), + model_http_policy: None, }, Some(route), );