Aether/apps/aether-gateway/src/state/integrations.rs

use std::time::Duration;

use aether_contracts::{ExecutionPlan, ExecutionResult, ProxySnapshot};
use aether_data_contracts::repository::candidates::{
    StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use aether_data_contracts::repository::global_models::{
    AdminGlobalModelListQuery, AdminProviderModelListQuery, StoredAdminGlobalModelPage,
    StoredAdminProviderModel, UpsertAdminProviderModelRecord,
};
use aether_data_contracts::repository::provider_catalog::{
    ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint,
    StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_data_contracts::DataLayerError;
use aether_model_fetch::{
    aggregate_models_for_cache, build_antigravity_load_code_assist_plan,
    fetch_models_from_transports, merge_upstream_metadata, model_fetch_interval_minutes,
    ModelFetchAssociationStore, ModelFetchTransportRuntime,
};
use aether_scheduler_core::SchedulerAffinityTarget;
use async_trait::async_trait;
use serde_json::Value;
use tracing::{debug, warn};

use super::{AppState, GatewayError};
use crate::clock::current_unix_secs;
use crate::model_fetch::{safe_model_fetch_error, CodexCatalogRuntime, ModelFetchRuntimeState};
use crate::provider_transport::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
use crate::request_candidate_runtime::{
    RequestCandidateRuntimeCapabilityReader, RequestCandidateRuntimeReader,
    RequestCandidateRuntimeWriter,
};
use crate::scheduler::state::SchedulerRuntimeState;
use crate::{execution_runtime, provider_transport};

const MODEL_FETCH_RESPONSE_BODY_LIMIT_BYTES: usize = 8 * 1024 * 1024;

#[async_trait]
impl provider_transport::ProviderTransportSnapshotSource for AppState {
    fn encryption_key(&self) -> Option<&str> {
        AppState::encryption_key(self)
    }

    async fn list_provider_catalog_providers_by_ids(
        &self,
        ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
        self.read_provider_catalog_providers_by_ids(ids)
            .await
            .map_err(provider_transport_snapshot_data_error)
    }

    async fn list_provider_catalog_endpoints_by_ids(
        &self,
        ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
        self.read_provider_catalog_endpoints_by_ids(ids)
            .await
            .map_err(provider_transport_snapshot_data_error)
    }

    async fn list_provider_catalog_keys_by_ids(
        &self,
        ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
        self.read_provider_catalog_keys_by_ids(ids)
            .await
            .map_err(provider_transport_snapshot_data_error)
    }
}

fn provider_transport_snapshot_data_error(error: GatewayError) -> DataLayerError {
    DataLayerError::UnexpectedValue(error.into_message())
}

impl AppState {
    pub(crate) async fn hydrate_antigravity_project_metadata_for_transport(
        &self,
        transport: &GatewayProviderTransportSnapshot,
    ) -> Option<GatewayProviderTransportSnapshot> {
        if !provider_transport::antigravity::is_antigravity_provider_transport(transport) {
            return None;
        }
        if matches!(
            provider_transport::antigravity::resolve_local_antigravity_request_auth(transport),
            provider_transport::antigravity::AntigravityRequestAuthSupport::Supported(_)
        ) {
            return Some(transport.clone());
        }

        let plan = match build_antigravity_load_code_assist_plan(self, transport).await {
            Ok(plan) => plan,
            Err(err) => {
                warn!(
                    provider_id = %transport.provider.id,
                    endpoint_id = %transport.endpoint.id,
                    key_id = %transport.key.id,
                    error = %err,
                    "antigravity project metadata hydration failed"
                );
                return None;
            }
        };
        let result =
            match execution_runtime::execute_execution_runtime_sync_plan(self, None, &plan).await {
                Ok(result) => result,
                Err(err) => {
                    warn!(
                        provider_id = %transport.provider.id,
                        endpoint_id = %transport.endpoint.id,
                        key_id = %transport.key.id,
                        error = ?err,
                        "antigravity project metadata hydration request failed"
                    );
                    return None;
                }
            };
        if !(200..300).contains(&result.status_code) {
            warn!(
                provider_id = %transport.provider.id,
                endpoint_id = %transport.endpoint.id,
                key_id = %transport.key.id,
                status_code = result.status_code,
                "antigravity project metadata hydration returned non-success status"
            );
            return None;
        }
        let Some(project_id) = result
            .body
            .as_ref()
            .and_then(|body| body.json_body.as_ref())
            .and_then(extract_antigravity_load_code_assist_project_id)
        else {
            warn!(
                provider_id = %transport.provider.id,
                endpoint_id = %transport.endpoint.id,
                key_id = %transport.key.id,
                "antigravity project metadata hydration response missing project"
            );
            return None;
        };
        let upstream_metadata = serde_json::json!({
            "antigravity": {
                "project_id": project_id,
                "updated_at": current_unix_secs(),
            }
        });
        let merged_metadata =
            merge_upstream_metadata(transport.key.upstream_metadata.as_ref(), &upstream_metadata);

        let mut hydrated = transport.clone();
        hydrated.key.upstream_metadata = Some(merged_metadata.clone());
        if !matches!(
            provider_transport::antigravity::resolve_local_antigravity_request_auth(&hydrated),
            provider_transport::antigravity::AntigravityRequestAuthSupport::Supported(_)
        ) {
            return None;
        }

        if let Err(err) = self
            .update_provider_catalog_key_upstream_metadata(
                &transport.key.id,
                Some(&merged_metadata),
                Some(current_unix_secs()),
            )
            .await
        {
            warn!(
                provider_id = %transport.provider.id,
                endpoint_id = %transport.endpoint.id,
                key_id = %transport.key.id,
                error = ?err,
                "antigravity project metadata hydration could not persist metadata"
            );
        }

        Some(hydrated)
    }

    pub(crate) async fn hydrate_gemini_cli_project_metadata_for_transport(
        &self,
        transport: &GatewayProviderTransportSnapshot,
    ) -> Option<GatewayProviderTransportSnapshot> {
        if !provider_transport::is_gemini_cli_provider_transport(transport) {
            return None;
        }
        if provider_transport::resolve_gemini_cli_project_id(transport).is_some() {
            return Some(transport.clone());
        }

        let outcome =
            match fetch_models_from_transports(self, std::slice::from_ref(transport)).await {
                Ok(outcome) => outcome,
                Err(err) => {
                    warn!(
                        provider_id = %transport.provider.id,
                        endpoint_id = %transport.endpoint.id,
                        key_id = %transport.key.id,
                        error = %safe_model_fetch_error(&err),
                        "gemini_cli project metadata hydration failed"
                    );
                    return None;
                }
            };
        let upstream_metadata = outcome.upstream_metadata.as_ref()?;
        let merged_metadata =
            merge_upstream_metadata(transport.key.upstream_metadata.as_ref(), upstream_metadata);

        let mut hydrated = transport.clone();
        hydrated.key.upstream_metadata = Some(merged_metadata.clone());
        if provider_transport::resolve_gemini_cli_project_id(&hydrated).is_none() {
            return None;
        }

        if let Err(err) = self
            .update_provider_catalog_key_upstream_metadata(
                &transport.key.id,
                Some(&merged_metadata),
                Some(current_unix_secs()),
            )
            .await
        {
            warn!(
                provider_id = %transport.provider.id,
                endpoint_id = %transport.endpoint.id,
                key_id = %transport.key.id,
                error = ?err,
                "gemini_cli project metadata hydration could not persist metadata"
            );
        }

        Some(hydrated)
    }
}

fn extract_antigravity_load_code_assist_project_id(value: &Value) -> Option<String> {
    let raw = value
        .get("cloudaicompanionProject")
        .or_else(|| value.get("cloudAiCompanionProject"))?;
    if let Some(project_id) = raw
        .as_str()
        .map(str::trim)
        .filter(|value| !value.is_empty())
    {
        return Some(project_id.to_string());
    }
    raw.as_object()
        .and_then(|object| {
            object
                .get("id")
                .or_else(|| object.get("project_id"))
                .or_else(|| object.get("projectId"))
        })
        .and_then(Value::as_str)
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned)
}

#[async_trait]
impl provider_transport::TransportTunnelAffinityLookup for AppState {
    async fn lookup_tunnel_attachment_owner(
        &self,
        node_id: &str,
    ) -> Result<Option<provider_transport::TransportTunnelAttachmentOwner>, String> {
        self.tunnel
            .lookup_attachment_owner(self.data.as_ref(), node_id)
            .await
            .map(|owner| {
                owner.map(|owner| provider_transport::TransportTunnelAttachmentOwner {
                    gateway_instance_id: owner.gateway_instance_id,
                    relay_base_url: owner.relay_base_url,
                    observed_at_unix_secs: owner.observed_at_unix_secs,
                })
            })
    }
}

#[async_trait]
impl provider_transport::VideoTaskTransportSnapshotLookup for AppState {
    async fn read_video_task_provider_transport_snapshot(
        &self,
        provider_id: &str,
        endpoint_id: &str,
        key_id: &str,
    ) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
        self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
            .await
            .map_err(GatewayError::into_message)
    }

    async fn resolve_video_task_proxy(
        &self,
        transport: &GatewayProviderTransportSnapshot,
    ) -> Option<ProxySnapshot> {
        self.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
            .await
    }
}

