refactor: carry normalized client session affinity

This commit is contained in:
fawney19
2026-05-05 21:34:17 +08:00
parent ddf3fb6f63
commit 98421126f2
10 changed files with 141 additions and 13 deletions

View File

@@ -115,6 +115,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)), request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json), original_request_body_json: Some(body_json),
original_request_body_base64: None, original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: body_json client_requested_stream: body_json
.get("stream") .get("stream")
.and_then(serde_json::Value::as_bool) .and_then(serde_json::Value::as_bool)

View File

@@ -6,10 +6,13 @@ use aether_ai_serving::{
provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type, provider_stream_event_api_format_for_provider_type as ai_provider_stream_event_api_format_for_provider_type,
AiExecutionReportContextParts, AiRequestOrigin, AiExecutionReportContextParts, AiRequestOrigin,
}; };
use aether_scheduler_core::SchedulerRankingOutcome; use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
use serde_json::{Map, Value}; use serde_json::{Map, Value};
use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin}; use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin};
use crate::client_session_affinity::{
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
use crate::orchestration::ExecutionAttemptIdentity; use crate::orchestration::ExecutionAttemptIdentity;
pub(crate) struct LocalExecutionReportContextParts<'a> { pub(crate) struct LocalExecutionReportContextParts<'a> {
@@ -40,6 +43,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) request_origin: Option<RequestOrigin>, pub(crate) request_origin: Option<RequestOrigin>,
pub(crate) original_request_body_json: Option<&'a Value>, pub(crate) original_request_body_json: Option<&'a Value>,
pub(crate) original_request_body_base64: Option<&'a str>, pub(crate) original_request_body_base64: Option<&'a str>,
pub(crate) client_session_affinity: Option<&'a ClientSessionAffinity>,
pub(crate) client_requested_stream: bool, pub(crate) client_requested_stream: bool,
pub(crate) upstream_is_stream: bool, pub(crate) upstream_is_stream: bool,
pub(crate) has_envelope: bool, pub(crate) has_envelope: bool,
@@ -61,6 +65,16 @@ pub(crate) fn build_local_execution_report_context(
parts.original_request_body_json, parts.original_request_body_json,
parts.original_request_body_base64, parts.original_request_body_base64,
); );
let mut extra_fields = parts.extra_fields;
if let Some(value) = parts
.client_session_affinity
.and_then(client_session_affinity_report_context_value)
{
extra_fields.insert(
CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD.to_string(),
value,
);
}
build_ai_execution_report_context(AiExecutionReportContextParts { build_ai_execution_report_context(AiExecutionReportContextParts {
auth_context: parts.auth_context, auth_context: parts.auth_context,
@@ -98,7 +112,7 @@ pub(crate) fn build_local_execution_report_context(
upstream_is_stream: parts.upstream_is_stream, upstream_is_stream: parts.upstream_is_stream,
has_envelope: parts.has_envelope, has_envelope: parts.has_envelope,
needs_conversion: parts.needs_conversion, needs_conversion: parts.needs_conversion,
extra_fields: parts.extra_fields, extra_fields,
}) })
} }
@@ -119,6 +133,7 @@ pub(crate) fn insert_provider_stream_event_api_format(
mod tests { mod tests {
use std::collections::BTreeMap; use std::collections::BTreeMap;
use aether_scheduler_core::ClientSessionAffinity;
use serde_json::{json, Map, Value}; use serde_json::{json, Map, Value};
use super::{ use super::{
@@ -154,7 +169,7 @@ mod tests {
} }
#[test] #[test]
fn local_execution_report_context_records_request_origin() { fn local_execution_report_context_records_request_origin_and_session_affinity() {
let auth_context = ExecutionRuntimeAuthContext { let auth_context = ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(), user_id: "user-1".to_string(),
api_key_id: "api-key-1".to_string(), api_key_id: "api-key-1".to_string(),
@@ -166,6 +181,10 @@ mod tests {
}; };
let original_headers = http::HeaderMap::new(); let original_headers = http::HeaderMap::new();
let provider_request_headers = BTreeMap::new(); let provider_request_headers = BTreeMap::new();
let client_session_affinity = ClientSessionAffinity::new(
Some("codex".to_string()),
Some("account=account-1;session=session-1".to_string()),
);
let report_context = let report_context =
build_local_execution_report_context(LocalExecutionReportContextParts { build_local_execution_report_context(LocalExecutionReportContextParts {
@@ -199,6 +218,7 @@ mod tests {
}), }),
original_request_body_json: Some(&json!({"model": "gpt-5"})), original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None, original_request_body_base64: None,
client_session_affinity: Some(&client_session_affinity),
client_requested_stream: false, client_requested_stream: false,
upstream_is_stream: false, upstream_is_stream: false,
has_envelope: false, has_envelope: false,
@@ -214,5 +234,12 @@ mod tests {
report_context["user_agent"], report_context["user_agent"],
Value::String("Claude-Code/1.0".to_string()) Value::String("Claude-Code/1.0".to_string())
); );
assert_eq!(
report_context["client_session_affinity"],
json!({
"client_family": "codex",
"session_key": "account=account-1;session=session-1"
})
);
} }
} }

View File

@@ -97,6 +97,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)), request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json), original_request_body_json: Some(body_json),
original_request_body_base64: resolved.provider_request_body_base64.as_deref(), original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming, client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream: spec_metadata.require_streaming, upstream_is_stream: spec_metadata.require_streaming,
has_envelope: false, has_envelope: false,

