Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions crates/switchyard-server/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 4 additions & 3 deletions crates/switchyard-server/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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| {
Expand Down
22 changes: 16 additions & 6 deletions crates/switchyard-server/src/routing_log.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down Expand Up @@ -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)
}),
}
}
}
Expand Down
89 changes: 84 additions & 5 deletions crates/switchyard-server/tests/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1489,8 +1489,10 @@ async fn all_inbound_formats_run_libsy_and_return_the_caller_format() -> TestRes
Ok(())
}

// Normalized metadata is authoritative when both ID forms are present;
// legacy-only callers remain supported for backward compatibility.
#[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");
Expand All @@ -1502,7 +1504,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"}]
Expand All @@ -1513,19 +1516,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()
Expand All @@ -1534,6 +1571,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]
Expand Down
Loading