refactor: carry normalized client session affinity

apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/payload.rs

@@ -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)),
                 original_request_body_json: Some(body_json),
                 original_request_body_base64: None,
+                client_session_affinity: input.client_session_affinity.as_ref(),
                 client_requested_stream: body_json
                     .get("stream")
                     .and_then(serde_json::Value::as_bool)

apps/aether-gateway/src/ai_serving/planner/report_context.rs

@@ -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,
     AiExecutionReportContextParts, AiRequestOrigin,
 };
-use aether_scheduler_core::SchedulerRankingOutcome;
+use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
 use serde_json::{Map, Value};
 
 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;
 
 pub(crate) struct LocalExecutionReportContextParts<'a> {
@@ -40,6 +43,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
     pub(crate) request_origin: Option<RequestOrigin>,
     pub(crate) original_request_body_json: Option<&'a Value>,
     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) upstream_is_stream: 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_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 {
         auth_context: parts.auth_context,
@@ -98,7 +112,7 @@ pub(crate) fn build_local_execution_report_context(
         upstream_is_stream: parts.upstream_is_stream,
         has_envelope: parts.has_envelope,
         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 {
     use std::collections::BTreeMap;
 
+    use aether_scheduler_core::ClientSessionAffinity;
     use serde_json::{json, Map, Value};
 
     use super::{
@@ -154,7 +169,7 @@ mod tests {
     }
 
     #[test]
-    fn local_execution_report_context_records_request_origin() {
+    fn local_execution_report_context_records_request_origin_and_session_affinity() {
         let auth_context = ExecutionRuntimeAuthContext {
             user_id: "user-1".to_string(),
             api_key_id: "api-key-1".to_string(),
@@ -166,6 +181,10 @@ mod tests {
         };
         let original_headers = http::HeaderMap::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 =
             build_local_execution_report_context(LocalExecutionReportContextParts {
@@ -199,6 +218,7 @@ mod tests {
                 }),
                 original_request_body_json: Some(&json!({"model": "gpt-5"})),
                 original_request_body_base64: None,
+                client_session_affinity: Some(&client_session_affinity),
                 client_requested_stream: false,
                 upstream_is_stream: false,
                 has_envelope: false,
@@ -214,5 +234,12 @@ mod tests {
             report_context["user_agent"],
             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"
+            })
+        );
     }
 }

apps/aether-gateway/src/ai_serving/planner/specialized/files/decision.rs

@@ -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)),
         original_request_body_json: Some(body_json),
         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,
         upstream_is_stream: spec_metadata.require_streaming,
         has_envelope: false,

apps/aether-gateway/src/ai_serving/planner/specialized/image/decision.rs

@@ -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)),
         original_request_body_json: Some(body_json),
         original_request_body_base64: body_base64,
+        client_session_affinity: input.client_session_affinity.as_ref(),
         client_requested_stream: spec_metadata.require_streaming,
         upstream_is_stream,
         has_envelope: false,

apps/aether-gateway/src/ai_serving/planner/specialized/video/decision.rs

@@ -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)),
         original_request_body_json: Some(body_json),
         original_request_body_base64: None,
+        client_session_affinity: input.client_session_affinity.as_ref(),
         client_requested_stream: false,
         upstream_is_stream: false,
         has_envelope: false,

apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs

@@ -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)),
                 original_request_body_json: Some(body_json),
                 original_request_body_base64: None,
+                client_session_affinity: input.client_session_affinity.as_ref(),
                 client_requested_stream: body_json
                     .get("stream")
                     .and_then(serde_json::Value::as_bool)

apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs

@@ -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)),
                 original_request_body_json: Some(body_json),
                 original_request_body_base64: None,
+                client_session_affinity: input.client_session_affinity.as_ref(),
                 client_requested_stream: body_json
                     .get("stream")
                     .and_then(serde_json::Value::as_bool)

apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/payload.rs

@@ -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)),
                 original_request_body_json: Some(body_json),
                 original_request_body_base64: None,
+                client_session_affinity: input.client_session_affinity.as_ref(),
                 client_requested_stream: body_json
                     .get("stream")
                     .and_then(serde_json::Value::as_bool)

apps/aether-gateway/src/client_session_affinity.rs

@@ -1,10 +1,11 @@
 use aether_scheduler_core::ClientSessionAffinity;
-use serde_json::Value;
+use serde_json::{Map, Value};
 
 use crate::headers::header_value_str;
 
 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 CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD: &str = "client_session_affinity";
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) enum ClientSessionSignalSource {
@@ -45,7 +46,11 @@ impl ClientSessionScope {
             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> {
@@ -113,6 +118,53 @@ pub(crate) fn client_session_scope_from_parts(
     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 {
     for adapter in specific_client_session_scope_adapters() {
         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();
-    match agent_id.map(str::trim).filter(|value| !value.is_empty()) {
-        Some(agent_id) => format!("session={root_session};agent={agent_id}"),
-        None => format!("session={root_session}"),
+    let mut parts = Vec::new();
+    if let Some(account_hint) = account_hint
+        .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> {
@@ -382,9 +446,11 @@ fn header_contains(headers: &http::HeaderMap, key: &str, needle: &str) -> bool {
 #[cfg(test)]
 mod tests {
     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,
     };
+    use aether_scheduler_core::ClientSessionAffinity;
     use http::{HeaderMap, HeaderValue};
     use serde_json::json;
 
@@ -479,10 +545,25 @@ mod tests {
         assert_eq!(affinity.client_family.as_deref(), Some("codex"));
         assert_eq!(
             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]
     fn claude_code_adapter_extracts_session_header() {
         let mut headers = HeaderMap::new();

apps/aether-gateway/src/orchestration/effects.rs

@@ -19,6 +19,9 @@ use super::{
     project_local_key_circuit_open, project_local_success_health, LocalFailoverClassification,
 };
 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::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
 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> {
     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 body_json = report_context
         .get("original_request_body")
@@ -823,9 +832,13 @@ mod tests {
             "api_key_id": "api-key-1",
             "client_api_format": "openai:chat",
             "model": "gpt-5",
+            "client_session_affinity": {
+                "client_family": "generic",
+                "session_key": "session=session-1;agent=coder"
+            },
             "original_headers": {
-                "x-aether-session-id": "session-1",
-                "x-aether-agent-id": "coder"
+                "x-aether-session-id": "raw-session",
+                "x-aether-agent-id": "raw-agent"
             },
             "original_request_body": {
                 "model": "gpt-5"