insights/rust/capture/tests/integration_ai_restrictions.rs

#[path = "common/integration_utils.rs"]
mod integration_utils;

use async_trait::async_trait;
use axum::http::StatusCode;
use axum::Router;
use axum_test_helper::TestClient;
use capture::ai_s3::{BlobStorage, MockBlobStorage};
use capture::api::CaptureError;
use capture::config::{AiRouting, CaptureMode};
use capture::event_restrictions::{
    EventRestrictionService, Pipeline, Restriction, RestrictionManager, RestrictionScope,
    RestrictionType,
};
use capture::quota_limiters::CaptureQuotaLimiter;
use capture::router::router;
use capture::sinks::Event;
use capture::team::StaticTeamResolver;
use capture::time::TimeSource;
use capture::v0_request::{DataType, OverflowReason, ProcessedEvent};
use chrono::{DateTime, Utc};
use common_redis::MockRedisClient;
use futures::StreamExt;
use integration_utils::{test_lifecycle_handlers, DEFAULT_CONFIG, DEFAULT_TEST_TIME};
use limiters::overflow::OverflowLimiter;
use limiters::token_dropper::TokenDropper;
use reqwest::multipart::{Form, Part};
use serde_json::{json, Value};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;

const TEST_BLOB_BUCKET: &str = "test-bucket";
const TEST_BLOB_PREFIX: &str = "llma/";

// The fixed team every token resolves to under StaticTeamResolver; AI blobs are
// partitioned by this validated team id instead of hash(token).
const TEST_TEAM_ID: i32 = 1;

fn create_mock_blob_storage() -> Arc<dyn BlobStorage> {
    Arc::new(MockBlobStorage::new(
        TEST_BLOB_BUCKET.to_string(),
        TEST_BLOB_PREFIX.to_string(),
    ))
}

struct FixedTime {
    pub time: DateTime<Utc>,
}

impl TimeSource for FixedTime {
    fn current_time(&self) -> DateTime<Utc> {
        self.time
    }
}

#[derive(Clone)]
struct CapturingSink {
    events: Arc<tokio::sync::Mutex<Vec<ProcessedEvent>>>,
}

impl CapturingSink {
    fn new() -> Self {
        Self {
            events: Arc::new(tokio::sync::Mutex::new(Vec::new())),
        }
    }

    async fn get_events(&self) -> Vec<ProcessedEvent> {
        self.events.lock().await.clone()
    }
}

#[async_trait]
impl Event for CapturingSink {
    async fn send(&self, event: ProcessedEvent) -> Result<(), CaptureError> {
        self.events.lock().await.push(event);
        Ok(())
    }

    async fn send_batch(&self, events: Vec<ProcessedEvent>) -> Result<(), CaptureError> {
        self.events.lock().await.extend(events);
        Ok(())
    }
}

async fn send_multipart_request(
    client: &TestClient,
    form: Form,
    auth_token: Option<&str>,
) -> axum_test_helper::TestResponse {
    let boundary = form.boundary().to_string();
    let content_type = format!("multipart/form-data; boundary={boundary}");

    let mut stream = form.into_stream();
    let mut body = Vec::new();

    while let Some(chunk) = stream.next().await {
        body.extend_from_slice(&chunk.unwrap());
    }

    let mut request = client
        .post("/i/v0/ai")
        .header("Content-Type", content_type)
        .body(body);

    if let Some(token) = auth_token {
        request = request.header("Authorization", format!("Bearer {token}"));
    }

    request.send().await
}

fn create_ai_event_form(event_name: &str, distinct_id: &str, properties: Value) -> Form {
    use uuid::Uuid;

    let event_data = json!({
        "uuid": Uuid::now_v7().to_string(),
        "event": event_name,
        "distinct_id": distinct_id
    });

    Form::new()
        .part(
            "event",
            Part::bytes(serde_json::to_vec(&event_data).unwrap())
                .mime_str("application/json")
                .unwrap(),
        )
        .part(
            "event.properties",
            Part::bytes(serde_json::to_vec(&properties).unwrap())
                .mime_str("application/json")
                .unwrap(),
        )
}