#[async_trait]
impl ModelFetchTransportRuntime for AppState {
    async fn resolve_local_oauth_request_auth(
        &self,
        transport: &GatewayProviderTransportSnapshot,
    ) -> Result<Option<LocalResolvedOAuthRequestAuth>, String> {
        AppState::resolve_local_oauth_request_auth(self, transport)
            .await
            .map_err(GatewayError::into_message)
    }

    async fn resolve_model_fetch_proxy(
        &self,
        transport: &GatewayProviderTransportSnapshot,
    ) -> Option<ProxySnapshot> {
        self.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
            .await
    }

    async fn execute_model_fetch_execution_plan(
        &self,
        plan: &ExecutionPlan,
    ) -> Result<ExecutionResult, String> {
        let bounded_plan = execution_runtime::transport::with_upstream_response_body_limit(
            plan,
            MODEL_FETCH_RESPONSE_BODY_LIMIT_BYTES,
        );
        execution_runtime::execute_execution_runtime_sync_plan(self, None, &bounded_plan)
            .await
            .map_err(GatewayError::into_message)
    }
}

#[async_trait]
impl CodexCatalogRuntime for AppState {
    fn codex_catalog_runtime_state(&self) -> &aether_runtime_state::RuntimeState {
        self.runtime_state.as_ref()
    }

