mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
9a84a6ff6c
Move protocol/request/response format modules under provider-oriented formats modules and update registry, transport, and architecture paths.
530 lines
18 KiB
Rust
530 lines
18 KiB
Rust
use std::collections::BTreeMap;
|
|
|
|
use serde_json::Value;
|
|
|
|
use crate::formats::shared::AiSurfaceFinalizeError;
|
|
use crate::provider_compat::kiro_stream::KiroToClaudeCliStreamState;
|
|
|
|
use super::surfaces::{
|
|
provider_adaptation_allows_sync_finalize_envelope, provider_adaptation_descriptor_for_envelope,
|
|
provider_adaptation_should_unwrap_stream_envelope, ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME,
|
|
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, KIRO_ENVELOPE_NAME,
|
|
};
|
|
|
|
pub fn provider_private_response_allows_sync_finalize(report_context: &Value) -> bool {
|
|
let has_envelope = report_context
|
|
.get("has_envelope")
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false);
|
|
if !has_envelope {
|
|
return true;
|
|
}
|
|
let envelope_name = report_context
|
|
.get("envelope_name")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
let provider_api_format = report_context
|
|
.get("provider_api_format")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
provider_adaptation_allows_sync_finalize_envelope(envelope_name, provider_api_format)
|
|
}
|
|
|
|
pub fn normalize_provider_private_report_context(report_context: Option<&Value>) -> Option<Value> {
|
|
let report_context = report_context?;
|
|
if !report_context
|
|
.get("has_envelope")
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false)
|
|
{
|
|
return Some(report_context.clone());
|
|
}
|
|
let envelope_name = report_context
|
|
.get("envelope_name")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
let provider_api_format = report_context
|
|
.get("provider_api_format")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
if provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format).is_none() {
|
|
return Some(report_context.clone());
|
|
}
|
|
Some(clear_private_envelope_context(report_context))
|
|
}
|
|
|
|
pub fn normalize_provider_private_response_value(
|
|
data: Value,
|
|
report_context: &Value,
|
|
) -> Option<Value> {
|
|
if !report_context
|
|
.get("has_envelope")
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false)
|
|
{
|
|
return Some(data);
|
|
}
|
|
let mut unwrapped = match report_context.get("envelope_name").and_then(Value::as_str) {
|
|
Some(KIRO_ENVELOPE_NAME) => data,
|
|
Some(GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME) => {
|
|
if let Some(response) = data
|
|
.get("response")
|
|
.and_then(Value::as_object)
|
|
.filter(|response| !response.contains_key("response"))
|
|
{
|
|
Value::Object(response.clone())
|
|
} else {
|
|
data
|
|
}
|
|
}
|
|
Some(ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME) => {
|
|
if let Some(response) = data
|
|
.get("response")
|
|
.and_then(Value::as_object)
|
|
.filter(|response| !response.contains_key("response"))
|
|
{
|
|
let mut unwrapped = response.clone();
|
|
if let Some(response_id) = data.get("responseId").cloned() {
|
|
unwrapped.insert("_v1internal_response_id".to_string(), response_id);
|
|
}
|
|
Value::Object(unwrapped)
|
|
} else {
|
|
data
|
|
}
|
|
}
|
|
_ => return None,
|
|
};
|
|
postprocess_private_response_value(&mut unwrapped, report_context);
|
|
Some(unwrapped)
|
|
}
|
|
|
|
pub fn transform_provider_private_stream_line(
|
|
report_context: &Value,
|
|
line: Vec<u8>,
|
|
) -> Result<Vec<u8>, serde_json::Error> {
|
|
let Ok(text) = std::str::from_utf8(&line) else {
|
|
return Ok(line);
|
|
};
|
|
let trimmed = text.trim_matches('\r').trim();
|
|
if trimmed.is_empty() || trimmed.starts_with(':') || trimmed.starts_with("event:") {
|
|
return Ok(Vec::new());
|
|
}
|
|
let Some(data_line) = trimmed.strip_prefix("data:") else {
|
|
return Ok(line);
|
|
};
|
|
let data_line = data_line.trim();
|
|
if data_line.is_empty() || data_line == "[DONE]" {
|
|
return Ok(line);
|
|
}
|
|
|
|
let body: Value = match serde_json::from_str(data_line) {
|
|
Ok(value) => value,
|
|
Err(_) => return Ok(line),
|
|
};
|
|
|
|
let envelope_name = report_context
|
|
.get("envelope_name")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
let provider_api_format = report_context
|
|
.get("provider_api_format")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
if !provider_adaptation_should_unwrap_stream_envelope(envelope_name, provider_api_format) {
|
|
return Ok(line);
|
|
}
|
|
let unwrapped = match envelope_name {
|
|
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME => body.get("response").cloned().unwrap_or(body),
|
|
ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME => {
|
|
let mut response = body.get("response").cloned().unwrap_or(body.clone());
|
|
if let Some(response_id) = body.get("responseId").cloned() {
|
|
if let Some(object) = response.as_object_mut() {
|
|
object
|
|
.entry("_v1internal_response_id".to_string())
|
|
.or_insert(response_id);
|
|
}
|
|
}
|
|
inject_antigravity_stream_tool_ids(&mut response);
|
|
response
|
|
}
|
|
_ => body,
|
|
};
|
|
|
|
let mut out = b"data: ".to_vec();
|
|
out.extend(serde_json::to_vec(&unwrapped)?);
|
|
out.extend_from_slice(b"\n\n");
|
|
Ok(out)
|
|
}
|
|
|
|
enum ProviderPrivateStreamNormalizeMode {
|
|
EnvelopeUnwrap,
|
|
KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
|
|
}
|
|
|
|
pub struct ProviderPrivateStreamNormalizer<'a> {
|
|
report_context: &'a Value,
|
|
buffered: Vec<u8>,
|
|
mode: ProviderPrivateStreamNormalizeMode,
|
|
}
|
|
|
|
pub fn maybe_build_provider_private_stream_normalizer<'a>(
|
|
report_context: Option<&'a Value>,
|
|
) -> Option<ProviderPrivateStreamNormalizer<'a>> {
|
|
let report_context = report_context?;
|
|
if !report_context
|
|
.get("has_envelope")
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false)
|
|
{
|
|
return None;
|
|
}
|
|
let envelope_name = report_context
|
|
.get("envelope_name")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
let provider_api_format = report_context
|
|
.get("provider_api_format")
|
|
.and_then(Value::as_str)
|
|
.unwrap_or_default();
|
|
let descriptor =
|
|
provider_adaptation_descriptor_for_envelope(envelope_name, provider_api_format)?;
|
|
let mode = if descriptor
|
|
.envelope_name
|
|
.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME)
|
|
{
|
|
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(Box::new(
|
|
KiroToClaudeCliStreamState::new(report_context),
|
|
))
|
|
} else if descriptor.unwraps_response_envelope {
|
|
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap
|
|
} else {
|
|
return None;
|
|
};
|
|
Some(ProviderPrivateStreamNormalizer {
|
|
report_context,
|
|
buffered: Vec::new(),
|
|
mode,
|
|
})
|
|
}
|
|
|
|
impl ProviderPrivateStreamNormalizer<'_> {
|
|
pub fn push_chunk(&mut self, chunk: &[u8]) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
|
match &mut self.mode {
|
|
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => {
|
|
state.push_chunk(self.report_context, chunk)
|
|
}
|
|
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => {
|
|
self.buffered.extend_from_slice(chunk);
|
|
let mut output = Vec::new();
|
|
while let Some(line_end) = self.buffered.iter().position(|byte| *byte == b'\n') {
|
|
let line = self.buffered.drain(..=line_end).collect::<Vec<_>>();
|
|
output.extend(
|
|
transform_provider_private_stream_line(self.report_context, line)
|
|
.map_err(AiSurfaceFinalizeError::from)?,
|
|
);
|
|
}
|
|
Ok(output)
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
|
|
match &mut self.mode {
|
|
ProviderPrivateStreamNormalizeMode::KiroToClaudeCli(state) => {
|
|
state.finish(self.report_context)
|
|
}
|
|
ProviderPrivateStreamNormalizeMode::EnvelopeUnwrap => {
|
|
if self.buffered.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
let line = std::mem::take(&mut self.buffered);
|
|
transform_provider_private_stream_line(self.report_context, line)
|
|
.map_err(AiSurfaceFinalizeError::from)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn stream_body_contains_error_event(body: &[u8]) -> bool {
|
|
let Ok(text) = std::str::from_utf8(body) else {
|
|
return false;
|
|
};
|
|
let mut current_event_type: Option<String> = None;
|
|
for raw_line in text.lines() {
|
|
let line = raw_line.trim_matches('\r').trim();
|
|
if line.is_empty() || line.starts_with(':') {
|
|
continue;
|
|
}
|
|
if let Some(event_name) = line.strip_prefix("event:") {
|
|
current_event_type = Some(event_name.trim().to_string());
|
|
continue;
|
|
}
|
|
let data_line = if let Some(rest) = line.strip_prefix("data:") {
|
|
rest.trim()
|
|
} else {
|
|
line
|
|
};
|
|
if data_line.is_empty() || data_line == "[DONE]" {
|
|
continue;
|
|
}
|
|
let Ok(mut event) = serde_json::from_str::<Value>(data_line) else {
|
|
continue;
|
|
};
|
|
if let Some(event_object) = event.as_object_mut() {
|
|
if !event_object.contains_key("type") {
|
|
if let Some(event_name) = current_event_type.take() {
|
|
event_object.insert("type".to_string(), Value::String(event_name));
|
|
}
|
|
}
|
|
}
|
|
if event
|
|
.get("type")
|
|
.and_then(Value::as_str)
|
|
.is_some_and(|value| value.eq_ignore_ascii_case("error"))
|
|
{
|
|
return true;
|
|
}
|
|
current_event_type = None;
|
|
}
|
|
false
|
|
}
|
|
|
|
fn clear_private_envelope_context(report_context: &Value) -> Value {
|
|
let mut normalized = report_context.clone();
|
|
if let Some(object) = normalized.as_object_mut() {
|
|
object.insert("has_envelope".to_string(), Value::Bool(false));
|
|
object.remove("envelope_name");
|
|
}
|
|
normalized
|
|
}
|
|
|
|
fn local_finalize_response_model(report_context: &Value) -> &str {
|
|
report_context
|
|
.get("mapped_model")
|
|
.and_then(Value::as_str)
|
|
.or_else(|| report_context.get("model").and_then(Value::as_str))
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
fn inject_antigravity_stream_tool_ids(value: &mut Value) {
|
|
let Some(candidates) = value.get_mut("candidates").and_then(Value::as_array_mut) else {
|
|
return;
|
|
};
|
|
|
|
for candidate in candidates {
|
|
let Some(parts) = candidate
|
|
.get_mut("content")
|
|
.and_then(Value::as_object_mut)
|
|
.and_then(|content| content.get_mut("parts"))
|
|
.and_then(Value::as_array_mut)
|
|
else {
|
|
continue;
|
|
};
|
|
|
|
let mut counters: BTreeMap<String, usize> = BTreeMap::new();
|
|
for part in parts {
|
|
let Some(function_call) = part.get_mut("functionCall").and_then(Value::as_object_mut)
|
|
else {
|
|
continue;
|
|
};
|
|
let has_id = function_call
|
|
.get("id")
|
|
.and_then(Value::as_str)
|
|
.is_some_and(|value| !value.is_empty());
|
|
if has_id {
|
|
continue;
|
|
}
|
|
let name = function_call
|
|
.get("name")
|
|
.and_then(Value::as_str)
|
|
.filter(|value| !value.is_empty())
|
|
.unwrap_or("unknown")
|
|
.to_string();
|
|
let index = counters.entry(name.clone()).or_insert(0);
|
|
function_call.insert(
|
|
"id".to_string(),
|
|
Value::String(format!("call_{name}_{index}")),
|
|
);
|
|
*index += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn inject_antigravity_sync_tool_ids(response: &mut Value, model: &str) {
|
|
if !model.to_ascii_lowercase().contains("claude") {
|
|
return;
|
|
}
|
|
|
|
let Some(candidates) = response.get_mut("candidates").and_then(Value::as_array_mut) else {
|
|
return;
|
|
};
|
|
|
|
for candidate in candidates {
|
|
let Some(parts) = candidate
|
|
.get_mut("content")
|
|
.and_then(Value::as_object_mut)
|
|
.and_then(|content| content.get_mut("parts"))
|
|
.and_then(Value::as_array_mut)
|
|
else {
|
|
continue;
|
|
};
|
|
|
|
let mut name_counters: BTreeMap<String, usize> = BTreeMap::new();
|
|
for part in parts {
|
|
let function_call = if let Some(function_call) =
|
|
part.get_mut("functionCall").and_then(Value::as_object_mut)
|
|
{
|
|
function_call
|
|
} else if let Some(function_call) =
|
|
part.get_mut("function_call").and_then(Value::as_object_mut)
|
|
{
|
|
function_call
|
|
} else {
|
|
continue;
|
|
};
|
|
let has_id = function_call
|
|
.get("id")
|
|
.and_then(Value::as_str)
|
|
.is_some_and(|value| !value.is_empty());
|
|
if has_id {
|
|
continue;
|
|
}
|
|
let function_name = function_call
|
|
.get("name")
|
|
.and_then(Value::as_str)
|
|
.filter(|value| !value.is_empty())
|
|
.unwrap_or("unknown")
|
|
.to_string();
|
|
let count = name_counters.entry(function_name.clone()).or_insert(0);
|
|
function_call.insert(
|
|
"id".to_string(),
|
|
Value::String(format!("call_{function_name}_{count}")),
|
|
);
|
|
*count += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
fn postprocess_private_response_value(data: &mut Value, report_context: &Value) {
|
|
if !matches!(
|
|
report_context.get("envelope_name").and_then(Value::as_str),
|
|
Some(ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME)
|
|
) {
|
|
return;
|
|
}
|
|
if let Some(object) = data.as_object_mut() {
|
|
if !object.contains_key("_v1internal_response_id") {
|
|
if let Some(response_id) = object.remove("responseId") {
|
|
object.insert("_v1internal_response_id".to_string(), response_id);
|
|
}
|
|
}
|
|
}
|
|
inject_antigravity_sync_tool_ids(data, local_finalize_response_model(report_context));
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use serde_json::json;
|
|
|
|
use super::{
|
|
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
|
|
normalize_provider_private_response_value, stream_body_contains_error_event,
|
|
transform_provider_private_stream_line,
|
|
};
|
|
|
|
#[test]
|
|
fn normalizes_supported_private_report_context() {
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"envelope_name": "antigravity:v1internal",
|
|
"provider_api_format": "gemini:generate_content",
|
|
});
|
|
let normalized = normalize_provider_private_report_context(Some(&report_context))
|
|
.expect("context should normalize");
|
|
assert_eq!(normalized["has_envelope"], json!(false));
|
|
assert!(normalized.get("envelope_name").is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn unwraps_antigravity_sync_response_and_injects_ids() {
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"provider_api_format": "gemini:generate_content",
|
|
"envelope_name": "antigravity:v1internal",
|
|
"mapped_model": "claude-sonnet-4-5",
|
|
});
|
|
let body = json!({
|
|
"response": {
|
|
"candidates": [{
|
|
"content": {
|
|
"parts": [{
|
|
"functionCall": {
|
|
"name": "get_weather",
|
|
"args": {"city": "SF"}
|
|
}
|
|
}]
|
|
}
|
|
}]
|
|
},
|
|
"responseId": "resp_123"
|
|
});
|
|
|
|
let normalized = normalize_provider_private_response_value(body, &report_context)
|
|
.expect("body should normalize");
|
|
assert_eq!(normalized["_v1internal_response_id"], json!("resp_123"));
|
|
assert_eq!(
|
|
normalized["candidates"][0]["content"]["parts"][0]["functionCall"]["id"],
|
|
json!("call_get_weather_0")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn unwraps_antigravity_stream_line_and_injects_ids() {
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"provider_api_format": "gemini:generate_content",
|
|
"client_api_format": "gemini:generate_content",
|
|
"envelope_name": "antigravity:v1internal",
|
|
"mapped_model": "claude-sonnet-4-5",
|
|
});
|
|
let output = transform_provider_private_stream_line(
|
|
&report_context,
|
|
b"data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{\"city\":\"SF\"}}}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"claude-sonnet-4-5\"},\"responseId\":\"resp_123\"}\n\n".to_vec(),
|
|
)
|
|
.expect("unwrap should succeed");
|
|
let output_text = String::from_utf8(output).expect("text should decode");
|
|
assert!(output_text.contains("\"_v1internal_response_id\":\"resp_123\""));
|
|
assert!(output_text.contains("\"id\":\"call_get_weather_0\""));
|
|
}
|
|
|
|
#[test]
|
|
fn private_stream_normalizer_unwraps_antigravity_stream() {
|
|
let report_context = json!({
|
|
"has_envelope": true,
|
|
"provider_api_format": "gemini:generate_content",
|
|
"client_api_format": "gemini:generate_content",
|
|
"envelope_name": "antigravity:v1internal",
|
|
"mapped_model": "claude-sonnet-4-5",
|
|
});
|
|
let mut normalizer = maybe_build_provider_private_stream_normalizer(Some(&report_context))
|
|
.expect("normalizer should exist");
|
|
let output = normalizer
|
|
.push_chunk(
|
|
b"data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{\"city\":\"SF\"}}}],\"role\":\"model\"},\"index\":0}],\"modelVersion\":\"claude-sonnet-4-5\"},\"responseId\":\"resp_123\"}\n\n",
|
|
)
|
|
.expect("unwrap should succeed");
|
|
let output_text = String::from_utf8(output).expect("text should decode");
|
|
assert!(output_text.contains("\"_v1internal_response_id\":\"resp_123\""));
|
|
assert!(output_text.contains("\"id\":\"call_get_weather_0\""));
|
|
}
|
|
|
|
#[test]
|
|
fn detects_sse_error_events_without_explicit_type_field() {
|
|
let body = br#"event: error
|
|
data: {"message":"bad"}
|
|
|
|
"#;
|
|
assert!(stream_body_contains_error_event(body));
|
|
}
|
|
}
|