async fn setup_ai_router_with_restriction(
    restriction_type: RestrictionType,
    token: &str,
) -> (Router, CapturingSink) {
    let (readiness, liveness, _monitor) = test_lifecycle_handlers();

    let sink = CapturingSink::new();
    let sink_clone = sink.clone();
    let timesource = FixedTime {
        time: DateTime::parse_from_rfc3339(DEFAULT_TEST_TIME)
            .expect("Invalid fixed time format")
            .with_timezone(&Utc),
    };
    let redis = Arc::new(MockRedisClient::new());

    let mut cfg = DEFAULT_CONFIG.clone();
    cfg.capture_mode = CaptureMode::Events;

    let quota_limiter =
        CaptureQuotaLimiter::new(&cfg, redis.clone(), Duration::from_secs(60 * 60 * 24 * 7));

    let service = EventRestrictionService::new(vec![Pipeline::Ai], Duration::from_secs(300));

    let mut manager = RestrictionManager::new();
    manager.insert_restrictions(
        Pipeline::Ai,
        token,
        vec![Restriction {
            restriction_type,
            scope: RestrictionScope::AllEvents,
            args: None,
        }],
    );
    service.update(manager).await;

    let router = router(
        timesource,
        readiness,
        liveness,
        Arc::new(sink),
        redis,
        None, // global_rate_limiter_token_distinctid
        quota_limiter,
        TokenDropper::default(),
        Some(service),
        None, // recorder_handle
        CaptureMode::Events,
        None,
        25 * 1024 * 1024,
        false,
        1_i64,
        false,
        0.0_f32,
        26_214_400,
        Some(create_mock_blob_storage()),
        None,
        256,                // body_read_chunk_size_kb
        Some(Arc::new(StaticTeamResolver::new(TEST_TEAM_ID))), // team_resolver
        Arc::new(capture::warnings::Warnings::default()),
        10 * 1024 * 1024,   // capture_v1_max_compressed_body_bytes
        50 * 1024 * 1024,   // capture_v1_max_decompressed_body_bytes
        None,               // overflow_limiter
        None,               // ai_events_overflow_limiter
        None,               // replay_overflow_limiter
        None,               // v1_sink_router
        8,                  // capture_v1_scatter_gather_min_batch
        None,               // ai_gateway_signing_secret
        AiRouting::Primary, // ai_routing
        false,              // ai_events_overflow_enabled
        None,               // ingestion_warning_emitter
    );

    (router, sink_clone)
}

struct ExpectedEvent<'a> {
    // CapturedEvent fields
    token: &'a str,
    distinct_id: &'a str,
    event_name: &'a str,
    // ProcessedEventMetadata fields
    data_type: DataType,
    force_overflow: bool,
    skip_person_processing: bool,
    redirect_to_dlq: bool,
    redirect_to_topic: Option<String>,
    // Properties to verify in the event data
    expected_properties: Option<Value>,
}

fn assert_event(event: &ProcessedEvent, expected: &ExpectedEvent) {
    // Assert CapturedEvent fields
    assert_eq!(event.event.token, expected.token, "token mismatch");
    assert_eq!(
        event.event.distinct_id, expected.distinct_id,
        "distinct_id mismatch"
    );
    assert_eq!(
        event.event.event, expected.event_name,
        "event name mismatch"
    );
    assert!(!event.event.ip.is_empty(), "ip should not be empty");
    assert!(!event.event.now.is_empty(), "now should not be empty");
    assert!(!event.event.data.is_empty(), "data should not be empty");

    // Assert ProcessedEventMetadata fields
    assert_eq!(
        event.metadata.data_type, expected.data_type,
        "data_type mismatch"
    );
    assert_eq!(
        event.metadata.event_name, expected.event_name,
        "metadata.event_name mismatch"
    );
    assert_eq!(
        event.metadata.force_overflow, expected.force_overflow,
        "force_overflow mismatch"
    );
    assert_eq!(
        event.metadata.skip_person_processing, expected.skip_person_processing,
        "skip_person_processing mismatch"
    );
    assert_eq!(
        event.metadata.redirect_to_dlq, expected.redirect_to_dlq,
        "redirect_to_dlq mismatch"
    );
    assert_eq!(
        event.metadata.redirect_to_topic, expected.redirect_to_topic,
        "redirect_to_topic mismatch"
    );

    // Assert properties in event data
    if let Some(expected_props) = &expected.expected_properties {
        let data: Value =
            serde_json::from_str(&event.event.data).expect("event.data should be valid JSON");
        let actual_props = data
            .get("properties")
            .expect("event data should have properties");
        for (key, expected_value) in expected_props.as_object().unwrap() {
            let actual_value = actual_props.get(key).unwrap_or(&Value::Null);
            assert_eq!(actual_value, expected_value, "property '{key}' mismatch");
        }
    }
}