    async fn read_codex_catalog_transport_snapshot(
        &self,
        provider_id: &str,
        endpoint_id: &str,
        key_id: &str,
    ) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
        self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
            .await
            .map_err(GatewayError::into_message)
    }

    async fn read_codex_catalog_credential_scope_strong(
        &self,
        provider_id: &str,
        key_id: &str,
    ) -> Result<Option<String>, String> {
        let Some(key) = self
            .list_provider_catalog_keys_by_ids_strong(&[key_id.to_string()])
            .await
            .map_err(GatewayError::into_message)?
            .into_iter()
            .find(|key| key.id == key_id && key.provider_id == provider_id && key.is_active)
        else {
            return Ok(None);
        };

        if let Some(scope) =
            crate::model_fetch::codex_catalog_credential_scope_from_stored_key(&key, None, None)
        {
            return Ok(Some(scope));
        }

        let decrypted_auth_config = self
            .decrypt_provider_catalog_key_auth_config(&key)
            .map_err(GatewayError::into_message)?;
        let decrypted_api_key = self
            .decrypt_provider_catalog_key_api_key(&key)
            .map_err(GatewayError::into_message)?;

        Ok(
            crate::model_fetch::codex_catalog_credential_scope_from_stored_key(
                &key,
                decrypted_auth_config.as_deref(),
                decrypted_api_key.as_deref(),
            ),
        )
    }
}

#[async_trait]
impl ModelFetchRuntimeState for AppState {
    fn has_provider_catalog_data_reader(&self) -> bool {
        AppState::has_provider_catalog_data_reader(self)
    }

    fn has_provider_catalog_data_writer(&self) -> bool {
        AppState::has_provider_catalog_data_writer(self)
    }

    async fn list_provider_catalog_providers(
        &self,
        active_only: bool,
    ) -> Result<Vec<StoredProviderCatalogProvider>, GatewayError> {
        AppState::list_provider_catalog_providers(self, active_only).await
    }

    async fn list_provider_catalog_providers_for_model_fetch(
        &self,
        active_only: bool,
    ) -> Result<Vec<StoredProviderCatalogProvider>, GatewayError> {
        self.data
            .list_provider_catalog_providers(active_only)
            .await
            .map_err(|err| GatewayError::Internal(err.to_string()))
    }

