From 51b191ad2c6ea0b58c11e4d27c05152051d495a4 Mon Sep 17 00:00:00 2001
From: RWDai <27391645+RWDai@users.noreply.github.com>
Date: Tue, 5 May 2026 11:23:54 +0800
Subject: [PATCH] feat: wire OpenAI chat session affinity

---
 .../planner/standard/openai/chat/decision/support.rs          | 3 +++
 .../src/ai_serving/planner/standard/openai/chat/mod.rs        | 4 ++--
 .../planner/standard/openai/chat/plans/candidates.rs          | 1 +
 .../ai_serving/planner/standard/openai/chat/plans/resolve.rs  | 3 +++
 .../ai_serving/planner/standard/openai/chat/plans/stream.rs   | 4 ++--
 .../src/ai_serving/planner/standard/openai/chat/plans/sync.rs | 4 ++--
 6 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs
index e59c9176a..9cf291186 100644
--- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs
+++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs
@@ -138,6 +138,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
         "openai:chat",
         Some(&input.requested_model),
         Some(&input.auth_snapshot),
+        input.client_session_affinity.as_ref(),
         input.required_capabilities.as_ref(),
         sticky_session_token.as_deref(),
         input.request_auth_channel.as_deref(),
@@ -217,6 +218,7 @@ pub(crate) async fn build_local_openai_chat_candidate_attempt_source<'a>(
         "openai:chat",
         Some(&input.requested_model),
         Some(&input.auth_snapshot),
+        input.client_session_affinity.as_ref(),
         input.required_capabilities.as_ref(),
         sticky_session_token.as_deref(),
         input.request_auth_channel.as_deref(),
@@ -294,6 +296,7 @@ pub(crate) async fn build_lazy_local_openai_chat_candidate_attempt_source<'a>(
         &input.requested_model,
         require_streaming,
         &input.auth_snapshot,
+        input.client_session_affinity.as_ref(),
         input.required_capabilities.as_ref(),
         sticky_session_token.as_deref(),
         input.request_auth_channel.as_deref(),
diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
index 1385db093..60db901aa 100644
--- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
+++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
@@ -139,7 +139,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
     }
 
     let Some(input) = resolve_local_openai_chat_decision_input(
-        state, trace_id, decision, body_json, plan_kind, false,
+        state, parts, trace_id, decision, body_json, plan_kind, false,
     )
     .await
     else {
@@ -205,7 +205,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
     }
 
     let Some(input) = resolve_local_openai_chat_decision_input(
-        state, trace_id, decision, body_json, plan_kind, false,
+        state, parts, trace_id, decision, body_json, plan_kind, false,
     )
     .await
     else {
diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs
index e1f574d63..a1985fd16 100644
--- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs
+++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/candidates.rs
@@ -26,6 +26,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
         require_streaming,
         input.required_capabilities.as_ref(),
         &input.auth_snapshot,
+        input.client_session_affinity.as_ref(),
         false,
         LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
     )
diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs
index 926533cc6..39c572dd5 100644
--- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs
+++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs
@@ -7,10 +7,12 @@ use crate::ai_serving::planner::decision_input::{
     build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
 };
 use crate::ai_serving::resolve_local_decision_execution_runtime_auth_context;
+use crate::client_session_affinity::client_session_affinity_from_parts;
 use crate::AppState;
 
 pub(crate) async fn resolve_local_openai_chat_decision_input(
     state: &AppState,
+    parts: &http::request::Parts,
     trace_id: &str,
     decision: &GatewayControlDecision,
     body_json: &serde_json::Value,
@@ -106,5 +108,6 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
 
     let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
     input.request_auth_channel = decision.request_auth_channel.clone();
+    input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
     Some(input)
 }
diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
index acc8fb54a..96226a381 100644
--- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
+++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
@@ -43,7 +43,7 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
     }
 
     let Some(input) = resolve_local_openai_chat_decision_input(
-        state, trace_id, decision, body_json, plan_kind, true,
+        state, parts, trace_id, decision, body_json, plan_kind, true,
     )
     .await
     else {
@@ -163,7 +163,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
     }
 
     let Some(input) = resolve_local_openai_chat_decision_input(
-        state, trace_id, decision, body_json, plan_kind, true,
+        state, parts, trace_id, decision, body_json, plan_kind, true,
     )
     .await
     else {
diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
index 032e5aa3f..620b9bf8f 100644
--- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
+++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
@@ -52,7 +52,7 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
     }
 
     let Some(input) = resolve_local_openai_chat_decision_input(
-        state, trace_id, decision, body_json, plan_kind, true,
+        state, parts, trace_id, decision, body_json, plan_kind, true,
     )
     .await
     else {
@@ -176,7 +176,7 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
     }
 
     let Some(input) = resolve_local_openai_chat_decision_input(
-        state, trace_id, decision, body_json, plan_kind, true,
+        state, parts, trace_id, decision, body_json, plan_kind, true,
     )
     .await
     else {