#[tokio::test]
async fn test_ai_drop_event_restriction() {
    let restricted_token = "pk-restricted_drop_token";
    let (router, sink) =
        setup_ai_router_with_restriction(RestrictionType::DropEvent, restricted_token).await;
    let test_client = TestClient::new(router);

    let properties = json!({
        "$ai_model": "gpt-4"
    });

    let form = create_ai_event_form("$ai_generation", "test_user", properties);

    let response = send_multipart_request(&test_client, form, Some(restricted_token)).await;

    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert!(
        events.is_empty(),
        "Event should be dropped by restriction, but {} events were published",
        events.len()
    );
}

#[tokio::test]
async fn test_ai_redirect_to_dlq_restriction() {
    let restricted_token = "pk-restricted_dlq_token";
    let (router, sink) =
        setup_ai_router_with_restriction(RestrictionType::RedirectToDlq, restricted_token).await;
    let test_client = TestClient::new(router);

    let properties = json!({
        "$ai_model": "gpt-4"
    });

    let form = create_ai_event_form("$ai_generation", "test_user", properties);

    let response = send_multipart_request(&test_client, form, Some(restricted_token)).await;

    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert_eq!(events.len(), 1);
    assert_event(
        &events[0],
        &ExpectedEvent {
            token: restricted_token,
            distinct_id: "test_user",
            event_name: "$ai_generation",
            data_type: DataType::AnalyticsMain,
            force_overflow: false,
            skip_person_processing: false,
            redirect_to_dlq: true,
            redirect_to_topic: None,
            expected_properties: Some(json!({
                "$ai_model": "gpt-4"
            })),
        },
    );
}

#[tokio::test]
async fn test_ai_force_overflow_restriction() {
    let restricted_token = "pk-restricted_overflow_token";
    let (router, sink) =
        setup_ai_router_with_restriction(RestrictionType::ForceOverflow, restricted_token).await;
    let test_client = TestClient::new(router);

    let properties = json!({
        "$ai_model": "gpt-4"
    });

    let form = create_ai_event_form("$ai_generation", "test_user", properties);

    let response = send_multipart_request(&test_client, form, Some(restricted_token)).await;

    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert_eq!(events.len(), 1);
    assert_event(
        &events[0],
        &ExpectedEvent {
            token: restricted_token,
            distinct_id: "test_user",
            event_name: "$ai_generation",
            data_type: DataType::AnalyticsMain,
            force_overflow: true,
            skip_person_processing: false,
            redirect_to_dlq: false,
            redirect_to_topic: None,
            expected_properties: Some(json!({
                "$ai_model": "gpt-4"
            })),
        },
    );
}

