i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-05 00:47:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(codex): add provider outbound policy boundary

Browse Source
This commit is contained in:
elky
2026-09-01 21:21:42 +08:00
parent 715f2773c3
commit d5f34b2ee2
12 changed files with 1313 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/codex_context.rs
+14 -19
View File
@@ -1,15 +1,15 @@
use std::sync::{Arc, OnceLock};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_provider_transport::CodexFingerprintConvergenceContext;
use http::{request::Parts, HeaderMap};
use serde_json::Value;
use uuid::Uuid;
use crate::ai_serving::transport::ProviderOutboundRequestContext;
use crate::client_session_affinity::codex_request_signals_from_request;
#[derive(Debug, Clone)]
pub(crate) struct CodexFingerprintContextSlot(Arc<OnceLock<CodexFingerprintConvergenceContext>>);
pub(crate) struct CodexFingerprintContextSlot(Arc<OnceLock<ProviderOutboundRequestContext>>);
impl Default for CodexFingerprintContextSlot {
fn default() -> Self {
@@ -18,11 +18,7 @@ impl Default for CodexFingerprintContextSlot {
}
impl CodexFingerprintContextSlot {
fn resolve(
&self,
headers: &HeaderMap,
body_json: &Value,
) -> CodexFingerprintConvergenceContext {
fn resolve(&self, headers: &HeaderMap, body_json: &Value) -> ProviderOutboundRequestContext {
self.0
.get_or_init(|| {
build_codex_fingerprint_context(headers, body_json, Uuid::now_v7().to_string())
@@ -34,10 +30,10 @@ impl CodexFingerprintContextSlot {
pub(crate) fn resolve_codex_fingerprint_context(
parts: &Parts,
body_json: &Value,
) -> CodexFingerprintConvergenceContext {
) -> ProviderOutboundRequestContext {
if let Some(context) = parts
.extensions
.get::<CodexFingerprintConvergenceContext>()
.get::<ProviderOutboundRequestContext>()
.cloned()
{
return context;
@@ -51,7 +47,7 @@ pub(crate) fn resolve_codex_fingerprint_context(
pub(crate) fn install_codex_fingerprint_context_slot(parts: &mut Parts) {
if parts
.extensions
.get::<CodexFingerprintConvergenceContext>()
.get::<ProviderOutboundRequestContext>()
.is_none()
&& parts
.extensions
@@ -67,11 +63,11 @@ pub(crate) fn install_codex_fingerprint_context_slot(parts: &mut Parts) {
pub(crate) fn ensure_codex_fingerprint_context(
parts: &mut Parts,
body_json: &Value,
) -> CodexFingerprintConvergenceContext {
) -> ProviderOutboundRequestContext {
let context = resolve_codex_fingerprint_context(parts, body_json);
if parts
.extensions
.get::<CodexFingerprintConvergenceContext>()
.get::<ProviderOutboundRequestContext>()
.is_none()
{
parts.extensions.remove::<CodexFingerprintContextSlot>();
@@ -84,7 +80,7 @@ pub(crate) fn attach_codex_logical_turn_context(
parts: &mut Parts,
body_json: &Value,
logical_turn_id: &str,
) -> CodexFingerprintConvergenceContext {
) -> ProviderOutboundRequestContext {
let context =
build_codex_fingerprint_context(&parts.headers, body_json, logical_turn_id.to_string());
parts.extensions.remove::<CodexFingerprintContextSlot>();
@@ -94,7 +90,7 @@ pub(crate) fn attach_codex_logical_turn_context(
pub(crate) fn restore_codex_logical_turn_context(
parts: &mut Parts,
context: &CodexFingerprintConvergenceContext,
context: &ProviderOutboundRequestContext,
) {
parts.extensions.remove::<CodexFingerprintContextSlot>();
parts.extensions.insert(context.clone());
@@ -104,10 +100,9 @@ fn build_codex_fingerprint_context(
headers: &HeaderMap,
body_json: &Value,
logical_turn_id: String,
) -> CodexFingerprintConvergenceContext {
) -> ProviderOutboundRequestContext {
let signals = codex_request_signals_from_request(headers, Some(body_json));
let mut context =
CodexFingerprintConvergenceContext::new(logical_turn_id, current_unix_millis());
let mut context = ProviderOutboundRequestContext::new(logical_turn_id, current_unix_millis());
if let Some(turn_id) = signals.turn_id {
context = context.with_original_turn_id(turn_id);
@@ -160,14 +155,14 @@ mod tests {
assert_eq!(context.original_client_session_id(), Some("header-thread"));
assert_eq!(context.original_prompt_cache_key(), Some("client-cache"));
assert_eq!(
parts.extensions.get::<CodexFingerprintConvergenceContext>(),
parts.extensions.get::<ProviderOutboundRequestContext>(),
Some(&context)
);
}
#[test]
fn restored_context_wins_over_retry_request_signals() {
let original = CodexFingerprintConvergenceContext::new("logical-turn", 1234)
let original = ProviderOutboundRequestContext::new("logical-turn", 1234)
.with_original_turn_id("original-turn")
.with_original_client_session_id("original-thread")
.with_original_prompt_cache_key("original-cache");
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
+89 -16
View File
@@ -13,7 +13,7 @@ use http::{HeaderMap, HeaderName, HeaderValue};
use serde_json::{json, Value};
use crate::ai_serving::planner::common::extract_standard_requested_model;
use crate::ai_serving::transport::CodexFingerprintConvergenceContext;
use crate::ai_serving::transport::ProviderOutboundRequestContext;
use crate::ai_serving::{
ClientSurface, ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot,
GatewayCredentialCarrier, GatewayProviderTransportSnapshot, PlannerAppState,
@@ -56,7 +56,7 @@ pub(crate) struct LocalRequestedModelDecisionInput {
pub(crate) client_surface: Option<ClientSurface>,
pub(crate) gateway_credential_carrier: Option<GatewayCredentialCarrier>,
pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
pub(crate) codex_fingerprint_context: Option<CodexFingerprintConvergenceContext>,
pub(crate) provider_outbound_context: Option<ProviderOutboundRequestContext>,
pub(crate) routing_policy: Option<ResolvedRoutingPolicy>,
pub(crate) routing_trace_seed: Option<RoutingDecisionTrace>,
pub(crate) routing_context: Option<LocalRoutingRequestContext>,
@@ -168,7 +168,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
provider_api_format.as_str(),
);
}
apply_codex_fingerprint_convergence_to_decision(
apply_provider_outbound_request_policies_to_decision(
input,
decision,
transport,
@@ -231,7 +231,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
provider_api_format.as_str(),
);
}
apply_codex_fingerprint_convergence_to_decision(
apply_provider_outbound_request_policies_to_decision(
input,
decision,
transport,
@@ -356,7 +356,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
if original_provider_request_body.is_some() {
decision.provider_request_body = Some(provider_request_body);
}
apply_codex_fingerprint_convergence_to_decision(
apply_provider_outbound_request_policies_to_decision(
input,
decision,
transport,
@@ -366,7 +366,7 @@ pub(crate) fn apply_provider_request_routing_policy_to_decision_with_websocket_m
Ok(())
}
fn apply_codex_fingerprint_convergence_to_decision(
fn apply_provider_outbound_request_policies_to_decision(
input: &LocalRequestedModelDecisionInput,
decision: &mut AiExecutionDecision,
transport: Option<&GatewayProviderTransportSnapshot>,
@@ -377,17 +377,17 @@ fn apply_codex_fingerprint_convergence_to_decision(
else {
return;
};
let Some(context) = input.codex_fingerprint_context.as_ref() else {
let Some(context) = input.provider_outbound_context.as_ref() else {
return;
};
let applied = crate::ai_serving::transport::apply_codex_fingerprint_convergence_with_context(
let results = crate::ai_serving::transport::apply_provider_outbound_request_policies(
transport,
provider_api_format,
context,
&mut decision.provider_request_headers,
provider_request_body,
);
if applied {
if results.iter().any(|result| result.was_applied()) {
decision.prompt_cache_key = provider_request_body
.get("prompt_cache_key")
.and_then(Value::as_str)
@@ -395,6 +395,31 @@ fn apply_codex_fingerprint_convergence_to_decision(
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
}
if results.is_empty() {
return;
}
for result in &results {
tracing::debug!(
event_name = "provider_outbound_policy_evaluated",
log_type = "event",
policy = ?result.policy,
outcome = ?result.outcome,
reason = ?result.reason,
mutation_scope = ?result.mutation_scope,
identity_scope = ?result.identity_scope,
"provider outbound request policy evaluated"
);
}
let Some(serde_json::Value::Object(report_context)) = decision.report_context.as_mut() else {
return;
};
report_context.insert(
"provider_outbound_policies".to_string(),
serde_json::json!({
"schema_version": 1,
"results": results,
}),
);
}
struct GatewayAuthenticatedDecisionInputPort<'a> {
@@ -483,7 +508,7 @@ pub(crate) fn build_local_requested_model_decision_input(
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
provider_outbound_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
@@ -498,7 +523,7 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
body_json: &Value,
client_api_format: &str,
) -> Result<(), GatewayError> {
input.codex_fingerprint_context =
input.provider_outbound_context =
Some(crate::ai_serving::codex_context::resolve_codex_fingerprint_context(parts, body_json));
let explicit_group = routing_header_value_str(&parts.headers, ROUTING_GROUP_HEADER);
let selected_group = match state.routing_group_read_repository() {
@@ -1325,7 +1350,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
provider_outbound_context: None,
routing_policy: None,
routing_trace_seed: None,
model_directive_policy: Default::default(),
@@ -1574,7 +1599,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
provider_outbound_context: None,
routing_policy: None,
routing_trace_seed: None,
model_directive_policy: Default::default(),
@@ -1644,7 +1669,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
provider_outbound_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
@@ -1713,6 +1738,35 @@ mod tests {
);
}
#[test]
fn non_codex_provider_outbound_policies_are_terminal_noop() {
let mut input = sample_decision_input();
input.routing_context = None;
input.provider_outbound_context = Some(ProviderOutboundRequestContext::new(
"logical-turn",
1_700_000_000_123,
));
let mut decision = sample_codex_fingerprint_decision();
decision.provider_type = Some("openai".to_string());
decision.provider_api_format = Some("openai:responses".to_string());
decision.client_api_format = Some("openai:responses".to_string());
let mut transport = sample_codex_fingerprint_transport();
transport.provider.provider_type = "openai".to_string();
let original_headers = decision.provider_request_headers.clone();
let original_body = decision.provider_request_body.clone();
apply_provider_request_routing_policy_to_decision(&input, &mut decision, Some(&transport))
.expect("non-Codex terminal finalization should succeed");
assert_eq!(decision.provider_request_headers, original_headers);
assert_eq!(decision.provider_request_body, original_body);
assert!(decision
.report_context
.as_ref()
.and_then(|context| context.get("provider_outbound_policies"))
.is_none());
}
#[test]
fn codex_fingerprint_convergence_runs_at_every_provider_routing_success_exit() {
let transport = sample_codex_fingerprint_transport();
@@ -1729,8 +1783,8 @@ mod tests {
});
let mut with_mutation = sample_decision_input();
for input in [&mut no_context, &mut empty_mutation, &mut with_mutation] {
input.codex_fingerprint_context = Some(
CodexFingerprintConvergenceContext::new(
input.provider_outbound_context = Some(
ProviderOutboundRequestContext::new(
uuid::Uuid::new_v4().to_string(),
1_756_668_000_000,
)
@@ -1799,6 +1853,25 @@ mod tests {
);
assert_eq!(body["client_metadata"]["x-codex-window-id"], window_id);
let policy_results = decision
.report_context
.as_ref()
.and_then(|context| context.get("provider_outbound_policies"))
.and_then(|policies| policies.get("results"))
.and_then(Value::as_array)
.expect("provider policy results");
assert_eq!(
policy_results.len(),
1,
"policy result count at {exit_name}"
);
assert_eq!(
policy_results[0]["policy"],
json!("codex_fingerprint_convergence")
);
assert_eq!(policy_results[0]["outcome"], json!("applied"));
assert_eq!(policy_results[0]["reason"], json!("applied"));
let header_metadata: Value =
serde_json::from_str(&decision.provider_request_headers["x-codex-turn-metadata"])
.expect("header turn metadata");
apps/aether-gateway/src/ai_serving/planner/standard/family/payload.rs
+1 -1
View File
@@ -377,7 +377,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
provider_outbound_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
+1 -1
View File
@@ -2182,7 +2182,7 @@ mod tests {
client_surface: None,
gateway_credential_carrier: None,
client_session_affinity: None,
codex_fingerprint_context: None,
provider_outbound_context: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
apps/aether-gateway/src/ai_serving/transport.rs
+8 -3
View File
@@ -62,7 +62,8 @@ pub(crate) use aether_provider_transport::{
append_transport_diagnostics_to_value, apply_codex_fingerprint_convergence,
apply_codex_fingerprint_convergence_with_context, apply_local_auth_config_header_overrides,
apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules,
apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules,
apply_local_header_rules_with_request_headers, apply_provider_outbound_request_policies,
apply_standard_provider_request_body_rules,
apply_standard_provider_request_body_rules_with_request_headers,
apply_transport_request_body_semantics, body_rules_are_locally_supported,
body_rules_handle_path, body_rules_have_enabled_rules,
@@ -112,7 +113,11 @@ pub(crate) use aether_provider_transport::{
GeminiCliRequestAuthSupport, GeminiCliRequestAuthUnsupportedReason,
GeminiCliRequestEnvelopeSupport, GeminiFilesHeadersInput, GeminiFilesRequestBodyError,
GeminiFilesRequestBodyParts, GrokHeaderInput, LocalResolvedOAuthRequestAuth,
ProviderOpenAiImageHeadersInput, ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
ProviderOpenAiImageHeadersInput, ProviderOutboundRequestContext,
ProviderOutboundRequestIdentityScope, ProviderOutboundRequestMutationScope,
ProviderOutboundRequestPolicy, ProviderOutboundRequestPolicyOutcome,
ProviderOutboundRequestPolicyReason, ProviderOutboundRequestPolicyResult,
ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
SameFormatProviderCompatibilityEdit, SameFormatProviderCompatibilityEditAction,
SameFormatProviderFamily, SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior,
SameFormatProviderRequestBehaviorParams, SameFormatProviderRequestBodyInput,
@@ -121,5 +126,5 @@ pub(crate) use aether_provider_transport::{
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
TransportRequestBodySemanticsError, TransportRequestUrlParams, GEMINI_CLI_USER_AGENT,
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, GROK_CHAT_PATH, GROK_INTERNAL_HEADER,
GROK_RATE_LIMITS_PATH, WINDSURF_ENVELOPE_NAME,
GROK_RATE_LIMITS_PATH, PROVIDER_OUTBOUND_CONTEXT_MAX_VALUE_BYTES, WINDSURF_ENVELOPE_NAME,
};
apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs
+8
View File
@@ -186,6 +186,7 @@ async fn handle_live_http(
client_model,
dialect,
None,
None,
)
.await
{
@@ -240,6 +241,12 @@ async fn handle_live_http(
);
};
let lease = LivePoolLeaseGuard::new(state, &candidate);
let provider_outbound_context = candidate
.execution
.provider_type
.as_deref()
.is_some_and(|provider_type| provider_type.trim().eq_ignore_ascii_case("codex"))
.then(|| candidate.provider_outbound_context.clone());
let binding = LiveCallBinding::from_candidate(&candidate);
let mut provider_session = offer.session.clone();
provider_session
@@ -525,6 +532,7 @@ async fn handle_live_http(
auth_context.api_key_id.as_str(),
call_id.as_str(),
&binding,
provider_outbound_context.as_ref(),
)
.await
{
apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs
+14 -10
View File
@@ -16,6 +16,7 @@ use serde_json::json;
use sha2::{Digest, Sha256};
use url::{form_urlencoded, Url};
use crate::ai_serving::transport::ProviderOutboundRequestContext;
use crate::ai_serving::{
build_standard_stream_plan_from_decision,
maybe_build_pinned_stream_local_same_format_provider_decision_payload, AiExecutionDecision,
@@ -49,8 +50,7 @@ pub(super) enum LiveAuthMode {
pub(super) struct PlannedLiveCandidate {
pub(super) execution: AiExecutionDecision,
pub(super) pinned_candidate: ResponsesWebSocketPinnedCandidate,
pub(super) codex_fingerprint_context:
aether_provider_transport::CodexFingerprintConvergenceContext,
pub(super) provider_outbound_context: ProviderOutboundRequestContext,
pub(super) client_model: String,
pub(super) provider_model: String,
pub(super) auth_mode: LiveAuthMode,
@@ -198,6 +198,7 @@ pub(super) async fn plan_live_candidate(
client_model: &str,
dialect: LiveRouteDialect,
pinned_candidate: Option<&ResponsesWebSocketPinnedCandidate>,
provider_outbound_context: Option<&ProviderOutboundRequestContext>,
) -> Result<LiveCandidatePlanningOutcome, GatewayError> {
let diagnostic_guard = LiveRuntimeMissDiagnosticGuard::new(state, trace_id);
let candidate = plan_live_candidate_inner(
@@ -209,6 +210,7 @@ pub(super) async fn plan_live_candidate(
client_model,
dialect,
pinned_candidate,
provider_outbound_context,
)
.await?;
Ok(LiveCandidatePlanningOutcome {
@@ -226,13 +228,18 @@ async fn plan_live_candidate_inner(
client_model: &str,
dialect: LiveRouteDialect,
pinned_candidate: Option<&ResponsesWebSocketPinnedCandidate>,
provider_outbound_context: Option<&ProviderOutboundRequestContext>,
) -> Result<Option<PlannedLiveCandidate>, GatewayError> {
if validate_model(client_model).is_err() || client_model.len() > MAX_LIVE_MODEL_BYTES {
return Ok(None);
}
let mut parts = build_live_planning_parts(headers, remote_addr);
let body = json!({"model": client_model, "input": []});
crate::ai_serving::codex_context::install_codex_fingerprint_context_slot(&mut parts);
if let Some(context) = provider_outbound_context {
crate::ai_serving::codex_context::restore_codex_logical_turn_context(&mut parts, context);
} else {
crate::ai_serving::codex_context::install_codex_fingerprint_context_slot(&mut parts);
}
let execution = maybe_build_pinned_stream_local_same_format_provider_decision_payload(
state,
&parts,
@@ -341,7 +348,7 @@ async fn plan_live_candidate_inner(
Ok(Some(PlannedLiveCandidate {
execution,
pinned_candidate,
codex_fingerprint_context:
provider_outbound_context:
crate::ai_serving::codex_context::resolve_codex_fingerprint_context(&parts, &body),
client_model: client_model.to_string(),
provider_model,
@@ -568,7 +575,7 @@ pub(super) fn build_live_stream_admission_attempt(
let mut parts = build_live_planning_parts(headers, remote_addr);
crate::ai_serving::codex_context::restore_codex_logical_turn_context(
&mut parts,
&candidate.codex_fingerprint_context,
&candidate.provider_outbound_context,
);
let body = json!({"model": candidate.client_model.as_str(), "input": []});
let mut execution = candidate.execution.clone();
@@ -578,6 +585,7 @@ pub(super) fn build_live_stream_admission_attempt(
}
fn live_auth_mode(provider_type: &str, effective_auth_type: &str) -> Option<LiveAuthMode> {
let provider_type = provider_type.trim();
match effective_auth_type.trim().to_ascii_lowercase().as_str() {
"api_key" | "bearer" => Some(LiveAuthMode::ApiKey),
"oauth" if provider_type.eq_ignore_ascii_case("codex") => Some(LiveAuthMode::ChatGptOauth),
@@ -931,11 +939,7 @@ mod tests {
"key-1",
)
.unwrap(),
codex_fingerprint_context:
aether_provider_transport::CodexFingerprintConvergenceContext::new(
"test-live-turn",
1,
),
provider_outbound_context: ProviderOutboundRequestContext::new("test-live-turn", 1),
client_model: "global-model".to_string(),
provider_model: "provider-model".to_string(),
auth_mode,
apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs
+514 -41
View File
@@ -13,6 +13,9 @@ use sha2::{Digest, Sha256};
use tokio::sync::{oneshot, watch};
use tokio::task::JoinHandle;
use crate::ai_serving::transport::{
ProviderOutboundRequestContext, PROVIDER_OUTBOUND_CONTEXT_MAX_VALUE_BYTES,
};
use crate::ai_serving::ResponsesWebSocketPinnedCandidate;
use super::planner::{LiveAuthMode, PlannedLiveCandidate};
@@ -21,6 +24,9 @@ use super::protocol::validate_call_id;
const SCHEMA_VERSION: u16 = 2;
const RECORD_PREFIX: &str = "codex_live:call:v2:";
const RECORD_DOMAIN: &[u8] = b"aether-codex-live-call-v2";
const CONTEXT_SCHEMA_VERSION: u16 = 1;
const CONTEXT_PREFIX: &str = "codex_live:call_context:v1:";
const CONTEXT_DOMAIN: &[u8] = b"aether-codex-live-call-context-v1";
const INDEX_PREFIX: &str = "codex_live:call_index:v2:";
const INDEX_DOMAIN: &[u8] = b"aether-codex-live-call-index-v2";
const LOCK_PREFIX: &str = "codex_live:call_lock:v2:";
@@ -38,6 +44,8 @@ const LOCK_OWNER: &str = "codex_live_call_registry";
const SIDEBAND_LOCK_OWNER: &str = "codex_live_sideband_attachment";
const MAX_RECORDS_PER_PRINCIPAL: usize = 64;
const MAX_SERIALIZED_RECORD_BYTES: usize = 4 * 1024;
const MAX_SERIALIZED_CONTEXT_BYTES: usize = 2 * 1024;
const MAX_CONTEXT_FIELD_BYTES: usize = PROVIDER_OUTBOUND_CONTEXT_MAX_VALUE_BYTES;
const MAX_PRINCIPAL_BYTES: usize = 256;
const MAX_RECORD_ID_BYTES: usize = 256;
const LIVE_LOG_TARGET: &str = "aether_gateway::handlers::proxy::codex_live";
@@ -51,7 +59,7 @@ enum RegisterCommitState {
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum LiveCallLookup {
Found(LiveCallBinding),
Found(LiveCallRecord),
Missing,
Expired,
}
@@ -277,6 +285,118 @@ pub(super) struct LiveCallBinding {
created_at_unix_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct LiveCallRecord {
binding: LiveCallBinding,
provider_outbound_context: Option<ProviderOutboundRequestContext>,
}
impl LiveCallRecord {
pub(super) fn binding(&self) -> &LiveCallBinding {
&self.binding
}
pub(super) fn provider_outbound_context(&self) -> Option<&ProviderOutboundRequestContext> {
self.provider_outbound_context.as_ref()
}
pub(super) fn matches_candidate(&self, candidate: &PlannedLiveCandidate) -> bool {
self.binding.matches_candidate(candidate)
&& self
.provider_outbound_context
.as_ref()
.is_none_or(|context| context == &candidate.provider_outbound_context)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct LiveProviderOutboundContextRecord {
schema_version: u16,
binding_created_at_unix_ms: u64,
logical_turn_id: String,
original_turn_id: Option<String>,
original_client_session_id: Option<String>,
original_prompt_cache_key: Option<String>,
turn_started_at_unix_ms: u64,
}
impl LiveProviderOutboundContextRecord {
fn from_context(
binding: &LiveCallBinding,
context: &ProviderOutboundRequestContext,
) -> Result<Self, LiveCallRegistryError> {
let record = Self {
schema_version: CONTEXT_SCHEMA_VERSION,
binding_created_at_unix_ms: binding.created_at_unix_ms,
logical_turn_id: context.logical_turn_id().to_string(),
original_turn_id: context.original_turn_id().map(ToOwned::to_owned),
original_client_session_id: context.original_client_session_id().map(ToOwned::to_owned),
original_prompt_cache_key: context.original_prompt_cache_key().map(ToOwned::to_owned),
turn_started_at_unix_ms: context.turn_started_at_unix_ms(),
};
record.validate(binding)?;
Ok(record)
}
fn into_context(
self,
binding: &LiveCallBinding,
) -> Result<ProviderOutboundRequestContext, LiveCallRegistryError> {
self.validate(binding)?;
let mut context =
ProviderOutboundRequestContext::new(self.logical_turn_id, self.turn_started_at_unix_ms);
if let Some(value) = self.original_turn_id {
context = context.with_original_turn_id(value);
}
if let Some(value) = self.original_client_session_id {
context = context.with_original_client_session_id(value);
}
if let Some(value) = self.original_prompt_cache_key {
context = context.with_original_prompt_cache_key(value);
}
Ok(context)
}
fn validate(&self, binding: &LiveCallBinding) -> Result<(), LiveCallRegistryError> {
if self.schema_version != CONTEXT_SCHEMA_VERSION {
return Err(LiveCallRegistryError::InvalidRecord(
"unsupported_context_schema_version",
));
}
if self.binding_created_at_unix_ms != binding.created_at_unix_ms {
return Err(LiveCallRegistryError::InvalidRecord(
"context_binding_mismatch",
));
}
if self.turn_started_at_unix_ms == 0 {
return Err(LiveCallRegistryError::InvalidRecord(
"invalid_context_started_at",
));
}
validate_context_field(
self.logical_turn_id.as_str(),
"invalid_context_logical_turn",
)?;
for (value, error) in [
(self.original_turn_id.as_deref(), "invalid_context_turn"),
(
self.original_client_session_id.as_deref(),
"invalid_context_session",
),
(
self.original_prompt_cache_key.as_deref(),
"invalid_context_prompt_cache",
),
] {
if let Some(value) = value {
validate_context_field(value, error)?;
}
}
Ok(())
}
}
impl LiveCallBinding {
pub(super) fn from_candidate(candidate: &PlannedLiveCandidate) -> Self {
Self {
@@ -370,9 +490,11 @@ impl LiveCallRegistry {
api_key_id: &str,
call_id: &str,
binding: &LiveCallBinding,
provider_outbound_context: Option<&ProviderOutboundRequestContext>,
) -> Result<(), LiveCallRegistryError> {
binding.validate()?;
let key = record_key(user_id, api_key_id, call_id)?;
let context_key = context_key(key.as_str())?;
let index = index_key(user_id, api_key_id)?;
let lock = lock_key(user_id, api_key_id)?;
let serialized =
@@ -380,12 +502,40 @@ impl LiveCallRegistry {
if serialized.len() > MAX_SERIALIZED_RECORD_BYTES {
return Err(LiveCallRegistryError::RecordTooLarge);
}
let context_record = provider_outbound_context
.map(|context| LiveProviderOutboundContextRecord::from_context(binding, context))
.transpose()?;
let serialized_context = context_record
.as_ref()
.map(serde_json::to_string)
.transpose()
.map_err(LiveCallRegistryError::Serialization)?;
if serialized_context
.as_ref()
.is_some_and(|serialized| serialized.len() > MAX_SERIALIZED_CONTEXT_BYTES)
{
return Err(LiveCallRegistryError::RecordTooLarge);
}
let lease = self.acquire_lock(lock.as_str()).await?;
let result = self
.register_locked(key.as_str(), index.as_str(), serialized, binding)
.register_locked(
key.as_str(),
context_key.as_str(),
index.as_str(),
serialized,
serialized_context,
binding,
context_record.as_ref(),
)
.await;
let exact_binding_committed = if result.is_err() {
self.exact_binding_is_committed(key.as_str(), binding).await
self.exact_binding_is_committed(
key.as_str(),
context_key.as_str(),
binding,
context_record.as_ref(),
)
.await
} else {
false
};
@@ -425,13 +575,13 @@ impl LiveCallRegistry {
user_id: &str,
api_key_id: &str,
call_id: &str,
) -> Result<Option<LiveCallBinding>, LiveCallRegistryError> {
) -> Result<Option<LiveCallRecord>, LiveCallRegistryError> {
Ok(
match self
.lookup_with_status(user_id, api_key_id, call_id)
.await?
{
LiveCallLookup::Found(binding) => Some(binding),
LiveCallLookup::Found(record) => Some(record),
LiveCallLookup::Missing | LiveCallLookup::Expired => None,
},
)
@@ -444,18 +594,38 @@ impl LiveCallRegistry {
call_id: &str,
) -> Result<LiveCallLookup, LiveCallRegistryError> {
let key = record_key(user_id, api_key_id, call_id)?;
if let Some(serialized) = self
let context_key = context_key(key.as_str())?;
let values = self
.runtime_state
.kv_get(key.as_str())
.kv_get_many(&[key.clone(), context_key])
.await
.map_err(LiveCallRegistryError::Storage)?
{
.map_err(LiveCallRegistryError::Storage)?;
if let Some(serialized) = values.first().and_then(Option::as_ref) {
let binding = serde_json::from_str::<LiveCallBinding>(serialized.as_str())
.map_err(LiveCallRegistryError::CorruptRecord)?;
binding.validate()?;
return Ok(LiveCallLookup::Found(binding));
// Missing companion data is a rolling-upgrade compatible legacy
// v2 binding. Sideband planning will use the existing fallback
// behavior until that short-lived binding expires.
let provider_outbound_context = values
.get(1)
.and_then(Option::as_ref)
.map(|serialized| {
serde_json::from_str::<LiveProviderOutboundContextRecord>(serialized.as_str())
.map_err(LiveCallRegistryError::CorruptRecord)?
.into_context(&binding)
})
.transpose()?;
return Ok(LiveCallLookup::Found(LiveCallRecord {
binding,
provider_outbound_context,
}));
}
// A companion may outlive the base record by the expiry grace period,
// or be visible briefly before registration commits the base. Do not
// delete it from a lock-free lookup; the bounded TTL and the next
// locked registration handle cleanup without racing the writer.
let index = index_key(user_id, api_key_id)?;
let indexed = self
.runtime_state
@@ -521,48 +691,121 @@ impl LiveCallRegistry {
}
}
async fn exact_binding_is_committed(&self, key: &str, expected: &LiveCallBinding) -> bool {
let stored =
match tokio::time::timeout(COMMIT_VERIFY_TIMEOUT, self.runtime_state.kv_get(key)).await
{
Ok(Ok(Some(stored))) => stored,
Ok(Ok(None) | Err(_)) | Err(_) => return false,
};
let Ok(actual) = serde_json::from_str::<LiveCallBinding>(stored.as_str()) else {
async fn exact_binding_is_committed(
&self,
key: &str,
context_storage_key: &str,
expected: &LiveCallBinding,
expected_context: Option<&LiveProviderOutboundContextRecord>,
) -> bool {
let keys = [key.to_string(), context_storage_key.to_string()];
let stored = match tokio::time::timeout(
COMMIT_VERIFY_TIMEOUT,
self.runtime_state.kv_get_many(&keys),
)
.await
{
Ok(Ok(stored)) => stored,
Ok(Err(_)) | Err(_) => return false,
};
let Some(Some(stored_binding)) = stored.first() else {
return false;
};
actual.validate().is_ok() && actual == *expected
let Ok(actual) = serde_json::from_str::<LiveCallBinding>(stored_binding.as_str()) else {
return false;
};
if actual.validate().is_err() || actual != *expected {
return false;
}
match (stored.get(1).and_then(Option::as_ref), expected_context) {
(None, None) => true,
(Some(stored), Some(expected)) => {
serde_json::from_str::<LiveProviderOutboundContextRecord>(stored.as_str())
.is_ok_and(|actual| actual == *expected)
}
_ => false,
}
}
async fn register_locked(
&self,
key: &str,
context_storage_key: &str,
index: &str,
serialized: String,
serialized_context: Option<String>,
binding: &LiveCallBinding,
context_record: Option<&LiveProviderOutboundContextRecord>,
) -> Result<(), LiveCallRegistryError> {
if let Some(existing) = self
let existing = self
.runtime_state
.kv_get(key)
.kv_get_many(&[key.to_string(), context_storage_key.to_string()])
.await
.map_err(LiveCallRegistryError::Storage)?
{
.map_err(LiveCallRegistryError::Storage)?;
let existing_binding = existing.first().and_then(Option::as_ref);
if let Some(existing) = existing_binding {
let existing = serde_json::from_str::<LiveCallBinding>(existing.as_str())
.map_err(LiveCallRegistryError::CorruptRecord)?;
if existing != *binding {
return Err(LiveCallRegistryError::OwnershipConflict);
}
}
self.runtime_state
let existing_context = existing.get(1).and_then(Option::as_ref);
if existing_binding.is_some() {
match (existing_context, context_record) {
(Some(existing), Some(expected)) => {
let existing = serde_json::from_str::<LiveProviderOutboundContextRecord>(
existing.as_str(),
)
.map_err(LiveCallRegistryError::CorruptRecord)?;
if existing != *expected {
return Err(LiveCallRegistryError::OwnershipConflict);
}
}
(Some(_), None) => return Err(LiveCallRegistryError::OwnershipConflict),
_ => {}
}
} else if existing_context.is_some() && serialized_context.is_none() {
// Old binaries and interrupted writes can leave a companion after
// the authoritative base record is gone. Clear it under the
// principal lock before committing a legacy/no-context binding.
self.runtime_state
.kv_delete(context_storage_key)
.await
.map_err(LiveCallRegistryError::Storage)?;
}
if let Some(serialized_context) = serialized_context {
self.runtime_state
// Keep the companion alive through the base record's expiry
// grace period. This prevents a normal TTL race from causing
// a sideband retry to mint a different provider identity.
.kv_set(
context_storage_key,
serialized_context,
Some(self.ttl.saturating_add(EXPIRED_LOOKUP_GRACE)),
)
.await
.map_err(LiveCallRegistryError::Storage)?;
}
if let Err(error) = self
.runtime_state
.kv_set(key, serialized, Some(self.ttl))
.await
.map_err(LiveCallRegistryError::Storage)?;
{
// Keep the companion until exact commit verification. A Redis
// timeout can happen after the base write committed; deleting the
// companion here would turn an idempotent retry into a split record.
// If the base did not commit, the bounded companion expires with
// the registry grace period and a later retry can safely replace it.
return Err(LiveCallRegistryError::Storage(error));
}
if let Err(error) = self
.runtime_state
.score_set(index, key, now_unix_ms() as f64)
.await
{
let _ = self.runtime_state.kv_delete(key).await;
let _ = self.runtime_state.kv_delete(context_storage_key).await;
return Err(LiveCallRegistryError::Storage(error));
}
if let Err(error) = self
@@ -572,6 +815,7 @@ impl LiveCallRegistry {
{
let _ = self.runtime_state.score_remove(index, key).await;
let _ = self.runtime_state.kv_delete(key).await;
let _ = self.runtime_state.kv_delete(context_storage_key).await;
return Err(LiveCallRegistryError::Storage(error));
}
let members = self
@@ -581,10 +825,15 @@ impl LiveCallRegistry {
.map_err(LiveCallRegistryError::Storage)?;
let overflow = members.len().saturating_sub(self.max_records_per_principal);
for oldest in members.into_iter().take(overflow) {
let oldest_context = context_key(oldest.as_str())?;
self.runtime_state
.kv_delete(oldest.as_str())
.await
.map_err(LiveCallRegistryError::Storage)?;
self.runtime_state
.kv_delete(oldest_context.as_str())
.await
.map_err(LiveCallRegistryError::Storage)?;
self.runtime_state
.score_remove(index, oldest.as_str())
.await
@@ -621,6 +870,21 @@ fn record_key(
))
}
fn context_key(record_key: &str) -> Result<String, LiveCallRegistryError> {
if !record_key.starts_with(RECORD_PREFIX)
|| record_key.len() != RECORD_PREFIX.len() + 64
|| !record_key[RECORD_PREFIX.len()..]
.bytes()
.all(|byte| byte.is_ascii_hexdigit())
{
return Err(LiveCallRegistryError::InvalidIdentity("invalid_record_key"));
}
Ok(format!(
"{CONTEXT_PREFIX}{}",
digest(CONTEXT_DOMAIN, &[record_key])
))
}
fn index_key(user_id: &str, api_key_id: &str) -> Result<String, LiveCallRegistryError> {
validate_principal(user_id, "invalid_user_id")?;
validate_principal(api_key_id, "invalid_api_key_id")?;
@@ -660,6 +924,19 @@ fn validate_principal(value: &str, error: &'static str) -> Result<(), LiveCallRe
Ok(())
}
fn validate_context_field(value: &str, error: &'static str) -> Result<(), LiveCallRegistryError> {
let encoded_len = serde_json::to_string(value)
.map(|encoded| encoded.len())
.unwrap_or(usize::MAX);
if value.trim().is_empty()
|| value.len() > MAX_CONTEXT_FIELD_BYTES
|| encoded_len > MAX_CONTEXT_FIELD_BYTES.saturating_add(2)
{
return Err(LiveCallRegistryError::InvalidRecord(error));
}
Ok(())
}
fn digest(domain: &[u8], components: &[&str]) -> String {
let mut digest = Sha256::new();
digest.update(domain);
@@ -704,6 +981,13 @@ mod tests {
}
}
fn provider_outbound_context(logical_turn_id: &str) -> ProviderOutboundRequestContext {
ProviderOutboundRequestContext::new(logical_turn_id, 1_700_000_000_123)
.with_original_turn_id("client-turn")
.with_original_client_session_id("client-session")
.with_original_prompt_cache_key("client-cache")
}
fn candidate_for_binding(binding: &LiveCallBinding) -> PlannedLiveCandidate {
let execution: crate::ai_serving::AiExecutionDecision =
serde_json::from_value(serde_json::json!({
@@ -718,11 +1002,7 @@ mod tests {
PlannedLiveCandidate {
execution,
pinned_candidate: binding.pinned_candidate.clone(),
codex_fingerprint_context:
aether_provider_transport::CodexFingerprintConvergenceContext::new(
"test-live-turn",
1,
),
provider_outbound_context: ProviderOutboundRequestContext::new("test-live-turn", 1),
client_model: binding.client_model.clone(),
provider_model: binding.provider_model.clone(),
auth_mode: binding.auth_mode,
@@ -735,7 +1015,13 @@ mod tests {
let state = runtime_state();
let registry = LiveCallRegistry::new(Arc::clone(&state));
registry
.register("user-1", "api-key-1", "rtc_secret", &binding("global"))
.register(
"user-1",
"api-key-1",
"rtc_secret",
&binding("global"),
None,
)
.await
.unwrap();
assert!(registry
@@ -755,14 +1041,117 @@ mod tests {
.is_none());
}
#[tokio::test]
async fn provider_outbound_context_round_trips_in_companion_record() {
let state = runtime_state();
let registry = LiveCallRegistry::new(Arc::clone(&state));
let binding = binding("global");
let context = provider_outbound_context("logical-live-turn");
registry
.register("user", "key", "rtc_context", &binding, Some(&context))
.await
.unwrap();
let record = registry
.lookup("user", "key", "rtc_context")
.await
.unwrap()
.expect("binding should exist");
assert_eq!(record.binding(), &binding);
assert_eq!(record.provider_outbound_context(), Some(&context));
let key = record_key("user", "key", "rtc_context").unwrap();
let serialized_binding = state.kv_get(key.as_str()).await.unwrap().unwrap();
let strict_v2 = serde_json::from_str::<LiveCallBinding>(&serialized_binding).unwrap();
assert_eq!(strict_v2, binding);
assert!(!serialized_binding.contains("provider_outbound_context"));
assert!(!serialized_binding.contains("logical_turn_id"));
}
#[test]
fn bounded_context_fields_fit_the_companion_record_budget() {
let binding = binding("global");
let value = "x".repeat(PROVIDER_OUTBOUND_CONTEXT_MAX_VALUE_BYTES);
let context = ProviderOutboundRequestContext::new(value.clone(), 1)
.with_original_turn_id(value.clone())
.with_original_client_session_id(value.clone())
.with_original_prompt_cache_key(value);
let record = LiveProviderOutboundContextRecord::from_context(&binding, &context)
.expect("bounded context should be valid");
let serialized = serde_json::to_string(&record).expect("context should serialize");
assert!(serialized.len() <= MAX_SERIALIZED_CONTEXT_BYTES);
}
#[tokio::test]
async fn legacy_binding_without_companion_context_still_loads() {
let state = runtime_state();
let registry = LiveCallRegistry::new(Arc::clone(&state));
registry
.register("user", "key", "rtc_legacy", &binding("global"), None)
.await
.unwrap();
let record = registry
.lookup("user", "key", "rtc_legacy")
.await
.unwrap()
.expect("legacy binding should exist");
assert!(record.provider_outbound_context().is_none());
}
#[tokio::test]
async fn corrupt_or_mismatched_companion_context_is_rejected() {
let state = runtime_state();
let registry = LiveCallRegistry::new(Arc::clone(&state));
let binding = binding("global");
let context = provider_outbound_context("logical-live-turn");
registry
.register("user", "key", "rtc_corrupt", &binding, Some(&context))
.await
.unwrap();
let key = record_key("user", "key", "rtc_corrupt").unwrap();
let context_key = context_key(key.as_str()).unwrap();
state
.kv_set(
context_key.as_str(),
"not-json",
Some(Duration::from_secs(60)),
)
.await
.unwrap();
assert!(matches!(
registry.lookup("user", "key", "rtc_corrupt").await,
Err(LiveCallRegistryError::CorruptRecord(_))
));
let mut mismatch = LiveProviderOutboundContextRecord::from_context(&binding, &context)
.expect("context should be valid");
mismatch.binding_created_at_unix_ms = mismatch.binding_created_at_unix_ms.saturating_add(1);
state
.kv_set(
context_key.as_str(),
serde_json::to_string(&mismatch).unwrap(),
Some(Duration::from_secs(60)),
)
.await
.unwrap();
assert!(matches!(
registry.lookup("user", "key", "rtc_corrupt").await,
Err(LiveCallRegistryError::InvalidRecord(
"context_binding_mismatch"
))
));
}
#[tokio::test]
async fn capacity_evicts_the_oldest_binding() {
let state = runtime_state();
let registry =
LiveCallRegistry::with_limits(Arc::clone(&state), Duration::from_secs(60), 2);
let context = provider_outbound_context("logical-live-turn");
for call_id in ["rtc_1", "rtc_2", "rtc_3"] {
registry
.register("user", "key", call_id, &binding(call_id))
.register("user", "key", call_id, &binding(call_id), Some(&context))
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(2)).await;
@@ -772,6 +1161,12 @@ mod tests {
.await
.unwrap()
.is_none());
let evicted_key = record_key("user", "key", "rtc_1").unwrap();
assert!(state
.kv_get(context_key(evicted_key.as_str()).unwrap().as_str())
.await
.unwrap()
.is_none());
assert!(registry
.lookup("user", "key", "rtc_2")
.await
@@ -789,8 +1184,15 @@ mod tests {
let state = runtime_state();
let registry =
LiveCallRegistry::with_limits(Arc::clone(&state), Duration::from_millis(5), 2);
let context = provider_outbound_context("logical-live-turn");
registry
.register("user", "key", "rtc_expiring", &binding("global"))
.register(
"user",
"key",
"rtc_expiring",
&binding("global"),
Some(&context),
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(20)).await;
@@ -806,6 +1208,53 @@ mod tests {
.unwrap(),
LiveCallLookup::Expired
);
let expired_key = record_key("user", "key", "rtc_expiring").unwrap();
assert!(state
.kv_get(context_key(expired_key.as_str()).unwrap().as_str())
.await
.unwrap()
.is_some());
}
#[tokio::test]
async fn orphan_companion_does_not_block_a_new_binding() {
let state = runtime_state();
let registry = LiveCallRegistry::new(Arc::clone(&state));
let old_binding = binding("old");
let old_context = provider_outbound_context("old-turn");
let key = record_key("user", "key", "rtc_reused").unwrap();
let context_storage_key = context_key(key.as_str()).unwrap();
let old_record =
LiveProviderOutboundContextRecord::from_context(&old_binding, &old_context).unwrap();
state
.kv_set(
context_storage_key.as_str(),
serde_json::to_string(&old_record).unwrap(),
Some(Duration::from_secs(60)),
)
.await
.unwrap();
let new_binding = binding("new");
let new_context = provider_outbound_context("new-turn");
registry
.register(
"user",
"key",
"rtc_reused",
&new_binding,
Some(&new_context),
)
.await
.unwrap();
let record = registry
.lookup("user", "key", "rtc_reused")
.await
.unwrap()
.expect("new binding should be committed");
assert_eq!(record.binding(), &new_binding);
assert_eq!(record.provider_outbound_context(), Some(&new_context));
}
#[tokio::test]
@@ -814,19 +1263,22 @@ mod tests {
let registry = LiveCallRegistry::new(Arc::clone(&state));
let original = binding("global-a");
registry
.register("user", "key", "rtc_shared", &original)
.register("user", "key", "rtc_shared", &original, None)
.await
.unwrap();
assert!(matches!(
registry
.register("user", "key", "rtc_shared", &binding("global-b"))
.register("user", "key", "rtc_shared", &binding("global-b"), None,)
.await,
Err(LiveCallRegistryError::OwnershipConflict)
));
assert_eq!(
registry.lookup("user", "key", "rtc_shared").await.unwrap(),
Some(original)
Some(LiveCallRecord {
binding: original,
provider_outbound_context: None,
})
);
}
@@ -947,7 +1399,7 @@ mod tests {
for call_id in [".", "..", "rtc/escape"] {
assert!(matches!(
registry
.register("user", "key", call_id, &binding("global"))
.register("user", "key", call_id, &binding("global"), None)
.await,
Err(LiveCallRegistryError::InvalidIdentity("invalid_call_id"))
));
@@ -992,11 +1444,17 @@ mod tests {
let state = runtime_state();
let registry = LiveCallRegistry::new(Arc::clone(&state));
let key = record_key("user", "key", "rtc_verify").unwrap();
let context_storage_key = context_key(key.as_str()).unwrap();
let expected = binding("global");
assert!(
!registry
.exact_binding_is_committed(key.as_str(), &expected)
.exact_binding_is_committed(
key.as_str(),
context_storage_key.as_str(),
&expected,
None
)
.await
);
@@ -1006,7 +1464,12 @@ mod tests {
.unwrap();
assert!(
!registry
.exact_binding_is_committed(key.as_str(), &expected)
.exact_binding_is_committed(
key.as_str(),
context_storage_key.as_str(),
&expected,
None
)
.await
);
@@ -1020,7 +1483,12 @@ mod tests {
.unwrap();
assert!(
!registry
.exact_binding_is_committed(key.as_str(), &expected)
.exact_binding_is_committed(
key.as_str(),
context_storage_key.as_str(),
&expected,
None
)
.await
);
@@ -1034,7 +1502,12 @@ mod tests {
.unwrap();
assert!(
registry
.exact_binding_is_committed(key.as_str(), &expected)
.exact_binding_is_committed(
key.as_str(),
context_storage_key.as_str(),
&expected,
None
)
.await
);
}
apps/aether-gateway/src/handlers/proxy/websocket/live/session.rs
+9 -6
View File
@@ -43,7 +43,7 @@ use super::protocol::{
LEGACY_LIVE_CALL_PATH, REALTIME_SIDEBAND_PATH,
};
use super::registry::{
LiveCallBinding, LiveCallLookup, LiveCallRegistry, LiveCallRegistryError, LiveSidebandLease,
LiveCallLookup, LiveCallRecord, LiveCallRegistry, LiveCallRegistryError, LiveSidebandLease,
LiveSidebandLeaseLoss,
};
@@ -302,8 +302,8 @@ pub(super) async fn prepare_live_websocket(
),
)
.await;
let binding = match lookup {
Ok(Ok(LiveCallLookup::Found(binding))) => binding,
let record = match lookup {
Ok(Ok(LiveCallLookup::Found(record))) => record,
Ok(Ok(LiveCallLookup::Missing)) => {
info!(
target: LIVE_LOG_TARGET,
@@ -360,7 +360,7 @@ pub(super) async fn prepare_live_websocket(
));
}
};
prepare_sideband_live_websocket(state, context, call_id, binding, dialect)
prepare_sideband_live_websocket(state, context, call_id, record, dialect)
.await
.map(PreparedLiveWebSocket::Sideband)
}
@@ -381,6 +381,7 @@ async fn prepare_direct_live_websocket(
client_model,
dialect,
None,
None,
)
.await
{
@@ -557,10 +558,11 @@ async fn prepare_sideband_live_websocket(
state: &AppState,
context: &WebSocketRequestContext,
call_id: String,
binding: LiveCallBinding,
record: LiveCallRecord,
dialect: LiveRouteDialect,
) -> Result<PreparedLiveSideband, LiveWebSocketPreflightRejection> {
let started_at = Instant::now();
let binding = record.binding();
let Some(auth) = context.decision.auth_context.as_ref() else {
return Err(preflight_rejection(
context,
@@ -635,6 +637,7 @@ async fn prepare_sideband_live_websocket(
binding.client_model(),
dialect,
Some(binding.pinned_candidate()),
record.provider_outbound_context(),
),
)
.await;
@@ -652,7 +655,7 @@ async fn prepare_sideband_live_websocket(
Ok(Ok(LiveCandidatePlanningOutcome {
candidate: Some(candidate),
..
})) if binding.matches_candidate(&candidate) => candidate,
})) if record.matches_candidate(&candidate) => candidate,
Ok(Ok(LiveCandidatePlanningOutcome {
candidate: Some(candidate),
..