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

feat(gateway): route OpenAI image streams through chat bridge

Browse Source
This commit is contained in:
ZheFox
2026-05-17 23:09:50 +08:00
parent d6c8c14de7
commit 680b617b00
12 changed files with 1699 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
View File

@@ -814,10 +814,10 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
}
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:image", provider_api_format);
ai_local_execution_contract_for_formats("openai:chat", provider_api_format);
Ok(Some(LocalOpenAiChatCandidatePayloadParts {
client_api_format: "openai:image".to_string(),
client_api_format: "openai:chat".to_string(),
auth_header: prepared_candidate.auth_header,
auth_value: prepared_candidate.auth_value,
mapped_model: prepared_candidate.mapped_model,
@@ -827,7 +827,7 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
upstream_url,
execution_strategy,
conversion_mode,
report_kind: "openai_image_stream_success".to_string(),
report_kind: "openai_chat_stream_success".to_string(),
envelope_name: None,
transport: Arc::clone(transport),
request_redacted: false,

10
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/support.rs
View File

@@ -386,7 +386,7 @@ pub(crate) async fn build_local_openai_chat_image_candidate_attempt_source<'a>(
Ok(build_local_execution_candidate_attempt_source_with_serving(
planner_state,
trace_id,
"openai:image",
"openai:chat",
Some(&input.requested_model),
Some(&input.auth_snapshot),
input.client_session_affinity.as_ref(),
@@ -400,12 +400,12 @@ pub(crate) async fn build_local_openai_chat_image_candidate_attempt_source<'a>(
|eligible| {
let provider_api_format = eligible.provider_api_format.clone();
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:image", &provider_api_format);
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
Some(build_local_execution_candidate_contract_metadata(
LocalExecutionCandidateMetadataParts {
eligible,
provider_api_format: provider_api_format.as_str(),
client_api_format: "openai:image",
client_api_format: "openai:chat",
extra_fields: serde_json::Map::new(),
},
execution_strategy,
@@ -426,13 +426,13 @@ pub(crate) async fn build_local_openai_chat_image_candidate_attempt_source<'a>(
.to_ascii_lowercase()
});
let (execution_strategy, conversion_mode) =
ai_local_execution_contract_for_formats("openai:image", &provider_api_format);
ai_local_execution_contract_for_formats("openai:chat", &provider_api_format);
skipped_candidate.extra_data = Some(
build_local_execution_candidate_contract_metadata_for_candidate(
&skipped_candidate.candidate,
skipped_candidate.transport_ref(),
provider_api_format.as_str(),
"openai:image",
"openai:chat",
serde_json::Map::new(),
execution_strategy,
conversion_mode,

8
apps/aether-gateway/src/tests/ai_execute/stream/image.rs
View File

@@ -1104,8 +1104,10 @@ async fn gateway_routes_openai_chat_stream_image_intent_to_openai_image_plan_wit
StatusCode::OK,
"{response_text}\n{stored_candidates:#?}"
);
assert!(response_text.contains("image_generation.completed"));
assert!(response_text.contains("aGVsbG8="));
assert!(response_text.contains("\"object\":\"chat.completion.chunk\""));
assert!(response_text.contains("![generated image](data:image/png;base64,aGVsbG8=)"));
assert!(response_text.contains("data: [DONE]"));
assert!(!response_text.contains("image_generation.completed"));
let seen_plan = seen_execution_plan
.lock()
@@ -1113,7 +1115,7 @@ async fn gateway_routes_openai_chat_stream_image_intent_to_openai_image_plan_wit
.clone()
.expect("execution plan should be captured");
assert_eq!(seen_plan.trace_id, "trace-chat-stream-image-bridge-123");
assert_eq!(seen_plan.client_api_format, "openai:image");
assert_eq!(seen_plan.client_api_format, "openai:chat");
assert_eq!(seen_plan.provider_api_format, "openai:image");
assert_eq!(seen_plan.url, "https://images.example.com/v1/responses");
assert!(seen_plan.plan_stream);

725
crates/aether-ai-formats/src/formats/openai/image/stream.rs
View File

@@ -1,9 +1,15 @@
use std::collections::BTreeSet;
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
use base64::Engine as _;
use serde_json::Value;
use serde_json::{Map, Value};
use crate::contracts::OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND;
use crate::formats::openai::responses::codex::CODEX_OPENAI_IMAGE_DEFAULT_OUTPUT_FORMAT;
use crate::formats::shared::sse::encode_json_sse;
use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
use crate::formats::shared::stream_core::common::{
build_openai_chat_chunk, build_openai_chat_finish_chunk, build_openai_chat_usage_chunk,
};
use crate::formats::shared::AiSurfaceFinalizeError;
#[derive(Default)]
@@ -20,6 +26,38 @@ struct OpenAiImageFrame {
b64_json: String,
}
#[derive(Default)]
pub struct OpenAiImageChatStreamState {
buffered: Vec<u8>,
response_id: Option<String>,
model: Option<String>,
latest_image: Option<OpenAiImageChatFrame>,
emitted_image_count: u64,
emitted_image_keys: BTreeSet<String>,
started: bool,
finished: bool,
emitted_failure: bool,
}
#[derive(Clone)]
struct OpenAiImageChatFrame {
b64_json: String,
output_format: Option<String>,
}
#[derive(Default)]
pub struct OpenAiImageStreamTerminalState {
event_name: Option<String>,
data_lines: Vec<String>,
response_id: Option<String>,
model: Option<String>,
image_count: u64,
image_keys: BTreeSet<String>,
usage: Option<Value>,
observed_finish: bool,
parser_error: Option<String>,
}
impl OpenAiImageStreamState {
pub fn push_chunk(
&mut self,
@@ -229,6 +267,657 @@ impl OpenAiImageStreamState {
}
}
impl OpenAiImageChatStreamState {
pub fn push_chunk(
&mut self,
report_context: &Value,
chunk: &[u8],
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.buffered.extend_from_slice(chunk);
let mut output = Vec::new();
while let Some(block_end) = find_sse_block_end(&self.buffered) {
let block = self.buffered.drain(..block_end).collect::<Vec<_>>();
output.extend(self.transform_block(report_context, &block)?);
drain_sse_separator(&mut self.buffered);
}
Ok(output)
}
pub fn finish(&mut self, report_context: &Value) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let mut output = if self.buffered.is_empty() {
Vec::new()
} else {
let block = std::mem::take(&mut self.buffered);
self.transform_block(report_context, &block)?
};
if !self.finished && !self.emitted_failure && self.latest_image.is_some() {
output.extend(self.emit_final(report_context, None)?);
}
Ok(output)
}
fn transform_block(
&mut self,
report_context: &Value,
block: &[u8],
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
let text = std::str::from_utf8(block)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
let mut event_name = None::<String>;
let mut data_lines = Vec::new();
for raw_line in text.lines() {
let line = raw_line.trim_end_matches('\r');
if let Some(value) = line.strip_prefix("event:") {
event_name = Some(value.trim().to_string());
} else if let Some(value) = line.strip_prefix("data:") {
data_lines.push(value.trim().to_string());
}
}
let data = data_lines.join("\n");
if data.is_empty() || data == "[DONE]" {
return Ok(Vec::new());
}
let event: Value = serde_json::from_str(&data)?;
let event_type = event
.get("type")
.and_then(Value::as_str)
.or(event_name.as_deref())
.unwrap_or_default();
match event_type {
"error" | "response.failed" | "image_generation.failed" | "image_edit.failed" => {
self.handle_failed(report_context, &event)
}
"response.image_generation_call.partial_image" => {
self.emit_empty_progress_chunk(report_context)
}
"response.output_item.done" => self.handle_output_item_done(report_context, &event),
"response.completed" | "response.done" => self.handle_completed(report_context, &event),
"image_generation.completed" | "image_edit.completed" => {
self.handle_image_completed(report_context, &event)
}
_ => Ok(Vec::new()),
}
}
fn handle_output_item_done(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.finished || self.emitted_failure {
return Ok(Vec::new());
}
let Some(item) = event.get("item").and_then(Value::as_object) else {
return Ok(Vec::new());
};
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return Ok(Vec::new());
}
if let Some(result) = item.get("result").and_then(Value::as_str).map(str::trim) {
if !result.is_empty() {
let key = image_chat_output_key(item, result);
if self.emitted_image_keys.insert(key) {
self.latest_image = Some(OpenAiImageChatFrame {
b64_json: result.to_string(),
output_format: item
.get("output_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
});
self.emitted_image_count = self.emitted_image_count.saturating_add(1);
}
}
}
self.ensure_started(report_context)
}
fn handle_completed(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.finished || self.emitted_failure {
return Ok(Vec::new());
}
if let Some(response) = event.get("response") {
self.update_identity_from_response(response);
if self.latest_image.is_none() {
if let Some(frame) = completed_response_image_chat_frame(response) {
self.latest_image = Some(frame);
self.emitted_image_count = self.emitted_image_count.saturating_add(1);
}
}
}
let usage = event
.get("response")
.and_then(Value::as_object)
.and_then(|response| {
response
.get("tool_usage")
.and_then(|value| value.get("image_gen"))
.cloned()
.or_else(|| response.get("usage").cloned())
});
self.emit_final(report_context, usage.as_ref())
}
fn handle_image_completed(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.finished || self.emitted_failure {
return Ok(Vec::new());
}
if let Some(result) = event
.get("b64_json")
.or_else(|| event.get("result"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
self.latest_image = Some(OpenAiImageChatFrame {
b64_json: result.to_string(),
output_format: event
.get("output_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
});
self.emitted_image_count = self.emitted_image_count.max(1);
}
self.emit_final(report_context, event.get("usage"))
}
fn handle_failed(
&mut self,
_report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.emitted_failure {
return Ok(Vec::new());
}
self.emitted_failure = true;
self.finished = true;
let mut output = encode_json_sse(
None,
&serde_json::json!({
"error": image_failure_error(event),
}),
)?;
output.extend(encode_done_sse());
Ok(output)
}
fn ensure_started(
&mut self,
report_context: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.started {
return Ok(Vec::new());
}
self.emit_empty_progress_chunk(report_context)
}
fn emit_empty_progress_chunk(
&mut self,
report_context: &Value,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
self.started = true;
let (response_id, model) = self.identity(report_context);
encode_json_sse(
None,
&build_openai_chat_chunk(&response_id, &model, String::new(), None, None),
)
}
fn emit_final(
&mut self,
report_context: &Value,
usage: Option<&Value>,
) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
if self.finished || self.emitted_failure {
return Ok(Vec::new());
}
let Some(latest_image) = self.latest_image.clone() else {
return self.ensure_started(report_context);
};
let mut output = self.ensure_started(report_context)?;
let (response_id, model) = self.identity(report_context);
output.extend(encode_json_sse(
None,
&build_openai_chat_chunk(
&response_id,
&model,
image_chat_markdown(&latest_image),
None,
None,
),
)?);
output.extend(encode_json_sse(
None,
&build_openai_chat_finish_chunk(&response_id, &model, Some("stop")),
)?);
if let Some((input_tokens, output_tokens, total_tokens, reasoning_tokens)) =
openai_image_chat_usage_counts(usage)
{
output.extend(encode_json_sse(
None,
&build_openai_chat_usage_chunk(
&response_id,
&model,
input_tokens,
output_tokens,
total_tokens,
reasoning_tokens,
),
)?);
}
output.extend(encode_done_sse());
self.finished = true;
Ok(output)
}
fn update_identity_from_response(&mut self, response: &Value) {
if let Some(id) = response
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
self.response_id = Some(id.replace("resp", "chatcmpl"));
}
if let Some(model) = response
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
self.model = Some(model.to_string());
}
}
fn identity(&self, report_context: &Value) -> (String, String) {
let response_id = self.response_id.clone().unwrap_or_else(|| {
report_context
.get("request_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| format!("chatcmpl-image-{value}"))
.unwrap_or_else(|| "chatcmpl-image".to_string())
});
let model = self
.model
.clone()
.or_else(|| {
report_context
.get("mapped_model")
.or_else(|| report_context.get("model"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
.unwrap_or_else(|| "gpt-image".to_string());
(response_id, model)
}
}
impl OpenAiImageStreamTerminalState {
pub fn push_line(
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<Option<ExecutionStreamTerminalSummary>, AiSurfaceFinalizeError> {
let text = std::str::from_utf8(&line)
.map_err(|err| AiSurfaceFinalizeError::new(err.to_string()))?;
let trimmed = text.trim_matches('\r').trim_matches('\n');
if trimmed.is_empty() {
self.flush_event(report_context)?;
return Ok(self.latest_summary(report_context));
}
if let Some(value) = trimmed.strip_prefix("event:") {
self.event_name = Some(value.trim().to_string());
} else if let Some(value) = trimmed.strip_prefix("data:") {
self.data_lines.push(value.trim().to_string());
}
Ok(self.latest_summary(report_context))
}
pub fn finish(
&mut self,
report_context: &Value,
) -> Result<Option<ExecutionStreamTerminalSummary>, AiSurfaceFinalizeError> {
self.flush_event(report_context)?;
if self.image_count > 0 && !self.observed_finish {
self.observed_finish = true;
}
Ok(self.latest_summary(report_context))
}
fn flush_event(&mut self, report_context: &Value) -> Result<(), AiSurfaceFinalizeError> {
if self.data_lines.is_empty() {
self.event_name = None;
return Ok(());
}
let data = std::mem::take(&mut self.data_lines).join("\n");
let event_name = self.event_name.take();
if data.is_empty() || data == "[DONE]" {
return Ok(());
}
let event = match serde_json::from_str::<Value>(&data) {
Ok(event) => event,
Err(err) => {
self.parser_error.get_or_insert_with(|| err.to_string());
return Ok(());
}
};
let event_type = event
.get("type")
.and_then(Value::as_str)
.or(event_name.as_deref())
.unwrap_or_default();
match event_type {
"response.output_item.done" => self.observe_output_item_done(&event),
"response.completed" | "response.done" => self.observe_completed(&event),
"image_generation.completed" | "image_edit.completed" => {
self.observe_image_completed(&event)
}
"error" | "response.failed" | "image_generation.failed" | "image_edit.failed" => {
self.parser_error
.get_or_insert_with(|| image_failure_error(&event).to_string());
self.observed_finish = true;
}
_ => {}
}
if self.model.is_none() {
self.model = image_bridge_model(Some(report_context));
}
Ok(())
}
fn observe_output_item_done(&mut self, event: &Value) {
let Some(item) = event.get("item").and_then(Value::as_object) else {
return;
};
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return;
}
let Some(result) = item
.get("result")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return;
};
let key = image_chat_output_key(item, result);
if self.image_keys.insert(key) {
self.image_count = self.image_count.saturating_add(1);
}
}
fn observe_completed(&mut self, event: &Value) {
self.observed_finish = true;
let Some(response) = event.get("response") else {
return;
};
self.update_identity_from_response(response);
if self.image_count == 0 {
self.image_count = completed_response_image_count(response);
}
self.usage = response
.get("tool_usage")
.and_then(|value| value.get("image_gen"))
.cloned()
.or_else(|| response.get("usage").cloned())
.or_else(|| self.usage.clone());
}
fn observe_image_completed(&mut self, event: &Value) {
self.observed_finish = true;
if self.image_count == 0 {
if event
.get("b64_json")
.or_else(|| event.get("result"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
{
self.image_count = 1;
}
}
self.usage = event.get("usage").cloned().or_else(|| self.usage.clone());
}
fn update_identity_from_response(&mut self, response: &Value) {
if let Some(id) = response
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
self.response_id = Some(id.to_string());
}
if let Some(model) = response
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
self.model = Some(model.to_string());
}
}
fn latest_summary(&self, report_context: &Value) -> Option<ExecutionStreamTerminalSummary> {
if self.image_count == 0
&& self.usage.is_none()
&& self.response_id.is_none()
&& self.model.is_none()
&& self.parser_error.is_none()
{
return None;
}
Some(ExecutionStreamTerminalSummary {
standardized_usage: openai_image_stream_standardized_usage(
self.usage.as_ref(),
Some(report_context),
self.image_count,
),
finish_reason: self.observed_finish.then(|| "stop".to_string()),
response_id: self.response_id.clone(),
model: self
.model
.clone()
.or_else(|| image_bridge_model(Some(report_context))),
observed_finish: self.observed_finish,
unknown_event_count: 0,
parser_error: self.parser_error.clone(),
})
}
}
fn completed_response_image_chat_frame(response: &Value) -> Option<OpenAiImageChatFrame> {
response
.get("output")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
.find_map(|item| {
let result = item.get("result").and_then(Value::as_str)?.trim();
if result.is_empty() {
return None;
}
Some(OpenAiImageChatFrame {
b64_json: result.to_string(),
output_format: item
.get("output_format")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
})
})
}
fn completed_response_image_count(response: &Value) -> u64 {
response
.get("output")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
.filter(|item| {
item.get("result")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
})
.count() as u64
}
fn image_chat_output_key(item: &Map<String, Value>, result: &str) -> String {
item.get("id")
.or_else(|| item.get("call_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| result.to_string())
}
fn openai_image_stream_standardized_usage(
usage: Option<&Value>,
report_context: Option<&Value>,
image_count: u64,
) -> Option<StandardizedUsage> {
let mut standardized_usage = usage
.and_then(openai_image_usage_to_standardized_usage)
.unwrap_or_else(StandardizedUsage::new);
if image_count > 0 {
standardized_usage.request_count = i64::try_from(image_count).unwrap_or(i64::MAX);
standardized_usage
.dimensions
.insert("image_count".to_string(), serde_json::json!(image_count));
}
if let Some(output_format) = image_request_output_format(report_context) {
standardized_usage.dimensions.insert(
"image_output_format".to_string(),
serde_json::json!(output_format),
);
}
if let Some(size) = image_request_size(report_context) {
standardized_usage
.dimensions
.insert("image_size".to_string(), serde_json::json!(size));
}
(standardized_usage.signal_score() > 0).then_some(standardized_usage)
}
fn openai_image_usage_to_standardized_usage(value: &Value) -> Option<StandardizedUsage> {
let usage = value.as_object()?;
let mut input_tokens = usage
.get("input_tokens")
.or_else(|| usage.get("prompt_tokens"))
.and_then(Value::as_i64)
.unwrap_or(0);
let output_tokens = usage
.get("output_tokens")
.or_else(|| usage.get("completion_tokens"))
.and_then(Value::as_i64)
.unwrap_or(0);
let cache_creation_tokens = usage
.get("cache_creation_input_tokens")
.and_then(Value::as_i64)
.or_else(|| {
usage
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_creation_tokens"))
.and_then(Value::as_i64)
})
.unwrap_or(0);
let cache_read_tokens = usage
.get("cache_read_input_tokens")
.and_then(Value::as_i64)
.or_else(|| {
usage
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_tokens"))
.and_then(Value::as_i64)
})
.unwrap_or(0);
let total_tokens = usage.get("total_tokens").and_then(Value::as_i64).unwrap_or(
input_tokens
.saturating_add(output_tokens)
.saturating_add(cache_creation_tokens)
.saturating_add(cache_read_tokens),
);
if input_tokens == 0 && total_tokens > output_tokens {
input_tokens = total_tokens.saturating_sub(output_tokens);
}
let mut standardized_usage = StandardizedUsage::new();
standardized_usage.input_tokens = input_tokens;
standardized_usage.output_tokens = output_tokens;
standardized_usage.cache_creation_tokens = cache_creation_tokens;
standardized_usage.cache_read_tokens = cache_read_tokens;
standardized_usage
.dimensions
.insert("total_tokens".to_string(), serde_json::json!(total_tokens));
Some(standardized_usage.normalize_cache_creation_breakdown())
}
fn image_chat_markdown(frame: &OpenAiImageChatFrame) -> String {
let mime_type = match frame
.output_format
.as_deref()
.unwrap_or("png")
.trim()
.to_ascii_lowercase()
.as_str()
{
"jpg" | "jpeg" => "image/jpeg".to_string(),
"webp" => "image/webp".to_string(),
"png" => "image/png".to_string(),
value if !value.is_empty() => format!("image/{value}"),
_ => "image/png".to_string(),
};
format!(
"![generated image](data:{mime_type};base64,{})",
frame.b64_json
)
}
fn openai_image_chat_usage_counts(usage: Option<&Value>) -> Option<(u64, u64, u64, u64)> {
let usage = usage.and_then(Value::as_object)?;
let mut input_tokens = usage
.get("input_tokens")
.or_else(|| usage.get("prompt_tokens"))
.and_then(Value::as_u64)
.unwrap_or(0);
let output_tokens = usage
.get("output_tokens")
.or_else(|| usage.get("completion_tokens"))
.and_then(Value::as_u64)
.unwrap_or(0);
let total_tokens = usage
.get("total_tokens")
.and_then(Value::as_u64)
.unwrap_or(input_tokens.saturating_add(output_tokens));
if input_tokens == 0 && total_tokens > output_tokens {
input_tokens = total_tokens.saturating_sub(output_tokens);
}
(total_tokens > 0).then_some((input_tokens, output_tokens, total_tokens, 0))
}
fn image_failure_error(event: &Value) -> Value {
let mut error = event
.get("error")
@@ -340,6 +1029,38 @@ fn image_request_operation(report_context: &Value) -> Option<&str> {
.filter(|value| !value.is_empty())
}
fn image_request_output_format(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|value| value.get("image_request"))
.and_then(|value| value.get("output_format"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn image_request_size(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|value| value.get("image_request"))
.and_then(|value| value.get("size"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn image_bridge_model(report_context: Option<&Value>) -> Option<String> {
report_context.and_then(|context| {
context
.get("mapped_model")
.or_else(|| context.get("model"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn find_sse_block_end(buffer: &[u8]) -> Option<usize> {
buffer
.windows(2)

125
crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs
View File

@@ -8,6 +8,7 @@ use crate::formats::openai::chat::stream::{
OpenAIChatClientEmitter, OpenAIChatProviderState, OpenAIResponsesClientEmitter,
OpenAIResponsesProviderState,
};
use crate::formats::openai::image::stream::OpenAiImageStreamTerminalState;
use crate::formats::shared::error_body::{
build_core_error_body_for_client_format, LocalCoreSyncErrorKind,
};
@@ -97,7 +98,7 @@ impl StreamingStandardFormatMatrix {
#[derive(Default)]
pub struct StreamingStandardTerminalObserver {
provider: Option<ProviderStreamParser>,
provider: Option<TerminalStreamParser>,
latest_summary: Option<ExecutionStreamTerminalSummary>,
}
@@ -111,8 +112,17 @@ impl StreamingStandardTerminalObserver {
let Some(provider) = self.provider.as_mut() else {
return Ok(());
};
let frames = provider.push_line(report_context, line)?;
self.observe_frames(frames);
match provider {
TerminalStreamParser::Standard(provider) => {
let frames = provider.push_line(report_context, line)?;
self.observe_frames(frames);
}
TerminalStreamParser::OpenAIImage(provider) => {
if let Some(summary) = provider.push_line(report_context, line)? {
self.latest_summary = Some(summary);
}
}
}
Ok(())
}
@@ -124,8 +134,17 @@ impl StreamingStandardTerminalObserver {
let Some(provider) = self.provider.as_mut() else {
return Ok(self.latest_summary.clone());
};
let frames = provider.finish(report_context)?;
self.observe_frames(frames);
match provider {
TerminalStreamParser::Standard(provider) => {
let frames = provider.finish(report_context)?;
self.observe_frames(frames);
}
TerminalStreamParser::OpenAIImage(provider) => {
if let Some(summary) = provider.finish(report_context)? {
self.latest_summary = Some(summary);
}
}
}
Ok(self.latest_summary.clone())
}
@@ -153,7 +172,7 @@ impl StreamingStandardTerminalObserver {
return;
}
let provider_api_format = provider_api_format_for_context(report_context);
self.provider = ProviderStreamParser::for_api_format(provider_api_format.as_str());
self.provider = TerminalStreamParser::for_api_format(provider_api_format.as_str());
}
fn observe_frames(&mut self, frames: Vec<CanonicalStreamFrame>) {
@@ -194,6 +213,23 @@ impl StreamingStandardTerminalObserver {
}
}
enum TerminalStreamParser {
Standard(ProviderStreamParser),
OpenAIImage(OpenAiImageStreamTerminalState),
}
impl TerminalStreamParser {
fn for_api_format(provider_api_format: &str) -> Option<Self> {
if provider_api_format
.trim()
.eq_ignore_ascii_case("openai:image")
{
return Some(Self::OpenAIImage(OpenAiImageStreamTerminalState::default()));
}
ProviderStreamParser::for_api_format(provider_api_format).map(Self::Standard)
}
}
enum ProviderStreamParser {
OpenAIChat(OpenAIChatProviderState),
OpenAIResponses(OpenAIResponsesProviderState),
@@ -936,4 +972,81 @@ mod tests {
assert_eq!(summary.unknown_event_count, 1);
assert!(!summary.observed_finish);
}
#[test]
fn terminal_observer_tracks_openai_image_stream_usage() {
let mut report_context = report_context("openai:image", "openai:chat");
report_context["image_request"] = json!({
"size": "1024x1024",
"output_format": "png",
});
let mut observer = StreamingStandardTerminalObserver::default();
observer
.push_line(
&report_context,
data_line(json!({
"type": "response.output_item.done",
"output_index": 0,
"item": {
"id": "ig_123",
"type": "image_generation_call",
"result": "aGVsbG8=",
},
})),
)
.expect("image output item should parse");
observer
.push_line(&report_context, b"\n".to_vec())
.expect("image output event should flush");
observer
.push_line(
&report_context,
data_line(json!({
"type": "response.completed",
"response": {
"id": "resp_image_123",
"model": "gpt-image-2",
"output": [],
"tool_usage": {
"image_gen": {
"input_tokens": 40,
"output_tokens": 60,
"total_tokens": 100,
},
},
},
})),
)
.expect("image completed should parse");
observer
.push_line(&report_context, b"\n".to_vec())
.expect("image completed event should flush");
let summary = observer
.finish(&report_context)
.expect("image summary should finish")
.expect("summary should exist");
let usage = summary
.standardized_usage
.expect("standardized usage should exist");
assert_eq!(summary.response_id.as_deref(), Some("resp_image_123"));
assert_eq!(summary.model.as_deref(), Some("gpt-image-2"));
assert_eq!(summary.finish_reason.as_deref(), Some("stop"));
assert!(summary.observed_finish);
assert_eq!(usage.input_tokens, 40);
assert_eq!(usage.output_tokens, 60);
assert_eq!(usage.request_count, 1);
assert_eq!(usage.dimensions.get("image_count"), Some(&json!(1)));
assert_eq!(usage.dimensions.get("total_tokens"), Some(&json!(100)));
assert_eq!(
usage.dimensions.get("image_size"),
Some(&json!("1024x1024"))
);
assert_eq!(
usage.dimensions.get("image_output_format"),
Some(&json!("png"))
);
}
}

82
crates/aether-ai-formats/src/formats/shared/stream_rewrite.rs
View File

@@ -1,6 +1,6 @@
use serde_json::Value;
use crate::formats::openai::image::stream::OpenAiImageStreamState;
use crate::formats::openai::image::stream::{OpenAiImageChatStreamState, OpenAiImageStreamState};
use crate::formats::shared::model_directives::model_directive_display_model_from_report_context;
use crate::formats::shared::stream_core::StreamingStandardFormatMatrix;
use crate::formats::shared::AiSurfaceFinalizeError;
@@ -15,6 +15,7 @@ pub enum FinalizeStreamRewriteMode {
EnvelopeUnwrap,
ModelDirectiveDisplay,
OpenAiImage,
OpenAiImageToOpenAiChat,
Standard,
KiroToClaudeCli,
KiroToClaudeCliThenStandard,
@@ -57,6 +58,14 @@ pub fn resolve_finalize_stream_rewrite_mode(
.then_some(FinalizeStreamRewriteMode::KiroToClaudeCliThenStandard);
}
if provider_api_format == "openai:image" && client_api_format == "openai:chat" {
return Some(FinalizeStreamRewriteMode::OpenAiImageToOpenAiChat);
}
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
return Some(FinalizeStreamRewriteMode::OpenAiImage);
}
if needs_conversion {
// CPA strategy: when provider and client share the same wire format
// (exact match or same family), pass through the stream verbatim.
@@ -73,10 +82,6 @@ pub fn resolve_finalize_stream_rewrite_mode(
.then_some(FinalizeStreamRewriteMode::Standard);
}
if provider_api_format == "openai:image" && client_api_format == "openai:image" {
return Some(FinalizeStreamRewriteMode::OpenAiImage);
}
if envelope_name.eq_ignore_ascii_case(KIRO_ENVELOPE_NAME) {
return (provider_api_format == "claude:messages"
&& client_api_format == "claude:messages")
@@ -106,6 +111,7 @@ enum AiSurfaceStreamRewriteState {
EnvelopeUnwrap,
ModelDirectiveDisplay,
OpenAiImage(Box<OpenAiImageStreamState>),
OpenAiImageToOpenAiChat(Box<OpenAiImageChatStreamState>),
Standard(Box<StreamingStandardFormatMatrix>),
KiroToClaudeCli(Box<KiroToClaudeCliStreamState>),
KiroToClaudeCliThenStandard {
@@ -132,6 +138,11 @@ pub fn maybe_build_ai_surface_stream_rewriter<'a>(
FinalizeStreamRewriteMode::OpenAiImage => {
AiSurfaceStreamRewriteState::OpenAiImage(Box::<OpenAiImageStreamState>::default())
}
FinalizeStreamRewriteMode::OpenAiImageToOpenAiChat => {
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(
Box::<OpenAiImageChatStreamState>::default(),
)
}
FinalizeStreamRewriteMode::Standard => {
AiSurfaceStreamRewriteState::Standard(Box::<StreamingStandardFormatMatrix>::default())
}
@@ -159,6 +170,9 @@ impl AiSurfaceStreamRewriter<'_> {
AiSurfaceStreamRewriteState::OpenAiImage(state) => {
state.push_chunk(self.report_context, chunk)
}
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => {
state.push_chunk(self.report_context, chunk)
}
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
state.push_chunk(self.report_context, chunk)
}
@@ -183,6 +197,9 @@ impl AiSurfaceStreamRewriter<'_> {
pub fn finish(&mut self) -> Result<Vec<u8>, AiSurfaceFinalizeError> {
match &mut self.state {
AiSurfaceStreamRewriteState::OpenAiImage(state) => state.finish(self.report_context),
AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(state) => {
state.finish(self.report_context)
}
AiSurfaceStreamRewriteState::KiroToClaudeCli(state) => {
state.finish(self.report_context)
}
@@ -228,6 +245,7 @@ impl AiSurfaceStreamRewriter<'_> {
transform_standard_line(state, self.report_context, line)
}
AiSurfaceStreamRewriteState::OpenAiImage(_)
| AiSurfaceStreamRewriteState::OpenAiImageToOpenAiChat(_)
| AiSurfaceStreamRewriteState::KiroToClaudeCli(_)
| AiSurfaceStreamRewriteState::KiroToClaudeCliThenStandard { .. } => Ok(Vec::new()),
}
@@ -686,4 +704,58 @@ data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"thinki
Some(FinalizeStreamRewriteMode::OpenAiImage)
);
}
#[test]
fn rewrites_openai_image_stream_to_openai_chat_final_chunk() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:chat",
"mapped_model": "gpt-image-2",
"request_id": "trace-image-chat-stream",
"needs_conversion": false,
});
assert_eq!(
resolve_finalize_stream_rewrite_mode(&report_context),
Some(FinalizeStreamRewriteMode::OpenAiImageToOpenAiChat)
);
let mut rewriter = maybe_build_ai_surface_stream_rewriter(Some(&report_context))
.expect("image to chat stream rewriter should exist");
let progress = rewriter
.push_chunk(
br#"event: response.image_generation_call.partial_image
data: {"type":"response.image_generation_call.partial_image","partial_image_b64":"cGFydGlhbA=="}
"#,
)
.expect("partial image should rewrite as progress");
let progress_text = String::from_utf8(progress).expect("progress output should be utf8");
assert!(progress_text.contains("\"object\":\"chat.completion.chunk\""));
assert!(!progress_text.contains("cGFydGlhbA=="));
let output_item = rewriter
.push_chunk(
br#"event: response.output_item.done
data: {"type":"response.output_item.done","item":{"type":"image_generation_call","id":"ig_1","result":"aGVsbG8=","output_format":"png"}}
"#,
)
.expect("output item should rewrite");
let output_item_text = String::from_utf8(output_item).expect("output item should be utf8");
assert!(output_item_text.is_empty());
let final_output = rewriter
.push_chunk(
br#"event: response.completed
data: {"type":"response.completed","response":{"id":"resp_123","model":"gpt-image-2","tool_usage":{"image_gen":{"total_tokens":0}},"output":[]}}
"#,
)
.expect("completed event should rewrite");
let final_text = String::from_utf8(final_output).expect("final output should be utf8");
assert!(final_text.contains("\"object\":\"chat.completion.chunk\""));
assert!(final_text.contains("![generated image](data:image/png;base64,aGVsbG8=)"));
assert!(final_text.contains("data: [DONE]"));
assert!(!final_text.contains("image_generation.completed"));
}
}

612
crates/aether-ai-formats/src/formats/shared/sync_to_stream.rs
View File

@@ -1,16 +1,21 @@
use std::borrow::Cow;
use aether_ai_formats::formats::conversion::response::{
convert_claude_response_to_openai_responses, convert_gemini_response_to_openai_responses,
convert_openai_chat_response_to_openai_responses,
};
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
use serde_json::{json, Value};
use serde_json::{json, Map, Value};
use crate::formats::claude::messages::stream::ClaudeClientEmitter;
use crate::formats::gemini::generate_content::stream::GeminiClientEmitter;
use crate::formats::openai::chat::stream::{
OpenAIChatClientEmitter, OpenAIResponsesClientEmitter, OpenAIResponsesProviderState,
};
use crate::formats::shared::sse::encode_json_sse;
use crate::formats::shared::sse::{encode_done_sse, encode_json_sse};
use crate::formats::shared::stream_core::common::{
build_openai_chat_chunk, build_openai_chat_finish_chunk, build_openai_chat_usage_chunk,
};
use crate::formats::shared::stream_core::CanonicalStreamFrame;
use crate::formats::shared::AiSurfaceFinalizeError;
@@ -27,12 +32,25 @@ pub fn maybe_bridge_standard_sync_json_to_stream(
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let provider_api_format = normalize_api_format(provider_api_format);
let client_api_format = normalize_api_format(client_api_format);
if client_api_format == "openai:image"
&& matches!(
provider_api_format.as_str(),
"openai:image" | "gemini:generate_content"
)
{
if provider_api_format == "openai:image" {
return match client_api_format.as_str() {
"openai:image" => {
maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context)
}
"openai:chat" => maybe_bridge_openai_image_sync_json_to_chat_stream(
provider_body_json,
report_context,
),
"openai:responses" | "openai:responses:compact" => {
maybe_bridge_openai_image_sync_json_to_responses_stream(
provider_body_json,
report_context,
)
}
_ => Ok(None),
};
}
if client_api_format == "openai:image" && provider_api_format == "gemini:generate_content" {
return maybe_bridge_openai_image_sync_json_to_stream(provider_body_json, report_context);
}
if !is_standard_api_format(provider_api_format.as_str())
@@ -72,49 +90,19 @@ fn maybe_bridge_openai_image_sync_json_to_stream(
provider_body_json: &Value,
report_context: Option<&Value>,
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let provider_api_format = report_context
.and_then(|value| value.get("provider_api_format"))
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or("openai:image");
let owned_response;
let provider_body_json = if provider_api_format == "gemini:generate_content" {
let Some(converted) =
crate::formats::shared::image_bridge::build_openai_image_response_from_gemini_response(
provider_body_json,
report_context,
)
else {
return Ok(None);
};
owned_response = converted;
&owned_response
} else if provider_body_json.get("output").is_some() && provider_body_json.get("data").is_none()
{
let Some(converted) = crate::formats::shared::image_bridge::build_openai_image_response_from_response_stream_sync_body(
provider_body_json,
report_context,
) else {
return Ok(None);
};
owned_response = converted;
&owned_response
} else {
provider_body_json
};
let Some(response) = provider_body_json.as_object() else {
return Ok(None);
};
let Some(image) = response
.get("data")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.find_map(extract_openai_image_sync_b64_json)
let Some(provider_body_json) =
normalize_openai_image_sync_response(provider_body_json, report_context)?
else {
return Ok(None);
};
let Some(response) = provider_body_json.as_ref().as_object() else {
return Ok(None);
};
let outputs = collect_openai_image_outputs(response, report_context);
let Some(image) = outputs.iter().find_map(OpenAiImageOutput::b64_json) else {
return Ok(None);
};
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
let usage = response.get("usage").cloned().unwrap_or(Value::Null);
let event_name = openai_image_completed_event_name(report_context);
let sse_body = encode_json_sse(
@@ -128,27 +116,446 @@ fn maybe_bridge_openai_image_sync_json_to_stream(
Ok(Some(SyncToStreamBridgeOutcome {
sse_body,
terminal_summary: Some(ExecutionStreamTerminalSummary {
standardized_usage: response
.get("usage")
.and_then(standardized_usage_from_openai_usage),
finish_reason: Some("stop".to_string()),
response_id: response
.get("id")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
model: response
.get("model")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| image_bridge_model(report_context)),
observed_finish: true,
unknown_event_count: 0,
parser_error: None,
}),
terminal_summary: Some(openai_image_terminal_summary(
response,
report_context,
image_count,
)),
}))
}
fn maybe_bridge_openai_image_sync_json_to_chat_stream(
provider_body_json: &Value,
report_context: Option<&Value>,
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let Some(provider_body_json) =
normalize_openai_image_sync_response(provider_body_json, report_context)?
else {
return Ok(None);
};
let Some(response) = provider_body_json.as_ref().as_object() else {
return Ok(None);
};
let outputs = collect_openai_image_outputs(response, report_context);
if outputs.is_empty() {
return Ok(None);
}
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
let summary = openai_image_terminal_summary(response, report_context, image_count);
let response_id = openai_image_bridge_response_id(response, report_context, "chatcmpl-image");
let model = openai_image_bridge_response_model(response, report_context);
let content = outputs
.iter()
.enumerate()
.map(|(index, output)| output.markdown(index))
.collect::<Vec<_>>()
.join("\n\n");
let mut sse_body = Vec::new();
sse_body.extend(encode_json_sse(
None,
&build_openai_chat_chunk(&response_id, &model, content, None, None),
)?);
sse_body.extend(encode_json_sse(
None,
&build_openai_chat_finish_chunk(&response_id, &model, Some("stop")),
)?);
if let Some((input_tokens, output_tokens, total_tokens, reasoning_tokens)) = summary
.standardized_usage
.as_ref()
.and_then(openai_chat_usage_counts)
{
sse_body.extend(encode_json_sse(
None,
&build_openai_chat_usage_chunk(
&response_id,
&model,
input_tokens,
output_tokens,
total_tokens,
reasoning_tokens,
),
)?);
}
sse_body.extend(encode_done_sse());
Ok(Some(SyncToStreamBridgeOutcome {
sse_body,
terminal_summary: Some(summary),
}))
}
fn maybe_bridge_openai_image_sync_json_to_responses_stream(
provider_body_json: &Value,
report_context: Option<&Value>,
) -> Result<Option<SyncToStreamBridgeOutcome>, AiSurfaceFinalizeError> {
let Some(provider_body_json) =
normalize_openai_image_sync_response(provider_body_json, report_context)?
else {
return Ok(None);
};
let Some(response) = provider_body_json.as_ref().as_object() else {
return Ok(None);
};
let outputs = collect_openai_image_outputs(response, report_context);
if outputs.is_empty() {
return Ok(None);
}
let response_id = openai_image_bridge_response_id(response, report_context, "resp-image");
let model = openai_image_bridge_response_model(response, report_context);
let mut response_output = Vec::new();
for (index, output) in outputs.iter().enumerate() {
response_output.push(output.responses_image_generation_item(&response_id, index));
}
let mut response_object = Map::new();
response_object.insert("id".to_string(), Value::String(response_id.clone()));
response_object.insert("object".to_string(), Value::String("response".to_string()));
response_object.insert("model".to_string(), Value::String(model));
response_object.insert("status".to_string(), Value::String("completed".to_string()));
response_object.insert("output".to_string(), Value::Array(response_output.clone()));
if let Some(created) = response.get("created").and_then(Value::as_i64) {
response_object.insert("created_at".to_string(), json!(created));
}
if let Some(usage) = response.get("usage").filter(|value| value.is_object()) {
response_object.insert("usage".to_string(), usage.clone());
}
let mut sse_body = Vec::new();
for (index, item) in response_output.iter().enumerate() {
sse_body.extend(encode_json_sse(
Some("response.output_item.done"),
&json!({
"type": "response.output_item.done",
"output_index": index,
"item": item,
}),
)?);
}
sse_body.extend(encode_json_sse(
Some("response.completed"),
&json!({
"type": "response.completed",
"response": Value::Object(response_object),
}),
)?);
let image_count = openai_image_response_image_count(response).max(outputs.len() as u64);
Ok(Some(SyncToStreamBridgeOutcome {
sse_body,
terminal_summary: Some(openai_image_terminal_summary(
response,
report_context,
image_count,
)),
}))
}
#[derive(Clone, Debug)]
struct OpenAiImageOutput {
b64_json: Option<String>,
url: Option<String>,
mime_type: String,
output_format: Option<String>,
revised_prompt: Option<String>,
}
impl OpenAiImageOutput {
fn b64_json(&self) -> Option<String> {
self.b64_json
.clone()
.or_else(|| self.url.as_deref().and_then(extract_base64_from_data_url))
}
fn source_url(&self) -> Option<String> {
self.url.clone().or_else(|| {
self.b64_json
.as_ref()
.map(|value| format!("data:{};base64,{value}", self.mime_type))
})
}
fn markdown(&self, index: usize) -> String {
let alt = if index == 0 {
"generated image".to_string()
} else {
format!("generated image {}", index + 1)
};
match self.source_url() {
Some(url) => format!("![{alt}]({url})"),
None => String::new(),
}
}
fn responses_image_generation_item(&self, response_id: &str, index: usize) -> Value {
let mut item = Map::new();
item.insert(
"id".to_string(),
Value::String(format!("{response_id}_img_{index}")),
);
item.insert(
"type".to_string(),
Value::String("image_generation_call".to_string()),
);
item.insert("status".to_string(), Value::String("completed".to_string()));
if let Some(result) = self.b64_json().or_else(|| self.url.clone()) {
item.insert("result".to_string(), Value::String(result));
}
if let Some(output_format) = self.output_format.as_ref() {
item.insert(
"output_format".to_string(),
Value::String(output_format.clone()),
);
}
if let Some(revised_prompt) = self.revised_prompt.as_ref() {
item.insert(
"revised_prompt".to_string(),
Value::String(revised_prompt.clone()),
);
}
Value::Object(item)
}
}
fn normalize_openai_image_sync_response<'a>(
provider_body_json: &'a Value,
report_context: Option<&Value>,
) -> Result<Option<Cow<'a, Value>>, AiSurfaceFinalizeError> {
let provider_api_format = report_context
.and_then(|value| value.get("provider_api_format"))
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or("openai:image");
if provider_api_format == "gemini:generate_content" {
let Some(converted) =
crate::formats::shared::image_bridge::build_openai_image_response_from_gemini_response(
provider_body_json,
report_context,
)
else {
return Ok(None);
};
return Ok(Some(Cow::Owned(converted)));
}
if provider_body_json.get("output").is_some() && provider_body_json.get("data").is_none() {
let Some(converted) = crate::formats::shared::image_bridge::build_openai_image_response_from_response_stream_sync_body(
provider_body_json,
report_context,
) else {
return Ok(None);
};
return Ok(Some(Cow::Owned(converted)));
}
Ok(Some(Cow::Borrowed(provider_body_json)))
}
fn collect_openai_image_outputs(
response: &Map<String, Value>,
report_context: Option<&Value>,
) -> Vec<OpenAiImageOutput> {
response
.get("data")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.filter_map(|item| openai_image_output_from_item(item, report_context))
.collect()
}
fn openai_image_output_from_item(
item: &Map<String, Value>,
report_context: Option<&Value>,
) -> Option<OpenAiImageOutput> {
let b64_json = extract_openai_image_sync_b64_json(item);
let url = item
.get("url")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
if b64_json.is_none() && url.is_none() {
return None;
}
let output_format = item
.get("output_format")
.or_else(|| item.get("format"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| image_request_output_format(report_context));
let mime_type = url
.as_deref()
.and_then(extract_mime_type_from_data_url)
.or_else(|| {
output_format
.as_deref()
.map(mime_type_from_image_output_format)
})
.unwrap_or_else(|| "image/png".to_string());
let revised_prompt = item
.get("revised_prompt")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
Some(OpenAiImageOutput {
b64_json,
url,
mime_type,
output_format,
revised_prompt,
})
}
fn openai_image_response_image_count(response: &Map<String, Value>) -> u64 {
response
.get("data")
.and_then(Value::as_array)
.map(|items| items.len() as u64)
.unwrap_or(0)
}
fn openai_image_terminal_summary(
response: &Map<String, Value>,
report_context: Option<&Value>,
image_count: u64,
) -> ExecutionStreamTerminalSummary {
ExecutionStreamTerminalSummary {
standardized_usage: openai_image_standardized_usage(
response.get("usage"),
report_context,
image_count,
),
finish_reason: Some("stop".to_string()),
response_id: response
.get("id")
.and_then(Value::as_str)
.map(ToOwned::to_owned),
model: response
.get("model")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| image_bridge_model(report_context)),
observed_finish: true,
unknown_event_count: 0,
parser_error: None,
}
}
fn openai_image_standardized_usage(
usage: Option<&Value>,
report_context: Option<&Value>,
image_count: u64,
) -> Option<StandardizedUsage> {
let mut standardized_usage = usage
.and_then(standardized_usage_from_openai_usage)
.unwrap_or_else(StandardizedUsage::new);
if image_count > 0 {
standardized_usage.request_count = i64::try_from(image_count).unwrap_or(i64::MAX);
standardized_usage
.dimensions
.insert("image_count".to_string(), json!(image_count));
}
if let Some(output_format) = image_request_output_format(report_context) {
standardized_usage
.dimensions
.insert("image_output_format".to_string(), json!(output_format));
}
if let Some(size) = image_request_size(report_context) {
standardized_usage
.dimensions
.insert("image_size".to_string(), json!(size));
}
(standardized_usage.signal_score() > 0).then_some(standardized_usage)
}
fn openai_chat_usage_counts(usage: &StandardizedUsage) -> Option<(u64, u64, u64, u64)> {
let input_tokens = usage.input_tokens.max(0) as u64;
let output_tokens = usage.output_tokens.max(0) as u64;
let reasoning_tokens = usage.reasoning_tokens.max(0) as u64;
let total_tokens = usage
.dimensions
.get("total_tokens")
.and_then(Value::as_u64)
.unwrap_or_else(|| {
input_tokens
.saturating_add(output_tokens)
.saturating_add(reasoning_tokens)
});
(total_tokens > 0).then_some((input_tokens, output_tokens, total_tokens, reasoning_tokens))
}
fn openai_image_bridge_response_id(
response: &Map<String, Value>,
report_context: Option<&Value>,
fallback_prefix: &str,
) -> String {
response
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
report_context
.and_then(|value| value.get("request_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| format!("{fallback_prefix}-{value}"))
})
.unwrap_or_else(|| fallback_prefix.to_string())
}
fn openai_image_bridge_response_model(
response: &Map<String, Value>,
report_context: Option<&Value>,
) -> String {
response
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| image_bridge_model(report_context))
.unwrap_or_else(|| "gpt-image".to_string())
}
fn image_request_output_format(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|value| value.get("image_request"))
.and_then(|value| value.get("output_format"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn image_request_size(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|value| value.get("image_request"))
.and_then(|value| value.get("size"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn mime_type_from_image_output_format(output_format: &str) -> String {
match output_format.trim().to_ascii_lowercase().as_str() {
"jpg" | "jpeg" => "image/jpeg".to_string(),
"webp" => "image/webp".to_string(),
"png" => "image/png".to_string(),
value if !value.is_empty() => format!("image/{value}"),
_ => "image/png".to_string(),
}
}
fn normalize_api_format(value: &str) -> String {
aether_ai_formats::normalize_api_format_alias(value)
}
@@ -186,6 +593,14 @@ fn extract_base64_from_data_url(value: &str) -> Option<String> {
(!payload.trim().is_empty()).then(|| payload.trim().to_string())
}
fn extract_mime_type_from_data_url(value: &str) -> Option<String> {
let trimmed = value.trim();
let (metadata, _) = trimmed.split_once(',')?;
let mime_type = metadata.strip_prefix("data:")?.strip_suffix(";base64")?;
let mime_type = mime_type.trim();
(!mime_type.is_empty()).then(|| mime_type.to_string())
}
fn openai_image_completed_event_name(report_context: Option<&Value>) -> &'static str {
if openai_image_request_operation(report_context) == Some("edit") {
"image_edit.completed"
@@ -565,6 +980,71 @@ mod tests {
.cloned(),
Some(json!(100))
);
assert_eq!(
summary
.standardized_usage
.as_ref()
.and_then(|usage| usage.dimensions.get("image_count"))
.cloned(),
Some(json!(1))
);
}
#[test]
fn bridges_openai_image_sync_json_to_openai_chat_sse() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:chat",
"mapped_model": "gpt-image-2",
"image_request": {
"operation": "generate",
"output_format": "png",
"size": "1024x1024"
}
});
let outcome = maybe_bridge_standard_sync_json_to_stream(
&json!({
"id": "img_123",
"created": 1776971267,
"model": "gpt-image-2",
"data": [
{"b64_json": "aGVsbG8="},
{"b64_json": "d29ybGQ="}
],
"usage": {
"total_tokens": 100,
"input_tokens": 50,
"output_tokens": 50
}
}),
"openai:image",
"openai:chat",
Some(&report_context),
)
.expect("bridge should succeed")
.expect("bridge should produce sse");
let output = utf8(outcome.sse_body);
assert!(output.contains("\"object\":\"chat.completion.chunk\""));
assert!(output.contains("![generated image](data:image/png;base64,aGVsbG8=)"));
assert!(output.contains("![generated image 2](data:image/png;base64,d29ybGQ=)"));
assert!(output.contains("\"finish_reason\":\"stop\""));
assert!(output.contains("data: [DONE]"));
assert!(!output.contains("image_generation.completed"));
let summary = outcome
.terminal_summary
.expect("terminal summary should exist");
let usage = summary
.standardized_usage
.as_ref()
.expect("standard usage should exist");
assert_eq!(usage.request_count, 2);
assert_eq!(usage.dimensions.get("image_count"), Some(&json!(2)));
assert_eq!(
usage.dimensions.get("image_size"),
Some(&json!("1024x1024"))
);
}
#[test]

1
crates/aether-billing/src/default_rule.rs
View File

@@ -86,6 +86,7 @@ impl DefaultBillingRuleGenerator {
),
("cache_read_tokens", "cache_read_tokens", json!(0)),
("request_count", "request_count", json!(1)),
("image_count", "image_count", json!(0)),
] {
dimension_mappings.insert(
name.to_string(),

167
crates/aether-billing/src/event_enrichment.rs
View File

@@ -125,24 +125,37 @@ fn calculate_billing_computation(
pricing: &BillingModelPricingSnapshot,
event: &UsageEvent,
) -> Result<BillingComputation, DataLayerError> {
let failed =
event.data.status_code.unwrap_or_default() >= 400 || event.data.error_message.is_some();
let is_image_usage = usage_event_is_image_usage(&event.data);
let image_count = if failed {
0
} else {
usage_event_image_count(&event.data).unwrap_or(0)
};
let request_count = if failed {
0
} else if is_image_usage && image_count > 0 {
image_count
} else {
1
};
let input = BillingUsageInput {
task_type: event
.data
.request_type
.clone()
.unwrap_or_else(|| "chat".to_string()),
task_type: if is_image_usage {
"image".to_string()
} else {
event
.data
.request_type
.clone()
.unwrap_or_else(|| "chat".to_string())
},
api_format: event
.data
.endpoint_api_format
.clone()
.or_else(|| event.data.api_format.clone()),
request_count: if event.data.status_code.unwrap_or_default() >= 400
|| event.data.error_message.is_some()
{
0
} else {
1
},
request_count,
input_tokens: event.data.input_tokens.unwrap_or_default() as i64,
output_tokens: event.data.output_tokens.unwrap_or_default() as i64,
cache_creation_tokens: event.data.cache_creation_input_tokens.unwrap_or_default() as i64,
@@ -155,6 +168,7 @@ fn calculate_billing_computation(
.cache_creation_ephemeral_1h_input_tokens
.unwrap_or_default() as i64,
cache_read_tokens: event.data.cache_read_input_tokens.unwrap_or_default() as i64,
image_count,
cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes,
};
@@ -165,6 +179,51 @@ fn calculate_billing_computation(
})
}
fn usage_event_is_image_usage(data: &aether_usage_runtime::UsageEventData) -> bool {
data.request_type
.as_deref()
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
|| api_format_endpoint_kind(data.endpoint_api_format.as_deref()) == Some("image")
|| api_format_endpoint_kind(data.api_format.as_deref()) == Some("image")
|| usage_event_image_count(data).is_some_and(|value| value > 0)
}
fn usage_event_image_count(data: &aether_usage_runtime::UsageEventData) -> Option<i64> {
metadata_dimension_i64(data.request_metadata.as_ref(), "dimensions", "image_count")
.or_else(|| {
metadata_dimension_i64(
data.request_metadata.as_ref(),
"billing_dimensions",
"image_count",
)
})
.filter(|value| *value > 0)
}
fn metadata_dimension_i64(
metadata: Option<&Value>,
bag_key: &str,
dimension_key: &str,
) -> Option<i64> {
metadata
.and_then(Value::as_object)
.and_then(|object| object.get(bag_key))
.and_then(Value::as_object)
.and_then(|object| object.get(dimension_key))
.and_then(|value| {
value
.as_i64()
.or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok()))
})
}
fn api_format_endpoint_kind(api_format: Option<&str>) -> Option<&str> {
api_format
.and_then(|value| value.split_once(':').map(|(_, kind)| kind))
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn apply_billing_computation(
event: &mut UsageEvent,
pricing: &BillingModelPricingSnapshot,
@@ -381,6 +440,90 @@ mod tests {
);
}
#[tokio::test]
async fn image_usage_uses_image_count_for_request_cost() {
let lookup = TestLookup {
name_context: Some(
StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
None,
None,
"global-image-1".to_string(),
"gpt-image-2".to_string(),
None,
Some(0.02),
None,
Some("model-image-1".to_string()),
Some("gpt-image-2".to_string()),
None,
None,
None,
)
.expect("billing context should build"),
),
model_id_context: None,
};
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-image-billing-1",
UsageEventData {
provider_name: "OpenAI Image".to_string(),
model: "gpt-image-2".to_string(),
provider_id: Some("provider-1".to_string()),
provider_api_key_id: Some("key-1".to_string()),
request_type: Some("chat".to_string()),
api_format: Some("openai:chat".to_string()),
endpoint_api_format: Some("openai:image".to_string()),
request_metadata: Some(json!({
"dimensions": {
"image_count": 3
}
})),
status_code: Some(200),
..UsageEventData::default()
},
);
enrich_usage_event_with_billing(&lookup, &mut event)
.await
.expect("billing should succeed");
assert_eq!(event.data.total_cost_usd, Some(0.06));
assert_eq!(event.data.actual_total_cost_usd, Some(0.06));
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_dimensions"))
.and_then(|value| value.get("request_count"))
.and_then(Value::as_i64),
Some(3)
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_dimensions"))
.and_then(|value| value.get("image_count"))
.and_then(Value::as_i64),
Some(3)
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_dimensions"))
.and_then(|value| value.get("effective_task_type"))
.and_then(Value::as_str),
Some("image")
);
}
#[tokio::test]
async fn enriches_cancelled_usage_event_with_billing_snapshot() {
let lookup = TestLookup {

2
crates/aether-billing/src/pricing.rs
View File

@@ -147,6 +147,7 @@ pub struct BillingUsageInput {
pub cache_creation_ephemeral_5m_tokens: i64,
pub cache_creation_ephemeral_1h_tokens: i64,
pub cache_read_tokens: i64,
pub image_count: i64,
pub cache_ttl_minutes: Option<i64>,
}
@@ -162,6 +163,7 @@ impl BillingUsageInput {
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 0,
image_count: 0,
cache_ttl_minutes: None,
}
}

5
crates/aether-billing/src/service.rs
View File

@@ -170,6 +170,7 @@ fn build_dimensions(input: &BillingUsageInput) -> BTreeMap<String, Value> {
"request_count".to_string(),
json!(input.request_count.max(0)),
),
("image_count".to_string(), json!(input.image_count.max(0))),
(
"total_input_context".to_string(),
json!(total_input_context),
@@ -256,6 +257,7 @@ mod tests {
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 100,
image_count: 0,
cache_ttl_minutes: Some(60),
},
)
@@ -282,6 +284,7 @@ mod tests {
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 800,
image_count: 0,
cache_ttl_minutes: Some(60),
},
)
@@ -351,6 +354,7 @@ mod tests {
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 100,
image_count: 0,
cache_ttl_minutes: Some(5),
},
)
@@ -420,6 +424,7 @@ mod tests {
cache_creation_ephemeral_5m_tokens: 0,
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 100,
image_count: 0,
cache_ttl_minutes: Some(60),
},
)

61
crates/aether-usage-runtime/src/write.rs
View File

@@ -243,7 +243,8 @@ pub fn build_lifecycle_usage_seed(
let model = context_string(context, "model")
.or_else(|| non_empty_str(plan.model_name.as_deref()))
.unwrap_or_else(|| "unknown".to_string());
let request_type = infer_request_type(api_format.as_deref());
let request_type =
infer_request_type_from_contracts(api_format.as_deref(), endpoint_api_format.as_deref());
let api_family = api_format
.as_deref()
.and_then(infer_api_family)
@@ -645,7 +646,10 @@ pub fn build_terminal_usage_context_seed(
.or_else(|| context_string(context, "provider_api_format"))
.or_else(|| non_empty_str(Some(plan.provider_api_format.as_str())))
.unwrap_or_default();
let request_type = infer_request_type(Some(client_contract.as_str()));
let request_type = infer_request_type_from_contracts(
Some(client_contract.as_str()),
Some(provider_contract.as_str()),
);
let has_format_conversion = resolve_has_format_conversion(
context,
client_contract.as_str(),
@@ -1237,7 +1241,10 @@ fn build_usage_event_data_seed_with_detail(
let provider_name = context_string(context, "provider_name")
.or_else(|| non_empty_str(plan.provider_name.as_deref()))
.unwrap_or_else(|| "unknown".to_string());
let request_type = Some(infer_request_type(api_format.as_deref()));
let request_type = Some(infer_request_type_from_contracts(
api_format.as_deref(),
endpoint_api_format.as_deref(),
));
let api_family = api_format
.as_deref()
.and_then(infer_api_family)
@@ -1858,6 +1865,19 @@ fn infer_request_type(api_format: Option<&str>) -> String {
}
}
fn infer_request_type_from_contracts(
client_api_format: Option<&str>,
provider_api_format: Option<&str>,
) -> String {
if matches!(
infer_endpoint_kind(provider_api_format.unwrap_or_default()),
Some("image")
) {
return "image".to_string();
}
infer_request_type(client_api_format)
}
fn infer_api_family(api_format: &str) -> Option<&str> {
api_format.split_once(':').map(|(family, _)| family)
}
@@ -1891,6 +1911,41 @@ fn apply_standardized_usage_seed(usage: &StandardizedUsage, data: &mut UsageEven
if total_tokens > 0 {
data.total_tokens = Some(total_tokens);
}
apply_standardized_usage_dimensions_seed(usage, data);
}
fn apply_standardized_usage_dimensions_seed(usage: &StandardizedUsage, data: &mut UsageEventData) {
if usage.dimensions.is_empty() && usage.request_count <= 0 {
return;
}
let mut dimensions = usage
.dimensions
.iter()
.map(|(key, value)| (key.clone(), value.clone()))
.collect::<Map<String, Value>>();
if usage.request_count > 0 {
dimensions
.entry("request_count".to_string())
.or_insert_with(|| json!(usage.request_count));
}
if dimensions.is_empty() {
return;
}
let mut metadata = match data.request_metadata.take() {
Some(Value::Object(object)) => object,
_ => Map::new(),
};
let mut existing_dimensions = match metadata.remove("dimensions") {
Some(Value::Object(object)) => object,
_ => Map::new(),
};
for (key, value) in dimensions {
existing_dimensions.insert(key, value);
}
metadata.insert("dimensions".to_string(), Value::Object(existing_dimensions));
data.request_metadata = Some(Value::Object(metadata));
}
fn standardized_usage_total_tokens(usage: &StandardizedUsage) -> u64 {