#[tokio::test]
async fn test_ai_skip_person_processing_restriction() {
    let restricted_token = "pk-restricted_skip_person_token";
    let (router, sink) =
        setup_ai_router_with_restriction(RestrictionType::SkipPersonProcessing, restricted_token)
            .await;
    let test_client = TestClient::new(router);

    let properties = json!({
        "$ai_model": "gpt-4"
    });

    let form = create_ai_event_form("$ai_generation", "test_user", properties);

    let response = send_multipart_request(&test_client, form, Some(restricted_token)).await;

    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert_eq!(events.len(), 1);
    assert_event(
        &events[0],
        &ExpectedEvent {
            token: restricted_token,
            distinct_id: "test_user",
            event_name: "$ai_generation",
            data_type: DataType::AnalyticsMain,
            force_overflow: false,
            skip_person_processing: true,
            redirect_to_dlq: false,
            redirect_to_topic: None,
            expected_properties: Some(json!({
                "$ai_model": "gpt-4"
            })),
        },
    );
}

#[tokio::test]
async fn test_ai_restriction_does_not_apply_to_other_tokens() {
    let restricted_token = "pk-restricted_token";
    let (router, sink) =
        setup_ai_router_with_restriction(RestrictionType::DropEvent, restricted_token).await;
    let test_client = TestClient::new(router);

    let properties = json!({
        "$ai_model": "gpt-4"
    });

    let form = create_ai_event_form("$ai_generation", "test_user", properties);

    let response =
        send_multipart_request(&test_client, form, Some("pk-not_restricted_token")).await;

    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert_eq!(
        events.len(),
        1,
        "Event should be published for non-restricted token"
    );
    assert_event(
        &events[0],
        &ExpectedEvent {
            token: "pk-not_restricted_token",
            distinct_id: "test_user",
            event_name: "$ai_generation",
            data_type: DataType::AnalyticsMain,
            force_overflow: false,
            skip_person_processing: false,
            redirect_to_dlq: false,
            redirect_to_topic: None,
            expected_properties: Some(json!({
                "$ai_model": "gpt-4"
            })),
        },
    );
}

async fn setup_ai_router_with_redirect_to_topic(
    token: &str,
    topic: &str,
) -> (Router, CapturingSink) {
    let (readiness, liveness, _monitor) = test_lifecycle_handlers();

    let sink = CapturingSink::new();
    let sink_clone = sink.clone();
    let timesource = FixedTime {
        time: DateTime::parse_from_rfc3339(DEFAULT_TEST_TIME)
            .expect("Invalid fixed time format")
            .with_timezone(&Utc),
    };
    let redis = Arc::new(MockRedisClient::new());

    let mut cfg = DEFAULT_CONFIG.clone();
    cfg.capture_mode = CaptureMode::Events;

    let quota_limiter =
        CaptureQuotaLimiter::new(&cfg, redis.clone(), Duration::from_secs(60 * 60 * 24 * 7));

    let service = EventRestrictionService::new(vec![Pipeline::Ai], Duration::from_secs(300));

    let mut manager = RestrictionManager::new();
    manager.insert_restrictions(
        Pipeline::Ai,
        token,
        vec![Restriction {
            restriction_type: RestrictionType::RedirectToTopic,
            scope: RestrictionScope::AllEvents,
            args: Some(json!({"topic": topic})),
        }],
    );
    service.update(manager).await;

    let router = router(
        timesource,
        readiness,
        liveness,
        Arc::new(sink),
        redis,
        None, // global_rate_limiter_token_distinctid
        quota_limiter,
        TokenDropper::default(),
        Some(service),
        None, // recorder_handle
        CaptureMode::Events,
        None,
        25 * 1024 * 1024,
        false,
        1_i64,
        false,
        0.0_f32,
        26_214_400,
        Some(create_mock_blob_storage()),
        None,
        256,                // body_read_chunk_size_kb
        10 * 1024 * 1024,   // capture_v1_max_compressed_body_bytes
        50 * 1024 * 1024,   // capture_v1_max_decompressed_body_bytes
        None,               // overflow_limiter
        None,               // ai_events_overflow_limiter
        None,               // replay_overflow_limiter
        None,               // v1_sink_router
        8,                  // capture_v1_scatter_gather_min_batch
        None,               // ai_gateway_signing_secret
        AiRouting::Primary, // ai_routing
        false,              // ai_events_overflow_enabled
        None,               // ingestion_warning_emitter
    );

    (router, sink_clone)
}

