mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 09:57:47 +08:00
feat(codex): stabilize identity across retries
This commit is contained in:
elky
@@ -23,6 +23,21 @@ pub(crate) struct ClientSessionScope {
|
||||
pub(crate) source: ClientSessionSignalSource,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub(crate) struct CodexRequestSignals {
|
||||
pub(crate) session_id: Option<String>,
|
||||
pub(crate) thread_id: Option<String>,
|
||||
pub(crate) turn_id: Option<String>,
|
||||
pub(crate) prompt_cache_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct CodexTurnMetadataSignals {
|
||||
session_id: Option<String>,
|
||||
thread_id: Option<String>,
|
||||
turn_id: Option<String>,
|
||||
}
|
||||
|
||||
impl ClientSessionScope {
|
||||
fn new(
|
||||
client_family: impl Into<String>,
|
||||
@@ -136,6 +151,13 @@ pub(crate) fn client_session_scope_from_request(
|
||||
.or_else(|| extract_scope_from_other_specific_adapters(&request, client_family.as_str()))
|
||||
}
|
||||
|
||||
pub(crate) fn codex_request_signals_from_request(
|
||||
headers: &http::HeaderMap,
|
||||
body_json: Option<&Value>,
|
||||
) -> CodexRequestSignals {
|
||||
extract_codex_request_signals(&ClientSessionRequest { headers, body_json })
|
||||
}
|
||||
|
||||
fn codex_search_session_scope(request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
|
||||
let session_id = request
|
||||
.body_json?
|
||||
@@ -408,29 +430,7 @@ impl ClientSessionScopeAdapter for CodexSessionScopeAdapter {
|
||||
}
|
||||
|
||||
fn extract_scope(&self, request: &ClientSessionRequest<'_>) -> Option<ClientSessionScope> {
|
||||
header_value_str(request.headers, "session-id")
|
||||
.or_else(|| header_value_str(request.headers, "thread-id"))
|
||||
.or_else(|| header_value_str(request.headers, "session_id"))
|
||||
.or_else(|| header_value_str(request.headers, "conversation_id"))
|
||||
.map(|root_session| {
|
||||
ClientSessionScope::new(
|
||||
self.family(),
|
||||
root_session,
|
||||
None,
|
||||
header_value_str(request.headers, "chatgpt-account-id"),
|
||||
ClientSessionSignalSource::Header,
|
||||
)
|
||||
})
|
||||
.or_else(|| {
|
||||
let body_session = GenericSessionScopeAdapter.extract_scope(request)?;
|
||||
Some(ClientSessionScope::new(
|
||||
self.family(),
|
||||
body_session.session_id,
|
||||
body_session.agent_id,
|
||||
header_value_str(request.headers, "chatgpt-account-id"),
|
||||
body_session.source,
|
||||
))
|
||||
})
|
||||
codex_request_session_scope_from_request(request)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -778,6 +778,162 @@ fn explicit_aether_session_scope(
|
||||
))
|
||||
}
|
||||
|
||||
fn extract_codex_request_signals(request: &ClientSessionRequest<'_>) -> CodexRequestSignals {
|
||||
let body_client_metadata = request
|
||||
.body_json
|
||||
.and_then(|body| body.get("client_metadata"))
|
||||
.and_then(Value::as_object);
|
||||
let body_turn_metadata = codex_turn_metadata_signals(
|
||||
body_client_metadata.and_then(|metadata| metadata.get("x-codex-turn-metadata")),
|
||||
);
|
||||
let header_turn_metadata = header_value_str(request.headers, "x-codex-turn-metadata")
|
||||
.map(|raw| parse_codex_turn_metadata(&raw))
|
||||
.unwrap_or_default();
|
||||
|
||||
let native_thread_id = header_value_str(request.headers, "thread-id")
|
||||
.or_else(|| {
|
||||
body_client_metadata
|
||||
.and_then(|metadata| value_at_map_path(metadata, "thread_id"))
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
.or_else(|| body_turn_metadata.thread_id.clone());
|
||||
let turn_id = body_client_metadata
|
||||
.and_then(|metadata| value_at_map_path(metadata, "turn_id"))
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| body_turn_metadata.turn_id.clone())
|
||||
.or_else(|| {
|
||||
request
|
||||
.body_json
|
||||
.and_then(|body| value_at_path(body, &["turn_id"]))
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
.or(header_turn_metadata.turn_id);
|
||||
let prompt_cache_key = request
|
||||
.body_json
|
||||
.and_then(|body| value_at_path(body, &["prompt_cache_key"]))
|
||||
.map(ToOwned::to_owned);
|
||||
let session_id =
|
||||
codex_request_session_scope(request, &body_turn_metadata).map(|scope| scope.session_id);
|
||||
let thread_id = native_thread_id.or_else(|| session_id.clone());
|
||||
|
||||
CodexRequestSignals {
|
||||
session_id,
|
||||
thread_id,
|
||||
turn_id,
|
||||
prompt_cache_key,
|
||||
}
|
||||
}
|
||||
|
||||
fn codex_request_session_scope_from_request(
|
||||
request: &ClientSessionRequest<'_>,
|
||||
) -> Option<ClientSessionScope> {
|
||||
let body_turn_metadata = codex_turn_metadata_signals(
|
||||
request
|
||||
.body_json
|
||||
.and_then(|body| body.get("client_metadata"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| metadata.get("x-codex-turn-metadata")),
|
||||
);
|
||||
codex_request_session_scope(request, &body_turn_metadata)
|
||||
}
|
||||
|
||||
fn codex_request_session_scope(
|
||||
request: &ClientSessionRequest<'_>,
|
||||
body_turn_metadata: &CodexTurnMetadataSignals,
|
||||
) -> Option<ClientSessionScope> {
|
||||
if let Some(scope) = explicit_aether_session_scope(request, CodexSessionScopeAdapter.family()) {
|
||||
return Some(scope);
|
||||
}
|
||||
|
||||
if let Some(root_session) = header_value_str(request.headers, "session-id")
|
||||
.or_else(|| header_value_str(request.headers, "thread-id"))
|
||||
.or_else(|| header_value_str(request.headers, "session_id"))
|
||||
.or_else(|| header_value_str(request.headers, "conversation_id"))
|
||||
.or_else(|| header_value_str(request.headers, "x-session-id"))
|
||||
{
|
||||
return Some(codex_session_scope(
|
||||
request,
|
||||
root_session,
|
||||
None,
|
||||
ClientSessionSignalSource::Header,
|
||||
));
|
||||
}
|
||||
|
||||
let body_client_metadata = request
|
||||
.body_json
|
||||
.and_then(|body| body.get("client_metadata"))
|
||||
.and_then(Value::as_object);
|
||||
if let Some(root_session) = body_client_metadata
|
||||
.and_then(|metadata| value_at_map_path(metadata, "session_id"))
|
||||
.or_else(|| {
|
||||
body_client_metadata.and_then(|metadata| value_at_map_path(metadata, "thread_id"))
|
||||
})
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| body_turn_metadata.session_id.clone())
|
||||
.or_else(|| body_turn_metadata.thread_id.clone())
|
||||
{
|
||||
return Some(codex_session_scope(
|
||||
request,
|
||||
root_session,
|
||||
None,
|
||||
ClientSessionSignalSource::Body,
|
||||
));
|
||||
}
|
||||
|
||||
let generic = GenericSessionScopeAdapter.extract_scope(request)?;
|
||||
Some(codex_session_scope(
|
||||
request,
|
||||
generic.session_id,
|
||||
generic.agent_id,
|
||||
generic.source,
|
||||
))
|
||||
}
|
||||
|
||||
fn codex_session_scope(
|
||||
request: &ClientSessionRequest<'_>,
|
||||
session_id: String,
|
||||
agent_id: Option<String>,
|
||||
source: ClientSessionSignalSource,
|
||||
) -> ClientSessionScope {
|
||||
ClientSessionScope::new(
|
||||
CodexSessionScopeAdapter.family(),
|
||||
session_id,
|
||||
agent_id,
|
||||
header_value_str(request.headers, "chatgpt-account-id"),
|
||||
source,
|
||||
)
|
||||
}
|
||||
|
||||
fn codex_turn_metadata_signals(value: Option<&Value>) -> CodexTurnMetadataSignals {
|
||||
match value {
|
||||
Some(Value::Object(metadata)) => codex_turn_metadata_signals_from_map(metadata),
|
||||
Some(Value::String(raw)) => parse_codex_turn_metadata(raw),
|
||||
_ => CodexTurnMetadataSignals::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_codex_turn_metadata(raw: &str) -> CodexTurnMetadataSignals {
|
||||
serde_json::from_str::<Map<String, Value>>(raw)
|
||||
.map(|metadata| codex_turn_metadata_signals_from_map(&metadata))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn codex_turn_metadata_signals_from_map(metadata: &Map<String, Value>) -> CodexTurnMetadataSignals {
|
||||
CodexTurnMetadataSignals {
|
||||
session_id: value_at_map_path(metadata, "session_id").map(ToOwned::to_owned),
|
||||
thread_id: value_at_map_path(metadata, "thread_id").map(ToOwned::to_owned),
|
||||
turn_id: value_at_map_path(metadata, "turn_id").map(ToOwned::to_owned),
|
||||
}
|
||||
}
|
||||
|
||||
fn value_at_map_path<'a>(object: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
|
||||
object
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn normalize_session_key(
|
||||
account_hint: Option<&str>,
|
||||
root_session: &str,
|
||||
@@ -831,12 +987,25 @@ mod tests {
|
||||
client_session_affinity_from_api_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,
|
||||
codex_request_signals_from_request, ClientSessionSignalSource, AETHER_AGENT_ID_HEADER,
|
||||
AETHER_SESSION_ID_HEADER,
|
||||
};
|
||||
use aether_scheduler_core::ClientSessionAffinity;
|
||||
use http::{HeaderMap, HeaderValue};
|
||||
use http::{HeaderMap, HeaderName, HeaderValue};
|
||||
use serde_json::json;
|
||||
|
||||
fn request_headers(values: &[(&str, &str)]) -> HeaderMap {
|
||||
values
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
(
|
||||
HeaderName::from_bytes(name.as_bytes()).expect("valid test header name"),
|
||||
HeaderValue::from_bytes(value.as_bytes()).expect("valid test header value"),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_adapter_extracts_body_session_and_agent() {
|
||||
let body = json!({
|
||||
@@ -933,6 +1102,276 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_request_signals_apply_session_precedence() {
|
||||
let cases = vec![
|
||||
(
|
||||
request_headers(&[
|
||||
(AETHER_SESSION_ID_HEADER, "aether-session"),
|
||||
("session-id", "header-session"),
|
||||
]),
|
||||
json!({"client_metadata": {"session_id": "body-session"}}),
|
||||
"aether-session",
|
||||
ClientSessionSignalSource::ExplicitAetherHeader,
|
||||
),
|
||||
(
|
||||
request_headers(&[
|
||||
("session-id", "header-session"),
|
||||
("thread-id", "header-thread"),
|
||||
("session_id", "header-session-underscore"),
|
||||
("conversation_id", "header-conversation"),
|
||||
]),
|
||||
json!({"client_metadata": {"session_id": "body-session"}}),
|
||||
"header-session",
|
||||
ClientSessionSignalSource::Header,
|
||||
),
|
||||
(
|
||||
request_headers(&[
|
||||
("thread-id", "header-thread"),
|
||||
("session_id", "header-session-underscore"),
|
||||
("conversation_id", "header-conversation"),
|
||||
]),
|
||||
json!({"client_metadata": {"session_id": "body-session"}}),
|
||||
"header-thread",
|
||||
ClientSessionSignalSource::Header,
|
||||
),
|
||||
(
|
||||
request_headers(&[
|
||||
("session_id", "header-session-underscore"),
|
||||
("conversation_id", "header-conversation"),
|
||||
]),
|
||||
json!({"client_metadata": {"session_id": "body-session"}}),
|
||||
"header-session-underscore",
|
||||
ClientSessionSignalSource::Header,
|
||||
),
|
||||
(
|
||||
request_headers(&[("conversation_id", "header-conversation")]),
|
||||
json!({"client_metadata": {"session_id": "body-session"}}),
|
||||
"header-conversation",
|
||||
ClientSessionSignalSource::Header,
|
||||
),
|
||||
(
|
||||
HeaderMap::new(),
|
||||
json!({
|
||||
"prompt_cache_key": "prompt-cache",
|
||||
"client_metadata": {
|
||||
"session_id": "body-session",
|
||||
"thread_id": "body-thread",
|
||||
"x-codex-turn-metadata": {
|
||||
"session_id": "nested-session",
|
||||
"thread_id": "nested-thread"
|
||||
}
|
||||
}
|
||||
}),
|
||||
"body-session",
|
||||
ClientSessionSignalSource::Body,
|
||||
),
|
||||
(
|
||||
HeaderMap::new(),
|
||||
json!({
|
||||
"prompt_cache_key": "prompt-cache",
|
||||
"client_metadata": {
|
||||
"thread_id": "body-thread",
|
||||
"x-codex-turn-metadata": {"session_id": "nested-session"}
|
||||
}
|
||||
}),
|
||||
"body-thread",
|
||||
ClientSessionSignalSource::Body,
|
||||
),
|
||||
(
|
||||
HeaderMap::new(),
|
||||
json!({
|
||||
"prompt_cache_key": "prompt-cache",
|
||||
"client_metadata": {
|
||||
"x-codex-turn-metadata": {
|
||||
"session_id": "nested-session",
|
||||
"thread_id": "nested-thread"
|
||||
}
|
||||
}
|
||||
}),
|
||||
"nested-session",
|
||||
ClientSessionSignalSource::Body,
|
||||
),
|
||||
(
|
||||
HeaderMap::new(),
|
||||
json!({
|
||||
"prompt_cache_key": "prompt-cache",
|
||||
"client_metadata": {
|
||||
"x-codex-turn-metadata": json!({
|
||||
"thread_id": "nested-thread"
|
||||
}).to_string()
|
||||
}
|
||||
}),
|
||||
"nested-thread",
|
||||
ClientSessionSignalSource::Body,
|
||||
),
|
||||
(
|
||||
HeaderMap::new(),
|
||||
json!({
|
||||
"prompt_cache_key": "prompt-cache",
|
||||
"conversation_id": "generic-conversation"
|
||||
}),
|
||||
"prompt-cache",
|
||||
ClientSessionSignalSource::Body,
|
||||
),
|
||||
(
|
||||
HeaderMap::new(),
|
||||
json!({"metadata": {"session_id": "generic-session"}}),
|
||||
"generic-session",
|
||||
ClientSessionSignalSource::Body,
|
||||
),
|
||||
];
|
||||
|
||||
for (headers, body, expected_session_id, expected_source) in cases {
|
||||
let signals = codex_request_signals_from_request(&headers, Some(&body));
|
||||
assert_eq!(signals.session_id.as_deref(), Some(expected_session_id));
|
||||
|
||||
let mut codex_headers = headers;
|
||||
codex_headers.insert(
|
||||
http::header::USER_AGENT,
|
||||
HeaderValue::from_static("codex_cli_rs/0.144.1"),
|
||||
);
|
||||
let scope = client_session_scope_from_request(&codex_headers, Some(&body))
|
||||
.expect("Codex scope should reuse the native signal precedence");
|
||||
assert_eq!(scope.client_family, "codex");
|
||||
assert_eq!(scope.session_id, expected_session_id);
|
||||
assert_eq!(scope.source, expected_source);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_request_signals_extract_thread_and_prompt_cache_independently() {
|
||||
let body = json!({
|
||||
"prompt_cache_key": "prompt-cache",
|
||||
"client_metadata": {
|
||||
"thread_id": "body-thread",
|
||||
"x-codex-turn-metadata": {"thread_id": "nested-thread"}
|
||||
}
|
||||
});
|
||||
let headers = request_headers(&[("thread-id", "header-thread")]);
|
||||
let header_signals = codex_request_signals_from_request(&headers, Some(&body));
|
||||
assert_eq!(header_signals.thread_id.as_deref(), Some("header-thread"));
|
||||
assert_eq!(
|
||||
header_signals.prompt_cache_key.as_deref(),
|
||||
Some("prompt-cache")
|
||||
);
|
||||
|
||||
let body_signals = codex_request_signals_from_request(&HeaderMap::new(), Some(&body));
|
||||
assert_eq!(body_signals.thread_id.as_deref(), Some("body-thread"));
|
||||
|
||||
let nested_body = json!({
|
||||
"client_metadata": {
|
||||
"x-codex-turn-metadata": json!({
|
||||
"thread_id": "nested-thread"
|
||||
}).to_string()
|
||||
}
|
||||
});
|
||||
let nested_signals =
|
||||
codex_request_signals_from_request(&HeaderMap::new(), Some(&nested_body));
|
||||
assert_eq!(nested_signals.thread_id.as_deref(), Some("nested-thread"));
|
||||
|
||||
let session_only_body = json!({"client_metadata": {"session_id": "body-session"}});
|
||||
let session_only_signals =
|
||||
codex_request_signals_from_request(&HeaderMap::new(), Some(&session_only_body));
|
||||
assert_eq!(
|
||||
session_only_signals.thread_id.as_deref(),
|
||||
Some("body-session")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_request_signals_use_live_session_header() {
|
||||
let headers = request_headers(&[("x-session-id", "live-session")]);
|
||||
let signals = codex_request_signals_from_request(&headers, None);
|
||||
|
||||
assert_eq!(signals.session_id.as_deref(), Some("live-session"));
|
||||
assert_eq!(signals.thread_id.as_deref(), Some("live-session"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_request_signals_prefer_responses_session_header_over_live_session_header() {
|
||||
let headers = request_headers(&[
|
||||
("session-id", "responses-session"),
|
||||
("x-session-id", "live-session"),
|
||||
]);
|
||||
let signals = codex_request_signals_from_request(&headers, None);
|
||||
|
||||
assert_eq!(signals.session_id.as_deref(), Some("responses-session"));
|
||||
assert_eq!(signals.thread_id.as_deref(), Some("responses-session"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_request_signals_apply_turn_precedence() {
|
||||
let headers = request_headers(&[("x-codex-turn-metadata", r#"{"turn_id":"header-turn"}"#)]);
|
||||
let direct_body = json!({
|
||||
"turn_id": "top-level-turn",
|
||||
"client_metadata": {
|
||||
"turn_id": "body-turn",
|
||||
"x-codex-turn-metadata": {"turn_id": "nested-turn"}
|
||||
}
|
||||
});
|
||||
assert_eq!(
|
||||
codex_request_signals_from_request(&headers, Some(&direct_body))
|
||||
.turn_id
|
||||
.as_deref(),
|
||||
Some("body-turn")
|
||||
);
|
||||
|
||||
let nested_object_body = json!({
|
||||
"turn_id": "top-level-turn",
|
||||
"client_metadata": {
|
||||
"x-codex-turn-metadata": {"turn_id": "nested-object-turn"}
|
||||
}
|
||||
});
|
||||
assert_eq!(
|
||||
codex_request_signals_from_request(&headers, Some(&nested_object_body))
|
||||
.turn_id
|
||||
.as_deref(),
|
||||
Some("nested-object-turn")
|
||||
);
|
||||
|
||||
let nested_string_body = json!({
|
||||
"turn_id": "top-level-turn",
|
||||
"client_metadata": {
|
||||
"x-codex-turn-metadata": json!({
|
||||
"turn_id": "nested-string-turn"
|
||||
}).to_string()
|
||||
}
|
||||
});
|
||||
assert_eq!(
|
||||
codex_request_signals_from_request(&headers, Some(&nested_string_body))
|
||||
.turn_id
|
||||
.as_deref(),
|
||||
Some("nested-string-turn")
|
||||
);
|
||||
|
||||
let top_level_body = json!({
|
||||
"turn_id": "top-level-turn",
|
||||
"client_metadata": {"x-codex-turn-metadata": "not-json"}
|
||||
});
|
||||
assert_eq!(
|
||||
codex_request_signals_from_request(&headers, Some(&top_level_body))
|
||||
.turn_id
|
||||
.as_deref(),
|
||||
Some("top-level-turn")
|
||||
);
|
||||
assert_eq!(
|
||||
codex_request_signals_from_request(&headers, None)
|
||||
.turn_id
|
||||
.as_deref(),
|
||||
Some("header-turn")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_request_signals_ignore_client_request_id() {
|
||||
let headers = request_headers(&[("x-client-request-id", "request-only-id")]);
|
||||
let signals =
|
||||
codex_request_signals_from_request(&headers, Some(&json!({"model": "gpt-5"})));
|
||||
|
||||
assert_eq!(signals, super::CodexRequestSignals::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn report_context_round_trips_normalized_session_affinity() {
|
||||
let affinity = ClientSessionAffinity::new(
|
||||
|
||||
Reference in New Issue
Block a user