mirror of
https://github.com/hanzoai/insights.git
synced 2026-10-30 16:15:31 +00:00
zeekay
19fdf55ceb
A key now opens with what it is allowed to touch. `pk-` is publishable and safe in a browser; `sk-` is secret and never leaves a server. The five marks it replaces all began `ph` -- the upstream's initials -- and they were the last of its brand a customer read back to us, in every SDK snippet. No word rule could reach them, because `ph` is not a word. The marks were also not telling the truth. The constant said `phc_` while every live project token said `hi_`, so the middleware that routes on the constant matched no real token at all. Two spellings existed and one was fiction. Minting now happens in exactly one function and reading in exactly one, so a second spelling has nowhere to live. `phx_` (personal) and `phs_` (project secret) both become `sk-`. The mark now says only that a value is secret; which store holds it is settled by asking. That was already how personal keys worked -- keys minted before the marks carry none, so the store always decided -- and both secret resolvers already returned None on a miss and fell through, so one shared extractor replaces the two that could drift apart. OAuth's access and refresh tokens stay distinguishable as `at-` and `rt-`, because the parser tells them apart before either is looked up, and the domain-verification token becomes `vt-` rather than keeping a sixth mint of its own. `bin/debrand --check` now refuses the marks as a shape, so a seventh cannot be coined. Two rules that were about to rename things that are not ours to spell are bounded: `Lemon` no longer eats `LemonSqueezy`, and the `hog_flow` rules no longer reach a Postgres table, a Datastore table or a route. Keys are stored hashed, so a mark cannot be rewritten in place: the three live project tokens have to be reminted.
3483 lines
111 KiB
Rust
3483 lines
111 KiB
Rust
#[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::quota_limiters::CaptureQuotaLimiter;
|
|
use capture::router::router;
|
|
use capture::sinks::Event;
|
|
use capture::team::StaticTeamResolver;
|
|
use capture::time::TimeSource;
|
|
use capture::v0_request::{OverflowReason, ProcessedEvent};
|
|
use chrono::{DateTime, TimeZone, Utc};
|
|
use common_ingestion_warnings::test_support::CollectingEmitter;
|
|
use common_ingestion_warnings::{WarningEmitter, WarningType, CAPTURE_AI_EVENTS};
|
|
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::io::Write;
|
|
use std::num::NonZeroU32;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
// Test constants for mock blob storage
|
|
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
|
|
// now 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(),
|
|
))
|
|
}
|
|
|
|
// Fixed time source for tests
|
|
struct FixedTime {
|
|
pub time: DateTime<Utc>,
|
|
}
|
|
|
|
impl TimeSource for FixedTime {
|
|
fn current_time(&self) -> DateTime<Utc> {
|
|
self.time
|
|
}
|
|
}
|
|
|
|
// Simple memory sink for tests
|
|
#[derive(Clone, Default)]
|
|
struct TestSink;
|
|
|
|
#[async_trait]
|
|
impl Event for TestSink {
|
|
async fn send(&self, _event: ProcessedEvent) -> Result<(), CaptureError> {
|
|
Ok(())
|
|
}
|
|
|
|
async fn send_batch(&self, _events: Vec<ProcessedEvent>) -> Result<(), CaptureError> {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
// Capturing sink for Kafka tests - stores events in memory
|
|
#[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(())
|
|
}
|
|
}
|
|
|
|
// Helper to build multipart form and send request
|
|
async fn send_multipart_request(
|
|
client: &TestClient,
|
|
form: Form,
|
|
auth_token: Option<&str>,
|
|
) -> axum_test_helper::TestResponse {
|
|
// Get the boundary from the form
|
|
let boundary = form.boundary().to_string();
|
|
let content_type = format!("multipart/form-data; boundary={boundary}");
|
|
|
|
// Use into_stream() to get the body bytes
|
|
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
|
|
}
|
|
|
|
/// Helper to create a basic AI event with properties in separate parts
|
|
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(),
|
|
)
|
|
}
|
|
|
|
// Helper to setup test router
|
|
fn setup_ai_test_router() -> Router {
|
|
let (readiness, liveness, _monitor) = test_lifecycle_handlers();
|
|
|
|
let sink = TestSink;
|
|
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));
|
|
|
|
router(
|
|
timesource,
|
|
readiness,
|
|
liveness,
|
|
Arc::new(sink),
|
|
redis,
|
|
None,
|
|
quota_limiter,
|
|
TokenDropper::default(),
|
|
None, // event_restriction_service
|
|
None, // recorder_handle
|
|
CaptureMode::Events,
|
|
None,
|
|
25 * 1024 * 1024,
|
|
false,
|
|
1_i64,
|
|
false,
|
|
0.0_f32,
|
|
26_214_400, // 25MB default for AI endpoint
|
|
Some(create_mock_blob_storage()), // ai_blob_storage
|
|
None, // body_chunk_read_timeout_ms
|
|
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
|
|
)
|
|
}
|
|
|
|
/// A router whose emit sites are live, plus the emitter to assert against.
|
|
/// Mirrors `setup_ai_test_router` and differs only in the last argument.
|
|
fn setup_ai_router_collecting_warnings() -> (Router, Arc<CollectingEmitter>) {
|
|
let (readiness, liveness, _monitor) = test_lifecycle_handlers();
|
|
|
|
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 emitter = Arc::new(CollectingEmitter::default());
|
|
let warning_emitter: Option<Arc<dyn WarningEmitter>> = Some(emitter.clone());
|
|
|
|
let app = router(
|
|
timesource,
|
|
readiness,
|
|
liveness,
|
|
Arc::new(TestSink),
|
|
redis,
|
|
None,
|
|
quota_limiter,
|
|
TokenDropper::default(),
|
|
None, // event_restriction_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,
|
|
50 * 1024 * 1024,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
8,
|
|
None,
|
|
AiRouting::Primary,
|
|
false,
|
|
warning_emitter,
|
|
);
|
|
|
|
(app, emitter)
|
|
}
|
|
|
|
// ============================================================================
|
|
// PHASE 1: HTTP ENDPOINT
|
|
// ============================================================================
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.1: HTTP Method Validation
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_get_returns_405() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let response = test_client
|
|
.get("/i/v0/ai")
|
|
.header("Authorization", "Bearer test_token")
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_put_returns_405() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let response = test_client
|
|
.put("/i/v0/ai")
|
|
.header("Authorization", "Bearer test_token")
|
|
.body("test")
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_delete_returns_405() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let response = test_client
|
|
.delete("/i/v0/ai")
|
|
.header("Authorization", "Bearer test_token")
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.1: Authentication Validation
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_no_auth_returns_401() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties);
|
|
|
|
let response = send_multipart_request(&test_client, form, None).await;
|
|
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.1: Content Type Validation
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_wrong_content_type_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let response = test_client
|
|
.post("/i/v0/ai")
|
|
.header("Content-Type", "application/json")
|
|
.header(
|
|
"Authorization",
|
|
"Bearer pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3",
|
|
)
|
|
.json(&event_data)
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_empty_body_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let response = test_client
|
|
.post("/i/v0/ai")
|
|
.header("Content-Type", "multipart/form-data; boundary=test")
|
|
.header(
|
|
"Authorization",
|
|
"Bearer pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3",
|
|
)
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.2: Multipart Parsing
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_multipart_parsing_with_multiple_blobs() {
|
|
use uuid::Uuid;
|
|
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"uuid": Uuid::now_v7().to_string(),
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test-multi-blob"
|
|
});
|
|
|
|
let input_blob = json!({"messages": [{"role": "user", "content": "Hello"}]});
|
|
let output_blob = json!({"choices": [{"message": {"content": "Hi there"}}]});
|
|
let metadata_blob = json!({"model_version": "1.0", "temperature": 0.7});
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_input",
|
|
Part::bytes(serde_json::to_vec(&input_blob).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_output",
|
|
Part::bytes(serde_json::to_vec(&output_blob).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_metadata",
|
|
Part::bytes(serde_json::to_vec(&metadata_blob).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let response_json: serde_json::Value = response.json::<serde_json::Value>().await;
|
|
let accepted_parts = response_json["accepted_parts"].as_array().unwrap();
|
|
assert_eq!(accepted_parts.len(), 5);
|
|
assert_eq!(accepted_parts[0]["name"], "event");
|
|
assert_eq!(accepted_parts[0]["length"].as_u64().unwrap(), 98); // UUID adds ~46 bytes
|
|
assert_eq!(accepted_parts[1]["name"], "event.properties");
|
|
assert_eq!(accepted_parts[1]["length"].as_u64().unwrap(), 31);
|
|
assert_eq!(accepted_parts[2]["name"], "event.properties.$ai_input");
|
|
assert_eq!(accepted_parts[2]["length"].as_u64().unwrap(), 48);
|
|
assert_eq!(accepted_parts[3]["name"], "event.properties.$ai_output");
|
|
assert_eq!(accepted_parts[3]["length"].as_u64().unwrap(), 48);
|
|
assert_eq!(accepted_parts[4]["name"], "event.properties.$ai_metadata");
|
|
assert_eq!(accepted_parts[4]["length"].as_u64().unwrap(), 41);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multipart_parsing_with_mixed_content_types() {
|
|
use uuid::Uuid;
|
|
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"uuid": Uuid::now_v7().to_string(),
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test-mixed-types"
|
|
});
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_json_blob",
|
|
Part::bytes(serde_json::to_vec(&json!({"type": "json"})).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_text_blob",
|
|
Part::bytes(b"This is plain text content".to_vec())
|
|
.mime_str("text/plain")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_binary_blob",
|
|
Part::bytes(vec![0x00, 0x01, 0x02, 0x03, 0x04, 0x05])
|
|
.mime_str("application/octet-stream")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let response_json: serde_json::Value = response.json::<serde_json::Value>().await;
|
|
let accepted_parts = response_json["accepted_parts"].as_array().unwrap();
|
|
assert_eq!(accepted_parts.len(), 5);
|
|
assert_eq!(accepted_parts[0]["length"].as_u64().unwrap(), 98); // UUID adds ~46 bytes
|
|
assert_eq!(accepted_parts[1]["length"].as_u64().unwrap(), 32);
|
|
assert_eq!(accepted_parts[2]["content-type"], "application/json");
|
|
assert_eq!(accepted_parts[2]["length"].as_u64().unwrap(), 15);
|
|
assert_eq!(accepted_parts[3]["content-type"], "text/plain");
|
|
assert_eq!(accepted_parts[3]["length"].as_u64().unwrap(), 26);
|
|
assert_eq!(
|
|
accepted_parts[4]["content-type"],
|
|
"application/octet-stream"
|
|
);
|
|
assert_eq!(accepted_parts[4]["length"].as_u64().unwrap(), 6);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multipart_parsing_with_large_blob() {
|
|
use uuid::Uuid;
|
|
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"uuid": Uuid::now_v7().to_string(),
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test-large"
|
|
});
|
|
|
|
// Create a large JSON blob (100KB)
|
|
let large_blob = json!({
|
|
"messages": (0..100).map(|_| json!({"role": "user", "content": "x".repeat(1000)})).collect::<Vec<_>>()
|
|
});
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_large_input",
|
|
Part::bytes(serde_json::to_vec(&large_blob).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let response_json: serde_json::Value = response.json::<serde_json::Value>().await;
|
|
let accepted_parts = response_json["accepted_parts"].as_array().unwrap();
|
|
assert_eq!(accepted_parts.len(), 3);
|
|
assert_eq!(accepted_parts[0]["length"].as_u64().unwrap(), 98); // UUID adds ~46 bytes
|
|
assert_eq!(accepted_parts[1]["length"].as_u64().unwrap(), 26);
|
|
assert_eq!(accepted_parts[2]["length"].as_u64().unwrap(), 102914);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multipart_parsing_with_empty_blob() {
|
|
use uuid::Uuid;
|
|
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"uuid": Uuid::now_v7().to_string(),
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test-empty"
|
|
});
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_empty",
|
|
Part::bytes(Vec::new())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
|
|
// Empty blobs should be rejected with 400 Bad Request
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.3: Boundary Validation
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_multipart_missing_boundary_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-missing-boundary"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
// Use into_stream() to get the body bytes
|
|
let mut stream = form.into_stream();
|
|
let mut body = Vec::new();
|
|
|
|
while let Some(chunk) = stream.next().await {
|
|
body.extend_from_slice(&chunk.unwrap());
|
|
}
|
|
|
|
// Send with missing boundary parameter in Content-Type
|
|
let response = test_client
|
|
.post("/i/v0/ai")
|
|
.header("Content-Type", "multipart/form-data")
|
|
.header(
|
|
"Authorization",
|
|
"Bearer pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3",
|
|
)
|
|
.body(body)
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multipart_corrupted_boundary_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-corrupted-boundary"
|
|
}
|
|
});
|
|
|
|
// Manually construct a request with corrupted boundary
|
|
let body = format!(
|
|
"------WebKitFormBoundary1234567890abcdef\r\nContent-Disposition: form-data; name=\"event\"; filename=\"event.json\"\r\nContent-Type: application/json\r\n\r\n{}\r\n------WebKitFormBoundary1234567890abcdef--\r\n",
|
|
serde_json::to_string(&event_data).unwrap()
|
|
);
|
|
|
|
let response = test_client
|
|
.post("/i/v0/ai")
|
|
.header(
|
|
"Content-Type",
|
|
"multipart/form-data; boundary=corrupted------WebKitFormBoundary1234567890abcdef",
|
|
)
|
|
.header(
|
|
"Authorization",
|
|
"Bearer pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3",
|
|
)
|
|
.body(body)
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.4: Event Processing Verification
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_multipart_event_not_first_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-event-not-first"
|
|
}
|
|
});
|
|
|
|
// Create multipart data with blob part first, then event part
|
|
let form = Form::new()
|
|
.part(
|
|
"event.properties.$ai_input",
|
|
Part::bytes(
|
|
serde_json::to_vec(&json!({"messages": [{"role": "user", "content": "test"}]}))
|
|
.unwrap(),
|
|
)
|
|
.file_name("input.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.5: Basic Validation
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_invalid_event_name_not_ai_prefix_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "invalid_event_name",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-invalid-name"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_invalid_event_name_regular_event_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$pageview",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"page": "/test"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_invalid_event_name_custom_event_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "button_clicked",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"button": "submit"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_all_allowed_ai_event_types_accepted() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let allowed_events = vec![
|
|
"$ai_generation",
|
|
"$ai_trace",
|
|
"$ai_span",
|
|
"$ai_embedding",
|
|
"$ai_metric",
|
|
"$ai_feedback",
|
|
];
|
|
|
|
for event_name in allowed_events {
|
|
let properties = json!({
|
|
"$ai_model": "test-model"
|
|
});
|
|
|
|
let form = create_ai_event_form(event_name, "test_user", properties);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(
|
|
response.status(),
|
|
StatusCode::OK,
|
|
"Event type {event_name} should be accepted"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_invalid_ai_event_type_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let invalid_events = vec!["$ai_unknown", "$ai_custom", "$ai_"];
|
|
|
|
for event_name in invalid_events {
|
|
let properties = json!({
|
|
"$ai_model": "test-model"
|
|
});
|
|
|
|
let form = create_ai_event_form(event_name, "test_user", properties);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(
|
|
response.status(),
|
|
StatusCode::BAD_REQUEST,
|
|
"Event type {event_name} should be rejected"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_missing_required_ai_properties_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"custom_property": "test_value"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_empty_event_name_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-empty-name"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_missing_distinct_id_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"properties": {
|
|
"$ai_model": "test-missing-distinct-id"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.file_name("event.json")
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_missing_uuid_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
// Event without UUID field
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-missing-uuid"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
|
|
let body = response.text().await;
|
|
assert!(
|
|
body.contains("UUID"),
|
|
"Error should mention UUID, got: {body}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_invalid_uuid_format_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let invalid_uuids = [
|
|
"not-a-uuid",
|
|
"12345",
|
|
"xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
|
|
"",
|
|
"550e8400-e29b-41d4-a716", // Truncated
|
|
];
|
|
|
|
for invalid_uuid in invalid_uuids {
|
|
let event_data = json!({
|
|
"uuid": invalid_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-invalid-uuid"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(
|
|
response.status(),
|
|
StatusCode::BAD_REQUEST,
|
|
"UUID '{invalid_uuid}' should be rejected"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_malicious_uuid_with_path_traversal_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
// These malicious values could be used for S3 path traversal if not validated
|
|
let malicious_uuids = [
|
|
"../../../etc/passwd",
|
|
"550e8400/../../../secret",
|
|
"uuid/with/slashes",
|
|
"..%2F..%2F..%2Fetc%2Fpasswd", // URL-encoded traversal
|
|
"550e8400-e29b-41d4-a716-446655440000/../other",
|
|
"\0null-byte-injection",
|
|
"uuid\nwith\nnewlines",
|
|
];
|
|
|
|
for malicious_uuid in malicious_uuids {
|
|
let event_data = json!({
|
|
"uuid": malicious_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test-malicious-uuid"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(
|
|
response.status(),
|
|
StatusCode::BAD_REQUEST,
|
|
"Malicious UUID '{malicious_uuid}' should be rejected"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_returns_200_for_valid_request() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let response_json: serde_json::Value = response.json::<serde_json::Value>().await;
|
|
assert!(response_json["accepted_parts"].is_array());
|
|
let accepted_parts = response_json["accepted_parts"].as_array().unwrap();
|
|
assert_eq!(accepted_parts.len(), 2);
|
|
assert_eq!(accepted_parts[0]["name"], "event");
|
|
assert_eq!(accepted_parts[0]["content-type"], "application/json");
|
|
assert_eq!(accepted_parts[0]["length"].as_u64().unwrap(), 98); // UUID adds ~46 bytes
|
|
assert_eq!(accepted_parts[1]["name"], "event.properties");
|
|
assert_eq!(accepted_parts[1]["content-type"], "application/json");
|
|
assert_eq!(accepted_parts[1]["length"].as_u64().unwrap(), 20);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Properties Handling Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_properties_in_event_part_only() {
|
|
use uuid::Uuid;
|
|
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"uuid": Uuid::now_v7().to_string(),
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "embedded-model",
|
|
"custom_field": "embedded-value"
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let response_json: serde_json::Value = response.json::<serde_json::Value>().await;
|
|
let accepted_parts = response_json["accepted_parts"].as_array().unwrap();
|
|
assert_eq!(accepted_parts.len(), 1);
|
|
assert_eq!(accepted_parts[0]["name"], "event");
|
|
assert_eq!(accepted_parts[0]["length"].as_u64().unwrap(), 174); // UUID adds ~46 bytes (128 + 46)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_properties_in_separate_part_only() {
|
|
use uuid::Uuid;
|
|
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"uuid": Uuid::now_v7().to_string(),
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "separate-model",
|
|
"custom_field": "separate-value"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let response_json: serde_json::Value = response.json::<serde_json::Value>().await;
|
|
let accepted_parts = response_json["accepted_parts"].as_array().unwrap();
|
|
assert_eq!(accepted_parts.len(), 2);
|
|
assert_eq!(accepted_parts[0]["name"], "event");
|
|
assert_eq!(accepted_parts[0]["length"].as_u64().unwrap(), 98); // UUID adds ~46 bytes
|
|
assert_eq!(accepted_parts[1]["name"], "event.properties");
|
|
assert_eq!(accepted_parts[1]["length"].as_u64().unwrap(), 62);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_properties_both_embedded_and_separate_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "embedded-model",
|
|
"custom_field": "embedded-value"
|
|
}
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "override-model",
|
|
"custom_field": "override-value"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Size Limit Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_event_exceeds_32kb_returns_413() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
// Create event with large properties that exceed 32KB
|
|
let large_value = "x".repeat(33 * 1024);
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"properties": {
|
|
"$ai_model": "test",
|
|
"large_field": large_value
|
|
}
|
|
});
|
|
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(serde_json::to_vec(&event_data).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_combined_event_properties_exceeds_960kb_returns_413() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
// Create event and properties that together exceed 960KB
|
|
// 961KB to ensure we exceed the limit
|
|
let large_value = "x".repeat(961 * 1024);
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test",
|
|
"large_field": large_value
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_sum_of_all_parts_exceeds_25mb_returns_413() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
// Create a blob that, combined with event and properties, exceeds 25MB
|
|
// Event ~52 bytes + properties ~24 bytes + blob ~25MB = exceeds limit
|
|
let large_blob = vec![0u8; 25 * 1024 * 1024];
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_input",
|
|
Part::bytes(large_blob)
|
|
.mime_str("application/octet-stream")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_request_body_exceeds_110_percent_limit_returns_413() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_data = json!({
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
// Create a blob that's just barely over 25MB
|
|
// This will result in a request body that exceeds 27.5MB (110% limit)
|
|
// once multipart overhead is added
|
|
let large_blob = vec![0u8; (25 * 1024 * 1024) + (3 * 1024 * 1024)]; // 28MB
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_input",
|
|
Part::bytes(large_blob)
|
|
.mime_str("application/octet-stream")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
|
|
// Axum's DefaultBodyLimit returns 413 Payload Too Large when the body exceeds the limit
|
|
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Content Type Validation Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_blob_with_application_octet_stream_content_type() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
let blob_data = vec![0u8, 1u8, 2u8, 3u8];
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties).part(
|
|
"event.properties.$ai_binary_data",
|
|
Part::bytes(blob_data)
|
|
.mime_str("application/octet-stream")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_blob_with_application_json_content_type() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
let blob_data = json!({"key": "value"}).to_string();
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties).part(
|
|
"event.properties.$ai_input",
|
|
Part::bytes(blob_data.into_bytes())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_blob_with_text_plain_content_type() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
let blob_data = "This is plain text data";
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties).part(
|
|
"event.properties.$ai_output",
|
|
Part::bytes(blob_data.as_bytes().to_vec())
|
|
.mime_str("text/plain")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_blob_with_invalid_content_type_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
let blob_data = vec![0u8, 1u8, 2u8, 3u8];
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties).part(
|
|
"event.properties.$ai_data",
|
|
Part::bytes(blob_data).mime_str("application/xml").unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_blob_without_content_type_returns_400() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test"
|
|
});
|
|
|
|
let blob_data = vec![0u8, 1u8, 2u8, 3u8];
|
|
|
|
// Create a part without Content-Type (reqwest doesn't set it if we don't call mime_str)
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties)
|
|
.part("event.properties.$ai_data", Part::bytes(blob_data));
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Compression Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_gzip_compressed_request() {
|
|
let router = setup_ai_test_router();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "test-gzip"
|
|
});
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties);
|
|
|
|
// Get the multipart body
|
|
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());
|
|
}
|
|
|
|
// Compress the body with gzip
|
|
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
|
|
encoder.write_all(&body).unwrap();
|
|
let compressed_body = encoder.finish().unwrap();
|
|
|
|
// Send compressed request
|
|
let response = test_client
|
|
.post("/i/v0/ai")
|
|
.header("Content-Type", content_type)
|
|
.header("Content-Encoding", "gzip")
|
|
.header(
|
|
"Authorization",
|
|
"Bearer pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3",
|
|
)
|
|
.body(compressed_body)
|
|
.send()
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify response (lengths should match uncompressed data)
|
|
let response_json: serde_json::Value = response.json::<serde_json::Value>().await;
|
|
let accepted_parts = response_json["accepted_parts"].as_array().unwrap();
|
|
assert_eq!(accepted_parts.len(), 2);
|
|
assert_eq!(accepted_parts[0]["name"], "event");
|
|
assert_eq!(accepted_parts[0]["length"].as_u64().unwrap(), 98); // UUID adds ~46 bytes
|
|
assert_eq!(accepted_parts[1]["name"], "event.properties");
|
|
assert_eq!(accepted_parts[1]["length"].as_u64().unwrap(), 25);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Scenario 1.6: Kafka Publishing and S3 Placeholders
|
|
// ----------------------------------------------------------------------------
|
|
|
|
// Helper to setup test router with CapturingSink
|
|
fn setup_ai_test_router_with_capturing_sink() -> (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 router = router(
|
|
timesource,
|
|
readiness,
|
|
liveness,
|
|
Arc::new(sink),
|
|
redis,
|
|
None,
|
|
quota_limiter,
|
|
TokenDropper::default(),
|
|
None, // event_restriction_service
|
|
None, // recorder_handle
|
|
CaptureMode::Events,
|
|
None,
|
|
25 * 1024 * 1024,
|
|
false,
|
|
1_i64,
|
|
false,
|
|
0.0_f32,
|
|
26_214_400, // 25MB default for AI endpoint
|
|
Some(create_mock_blob_storage()), // ai_blob_storage
|
|
None, // body_chunk_read_timeout_ms
|
|
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)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_published_to_kafka() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "550e8400-e29b-41d4-a716-446655440000";
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify event was published to Kafka
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
1,
|
|
"Expected exactly one event to be published"
|
|
);
|
|
|
|
let event = &events[0];
|
|
|
|
// Verify event UUID
|
|
let event_json: serde_json::Value = serde_json::from_str(&event.event.data).unwrap();
|
|
assert_eq!(event_json["uuid"], event_uuid);
|
|
assert_eq!(event_json["event"], "$ai_generation");
|
|
assert_eq!(event_json["distinct_id"], "test_user");
|
|
|
|
// Verify properties
|
|
let props = event_json["properties"].as_object().unwrap();
|
|
assert_eq!(props["$ai_model"], "gpt-4");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_blobs_published_with_s3_placeholders() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "650e8400-e29b-41d4-a716-446655440001";
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let input_blob = json!({"messages": [{"role": "user", "content": "Hello"}]});
|
|
let output_blob = json!({"choices": [{"message": {"content": "Hi there"}}]});
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_input",
|
|
Part::bytes(serde_json::to_vec(&input_blob).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_output",
|
|
Part::bytes(serde_json::to_vec(&output_blob).unwrap())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify event was published to Kafka
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
1,
|
|
"Expected exactly one event to be published"
|
|
);
|
|
|
|
let event = &events[0];
|
|
let event_json: serde_json::Value = serde_json::from_str(&event.event.data).unwrap();
|
|
|
|
// Verify properties contain S3 URLs from mock blob storage
|
|
let props = event_json["properties"].as_object().unwrap();
|
|
|
|
// Both blobs should have S3 URLs
|
|
let input_url = props["$ai_input"].as_str().unwrap();
|
|
let output_url = props["$ai_output"].as_str().unwrap();
|
|
|
|
// Verify S3 URLs point to same file with different ranges
|
|
// URL format: s3://test-bucket/llma/<team_hash>/<uuid>?range=...
|
|
// Partitioned by the validated team id (StaticTeamResolver → TEST_TEAM_ID),
|
|
// hashed (first 16 chars of SHA-256) by the blob store's path sanitizer.
|
|
let team_hash = "6b86b273ff34fce1"; // SHA-256 of "1" (TEST_TEAM_ID)
|
|
let expected_prefix = format!("s3://{TEST_BLOB_BUCKET}/{TEST_BLOB_PREFIX}{team_hash}/");
|
|
assert!(
|
|
input_url.starts_with(&expected_prefix),
|
|
"Input URL should start with {expected_prefix}, got {input_url}"
|
|
);
|
|
assert!(
|
|
output_url.starts_with(&expected_prefix),
|
|
"Output URL should start with {expected_prefix}, got {output_url}"
|
|
);
|
|
|
|
// Extract UUIDs from both URLs (should be the same)
|
|
// URL format: s3://test-bucket/llma/<token_hash>/<uuid>?range=...
|
|
let input_uuid = input_url
|
|
.split('/')
|
|
.nth(5)
|
|
.unwrap()
|
|
.split('?')
|
|
.next()
|
|
.unwrap();
|
|
let output_uuid = output_url
|
|
.split('/')
|
|
.nth(5)
|
|
.unwrap()
|
|
.split('?')
|
|
.next()
|
|
.unwrap();
|
|
assert_eq!(
|
|
input_uuid, output_uuid,
|
|
"Both blobs should point to same file"
|
|
);
|
|
assert_eq!(input_uuid, event_uuid, "UUID should match event UUID");
|
|
|
|
// Verify ranges are sequential
|
|
assert!(
|
|
input_url.contains("range="),
|
|
"Input URL should have range parameter"
|
|
);
|
|
assert!(
|
|
output_url.contains("range="),
|
|
"Output URL should have range parameter"
|
|
);
|
|
|
|
// Extract range values and verify they are non-overlapping and sequential
|
|
let input_range = input_url.split("range=").nth(1).unwrap();
|
|
let output_range = output_url.split("range=").nth(1).unwrap();
|
|
|
|
// Parse ranges: format is "start-end"
|
|
let input_parts: Vec<usize> = input_range.split('-').map(|s| s.parse().unwrap()).collect();
|
|
let output_parts: Vec<usize> = output_range
|
|
.split('-')
|
|
.map(|s| s.parse().unwrap())
|
|
.collect();
|
|
|
|
let (input_start, input_end) = (input_parts[0], input_parts[1]);
|
|
let (output_start, output_end) = (output_parts[0], output_parts[1]);
|
|
|
|
// Verify ranges are valid (end > start)
|
|
assert!(input_end > input_start, "Input range should be valid");
|
|
assert!(output_end > output_start, "Output range should be valid");
|
|
|
|
// Verify ranges don't overlap (output starts after input ends)
|
|
assert!(
|
|
output_start > input_end,
|
|
"Output blob range should start after input blob range ends"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_multiple_blobs_sequential_ranges() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "750e8400-e29b-41d4-a716-446655440002";
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "claude-3"
|
|
});
|
|
|
|
// Create blobs of different sizes
|
|
let blob1 = "A".repeat(100); // 100 bytes
|
|
let blob2 = "B".repeat(200); // 200 bytes
|
|
let blob3 = "C".repeat(150); // 150 bytes
|
|
|
|
let form = 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(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_blob1",
|
|
Part::bytes(blob1.into_bytes())
|
|
.mime_str("text/plain")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_blob2",
|
|
Part::bytes(blob2.into_bytes())
|
|
.mime_str("text/plain")
|
|
.unwrap(),
|
|
)
|
|
.part(
|
|
"event.properties.$ai_blob3",
|
|
Part::bytes(blob3.into_bytes())
|
|
.mime_str("text/plain")
|
|
.unwrap(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify event was published to Kafka
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let event_json: serde_json::Value = serde_json::from_str(&event.event.data).unwrap();
|
|
let props = event_json["properties"].as_object().unwrap();
|
|
|
|
// Extract all URLs
|
|
let url1 = props["$ai_blob1"].as_str().unwrap();
|
|
let url2 = props["$ai_blob2"].as_str().unwrap();
|
|
let url3 = props["$ai_blob3"].as_str().unwrap();
|
|
|
|
// Verify all point to same file
|
|
// URL format: s3://test-bucket/llma/<token>/<uuid>?range=...
|
|
let uuid1 = url1.split('/').nth(5).unwrap().split('?').next().unwrap();
|
|
let uuid2 = url2.split('/').nth(5).unwrap().split('?').next().unwrap();
|
|
let uuid3 = url3.split('/').nth(5).unwrap().split('?').next().unwrap();
|
|
assert_eq!(uuid1, uuid2);
|
|
assert_eq!(uuid2, uuid3);
|
|
assert_eq!(uuid1, event_uuid);
|
|
|
|
// Extract and parse ranges
|
|
let range1 = url1.split("range=").nth(1).unwrap();
|
|
let range2 = url2.split("range=").nth(1).unwrap();
|
|
let range3 = url3.split("range=").nth(1).unwrap();
|
|
|
|
let parts1: Vec<usize> = range1.split('-').map(|s| s.parse().unwrap()).collect();
|
|
let parts2: Vec<usize> = range2.split('-').map(|s| s.parse().unwrap()).collect();
|
|
let parts3: Vec<usize> = range3.split('-').map(|s| s.parse().unwrap()).collect();
|
|
|
|
let (start1, end1) = (parts1[0], parts1[1]);
|
|
let (start2, end2) = (parts2[0], parts2[1]);
|
|
let (start3, end3) = (parts3[0], parts3[1]);
|
|
|
|
// Verify ranges are valid
|
|
assert!(end1 > start1, "First range should be valid");
|
|
assert!(end2 > start2, "Second range should be valid");
|
|
assert!(end3 > start3, "Third range should be valid");
|
|
|
|
// Verify ranges are sequential (non-overlapping)
|
|
assert!(start2 > end1, "Second blob should start after first ends");
|
|
assert!(start3 > end2, "Third blob should start after second ends");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_metadata_preserved_in_kafka() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "850e8400-e29b-41d4-a716-446655440003";
|
|
let distinct_id = "user_12345";
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_trace",
|
|
"distinct_id": distinct_id,
|
|
"timestamp": "2024-01-15T10:30:00Z"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4-turbo",
|
|
"$ai_trace_id": "trace_abc123",
|
|
"custom_metadata": "test_value"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify event metadata
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let event_json: serde_json::Value = serde_json::from_str(&event.event.data).unwrap();
|
|
|
|
// Verify all metadata is preserved
|
|
assert_eq!(event_json["uuid"], event_uuid);
|
|
assert_eq!(event_json["event"], "$ai_trace");
|
|
assert_eq!(event_json["distinct_id"], distinct_id);
|
|
assert_eq!(event_json["timestamp"], "2024-01-15T10:30:00Z");
|
|
|
|
let props = event_json["properties"].as_object().unwrap();
|
|
assert_eq!(props["$ai_model"], "gpt-4-turbo");
|
|
assert_eq!(props["$ai_trace_id"], "trace_abc123");
|
|
assert_eq!(props["custom_metadata"], "test_value");
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Timestamp and $ignore_sent_at Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_ignore_sent_at_true() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "950e8400-e29b-41d4-a716-446655440004";
|
|
let event_timestamp = "2024-01-15T10:00:00Z";
|
|
let sent_at = "2024-01-15T10:05:00Z"; // 5 minutes later
|
|
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"timestamp": event_timestamp,
|
|
"sent_at": sent_at
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4",
|
|
"$ignore_sent_at": true // Should ignore clock skew correction
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify timestamp was not adjusted (sent_at was ignored)
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let computed_timestamp = event.metadata.computed_timestamp.unwrap();
|
|
|
|
// Should use the event timestamp directly without clock skew correction
|
|
let expected = DateTime::parse_from_rfc3339(event_timestamp)
|
|
.unwrap()
|
|
.with_timezone(&chrono::Utc);
|
|
assert_eq!(
|
|
computed_timestamp, expected,
|
|
"With $ignore_sent_at=true, timestamp should match event timestamp exactly"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_ignore_sent_at_false() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "a50e8400-e29b-41d4-a716-446655440005";
|
|
let event_timestamp = "2024-01-01T11:59:55Z";
|
|
let sent_at = "2024-01-01T12:00:05Z"; // 10 seconds later
|
|
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"timestamp": event_timestamp,
|
|
"sent_at": sent_at
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4",
|
|
"$ignore_sent_at": false // Should apply clock skew correction
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify timestamp was adjusted using clock skew correction
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let computed_timestamp = event.metadata.computed_timestamp.unwrap();
|
|
|
|
// With clock skew correction:
|
|
// computed = now + (timestamp - sent_at)
|
|
// now = DEFAULT_TEST_TIME (2025-07-01T11:00:00Z)
|
|
// timestamp - sent_at = 11:59:55 - 12:00:05 = -10 seconds
|
|
// computed = 2025-07-01T11:00:00Z - 10s = 2025-07-01T10:59:50Z
|
|
let expected = chrono::Utc
|
|
.with_ymd_and_hms(2025, 7, 1, 10, 59, 50)
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
computed_timestamp, expected,
|
|
"With $ignore_sent_at=false, clock skew correction should be applied"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_without_ignore_sent_at_defaults_to_false() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "b50e8400-e29b-41d4-a716-446655440006";
|
|
let event_timestamp = "2024-01-01T11:59:55Z";
|
|
let sent_at = "2024-01-01T12:00:05Z"; // 10 seconds later
|
|
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"timestamp": event_timestamp,
|
|
"sent_at": sent_at
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
// No $ignore_sent_at property - should default to false
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify timestamp was adjusted (default behavior is to apply clock skew correction)
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let computed_timestamp = event.metadata.computed_timestamp.unwrap();
|
|
|
|
// Should apply clock skew correction by default
|
|
let expected = chrono::Utc
|
|
.with_ymd_and_hms(2025, 7, 1, 10, 59, 50)
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
computed_timestamp, expected,
|
|
"Without $ignore_sent_at property, clock skew correction should be applied (default behavior)"
|
|
);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// IP Address Extraction and Redaction Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_ip_defaults_to_localhost_in_tests() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "c50e8400-e29b-41d4-a716-446655440007";
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify IP address defaults to 127.0.0.1 in test environment
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
assert_eq!(
|
|
event.event.ip, "127.0.0.1",
|
|
"IP address should default to 127.0.0.1 when ConnectInfo is not available"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_ip_redacted_for_internal_events() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "d50e8400-e29b-41d4-a716-446655440008";
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user"
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4",
|
|
"capture_internal": true // Mark as internal event - IP should be redacted
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify IP address is redacted to 127.0.0.1 for internal events
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
assert_eq!(
|
|
event.event.ip, "127.0.0.1",
|
|
"IP address should be redacted to 127.0.0.1 for events with capture_internal property"
|
|
);
|
|
|
|
// Verify the capture_internal property is preserved
|
|
let event_json: serde_json::Value = serde_json::from_str(&event.event.data).unwrap();
|
|
let props = event_json["properties"].as_object().unwrap();
|
|
assert_eq!(props["capture_internal"], true);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Timestamp Conversion Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_invalid_sent_at_skips_clock_skew_correction() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "e50e8400-e29b-41d4-a716-446655440009";
|
|
let event_timestamp = "2024-01-01T11:59:55Z";
|
|
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"timestamp": event_timestamp,
|
|
"sent_at": "invalid-timestamp-format" // Invalid sent_at that can't be parsed
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify timestamp was not adjusted (sent_at was invalid, so treated as missing)
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let computed_timestamp = event.metadata.computed_timestamp.unwrap();
|
|
|
|
// Should use the event timestamp directly without clock skew correction
|
|
// (same behavior as if sent_at was missing)
|
|
let expected = DateTime::parse_from_rfc3339(event_timestamp)
|
|
.unwrap()
|
|
.with_timezone(&chrono::Utc);
|
|
assert_eq!(
|
|
computed_timestamp, expected,
|
|
"With invalid sent_at, timestamp should match event timestamp exactly (no clock skew correction)"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_without_sent_at_uses_event_timestamp() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "f50e8400-e29b-41d4-a716-446655440010";
|
|
let event_timestamp = "2024-01-01T11:59:55Z";
|
|
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"timestamp": event_timestamp
|
|
// No sent_at field
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify timestamp uses event timestamp without clock skew correction
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let computed_timestamp = event.metadata.computed_timestamp.unwrap();
|
|
|
|
// Should use the event timestamp directly
|
|
let expected = DateTime::parse_from_rfc3339(event_timestamp)
|
|
.unwrap()
|
|
.with_timezone(&chrono::Utc);
|
|
assert_eq!(
|
|
computed_timestamp, expected,
|
|
"Without sent_at, timestamp should match event timestamp exactly"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_valid_sent_at_applies_clock_skew_correction() {
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let event_uuid = "050e8400-e29b-41d4-a716-446655440011";
|
|
let event_timestamp = "2024-01-01T11:59:55Z";
|
|
let sent_at = "2024-01-01T12:00:05Z"; // 10 seconds later
|
|
|
|
let event_data = json!({
|
|
"uuid": event_uuid,
|
|
"event": "$ai_generation",
|
|
"distinct_id": "test_user",
|
|
"timestamp": event_timestamp,
|
|
"sent_at": sent_at // Valid sent_at
|
|
});
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = 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(),
|
|
);
|
|
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify timestamp was adjusted using clock skew correction
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
|
|
let event = &events[0];
|
|
let computed_timestamp = event.metadata.computed_timestamp.unwrap();
|
|
|
|
// With clock skew correction:
|
|
// computed = now + (timestamp - sent_at)
|
|
// now = DEFAULT_TEST_TIME (2025-07-01T11:00:00Z)
|
|
// timestamp - sent_at = 11:59:55 - 12:00:05 = -10 seconds
|
|
// computed = 2025-07-01T11:00:00Z - 10s = 2025-07-01T10:59:50Z
|
|
let expected = chrono::Utc
|
|
.with_ymd_and_hms(2025, 7, 1, 10, 59, 50)
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
computed_timestamp, expected,
|
|
"With valid sent_at, clock skew correction should be applied"
|
|
);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Token Dropper Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
// Helper to setup test router with custom TokenDropper and CapturingSink
|
|
fn setup_ai_test_router_with_token_dropper(token_dropper: TokenDropper) -> (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 router = router(
|
|
timesource,
|
|
readiness,
|
|
liveness,
|
|
Arc::new(sink),
|
|
redis,
|
|
None,
|
|
quota_limiter,
|
|
token_dropper,
|
|
None, // event_restriction_service
|
|
None, // recorder_handle
|
|
CaptureMode::Events,
|
|
None, // concurrency_limit
|
|
25 * 1024 * 1024, // event_size_limit
|
|
false, // enable_historical_rerouting
|
|
1, // historical_rerouting_threshold_days
|
|
false, // is_mirror_deploy
|
|
0.0, // verbose_sample_percent
|
|
26_214_400, // ai_max_sum_of_parts_bytes
|
|
Some(create_mock_blob_storage()), // ai_blob_storage
|
|
None, // body_chunk_read_timeout_ms
|
|
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)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_token_dropper_drops_matching_token() {
|
|
// Configure token dropper to drop all events for a specific token
|
|
let token_dropper = TokenDropper::new("pk-dropped_token");
|
|
let (router, sink) = setup_ai_test_router_with_token_dropper(token_dropper);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties);
|
|
|
|
// Use the dropped token
|
|
let response = send_multipart_request(&test_client, form, Some("pk-dropped_token")).await;
|
|
|
|
// Should return 200 OK (silently drops to avoid alerting clients)
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify no event was published to Kafka
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
0,
|
|
"Event should be dropped when token matches dropper configuration"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_token_dropper_allows_non_matching_token() {
|
|
// Configure token dropper to drop a different token
|
|
let token_dropper = TokenDropper::new("pk-other_token");
|
|
let (router, sink) = setup_ai_test_router_with_token_dropper(token_dropper);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties);
|
|
|
|
// Use a different token that is NOT in the dropper
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify event WAS published to Kafka
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
1,
|
|
"Event should be published when token does not match dropper configuration"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_token_dropper_drops_matching_token_and_distinct_id() {
|
|
// Configure token dropper to drop events for specific token:distinct_id combination
|
|
let token_dropper = TokenDropper::new("pk-specific_token:blocked_user");
|
|
let (router, sink) = setup_ai_test_router_with_token_dropper(token_dropper);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
// Create event with matching distinct_id
|
|
let form = create_ai_event_form("$ai_generation", "blocked_user", properties);
|
|
|
|
let response = send_multipart_request(&test_client, form, Some("pk-specific_token")).await;
|
|
|
|
// Should return 200 OK (silently drops to avoid alerting clients)
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify no event was published
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
0,
|
|
"Event should be dropped when token:distinct_id matches dropper configuration"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_token_dropper_allows_different_distinct_id() {
|
|
// Configure token dropper to drop events for specific token:distinct_id combination
|
|
let token_dropper = TokenDropper::new("pk-specific_token:blocked_user");
|
|
let (router, sink) = setup_ai_test_router_with_token_dropper(token_dropper);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
// Create event with DIFFERENT distinct_id
|
|
let form = create_ai_event_form("$ai_generation", "allowed_user", properties);
|
|
|
|
let response = send_multipart_request(&test_client, form, Some("pk-specific_token")).await;
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify event WAS published (different distinct_id)
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
1,
|
|
"Event should be published when distinct_id does not match dropper configuration"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_token_dropper_returns_success_with_empty_accepted_parts() {
|
|
// Configure token dropper to drop all events for a specific token
|
|
let token_dropper = TokenDropper::new("pk-dropped_token");
|
|
let (router, _sink) = setup_ai_test_router_with_token_dropper(token_dropper);
|
|
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-dropped_token")).await;
|
|
|
|
// Token dropper silently drops - returns 200 OK to avoid alerting clients
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify response body contains empty accepted_parts
|
|
let response_text = response.text().await;
|
|
assert!(
|
|
response_text.contains(r#""accepted_parts":[]"#),
|
|
"Response should contain empty 'accepted_parts' field, got: {response_text}"
|
|
);
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Quota Limiter Tests
|
|
// ----------------------------------------------------------------------------
|
|
|
|
use capture::quota_limiters::is_llm_event;
|
|
use limiters::redis::{QuotaResource, QUOTA_LIMITER_CACHE_KEY};
|
|
|
|
// Helper to setup test router with quota limiter configured to limit AI events
|
|
fn setup_ai_test_router_with_llm_quota_limited(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),
|
|
};
|
|
|
|
// Set up mock Redis to return this token as limited for LLM events
|
|
let llm_key = format!(
|
|
"{}{}",
|
|
QUOTA_LIMITER_CACHE_KEY,
|
|
QuotaResource::LLMEvents.as_str()
|
|
);
|
|
let redis =
|
|
Arc::new(MockRedisClient::new().zrangebyscore_ret(&llm_key, vec![token.to_string()]));
|
|
|
|
let mut cfg = DEFAULT_CONFIG.clone();
|
|
cfg.capture_mode = CaptureMode::Events;
|
|
|
|
let quota_limiter = CaptureQuotaLimiter::new(&cfg, redis.clone(), Duration::from_secs(60))
|
|
.add_scoped_limiter(QuotaResource::LLMEvents, is_llm_event);
|
|
|
|
let router = router(
|
|
timesource,
|
|
readiness,
|
|
liveness,
|
|
Arc::new(sink),
|
|
redis,
|
|
None,
|
|
quota_limiter,
|
|
TokenDropper::default(),
|
|
None, // event_restriction_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()), // ai_blob_storage
|
|
None, // body_chunk_read_timeout_ms
|
|
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)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_quota_limiter_drops_when_over_quota() {
|
|
let limited_token = "pk-limited_token_for_quota";
|
|
let (router, sink) = setup_ai_test_router_with_llm_quota_limited(limited_token);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties);
|
|
|
|
// Use the token that is over quota
|
|
let response = send_multipart_request(&test_client, form, Some(limited_token)).await;
|
|
|
|
// Should return 429 Too Many Requests (billing limit)
|
|
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
|
|
|
// Verify no event was published to Kafka
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
0,
|
|
"Event should be dropped when token is over LLM quota"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_quota_limiter_allows_when_not_over_quota() {
|
|
// Set up quota limiter with a different token limited
|
|
let limited_token = "pk-other_limited_token";
|
|
let (router, sink) = setup_ai_test_router_with_llm_quota_limited(limited_token);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({
|
|
"$ai_model": "gpt-4"
|
|
});
|
|
|
|
let form = create_ai_event_form("$ai_generation", "test_user", properties);
|
|
|
|
// Use a different token that is NOT over quota
|
|
let response = send_multipart_request(
|
|
&test_client,
|
|
form,
|
|
Some("pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3"),
|
|
)
|
|
.await;
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
// Verify event WAS published to Kafka
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
1,
|
|
"Event should be published when token is not over quota"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_quota_limiter_returns_billing_limit_error_message() {
|
|
let limited_token = "pk-quota_limited_token";
|
|
let (router, _sink) = setup_ai_test_router_with_llm_quota_limited(limited_token);
|
|
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(limited_token)).await;
|
|
|
|
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
|
|
|
// Verify response body contains billing limit error message
|
|
let response_text = response.text().await;
|
|
assert!(
|
|
response_text.contains("billing limit"),
|
|
"Response should contain 'billing limit' error message, got: {response_text}"
|
|
);
|
|
}
|
|
|
|
// ============================================================================
|
|
// OverflowLimiter coverage for the AI endpoint
|
|
// ============================================================================
|
|
//
|
|
// These tests exercise the shared `stamp_overflow_reason` helper invoked in
|
|
// `ai_handler` immediately after `build_kafka_event`. Pre-refactor the
|
|
// OverflowLimiter check happened inside the kafka sink; now it's a pipeline
|
|
// concern stamped into `ProcessedEventMetadata::overflow_reason`. Since AI
|
|
// bypasses `events::analytics::process_events`, these tests are the
|
|
// regression guard for `capture-ai-prod-us` (which runs with
|
|
// `OVERFLOW_ENABLED=true` + `OVERFLOW_PRESERVE_PARTITION_LOCALITY=true`).
|
|
|
|
const AI_OVERFLOW_TEST_TOKEN: &str = "pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3";
|
|
|
|
/// Variant of `setup_ai_test_router_with_capturing_sink` that wires a real
|
|
/// `OverflowLimiter` into the router. Existing helpers still pass `None`; this
|
|
/// one opts in to exercise the governor path.
|
|
fn setup_ai_test_router_with_overflow_limiter(
|
|
overflow_limiter: Arc<OverflowLimiter>,
|
|
) -> (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 router = router(
|
|
timesource,
|
|
readiness,
|
|
liveness,
|
|
Arc::new(sink),
|
|
redis,
|
|
None,
|
|
quota_limiter,
|
|
TokenDropper::default(),
|
|
None, // event_restriction_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),
|
|
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)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_hot_key_stamps_force_limited_reason() {
|
|
// Configure the OverflowLimiter to force-route the specific
|
|
// `token:distinct_id` key. `per_second`/`burst` are generous so only the
|
|
// explicit `keys_to_reroute` entry triggers — this isolates the
|
|
// ForceLimited branch from the RateLimited branch.
|
|
let hot_distinct_id = "user_hot_key";
|
|
let hot_key = format!("{AI_OVERFLOW_TEST_TOKEN}:{hot_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 (matches capture-ai-prod-us config)
|
|
));
|
|
|
|
let (router, sink) = setup_ai_test_router_with_overflow_limiter(overflow_limiter);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({"$ai_model": "gpt-4"});
|
|
let form = create_ai_event_form("$ai_generation", hot_distinct_id, properties);
|
|
|
|
let response = send_multipart_request(&test_client, form, Some(AI_OVERFLOW_TEST_TOKEN)).await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1, "AI event should still reach the sink");
|
|
|
|
let event = &events[0];
|
|
assert_eq!(
|
|
event.metadata.overflow_reason,
|
|
Some(OverflowReason::ForceLimited),
|
|
"hot key must be stamped ForceLimited so the sink routes to overflow"
|
|
);
|
|
assert!(
|
|
event.metadata.skip_person_processing,
|
|
"ForceLimited implies skip_person_processing so the header is set"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_event_with_cold_key_leaves_overflow_reason_none() {
|
|
// Governor is wide open and no keys_to_reroute — the helper must leave
|
|
// `overflow_reason` as None so the sink uses the default topic.
|
|
let overflow_limiter = Arc::new(OverflowLimiter::new(
|
|
NonZeroU32::new(1_000).unwrap(),
|
|
NonZeroU32::new(1_000).unwrap(),
|
|
None,
|
|
true,
|
|
));
|
|
|
|
let (router, sink) = setup_ai_test_router_with_overflow_limiter(overflow_limiter);
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({"$ai_model": "gpt-4"});
|
|
let form = create_ai_event_form("$ai_generation", "cold_user", properties);
|
|
|
|
let response = send_multipart_request(&test_client, form, Some(AI_OVERFLOW_TEST_TOKEN)).await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
assert_eq!(events[0].metadata.overflow_reason, None);
|
|
assert!(!events[0].metadata.skip_person_processing);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_without_overflow_limiter_is_parity_with_pre_refactor() {
|
|
// `capture-ai-*` with `OVERFLOW_ENABLED=false` (or setups without a
|
|
// limiter configured at all) must behave exactly as before: reason stays
|
|
// None. This pins the `None` short-circuit inside `stamp_overflow_reason`.
|
|
let (router, sink) = setup_ai_test_router_with_capturing_sink();
|
|
let test_client = TestClient::new(router);
|
|
|
|
let properties = json!({"$ai_model": "gpt-4"});
|
|
let form = create_ai_event_form("$ai_generation", "any_user", properties);
|
|
|
|
let response = send_multipart_request(&test_client, form, Some(AI_OVERFLOW_TEST_TOKEN)).await;
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
assert_eq!(events[0].metadata.overflow_reason, None);
|
|
}
|
|
|
|
// ============================================================================
|
|
// AI-gateway provenance on /i/v0/ai
|
|
// ============================================================================
|
|
|
|
const GW_SECRET: &str = "test-signing-secret";
|
|
|
|
// Shared router builder for the provenance tests, with the signing secret set;
|
|
// callers supply only what differs (the redis client and quota limiter).
|
|
fn ai_router(
|
|
redis: Arc<MockRedisClient>,
|
|
quota_limiter: CaptureQuotaLimiter,
|
|
) -> (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)
|
|
.unwrap()
|
|
.with_timezone(&Utc),
|
|
};
|
|
let router = router(
|
|
timesource,
|
|
readiness,
|
|
liveness,
|
|
Arc::new(sink),
|
|
redis,
|
|
None,
|
|
quota_limiter,
|
|
TokenDropper::default(),
|
|
None,
|
|
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,
|
|
50 * 1024 * 1024,
|
|
None,
|
|
None, // ai_events_overflow_limiter
|
|
None,
|
|
None,
|
|
8,
|
|
Some(GW_SECRET.to_string()),
|
|
AiRouting::Primary, // ai_routing
|
|
false, // ai_events_overflow_enabled
|
|
None, // ingestion_warning_emitter
|
|
);
|
|
(router, sink_clone)
|
|
}
|
|
|
|
// LLM-quota-limited (the token reads as over quota) with the signing secret set,
|
|
// so a verified event can prove it bypasses the limiter.
|
|
fn setup_ai_router_quota_limited_with_secret(token: &str) -> (Router, CapturingSink) {
|
|
let llm_key = format!(
|
|
"{}{}",
|
|
QUOTA_LIMITER_CACHE_KEY,
|
|
QuotaResource::LLMEvents.as_str()
|
|
);
|
|
let redis =
|
|
Arc::new(MockRedisClient::new().zrangebyscore_ret(&llm_key, vec![token.to_string()]));
|
|
let mut cfg = DEFAULT_CONFIG.clone();
|
|
cfg.capture_mode = CaptureMode::Events;
|
|
let quota_limiter = CaptureQuotaLimiter::new(&cfg, redis.clone(), Duration::from_secs(60))
|
|
.add_scoped_limiter(QuotaResource::LLMEvents, is_llm_event);
|
|
ai_router(redis, quota_limiter)
|
|
}
|
|
|
|
// Signing secret set, no quota limit, so the published event is observable
|
|
// (used to assert a forged marker is stripped).
|
|
fn setup_ai_router_with_secret() -> (Router, CapturingSink) {
|
|
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));
|
|
ai_router(redis, quota_limiter)
|
|
}
|
|
|
|
// Sends an AI multipart request with the gateway provenance headers attached.
|
|
async fn send_signed_ai_request(
|
|
client: &TestClient,
|
|
form: Form,
|
|
token: &str,
|
|
signature: &str,
|
|
signed_at: &str,
|
|
request_id: &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());
|
|
}
|
|
client
|
|
.post("/i/v0/ai")
|
|
.header("Content-Type", content_type)
|
|
.header("Authorization", format!("Bearer {token}"))
|
|
.header("Insights-Ai-Gateway-Signature", signature)
|
|
.header("Insights-Ai-Gateway-Signed-At", signed_at)
|
|
.header("Insights-Ai-Gateway-Request-Id", request_id)
|
|
.body(body)
|
|
.send()
|
|
.await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_verified_gateway_event_bypasses_quota_and_is_stamped() {
|
|
let token = "pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3";
|
|
let (router, sink) = setup_ai_router_quota_limited_with_secret(token);
|
|
let client = TestClient::new(router);
|
|
|
|
let distinct_id = "user-7";
|
|
let request_id = "req-verified-1";
|
|
let sig = capture::v1::gateway_provenance::sign_for_test(
|
|
GW_SECRET.as_bytes(),
|
|
token,
|
|
distinct_id,
|
|
request_id,
|
|
DEFAULT_TEST_TIME,
|
|
);
|
|
let form = create_ai_event_form(
|
|
"$ai_generation",
|
|
distinct_id,
|
|
json!({"$ai_model": "claude"}),
|
|
);
|
|
let resp =
|
|
send_signed_ai_request(&client, form, token, &sig, DEFAULT_TEST_TIME, request_id).await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
// The token is LLM-quota-limited, so an unverified event would be dropped. A
|
|
// verified gateway event must be exempt (published) and stamped.
|
|
let events = sink.get_events().await;
|
|
assert_eq!(
|
|
events.len(),
|
|
1,
|
|
"verified gateway event must bypass the LLM quota limiter"
|
|
);
|
|
let data: Value = serde_json::from_str(&events[0].event.data).unwrap();
|
|
assert_eq!(data["properties"]["$ai_gateway_verified"], json!(true));
|
|
assert_eq!(
|
|
data["properties"]["$ai_gateway_request_id"],
|
|
json!(request_id)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_forged_marker_is_stripped() {
|
|
let token = "pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3";
|
|
let (router, sink) = setup_ai_router_with_secret();
|
|
let client = TestClient::new(router);
|
|
|
|
// Client forges the verified marker and sends an invalid signature.
|
|
let form = create_ai_event_form(
|
|
"$ai_generation",
|
|
"user-7",
|
|
json!({"$ai_model": "claude", "$ai_gateway_verified": true}),
|
|
);
|
|
let resp = send_signed_ai_request(
|
|
&client,
|
|
form,
|
|
token,
|
|
"deadbeef",
|
|
DEFAULT_TEST_TIME,
|
|
"req-forged",
|
|
)
|
|
.await;
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
|
|
let events = sink.get_events().await;
|
|
assert_eq!(events.len(), 1);
|
|
let data: Value = serde_json::from_str(&events[0].event.data).unwrap();
|
|
assert!(
|
|
data["properties"].get("$ai_gateway_verified").is_none(),
|
|
"a forged $ai_gateway_verified with an invalid signature must be stripped"
|
|
);
|
|
}
|
|
|
|
// Over the global Events quota (the team-wide cap, not the scoped LLM one), with
|
|
// the signing secret set.
|
|
fn setup_ai_router_global_quota_limited_with_secret(token: &str) -> (Router, CapturingSink) {
|
|
let events_key = format!(
|
|
"{}{}",
|
|
QUOTA_LIMITER_CACHE_KEY,
|
|
QuotaResource::Events.as_str()
|
|
);
|
|
let redis =
|
|
Arc::new(MockRedisClient::new().zrangebyscore_ret(&events_key, vec![token.to_string()]));
|
|
let mut cfg = DEFAULT_CONFIG.clone();
|
|
cfg.capture_mode = CaptureMode::Events;
|
|
let quota_limiter = CaptureQuotaLimiter::new(&cfg, redis.clone(), Duration::from_secs(60));
|
|
ai_router(redis, quota_limiter)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_ai_endpoint_verified_gateway_event_still_subject_to_global_quota() {
|
|
let token = "pk-VXRzc3poSG9GZm1JenRianJ6TTJFZGh4OWY2QXzx9f3";
|
|
let (router, sink) = setup_ai_router_global_quota_limited_with_secret(token);
|
|
let client = TestClient::new(router);
|
|
|
|
let distinct_id = "user-7";
|
|
let request_id = "req-global-limited";
|
|
let sig = capture::v1::gateway_provenance::sign_for_test(
|
|
GW_SECRET.as_bytes(),
|
|
token,
|
|
distinct_id,
|
|
request_id,
|
|
DEFAULT_TEST_TIME,
|
|
);
|
|
let form = create_ai_event_form(
|
|
"$ai_generation",
|
|
distinct_id,
|
|
json!({"$ai_model": "claude"}),
|
|
);
|
|
let resp =
|
|
send_signed_ai_request(&client, form, token, &sig, DEFAULT_TEST_TIME, request_id).await;
|
|
|
|
// Verified gateway events are exempt from the scoped LLM quota but not the
|
|
// global Events quota, so an over-global-quota token is still rejected.
|
|
assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
|
|
assert_eq!(
|
|
sink.get_events().await.len(),
|
|
0,
|
|
"verified gateway event must still obey the global Events quota"
|
|
);
|
|
}
|
|
|
|
// ============================================================================
|
|
// PHASE 11: INGESTION WARNINGS
|
|
// ============================================================================
|
|
//
|
|
// The mapping, details, and count are pinned at the helper level in
|
|
// `ingestion_warnings::ai`. These cover what a unit test can't: that the emit
|
|
// seam is wired into the handler at all, that attribution reaches it from the
|
|
// event part, and that a rejection still produces its original HTTP response.
|
|
|
|
#[tokio::test]
|
|
async fn non_ai_event_name_warns_and_still_returns_400() {
|
|
let (app, emitter) = setup_ai_router_collecting_warnings();
|
|
let client = TestClient::new(app);
|
|
|
|
// Analytics sent to the AI path: the mistake invalid_ai_event exists for.
|
|
let form = Form::new().part(
|
|
"event",
|
|
Part::bytes(
|
|
serde_json::to_vec(&json!({
|
|
"uuid": uuid::Uuid::now_v7().to_string(),
|
|
"event": "$pageview",
|
|
"distinct_id": "user-1",
|
|
"properties": {"$ai_model": "gpt-4", "$lib": "insights-python", "$lib_version": "3.1.0"}
|
|
}))
|
|
.unwrap(),
|
|
)
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let resp = send_multipart_request(&client, form, Some("test_token")).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
let emitted = emitter.emitted();
|
|
assert_eq!(emitted.len(), 1);
|
|
assert_eq!(emitted[0].warning, WarningType::InvalidAiEvent);
|
|
assert_eq!(emitted[0].source, CAPTURE_AI_EVENTS);
|
|
assert_eq!(emitted[0].token, "test_token");
|
|
assert_eq!(emitted[0].count, 1);
|
|
assert_eq!(
|
|
emitted[0].extra_details.get("eventName"),
|
|
Some(&json!("$pageview"))
|
|
);
|
|
// Attribution came off the event part's own properties.
|
|
assert_eq!(
|
|
emitted[0].extra_details.get("lib"),
|
|
Some(&json!("insights-python"))
|
|
);
|
|
assert_eq!(
|
|
emitted[0].extra_details.get("libVersion"),
|
|
Some(&json!("3.1.0"))
|
|
);
|
|
assert_eq!(
|
|
emitted[0].extra_details.get("path"),
|
|
Some(&json!("/i/v0/ai"))
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn missing_ai_model_warns_as_invalid_ai_event() {
|
|
let (app, emitter) = setup_ai_router_collecting_warnings();
|
|
let client = TestClient::new(app);
|
|
|
|
let form = create_ai_event_form("$ai_generation", "user-2", json!({"foo": "bar"}));
|
|
let resp = send_multipart_request(&client, form, Some("test_token")).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
let emitted = emitter.emitted();
|
|
assert_eq!(emitted.len(), 1);
|
|
assert_eq!(emitted[0].warning, WarningType::InvalidAiEvent);
|
|
// No $lib in the properties, so attribution falls back rather than dropping keys.
|
|
assert_eq!(emitted[0].extra_details.get("lib"), Some(&json!("unknown")));
|
|
}
|
|
|
|
// Structural failures happen before the event part parses, so this also covers
|
|
// the pre-event-part attribution fallback on the seam.
|
|
#[tokio::test]
|
|
async fn unknown_multipart_field_warns_as_invalid_ai_payload() {
|
|
let (app, emitter) = setup_ai_router_collecting_warnings();
|
|
let client = TestClient::new(app);
|
|
|
|
let form = create_ai_event_form("$ai_generation", "user-3", json!({"$ai_model": "gpt-4"}))
|
|
.part(
|
|
"surprise",
|
|
Part::bytes(b"data".to_vec())
|
|
.mime_str("application/json")
|
|
.unwrap(),
|
|
);
|
|
|
|
let resp = send_multipart_request(&client, form, Some("test_token")).await;
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
let emitted = emitter.emitted();
|
|
assert_eq!(emitted.len(), 1);
|
|
assert_eq!(emitted[0].warning, WarningType::InvalidAiPayload);
|
|
assert_eq!(
|
|
emitted[0].extra_details.get("part"),
|
|
Some(&json!("surprise"))
|
|
);
|
|
}
|
|
|
|
// A missing Bearer header is pre-token, so there is no team to attribute to and
|
|
// the seam must stay silent. This is the property that keeps unattributable
|
|
// failures out of the warnings table.
|
|
#[tokio::test]
|
|
async fn unauthenticated_request_never_warns() {
|
|
let (app, emitter) = setup_ai_router_collecting_warnings();
|
|
let client = TestClient::new(app);
|
|
|
|
let form = create_ai_event_form("$ai_generation", "user-4", json!({"$ai_model": "gpt-4"}));
|
|
let resp = send_multipart_request(&client, form, None).await;
|
|
|
|
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
|
assert!(emitter.emitted().is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn valid_ai_event_never_warns() {
|
|
let (app, emitter) = setup_ai_router_collecting_warnings();
|
|
let client = TestClient::new(app);
|
|
|
|
let form = create_ai_event_form("$ai_generation", "user-5", json!({"$ai_model": "gpt-4"}));
|
|
let resp = send_multipart_request(&client, form, Some("test_token")).await;
|
|
|
|
assert_eq!(resp.status(), StatusCode::OK);
|
|
assert!(emitter.emitted().is_empty());
|
|
}
|