Aether/crates/aether-ai/serving/src/report_context.rs

use std::collections::BTreeMap;

use aether_ai_formats::api::ExecutionRuntimeAuthContext;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_scheduler_core::SchedulerRankingOutcome;
use serde_json::{Map, Value};

use crate::append_ai_ranking_metadata_to_object;

#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AiRequestOrigin {
    pub client_ip: Option<String>,
    pub user_agent: Option<String>,
}

pub struct AiExecutionReportContextParts<'a> {
    pub auth_context: &'a ExecutionRuntimeAuthContext,
    pub request_id: &'a str,
    pub candidate_id: &'a str,
    pub candidate_index: u32,
    pub retry_index: u32,
    pub pool_key_index: Option<u32>,
    pub model: &'a str,
    pub provider_name: &'a str,
    pub provider_id: &'a str,
    pub endpoint_id: &'a str,
    pub key_id: &'a str,
    pub key_name: Option<&'a str>,
    pub model_id: Option<&'a str>,
    pub global_model_id: Option<&'a str>,
    pub global_model_name: Option<&'a str>,
    pub provider_api_format: &'a str,
    pub client_api_format: &'a str,
    pub mapped_model: Option<&'a str>,
    pub candidate_group_id: Option<&'a str>,
    pub ranking: Option<&'a SchedulerRankingOutcome>,
    pub upstream_url: Option<&'a str>,
    pub header_rules: Option<&'a Value>,
    pub body_rules: Option<&'a Value>,
    pub provider_request_method: Option<Value>,
    pub provider_request_headers: Option<&'a BTreeMap<String, String>>,
    pub original_headers: &'a BTreeMap<String, String>,
    pub original_request_body: Option<Value>,
    pub request_origin: AiRequestOrigin,
    pub client_requested_stream: bool,
    pub upstream_is_stream: bool,
    pub has_envelope: bool,
    pub needs_conversion: bool,
    pub extra_fields: Map<String, Value>,
}

