From cec2b682831f4f9dde3f35bbe7622084ecf217c6 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 3 Jun 2026 15:33:09 -0500 Subject: [PATCH 1/3] Add configurable backend timeouts --- .env.example | 8 +- Cargo.toml | 2 +- Dockerfile | 4 +- README.md | 2 + docker-compose.yml | 4 +- justfile | 6 ++ src/config.rs | 85 ++++++++++++++++++- src/proxy.rs | 188 ++++++++++++++++++++++++++++++++----------- tests/health_test.rs | 4 + 9 files changed, 249 insertions(+), 54 deletions(-) diff --git a/.env.example b/.env.example index 2fee4ff..b94479e 100644 --- a/.env.example +++ b/.env.example @@ -22,5 +22,11 @@ MAPLE_DEBUG=false # Recommended: true for Docker deployments MAPLE_ENABLE_CORS=true +# Timeouts +# Backend request setup and non-streaming response timeout, in seconds +MAPLE_REQUEST_TIMEOUT_SECS=300 +# Maximum idle time between streaming chunks, in seconds +MAPLE_STREAM_IDLE_TIMEOUT_SECS=300 + # Rust Logging (optional) -# RUST_LOG=info,maple_proxy=debug \ No newline at end of file +# RUST_LOG=info,maple_proxy=debug diff --git a/Cargo.toml b/Cargo.toml index ba8f00f..28f2601 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,7 +24,7 @@ opensecret = "3.1.1" # Web server axum = { version = "0.8.4", features = ["http2", "macros"] } -tokio = { version = "1.47", features = ["net", "rt-multi-thread", "macros", "sync"] } +tokio = { version = "1.47", features = ["net", "rt-multi-thread", "macros", "sync", "time"] } tower = { version = "0.5.2", features = ["full"] } tower-http = { version = "0.6.6", features = ["cors", "trace"] } diff --git a/Dockerfile b/Dockerfile index cc9aa77..3be180f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -56,6 +56,8 @@ ENV MAPLE_HOST=0.0.0.0 \ MAPLE_BACKEND_URL=https://enclave.trymaple.ai \ MAPLE_DEBUG=false \ MAPLE_ENABLE_CORS=true \ + MAPLE_REQUEST_TIMEOUT_SECS=300 \ + MAPLE_STREAM_IDLE_TIMEOUT_SECS=300 \ RUST_LOG=info # Expose the port @@ -67,4 +69,4 @@ HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \ CMD curl -f http://localhost:8080/health || exit 1 # Run the binary -ENTRYPOINT ["/usr/local/bin/maple-proxy"] \ No newline at end of file +ENTRYPOINT ["/usr/local/bin/maple-proxy"] diff --git a/README.md b/README.md index 956cc59..a2116e7 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,8 @@ export MAPLE_BACKEND_URL=http://localhost:3000 # Maple backend URL (prod export MAPLE_API_KEY=your-maple-api-key # Default API key (optional) export MAPLE_DEBUG=true # Enable debug logging export MAPLE_ENABLE_CORS=true # Enable CORS +export MAPLE_REQUEST_TIMEOUT_SECS=300 # Backend request timeout +export MAPLE_STREAM_IDLE_TIMEOUT_SECS=300 # Streaming idle timeout between chunks ``` Or use CLI arguments: diff --git a/docker-compose.yml b/docker-compose.yml index 32df0c8..113fdcc 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -27,6 +27,8 @@ services: # Optional configurations - MAPLE_DEBUG=${MAPLE_DEBUG:-false} - MAPLE_ENABLE_CORS=${MAPLE_ENABLE_CORS:-true} + - MAPLE_REQUEST_TIMEOUT_SECS=${MAPLE_REQUEST_TIMEOUT_SECS:-300} + - MAPLE_STREAM_IDLE_TIMEOUT_SECS=${MAPLE_STREAM_IDLE_TIMEOUT_SECS:-300} - RUST_LOG=${RUST_LOG:-info} # Production-grade restart policy @@ -62,4 +64,4 @@ services: networks: maple-network: - driver: bridge \ No newline at end of file + driver: bridge diff --git a/justfile b/justfile index 39999bf..04251b8 100644 --- a/justfile +++ b/justfile @@ -166,6 +166,8 @@ env: @echo "MAPLE_API_KEY: ${MAPLE_API_KEY:-[not set]}" @echo "MAPLE_DEBUG: ${MAPLE_DEBUG:-false}" @echo "MAPLE_ENABLE_CORS: ${MAPLE_ENABLE_CORS:-false}" + @echo "MAPLE_REQUEST_TIMEOUT_SECS: ${MAPLE_REQUEST_TIMEOUT_SECS:-300}" + @echo "MAPLE_STREAM_IDLE_TIMEOUT_SECS: ${MAPLE_STREAM_IDLE_TIMEOUT_SECS:-300}" # Build Docker image docker-build: @@ -182,6 +184,8 @@ docker-run: -e MAPLE_BACKEND_URL=${MAPLE_BACKEND_URL:-https://enclave.trymaple.ai} \ -e MAPLE_DEBUG=${MAPLE_DEBUG:-false} \ -e MAPLE_ENABLE_CORS=${MAPLE_ENABLE_CORS:-true} \ + -e MAPLE_REQUEST_TIMEOUT_SECS=${MAPLE_REQUEST_TIMEOUT_SECS:-300} \ + -e MAPLE_STREAM_IDLE_TIMEOUT_SECS=${MAPLE_STREAM_IDLE_TIMEOUT_SECS:-300} \ maple-proxy:latest # Run Docker container in detached mode @@ -194,6 +198,8 @@ docker-run-detached: -e MAPLE_BACKEND_URL=${MAPLE_BACKEND_URL:-https://enclave.trymaple.ai} \ -e MAPLE_DEBUG=${MAPLE_DEBUG:-false} \ -e MAPLE_ENABLE_CORS=${MAPLE_ENABLE_CORS:-true} \ + -e MAPLE_REQUEST_TIMEOUT_SECS=${MAPLE_REQUEST_TIMEOUT_SECS:-300} \ + -e MAPLE_STREAM_IDLE_TIMEOUT_SECS=${MAPLE_STREAM_IDLE_TIMEOUT_SECS:-300} \ maple-proxy:latest @echo "✅ Container started. Use 'just docker-stop' to stop it." diff --git a/src/config.rs b/src/config.rs index d5b4fa8..4c959ca 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,9 @@ use clap::Parser; use serde::{Deserialize, Serialize}; -use std::net::SocketAddr; +use std::{net::SocketAddr, time::Duration}; + +pub const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 300; +pub const DEFAULT_STREAM_IDLE_TIMEOUT_SECS: u64 = 300; #[derive(Parser, Debug, Clone)] #[command(name = "maple-proxy")] @@ -33,6 +36,22 @@ pub struct Config { /// Enable CORS for all origins (useful for web clients) #[arg(long, env = "MAPLE_ENABLE_CORS")] pub enable_cors: bool, + + /// Timeout for backend request setup and non-streaming responses, in seconds + #[arg( + long, + env = "MAPLE_REQUEST_TIMEOUT_SECS", + default_value_t = DEFAULT_REQUEST_TIMEOUT_SECS + )] + pub request_timeout_secs: u64, + + /// Maximum time to wait between streaming response chunks, in seconds + #[arg( + long, + env = "MAPLE_STREAM_IDLE_TIMEOUT_SECS", + default_value_t = DEFAULT_STREAM_IDLE_TIMEOUT_SECS + )] + pub stream_idle_timeout_secs: u64, } impl Config { @@ -58,9 +77,19 @@ impl Config { default_api_key: None, debug: false, enable_cors: false, + request_timeout_secs: DEFAULT_REQUEST_TIMEOUT_SECS, + stream_idle_timeout_secs: DEFAULT_STREAM_IDLE_TIMEOUT_SECS, } } + pub fn request_timeout(&self) -> Duration { + Duration::from_secs(self.request_timeout_secs) + } + + pub fn stream_idle_timeout(&self) -> Duration { + Duration::from_secs(self.stream_idle_timeout_secs) + } + /// Builder-style method to set the API key pub fn with_api_key(mut self, api_key: String) -> Self { self.default_api_key = Some(api_key); @@ -78,6 +107,18 @@ impl Config { self.enable_cors = enable_cors; self } + + /// Builder-style method to set the backend request timeout + pub fn with_request_timeout_secs(mut self, request_timeout_secs: u64) -> Self { + self.request_timeout_secs = request_timeout_secs; + self + } + + /// Builder-style method to set the streaming idle timeout + pub fn with_stream_idle_timeout_secs(mut self, stream_idle_timeout_secs: u64) -> Self { + self.stream_idle_timeout_secs = stream_idle_timeout_secs; + self + } } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -114,3 +155,45 @@ impl OpenAIError { Self::new(message, "server_error") } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn config_new_uses_timeout_defaults() { + let config = Config::new( + "127.0.0.1".to_string(), + 8080, + "https://enclave.trymaple.ai".to_string(), + ); + + assert_eq!(config.request_timeout_secs, DEFAULT_REQUEST_TIMEOUT_SECS); + assert_eq!( + config.stream_idle_timeout_secs, + DEFAULT_STREAM_IDLE_TIMEOUT_SECS + ); + assert_eq!( + config.request_timeout(), + Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS) + ); + assert_eq!( + config.stream_idle_timeout(), + Duration::from_secs(DEFAULT_STREAM_IDLE_TIMEOUT_SECS) + ); + } + + #[test] + fn timeout_builder_methods_override_defaults() { + let config = Config::new( + "127.0.0.1".to_string(), + 8080, + "https://enclave.trymaple.ai".to_string(), + ) + .with_request_timeout_secs(45) + .with_stream_idle_timeout_secs(15); + + assert_eq!(config.request_timeout(), Duration::from_secs(45)); + assert_eq!(config.stream_idle_timeout(), Duration::from_secs(15)); + } +} diff --git a/src/proxy.rs b/src/proxy.rs index 2483921..4fc6c2f 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -81,6 +81,7 @@ impl ProxyState { let cache_key = api_key.to_string(); let client_entry = self.client_entry_for_api_key(&cache_key); let backend_url = self.config.backend_url.clone(); + let request_timeout = self.config.request_timeout(); let init_api_key = cache_key.clone(); let client = client_entry @@ -90,7 +91,7 @@ impl ProxyState { "Creating OpenSecret client for API key: {}...", &init_api_key[..8.min(init_api_key.len())] ); - create_client_with_auth(&backend_url, &init_api_key) + create_client_with_auth(&backend_url, &init_api_key, request_timeout) .await .map(Arc::new) }) @@ -164,19 +165,48 @@ fn extract_api_key( async fn create_client_with_auth( backend_url: &str, api_key: &str, + request_timeout: Duration, ) -> Result { let client = OpenSecretClient::new_with_api_key(backend_url, api_key.to_string()) .map_err(|e| OpenAIError::server_error(format!("Failed to create client: {}", e)))?; // Perform attestation handshake - client.perform_attestation_handshake().await.map_err(|e| { - error!("Attestation handshake failed: {}", e); - OpenAIError::server_error("Failed to establish secure connection with Maple backend") - })?; + tokio::time::timeout(request_timeout, client.perform_attestation_handshake()) + .await + .map_err(|_| { + error!( + "Attestation handshake timed out after {} seconds", + request_timeout.as_secs() + ); + OpenAIError::server_error(format!( + "Timed out establishing secure connection with Maple backend after {} seconds", + request_timeout.as_secs() + )) + })? + .map_err(|e| { + error!("Attestation handshake failed: {}", e); + OpenAIError::server_error("Failed to establish secure connection with Maple backend") + })?; Ok(client) } +fn timeout_response(operation: &str, timeout: Duration) -> (StatusCode, Json) { + error!( + "{} timed out after {} seconds", + operation, + timeout.as_secs() + ); + ( + StatusCode::GATEWAY_TIMEOUT, + Json(OpenAIError::server_error(format!( + "{} timed out after {} seconds", + operation, + timeout.as_secs() + ))), + ) +} + pub async fn list_models( State(state): State>, headers: HeaderMap, @@ -194,16 +224,20 @@ pub async fn list_models( .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, Json(e)))?; - let models = client.get_models().await.map_err(|e| { - error!("Failed to get models: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to retrieve models: {}", - e - ))), - ) - })?; + let request_timeout = state.config.request_timeout(); + let models = tokio::time::timeout(request_timeout, client.get_models()) + .await + .map_err(|_| timeout_response("List models request", request_timeout))? + .map_err(|e| { + error!("Failed to get models: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(OpenAIError::server_error(format!( + "Failed to retrieve models: {}", + e + ))), + ) + })?; debug!("Successfully retrieved {} models", models.data.len()); Ok(Json(models)) @@ -231,37 +265,46 @@ pub async fn create_chat_completion( // Check if streaming is requested if request.stream.unwrap_or(false) { // Handle streaming response - let stream = client - .create_chat_completion_stream(request) - .await - .map_err(|e| { - error!("Failed to create streaming chat completion: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to create streaming completion: {}", - e - ))), - ) - })?; - - let sse_stream = create_sse_stream(stream); - Ok(Sse::new(sse_stream).into_response()) - } else { - // Handle non-streaming response - request.stream = Some(false); // Ensure it's explicitly false - - let response = client.create_chat_completion(request).await.map_err(|e| { - error!("Failed to create chat completion: {}", e); + let request_timeout = state.config.request_timeout(); + let stream = tokio::time::timeout( + request_timeout, + client.create_chat_completion_stream(request), + ) + .await + .map_err(|_| timeout_response("Streaming chat completion request", request_timeout))? + .map_err(|e| { + error!("Failed to create streaming chat completion: {}", e); ( StatusCode::INTERNAL_SERVER_ERROR, Json(OpenAIError::server_error(format!( - "Failed to create completion: {}", + "Failed to create streaming completion: {}", e ))), ) })?; + let sse_stream = create_sse_stream(stream, state.config.stream_idle_timeout()); + Ok(Sse::new(sse_stream).into_response()) + } else { + // Handle non-streaming response + request.stream = Some(false); // Ensure it's explicitly false + + let request_timeout = state.config.request_timeout(); + let response = + tokio::time::timeout(request_timeout, client.create_chat_completion(request)) + .await + .map_err(|_| timeout_response("Chat completion request", request_timeout))? + .map_err(|e| { + error!("Failed to create chat completion: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(OpenAIError::server_error(format!( + "Failed to create completion: {}", + e + ))), + ) + })?; + debug!("Successfully created chat completion: {}", response.id); Ok(Json(response).into_response()) } @@ -269,11 +312,30 @@ pub async fn create_chat_completion( fn create_sse_stream( mut stream: std::pin::Pin> + Send>>, + stream_idle_timeout: Duration, ) -> impl Stream> { async_stream::stream! { use futures::StreamExt; - while let Some(chunk_result) = stream.next().await { + loop { + let chunk_result = match tokio::time::timeout(stream_idle_timeout, stream.next()).await { + Ok(Some(chunk_result)) => chunk_result, + Ok(None) => break, + Err(_) => { + error!( + "Streaming chat completion idle timed out after {} seconds", + stream_idle_timeout.as_secs() + ); + let error_event = axum::response::sse::Event::default() + .data(format!( + r#"{{"error": "Stream idle timeout after {} seconds"}}"#, + stream_idle_timeout.as_secs() + )); + yield Ok(error_event); + break; + } + }; + match chunk_result { Ok(chunk) => { match serde_json::to_string(&chunk) { @@ -323,16 +385,20 @@ pub async fn create_embeddings( .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, Json(e)))?; - let response = client.create_embeddings(request).await.map_err(|e| { - error!("Failed to create embeddings: {}", e); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(OpenAIError::server_error(format!( - "Failed to create embeddings: {}", - e - ))), - ) - })?; + let request_timeout = state.config.request_timeout(); + let response = tokio::time::timeout(request_timeout, client.create_embeddings(request)) + .await + .map_err(|_| timeout_response("Embeddings request", request_timeout))? + .map_err(|e| { + error!("Failed to create embeddings: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(OpenAIError::server_error(format!( + "Failed to create embeddings: {}", + e + ))), + ) + })?; debug!( "Successfully created embeddings with {} vectors", @@ -353,6 +419,8 @@ mod tests { default_api_key: None, debug: false, enable_cors: false, + request_timeout_secs: 300, + stream_idle_timeout_secs: 300, } } @@ -419,4 +487,26 @@ mod tests { assert!(!state.clients.contains_key("key-a")); } + + #[tokio::test] + async fn sse_stream_emits_error_when_idle_timeout_expires() { + use futures::{stream, StreamExt}; + + let backend_stream = Box::pin(stream::pending::>()); + let sse_stream = create_sse_stream(backend_stream, Duration::from_millis(1)); + futures::pin_mut!(sse_stream); + + assert!( + tokio::time::timeout(Duration::from_secs(1), sse_stream.next()) + .await + .unwrap() + .is_some() + ); + assert!( + tokio::time::timeout(Duration::from_secs(1), sse_stream.next()) + .await + .unwrap() + .is_some() + ); + } } diff --git a/tests/health_test.rs b/tests/health_test.rs index 6ddd3cd..073dced 100644 --- a/tests/health_test.rs +++ b/tests/health_test.rs @@ -13,6 +13,8 @@ async fn test_health_check_endpoint() { default_api_key: None, debug: false, enable_cors: false, + request_timeout_secs: 300, + stream_idle_timeout_secs: 300, }; // Create test server @@ -38,6 +40,8 @@ async fn test_root_health_check() { default_api_key: None, debug: false, enable_cors: false, + request_timeout_secs: 300, + stream_idle_timeout_secs: 300, }; let app = create_app(config); From 499129a10b0297f8656c6c7b36cac78affdde4f1 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 3 Jun 2026 15:44:08 -0500 Subject: [PATCH 2/3] Address timeout review feedback --- README.md | 4 +++ src/config.rs | 18 +++++++++-- src/proxy.rs | 73 ++++++++++++++++++++++---------------------- tests/health_test.rs | 30 ++++++------------ 4 files changed, 66 insertions(+), 59 deletions(-) diff --git a/README.md b/README.md index a2116e7..7761615 100644 --- a/README.md +++ b/README.md @@ -244,6 +244,8 @@ docker pull ghcr.io/opensecretcloud/maple-proxy:latest # Run with your API key docker run -p 8080:8080 \ -e MAPLE_BACKEND_URL=https://enclave.trymaple.ai \ + -e MAPLE_REQUEST_TIMEOUT_SECS=300 \ + -e MAPLE_STREAM_IDLE_TIMEOUT_SECS=300 \ ghcr.io/opensecretcloud/maple-proxy:latest ``` @@ -342,6 +344,8 @@ The Docker image: environment: - MAPLE_BACKEND_URL=https://enclave.trymaple.ai # Production backend - MAPLE_ENABLE_CORS=true # Enable for web apps + - MAPLE_REQUEST_TIMEOUT_SECS=300 # Backend request timeout + - MAPLE_STREAM_IDLE_TIMEOUT_SECS=300 # Streaming idle timeout - RUST_LOG=info # Logging level # - MAPLE_API_KEY=xxx # Only for private deployments! ``` diff --git a/src/config.rs b/src/config.rs index 4c959ca..49c1900 100644 --- a/src/config.rs +++ b/src/config.rs @@ -41,7 +41,8 @@ pub struct Config { #[arg( long, env = "MAPLE_REQUEST_TIMEOUT_SECS", - default_value_t = DEFAULT_REQUEST_TIMEOUT_SECS + default_value_t = DEFAULT_REQUEST_TIMEOUT_SECS, + value_parser = clap::value_parser!(u64).range(1..) )] pub request_timeout_secs: u64, @@ -49,7 +50,8 @@ pub struct Config { #[arg( long, env = "MAPLE_STREAM_IDLE_TIMEOUT_SECS", - default_value_t = DEFAULT_STREAM_IDLE_TIMEOUT_SECS + default_value_t = DEFAULT_STREAM_IDLE_TIMEOUT_SECS, + value_parser = clap::value_parser!(u64).range(1..) )] pub stream_idle_timeout_secs: u64, } @@ -159,6 +161,7 @@ impl OpenAIError { #[cfg(test)] mod tests { use super::*; + use clap::{error::ErrorKind, Parser}; #[test] fn config_new_uses_timeout_defaults() { @@ -196,4 +199,15 @@ mod tests { assert_eq!(config.request_timeout(), Duration::from_secs(45)); assert_eq!(config.stream_idle_timeout(), Duration::from_secs(15)); } + + #[test] + fn timeout_cli_values_must_be_positive() { + let request_timeout_error = + Config::try_parse_from(["maple-proxy", "--request-timeout-secs", "0"]).unwrap_err(); + assert_eq!(request_timeout_error.kind(), ErrorKind::ValueValidation); + + let stream_idle_timeout_error = + Config::try_parse_from(["maple-proxy", "--stream-idle-timeout-secs", "0"]).unwrap_err(); + assert_eq!(stream_idle_timeout_error.kind(), ErrorKind::ValueValidation); + } } diff --git a/src/proxy.rs b/src/proxy.rs index 4fc6c2f..be38a8a 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -22,6 +22,8 @@ use tracing::{debug, error}; const CLIENT_CACHE_MAX_ENTRIES: usize = 1024; const CLIENT_CACHE_ENTRY_TTL: Duration = Duration::from_secs(60 * 60); +type ProxyError = (StatusCode, Json); + struct CachedClientEntry { cell: OnceCell>, created_at: Instant, @@ -74,10 +76,7 @@ impl ProxyState { .clone() } - async fn client_for_api_key( - &self, - api_key: &str, - ) -> Result, OpenAIError> { + async fn client_for_api_key(&self, api_key: &str) -> Result, ProxyError> { let cache_key = api_key.to_string(); let client_entry = self.client_entry_for_api_key(&cache_key); let backend_url = self.config.backend_url.clone(); @@ -166,32 +165,36 @@ async fn create_client_with_auth( backend_url: &str, api_key: &str, request_timeout: Duration, -) -> Result { - let client = OpenSecretClient::new_with_api_key(backend_url, api_key.to_string()) - .map_err(|e| OpenAIError::server_error(format!("Failed to create client: {}", e)))?; +) -> Result { + let client = + OpenSecretClient::new_with_api_key(backend_url, api_key.to_string()).map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(OpenAIError::server_error(format!( + "Failed to create client: {}", + e + ))), + ) + })?; // Perform attestation handshake tokio::time::timeout(request_timeout, client.perform_attestation_handshake()) .await - .map_err(|_| { - error!( - "Attestation handshake timed out after {} seconds", - request_timeout.as_secs() - ); - OpenAIError::server_error(format!( - "Timed out establishing secure connection with Maple backend after {} seconds", - request_timeout.as_secs() - )) - })? + .map_err(|_| timeout_response("Attestation handshake", request_timeout))? .map_err(|e| { error!("Attestation handshake failed: {}", e); - OpenAIError::server_error("Failed to establish secure connection with Maple backend") + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(OpenAIError::server_error( + "Failed to establish secure connection with Maple backend", + )), + ) })?; Ok(client) } -fn timeout_response(operation: &str, timeout: Duration) -> (StatusCode, Json) { +fn timeout_response(operation: &str, timeout: Duration) -> ProxyError { error!( "{} timed out after {} seconds", operation, @@ -219,10 +222,7 @@ pub async fn list_models( &api_key[..8.min(api_key.len())] ); - let client = state - .client_for_api_key(&api_key) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, Json(e)))?; + let client = state.client_for_api_key(&api_key).await?; let request_timeout = state.config.request_timeout(); let models = tokio::time::timeout(request_timeout, client.get_models()) @@ -257,10 +257,7 @@ pub async fn create_chat_completion( request.stream.unwrap_or(false) ); - let client = state - .client_for_api_key(&api_key) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, Json(e)))?; + let client = state.client_for_api_key(&api_key).await?; // Check if streaming is requested if request.stream.unwrap_or(false) { @@ -316,11 +313,15 @@ fn create_sse_stream( ) -> impl Stream> { async_stream::stream! { use futures::StreamExt; + let mut completed_normally = false; loop { let chunk_result = match tokio::time::timeout(stream_idle_timeout, stream.next()).await { Ok(Some(chunk_result)) => chunk_result, - Ok(None) => break, + Ok(None) => { + completed_normally = true; + break; + } Err(_) => { error!( "Streaming chat completion idle timed out after {} seconds", @@ -363,10 +364,11 @@ fn create_sse_stream( } } - // Send [DONE] event to indicate end of stream - let done_event = axum::response::sse::Event::default() - .data("[DONE]"); - yield Ok(done_event); + if completed_normally { + let done_event = axum::response::sse::Event::default() + .data("[DONE]"); + yield Ok(done_event); + } } } @@ -380,10 +382,7 @@ pub async fn create_embeddings( debug!("Embeddings request for model: {}", request.model); - let client = state - .client_for_api_key(&api_key) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, Json(e)))?; + let client = state.client_for_api_key(&api_key).await?; let request_timeout = state.config.request_timeout(); let response = tokio::time::timeout(request_timeout, client.create_embeddings(request)) @@ -506,7 +505,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(1), sse_stream.next()) .await .unwrap() - .is_some() + .is_none() ); } } diff --git a/tests/health_test.rs b/tests/health_test.rs index 073dced..fe914bc 100644 --- a/tests/health_test.rs +++ b/tests/health_test.rs @@ -6,16 +6,11 @@ use serde_json::Value; #[tokio::test] async fn test_health_check_endpoint() { // Create test config - let config = Config { - host: "127.0.0.1".to_string(), - port: 0, // Use random port for testing - backend_url: "http://localhost:3000".to_string(), - default_api_key: None, - debug: false, - enable_cors: false, - request_timeout_secs: 300, - stream_idle_timeout_secs: 300, - }; + let config = Config::new( + "127.0.0.1".to_string(), + 0, // Use random port for testing + "http://localhost:3000".to_string(), + ); // Create test server let app = create_app(config); @@ -33,16 +28,11 @@ async fn test_health_check_endpoint() { #[tokio::test] async fn test_root_health_check() { - let config = Config { - host: "127.0.0.1".to_string(), - port: 0, - backend_url: "http://localhost:3000".to_string(), - default_api_key: None, - debug: false, - enable_cors: false, - request_timeout_secs: 300, - stream_idle_timeout_secs: 300, - }; + let config = Config::new( + "127.0.0.1".to_string(), + 0, + "http://localhost:3000".to_string(), + ); let app = create_app(config); let server = TestServer::new(app).unwrap(); From 82cfecb4d12999f6fe78071fbd60f6d12cff0e50 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Wed, 3 Jun 2026 15:49:15 -0500 Subject: [PATCH 3/3] Document timeout env vars in guidance --- CLAUDE.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index e07d45a..4ea10ad 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -80,6 +80,8 @@ Environment variables (can be set in .env file): - `MAPLE_API_KEY` - Default API key (optional) - `MAPLE_DEBUG` - Enable debug logging - `MAPLE_ENABLE_CORS` - Enable CORS for web clients +- `MAPLE_REQUEST_TIMEOUT_SECS` - Backend request timeout in seconds (default: 300) +- `MAPLE_STREAM_IDLE_TIMEOUT_SECS` - Streaming idle timeout in seconds (default: 300) ## Testing @@ -96,4 +98,4 @@ Key dependencies: - **tokio** - Async runtime - **tower/tower-http** - Middleware for CORS and tracing - **clap** - CLI argument parsing -- **dotenvy** - .env file support \ No newline at end of file +- **dotenvy** - .env file support