Aether/apps/aether-gateway/src/ai_serving/planner/route.rs

use crate::ai_serving::GatewayControlDecision;
use crate::ai_serving::{
    is_matching_stream_http_request as is_matching_stream_http_request_impl,
    resolve_execution_runtime_stream_plan_kind_with_client_surface as resolve_execution_runtime_stream_plan_kind_impl,
    resolve_execution_runtime_sync_plan_kind_with_client_surface as resolve_execution_runtime_sync_plan_kind_impl,
    supports_stream_execution_decision_kind as supports_stream_execution_decision_kind_impl,
    supports_sync_execution_decision_kind as supports_sync_execution_decision_kind_impl,
};

pub(crate) fn resolve_execution_runtime_stream_plan_kind(
    parts: &http::request::Parts,
    decision: &GatewayControlDecision,
) -> Option<&'static str> {
    let plan_kind = resolve_execution_runtime_stream_plan_kind_impl(
        decision.route_class.as_deref(),
        decision.route_family.as_deref(),
        decision.route_kind.as_deref(),
        decision.client_surface,
        decision.request_auth_channel.as_deref(),
        &parts.method,
        parts.uri.path(),
    )?;
    crate::ai_serving::plan_kind_matches_api_operation(plan_kind, true, decision.api_operation)
        .then_some(plan_kind)
}

pub(crate) fn resolve_execution_runtime_sync_plan_kind(
    parts: &http::request::Parts,
    decision: &GatewayControlDecision,
) -> Option<&'static str> {
    let plan_kind = resolve_execution_runtime_sync_plan_kind_impl(
        decision.route_class.as_deref(),
        decision.route_family.as_deref(),
        decision.route_kind.as_deref(),
        decision.client_surface,
        decision.request_auth_channel.as_deref(),
        &parts.method,
        parts.uri.path(),
    )?;
    crate::ai_serving::plan_kind_matches_api_operation(plan_kind, false, decision.api_operation)
        .then_some(plan_kind)
}

pub(crate) fn is_matching_stream_request(
    plan_kind: &str,
    parts: &http::request::Parts,
    body_json: &serde_json::Value,
    body_base64: Option<&str>,
) -> bool {
    is_matching_stream_http_request_impl(plan_kind, parts, body_json, body_base64)
}

pub(crate) fn supports_sync_execution_decision_kind(plan_kind: &str) -> bool {
    supports_sync_execution_decision_kind_impl(plan_kind)
}

pub(crate) fn supports_stream_execution_decision_kind(plan_kind: &str) -> bool {
    supports_stream_execution_decision_kind_impl(plan_kind)
}

#[cfg(test)]
mod tests {
    use axum::http::{Method, Request};
    use base64::Engine as _;

    use super::{
        is_matching_stream_request, resolve_execution_runtime_stream_plan_kind,
        resolve_execution_runtime_sync_plan_kind, supports_stream_execution_decision_kind,
        supports_sync_execution_decision_kind,
    };
    use crate::ai_serving::{ApiOperation, ClientSurface, GatewayControlDecision};

    fn sample_decision(route_family: &str, route_kind: &str) -> GatewayControlDecision {
        GatewayControlDecision {
            public_path: "/".to_string(),
            public_query_string: None,
            route_class: Some("ai_public".to_string()),
            route_family: Some(route_family.to_string()),
            route_kind: Some(route_kind.to_string()),
            client_surface: None,
            api_operation: None,
            gateway_credential_carrier: None,
            request_auth_channel: None,
            auth_context: None,
            admin_principal: None,
            auth_endpoint_signature: None,
            execution_runtime_candidate: true,
            local_auth_rejection: None,
            model_directive_policy: Default::default(),
        }
    }

    fn sample_decision_with_auth_channel(
        route_family: &str,
        route_kind: &str,
        request_auth_channel: &str,
    ) -> GatewayControlDecision {
        let mut decision = sample_decision(route_family, route_kind);
        decision.request_auth_channel = Some(request_auth_channel.to_string());
        decision
    }

    #[test]
    fn resolves_openai_chat_plan_kinds_via_format_crate() {
        let request = Request::builder()
            .method(Method::POST)
            .uri("/v1/chat/completions")
            .body(())
            .expect("request should build");
        let (parts, _) = request.into_parts();
        let decision = sample_decision("openai", "chat");

        assert_eq!(
            resolve_execution_runtime_sync_plan_kind(&parts, &decision),
            Some("openai_chat_sync")
        );
        assert_eq!(
            resolve_execution_runtime_stream_plan_kind(&parts, &decision),
            Some("openai_chat_stream")
        );
    }