pub fn build_ai_execution_report_context(parts: AiExecutionReportContextParts<'_>) -> Value {
    let mut object = Map::new();
    object.insert(
        "user_id".to_string(),
        Value::String(parts.auth_context.user_id.clone()),
    );
    object.insert(
        "api_key_id".to_string(),
        Value::String(parts.auth_context.api_key_id.clone()),
    );
    object.insert(
        "api_key_is_standalone".to_string(),
        Value::Bool(parts.auth_context.api_key_is_standalone),
    );
    object.insert(
        "analytics_attribution".into(),
        serde_json::json!({
            "is_standalone": parts.auth_context.api_key_is_standalone,
            "record_kind": "request"
        }),
    );
    object.insert(
        "username".to_string(),
        parts
            .auth_context
            .username
            .clone()
            .map(Value::String)
            .unwrap_or(Value::Null),
    );
    object.insert(
        "api_key_name".to_string(),
        parts
            .auth_context
            .api_key_name
            .clone()
            .map(Value::String)
            .unwrap_or(Value::Null),
    );
    object.insert(
        "request_id".to_string(),
        Value::String(parts.request_id.to_string()),
    );
    object.insert(
        "candidate_id".to_string(),
        Value::String(parts.candidate_id.to_string()),
    );
    object.insert(
        "candidate_index".to_string(),
        Value::Number(parts.candidate_index.into()),
    );
    object.insert(
        "retry_index".to_string(),
        Value::Number(parts.retry_index.into()),
    );
    object.insert("model".to_string(), Value::String(parts.model.to_string()));
    object.insert(
        "provider_name".to_string(),
        Value::String(parts.provider_name.to_string()),
    );
    object.insert(
        "provider_id".to_string(),
        Value::String(parts.provider_id.to_string()),
    );
    object.insert(
        "endpoint_id".to_string(),
        Value::String(parts.endpoint_id.to_string()),
    );
    object.insert(
        "key_id".to_string(),
        Value::String(parts.key_id.to_string()),
    );
    object.insert(
        "provider_api_format".to_string(),
        Value::String(parts.provider_api_format.to_string()),
    );
    object.insert(
        "client_api_format".to_string(),
        Value::String(parts.client_api_format.to_string()),
    );
    object.insert(
        "original_headers".to_string(),
        serde_json::to_value(parts.original_headers).expect("control headers should serialize"),
    );
    object.insert(
        "original_request_body".to_string(),
        parts.original_request_body.unwrap_or(Value::Null),
    );
    if let Some(client_ip) = parts.request_origin.client_ip {
        object.insert("client_ip".to_string(), Value::String(client_ip));
    }
    if let Some(user_agent) = parts.request_origin.user_agent {
        object.insert("user_agent".to_string(), Value::String(user_agent));
    }
    object.insert(
        "client_requested_stream".to_string(),
        Value::Bool(parts.client_requested_stream),
    );
    object.insert(
        UPSTREAM_IS_STREAM_KEY.to_string(),
        Value::Bool(parts.upstream_is_stream),
    );
    object.insert("has_envelope".to_string(), Value::Bool(parts.has_envelope));
    object.insert(
        "needs_conversion".to_string(),
        Value::Bool(parts.needs_conversion),
    );

    if let Some(key_name) = parts.key_name {
        object.insert("key_name".to_string(), Value::String(key_name.to_string()));
    }
    if let Some(model_id) = parts.model_id {
        object.insert("model_id".to_string(), Value::String(model_id.to_string()));
    }
    if let Some(global_model_id) = parts.global_model_id {
        object.insert(
            "global_model_id".to_string(),
            Value::String(global_model_id.to_string()),
        );
    }
    if let Some(global_model_name) = parts.global_model_name {
        object.insert(
            "global_model_name".to_string(),
            Value::String(global_model_name.to_string()),
        );
    }
    if let Some(mapped_model) = parts.mapped_model {
        object.insert(
            "mapped_model".to_string(),
            Value::String(mapped_model.to_string()),
        );
    }
    if let Some(candidate_group_id) = parts.candidate_group_id {
        object.insert(
            "candidate_group_id".to_string(),
            Value::String(candidate_group_id.to_string()),
        );
    }
    if let Some(ranking) = parts.ranking {
        append_ai_ranking_metadata_to_object(&mut object, ranking);
    }
    if let Some(upstream_url) = parts.upstream_url {
        object.insert(
            "upstream_url".to_string(),
            Value::String(upstream_url.to_string()),
        );
    }
    if let Some(header_rules) = parts.header_rules {
        object.insert("header_rules".to_string(), header_rules.clone());
    }
    if let Some(body_rules) = parts.body_rules {
        object.insert("body_rules".to_string(), body_rules.clone());
    }
    if let Some(provider_request_method) = parts.provider_request_method {
        object.insert(
            "provider_request_method".to_string(),
            provider_request_method,
        );
    }
    if let Some(provider_request_headers) = parts.provider_request_headers {
        object.insert(
            "provider_request_headers".to_string(),
            serde_json::to_value(provider_request_headers)
                .expect("provider request headers should serialize"),
        );
    }
    if let Some(pool_key_index) = parts.pool_key_index {
        object.insert(
            "pool_key_index".to_string(),
            Value::Number(pool_key_index.into()),
        );
    }

    object.extend(parts.extra_fields.into_iter().filter(|(key, _)| {
        !matches!(
            key.as_str(),
            "analytics_attribution"
                | "analytics_failure"
                | "analytics_measurement"
                | "usage_token_source"
        )
    }));
    if uses_estimated_token_adapter(&object) {
        // Bind the adapter's provenance to the planner context. This is a hint,
        // not a measurement: usage records materialize it only when tokens exist.
        object.insert(
            "usage_token_source".into(),
            Value::String("estimated".into()),
        );
    }
    Value::Object(object)
}