    async fn list_provider_catalog_endpoints_by_provider_ids(
        &self,
        provider_ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogEndpoint>, GatewayError> {
        AppState::list_provider_catalog_endpoints_by_provider_ids(self, provider_ids).await
    }

    async fn list_provider_catalog_endpoints_for_model_fetch(
        &self,
        provider_ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogEndpoint>, GatewayError> {
        self.data
            .list_provider_catalog_endpoints_by_provider_ids(provider_ids)
            .await
            .map_err(|err| GatewayError::Internal(err.to_string()))
    }

    async fn list_provider_catalog_keys_for_model_fetch(
        &self,
        provider_ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogKey>, String> {
        self.data
            .list_provider_catalog_keys_by_provider_ids(provider_ids)
            .await
            .map_err(|err| err.to_string())
    }

    async fn read_provider_transport_snapshot(
        &self,
        provider_id: &str,
        endpoint_id: &str,
        key_id: &str,
    ) -> Result<Option<GatewayProviderTransportSnapshot>, GatewayError> {
        AppState::read_provider_transport_snapshot(self, provider_id, endpoint_id, key_id).await
    }

    async fn execute_execution_runtime_sync_plan(
        &self,
        plan: &ExecutionPlan,
    ) -> Result<ExecutionResult, GatewayError> {
        execution_runtime::execute_execution_runtime_sync_plan(self, None, plan).await
    }

    async fn read_recent_codex_catalog_client_version(
        &self,
        provider_id: &str,
        key_id: &str,
    ) -> Option<String> {
        crate::model_fetch::read_codex_management_catalog(self, provider_id, key_id)
            .await
            .map(|catalog| catalog.client_version)
    }

    async fn update_provider_catalog_key_model_fetch_state(
        &self,
        key_id: &str,
        allowed_models: Option<&Value>,
        last_models_fetch_at_unix_secs: Option<u64>,
        last_models_fetch_error: Option<&str>,
        updated_at_unix_secs: Option<u64>,
    ) -> Result<(), GatewayError> {
        AppState::update_provider_catalog_key_model_fetch_state(
            self,
            key_id,
            allowed_models,
            last_models_fetch_at_unix_secs,
            last_models_fetch_error,
            updated_at_unix_secs,
        )
        .await?;
        Ok(())
    }

    async fn update_provider_catalog_key_model_fetch_success(
        &self,
        key_id: &str,
        allowed_models: Option<&Value>,
        last_models_fetch_at_unix_secs: u64,
        upstream_metadata_updates: &[ProviderCatalogUpstreamMetadataNamespaceUpdate],
        updated_at_unix_secs: Option<u64>,
    ) -> Result<(), GatewayError> {
        AppState::update_provider_catalog_key_model_fetch_success(
            self,
            key_id,
            allowed_models,
            last_models_fetch_at_unix_secs,
            upstream_metadata_updates,
            updated_at_unix_secs,
        )
        .await?;
        Ok(())
    }

    async fn write_upstream_models_cache(
        &self,
        provider_id: &str,
        key_id: &str,
        cached_models: &[Value],
    ) {
        let models = aggregate_models_for_cache(cached_models);
        if models.is_empty() {
            return;
        }
        let Ok(serialized) = serde_json::to_string(&models) else {
            return;
        };
        let cache_key = format!("upstream_models:{provider_id}:{key_id}");
        if let Err(err) = self
            .runtime_state
            .kv_set(
                &cache_key,
                serialized,
                Some(std::time::Duration::from_secs(
                    model_fetch_interval_minutes().saturating_mul(60),
                )),
            )
            .await
        {
            debug!(
                provider_id = %provider_id,
                key_id = %key_id,
                error = %err,
                "gateway model fetch cache write failed"
            );
        }
    }
}

#[async_trait]
impl ModelFetchAssociationStore for AppState {
    type Error = String;

    fn has_global_model_reader(&self) -> bool {
        self.data.has_global_model_reader()
    }

    fn has_global_model_writer(&self) -> bool {
        self.data.has_global_model_writer()
    }

    fn model_fetch_internal_error(&self, message: String) -> Self::Error {
        message
    }

    async fn list_admin_provider_models(
        &self,
        query: &AdminProviderModelListQuery,
    ) -> Result<Vec<StoredAdminProviderModel>, Self::Error> {
        AppState::list_admin_provider_models(self, query)
            .await
            .map_err(|err| format!("{err:?}"))
    }

    async fn list_admin_global_models(
        &self,
        query: &AdminGlobalModelListQuery,
    ) -> Result<StoredAdminGlobalModelPage, Self::Error> {
        AppState::list_admin_global_models(self, query)
            .await
            .map_err(|err| format!("{err:?}"))
    }

    async fn create_admin_provider_model(
        &self,
        record: &UpsertAdminProviderModelRecord,
    ) -> Result<Option<StoredAdminProviderModel>, Self::Error> {
        AppState::create_admin_provider_model(self, record)
            .await
            .map_err(|err| format!("{err:?}"))
    }

    async fn list_provider_catalog_keys_by_provider_ids(
        &self,
        provider_ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogKey>, Self::Error> {
        AppState::list_provider_catalog_keys_by_provider_ids(self, provider_ids)
            .await
            .map_err(|err| format!("{err:?}"))
    }
}

#[async_trait]
impl RequestCandidateRuntimeReader for AppState {
    async fn read_request_candidates_by_request_id(
        &self,
        request_id: &str,
    ) -> Result<Vec<StoredRequestCandidate>, GatewayError> {
        AppState::read_request_candidates_by_request_id(self, request_id).await
    }
}

#[async_trait]
impl RequestCandidateRuntimeCapabilityReader for AppState {
    async fn read_request_candidate_user_model_capability_settings(
        &self,
        user_id: &str,
    ) -> Result<Option<Value>, GatewayError> {
        AppState::read_user_model_capability_settings(self, user_id).await
    }

    async fn read_request_candidate_api_key_force_capabilities(
        &self,
        user_id: &str,
        api_key_id: &str,
    ) -> Result<Option<Value>, GatewayError> {
        AppState::read_auth_api_key_force_capabilities(self, user_id, api_key_id).await
    }
}

#[async_trait]
impl RequestCandidateRuntimeWriter for AppState {
    fn has_request_candidate_data_writer(&self) -> bool {
        AppState::has_request_candidate_data_writer(self)
    }

    async fn upsert_request_candidate(
        &self,
        candidate: UpsertRequestCandidateRecord,
    ) -> Result<Option<StoredRequestCandidate>, GatewayError> {
        AppState::upsert_request_candidate(self, candidate).await
    }

    async fn enqueue_request_candidate_status(
        &self,
        candidate: UpsertRequestCandidateRecord,
    ) -> Result<Option<()>, GatewayError> {
        AppState::enqueue_request_candidate_status(self, candidate).await
    }

    fn try_enqueue_request_candidate_status(
        &self,
        candidate: UpsertRequestCandidateRecord,
    ) -> Result<(), UpsertRequestCandidateRecord> {
        AppState::try_enqueue_request_candidate_status(self, candidate)
    }
}

#[async_trait]
impl SchedulerRuntimeState for AppState {
    async fn read_provider_quota_snapshot(
        &self,
        provider_id: &str,
    ) -> Result<Option<StoredProviderQuotaSnapshot>, GatewayError> {
        AppState::read_provider_quota_snapshot(self, provider_id).await
    }

    async fn read_provider_catalog_providers_by_ids(
        &self,
        provider_ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogProvider>, GatewayError> {
        AppState::read_provider_catalog_providers_by_ids(self, provider_ids).await
    }

    async fn read_provider_catalog_keys_by_ids(
        &self,
        key_ids: &[String],
    ) -> Result<Vec<StoredProviderCatalogKey>, GatewayError> {
        AppState::read_provider_catalog_keys_by_ids(self, key_ids).await
    }

    async fn read_recent_request_candidates(
        &self,
        limit: usize,
    ) -> Result<Vec<StoredRequestCandidate>, GatewayError> {
        AppState::read_recent_runtime_request_candidates(self, limit).await
    }

    fn provider_key_rpm_reset_at(&self, key_id: &str, now_unix_secs: u64) -> Option<u64> {
        AppState::provider_key_rpm_reset_at(self, key_id, now_unix_secs)
    }

    fn read_cached_scheduler_affinity_target(
        &self,
        cache_key: &str,
        ttl: Duration,
    ) -> Option<SchedulerAffinityTarget> {
        AppState::read_scheduler_affinity_target(self, cache_key, ttl)
    }

    fn scheduler_affinity_epoch(&self) -> u64 {
        AppState::scheduler_affinity_epoch(self)
    }

    fn remember_scheduler_affinity_target(
        &self,
        cache_key: &str,
        target: SchedulerAffinityTarget,
        ttl: Duration,
        max_entries: usize,
    ) {
        AppState::remember_scheduler_affinity_target(self, cache_key, target, ttl, max_entries);
    }

    fn remember_scheduler_affinity_target_for_epoch(
        &self,
        cache_key: &str,
        target: SchedulerAffinityTarget,
        ttl: Duration,
        max_entries: usize,
        expected_epoch: Option<u64>,
    ) -> bool {
        AppState::remember_scheduler_affinity_target_for_epoch(
            self,
            cache_key,
            target,
            ttl,
            max_entries,
            expected_epoch,
        )
    }
}