    #[test]
    fn resolves_endpoint_route_kinds_by_request_auth_channel_via_format_crate() {
        let claude_request = Request::builder()
            .method(Method::POST)
            .uri("/v1/messages")
            .body(())
            .expect("request should build");
        let (claude_parts, _) = claude_request.into_parts();

        let claude_api_key = sample_decision_with_auth_channel("claude", "messages", "api_key");
        let mut claude_bearer =
            sample_decision_with_auth_channel("claude", "messages", "bearer_like");
        claude_bearer.client_surface = Some(ClientSurface::ClaudeCode);
        assert_eq!(
            resolve_execution_runtime_sync_plan_kind(&claude_parts, &claude_api_key),
            Some("claude_chat_sync")
        );
        assert_eq!(
            resolve_execution_runtime_stream_plan_kind(&claude_parts, &claude_bearer),
            Some("claude_cli_stream")
        );

        let claude_sdk_bearer =
            sample_decision_with_auth_channel("claude", "messages", "bearer_like");
        assert_eq!(
            resolve_execution_runtime_sync_plan_kind(&claude_parts, &claude_sdk_bearer),
            Some("claude_chat_sync")
        );

        let gemini_request = Request::builder()
            .method(Method::POST)
            .uri("/v1beta/models/gemini-2.5-pro:generateContent")
            .body(())
            .expect("request should build");
        let (gemini_parts, _) = gemini_request.into_parts();

        let gemini_api_key =
            sample_decision_with_auth_channel("gemini", "generate_content", "api_key");
        let gemini_bearer =
            sample_decision_with_auth_channel("gemini", "generate_content", "bearer_like");
        assert_eq!(
            resolve_execution_runtime_sync_plan_kind(&gemini_parts, &gemini_api_key),
            Some("gemini_chat_sync")
        );
        assert_eq!(
            resolve_execution_runtime_sync_plan_kind(&gemini_parts, &gemini_bearer),
            Some("gemini_cli_sync")
        );
    }

    #[test]
    fn resolves_claude_count_tokens_as_native_sync_operation() {
        let request = Request::builder()
            .method(Method::POST)
            .uri("/v1/messages/count_tokens")
            .body(())
            .expect("request should build");
        let (parts, _) = request.into_parts();
        let mut decision = sample_decision("claude", "count_tokens");
        decision.api_operation = Some(ApiOperation::ClaudeCountTokens);

        assert_eq!(
            resolve_execution_runtime_sync_plan_kind(&parts, &decision),
            Some("claude_count_tokens_sync")
        );
        assert!(supports_sync_execution_decision_kind(
            "claude_count_tokens_sync"
        ));

        decision.api_operation = Some(ApiOperation::ClaudeMessagesCreate);
        assert_eq!(
            resolve_execution_runtime_sync_plan_kind(&parts, &decision),
            None
        );
        assert_eq!(
            resolve_execution_runtime_stream_plan_kind(&parts, &decision),
            None
        );
    }

    #[test]
    fn stream_matching_uses_surface_route_logic() {
        let request = Request::builder()
            .method(Method::POST)
            .uri("/v1/chat/completions")
            .body(())
            .expect("request should build");
        let (parts, _) = request.into_parts();

        assert!(!is_matching_stream_request(
            "openai_chat_stream",
            &parts,
            &serde_json::json!({"stream": false}),
            None,
        ));
        assert!(is_matching_stream_request(
            "openai_chat_stream",
            &parts,
            &serde_json::json!({"stream": true}),
            None,
        ));
        assert!(supports_sync_execution_decision_kind("openai_chat_sync"));
        assert!(supports_stream_execution_decision_kind(
            "openai_chat_stream"
        ));
    }

    #[test]
    fn image_stream_matching_parses_multipart_stream_flag() {
        let request = Request::builder()
            .method(Method::POST)
            .uri("/v1/images/edits")
            .header(
                http::header::CONTENT_TYPE,
                "multipart/form-data; boundary=image-stream-boundary",
            )
            .body(())
            .expect("request should build");
        let (parts, _) = request.into_parts();
        let body = concat!(
            "--image-stream-boundary\r\n",
            "Content-Disposition: form-data; name=\"stream\"\r\n\r\n",
            "true\r\n",
            "--image-stream-boundary--\r\n"
        );
        let body_base64 = base64::engine::general_purpose::STANDARD.encode(body.as_bytes());

        assert!(is_matching_stream_request(
            "openai_image_stream",
            &parts,
            &serde_json::json!({}),
            Some(body_base64.as_str()),
        ));
    }
}