/// Sets up an AI router that has BOTH a `ForceOverflow` event restriction AND
/// a configured `OverflowLimiter`. Used to verify the interaction ordering:
/// `force_overflow` must short-circuit the limiter so
/// `ProcessedEventMetadata::overflow_reason` stays `None`. This mirrors the
/// analytics `force_overflow` path in `events/analytics.rs` — same semantics,
/// different entry point.
async fn setup_ai_router_with_force_overflow_and_limiter(
    token: &str,
    overflow_limiter: Arc<OverflowLimiter>,
    ai_routing: AiRouting,
    ai_events_overflow_enabled: bool,
) -> (Router, CapturingSink) {
    let (readiness, liveness, _monitor) = test_lifecycle_handlers();

    let sink = CapturingSink::new();
    let sink_clone = sink.clone();
    let timesource = FixedTime {
        time: DateTime::parse_from_rfc3339(DEFAULT_TEST_TIME)
            .expect("Invalid fixed time format")
            .with_timezone(&Utc),
    };
    let redis = Arc::new(MockRedisClient::new());

    let mut cfg = DEFAULT_CONFIG.clone();
    cfg.capture_mode = CaptureMode::Events;

    let quota_limiter =
        CaptureQuotaLimiter::new(&cfg, redis.clone(), Duration::from_secs(60 * 60 * 24 * 7));

    let service = EventRestrictionService::new(vec![Pipeline::Ai], Duration::from_secs(300));

    let mut manager = RestrictionManager::new();
    manager.insert_restrictions(
        Pipeline::Ai,
        token,
        vec![Restriction {
            restriction_type: RestrictionType::ForceOverflow,
            scope: RestrictionScope::AllEvents,
            args: None,
        }],
    );
    service.update(manager).await;

    let router = router(
        timesource,
        readiness,
        liveness,
        Arc::new(sink),
        redis,
        None,
        quota_limiter,
        TokenDropper::default(),
        Some(service),
        None, // recorder_handle
        CaptureMode::Events,
        None,
        25 * 1024 * 1024,
        false,
        1_i64,
        false,
        0.0_f32,
        26_214_400,
        Some(create_mock_blob_storage()),
        None,
        256,
        10 * 1024 * 1024,       // capture_v1_max_compressed_body_bytes
        50 * 1024 * 1024,       // capture_v1_max_decompressed_body_bytes
        Some(overflow_limiter), // overflow_limiter
        None,                   // ai_events_overflow_limiter
        None,                   // replay_overflow_limiter
        None,                   // v1_sink_router
        8,                      // capture_v1_scatter_gather_min_batch
        None,                   // ai_gateway_signing_secret
        ai_routing,
        ai_events_overflow_enabled,
        None, // ingestion_warning_emitter
    );

    (router, sink_clone)
}

#[tokio::test]
async fn test_ai_force_overflow_restriction_wins_over_overflow_limiter() {
    // The OverflowLimiter is configured to ForceLimited this exact key, so if
    // the short-circuit fails we'd see `overflow_reason = ForceLimited` AND
    // `skip_person_processing = true`. The restriction-driven code path must
    // fire first, leaving the limiter untouched: force_overflow = true,
    // overflow_reason = None, skip_person_processing stays false (because the
    // restriction only sets force_overflow, not skip_person_processing).
    let restricted_token = "pk-force_overflow_with_limiter_token";
    let distinct_id = "test_user";
    let hot_key = format!("{restricted_token}:{distinct_id}");

    let overflow_limiter = Arc::new(OverflowLimiter::new(
        NonZeroU32::new(1_000).unwrap(),
        NonZeroU32::new(1_000).unwrap(),
        Some(hot_key),
        true, // preserve_locality
    ));

    let (router, sink) = setup_ai_router_with_force_overflow_and_limiter(
        restricted_token,
        overflow_limiter,
        AiRouting::Primary,
        false,
    )
    .await;
    let test_client = TestClient::new(router);

    let properties = json!({"$ai_model": "gpt-4"});
    let form = create_ai_event_form("$ai_generation", distinct_id, properties);

    let response = send_multipart_request(&test_client, form, Some(restricted_token)).await;
    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert_eq!(events.len(), 1);

    assert_event(
        &events[0],
        &ExpectedEvent {
            token: restricted_token,
            distinct_id,
            event_name: "$ai_generation",
            data_type: DataType::AnalyticsMain,
            force_overflow: true,
            skip_person_processing: false,
            redirect_to_dlq: false,
            redirect_to_topic: None,
            expected_properties: Some(json!({"$ai_model": "gpt-4"})),
        },
    );

    assert_eq!(
        events[0].metadata.overflow_reason, None,
        "force_overflow must short-circuit the limiter; overflow_reason stays None so the \
         sink routes via the event_restriction branch (not the limiter branches)"
    );
}

