feat(gateway): persist OpenAI Responses continuation history

This commit is contained in:
ZheFox
2026-07-30 19:26:52 +08:00
parent 1ab4f079c9
commit 118f441029
25 changed files with 1881 additions and 84 deletions
@@ -1,5 +1,7 @@
use serde_json::Value;
use aether_ai_formats::api::ResponseHistoryRecord;
use crate::ai_serving::{
maybe_build_ai_surface_stream_rewriter, AiSurfaceFinalizeError, AiSurfaceStreamRewriter,
};
@@ -24,6 +26,10 @@ impl LocalStreamRewriter<'_> {
pub(crate) fn finish(&mut self) -> Result<Vec<u8>, GatewayError> {
self.inner.finish().map_err(map_surface_error)
}
pub(crate) fn take_response_history_record(&mut self) -> Option<ResponseHistoryRecord> {
self.inner.take_response_history_record()
}
}
fn map_surface_error(error: AiSurfaceFinalizeError) -> GatewayError {
@@ -3,6 +3,7 @@ pub(crate) mod api;
mod finalize;
mod planner;
mod pure;
mod response_history;
pub(crate) mod transport;
use axum::body::Body;
@@ -64,6 +65,10 @@ pub(crate) use self::planner::{
SkippedLocalExecutionCandidate,
};
pub(crate) use self::pure::*;
pub(crate) use self::response_history::{
hydrate_openai_response_history, persist_converted_response_history,
persist_response_history_record,
};
pub(crate) use self::transport::{
append_transport_diagnostics_to_value, build_request_trace_proxy_value,
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
@@ -584,6 +584,14 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
return Ok(None);
}
};
crate::ai_serving::hydrate_openai_response_history(
state.runtime_state(),
body_json,
spec_metadata.api_format,
provider_api_format,
input.auth_context.api_key_id.as_str(),
)
.await?;
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
@@ -3,7 +3,7 @@ use serde_json::Value;
use crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers;
use crate::ai_serving::{
apply_openai_responses_compact_special_body_edits,
build_cross_format_openai_responses_request_body_with_model_directives as surface_build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope as surface_build_cross_format_openai_responses_request_body,
build_local_openai_responses_request_body_with_model_directives as surface_build_local_openai_responses_request_body,
GatewayProviderTransportSnapshot,
};
@@ -105,7 +105,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
force_body_stream_field: bool,
provider_type: &str,
body_rules: Option<&Value>,
_user_api_key_id: Option<&str>,
user_api_key_id: Option<&str>,
request_headers: &http::HeaderMap,
enable_model_directives: bool,
) -> Option<Value> {
@@ -119,6 +119,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body(
provider_type,
body_rules,
request_headers,
user_api_key_id,
None,
enable_model_directives,
)
@@ -134,6 +135,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body_with_codex_model_
provider_type: &str,
body_rules: Option<&Value>,
request_headers: &http::HeaderMap,
history_scope: Option<&str>,
model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>,
enable_model_directives: bool,
) -> Option<Value> {
@@ -144,6 +146,7 @@ pub(crate) fn build_cross_format_openai_responses_request_body_with_codex_model_
provider_api_format,
upstream_is_stream,
enable_model_directives,
history_scope,
)?;
let mut provider_request_body =
apply_standard_provider_request_body_rules_with_request_headers(
@@ -82,6 +82,96 @@ fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
)
}
fn response_function_tool_names(body: &Value) -> Vec<String> {
body.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|tool| tool.get("type").and_then(Value::as_str) == Some("function"))
.filter_map(|tool| tool.get("name").and_then(Value::as_str))
.map(ToOwned::to_owned)
.collect()
}
fn chat_function_tool_names(body: &Value) -> Vec<String> {
body.get("tools")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|tool| {
tool.get("function")
.and_then(|function| function.get("name"))
.and_then(Value::as_str)
})
.map(ToOwned::to_owned)
.collect()
}
fn response_input_call_ids(body: &Value) -> Vec<String> {
body.get("input")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|item| {
matches!(
item.get("type").and_then(Value::as_str),
Some("function_call" | "function_call_output")
)
})
.filter_map(|item| item.get("call_id").and_then(Value::as_str))
.map(ToOwned::to_owned)
.collect()
}
fn chat_message_call_ids(body: &Value) -> Vec<String> {
let mut call_ids = Vec::new();
for message in body
.get("messages")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
if let Some(tool_calls) = message.get("tool_calls").and_then(Value::as_array) {
call_ids.extend(
tool_calls
.iter()
.filter_map(|tool_call| tool_call.get("id").and_then(Value::as_str))
.map(ToOwned::to_owned),
);
}
if let Some(tool_call_id) = message.get("tool_call_id").and_then(Value::as_str) {
call_ids.push(tool_call_id.to_string());
}
}
call_ids
}
fn log_responses_to_chat_tool_conversion(trace_id: &str, inbound: &Value, outbound: &Value) {
let inbound_tool_names = response_function_tool_names(inbound);
let outbound_tool_names = chat_function_tool_names(outbound);
let inbound_call_ids = response_input_call_ids(inbound);
let outbound_call_ids = chat_message_call_ids(outbound);
let previous_response_id = inbound
.get("previous_response_id")
.and_then(Value::as_str)
.unwrap_or_default();
debug!(
event_name = "openai_responses_to_chat_tool_conversion",
log_type = "debug",
trace_id = %trace_id,
inbound_tool_count = inbound_tool_names.len(),
outbound_tool_count = outbound_tool_names.len(),
inbound_tool_names = ?inbound_tool_names,
outbound_tool_names = ?outbound_tool_names,
previous_response_id = %previous_response_id,
inbound_call_ids = ?inbound_call_ids,
outbound_call_ids = ?outbound_call_ids,
history_recovered = !previous_response_id.is_empty()
&& outbound_call_ids.iter().any(|call_id| inbound_call_ids.contains(call_id)),
"converted OpenAI Responses tools and continuation context to OpenAI Chat"
);
}
pub(crate) struct LocalOpenAiResponsesCandidatePayloadParts {
pub(super) auth_header: String,
pub(super) auth_value: String,
@@ -306,6 +396,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
return Ok(None);
}
};
crate::ai_serving::hydrate_openai_response_history(
state.runtime_state(),
body_json,
spec_metadata.api_format,
provider_api_format,
input.auth_context.api_key_id.as_str(),
)
.await?;
let redaction = resolve_provider_chat_pii_redaction(
state,
parts,
@@ -367,6 +465,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
transport.endpoint.body_rules.as_ref()
},
effective_headers,
Some(input.auth_context.api_key_id.as_str()),
codex_model_capabilities.as_ref(),
false,
)
@@ -471,6 +570,11 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
.await;
return Ok(None);
}
if needs_bidirectional_conversion
&& crate::ai_serving::api_format_alias_matches(provider_api_format, "openai:chat")
{
log_responses_to_chat_tool_conversion(trace_id, body_json, &base_provider_request_body);
}
let provider_request_body = base_provider_request_body;
if let Some(kiro_auth) = kiro_auth.as_ref() {
@@ -17,6 +17,7 @@ pub(crate) use aether_ai_formats::api::{
build_cross_format_openai_chat_request_body_with_model_directives,
build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_request_body_with_model_directives,
build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope,
build_gemini_image_request_body_from_openai_image_request,
build_gemini_image_response_from_openai_image_response,
build_gemini_image_response_from_openai_responses_image_response, build_generated_tool_call_id,
@@ -0,0 +1,96 @@
use aether_ai_formats::api::{
hydrate_response_history, normalize_api_format_alias, record_converted_response_history,
response_history_is_loaded, response_history_storage_key, ResponseHistoryRecord,
};
use aether_runtime_state::RuntimeState;
use serde_json::Value;
use tracing::warn;
use crate::GatewayError;
pub(crate) async fn hydrate_openai_response_history(
runtime_state: &RuntimeState,
request: &Value,
client_api_format: &str,
provider_api_format: &str,
history_scope: &str,
) -> Result<(), GatewayError> {
if normalize_api_format_alias(client_api_format) != "openai:responses"
|| normalize_api_format_alias(provider_api_format) != "openai:chat"
{
return Ok(());
}
let Some(previous_response_id) = request
.get("previous_response_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(());
};
if response_history_is_loaded(previous_response_id, Some(history_scope)) {
return Ok(());
}
let storage_key = response_history_storage_key(previous_response_id, Some(history_scope));
let payload = runtime_state.kv_get(&storage_key).await.map_err(|error| {
warn!(
event_name = "openai_response_history_read_failed",
log_type = "ops",
backend = runtime_state.backend_kind().as_str(),
error = ?error,
"gateway failed to read shared OpenAI response history"
);
GatewayError::Internal("OpenAI response history lookup failed".to_string())
})?;
let Some(payload) = payload else {
return Ok(());
};
if let Err(error) =
hydrate_response_history(previous_response_id, Some(history_scope), &payload)
{
let _ = runtime_state.kv_delete(&storage_key).await;
warn!(
event_name = "openai_response_history_invalid",
log_type = "ops",
backend = runtime_state.backend_kind().as_str(),
error = %error,
"gateway rejected invalid shared OpenAI response history"
);
return Err(GatewayError::Internal(
"OpenAI response history validation failed".to_string(),
));
}
Ok(())
}
pub(crate) async fn persist_response_history_record(
runtime_state: &RuntimeState,
record: ResponseHistoryRecord,
) {
if let Err(error) = runtime_state
.kv_set(&record.storage_key, record.payload, Some(record.ttl))
.await
{
warn!(
event_name = "openai_response_history_write_failed",
log_type = "ops",
backend = runtime_state.backend_kind().as_str(),
error = ?error,
"gateway failed to persist shared OpenAI response history"
);
}
}
pub(crate) async fn persist_converted_response_history(
runtime_state: &RuntimeState,
report_context: &Value,
response: Option<&Value>,
) {
let Some(response) = response else {
return;
};
if let Some(record) = record_converted_response_history(report_context, response) {
persist_response_history_record(runtime_state, record).await;
}
}
@@ -6366,6 +6366,13 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
report_context.as_ref(),
) {
Ok(Some(outcome)) => {
if let Some(record) = outcome.response_history_record {
crate::ai_serving::persist_response_history_record(
state.runtime_state(),
record,
)
.await;
}
headers.remove("content-encoding");
headers.remove("content-length");
headers.insert(
@@ -6570,6 +6577,15 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
if stream_commit_gate.is_uncommitted() {
stream_commit_gate.commit();
}
let prefetched_response_history_persisted = if let Some(record) = local_stream_rewriter
.as_mut()
.and_then(|rewriter| rewriter.take_response_history_record())
{
crate::ai_serving::persist_response_history_record(state.runtime_state(), record).await;
true
} else {
false
};
drop(private_stream_normalizer);
drop(local_stream_rewriter);
@@ -6906,6 +6922,11 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
}
}
}
if prefetched_response_history_persisted {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
let _ = rewriter.take_response_history_record();
}
}
}
if terminal_failure.is_none() && !reached_eof {
@@ -7159,6 +7180,19 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
normalized_chunk
};
if provider_private_error_body_json.is_none() {
if let Some(record) = local_stream_rewriter
.as_mut()
.and_then(|rewriter| rewriter.take_response_history_record())
{
crate::ai_serving::persist_response_history_record(
state_for_report.runtime_state(),
record,
)
.await;
}
}
if rewritten_chunk.is_empty() {
if let Some(error_body_json) = provider_private_error_body_json {
let error_status_code = resolve_provider_stream_error_status_code(
@@ -7350,6 +7384,18 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
} else {
normalized_chunk
};
if provider_private_error_body_json.is_none() {
if let Some(record) = local_stream_rewriter
.as_mut()
.and_then(|rewriter| rewriter.take_response_history_record())
{
crate::ai_serving::persist_response_history_record(
state_for_report.runtime_state(),
record,
)
.await;
}
}
if !rewritten_chunk.is_empty() {
append_stream_capture_bytes(
&mut buffered_body,
@@ -7423,7 +7469,15 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
}
if !downstream_dropped && terminal_failure.is_none() {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
match rewriter.finish() {
let finish_result = rewriter.finish();
if let Some(record) = rewriter.take_response_history_record() {
crate::ai_serving::persist_response_history_record(
state_for_report.runtime_state(),
record,
)
.await;
}
match finish_result {
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
append_stream_capture_bytes(
&mut buffered_body,
@@ -7481,6 +7535,19 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
}
}
if terminal_failure.is_none() {
if let Some(record) = local_stream_rewriter
.as_mut()
.and_then(|rewriter| rewriter.take_response_history_record())
{
crate::ai_serving::persist_response_history_record(
state_for_report.runtime_state(),
record,
)
.await;
}
}
if !downstream_dropped {
if let Some(failure) = terminal_failure.as_ref() {
let terminal_event = if is_openai_image_stream_for_report {
@@ -1789,6 +1789,17 @@ async fn apply_sync_success_effects(
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
) {
if let Some(report_context) = report_context {
crate::ai_serving::persist_converted_response_history(
state.runtime_state(),
report_context,
payload
.client_body_json
.as_ref()
.or(payload.body_json.as_ref()),
)
.await;
}
apply_local_execution_effect(
state,
LocalExecutionEffectContext {