ensure that request id is consistent

crates/brightstaff/src/handlers/agent_chat_completions.rs

@@ -157,7 +157,33 @@ async fn handle_agent_chat(
         .strip_prefix("/agents")
         .unwrap()
         .to_string();
-    let request_headers = request.headers().clone();
+
+    let (request_headers, request_id) = {
+        let mut headers = request.headers().clone();
+        headers.remove(common::consts::ENVOY_ORIGINAL_PATH_HEADER);
+
+        if !headers.contains_key(common::consts::REQUEST_ID_HEADER) {
+            let request_id = uuid::Uuid::new_v4().to_string();
+            info!(
+                "Request id not found in headers, generated new request id: {}",
+                request_id
+            );
+            headers.insert(
+                common::consts::REQUEST_ID_HEADER,
+                hyper::header::HeaderValue::from_str(&request_id).unwrap(),
+            );
+        }
+        let request_id = headers
+            .get(common::consts::REQUEST_ID_HEADER)
+            .and_then(|v| v.to_str().ok())
+            .unwrap()
+            .to_string();
+
+        (headers, request_id)
+    };
+
+    info!("Processing request with Request ID: {}", request_id);
+
     let chat_request_bytes = request.collect().await?.to_bytes();
 
     debug!(

crates/brightstaff/src/handlers/pipeline_processor.rs

@@ -386,7 +386,7 @@ impl PipelineProcessor {
     async fn send_mcp_request(
         &self,
         json_rpc_request: &JsonRpcRequest,
-        headers: HeaderMap,
+        headers: &HeaderMap,
         agent_id: &str,
     ) -> Result<reqwest::Response, PipelineError> {
         let request_body = serde_json::to_string(json_rpc_request)?;
@@ -399,7 +399,7 @@ impl PipelineProcessor {
         let response = self
             .client
             .post(format!("{}/mcp", self.url))
-            .headers(headers)
+            .headers(headers.clone())
             .body(request_body)
             .send()
             .await?;
@@ -443,7 +443,12 @@ impl PipelineProcessor {
             session_id.clone()
         } else {
             let session_id = self
-                .get_new_session_id(&agent.id, trace_id.clone(), filter_span_id.clone())
+                .get_new_session_id(
+                    &agent.id,
+                    trace_id.clone(),
+                    filter_span_id.clone(),
+                    request_headers,
+                )
                 .await;
             self.agent_id_session_map
                 .insert(agent.id.clone(), session_id.clone());
@@ -476,7 +481,7 @@ impl PipelineProcessor {
         let start_instant = Instant::now();
 
         let response = self
-            .send_mcp_request(&json_rpc_request, agent_headers, &agent.id)
+            .send_mcp_request(&json_rpc_request, &agent_headers, &agent.id)
             .await?;
         let http_status = response.status();
         let response_bytes = response.bytes().await?;
@@ -608,6 +613,7 @@ impl PipelineProcessor {
         session_id: &str,
         trace_id: String,
         parent_span_id: String,
+        request_headers: &HeaderMap,
     ) -> Result<(), PipelineError> {
         let initialized_notification = JsonRpcNotification {
             jsonrpc: JSON_RPC_VERSION.to_string(),
@@ -619,7 +625,7 @@ impl PipelineProcessor {
         debug!("Sending initialized notification for agent {}", agent_id);
 
         let headers = self.build_mcp_headers(
-            &HeaderMap::new(),
+            request_headers,
             agent_id,
             Some(session_id),
             trace_id.clone(),
@@ -647,13 +653,14 @@ impl PipelineProcessor {
         agent_id: &str,
         trace_id: String,
         parent_span_id: String,
+        request_headers: &HeaderMap,
     ) -> String {
         info!("Initializing MCP session for agent {}", agent_id);
 
         let initialize_request = self.build_initialize_request();
         let headers = self
             .build_mcp_headers(
-                &HeaderMap::new(),
+                request_headers,
                 agent_id,
                 None,
                 trace_id.clone(),
@@ -661,8 +668,10 @@ impl PipelineProcessor {
             )
             .expect("Failed to build headers for initialization");
 
+        info!("Initialize request headers: {:?}", headers);
+
         let response = self
-            .send_mcp_request(&initialize_request, headers, agent_id)
+            .send_mcp_request(&initialize_request, &headers, agent_id)
             .await
             .expect("Failed to initialize MCP session");
 
@@ -686,6 +695,7 @@ impl PipelineProcessor {
             &session_id,
             trace_id.clone(),
             parent_span_id.clone(),
+            &headers,
         )
         .await
         .expect("Failed to send initialized notification");

crates/common/src/consts.rs

@@ -22,6 +22,7 @@ pub const X_ARCH_TOOL_CALL: &str = "x-arch-tool-call-message";
 pub const X_ARCH_FC_MODEL_RESPONSE: &str = "x-arch-fc-model-response";
 pub const ARCH_FC_MODEL_NAME: &str = "Arch-Function";
 pub const REQUEST_ID_HEADER: &str = "x-request-id";
+pub const ENVOY_ORIGINAL_PATH_HEADER: &str = "x-envoy-original-path";
 pub const TRACE_PARENT_HEADER: &str = "traceparent";
 pub const ARCH_INTERNAL_CLUSTER_NAME: &str = "arch_internal";
 pub const ARCH_UPSTREAM_HOST_HEADER: &str = "x-arch-upstream";