more fixes

crates/llm_gateway/src/stream_context.rs

@@ -159,11 +159,11 @@ impl StreamContext {
         // Tokenize and record token count.
         let token_count = tokenizer::token_count(model, json_string).unwrap_or(0);
 
+        trace!("Recorded input token count: {}", token_count);
         // Record the token count to metrics.
         self.metrics
             .input_sequence_length
             .record(token_count as u64);
-        trace!("Recorded input token count: {}", token_count);
 
         // Check if rate limiting needs to be applied.
         if let Some(selector) = self.ratelimit_selector.take() {

crates/llm_gateway/tests/integration.rs

@@ -20,6 +20,11 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) {
         .call_proxy_on_request_headers(http_context, 0, false)
         .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path"))
         .returning(Some("/v1/chat/completions"))
+        .expect_get_header_map_value(
+            Some(MapType::HttpRequestHeaders),
+            Some("x-arch-llm-provider"),
+        )
+        .returning(None)
         .expect_get_header_map_value(
             Some(MapType::HttpRequestHeaders),
             Some("x-arch-llm-provider-hint"),
@@ -36,6 +41,7 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) {
             Some("Authorization"),
             Some("Bearer secret_key"),
         )
+        .expect_remove_header_map_value(Some(MapType::HttpRequestHeaders), Some("content-length"))
         .expect_get_header_map_value(
             Some(MapType::HttpRequestHeaders),
             Some("x-arch-llm-provider-hint"),
@@ -48,8 +54,6 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) {
         .returning(Some("selector-key"))
         .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("selector-key"))
         .returning(Some("selector-value"))
-        .expect_get_header_map_pairs(Some(MapType::HttpRequestHeaders))
-        .returning(None)
         .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path"))
         .returning(Some("/v1/chat/completions"))
         .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("x-request-id"))
@@ -223,8 +227,8 @@ fn llm_gateway_successful_request_to_open_ai_chat_completions() {
         .returning(Some(chat_completions_request_body))
         .expect_log(Some(LogLevel::Trace), None)
         .expect_log(Some(LogLevel::Trace), None)
-        .expect_metric_record("input_sequence_length", 21)
         .expect_log(Some(LogLevel::Trace), None)
+        .expect_metric_record("input_sequence_length", 21)
         .expect_log(Some(LogLevel::Debug), None)
         .expect_log(Some(LogLevel::Debug), None)
         .expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None)
@@ -266,7 +270,7 @@ fn llm_gateway_bad_request_to_open_ai_chat_completions() {
     {\
         \"messages\": [\
         {\
-            \"role\": \"system\",\
+            \"role\": \"system\"\
         },\
         {\
             \"role\": \"user\",\
@@ -283,14 +287,19 @@ fn llm_gateway_bad_request_to_open_ai_chat_completions() {
         )
         .expect_get_buffer_bytes(Some(BufferType::HttpRequestBody))
         .returning(Some(incomplete_chat_completions_request_body))
-        .expect_log(Some(LogLevel::Debug), None)
+        .expect_log(Some(LogLevel::Trace), None)
         .expect_send_local_response(
             Some(StatusCode::BAD_REQUEST.as_u16().into()),
             None,
             None,
             None,
         )
-        .execute_and_expect(ReturnType::Action(Action::Pause))
+        .expect_log(Some(LogLevel::Trace), None)
+        .expect_log(Some(LogLevel::Trace), None)
+        .expect_metric_record("input_sequence_length", 14)
+        .expect_log(Some(LogLevel::Debug), None)
+        .expect_log(Some(LogLevel::Debug), None)
+        .execute_and_expect(ReturnType::Action(Action::Continue))
         .unwrap();
 }
 

crates/prompt_gateway/tests/integration.rs

@@ -352,10 +352,10 @@ fn prompt_gateway_request_to_llm_gateway() {
                     tool_type: ToolType::Function,
                     function: FunctionCallDetail {
                         name: String::from("weather_forecast"),
-                        arguments: HashMap::from([(
+                        arguments: Some(HashMap::from([(
                             String::from("city"),
                             Value::String(String::from("seattle")),
-                        )]),
+                        )])),
                     },
                 }]),
                 model: None,
@@ -381,8 +381,8 @@ fn prompt_gateway_request_to_llm_gateway() {
         .expect_log(Some(LogLevel::Debug), None)
         .expect_log(Some(LogLevel::Trace), None)
         .expect_log(Some(LogLevel::Debug), None)
-        .expect_log(Some(LogLevel::Trace), None)
         .expect_log(Some(LogLevel::Debug), None)
+        .expect_log(Some(LogLevel::Trace), None)
         .expect_http_call(
             Some("arch_internal"),
             Some(vec![
@@ -400,6 +400,7 @@ fn prompt_gateway_request_to_llm_gateway() {
         )
         .returning(Some(2))
         .expect_metric_increment("active_http_calls", 1)
+        .expect_log(Some(LogLevel::Trace), None)
         .execute_and_expect(ReturnType::None)
         .unwrap();
 
@@ -499,10 +500,10 @@ fn prompt_gateway_request_no_intent_match() {
                     tool_type: ToolType::Function,
                     function: FunctionCallDetail {
                         name: String::from("weather_forecast"),
-                        arguments: HashMap::from([(
+                        arguments: Some(HashMap::from([(
                             String::from("city"),
                             Value::String(String::from("seattle")),
-                        )]),
+                        )])),
                     },
                 }]),
                 model: None,
@@ -655,10 +656,10 @@ fn prompt_gateway_request_no_intent_match_default_target() {
                     tool_type: ToolType::Function,
                     function: FunctionCallDetail {
                         name: String::from("weather_forecast"),
-                        arguments: HashMap::from([(
+                        arguments: Some(HashMap::from([(
                             String::from("city"),
                             Value::String(String::from("seattle")),
-                        )]),
+                        )])),
                     },
                 }]),
                 model: None,

demos/shared/test_runner/run_demo_tests.sh

@@ -15,7 +15,7 @@ do
   echo "starting hurl tests"
   pwd
   ls
-  hurl --test hurl_tests
+  hurl --test hurl_tests/*.hurl
   echo "stopping docker containers and archgw"
   archgw down
   docker compose down -v