View File

@@ -92,6 +92,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)), request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json), original_request_body_json: Some(body_json),
original_request_body_base64: body_base64, original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: spec_metadata.require_streaming, client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream, upstream_is_stream,
has_envelope: false, has_envelope: false,

View File

@@ -78,6 +78,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)), request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json), original_request_body_json: Some(body_json),
original_request_body_base64: None, original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: false, client_requested_stream: false,
upstream_is_stream: false, upstream_is_stream: false,
has_envelope: false, has_envelope: false,

View File

@@ -123,6 +123,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)), request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json), original_request_body_json: Some(body_json),
original_request_body_base64: None, original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: body_json client_requested_stream: body_json
.get("stream") .get("stream")
.and_then(serde_json::Value::as_bool) .and_then(serde_json::Value::as_bool)

View File

@@ -110,6 +110,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)), request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json), original_request_body_json: Some(body_json),
original_request_body_base64: None, original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: body_json client_requested_stream: body_json
.get("stream") .get("stream")
.and_then(serde_json::Value::as_bool) .and_then(serde_json::Value::as_bool)

View File

@@ -108,6 +108,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)), request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json), original_request_body_json: Some(body_json),
original_request_body_base64: None, original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
client_requested_stream: body_json client_requested_stream: body_json
.get("stream") .get("stream")
.and_then(serde_json::Value::as_bool) .and_then(serde_json::Value::as_bool)

View File

