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

crates/aether-billing/src/default_rule.rs

@@ -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(),

crates/aether-billing/src/event_enrichment.rs

@@ -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 {

crates/aether-billing/src/pricing.rs

@@ -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,
         }
     }

crates/aether-billing/src/service.rs

@@ -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),
                 },
             )