fn uses_estimated_token_adapter(context: &Map<String, Value>) -> bool {
    let grok = context
        .get("provider_type")
        .and_then(Value::as_str)
        .is_some_and(|value| value.eq_ignore_ascii_case("grok"))
        || context
            .get("provider_request_headers")
            .and_then(Value::as_object)
            .is_some_and(|headers| {
                headers.iter().any(|(name, value)| {
                    name.eq_ignore_ascii_case("x-aether-grok-runtime")
                        && value.as_str() == Some("1")
                })
            });
    let matches = |key: &str, expected: &str| {
        context
            .get(key)
            .and_then(Value::as_str)
            .is_some_and(|value| value.trim().eq_ignore_ascii_case(expected))
    };
    let kiro = context.get("has_envelope").and_then(Value::as_bool) == Some(true)
        && matches("envelope_name", "kiro:generateAssistantResponse")
        && matches("provider_api_format", "claude:messages")
        && !matches("client_envelope_name", "kiro:generateAssistantResponse");
    grok || kiro
}

pub fn provider_stream_event_api_format_for_provider_type(
    provider_type: &str,
) -> Option<&'static str> {
    match provider_type.trim().to_ascii_lowercase().as_str() {
        "codex" => Some("openai:responses"),
        _ => None,
    }
}

pub fn insert_provider_stream_event_api_format(
    extra_fields: &mut Map<String, Value>,
    provider_type: &str,
) {
    if let Some(api_format) = provider_stream_event_api_format_for_provider_type(provider_type) {
        extra_fields.insert(
            "provider_stream_event_api_format".to_string(),
            Value::String(api_format.to_string()),
        );
    }
}

pub fn build_ai_report_context_original_request_echo(
    body_json: Option<&Value>,
    body_bytes_b64: Option<&str>,
) -> Option<Value> {
    if let Some(body_bytes_b64) = body_bytes_b64
        .map(str::trim)
        .filter(|value| !value.is_empty())
    {
        return Some(serde_json::json!({ "body_bytes_b64": body_bytes_b64 }));
    }

    body_json.filter(|body| !body.is_null()).cloned()
}

#[cfg(test)]
mod tests {
    use super::*;
    use aether_scheduler_core::{SchedulerPriorityMode, SchedulerRankingMode};
    use serde_json::json;

    fn sample_auth_context() -> ExecutionRuntimeAuthContext {
        ExecutionRuntimeAuthContext {
            user_id: "user-1".to_string(),
            api_key_id: "key-1".to_string(),
            username: Some("alice".to_string()),
            api_key_name: Some("primary".to_string()),
            balance_remaining: Some(42.0),
            access_allowed: true,
            api_key_is_standalone: false,
        }
    }

