mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 01:47:47 +08:00
feat(gateway): persist OpenAI Responses continuation history
This commit is contained in:
@@ -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(
|
||||
|
||||
+104
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user