fix(server): use normalized session IDs in routing stats

crates/switchyard-server/README.md

@@ -88,8 +88,10 @@ for equal weighting. The optional `seed` reproduces the selection sequence for t
 Pass `--routing-log-file PATH` to append one JSON record after each completed routed response.
 Streaming responses are recorded after the stream drains. When enabled,
 `GET /v1/routing/session-stats?session_id=ID` rescans the durable log and returns call and token
-totals for that exact `proxy_x_session_id`, grouped by served model. The endpoint returns `404` when
-the session has no records and is not registered when routing logging is disabled.
+totals for that normalized session ID, normally supplied as `x-switchyard-session-id`, grouped by
+served model. The legacy `proxy_x_session_id` remains a fallback when no normalized session ID is
+present. The endpoint returns `404` when the session has no records and is not registered when
+routing logging is disabled.
 
 An `llm_classifier` route sends each task to `classifier_target` for a capability verdict, then
 routes to `weak_target` or `strong_target`. Beyond the three targets it accepts these keys; only

crates/switchyard-server/src/lib.rs

@@ -580,11 +580,11 @@ async fn handle_endpoint_inner(
     body: std::result::Result<Json<Value>, JsonRejection>,
     wire_format: WireFormat,
 ) -> Response {
+    let metadata = metadata_from_headers(headers);
     let routing_log_context = state
         .routing_log
         .as_ref()
-        .map(|_| routing_log::RoutingLogContext::from_headers(&headers));
-    let metadata = metadata_from_headers(headers);
+        .map(|_| routing_log::RoutingLogContext::from_metadata(&metadata));
     let request_log = RequestLogContext {
         started: started.0,
         wire_format,
@@ -1318,7 +1318,8 @@ mod tests {
         let log = SharedRoutingLog::new(dir.path().join("routing.jsonl")).expect("routing log");
         let mut headers = HeaderMap::new();
         headers.insert("proxy_x_session_id", "session-1".parse().expect("header"));
-        let context = routing_log::RoutingLogContext::from_headers(&headers);
+        let metadata = metadata_from_headers(headers);
+        let context = routing_log::RoutingLogContext::from_metadata(&metadata);
         let observer = stats_observer(StatsAccumulator::default(), Some((log.clone(), context)));
 
         let call = |model: &str, is_answer_call: bool| {

crates/switchyard-server/src/routing_log.rs

@@ -12,12 +12,12 @@ use std::time::SystemTime;
 
 use humantime::format_rfc3339_millis;
 use serde::{Deserialize, Serialize};
-use switchyard_protocol::{ModelId, Usage};
+use switchyard_protocol::{Metadata, ModelId, Usage};
 
 use crate::usage_metrics::token_usage;
 use crate::{ServerError, ServerResult};
 
-const SESSION_ID_HEADER: &str = "proxy_x_session_id";
+const LEGACY_SESSION_ID_HEADER: &str = "proxy_x_session_id";
 const TASK_HEADER: &str = "x-switchyard-intake-task";
 const TRIAL_ID_HEADER: &str = "x-switchyard-trial-id";
 
@@ -104,11 +104,21 @@ pub(crate) struct RoutingLogContext {
 }
 
 impl RoutingLogContext {
-    pub(crate) fn from_headers(headers: &http::HeaderMap) -> Self {
+    /// Captures the normalized session ID, with the legacy log-only header as a fallback.
+    pub(crate) fn from_metadata(metadata: &Metadata) -> Self {
+        let headers = metadata.http_headers.as_ref();
         Self {
-            task: nonempty_header(headers, TASK_HEADER).map(|s| s.to_string()),
-            trial_id: nonempty_header(headers, TRIAL_ID_HEADER).map(|s| s.to_string()),
-            session_id: nonempty_header(headers, SESSION_ID_HEADER).map(|s| s.to_string()),
+            task: headers
+                .and_then(|headers| nonempty_header(headers, TASK_HEADER))
+                .map(str::to_string),
+            trial_id: headers
+                .and_then(|headers| nonempty_header(headers, TRIAL_ID_HEADER))
+                .map(str::to_string),
+            session_id: metadata.session_id.clone().or_else(|| {
+                headers
+                    .and_then(|headers| nonempty_header(headers, LEGACY_SESSION_ID_HEADER))
+                    .map(str::to_string)
+            }),
         }
     }
 }

crates/switchyard-server/tests/server.rs

@@ -1490,7 +1490,7 @@ async fn all_inbound_formats_run_libsy_and_return_the_caller_format() -> TestRes
 }
 
 #[tokio::test]
-async fn routing_log_exposes_session_stats() -> TestResult {
+async fn routing_log_prefers_canonical_and_preserves_legacy_fallback() -> TestResult {
     let upstream = MockUpstream::start().await?;
     let temp_dir = tempfile::tempdir()?;
     let log_path = temp_dir.path().join("routing.jsonl");
@@ -1502,7 +1502,8 @@ async fn routing_log_exposes_session_stats() -> TestResult {
         .method("POST")
         .uri("/v1/chat/completions")
         .header("content-type", "application/json")
-        .header("proxy_x_session_id", "session-1")
+        .header("x-switchyard-session-id", "canonical-session")
+        .header("proxy_x_session_id", "legacy-session")
         .body(Body::from(serde_json::to_vec(&json!({
             "model": ROUTE_MODEL,
             "messages": [{"role": "user", "content": "hello"}]
@@ -1513,19 +1514,53 @@ async fn routing_log_exposes_session_stats() -> TestResult {
     let stats = send(
         &app,
         "GET",
-        "/v1/routing/session-stats?session_id=session-1",
+        "/v1/routing/session-stats?session_id=canonical-session",
         None,
     )
-    .await?
-    .json()?;
+    .await?;
+    assert_eq!(stats.status, StatusCode::OK);
+    let stats = stats.json()?;
     assert_eq!(stats["total_calls"], 1);
     assert_eq!(stats["total_prompt_tokens"], 10);
     assert_eq!(stats["total_cached_tokens"], 7);
     assert_eq!(stats["models"]["model/a"]["completion_tokens"], 2);
 
+    let legacy = send(
+        &app,
+        "GET",
+        "/v1/routing/session-stats?session_id=legacy-session",
+        None,
+    )
+    .await?;
+    assert_eq!(legacy.status, StatusCode::NOT_FOUND);
+
+    let legacy_only = send_with_headers(
+        &app,
+        "POST",
+        "/v1/chat/completions",
+        Some(json!({
+            "model": ROUTE_MODEL,
+            "messages": [{"role": "user", "content": "hello"}]
+        })),
+        &[("proxy_x_session_id", "legacy-only-session")],
+    )
+    .await?;
+    assert_eq!(legacy_only.status, StatusCode::OK);
+
+    let legacy_stats = send(
+        &app,
+        "GET",
+        "/v1/routing/session-stats?session_id=legacy-only-session",
+        None,
+    )
+    .await?;
+    assert_eq!(legacy_stats.status, StatusCode::OK);
+    assert_eq!(legacy_stats.json()?["total_calls"], 1);
+
     let records = std::fs::read_to_string(log_path)?;
     let first: Value =
         serde_json::from_str(records.lines().next().ok_or("routing log was empty")?)?;
+    assert_eq!(first["session_id"], "canonical-session");
     assert!(
         first["ts"]
             .as_str()
@@ -1534,6 +1569,48 @@ async fn routing_log_exposes_session_stats() -> TestResult {
     Ok(())
 }
 
+#[tokio::test]
+async fn routing_log_keeps_the_canonical_session_id_until_a_stream_drains() -> TestResult {
+    let upstream = MockUpstream::start().await?;
+    let temp_dir = tempfile::tempdir()?;
+    let state = random_state(&upstream.base_url, &[(ROUTE_MODEL, &["model/a"])])?
+        .with_routing_log(temp_dir.path().join("routing.jsonl"))?;
+    let app = build_switchyard_router(state);
+
+    // `send_with_headers` collects the response body, so the stream wrapper reaches
+    // its terminal usage record before the stats query runs.
+    let response = send_with_headers(
+        &app,
+        "POST",
+        "/v1/chat/completions",
+        Some(json!({
+            "model": ROUTE_MODEL,
+            "messages": [{"role": "user", "content": "hello"}],
+            "stream": true
+        })),
+        &[("x-switchyard-session-id", "streaming-session")],
+    )
+    .await?;
+    assert_eq!(response.status, StatusCode::OK);
+    assert!(response.text()?.contains("data: [DONE]"));
+
+    let stats = send(
+        &app,
+        "GET",
+        "/v1/routing/session-stats?session_id=streaming-session",
+        None,
+    )
+    .await?;
+    assert_eq!(stats.status, StatusCode::OK);
+    let stats = stats.json()?;
+    assert_eq!(stats["total_calls"], 1);
+    assert_eq!(stats["total_prompt_tokens"], 12);
+    assert_eq!(stats["total_cached_tokens"], 7);
+    assert_eq!(stats["total_cache_creation_tokens"], 2);
+    assert_eq!(stats["total_completion_tokens"], 5);
+    Ok(())
+}
+
 // Overflow history is isolated per child, cleared with the session, and not retained when a
 // child lacks an agent ID.
 #[tokio::test]