    #[test]
    fn report_context_builds_core_execution_fields() {
        let auth_context = sample_auth_context();
        let original_headers = BTreeMap::from([("x-trace-id".to_string(), "trace-a".to_string())]);
        let provider_headers =
            BTreeMap::from([("authorization".to_string(), "Bearer token".to_string())]);
        let mut extra_fields = Map::new();
        extra_fields.insert("extra".to_string(), json!("value"));
        extra_fields.insert(
            "analytics_attribution".into(),
            json!({"is_standalone":true,"actor_user_id":"forged"}),
        );
        extra_fields.insert("analytics_failure".into(), json!({"origin":"client"}));
        extra_fields.insert("analytics_measurement".into(), json!({"source":"reported"}));
        extra_fields.insert("usage_token_source".into(), json!("estimated"));
        let ranking = SchedulerRankingOutcome {
            original_index: 2,
            ranking_index: 1,
            priority_mode: SchedulerPriorityMode::Provider,
            ranking_mode: SchedulerRankingMode::CacheAffinity,
            priority_slot: 7,
            promoted_by: Some("cached_affinity"),
            demoted_by: None,
        };

        let build_report = |provider_headers: &BTreeMap<String, String>,
                            extra_fields: Map<String, Value>| {
            build_ai_execution_report_context(AiExecutionReportContextParts {
                auth_context: &auth_context,
                request_id: "trace-a",
                candidate_id: "candidate-a",
                candidate_index: 3,
                retry_index: 1,
                pool_key_index: Some(0),
                model: "gpt-5",
                provider_name: "RightCode",
                provider_id: "provider-1",
                endpoint_id: "endpoint-1",
                key_id: "key-1",
                key_name: Some("primary"),
                model_id: Some("model-1"),
                global_model_id: Some("global-1"),
                global_model_name: Some("GPT-5"),
                provider_api_format: "openai:responses",
                client_api_format: "openai:chat",
                mapped_model: Some("gpt-5"),
                candidate_group_id: Some("group-1"),
                ranking: Some(&ranking),
                upstream_url: Some("https://example.com/v1/responses"),
                header_rules: Some(&json!({"set": []})),
                body_rules: None,
                provider_request_method: Some(json!("POST")),
                provider_request_headers: Some(provider_headers),
                original_headers: &original_headers,
                original_request_body: Some(json!({"model": "gpt-5"})),
                request_origin: AiRequestOrigin {
                    client_ip: Some("127.0.0.1".to_string()),
                    user_agent: Some("test-agent".to_string()),
                },
                client_requested_stream: false,
                upstream_is_stream: true,
                has_envelope: false,
                needs_conversion: true,
                extra_fields,
            })
        };
        let report = build_report(&provider_headers, extra_fields.clone());

        assert_eq!(report["user_id"], "user-1");
        assert_eq!(report["candidate_index"], 3);
        assert_eq!(report["retry_index"], 1);
        assert_eq!(report["pool_key_index"], 0);
        assert_eq!(report["original_headers"]["x-trace-id"], "trace-a");
        assert_eq!(report["original_request_body"]["model"], "gpt-5");
        assert_eq!(report["ranking_index"], 1);
        assert_eq!(
            report["provider_request_headers"]["authorization"],
            "Bearer token"
        );
        assert_eq!(report["extra"], "value");
        assert_eq!(
            report["analytics_attribution"],
            json!({"is_standalone":false,"record_kind":"request"})
        );
        assert!(report["analytics_attribution"]
            .get("actor_user_id")
            .is_none());
        assert!(report.get("analytics_failure").is_none());
        assert!(report.get("analytics_measurement").is_none());
        assert!(report.get("usage_token_source").is_none());

        let grok_headers = BTreeMap::from([("X-Aether-Grok-Runtime".into(), "1".into())]);
        assert_eq!(
            build_report(&grok_headers, extra_fields.clone())["usage_token_source"],
            "estimated"
        );
        let mut native = extra_fields.clone();
        native.insert("provider_type".into(), json!("Grok"));
        assert_eq!(
            build_report(&provider_headers, native)["usage_token_source"],
            "estimated"
        );

        let mut kiro = extra_fields;
        kiro.insert("has_envelope".into(), json!(true));
        kiro.insert(
            "envelope_name".into(),
            json!("kiro:generateAssistantResponse"),
        );
        kiro.insert("provider_api_format".into(), json!("claude:messages"));
        assert_eq!(
            build_report(&provider_headers, kiro.clone())["usage_token_source"],
            "estimated"
        );
        kiro.insert(
            "client_envelope_name".into(),
            json!("kiro:generateAssistantResponse"),
        );
        assert!(build_report(&provider_headers, kiro)
            .get("usage_token_source")
            .is_none());
    }

    #[test]
    fn provider_stream_event_api_format_is_codex_only() {
        assert_eq!(
            provider_stream_event_api_format_for_provider_type("codex"),
            Some("openai:responses")
        );
        assert_eq!(
            provider_stream_event_api_format_for_provider_type(" CODEX "),
            Some("openai:responses")
        );
        assert_eq!(
            provider_stream_event_api_format_for_provider_type("openai"),
            None
        );
    }

    #[test]
    fn original_request_echo_preserves_full_request_body() {
        let body = json!({
            "messages": [{"role": "user", "content": "large payload should be omitted"}],
            "service_tier": "default",
            "instructions": "Be concise.",
            "thinking": {"type": "enabled", "budget_tokens": 512},
            "metadata": {"trace": "keep"},
            "body_bytes_b64": "aGVsbG8=",
        });

        let echo = build_ai_report_context_original_request_echo(Some(&body), None)
            .expect("echo should be produced");

        assert_eq!(echo, body);
    }

    #[test]
    fn original_request_echo_prefers_binary_body_bytes() {
        let echo = build_ai_report_context_original_request_echo(
            Some(&json!({"ignored": true})),
            Some("aGVsbG8="),
        )
        .expect("echo should be produced");

        assert_eq!(echo, json!({"body_bytes_b64": "aGVsbG8="}));
    }
}