From be3dfdf66a57f5ddfa0803bfd2c6be789785a49c Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 31 Oct 2026 15:25:25 +0530 Subject: [PATCH 1/6] feat: add identity persistence seam --- src/identities/mod.rs | 3 + src/identities/persistence.rs | 30 +++++ src/lib.rs | 1 + src/main.rs | 2 +- src/routes/mod.rs | 12 +- src/server.rs | 9 +- src/services/environment.rs | 65 +++++++++- tests/fixtures.rs | 2 +- tests/test_cors.rs | 2 +- tests/test_health.rs | 10 +- tests/test_identity_persistence.rs | 196 +++++++++++++++++++++++++++++ tests/test_server.rs | 4 +- 12 files changed, 317 insertions(+), 19 deletions(-) create mode 100644 src/identities/mod.rs create mode 100644 src/identities/persistence.rs create mode 100644 tests/test_identity_persistence.rs diff --git a/src/identities/mod.rs b/src/identities/mod.rs new file mode 100644 index 0000000..e715637 --- /dev/null +++ b/src/identities/mod.rs @@ -0,0 +1,3 @@ +pub mod persistence; + +pub use persistence::{IdentityPersistence, IdentityTraits, RequestOrigin}; diff --git a/src/identities/persistence.rs b/src/identities/persistence.rs new file mode 100644 index 0000000..c8d06d0 --- /dev/null +++ b/src/identities/persistence.rs @@ -0,0 +1,30 @@ +use crate::models::{IdentityWithTraits, TraitModel}; +use async_trait::async_trait; + +#[async_trait] +pub trait IdentityPersistence: Send + Sync { + async fn load( + &self, + environment: &serde_json::Value, + identity: &IdentityWithTraits, + ) -> IdentityTraits; + + fn persist( + &self, + environment: &serde_json::Value, + identity: &IdentityWithTraits, + origin: RequestOrigin, + ); +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct IdentityTraits { + pub evaluation: Vec, + pub response: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RequestOrigin { + Client, + Server, +} diff --git a/src/lib.rs b/src/lib.rs index c13d760..4054243 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ pub mod cache; pub mod config; pub mod environments; pub mod error; +pub mod identities; pub mod models; pub mod routes; pub mod server; diff --git a/src/main.rs b/src/main.rs index aa74557..f8c213b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,5 +2,5 @@ use edge_proxy::config::settings::get_settings; #[tokio::main] async fn main() -> anyhow::Result<()> { - edge_proxy::run(get_settings()?).await + edge_proxy::run(get_settings()?, None).await } diff --git a/src/routes/mod.rs b/src/routes/mod.rs index fc419b0..c0bda5b 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -6,6 +6,7 @@ pub mod health; pub mod identities; use crate::config::AppSettings; +use crate::identities::IdentityPersistence; use crate::services::EnvironmentService; use axum::{ Router, @@ -15,9 +16,16 @@ use axum::{ use std::sync::Arc; use tower_http::{compression::CompressionLayer, normalize_path::NormalizePath, trace::TraceLayer}; -pub fn create_router(settings: AppSettings) -> (Router, Arc) { +pub fn create_router( + settings: AppSettings, + persistence: Option>, +) -> (Router, Arc) { let cors = cors::layer(&settings.allow_origins); - let environment_service = Arc::new(EnvironmentService::new(settings)); + let environment_service = EnvironmentService::new(settings); + let environment_service = Arc::new(match persistence { + Some(persistence) => environment_service.with_identity_persistence(persistence), + None => environment_service, + }); let router = Router::new() // Health check routes diff --git a/src/server.rs b/src/server.rs index 0ccba09..0442b5c 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,12 +1,17 @@ use crate::config::AppSettings; use crate::config::logging::setup_logging; +use crate::identities::IdentityPersistence; use crate::routes::create_router; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; use tracing::{info, warn}; const DEFAULT_HOST: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED); -pub async fn run(settings: AppSettings) -> anyhow::Result<()> { +pub async fn run( + settings: AppSettings, + persistence: Option>, +) -> anyhow::Result<()> { setup_logging(&settings.logging); info!("Starting Edge Proxy server..."); @@ -27,7 +32,7 @@ pub async fn run(settings: AppSettings) -> anyhow::Result<()> { ); } - let (app, environment_service) = create_router(settings.clone()); + let (app, environment_service) = create_router(settings.clone(), persistence); // Refreshes must never overlap: a delayed older poll finishing after a // newer one could restore removed environments or rotated keys. The diff --git a/src/services/environment.rs b/src/services/environment.rs index 5ef66dd..f13348f 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -2,8 +2,9 @@ use crate::cache::{EnvironmentsCache, LocalMemEnvironmentsCache}; use crate::config::settings::AppSettings; use crate::environments::{EnvironmentIndex, EnvironmentKeys}; use crate::error::{EdgeProxyError, Result}; +use crate::identities::{IdentityPersistence, RequestOrigin}; use crate::models::{ - APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, + APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, TraitModel, }; use crate::services::feature_utils::filter_out_server_key_only_flag_results; use chrono::{DateTime, Utc}; @@ -23,6 +24,7 @@ pub struct EnvironmentService { pub settings: AppSettings, pub last_updated_at: Arc>>>, environments: EnvironmentIndex, + persistence: Option>, } impl EnvironmentService { @@ -41,6 +43,7 @@ impl EnvironmentService { settings, last_updated_at: Arc::new(RwLock::new(None)), environments, + persistence: None, } } @@ -50,6 +53,11 @@ impl EnvironmentService { service } + pub fn with_identity_persistence(mut self, persistence: Arc) -> Self { + self.persistence = Some(persistence); + self + } + pub async fn refresh_environment_caches(&self) -> bool { // Sync first so an environment added to the proxy config gets its // document fetched in the same pass. @@ -342,6 +350,54 @@ impl EnvironmentService { ) -> Result { let keys = self.resolve_key(environment_key)?; + let Some(persistence) = &self.persistence else { + return self + .evaluate_identity( + &keys, + environment_key, + &identity.identifier, + &identity.traits, + &identity.traits, + ) + .await; + }; + + let environment = self + .cache + .get_environment(&keys.client_key) + .await + .ok_or_else(|| { + EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string()) + })?; + let traits = persistence.load(&environment, identity).await; + let response = self + .evaluate_identity( + &keys, + environment_key, + &identity.identifier, + &traits.evaluation, + &traits.response, + ) + .await?; + + let origin = if environment_key == keys.client_key { + RequestOrigin::Client + } else { + RequestOrigin::Server + }; + persistence.persist(&environment, identity, origin); + + Ok(response) + } + + async fn evaluate_identity( + &self, + keys: &EnvironmentKeys, + environment_key: &str, + identifier: &str, + evaluation_traits: &[TraitModel], + response_traits: &[TraitModel], + ) -> Result { // Get pre-computed context from cache let context = self .cache @@ -352,10 +408,10 @@ impl EnvironmentService { })?; let flagsmith_traits: Vec = - identity.traits.iter().map(Into::into).collect(); + evaluation_traits.iter().map(Into::into).collect(); let context_with_identity = - add_identity_to_context(&context, &identity.identifier, &flagsmith_traits); + add_identity_to_context(&context, identifier, &flagsmith_traits); let evaluation_result = get_evaluation_result(&context_with_identity); @@ -373,8 +429,7 @@ impl EnvironmentService { Ok(IdentityResponse { flags: flag_results.iter().map(Into::into).collect(), - traits: identity - .traits + traits: response_traits .iter() .map(|t| t.to_response_json()) .collect(), diff --git a/tests/fixtures.rs b/tests/fixtures.rs index 59781a1..d6ae472 100644 --- a/tests/fixtures.rs +++ b/tests/fixtures.rs @@ -89,7 +89,7 @@ pub fn environment_1() -> serde_json::Value { { "value": "test", "operator": "EQUAL", - "property": "first_name", + "property_": "first_name", } ], "type": "ANY", diff --git a/tests/test_cors.rs b/tests/test_cors.rs index c7fa68b..c18a543 100644 --- a/tests/test_cors.rs +++ b/tests/test_cors.rs @@ -17,7 +17,7 @@ async fn server_allowing(allow_origins: &[&str]) -> TestServer { allow_origins: allow_origins.iter().map(|o| o.to_string()).collect(), ..AppSettings::new() }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); service .cache .put_environment(&environment_1_api_key(), environment_1()) diff --git a/tests/test_health.rs b/tests/test_health.rs index 8016c49..d4340d2 100644 --- a/tests/test_health.rs +++ b/tests/test_health.rs @@ -7,7 +7,7 @@ use edge_proxy::routes::create_router; async fn test_liveness_check() { // Given let settings = AppSettings::new(); - let (app, _) = create_router(settings); + let (app, _) = create_router(settings, None); let server = TestServer::new(app).unwrap(); // When @@ -21,7 +21,7 @@ async fn test_liveness_check() { async fn test_health_check_returns_200_if_cache_was_updated_recently() { // Given let settings = AppSettings::new(); - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); { let mut last_updated = service.last_updated_at.write().await; *last_updated = Some(Utc::now()); @@ -42,7 +42,7 @@ async fn test_health_check_returns_200_if_cache_was_updated_recently() { async fn test_health_check_returns_503_if_cache_was_not_updated() { // Given let settings = AppSettings::new(); - let (app, _service) = create_router(settings); + let (app, _service) = create_router(settings, None); let server = TestServer::new(app).unwrap(); // last_updated_at is None by default @@ -62,7 +62,7 @@ async fn test_health_check_returns_503_if_cache_was_not_updated() { async fn test_health_check_returns_503_if_cache_is_stale() { // Given let settings = AppSettings::new(); - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); { let mut last_updated = service.last_updated_at.write().await; *last_updated = Some(Utc::now() - chrono::Duration::days(10)); @@ -88,7 +88,7 @@ async fn test_health_check_returns_200_if_cache_is_never_stale() { settings.health_check = HealthCheckSettings { environment_update_grace_period_seconds: None, }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); let last_update_time = Utc::now() - chrono::Duration::days(10); { let mut last_updated = service.last_updated_at.write().await; diff --git a/tests/test_identity_persistence.rs b/tests/test_identity_persistence.rs new file mode 100644 index 0000000..f92e438 --- /dev/null +++ b/tests/test_identity_persistence.rs @@ -0,0 +1,196 @@ +mod fixtures; + +use async_trait::async_trait; +use axum_test::TestServer; +use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; +use edge_proxy::identities::{IdentityPersistence, IdentityTraits, RequestOrigin}; +use edge_proxy::models::{IdentityWithTraits, TraitModel}; +use edge_proxy::routes::create_router; +use fixtures::*; +use rstest::rstest; +use serde_json::{Value, json}; +use std::sync::{Arc, Mutex}; + +#[derive(Debug, PartialEq)] +struct PersistCall { + environment_api_key: Value, + identifier: String, + origin: RequestOrigin, +} + +struct RecordingPersistence { + traits: IdentityTraits, + persist_calls: Mutex>, +} + +impl RecordingPersistence { + fn new(evaluation: Vec, response: Vec) -> Self { + Self { + traits: IdentityTraits { + evaluation, + response, + }, + persist_calls: Mutex::new(Vec::new()), + } + } +} + +#[async_trait] +impl IdentityPersistence for RecordingPersistence { + async fn load(&self, _environment: &Value, _identity: &IdentityWithTraits) -> IdentityTraits { + self.traits.clone() + } + + fn persist(&self, environment: &Value, identity: &IdentityWithTraits, origin: RequestOrigin) { + self.persist_calls.lock().unwrap().push(PersistCall { + environment_api_key: environment["api_key"].clone(), + identifier: identity.identifier.clone(), + origin, + }); + } +} + +fn trait_model(trait_key: &str, trait_value: &str) -> TraitModel { + TraitModel { + trait_key: trait_key.to_string(), + trait_value: json!(trait_value), + transient: false, + } +} + +async fn setup_test_server(persistence: Option>) -> TestServer { + let settings = AppSettings { + environment_key_pairs: vec![EnvironmentKeyPair { + server_side_key: environment_1_server_key(), + client_side_key: environment_1_api_key(), + }], + api_url: "http://test".to_string(), + ..AppSettings::new() + }; + + let (app, service) = create_router(settings, persistence); + + service + .cache + .put_environment(&environment_1_api_key(), environment_1()) + .await; + + TestServer::new(app).unwrap() +} + +async fn post_identity(server: &TestServer, environment_key: &str, traits: Value) -> Value { + let response = server + .post("/api/v1/identities/") + .add_header("X-Environment-Key", environment_key) + .json(&json!({"identifier": "identity_1", "traits": traits})) + .await; + response.assert_status_ok(); + response.json() +} + +fn feature_2_value(body: &Value) -> &Value { + &body["flags"] + .as_array() + .unwrap() + .iter() + .find(|f| f["feature"]["name"] == "feature_2") + .unwrap()["feature_state_value"] +} + +#[rstest] +#[case::segment_match("test", "segment_override")] +#[case::segment_miss("other", "2.3")] +#[tokio::test] +async fn test_identity_without_persistence_evaluates_and_returns_request_traits( + #[case] first_name: &str, + #[case] expected_feature_2_value: &str, +) { + // Given + let server = setup_test_server(None).await; + let request_traits = json!([{"trait_key": "first_name", "trait_value": first_name}]); + + // When + let body = post_identity(&server, &environment_1_api_key(), request_traits.clone()).await; + + // Then + assert_eq!(feature_2_value(&body), expected_feature_2_value); + assert_eq!(body["traits"], request_traits); +} + +#[rstest] +#[case::loaded_traits_match_segment("other", "test", "segment_override")] +#[case::request_traits_ignored("test", "other", "2.3")] +#[tokio::test] +async fn test_identity_with_persistence_evaluates_loaded_traits( + #[case] request_first_name: &str, + #[case] loaded_first_name: &str, + #[case] expected_feature_2_value: &str, +) { + // Given + let persistence = Arc::new(RecordingPersistence::new( + vec![trait_model("first_name", loaded_first_name)], + Vec::new(), + )); + let server = setup_test_server(Some(persistence)).await; + + // When + let body = post_identity( + &server, + &environment_1_api_key(), + json!([{"trait_key": "first_name", "trait_value": request_first_name}]), + ) + .await; + + // Then + assert_eq!(feature_2_value(&body), expected_feature_2_value); +} + +#[tokio::test] +async fn test_identity_with_persistence_returns_loaded_response_traits() { + // Given + let persistence = Arc::new(RecordingPersistence::new( + Vec::new(), + vec![trait_model("stored_key", "stored_value")], + )); + let server = setup_test_server(Some(persistence)).await; + + // When + let body = post_identity( + &server, + &environment_1_api_key(), + json!([{"trait_key": "request_key", "trait_value": "request_value"}]), + ) + .await; + + // Then + assert_eq!( + body["traits"], + json!([{"trait_key": "stored_key", "trait_value": "stored_value"}]) + ); +} + +#[rstest] +#[case::client_key(environment_1_api_key(), RequestOrigin::Client)] +#[case::server_key(environment_1_server_key(), RequestOrigin::Server)] +#[tokio::test] +async fn test_identity_with_persistence_persists_once_with_origin( + #[case] environment_key: String, + #[case] expected_origin: RequestOrigin, +) { + // Given + let persistence = Arc::new(RecordingPersistence::new(Vec::new(), Vec::new())); + let server = setup_test_server(Some(persistence.clone())).await; + + // When + post_identity(&server, &environment_key, json!([])).await; + + // Then + assert_eq!( + *persistence.persist_calls.lock().unwrap(), + vec![PersistCall { + environment_api_key: json!(environment_1_api_key()), + identifier: "identity_1".to_string(), + origin: expected_origin, + }] + ); +} diff --git a/tests/test_server.rs b/tests/test_server.rs index 5853ba5..da45a36 100644 --- a/tests/test_server.rs +++ b/tests/test_server.rs @@ -17,7 +17,7 @@ async fn setup_test_server() -> TestServer { ..AppSettings::new() }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); service .cache @@ -238,7 +238,7 @@ async fn test_get_environment_document() { }], ..AppSettings::new() }; - let (app, service) = create_router(settings); + let (app, service) = create_router(settings, None); service .cache .put_environment("client_key", environment_1()) From c0c5235e4e28e8998a6b228a5357c46d5f400f42 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 31 Oct 2026 16:23:40 +0530 Subject: [PATCH 2/6] refactor: rename identity persistence seam to IdentityStore --- src/identities/mod.rs | 4 +- src/identities/{persistence.rs => store.rs} | 6 +- src/routes/mod.rs | 8 +-- src/server.rs | 6 +- src/services/environment.rs | 16 +++--- ..._persistence.rs => test_identity_store.rs} | 57 +++++++++++-------- 6 files changed, 53 insertions(+), 44 deletions(-) rename src/identities/{persistence.rs => store.rs} (87%) rename tests/{test_identity_persistence.rs => test_identity_store.rs} (73%) diff --git a/src/identities/mod.rs b/src/identities/mod.rs index e715637..aa867fd 100644 --- a/src/identities/mod.rs +++ b/src/identities/mod.rs @@ -1,3 +1,3 @@ -pub mod persistence; +pub mod store; -pub use persistence::{IdentityPersistence, IdentityTraits, RequestOrigin}; +pub use store::{IdentityStore, IdentityTraits, RequestOrigin}; diff --git a/src/identities/persistence.rs b/src/identities/store.rs similarity index 87% rename from src/identities/persistence.rs rename to src/identities/store.rs index c8d06d0..fdfaacb 100644 --- a/src/identities/persistence.rs +++ b/src/identities/store.rs @@ -2,14 +2,14 @@ use crate::models::{IdentityWithTraits, TraitModel}; use async_trait::async_trait; #[async_trait] -pub trait IdentityPersistence: Send + Sync { - async fn load( +pub trait IdentityStore: Send + Sync { + async fn load_traits( &self, environment: &serde_json::Value, identity: &IdentityWithTraits, ) -> IdentityTraits; - fn persist( + fn save_traits( &self, environment: &serde_json::Value, identity: &IdentityWithTraits, diff --git a/src/routes/mod.rs b/src/routes/mod.rs index c0bda5b..b00ec37 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -6,7 +6,7 @@ pub mod health; pub mod identities; use crate::config::AppSettings; -use crate::identities::IdentityPersistence; +use crate::identities::IdentityStore; use crate::services::EnvironmentService; use axum::{ Router, @@ -18,12 +18,12 @@ use tower_http::{compression::CompressionLayer, normalize_path::NormalizePath, t pub fn create_router( settings: AppSettings, - persistence: Option>, + identity_store: Option>, ) -> (Router, Arc) { let cors = cors::layer(&settings.allow_origins); let environment_service = EnvironmentService::new(settings); - let environment_service = Arc::new(match persistence { - Some(persistence) => environment_service.with_identity_persistence(persistence), + let environment_service = Arc::new(match identity_store { + Some(identity_store) => environment_service.with_identity_store(identity_store), None => environment_service, }); diff --git a/src/server.rs b/src/server.rs index 0442b5c..9c5d95d 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,6 +1,6 @@ use crate::config::AppSettings; use crate::config::logging::setup_logging; -use crate::identities::IdentityPersistence; +use crate::identities::IdentityStore; use crate::routes::create_router; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::sync::Arc; @@ -10,7 +10,7 @@ const DEFAULT_HOST: IpAddr = IpAddr::V4(Ipv4Addr::UNSPECIFIED); pub async fn run( settings: AppSettings, - persistence: Option>, + identity_store: Option>, ) -> anyhow::Result<()> { setup_logging(&settings.logging); @@ -32,7 +32,7 @@ pub async fn run( ); } - let (app, environment_service) = create_router(settings.clone(), persistence); + let (app, environment_service) = create_router(settings.clone(), identity_store); // Refreshes must never overlap: a delayed older poll finishing after a // newer one could restore removed environments or rotated keys. The diff --git a/src/services/environment.rs b/src/services/environment.rs index f13348f..6ac9970 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -2,7 +2,7 @@ use crate::cache::{EnvironmentsCache, LocalMemEnvironmentsCache}; use crate::config::settings::AppSettings; use crate::environments::{EnvironmentIndex, EnvironmentKeys}; use crate::error::{EdgeProxyError, Result}; -use crate::identities::{IdentityPersistence, RequestOrigin}; +use crate::identities::{IdentityStore, RequestOrigin}; use crate::models::{ APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, TraitModel, }; @@ -24,7 +24,7 @@ pub struct EnvironmentService { pub settings: AppSettings, pub last_updated_at: Arc>>>, environments: EnvironmentIndex, - persistence: Option>, + identity_store: Option>, } impl EnvironmentService { @@ -43,7 +43,7 @@ impl EnvironmentService { settings, last_updated_at: Arc::new(RwLock::new(None)), environments, - persistence: None, + identity_store: None, } } @@ -53,8 +53,8 @@ impl EnvironmentService { service } - pub fn with_identity_persistence(mut self, persistence: Arc) -> Self { - self.persistence = Some(persistence); + pub fn with_identity_store(mut self, identity_store: Arc) -> Self { + self.identity_store = Some(identity_store); self } @@ -350,7 +350,7 @@ impl EnvironmentService { ) -> Result { let keys = self.resolve_key(environment_key)?; - let Some(persistence) = &self.persistence else { + let Some(identity_store) = &self.identity_store else { return self .evaluate_identity( &keys, @@ -369,7 +369,7 @@ impl EnvironmentService { .ok_or_else(|| { EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string()) })?; - let traits = persistence.load(&environment, identity).await; + let traits = identity_store.load_traits(&environment, identity).await; let response = self .evaluate_identity( &keys, @@ -385,7 +385,7 @@ impl EnvironmentService { } else { RequestOrigin::Server }; - persistence.persist(&environment, identity, origin); + identity_store.save_traits(&environment, identity, origin); Ok(response) } diff --git a/tests/test_identity_persistence.rs b/tests/test_identity_store.rs similarity index 73% rename from tests/test_identity_persistence.rs rename to tests/test_identity_store.rs index f92e438..3633eb0 100644 --- a/tests/test_identity_persistence.rs +++ b/tests/test_identity_store.rs @@ -3,7 +3,7 @@ mod fixtures; use async_trait::async_trait; use axum_test::TestServer; use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; -use edge_proxy::identities::{IdentityPersistence, IdentityTraits, RequestOrigin}; +use edge_proxy::identities::{IdentityStore, IdentityTraits, RequestOrigin}; use edge_proxy::models::{IdentityWithTraits, TraitModel}; use edge_proxy::routes::create_router; use fixtures::*; @@ -12,37 +12,46 @@ use serde_json::{Value, json}; use std::sync::{Arc, Mutex}; #[derive(Debug, PartialEq)] -struct PersistCall { +struct SaveCall { environment_api_key: Value, identifier: String, origin: RequestOrigin, } -struct RecordingPersistence { +struct RecordingIdentityStore { traits: IdentityTraits, - persist_calls: Mutex>, + save_calls: Mutex>, } -impl RecordingPersistence { +impl RecordingIdentityStore { fn new(evaluation: Vec, response: Vec) -> Self { Self { traits: IdentityTraits { evaluation, response, }, - persist_calls: Mutex::new(Vec::new()), + save_calls: Mutex::new(Vec::new()), } } } #[async_trait] -impl IdentityPersistence for RecordingPersistence { - async fn load(&self, _environment: &Value, _identity: &IdentityWithTraits) -> IdentityTraits { +impl IdentityStore for RecordingIdentityStore { + async fn load_traits( + &self, + _environment: &Value, + _identity: &IdentityWithTraits, + ) -> IdentityTraits { self.traits.clone() } - fn persist(&self, environment: &Value, identity: &IdentityWithTraits, origin: RequestOrigin) { - self.persist_calls.lock().unwrap().push(PersistCall { + fn save_traits( + &self, + environment: &Value, + identity: &IdentityWithTraits, + origin: RequestOrigin, + ) { + self.save_calls.lock().unwrap().push(SaveCall { environment_api_key: environment["api_key"].clone(), identifier: identity.identifier.clone(), origin, @@ -58,7 +67,7 @@ fn trait_model(trait_key: &str, trait_value: &str) -> TraitModel { } } -async fn setup_test_server(persistence: Option>) -> TestServer { +async fn setup_test_server(identity_store: Option>) -> TestServer { let settings = AppSettings { environment_key_pairs: vec![EnvironmentKeyPair { server_side_key: environment_1_server_key(), @@ -68,7 +77,7 @@ async fn setup_test_server(persistence: Option>) -> ..AppSettings::new() }; - let (app, service) = create_router(settings, persistence); + let (app, service) = create_router(settings, identity_store); service .cache @@ -101,7 +110,7 @@ fn feature_2_value(body: &Value) -> &Value { #[case::segment_match("test", "segment_override")] #[case::segment_miss("other", "2.3")] #[tokio::test] -async fn test_identity_without_persistence_evaluates_and_returns_request_traits( +async fn test_identity_without_store_evaluates_and_returns_request_traits( #[case] first_name: &str, #[case] expected_feature_2_value: &str, ) { @@ -121,17 +130,17 @@ async fn test_identity_without_persistence_evaluates_and_returns_request_traits( #[case::loaded_traits_match_segment("other", "test", "segment_override")] #[case::request_traits_ignored("test", "other", "2.3")] #[tokio::test] -async fn test_identity_with_persistence_evaluates_loaded_traits( +async fn test_identity_with_store_evaluates_loaded_traits( #[case] request_first_name: &str, #[case] loaded_first_name: &str, #[case] expected_feature_2_value: &str, ) { // Given - let persistence = Arc::new(RecordingPersistence::new( + let identity_store = Arc::new(RecordingIdentityStore::new( vec![trait_model("first_name", loaded_first_name)], Vec::new(), )); - let server = setup_test_server(Some(persistence)).await; + let server = setup_test_server(Some(identity_store)).await; // When let body = post_identity( @@ -146,13 +155,13 @@ async fn test_identity_with_persistence_evaluates_loaded_traits( } #[tokio::test] -async fn test_identity_with_persistence_returns_loaded_response_traits() { +async fn test_identity_with_store_returns_loaded_response_traits() { // Given - let persistence = Arc::new(RecordingPersistence::new( + let identity_store = Arc::new(RecordingIdentityStore::new( Vec::new(), vec![trait_model("stored_key", "stored_value")], )); - let server = setup_test_server(Some(persistence)).await; + let server = setup_test_server(Some(identity_store)).await; // When let body = post_identity( @@ -173,21 +182,21 @@ async fn test_identity_with_persistence_returns_loaded_response_traits() { #[case::client_key(environment_1_api_key(), RequestOrigin::Client)] #[case::server_key(environment_1_server_key(), RequestOrigin::Server)] #[tokio::test] -async fn test_identity_with_persistence_persists_once_with_origin( +async fn test_identity_with_store_saves_once_with_origin( #[case] environment_key: String, #[case] expected_origin: RequestOrigin, ) { // Given - let persistence = Arc::new(RecordingPersistence::new(Vec::new(), Vec::new())); - let server = setup_test_server(Some(persistence.clone())).await; + let identity_store = Arc::new(RecordingIdentityStore::new(Vec::new(), Vec::new())); + let server = setup_test_server(Some(identity_store.clone())).await; // When post_identity(&server, &environment_key, json!([])).await; // Then assert_eq!( - *persistence.persist_calls.lock().unwrap(), - vec![PersistCall { + *identity_store.save_calls.lock().unwrap(), + vec![SaveCall { environment_api_key: json!(environment_1_api_key()), identifier: "identity_1".to_string(), origin: expected_origin, From 34ae579e832661699756b8d4db206ca517012e8d Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 31 Oct 2026 16:31:39 +0530 Subject: [PATCH 3/6] refactor: split identity traits into traits and system traits --- src/identities/store.rs | 4 ++-- src/services/environment.rs | 23 ++++++++++++----------- tests/test_identity_store.rs | 32 +++++++++++++++++--------------- 3 files changed, 31 insertions(+), 28 deletions(-) diff --git a/src/identities/store.rs b/src/identities/store.rs index fdfaacb..24ce1ad 100644 --- a/src/identities/store.rs +++ b/src/identities/store.rs @@ -19,8 +19,8 @@ pub trait IdentityStore: Send + Sync { #[derive(Debug, Clone, PartialEq, Eq)] pub struct IdentityTraits { - pub evaluation: Vec, - pub response: Vec, + pub traits: Vec, + pub system_traits: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/services/environment.rs b/src/services/environment.rs index 6ac9970..d82f967 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -357,7 +357,7 @@ impl EnvironmentService { environment_key, &identity.identifier, &identity.traits, - &identity.traits, + &[], ) .await; }; @@ -375,8 +375,8 @@ impl EnvironmentService { &keys, environment_key, &identity.identifier, - &traits.evaluation, - &traits.response, + &traits.traits, + &traits.system_traits, ) .await?; @@ -395,8 +395,8 @@ impl EnvironmentService { keys: &EnvironmentKeys, environment_key: &str, identifier: &str, - evaluation_traits: &[TraitModel], - response_traits: &[TraitModel], + traits: &[TraitModel], + system_traits: &[TraitModel], ) -> Result { // Get pre-computed context from cache let context = self @@ -407,8 +407,12 @@ impl EnvironmentService { EdgeProxyError::ServiceUnavailable("Environment not loaded".to_string()) })?; - let flagsmith_traits: Vec = - evaluation_traits.iter().map(Into::into).collect(); + let flagsmith_traits: Vec = traits + .iter() + .filter(|t| !system_traits.iter().any(|s| s.trait_key == t.trait_key)) + .chain(system_traits) + .map(Into::into) + .collect(); let context_with_identity = add_identity_to_context(&context, identifier, &flagsmith_traits); @@ -429,10 +433,7 @@ impl EnvironmentService { Ok(IdentityResponse { flags: flag_results.iter().map(Into::into).collect(), - traits: response_traits - .iter() - .map(|t| t.to_response_json()) - .collect(), + traits: traits.iter().map(|t| t.to_response_json()).collect(), }) } diff --git a/tests/test_identity_store.rs b/tests/test_identity_store.rs index 3633eb0..6c0aa6c 100644 --- a/tests/test_identity_store.rs +++ b/tests/test_identity_store.rs @@ -24,11 +24,11 @@ struct RecordingIdentityStore { } impl RecordingIdentityStore { - fn new(evaluation: Vec, response: Vec) -> Self { + fn new(traits: Vec, system_traits: Vec) -> Self { Self { traits: IdentityTraits { - evaluation, - response, + traits, + system_traits, }, save_calls: Mutex::new(Vec::new()), } @@ -127,26 +127,28 @@ async fn test_identity_without_store_evaluates_and_returns_request_traits( } #[rstest] -#[case::loaded_traits_match_segment("other", "test", "segment_override")] -#[case::request_traits_ignored("test", "other", "2.3")] +#[case::loaded_trait_matches_segment(vec![trait_model("first_name", "test")], vec![], "segment_override")] +#[case::loaded_trait_misses_segment(vec![trait_model("first_name", "other")], vec![], "2.3")] +#[case::system_trait_wins_over_loaded_trait( + vec![trait_model("first_name", "other")], + vec![trait_model("first_name", "test")], + "segment_override" +)] #[tokio::test] -async fn test_identity_with_store_evaluates_loaded_traits( - #[case] request_first_name: &str, - #[case] loaded_first_name: &str, +async fn test_identity_with_store_evaluates_loaded_and_system_traits( + #[case] traits: Vec, + #[case] system_traits: Vec, #[case] expected_feature_2_value: &str, ) { // Given - let identity_store = Arc::new(RecordingIdentityStore::new( - vec![trait_model("first_name", loaded_first_name)], - Vec::new(), - )); + let identity_store = Arc::new(RecordingIdentityStore::new(traits, system_traits)); let server = setup_test_server(Some(identity_store)).await; // When let body = post_identity( &server, &environment_1_api_key(), - json!([{"trait_key": "first_name", "trait_value": request_first_name}]), + json!([{"trait_key": "first_name", "trait_value": "ignored"}]), ) .await; @@ -155,11 +157,11 @@ async fn test_identity_with_store_evaluates_loaded_traits( } #[tokio::test] -async fn test_identity_with_store_returns_loaded_response_traits() { +async fn test_identity_with_store_returns_loaded_traits_without_system_traits() { // Given let identity_store = Arc::new(RecordingIdentityStore::new( - Vec::new(), vec![trait_model("stored_key", "stored_value")], + vec![trait_model("system_key", "system_value")], )); let server = setup_test_server(Some(identity_store)).await; From ef8b56821b28765b3eb39dd7046012a01a4e3926 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 31 Oct 2026 16:34:03 +0530 Subject: [PATCH 4/6] refactor: rename RequestOrigin to EnvironmentKeyType --- src/identities/mod.rs | 2 +- src/identities/store.rs | 4 ++-- src/services/environment.rs | 10 +++++----- tests/test_identity_store.rs | 18 +++++++++--------- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/src/identities/mod.rs b/src/identities/mod.rs index aa867fd..b2f3715 100644 --- a/src/identities/mod.rs +++ b/src/identities/mod.rs @@ -1,3 +1,3 @@ pub mod store; -pub use store::{IdentityStore, IdentityTraits, RequestOrigin}; +pub use store::{EnvironmentKeyType, IdentityStore, IdentityTraits}; diff --git a/src/identities/store.rs b/src/identities/store.rs index 24ce1ad..6f68997 100644 --- a/src/identities/store.rs +++ b/src/identities/store.rs @@ -13,7 +13,7 @@ pub trait IdentityStore: Send + Sync { &self, environment: &serde_json::Value, identity: &IdentityWithTraits, - origin: RequestOrigin, + key_type: EnvironmentKeyType, ); } @@ -24,7 +24,7 @@ pub struct IdentityTraits { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum RequestOrigin { +pub enum EnvironmentKeyType { Client, Server, } diff --git a/src/services/environment.rs b/src/services/environment.rs index d82f967..4f78de4 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -2,7 +2,7 @@ use crate::cache::{EnvironmentsCache, LocalMemEnvironmentsCache}; use crate::config::settings::AppSettings; use crate::environments::{EnvironmentIndex, EnvironmentKeys}; use crate::error::{EdgeProxyError, Result}; -use crate::identities::{IdentityStore, RequestOrigin}; +use crate::identities::{EnvironmentKeyType, IdentityStore}; use crate::models::{ APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, TraitModel, }; @@ -380,12 +380,12 @@ impl EnvironmentService { ) .await?; - let origin = if environment_key == keys.client_key { - RequestOrigin::Client + let key_type = if environment_key == keys.client_key { + EnvironmentKeyType::Client } else { - RequestOrigin::Server + EnvironmentKeyType::Server }; - identity_store.save_traits(&environment, identity, origin); + identity_store.save_traits(&environment, identity, key_type); Ok(response) } diff --git a/tests/test_identity_store.rs b/tests/test_identity_store.rs index 6c0aa6c..f362586 100644 --- a/tests/test_identity_store.rs +++ b/tests/test_identity_store.rs @@ -3,7 +3,7 @@ mod fixtures; use async_trait::async_trait; use axum_test::TestServer; use edge_proxy::config::settings::{AppSettings, EnvironmentKeyPair}; -use edge_proxy::identities::{IdentityStore, IdentityTraits, RequestOrigin}; +use edge_proxy::identities::{EnvironmentKeyType, IdentityStore, IdentityTraits}; use edge_proxy::models::{IdentityWithTraits, TraitModel}; use edge_proxy::routes::create_router; use fixtures::*; @@ -15,7 +15,7 @@ use std::sync::{Arc, Mutex}; struct SaveCall { environment_api_key: Value, identifier: String, - origin: RequestOrigin, + key_type: EnvironmentKeyType, } struct RecordingIdentityStore { @@ -49,12 +49,12 @@ impl IdentityStore for RecordingIdentityStore { &self, environment: &Value, identity: &IdentityWithTraits, - origin: RequestOrigin, + key_type: EnvironmentKeyType, ) { self.save_calls.lock().unwrap().push(SaveCall { environment_api_key: environment["api_key"].clone(), identifier: identity.identifier.clone(), - origin, + key_type, }); } } @@ -181,12 +181,12 @@ async fn test_identity_with_store_returns_loaded_traits_without_system_traits() } #[rstest] -#[case::client_key(environment_1_api_key(), RequestOrigin::Client)] -#[case::server_key(environment_1_server_key(), RequestOrigin::Server)] +#[case::client_key(environment_1_api_key(), EnvironmentKeyType::Client)] +#[case::server_key(environment_1_server_key(), EnvironmentKeyType::Server)] #[tokio::test] -async fn test_identity_with_store_saves_once_with_origin( +async fn test_identity_with_store_saves_once_with_key_type( #[case] environment_key: String, - #[case] expected_origin: RequestOrigin, + #[case] expected_key_type: EnvironmentKeyType, ) { // Given let identity_store = Arc::new(RecordingIdentityStore::new(Vec::new(), Vec::new())); @@ -201,7 +201,7 @@ async fn test_identity_with_store_saves_once_with_origin( vec![SaveCall { environment_api_key: json!(environment_1_api_key()), identifier: "identity_1".to_string(), - origin: expected_origin, + key_type: expected_key_type, }] ); } From f2ca4b45a527687b2b423fdbe50b4c2a05ea725a Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 31 Oct 2026 16:58:16 +0530 Subject: [PATCH 5/6] test: rename identity store tests --- tests/test_identity_store.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_identity_store.rs b/tests/test_identity_store.rs index f362586..f76b04c 100644 --- a/tests/test_identity_store.rs +++ b/tests/test_identity_store.rs @@ -110,7 +110,7 @@ fn feature_2_value(body: &Value) -> &Value { #[case::segment_match("test", "segment_override")] #[case::segment_miss("other", "2.3")] #[tokio::test] -async fn test_identity_without_store_evaluates_and_returns_request_traits( +async fn test_post_identities_without_store_uses_request_traits( #[case] first_name: &str, #[case] expected_feature_2_value: &str, ) { @@ -127,15 +127,15 @@ async fn test_identity_without_store_evaluates_and_returns_request_traits( } #[rstest] -#[case::loaded_trait_matches_segment(vec![trait_model("first_name", "test")], vec![], "segment_override")] -#[case::loaded_trait_misses_segment(vec![trait_model("first_name", "other")], vec![], "2.3")] -#[case::system_trait_wins_over_loaded_trait( +#[case::trait_matches_segment(vec![trait_model("first_name", "test")], vec![], "segment_override")] +#[case::trait_misses_segment(vec![trait_model("first_name", "other")], vec![], "2.3")] +#[case::system_trait_overrides_trait( vec![trait_model("first_name", "other")], vec![trait_model("first_name", "test")], "segment_override" )] #[tokio::test] -async fn test_identity_with_store_evaluates_loaded_and_system_traits( +async fn test_post_identities_with_store_evaluates_store_traits_not_request_traits( #[case] traits: Vec, #[case] system_traits: Vec, #[case] expected_feature_2_value: &str, @@ -157,7 +157,7 @@ async fn test_identity_with_store_evaluates_loaded_and_system_traits( } #[tokio::test] -async fn test_identity_with_store_returns_loaded_traits_without_system_traits() { +async fn test_post_identities_with_store_returns_traits_not_request_or_system_traits() { // Given let identity_store = Arc::new(RecordingIdentityStore::new( vec![trait_model("stored_key", "stored_value")], @@ -184,7 +184,7 @@ async fn test_identity_with_store_returns_loaded_traits_without_system_traits() #[case::client_key(environment_1_api_key(), EnvironmentKeyType::Client)] #[case::server_key(environment_1_server_key(), EnvironmentKeyType::Server)] #[tokio::test] -async fn test_identity_with_store_saves_once_with_key_type( +async fn test_post_identities_with_store_saves_traits_once_with_key_type( #[case] environment_key: String, #[case] expected_key_type: EnvironmentKeyType, ) { From a73cffe9c8ab35d79538a478d9b9264e4c60d3e6 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Sat, 31 Oct 2026 17:44:35 +0530 Subject: [PATCH 6/6] refactor: extract identity store path from get_identity_response_data --- src/environments.rs | 9 ++++++++ src/services/environment.rs | 42 +++++++++++++++++++++++-------------- 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/src/environments.rs b/src/environments.rs index ac8a7fa..87d4f80 100644 --- a/src/environments.rs +++ b/src/environments.rs @@ -5,6 +5,7 @@ use chrono::{DateTime, Utc}; use parking_lot::RwLock; use crate::config::settings::EnvironmentKeyPair; +use crate::identities::EnvironmentKeyType; /// A server-side (`ser.`) key together with the validity metadata the /// proxy config endpoint reports. Statically configured keys carry no @@ -32,6 +33,14 @@ pub struct EnvironmentKeys { } impl EnvironmentKeys { + pub fn key_type(&self, environment_key: &str) -> EnvironmentKeyType { + if environment_key == self.client_key { + EnvironmentKeyType::Client + } else { + EnvironmentKeyType::Server + } + } + /// The first server-side key still usable for upstream fetches. pub fn valid_server_key(&self) -> Option<&ServerKey> { self.server_keys.iter().find(|key| key.is_valid()) diff --git a/src/services/environment.rs b/src/services/environment.rs index 4f78de4..d6c2628 100644 --- a/src/services/environment.rs +++ b/src/services/environment.rs @@ -2,7 +2,7 @@ use crate::cache::{EnvironmentsCache, LocalMemEnvironmentsCache}; use crate::config::settings::AppSettings; use crate::environments::{EnvironmentIndex, EnvironmentKeys}; use crate::error::{EdgeProxyError, Result}; -use crate::identities::{EnvironmentKeyType, IdentityStore}; +use crate::identities::IdentityStore; use crate::models::{ APIFeatureState, IdentityResponse, IdentityWithTraits, ProxyConfigEnvironment, TraitModel, }; @@ -349,19 +349,36 @@ impl EnvironmentService { environment_key: &str, ) -> Result { let keys = self.resolve_key(environment_key)?; - - let Some(identity_store) = &self.identity_store else { - return self - .evaluate_identity( + match &self.identity_store { + None => { + self.evaluate_identity( &keys, environment_key, &identity.identifier, &identity.traits, &[], ) - .await; - }; + .await + } + Some(identity_store) => { + self.evaluate_with_identity_store( + identity_store.as_ref(), + &keys, + environment_key, + identity, + ) + .await + } + } + } + async fn evaluate_with_identity_store( + &self, + identity_store: &dyn IdentityStore, + keys: &EnvironmentKeys, + environment_key: &str, + identity: &IdentityWithTraits, + ) -> Result { let environment = self .cache .get_environment(&keys.client_key) @@ -372,21 +389,14 @@ impl EnvironmentService { let traits = identity_store.load_traits(&environment, identity).await; let response = self .evaluate_identity( - &keys, + keys, environment_key, &identity.identifier, &traits.traits, &traits.system_traits, ) .await?; - - let key_type = if environment_key == keys.client_key { - EnvironmentKeyType::Client - } else { - EnvironmentKeyType::Server - }; - identity_store.save_traits(&environment, identity, key_type); - + identity_store.save_traits(&environment, identity, keys.key_type(environment_key)); Ok(response) }