diff --git a/crates/brightstaff/src/handlers/agent_chat_completions.rs b/crates/brightstaff/src/handlers/agent_chat_completions.rs index 345ec69a..bac9607b 100644 --- a/crates/brightstaff/src/handlers/agent_chat_completions.rs +++ b/crates/brightstaff/src/handlers/agent_chat_completions.rs @@ -206,7 +206,7 @@ async fn handle_agent_chat( let message: Vec = client_request.get_messages(); // Extract trace parent for routing - let trace_parent = request_headers + let traceparent = request_headers .iter() .find(|(key, _)| key.as_str() == TRACE_PARENT_HEADER) .map(|(_, value)| value.to_str().unwrap_or_default().to_string()); @@ -224,7 +224,7 @@ async fn handle_agent_chat( }; // Parse trace parent to get trace_id and parent_span_id - let (trace_id, parent_span_id) = if let Some(ref tp) = trace_parent { + let (trace_id, parent_span_id) = if let Some(ref tp) = traceparent { parse_traceparent(tp) } else { (String::new(), None) @@ -236,12 +236,7 @@ async fn handle_agent_chat( let selection_start_instant = Instant::now(); let selected_agents = agent_selector - .select_agents( - &message, - &listener, - trace_parent.clone(), - request_id.clone(), - ) + .select_agents(&message, &listener, traceparent.clone(), request_id.clone()) .await?; // Record agent selection span diff --git a/crates/brightstaff/src/handlers/llm.rs b/crates/brightstaff/src/handlers/llm.rs index 37cb008e..4b7d2f70 100644 --- a/crates/brightstaff/src/handlers/llm.rs +++ b/crates/brightstaff/src/handlers/llm.rs @@ -44,22 +44,40 @@ pub async fn llm_chat( ) -> Result>, hyper::Error> { let request_path = request.uri().path().to_string(); let request_headers = request.headers().clone(); - let request_id = request_headers + let request_id: String = match request_headers .get(REQUEST_ID_HEADER) .and_then(|h| h.to_str().ok()) .map(|s| s.to_string()) - .unwrap_or_else(|| "unknown".to_string()); + { + Some(id) => id, + None => { + let generated_id = uuid::Uuid::new_v4().to_string(); + warn!( + "[PLANO_REQ_ID:{}] | REQUEST_ID header missing, generated new ID", + generated_id + ); + generated_id + } + }; // Extract or generate traceparent - this establishes the trace context for all spans - let traceparent: String = request_headers + let traceparent: String = match request_headers .get(TRACE_PARENT_HEADER) .and_then(|h| h.to_str().ok()) .map(|s| s.to_string()) - .unwrap_or_else(|| { + { + Some(tp) => tp, + None => { use uuid::Uuid; let trace_id = Uuid::new_v4().to_string().replace("-", ""); - format!("00-{}-0000000000000000-01", trace_id) - }); + let generated_tp = format!("00-{}-0000000000000000-01", trace_id); + warn!( + "[PLANO_REQ_ID:{}] | TRACE_PARENT header missing, generated new traceparent: {}", + request_id, generated_tp + ); + generated_tp + } + }; let mut request_headers = request_headers; let chat_request_bytes = request.collect().await?.to_bytes(); @@ -204,11 +222,10 @@ pub async fn llm_chat( let routing_result = match router_chat_get_upstream_model( router_service, client_request, // Pass the original request - router_chat will convert it - &request_headers, trace_collector.clone(), &traceparent, &request_path, - Some(request_id.to_string()), + &request_id, ) .await { @@ -318,7 +335,7 @@ pub async fn llm_chat( is_streaming_request, false, // Not OpenAI upstream since should_manage_state is true content_encoding, - request_id.clone(), + request_id, ); create_streaming_response(byte_stream, state_processor, 16) } else { diff --git a/crates/brightstaff/src/handlers/router_chat.rs b/crates/brightstaff/src/handlers/router_chat.rs index b4045095..67e25338 100644 --- a/crates/brightstaff/src/handlers/router_chat.rs +++ b/crates/brightstaff/src/handlers/router_chat.rs @@ -36,17 +36,14 @@ impl RoutingError { pub async fn router_chat_get_upstream_model( router_service: Arc, client_request: ProviderRequestType, - _request_headers: &hyper::HeaderMap, trace_collector: Arc, traceparent: &str, request_path: &str, - request_id: Option, + request_id: &str, ) -> Result { // Clone metadata for routing before converting (which consumes client_request) let routing_metadata = client_request.metadata().clone(); - let request_id = request_id.unwrap_or_else(|| "unknown".to_string()); - // Convert to ChatCompletionsRequest for routing (regardless of input type) let chat_request = match ProviderRequestType::try_from(( client_request, @@ -128,9 +125,9 @@ pub async fn router_chat_get_upstream_model( let routing_result = router_service .determine_route( &chat_request.messages, - Some(traceparent.to_string()), + traceparent, usage_preferences, - Some(request_id.clone()), + request_id, ) .await; diff --git a/crates/brightstaff/src/router/llm_router.rs b/crates/brightstaff/src/router/llm_router.rs index 6d9711f2..e99652d3 100644 --- a/crates/brightstaff/src/router/llm_router.rs +++ b/crates/brightstaff/src/router/llm_router.rs @@ -77,9 +77,9 @@ impl RouterService { pub async fn determine_route( &self, messages: &[Message], - trace_parent: Option, + traceparent: &str, usage_preferences: Option>, - request_id: Option, + request_id: &str, ) -> Result> { if messages.is_empty() { return Ok(None); @@ -117,19 +117,15 @@ impl RouterService { header::HeaderValue::from_str(&self.routing_provider_name).unwrap(), ); - if let Some(trace_parent) = trace_parent { - llm_route_request_headers.insert( - header::HeaderName::from_static(TRACE_PARENT_HEADER), - header::HeaderValue::from_str(&trace_parent).unwrap(), - ); - } + llm_route_request_headers.insert( + header::HeaderName::from_static(TRACE_PARENT_HEADER), + header::HeaderValue::from_str(traceparent).unwrap(), + ); - if let Some(request_id) = request_id { - llm_route_request_headers.insert( - header::HeaderName::from_static(REQUEST_ID_HEADER), - header::HeaderValue::from_str(&request_id).unwrap(), - ); - } + llm_route_request_headers.insert( + header::HeaderName::from_static(REQUEST_ID_HEADER), + header::HeaderValue::from_str(request_id).unwrap(), + ); llm_route_request_headers.insert( header::HeaderName::from_static("model"),