@@ -1,10 +1,11 @@
use aether_scheduler_core::ClientSessionAffinity; use aether_scheduler_core::ClientSessionAffinity;
use serde_json::Value; use serde_json::{Map, Value};
use crate::headers::header_value_str; use crate::headers::header_value_str;
pub(crate) const AETHER_SESSION_ID_HEADER: &str = "x-aether-session-id"; pub(crate) const AETHER_SESSION_ID_HEADER: &str = "x-aether-session-id";
pub(crate) const AETHER_AGENT_ID_HEADER: &str = "x-aether-agent-id"; pub(crate) const AETHER_AGENT_ID_HEADER: &str = "x-aether-agent-id";
pub(crate) const CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD: &str = "client_session_affinity";
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ClientSessionSignalSource { pub(crate) enum ClientSessionSignalSource {
@@ -45,7 +46,11 @@ impl ClientSessionScope {
return None; return None;
} }
Some(normalize_session_key(session_id, self.agent_id.as_deref())) Some(normalize_session_key(
self.account_hint.as_deref(),
session_id,
self.agent_id.as_deref(),
))
} }
pub(crate) fn scheduler_affinity(&self) -> Option<ClientSessionAffinity> { pub(crate) fn scheduler_affinity(&self) -> Option<ClientSessionAffinity> {
@@ -113,6 +118,53 @@ pub(crate) fn client_session_scope_from_parts(
client_session_scope_from_request(&parts.headers, body_json) client_session_scope_from_request(&parts.headers, body_json)
} }
pub(crate) fn client_session_affinity_report_context_value(
affinity: &ClientSessionAffinity,
) -> Option<Value> {
let session_key = affinity
.session_key
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())?;
let mut object = Map::new();
if let Some(client_family) = affinity
.client_family
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
object.insert(
"client_family".to_string(),
Value::String(client_family.to_ascii_lowercase()),
);
}
object.insert(
"session_key".to_string(),
Value::String(session_key.to_string()),
);
Some(Value::Object(object))
}
pub(crate) fn client_session_affinity_from_report_context_value(
value: Option<&Value>,
) -> Option<ClientSessionAffinity> {
let object = value?.as_object()?;
let session_key = object
.get("session_key")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?
.to_string();
let client_family = object
.get("client_family")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_ascii_lowercase);
Some(ClientSessionAffinity::new(client_family, Some(session_key)))
}
fn detect_client_family(request: &ClientSessionRequest<'_>) -> String { fn detect_client_family(request: &ClientSessionRequest<'_>) -> String {
for adapter in specific_client_session_scope_adapters() { for adapter in specific_client_session_scope_adapters() {
if adapter.detect(request) { if adapter.detect(request) {
@@ -350,12 +402,24 @@ fn explicit_aether_session_scope(
)) ))
} }
fn normalize_session_key(root_session: &str, agent_id: Option<&str>) -> String { fn normalize_session_key(
account_hint: Option<&str>,
root_session: &str,
agent_id: Option<&str>,
) -> String {
let root_session = root_session.trim(); let root_session = root_session.trim();
match agent_id.map(str::trim).filter(|value| !value.is_empty()) { let mut parts = Vec::new();
Some(agent_id) => format!("session={root_session};agent={agent_id}"), if let Some(account_hint) = account_hint
None => format!("session={root_session}"), .map(str::trim)
.filter(|value| !value.is_empty())
{
parts.push(format!("account={account_hint}"));
} }
parts.push(format!("session={root_session}"));
if let Some(agent_id) = agent_id.map(str::trim).filter(|value| !value.is_empty()) {
parts.push(format!("agent={agent_id}"));
}
parts.join(";")
} }
fn value_at_paths<'a>(body: &'a Value, paths: &[&[&str]]) -> Option<&'a str> { fn value_at_paths<'a>(body: &'a Value, paths: &[&[&str]]) -> Option<&'a str> {
@@ -382,9 +446,11 @@ fn header_contains(headers: &http::HeaderMap, key: &str, needle: &str) -> bool {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::{ use super::{
client_session_affinity_from_request, client_session_scope_from_request, client_session_affinity_from_report_context_value, client_session_affinity_from_request,
client_session_affinity_report_context_value, client_session_scope_from_request,
ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER, ClientSessionSignalSource, AETHER_AGENT_ID_HEADER, AETHER_SESSION_ID_HEADER,
}; };
use aether_scheduler_core::ClientSessionAffinity;
use http::{HeaderMap, HeaderValue}; use http::{HeaderMap, HeaderValue};
use serde_json::json; use serde_json::json;
@@ -479,10 +545,25 @@ mod tests {
assert_eq!(affinity.client_family.as_deref(), Some("codex")); assert_eq!(affinity.client_family.as_deref(), Some("codex"));
assert_eq!( assert_eq!(
affinity.session_key.as_deref(), affinity.session_key.as_deref(),
Some("session=prompt-session-1") Some("account=account-1;session=prompt-session-1")
); );
} }
#[test]
fn report_context_round_trips_normalized_session_affinity() {
let affinity = ClientSessionAffinity::new(
Some("codex".to_string()),
Some("account=account-1;session=session-1".to_string()),
);
let value = client_session_affinity_report_context_value(&affinity)
.expect("report context value should build");
let parsed = client_session_affinity_from_report_context_value(Some(&value))
.expect("report context value should parse");
assert_eq!(parsed, affinity);
}
#[test] #[test]
fn claude_code_adapter_extracts_session_header() { fn claude_code_adapter_extracts_session_header() {
let mut headers = HeaderMap::new(); let mut headers = HeaderMap::new();

View File

@@ -19,6 +19,9 @@ use super::{
project_local_key_circuit_open, project_local_success_health, LocalFailoverClassification, project_local_key_circuit_open, project_local_success_health, LocalFailoverClassification,
}; };
use crate::ai_serving::extract_pool_sticky_session_token; use crate::ai_serving::extract_pool_sticky_session_token;
use crate::client_session_affinity::{
client_session_affinity_from_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
use crate::clock::current_unix_secs; use crate::clock::current_unix_secs;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value; use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::handlers::shared::provider_pool::{ use crate::handlers::shared::provider_pool::{
@@ -163,6 +166,12 @@ fn local_scheduler_affinity_cache_key(report_context: Option<&Value>) -> Option<
fn local_client_session_affinity(report_context: Option<&Value>) -> Option<ClientSessionAffinity> { fn local_client_session_affinity(report_context: Option<&Value>) -> Option<ClientSessionAffinity> {
let report_context = report_context?; let report_context = report_context?;
if let Some(affinity) = client_session_affinity_from_report_context_value(
report_context.get(CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD),
) {
return Some(affinity);
}
let headers = header_map_from_report_context(report_context.get("original_headers")); let headers = header_map_from_report_context(report_context.get("original_headers"));
let body_json = report_context let body_json = report_context
.get("original_request_body") .get("original_request_body")
@@ -823,9 +832,13 @@ mod tests {
"api_key_id": "api-key-1", "api_key_id": "api-key-1",
"client_api_format": "openai:chat", "client_api_format": "openai:chat",
"model": "gpt-5", "model": "gpt-5",
"client_session_affinity": {
"client_family": "generic",
"session_key": "session=session-1;agent=coder"
},
"original_headers": { "original_headers": {
"x-aether-session-id": "session-1", "x-aether-session-id": "raw-session",
"x-aether-agent-id": "coder" "x-aether-agent-id": "raw-agent"
}, },
"original_request_body": { "original_request_body": {
"model": "gpt-5" "model": "gpt-5"