#[tokio::test]
async fn test_ai_endpoint_ignores_analytics_ai_routing_config() {
    // Deployment-config cross-contamination guard: the CAPTURE_ANALYTICS_AI_EVENTS_*
    // family configures $ai_* routing on capture-analytics deployments. If it
    // leaks onto a capture-ai deployment (whose main topic already IS the AI
    // topic), the AI endpoint must behave exactly as without it: events stay
    // on the AnalyticsMain lane and the OverflowLimiter applies by analytics
    // rules. Catches a regression that types AI-endpoint events as
    // DataType::AiEvents or gates their overflow on the analytics valve.
    let token = "pk-ai_endpoint_routing_leak_token";
    let distinct_id = "test_user";
    let hot_key = format!("{token}:{distinct_id}");

    let overflow_limiter = Arc::new(OverflowLimiter::new(
        NonZeroU32::new(1_000).unwrap(),
        NonZeroU32::new(1_000).unwrap(),
        Some(hot_key),
        false, // preserve_locality
    ));

    // Restriction is registered for a different token, so only the limiter acts.
    let (router, sink) = setup_ai_router_with_force_overflow_and_limiter(
        "pk-some_other_token",
        overflow_limiter,
        AiRouting::Secondary,
        true,
    )
    .await;
    let test_client = TestClient::new(router);

    let properties = json!({"$ai_model": "gpt-4"});
    let form = create_ai_event_form("$ai_generation", distinct_id, properties);

    let response = send_multipart_request(&test_client, form, Some(token)).await;
    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert_eq!(events.len(), 1);
    assert_eq!(
        events[0].metadata.data_type,
        DataType::AnalyticsMain,
        "analytics-deployment routing config must not divert AI-endpoint events"
    );
    assert_eq!(
        events[0].metadata.overflow_reason,
        Some(OverflowReason::ForceLimited),
        "the limiter must keep applying by AnalyticsMain rules, valve or not"
    );
}

#[tokio::test]
async fn test_ai_redirect_to_topic_restriction() {
    let restricted_token = "pk-restricted_redirect_topic_token";
    let target_topic = "custom_ai_topic";
    let (router, sink) =
        setup_ai_router_with_redirect_to_topic(restricted_token, target_topic).await;
    let test_client = TestClient::new(router);

    let properties = json!({
        "$ai_model": "gpt-4"
    });

    let form = create_ai_event_form("$ai_generation", "test_user", properties);

    let response = send_multipart_request(&test_client, form, Some(restricted_token)).await;

    assert_eq!(response.status(), StatusCode::OK);

    let events = sink.get_events().await;
    assert_eq!(events.len(), 1);
    assert_event(
        &events[0],
        &ExpectedEvent {
            token: restricted_token,
            distinct_id: "test_user",
            event_name: "$ai_generation",
            data_type: DataType::AnalyticsMain,
            force_overflow: false,
            skip_person_processing: false,
            redirect_to_dlq: false,
            redirect_to_topic: Some(target_topic.to_string()),
            expected_properties: Some(json!({
                "$ai_model": "gpt-4"
            })),
        },
    );
}