chore: publish clean community baseline
This commit is contained in:
@@ -0,0 +1,292 @@
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{RuntimeError, WindowExecutionResult};
|
||||
|
||||
pub fn collect_window_result(
|
||||
response: &Value,
|
||||
config: &crank_core::StreamingConfig,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
let mut items = extract_path_value(response, config.items_path.as_deref())
|
||||
.and_then(|value| value.as_array().cloned())
|
||||
.unwrap_or_default();
|
||||
let cursor = extract_path_value(response, config.cursor_path.as_deref()).cloned();
|
||||
let done = extract_path_value(response, config.done_path.as_deref())
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(true);
|
||||
|
||||
let mut truncated = false;
|
||||
let mut has_more = !done;
|
||||
|
||||
if let Some(max_items) = config.max_items.map(|value| value as usize) {
|
||||
if items.len() > max_items {
|
||||
items.truncate(max_items);
|
||||
truncated = true;
|
||||
has_more = true;
|
||||
}
|
||||
}
|
||||
|
||||
let mut summary = build_summary(response, &items, config);
|
||||
redact_and_truncate(&mut summary, &mut items, config);
|
||||
|
||||
if let Some(max_bytes) = config.max_bytes.map(|value| value as usize) {
|
||||
while payload_size_bytes(&summary, &items, cursor.as_ref()) > max_bytes && !items.is_empty()
|
||||
{
|
||||
items.pop();
|
||||
truncated = true;
|
||||
has_more = true;
|
||||
}
|
||||
|
||||
if payload_size_bytes(&summary, &items, cursor.as_ref()) > max_bytes {
|
||||
summary = json!({
|
||||
"notice": "window payload exceeded byte limit",
|
||||
"items_returned": items.len()
|
||||
});
|
||||
truncated = true;
|
||||
has_more = true;
|
||||
}
|
||||
}
|
||||
|
||||
if matches!(
|
||||
config.aggregation_mode,
|
||||
crank_core::AggregationMode::SummaryOnly
|
||||
) {
|
||||
items.clear();
|
||||
}
|
||||
|
||||
Ok(WindowExecutionResult {
|
||||
summary,
|
||||
items,
|
||||
cursor,
|
||||
window_complete: !has_more,
|
||||
truncated,
|
||||
has_more,
|
||||
})
|
||||
}
|
||||
|
||||
fn build_summary(response: &Value, items: &[Value], config: &crank_core::StreamingConfig) -> Value {
|
||||
match config.aggregation_mode {
|
||||
crank_core::AggregationMode::RawItems => Value::Null,
|
||||
crank_core::AggregationMode::SummaryOnly
|
||||
| crank_core::AggregationMode::SummaryPlusSamples => {
|
||||
extract_path_value(response, config.summary_path.as_deref())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| stats_summary(items))
|
||||
}
|
||||
crank_core::AggregationMode::Stats => stats_summary(items),
|
||||
crank_core::AggregationMode::LatestState => items.last().cloned().unwrap_or(Value::Null),
|
||||
}
|
||||
}
|
||||
|
||||
fn stats_summary(items: &[Value]) -> Value {
|
||||
let mut summary = Map::new();
|
||||
summary.insert("items_total".to_owned(), Value::from(items.len() as u64));
|
||||
summary.insert("items_sampled".to_owned(), Value::from(items.len() as u64));
|
||||
Value::Object(summary)
|
||||
}
|
||||
|
||||
fn redact_and_truncate(
|
||||
summary: &mut Value,
|
||||
items: &mut [Value],
|
||||
config: &crank_core::StreamingConfig,
|
||||
) {
|
||||
crate::redaction::redact_paths(summary, &config.redacted_paths);
|
||||
|
||||
for item in items.iter_mut() {
|
||||
crate::redaction::redact_paths(item, &config.redacted_paths);
|
||||
}
|
||||
|
||||
if config.truncate_item_fields {
|
||||
if let Some(max_len) = config.max_field_length.map(|value| value as usize) {
|
||||
for item in items.iter_mut() {
|
||||
crate::redaction::truncate_item_fields(item, max_len);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn payload_size_bytes(summary: &Value, items: &[Value], cursor: Option<&Value>) -> usize {
|
||||
serde_json::to_vec(&json!({
|
||||
"summary": summary,
|
||||
"items": items,
|
||||
"cursor": cursor
|
||||
}))
|
||||
.map(|value| value.len())
|
||||
.unwrap_or(usize::MAX)
|
||||
}
|
||||
|
||||
pub fn extract_path_value<'a>(value: &'a Value, path: Option<&str>) -> Option<&'a Value> {
|
||||
let path = path?;
|
||||
let path = path.strip_prefix("$.").or_else(|| path.strip_prefix('$'))?;
|
||||
|
||||
if path.is_empty() {
|
||||
return Some(value);
|
||||
}
|
||||
|
||||
let mut current = value;
|
||||
for segment in path.split('.') {
|
||||
current = current.get(segment)?;
|
||||
}
|
||||
|
||||
Some(current)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crank_core::{
|
||||
AggregationMode, ExecutionMode, StreamingConfig, ToolFamilyConfig, TransportBehavior,
|
||||
};
|
||||
|
||||
use super::collect_window_result;
|
||||
|
||||
fn window_config() -> StreamingConfig {
|
||||
StreamingConfig {
|
||||
mode: ExecutionMode::Window,
|
||||
transport_behavior: TransportBehavior::ServerStream,
|
||||
window_duration_ms: Some(3_000),
|
||||
poll_interval_ms: None,
|
||||
upstream_timeout_ms: Some(5_000),
|
||||
idle_timeout_ms: None,
|
||||
max_session_lifetime_ms: None,
|
||||
max_items: Some(10),
|
||||
max_bytes: Some(10_000),
|
||||
aggregation_mode: AggregationMode::RawItems,
|
||||
summary_path: Some("$.summary".to_owned()),
|
||||
items_path: Some("$.items".to_owned()),
|
||||
cursor_path: Some("$.cursor".to_owned()),
|
||||
status_path: None,
|
||||
done_path: Some("$.done".to_owned()),
|
||||
redacted_paths: Vec::new(),
|
||||
truncate_item_fields: false,
|
||||
max_field_length: None,
|
||||
drop_duplicates: false,
|
||||
sampling_rate: None,
|
||||
tool_family: ToolFamilyConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collects_raw_items_mode() {
|
||||
let config = window_config();
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }, { "message": "two" }],
|
||||
"summary": { "count": 2 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.summary, Value::Null);
|
||||
assert_eq!(result.items.len(), 2);
|
||||
assert!(result.window_complete);
|
||||
assert!(!result.truncated);
|
||||
assert!(!result.has_more);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_summary_only_mode() {
|
||||
let mut config = window_config();
|
||||
config.aggregation_mode = AggregationMode::SummaryOnly;
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }],
|
||||
"summary": { "count": 1 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.summary, json!({ "count": 1 }));
|
||||
assert!(result.items.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_summary_plus_samples() {
|
||||
let mut config = window_config();
|
||||
config.aggregation_mode = AggregationMode::SummaryPlusSamples;
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }, { "message": "two" }],
|
||||
"summary": { "count": 2 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.summary, json!({ "count": 2 }));
|
||||
assert_eq!(result.items.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_by_item_count() {
|
||||
let mut config = window_config();
|
||||
config.max_items = Some(1);
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }, { "message": "two" }],
|
||||
"cursor": "next",
|
||||
"done": false
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.items.len(), 1);
|
||||
assert!(result.truncated);
|
||||
assert!(result.has_more);
|
||||
assert!(!result.window_complete);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_by_byte_limit() {
|
||||
let mut config = window_config();
|
||||
config.max_bytes = Some(120);
|
||||
config.aggregation_mode = AggregationMode::SummaryPlusSamples;
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [
|
||||
{ "message": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" },
|
||||
{ "message": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" }
|
||||
],
|
||||
"summary": { "count": 2 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(result.truncated);
|
||||
assert!(result.items.len() < 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redacts_and_truncates_values() {
|
||||
let mut config = window_config();
|
||||
config.redacted_paths = vec!["$.secret".to_owned()];
|
||||
config.truncate_item_fields = true;
|
||||
config.max_field_length = Some(4);
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [
|
||||
{ "message": "abcdefghijklmnop", "secret": "token" }
|
||||
],
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.items[0]["secret"], json!("[REDACTED]"));
|
||||
assert_eq!(result.items[0]["message"], json!("abcd..."));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use crank_core::{AuthConfig, AuthProfile, SecretId};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{PreparedRequest, RuntimeError};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ResolvedAuth {
|
||||
Bearer { header_name: String, token: String },
|
||||
Basic { username: String, password: String },
|
||||
ApiKeyHeader { header_name: String, value: String },
|
||||
ApiKeyQuery { param_name: String, value: String },
|
||||
}
|
||||
|
||||
impl ResolvedAuth {
|
||||
pub fn from_profile(
|
||||
profile: &AuthProfile,
|
||||
secrets: &BTreeMap<SecretId, Value>,
|
||||
) -> Result<Self, RuntimeError> {
|
||||
match &profile.config {
|
||||
AuthConfig::Bearer(config) => Ok(Self::Bearer {
|
||||
header_name: config.header_name.clone(),
|
||||
token: extract_secret_string(secrets, &config.secret_id, &["token", "value"])?,
|
||||
}),
|
||||
AuthConfig::Basic(config) => Ok(Self::Basic {
|
||||
username: extract_secret_string(
|
||||
secrets,
|
||||
&config.username_secret_id,
|
||||
&["username", "value"],
|
||||
)?,
|
||||
password: extract_secret_string(
|
||||
secrets,
|
||||
&config.password_secret_id,
|
||||
&["password", "value"],
|
||||
)?,
|
||||
}),
|
||||
AuthConfig::ApiKeyHeader(config) => Ok(Self::ApiKeyHeader {
|
||||
header_name: config.header_name.clone(),
|
||||
value: extract_secret_string(secrets, &config.secret_id, &["token", "value"])?,
|
||||
}),
|
||||
AuthConfig::ApiKeyQuery(config) => Ok(Self::ApiKeyQuery {
|
||||
param_name: config.param_name.clone(),
|
||||
value: extract_secret_string(secrets, &config.secret_id, &["token", "value"])?,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn apply(&self, prepared_request: &mut PreparedRequest) {
|
||||
match self {
|
||||
Self::Bearer { header_name, token } => {
|
||||
prepared_request
|
||||
.headers
|
||||
.insert(header_name.clone(), format!("Bearer {token}"));
|
||||
}
|
||||
Self::Basic { username, password } => {
|
||||
let credentials = STANDARD.encode(format!("{username}:{password}"));
|
||||
prepared_request
|
||||
.headers
|
||||
.insert("Authorization".to_owned(), format!("Basic {credentials}"));
|
||||
}
|
||||
Self::ApiKeyHeader { header_name, value } => {
|
||||
prepared_request
|
||||
.headers
|
||||
.insert(header_name.clone(), value.clone());
|
||||
}
|
||||
Self::ApiKeyQuery { param_name, value } => {
|
||||
prepared_request
|
||||
.query_params
|
||||
.insert(param_name.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_secret_string(
|
||||
secrets: &BTreeMap<SecretId, Value>,
|
||||
secret_id: &SecretId,
|
||||
preferred_fields: &[&str],
|
||||
) -> Result<String, RuntimeError> {
|
||||
let Some(value) = secrets.get(secret_id) else {
|
||||
return Err(RuntimeError::MissingSecret {
|
||||
secret_id: secret_id.as_str().to_owned(),
|
||||
});
|
||||
};
|
||||
|
||||
if let Some(text) = value.as_str() {
|
||||
return Ok(text.to_owned());
|
||||
}
|
||||
|
||||
let Some(object) = value.as_object() else {
|
||||
return Err(RuntimeError::InvalidAuthSecretValue {
|
||||
secret_id: secret_id.as_str().to_owned(),
|
||||
reason: "secret payload must be a string or object".to_owned(),
|
||||
});
|
||||
};
|
||||
|
||||
for field in preferred_fields {
|
||||
if let Some(text) = object.get(*field).and_then(Value::as_str) {
|
||||
return Ok(text.to_owned());
|
||||
}
|
||||
}
|
||||
|
||||
Err(RuntimeError::InvalidAuthSecretValue {
|
||||
secret_id: secret_id.as_str().to_owned(),
|
||||
reason: format!(
|
||||
"secret payload must contain one of: {}",
|
||||
preferred_fields.join(", ")
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crank_core::{
|
||||
ApiKeyHeaderAuthConfig, ApiKeyQueryAuthConfig, AuthConfig, AuthKind, AuthProfile,
|
||||
BasicAuthConfig, BearerAuthConfig, SecretId, WorkspaceId,
|
||||
};
|
||||
use serde_json::json;
|
||||
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
|
||||
|
||||
use crate::{PreparedRequest, auth::ResolvedAuth};
|
||||
|
||||
fn timestamp(value: &str) -> OffsetDateTime {
|
||||
OffsetDateTime::parse(value, &Rfc3339).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_bearer_auth_to_headers() {
|
||||
let auth = ResolvedAuth::from_profile(
|
||||
&AuthProfile {
|
||||
id: "auth_01".into(),
|
||||
workspace_id: WorkspaceId::new("ws_default"),
|
||||
name: "bearer".to_owned(),
|
||||
kind: AuthKind::Bearer,
|
||||
config: AuthConfig::Bearer(BearerAuthConfig {
|
||||
header_name: "Authorization".to_owned(),
|
||||
secret_id: SecretId::new("secret_token"),
|
||||
}),
|
||||
created_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
updated_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
},
|
||||
&BTreeMap::from([(SecretId::new("secret_token"), json!({"token":"abc"}))]),
|
||||
)
|
||||
.unwrap();
|
||||
let mut request = PreparedRequest::default();
|
||||
|
||||
auth.apply(&mut request);
|
||||
|
||||
assert_eq!(
|
||||
request.headers.get("Authorization").map(String::as_str),
|
||||
Some("Bearer abc")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_basic_auth_to_headers() {
|
||||
let auth = ResolvedAuth::from_profile(
|
||||
&AuthProfile {
|
||||
id: "auth_01".into(),
|
||||
workspace_id: WorkspaceId::new("ws_default"),
|
||||
name: "basic".to_owned(),
|
||||
kind: AuthKind::Basic,
|
||||
config: AuthConfig::Basic(BasicAuthConfig {
|
||||
username_secret_id: SecretId::new("secret_user"),
|
||||
password_secret_id: SecretId::new("secret_pass"),
|
||||
}),
|
||||
created_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
updated_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
},
|
||||
&BTreeMap::from([
|
||||
(SecretId::new("secret_user"), json!({"username":"demo"})),
|
||||
(SecretId::new("secret_pass"), json!({"password":"s3cr3t"})),
|
||||
]),
|
||||
)
|
||||
.unwrap();
|
||||
let mut request = PreparedRequest::default();
|
||||
|
||||
auth.apply(&mut request);
|
||||
|
||||
assert_eq!(
|
||||
request.headers.get("Authorization").map(String::as_str),
|
||||
Some("Basic ZGVtbzpzM2NyM3Q=")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_api_key_header_auth() {
|
||||
let auth = ResolvedAuth::from_profile(
|
||||
&AuthProfile {
|
||||
id: "auth_01".into(),
|
||||
workspace_id: WorkspaceId::new("ws_default"),
|
||||
name: "api-header".to_owned(),
|
||||
kind: AuthKind::ApiKeyHeader,
|
||||
config: AuthConfig::ApiKeyHeader(ApiKeyHeaderAuthConfig {
|
||||
header_name: "X-Api-Key".to_owned(),
|
||||
secret_id: SecretId::new("secret_api_key"),
|
||||
}),
|
||||
created_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
updated_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
},
|
||||
&BTreeMap::from([(SecretId::new("secret_api_key"), json!("key-123"))]),
|
||||
)
|
||||
.unwrap();
|
||||
let mut request = PreparedRequest::default();
|
||||
|
||||
auth.apply(&mut request);
|
||||
|
||||
assert_eq!(
|
||||
request.headers.get("X-Api-Key").map(String::as_str),
|
||||
Some("key-123")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_api_key_query_auth() {
|
||||
let auth = ResolvedAuth::from_profile(
|
||||
&AuthProfile {
|
||||
id: "auth_01".into(),
|
||||
workspace_id: WorkspaceId::new("ws_default"),
|
||||
name: "api-query".to_owned(),
|
||||
kind: AuthKind::ApiKeyQuery,
|
||||
config: AuthConfig::ApiKeyQuery(ApiKeyQueryAuthConfig {
|
||||
param_name: "api_key".to_owned(),
|
||||
secret_id: SecretId::new("secret_api_key"),
|
||||
}),
|
||||
created_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
updated_at: timestamp("2026-04-07T00:00:00Z"),
|
||||
},
|
||||
&BTreeMap::from([(SecretId::new("secret_api_key"), json!({"value":"key-123"}))]),
|
||||
)
|
||||
.unwrap();
|
||||
let mut request = PreparedRequest::default();
|
||||
|
||||
auth.apply(&mut request);
|
||||
|
||||
assert_eq!(
|
||||
request.query_params.get("api_key").map(String::as_str),
|
||||
Some("key-123")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,903 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
env,
|
||||
num::ParseIntError,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use crank_core::{
|
||||
CacheBackend, CacheScope, CacheStoreError, CachedResponse, CoordinationStateStore,
|
||||
CoordinationStateValue, RateLimitBucketState, RateLimitStateStore, ReplayGuardStatus,
|
||||
ReplayGuardStore, ResponseCacheStore,
|
||||
};
|
||||
use redis::{Client, aio::ConnectionManager};
|
||||
use serde::de::DeserializeOwned;
|
||||
use thiserror::Error;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RuntimeCacheConfig {
|
||||
pub backend: CacheBackend,
|
||||
pub url: Option<String>,
|
||||
pub default_ttl_ms: Option<u64>,
|
||||
}
|
||||
|
||||
impl Default for RuntimeCacheConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
backend: CacheBackend::Memory,
|
||||
url: None,
|
||||
default_ttl_ms: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RuntimeCacheConfig {
|
||||
pub fn from_env() -> Result<Self, RuntimeCacheConfigError> {
|
||||
let backend = parse_backend()?;
|
||||
let url = parse_optional_string("CRANK_CACHE_URL")?;
|
||||
let default_ttl_ms = parse_optional_u64("CRANK_CACHE_DEFAULT_TTL_MS")?;
|
||||
|
||||
if backend.is_external() && url.is_none() {
|
||||
return Err(RuntimeCacheConfigError::MissingUrl { backend });
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
backend,
|
||||
url,
|
||||
default_ttl_ms,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RuntimeCacheStores {
|
||||
pub backend: CacheBackend,
|
||||
pub response: Arc<dyn ResponseCacheStore>,
|
||||
pub rate_limit: Arc<dyn RateLimitStateStore>,
|
||||
pub replay_guard: Arc<dyn ReplayGuardStore>,
|
||||
pub coordination: Arc<dyn CoordinationStateStore>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RedisCacheStore {
|
||||
backend: CacheBackend,
|
||||
connection_manager: ConnectionManager,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct InMemoryResponseCacheStore {
|
||||
entries: Arc<RwLock<HashMap<String, ExpiringValue<CachedResponse>>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct InMemoryRateLimitStateStore {
|
||||
entries: Arc<RwLock<HashMap<String, ExpiringValue<RateLimitBucketState>>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct InMemoryReplayGuardStore {
|
||||
entries: Arc<RwLock<HashMap<String, Instant>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct InMemoryCoordinationStateStore {
|
||||
entries: Arc<RwLock<HashMap<String, ExpiringValue<CoordinationStateValue>>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct ExpiringValue<T> {
|
||||
value: T,
|
||||
expires_at: Instant,
|
||||
}
|
||||
|
||||
impl RuntimeCacheStores {
|
||||
pub async fn from_config(
|
||||
config: &RuntimeCacheConfig,
|
||||
) -> Result<Self, RuntimeCacheStoreInitError> {
|
||||
match config.backend {
|
||||
CacheBackend::Memory => {
|
||||
let response = Arc::new(InMemoryResponseCacheStore::default());
|
||||
let rate_limit = Arc::new(InMemoryRateLimitStateStore::default());
|
||||
let replay_guard = Arc::new(InMemoryReplayGuardStore::default());
|
||||
let coordination = Arc::new(InMemoryCoordinationStateStore::default());
|
||||
Ok(Self {
|
||||
backend: CacheBackend::Memory,
|
||||
response,
|
||||
rate_limit,
|
||||
replay_guard,
|
||||
coordination,
|
||||
})
|
||||
}
|
||||
CacheBackend::Valkey | CacheBackend::Redis => {
|
||||
let url = config
|
||||
.url
|
||||
.as_deref()
|
||||
.ok_or(RuntimeCacheStoreInitError::MissingUrl {
|
||||
backend: config.backend,
|
||||
})?;
|
||||
let store = Arc::new(RedisCacheStore::connect(config.backend, url).await?);
|
||||
Ok(Self {
|
||||
backend: config.backend,
|
||||
response: store.clone(),
|
||||
rate_limit: store.clone(),
|
||||
replay_guard: store.clone(),
|
||||
coordination: store,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisCacheStore {
|
||||
pub async fn connect(
|
||||
backend: CacheBackend,
|
||||
url: &str,
|
||||
) -> Result<Self, RuntimeCacheStoreInitError> {
|
||||
let client =
|
||||
Client::open(url).map_err(|source| RuntimeCacheStoreInitError::InvalidUrl {
|
||||
backend,
|
||||
url: url.to_owned(),
|
||||
details: source.to_string(),
|
||||
})?;
|
||||
let connection_manager = client.get_connection_manager().await.map_err(|source| {
|
||||
RuntimeCacheStoreInitError::ConnectFailed {
|
||||
backend,
|
||||
details: source.to_string(),
|
||||
}
|
||||
})?;
|
||||
Ok(Self {
|
||||
backend,
|
||||
connection_manager,
|
||||
})
|
||||
}
|
||||
|
||||
fn prefixed_key(&self, kind: &str, key: &str) -> String {
|
||||
format!("crank:{kind}:{key}")
|
||||
}
|
||||
|
||||
fn serialize_value<T: serde::Serialize>(&self, value: &T) -> Result<Vec<u8>, CacheStoreError> {
|
||||
serde_json::to_vec(value).map_err(|source| CacheStoreError::Serialization {
|
||||
message: source.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn deserialize_value<T: DeserializeOwned>(&self, value: &[u8]) -> Result<T, CacheStoreError> {
|
||||
serde_json::from_slice(value).map_err(|source| CacheStoreError::Serialization {
|
||||
message: source.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn ttl_ms(&self, ttl: Duration) -> Result<u64, CacheStoreError> {
|
||||
if ttl.is_zero() {
|
||||
return Err(CacheStoreError::InvalidKey {
|
||||
message: "cache ttl must be greater than zero".to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX))
|
||||
}
|
||||
|
||||
fn unavailable(&self, source: redis::RedisError) -> CacheStoreError {
|
||||
CacheStoreError::Unavailable {
|
||||
message: format!("{} backend error: {source}", self.backend),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_json<T: DeserializeOwned>(
|
||||
&self,
|
||||
kind: &str,
|
||||
key: &str,
|
||||
) -> Result<Option<T>, CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let storage_key = self.prefixed_key(kind, key);
|
||||
let mut connection = self.connection_manager.clone();
|
||||
let encoded: Option<Vec<u8>> = redis::cmd("GET")
|
||||
.arg(storage_key)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_err(|source| self.unavailable(source))?;
|
||||
encoded
|
||||
.map(|bytes| self.deserialize_value(&bytes))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn put_json<T: serde::Serialize>(
|
||||
&self,
|
||||
kind: &str,
|
||||
key: &str,
|
||||
value: &T,
|
||||
ttl: Duration,
|
||||
) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let storage_key = self.prefixed_key(kind, key);
|
||||
let encoded = self.serialize_value(value)?;
|
||||
let ttl_ms = self.ttl_ms(ttl)?;
|
||||
let mut connection = self.connection_manager.clone();
|
||||
redis::cmd("PSETEX")
|
||||
.arg(storage_key)
|
||||
.arg(ttl_ms)
|
||||
.arg(encoded)
|
||||
.query_async::<()>(&mut connection)
|
||||
.await
|
||||
.map_err(|source| self.unavailable(source))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_kind_key(&self, kind: &str, key: &str) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let storage_key = self.prefixed_key(kind, key);
|
||||
let mut connection = self.connection_manager.clone();
|
||||
redis::cmd("DEL")
|
||||
.arg(storage_key)
|
||||
.query_async::<()>(&mut connection)
|
||||
.await
|
||||
.map_err(|source| self.unavailable(source))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn coordination_kind(scope: CacheScope) -> &'static str {
|
||||
match scope {
|
||||
CacheScope::Response => "coordination:response",
|
||||
CacheScope::RateLimit => "coordination:rate_limit",
|
||||
CacheScope::ReplayGuard => "coordination:replay_guard",
|
||||
CacheScope::Coordination => "coordination:coordination",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ResponseCacheStore for InMemoryResponseCacheStore {
|
||||
async fn get(&self, key: &str) -> Result<Option<CachedResponse>, CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let now = Instant::now();
|
||||
let mut entries = self.entries.write().await;
|
||||
retain_unexpired(&mut entries, now);
|
||||
Ok(entries.get(key).map(|entry| entry.value.clone()))
|
||||
}
|
||||
|
||||
async fn put(
|
||||
&self,
|
||||
key: &str,
|
||||
value: CachedResponse,
|
||||
ttl: Duration,
|
||||
) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let expires_at = expiry_from_ttl(ttl)?;
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.insert(key.to_owned(), ExpiringValue { value, expires_at });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete(&self, key: &str) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.remove(key);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RateLimitStateStore for InMemoryRateLimitStateStore {
|
||||
async fn get_bucket(&self, key: &str) -> Result<Option<RateLimitBucketState>, CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let now = Instant::now();
|
||||
let mut entries = self.entries.write().await;
|
||||
retain_unexpired(&mut entries, now);
|
||||
Ok(entries.get(key).map(|entry| entry.value))
|
||||
}
|
||||
|
||||
async fn put_bucket(
|
||||
&self,
|
||||
key: &str,
|
||||
value: RateLimitBucketState,
|
||||
ttl: Duration,
|
||||
) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let expires_at = expiry_from_ttl(ttl)?;
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.insert(key.to_owned(), ExpiringValue { value, expires_at });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_bucket(&self, key: &str) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.remove(key);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ReplayGuardStore for InMemoryReplayGuardStore {
|
||||
async fn mark_seen(
|
||||
&self,
|
||||
key: &str,
|
||||
ttl: Duration,
|
||||
) -> Result<ReplayGuardStatus, CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let expires_at = expiry_from_ttl(ttl)?;
|
||||
let now = Instant::now();
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.retain(|_, entry_expires_at| *entry_expires_at > now);
|
||||
if entries.get(key).is_some() {
|
||||
return Ok(ReplayGuardStatus::AlreadySeen);
|
||||
}
|
||||
entries.insert(key.to_owned(), expires_at);
|
||||
Ok(ReplayGuardStatus::Fresh)
|
||||
}
|
||||
|
||||
async fn clear(&self, key: &str) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.remove(key);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CoordinationStateStore for InMemoryCoordinationStateStore {
|
||||
async fn get_value(
|
||||
&self,
|
||||
scope: CacheScope,
|
||||
key: &str,
|
||||
) -> Result<Option<CoordinationStateValue>, CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let storage_key = scoped_key(scope, key);
|
||||
let now = Instant::now();
|
||||
let mut entries = self.entries.write().await;
|
||||
retain_unexpired(&mut entries, now);
|
||||
Ok(entries.get(&storage_key).map(|entry| entry.value.clone()))
|
||||
}
|
||||
|
||||
async fn put_value(
|
||||
&self,
|
||||
scope: CacheScope,
|
||||
key: &str,
|
||||
value: CoordinationStateValue,
|
||||
ttl: Duration,
|
||||
) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let expires_at = expiry_from_ttl(ttl)?;
|
||||
let storage_key = scoped_key(scope, key);
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.insert(storage_key, ExpiringValue { value, expires_at });
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_value(&self, scope: CacheScope, key: &str) -> Result<(), CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let storage_key = scoped_key(scope, key);
|
||||
let mut entries = self.entries.write().await;
|
||||
entries.remove(&storage_key);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ResponseCacheStore for RedisCacheStore {
|
||||
async fn get(&self, key: &str) -> Result<Option<CachedResponse>, CacheStoreError> {
|
||||
self.get_json("response", key).await
|
||||
}
|
||||
|
||||
async fn put(
|
||||
&self,
|
||||
key: &str,
|
||||
value: CachedResponse,
|
||||
ttl: Duration,
|
||||
) -> Result<(), CacheStoreError> {
|
||||
self.put_json("response", key, &value, ttl).await
|
||||
}
|
||||
|
||||
async fn delete(&self, key: &str) -> Result<(), CacheStoreError> {
|
||||
self.delete_kind_key("response", key).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RateLimitStateStore for RedisCacheStore {
|
||||
async fn get_bucket(&self, key: &str) -> Result<Option<RateLimitBucketState>, CacheStoreError> {
|
||||
self.get_json("rate_limit", key).await
|
||||
}
|
||||
|
||||
async fn put_bucket(
|
||||
&self,
|
||||
key: &str,
|
||||
value: RateLimitBucketState,
|
||||
ttl: Duration,
|
||||
) -> Result<(), CacheStoreError> {
|
||||
self.put_json("rate_limit", key, &value, ttl).await
|
||||
}
|
||||
|
||||
async fn delete_bucket(&self, key: &str) -> Result<(), CacheStoreError> {
|
||||
self.delete_kind_key("rate_limit", key).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ReplayGuardStore for RedisCacheStore {
|
||||
async fn mark_seen(
|
||||
&self,
|
||||
key: &str,
|
||||
ttl: Duration,
|
||||
) -> Result<ReplayGuardStatus, CacheStoreError> {
|
||||
validate_key(key)?;
|
||||
let storage_key = self.prefixed_key("replay_guard", key);
|
||||
let ttl_ms = self.ttl_ms(ttl)?;
|
||||
let mut connection = self.connection_manager.clone();
|
||||
let result: Option<String> = redis::cmd("SET")
|
||||
.arg(storage_key)
|
||||
.arg("1")
|
||||
.arg("PX")
|
||||
.arg(ttl_ms)
|
||||
.arg("NX")
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_err(|source| self.unavailable(source))?;
|
||||
Ok(if result.is_some() {
|
||||
ReplayGuardStatus::Fresh
|
||||
} else {
|
||||
ReplayGuardStatus::AlreadySeen
|
||||
})
|
||||
}
|
||||
|
||||
async fn clear(&self, key: &str) -> Result<(), CacheStoreError> {
|
||||
self.delete_kind_key("replay_guard", key).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CoordinationStateStore for RedisCacheStore {
|
||||
async fn get_value(
|
||||
&self,
|
||||
scope: CacheScope,
|
||||
key: &str,
|
||||
) -> Result<Option<CoordinationStateValue>, CacheStoreError> {
|
||||
self.get_json(Self::coordination_kind(scope), key).await
|
||||
}
|
||||
|
||||
async fn put_value(
|
||||
&self,
|
||||
scope: CacheScope,
|
||||
key: &str,
|
||||
value: CoordinationStateValue,
|
||||
ttl: Duration,
|
||||
) -> Result<(), CacheStoreError> {
|
||||
self.put_json(Self::coordination_kind(scope), key, &value, ttl)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn delete_value(&self, scope: CacheScope, key: &str) -> Result<(), CacheStoreError> {
|
||||
self.delete_kind_key(Self::coordination_kind(scope), key)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_key(key: &str) -> Result<(), CacheStoreError> {
|
||||
if key.trim().is_empty() {
|
||||
return Err(CacheStoreError::InvalidKey {
|
||||
message: "cache key must not be empty".to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn expiry_from_ttl(ttl: Duration) -> Result<Instant, CacheStoreError> {
|
||||
if ttl.is_zero() {
|
||||
return Err(CacheStoreError::InvalidKey {
|
||||
message: "cache ttl must be greater than zero".to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(Instant::now() + ttl)
|
||||
}
|
||||
|
||||
fn retain_unexpired<T>(entries: &mut HashMap<String, ExpiringValue<T>>, now: Instant) {
|
||||
entries.retain(|_, entry| entry.expires_at > now);
|
||||
}
|
||||
|
||||
fn scoped_key(scope: CacheScope, key: &str) -> String {
|
||||
format!("{scope:?}:{key}")
|
||||
}
|
||||
|
||||
fn parse_backend() -> Result<CacheBackend, RuntimeCacheConfigError> {
|
||||
match env::var("CRANK_CACHE_BACKEND") {
|
||||
Ok(raw) => raw
|
||||
.parse::<CacheBackend>()
|
||||
.map_err(|source| RuntimeCacheConfigError::InvalidBackend { value: raw, source }),
|
||||
Err(env::VarError::NotPresent) => Ok(CacheBackend::Memory),
|
||||
Err(env::VarError::NotUnicode(_)) => Err(RuntimeCacheConfigError::InvalidUnicode {
|
||||
name: "CRANK_CACHE_BACKEND",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_optional_string(name: &'static str) -> Result<Option<String>, RuntimeCacheConfigError> {
|
||||
match env::var(name) {
|
||||
Ok(raw) => {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(trimmed.to_owned()))
|
||||
}
|
||||
}
|
||||
Err(env::VarError::NotPresent) => Ok(None),
|
||||
Err(env::VarError::NotUnicode(_)) => Err(RuntimeCacheConfigError::InvalidUnicode { name }),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_optional_u64(name: &'static str) -> Result<Option<u64>, RuntimeCacheConfigError> {
|
||||
match env::var(name) {
|
||||
Ok(raw) => {
|
||||
let value = raw
|
||||
.parse::<u64>()
|
||||
.map_err(|source| RuntimeCacheConfigError::InvalidTtl { value: raw, source })?;
|
||||
if value == 0 {
|
||||
return Err(RuntimeCacheConfigError::ZeroTtl { name });
|
||||
}
|
||||
Ok(Some(value))
|
||||
}
|
||||
Err(env::VarError::NotPresent) => Ok(None),
|
||||
Err(env::VarError::NotUnicode(_)) => Err(RuntimeCacheConfigError::InvalidUnicode { name }),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum RuntimeCacheConfigError {
|
||||
#[error("{name} must contain valid UTF-8")]
|
||||
InvalidUnicode { name: &'static str },
|
||||
#[error("CRANK_CACHE_BACKEND must be one of memory, valkey, redis, got {value}")]
|
||||
InvalidBackend {
|
||||
value: String,
|
||||
source: crank_core::ParseCacheBackendError,
|
||||
},
|
||||
#[error("CRANK_CACHE_DEFAULT_TTL_MS must be a positive integer, got {value}")]
|
||||
InvalidTtl {
|
||||
value: String,
|
||||
source: ParseIntError,
|
||||
},
|
||||
#[error("{name} must be greater than zero")]
|
||||
ZeroTtl { name: &'static str },
|
||||
#[error("{backend} backend requires CRANK_CACHE_URL")]
|
||||
MissingUrl { backend: CacheBackend },
|
||||
}
|
||||
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum RuntimeCacheStoreInitError {
|
||||
#[error("{backend} backend requires CRANK_CACHE_URL")]
|
||||
MissingUrl { backend: CacheBackend },
|
||||
#[error("invalid {backend} cache url {url}: {details}")]
|
||||
InvalidUrl {
|
||||
backend: CacheBackend,
|
||||
url: String,
|
||||
details: String,
|
||||
},
|
||||
#[error("failed to connect to {backend} cache backend: {details}")]
|
||||
ConnectFailed {
|
||||
backend: CacheBackend,
|
||||
details: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use crank_core::CacheBackend;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
InMemoryCoordinationStateStore, InMemoryRateLimitStateStore, InMemoryReplayGuardStore,
|
||||
InMemoryResponseCacheStore, RedisCacheStore, RuntimeCacheConfig, RuntimeCacheConfigError,
|
||||
RuntimeCacheStoreInitError, RuntimeCacheStores,
|
||||
};
|
||||
use crank_core::{
|
||||
CacheScope, CacheStoreError, CachedHeader, CachedResponse, CoordinationStateStore,
|
||||
CoordinationStateValue, RateLimitBucketState, RateLimitStateStore, ReplayGuardStatus,
|
||||
ReplayGuardStore, ResponseCacheStore,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn defaults_to_in_memory_cache_without_url() {
|
||||
let config = RuntimeCacheConfig::default();
|
||||
|
||||
assert_eq!(config.backend, CacheBackend::Memory);
|
||||
assert_eq!(config.url, None);
|
||||
assert_eq!(config.default_ttl_ms, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn loads_valkey_config_from_env() {
|
||||
unsafe {
|
||||
std::env::set_var("CRANK_CACHE_BACKEND", "valkey");
|
||||
std::env::set_var("CRANK_CACHE_URL", "redis://cache:6379/0");
|
||||
std::env::set_var("CRANK_CACHE_DEFAULT_TTL_MS", "15000");
|
||||
}
|
||||
|
||||
let config = RuntimeCacheConfig::from_env().unwrap();
|
||||
|
||||
assert_eq!(config.backend, CacheBackend::Valkey);
|
||||
assert_eq!(config.url.as_deref(), Some("redis://cache:6379/0"));
|
||||
assert_eq!(config.default_ttl_ms, Some(15_000));
|
||||
|
||||
unsafe {
|
||||
std::env::remove_var("CRANK_CACHE_BACKEND");
|
||||
std::env::remove_var("CRANK_CACHE_URL");
|
||||
std::env::remove_var("CRANK_CACHE_DEFAULT_TTL_MS");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_external_backend_without_url() {
|
||||
unsafe {
|
||||
std::env::set_var("CRANK_CACHE_BACKEND", "redis");
|
||||
std::env::remove_var("CRANK_CACHE_URL");
|
||||
}
|
||||
|
||||
let error = RuntimeCacheConfig::from_env().unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
RuntimeCacheConfigError::MissingUrl {
|
||||
backend: CacheBackend::Redis
|
||||
}
|
||||
));
|
||||
|
||||
unsafe {
|
||||
std::env::remove_var("CRANK_CACHE_BACKEND");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_ttl() {
|
||||
unsafe {
|
||||
std::env::set_var("CRANK_CACHE_DEFAULT_TTL_MS", "0");
|
||||
}
|
||||
|
||||
let error = RuntimeCacheConfig::from_env().unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
RuntimeCacheConfigError::ZeroTtl {
|
||||
name: "CRANK_CACHE_DEFAULT_TTL_MS"
|
||||
}
|
||||
));
|
||||
|
||||
unsafe {
|
||||
std::env::remove_var("CRANK_CACHE_DEFAULT_TTL_MS");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn in_memory_response_cache_roundtrips_values() {
|
||||
let store = InMemoryResponseCacheStore::default();
|
||||
let value = CachedResponse {
|
||||
status: 200,
|
||||
headers: vec![CachedHeader {
|
||||
name: "content-type".to_owned(),
|
||||
value: "application/json".to_owned(),
|
||||
}],
|
||||
body: br#"{"ok":true}"#.to_vec(),
|
||||
data: br#"{"ok":true}"#.to_vec(),
|
||||
};
|
||||
|
||||
store
|
||||
.put("response:crm:list", value.clone(), Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.get("response:crm:list").await.unwrap(),
|
||||
Some(value.clone())
|
||||
);
|
||||
|
||||
store.delete("response:crm:list").await.unwrap();
|
||||
|
||||
assert_eq!(store.get("response:crm:list").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn in_memory_rate_limit_store_roundtrips_bucket_state() {
|
||||
let store = InMemoryRateLimitStateStore::default();
|
||||
let bucket = RateLimitBucketState {
|
||||
tokens_micros: 1_500_000,
|
||||
last_refill_unix_ms: 1_735_689_000_000,
|
||||
};
|
||||
|
||||
store
|
||||
.put_bucket("tenant:alpha", bucket, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.get_bucket("tenant:alpha").await.unwrap(),
|
||||
Some(bucket)
|
||||
);
|
||||
|
||||
store.delete_bucket("tenant:alpha").await.unwrap();
|
||||
|
||||
assert_eq!(store.get_bucket("tenant:alpha").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn in_memory_replay_guard_marks_key_only_once_until_cleared() {
|
||||
let store = InMemoryReplayGuardStore::default();
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.mark_seen("token:nonce:1", Duration::from_secs(10))
|
||||
.await
|
||||
.unwrap(),
|
||||
ReplayGuardStatus::Fresh
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.mark_seen("token:nonce:1", Duration::from_secs(10))
|
||||
.await
|
||||
.unwrap(),
|
||||
ReplayGuardStatus::AlreadySeen
|
||||
);
|
||||
|
||||
store.clear("token:nonce:1").await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.mark_seen("token:nonce:1", Duration::from_secs(10))
|
||||
.await
|
||||
.unwrap(),
|
||||
ReplayGuardStatus::Fresh
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn in_memory_coordination_store_scopes_keys() {
|
||||
let store = InMemoryCoordinationStateStore::default();
|
||||
let response_value = CoordinationStateValue {
|
||||
payload: json!({ "cursor": "abc" }),
|
||||
};
|
||||
let session_value = CoordinationStateValue {
|
||||
payload: json!({ "cursor": "xyz" }),
|
||||
};
|
||||
|
||||
store
|
||||
.put_value(
|
||||
CacheScope::Response,
|
||||
"job-1",
|
||||
response_value.clone(),
|
||||
Duration::from_secs(20),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.put_value(
|
||||
CacheScope::Coordination,
|
||||
"job-1",
|
||||
session_value.clone(),
|
||||
Duration::from_secs(20),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.get_value(CacheScope::Response, "job-1")
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(response_value)
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.get_value(CacheScope::Coordination, "job-1")
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(session_value)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn in_memory_stores_reject_empty_keys() {
|
||||
let response_store = InMemoryResponseCacheStore::default();
|
||||
let replay_store = InMemoryReplayGuardStore::default();
|
||||
|
||||
let error = response_store.get("").await.unwrap_err();
|
||||
assert!(matches!(error, CacheStoreError::InvalidKey { .. }));
|
||||
|
||||
let error = replay_store
|
||||
.mark_seen("", Duration::from_secs(1))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(error, CacheStoreError::InvalidKey { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn runtime_cache_stores_default_to_memory_backend() {
|
||||
let stores = RuntimeCacheStores::from_config(&RuntimeCacheConfig::default())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(stores.backend, CacheBackend::Memory);
|
||||
|
||||
stores
|
||||
.response
|
||||
.put(
|
||||
"response:health",
|
||||
CachedResponse {
|
||||
status: 200,
|
||||
headers: vec![],
|
||||
body: b"ok".to_vec(),
|
||||
data: br#"null"#.to_vec(),
|
||||
},
|
||||
Duration::from_secs(5),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
stores
|
||||
.response
|
||||
.get("response:health")
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redis_cache_store_scopes_keys_by_kind() {
|
||||
let key = "agent:primary";
|
||||
|
||||
let response_key = format!("crank:{}:{}", "response", key);
|
||||
let coordination_key = format!(
|
||||
"crank:{}:{}",
|
||||
RedisCacheStore::coordination_kind(CacheScope::Coordination),
|
||||
key
|
||||
);
|
||||
|
||||
assert_eq!(response_key, "crank:response:agent:primary");
|
||||
assert_eq!(
|
||||
coordination_key,
|
||||
"crank:coordination:coordination:agent:primary"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redis_cache_store_roundtrips_serialized_values() {
|
||||
let response = CachedResponse {
|
||||
status: 202,
|
||||
headers: vec![CachedHeader {
|
||||
name: "x-cache".to_owned(),
|
||||
value: "hit".to_owned(),
|
||||
}],
|
||||
body: br#"{"queued":true}"#.to_vec(),
|
||||
data: br#"null"#.to_vec(),
|
||||
};
|
||||
|
||||
let encoded = serde_json::to_vec(&response).unwrap();
|
||||
let decoded: CachedResponse = serde_json::from_slice(&encoded).unwrap();
|
||||
|
||||
assert_eq!(decoded, response);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_cache_stores_report_missing_external_url() {
|
||||
let future = RuntimeCacheStores::from_config(&RuntimeCacheConfig {
|
||||
backend: CacheBackend::Valkey,
|
||||
url: None,
|
||||
default_ttl_ms: None,
|
||||
});
|
||||
|
||||
let runtime = tokio::runtime::Runtime::new().unwrap();
|
||||
let error = match runtime.block_on(future) {
|
||||
Ok(_) => panic!("expected missing external cache url error"),
|
||||
Err(error) => error,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
RuntimeCacheStoreInitError::MissingUrl {
|
||||
backend: CacheBackend::Valkey
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::{RuntimeCacheConfig, RuntimeCacheStoreInitError, RuntimeCacheStores};
|
||||
|
||||
#[async_trait]
|
||||
pub trait CacheBackendFactory: Send + Sync {
|
||||
async fn build(
|
||||
&self,
|
||||
config: &RuntimeCacheConfig,
|
||||
) -> Result<RuntimeCacheStores, RuntimeCacheStoreInitError>;
|
||||
}
|
||||
|
||||
pub type SharedCacheBackendFactory = Arc<dyn CacheBackendFactory>;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct BuiltinCacheBackendFactory;
|
||||
|
||||
#[async_trait]
|
||||
impl CacheBackendFactory for BuiltinCacheBackendFactory {
|
||||
async fn build(
|
||||
&self,
|
||||
config: &RuntimeCacheConfig,
|
||||
) -> Result<RuntimeCacheStores, RuntimeCacheStoreInitError> {
|
||||
RuntimeCacheStores::from_config(config).await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::{CacheBackendFactory, RuntimeCacheConfig};
|
||||
|
||||
use super::BuiltinCacheBackendFactory;
|
||||
|
||||
#[tokio::test]
|
||||
async fn builtin_factory_builds_default_memory_stores() {
|
||||
let stores = BuiltinCacheBackendFactory
|
||||
.build(&RuntimeCacheConfig::default())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(stores.backend, crank_core::CacheBackend::Memory);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
use crank_adapter_rest::RestAdapterError;
|
||||
use crank_core::{ExecutionMode, Protocol};
|
||||
use crank_mapping::MappingError;
|
||||
use crank_schema::SchemaError;
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum RuntimeError {
|
||||
#[error(transparent)]
|
||||
Schema(#[from] SchemaError),
|
||||
#[error(transparent)]
|
||||
Mapping(#[from] MappingError),
|
||||
#[error("{0}")]
|
||||
GraphqlAdapter(String),
|
||||
#[error("{0}")]
|
||||
GrpcAdapter(String),
|
||||
#[error(transparent)]
|
||||
RestAdapter(#[from] RestAdapterError),
|
||||
#[error("{0}")]
|
||||
ProtocolAdapter(String),
|
||||
#[error("{0}")]
|
||||
SoapAdapter(String),
|
||||
#[error("{0}")]
|
||||
WebsocketAdapter(String),
|
||||
#[error("protocol {protocol:?} is not supported by runtime")]
|
||||
UnsupportedProtocol { protocol: Protocol },
|
||||
#[error("operation {operation_id} does not define streaming config")]
|
||||
MissingStreamingConfig { operation_id: String },
|
||||
#[error("operation {operation_id} does not support requested execution mode {mode:?}")]
|
||||
UnsupportedExecutionMode {
|
||||
operation_id: String,
|
||||
mode: ExecutionMode,
|
||||
},
|
||||
#[error("runtime concurrency limit exceeded for {kind} executions (limit {limit})")]
|
||||
ConcurrencyLimitExceeded { kind: &'static str, limit: usize },
|
||||
#[error("invalid prepared request at {field}: {reason}")]
|
||||
InvalidPreparedRequest { field: String, reason: String },
|
||||
#[error("invalid streaming payload at {field}: {reason}")]
|
||||
InvalidStreamingPayload { field: String, reason: String },
|
||||
#[error("auth profile {auth_profile_id} was not found")]
|
||||
MissingAuthProfile { auth_profile_id: String },
|
||||
#[error("secret {secret_id} was not found")]
|
||||
MissingSecret { secret_id: String },
|
||||
#[error("secret {secret_id} does not have current version {version}")]
|
||||
MissingSecretVersion { secret_id: String, version: u32 },
|
||||
#[error("invalid secret payload for {secret_id}: {reason}")]
|
||||
InvalidAuthSecretValue { secret_id: String, reason: String },
|
||||
#[error("secret crypto error during {operation}: {details}")]
|
||||
SecretCrypto {
|
||||
operation: &'static str,
|
||||
details: String,
|
||||
},
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,67 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crank_adapter_rest::RestAdapter;
|
||||
use crank_core::{
|
||||
AdapterRegistry, MeteringSink, NoopMeteringSink, ResponseCacheStore, SharedMeteringSink,
|
||||
SharedProtocolAdapter,
|
||||
};
|
||||
|
||||
use crate::{RuntimeExecutor, RuntimeLimits};
|
||||
|
||||
pub struct RuntimeExecutorBuilder {
|
||||
limits: RuntimeLimits,
|
||||
adapters: AdapterRegistry,
|
||||
response_cache: Option<Arc<dyn ResponseCacheStore>>,
|
||||
metering_sink: SharedMeteringSink,
|
||||
}
|
||||
|
||||
impl RuntimeExecutorBuilder {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
limits: RuntimeLimits::default(),
|
||||
adapters: AdapterRegistry::empty(),
|
||||
response_cache: None,
|
||||
metering_sink: Arc::new(NoopMeteringSink) as Arc<dyn MeteringSink>,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_limits(mut self, limits: RuntimeLimits) -> Self {
|
||||
self.limits = limits;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn register_adapter(mut self, adapter: SharedProtocolAdapter) -> Self {
|
||||
self.adapters = self.adapters.register(adapter);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_response_cache(mut self, store: Arc<dyn ResponseCacheStore>) -> Self {
|
||||
self.response_cache = Some(store);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_metering_sink(mut self, sink: SharedMeteringSink) -> Self {
|
||||
self.metering_sink = sink;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> RuntimeExecutor {
|
||||
RuntimeExecutor::from_builder_parts(
|
||||
self.limits,
|
||||
self.adapters,
|
||||
self.response_cache,
|
||||
self.metering_sink,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RuntimeExecutorBuilder {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn community_default() -> RuntimeExecutorBuilder {
|
||||
RuntimeExecutorBuilder::new()
|
||||
.register_adapter(Arc::new(RestAdapter::new()) as SharedProtocolAdapter)
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
mod aggregation;
|
||||
mod auth;
|
||||
mod cache;
|
||||
mod cache_factory;
|
||||
mod error;
|
||||
mod executor;
|
||||
mod executor_builder;
|
||||
mod limits;
|
||||
mod model;
|
||||
mod rate_limit;
|
||||
mod redaction;
|
||||
mod request_context;
|
||||
mod secret_crypto;
|
||||
mod streaming;
|
||||
|
||||
pub use auth::ResolvedAuth;
|
||||
pub use cache::{
|
||||
InMemoryCoordinationStateStore, InMemoryRateLimitStateStore, InMemoryReplayGuardStore,
|
||||
InMemoryResponseCacheStore, RedisCacheStore, RuntimeCacheConfig, RuntimeCacheConfigError,
|
||||
RuntimeCacheStoreInitError, RuntimeCacheStores,
|
||||
};
|
||||
pub use cache_factory::{
|
||||
BuiltinCacheBackendFactory, CacheBackendFactory, SharedCacheBackendFactory,
|
||||
};
|
||||
pub use error::RuntimeError;
|
||||
pub use executor::RuntimeExecutor;
|
||||
pub use executor_builder::{RuntimeExecutorBuilder, community_default};
|
||||
pub use limits::{RuntimeLimits, RuntimeLimitsConfigError};
|
||||
pub use model::{AdapterResponse, PreparedRequest, RuntimeOperation};
|
||||
pub use rate_limit::{
|
||||
RateLimitRejection, RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter,
|
||||
};
|
||||
pub use request_context::{MeteringContext, ResponseCacheScope, RuntimeRequestContext};
|
||||
pub use secret_crypto::SecretCrypto;
|
||||
pub use streaming::WindowExecutionResult;
|
||||
@@ -0,0 +1,119 @@
|
||||
use std::{env, num::ParseIntError};
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
const DEFAULT_MAX_CONCURRENT_UNARY: usize = 64;
|
||||
const DEFAULT_MAX_CONCURRENT_WINDOW: usize = 16;
|
||||
const DEFAULT_MAX_CONCURRENT_SESSIONS: usize = 16;
|
||||
const DEFAULT_MAX_CONCURRENT_JOBS: usize = 16;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct RuntimeLimits {
|
||||
pub max_concurrent_unary: usize,
|
||||
pub max_concurrent_window: usize,
|
||||
pub max_concurrent_sessions: usize,
|
||||
pub max_concurrent_jobs: usize,
|
||||
}
|
||||
|
||||
impl Default for RuntimeLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_concurrent_unary: DEFAULT_MAX_CONCURRENT_UNARY,
|
||||
max_concurrent_window: DEFAULT_MAX_CONCURRENT_WINDOW,
|
||||
max_concurrent_sessions: DEFAULT_MAX_CONCURRENT_SESSIONS,
|
||||
max_concurrent_jobs: DEFAULT_MAX_CONCURRENT_JOBS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RuntimeLimits {
|
||||
pub fn from_env() -> Result<Self, RuntimeLimitsConfigError> {
|
||||
Ok(Self {
|
||||
max_concurrent_unary: parse_limit(
|
||||
"CRANK_RUNTIME_MAX_CONCURRENT_UNARY",
|
||||
DEFAULT_MAX_CONCURRENT_UNARY,
|
||||
)?,
|
||||
max_concurrent_window: parse_limit(
|
||||
"CRANK_RUNTIME_MAX_CONCURRENT_WINDOW",
|
||||
DEFAULT_MAX_CONCURRENT_WINDOW,
|
||||
)?,
|
||||
max_concurrent_sessions: parse_limit(
|
||||
"CRANK_RUNTIME_MAX_CONCURRENT_SESSIONS",
|
||||
DEFAULT_MAX_CONCURRENT_SESSIONS,
|
||||
)?,
|
||||
max_concurrent_jobs: parse_limit(
|
||||
"CRANK_RUNTIME_MAX_CONCURRENT_JOBS",
|
||||
DEFAULT_MAX_CONCURRENT_JOBS,
|
||||
)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_limit(name: &'static str, default: usize) -> Result<usize, RuntimeLimitsConfigError> {
|
||||
match env::var(name) {
|
||||
Ok(raw) => {
|
||||
let value =
|
||||
raw.parse::<usize>()
|
||||
.map_err(|source| RuntimeLimitsConfigError::InvalidValue {
|
||||
name,
|
||||
value: raw,
|
||||
source,
|
||||
})?;
|
||||
if value == 0 {
|
||||
return Err(RuntimeLimitsConfigError::ZeroValue { name });
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
Err(env::VarError::NotPresent) => Ok(default),
|
||||
Err(env::VarError::NotUnicode(_)) => Err(RuntimeLimitsConfigError::InvalidUnicode { name }),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum RuntimeLimitsConfigError {
|
||||
#[error("{name} must contain valid UTF-8")]
|
||||
InvalidUnicode { name: &'static str },
|
||||
#[error("{name} must be a positive integer, got {value}")]
|
||||
InvalidValue {
|
||||
name: &'static str,
|
||||
value: String,
|
||||
source: ParseIntError,
|
||||
},
|
||||
#[error("{name} must be greater than zero")]
|
||||
ZeroValue { name: &'static str },
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{RuntimeLimits, RuntimeLimitsConfigError};
|
||||
|
||||
#[test]
|
||||
fn defaults_are_positive() {
|
||||
let limits = RuntimeLimits::default();
|
||||
|
||||
assert!(limits.max_concurrent_unary > 0);
|
||||
assert!(limits.max_concurrent_window > 0);
|
||||
assert!(limits.max_concurrent_sessions > 0);
|
||||
assert!(limits.max_concurrent_jobs > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_limit_values() {
|
||||
unsafe {
|
||||
std::env::set_var("CRANK_RUNTIME_MAX_CONCURRENT_UNARY", "0");
|
||||
}
|
||||
|
||||
let error = RuntimeLimits::from_env().unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
RuntimeLimitsConfigError::ZeroValue {
|
||||
name: "CRANK_RUNTIME_MAX_CONCURRENT_UNARY"
|
||||
}
|
||||
));
|
||||
|
||||
unsafe {
|
||||
std::env::remove_var("CRANK_RUNTIME_MAX_CONCURRENT_UNARY");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crank_core::{ExecutionConfig, Operation, OperationId, Protocol, Target, ToolDescription};
|
||||
use crank_mapping::MappingSet;
|
||||
use crank_schema::Schema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct RuntimeOperation {
|
||||
pub operation_id: OperationId,
|
||||
pub operation_version: u32,
|
||||
pub tool_name: String,
|
||||
pub protocol: Protocol,
|
||||
pub target: Target,
|
||||
pub input_schema: Schema,
|
||||
pub output_schema: Schema,
|
||||
pub input_mapping: MappingSet,
|
||||
pub output_mapping: MappingSet,
|
||||
pub execution_config: ExecutionConfig,
|
||||
pub tool_description: ToolDescription,
|
||||
}
|
||||
|
||||
impl From<Operation<Schema, MappingSet>> for RuntimeOperation {
|
||||
fn from(value: Operation<Schema, MappingSet>) -> Self {
|
||||
Self {
|
||||
operation_id: value.id,
|
||||
operation_version: value.version,
|
||||
tool_name: value.name,
|
||||
protocol: value.protocol,
|
||||
target: value.target,
|
||||
input_schema: value.input_schema,
|
||||
output_schema: value.output_schema,
|
||||
input_mapping: value.input_mapping,
|
||||
output_mapping: value.output_mapping,
|
||||
execution_config: value.execution_config,
|
||||
tool_description: value.tool_description,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
|
||||
pub struct PreparedRequest {
|
||||
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
|
||||
pub path_params: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
|
||||
pub query_params: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
|
||||
pub headers: BTreeMap<String, String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub grpc: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub variables: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub body: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub timeout_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AdapterResponse {
|
||||
pub status_code: u16,
|
||||
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
|
||||
pub headers: BTreeMap<String, String>,
|
||||
pub body: Value,
|
||||
pub data: Value,
|
||||
}
|
||||
|
||||
impl From<PreparedRequest> for crank_core::PreparedRequest {
|
||||
fn from(value: PreparedRequest) -> Self {
|
||||
Self {
|
||||
path_params: value.path_params,
|
||||
query_params: value.query_params,
|
||||
headers: value.headers,
|
||||
grpc: value.grpc,
|
||||
variables: value.variables,
|
||||
body: value.body,
|
||||
timeout_ms: value.timeout_ms,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<crank_core::PreparedRequest> for PreparedRequest {
|
||||
fn from(value: crank_core::PreparedRequest) -> Self {
|
||||
Self {
|
||||
path_params: value.path_params,
|
||||
query_params: value.query_params,
|
||||
headers: value.headers,
|
||||
grpc: value.grpc,
|
||||
variables: value.variables,
|
||||
body: value.body,
|
||||
timeout_ms: value.timeout_ms,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<crank_core::AdapterResponse> for AdapterResponse {
|
||||
fn from(value: crank_core::AdapterResponse) -> Self {
|
||||
Self {
|
||||
status_code: value.status_code,
|
||||
headers: value.headers,
|
||||
body: value.body,
|
||||
data: value.data,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,269 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, Mutex},
|
||||
time::{Duration, Instant, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use crank_core::{RateLimitBucketState, RateLimitStateStore};
|
||||
use thiserror::Error;
|
||||
|
||||
const STALE_KEY_TTL: Duration = Duration::from_secs(300);
|
||||
const TOKEN_SCALE: u64 = 1_000_000;
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct RequestRateLimitConfig {
|
||||
pub requests_per_second: u32,
|
||||
pub burst: u32,
|
||||
}
|
||||
|
||||
impl RequestRateLimitConfig {
|
||||
pub fn new(requests_per_second: u32, burst: u32) -> Result<Self, RequestRateLimitConfigError> {
|
||||
if requests_per_second == 0 {
|
||||
return Err(RequestRateLimitConfigError::InvalidRequestsPerSecond(
|
||||
requests_per_second,
|
||||
));
|
||||
}
|
||||
if burst == 0 {
|
||||
return Err(RequestRateLimitConfigError::InvalidBurst(burst));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
requests_per_second,
|
||||
burst,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum RequestRateLimitConfigError {
|
||||
#[error("requests_per_second must be greater than zero, got {0}")]
|
||||
InvalidRequestsPerSecond(u32),
|
||||
#[error("burst must be greater than zero, got {0}")]
|
||||
InvalidBurst(u32),
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct RateLimitRejection {
|
||||
pub retry_after_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RequestRateLimiter {
|
||||
config: RequestRateLimitConfig,
|
||||
backend: RequestRateLimiterBackend,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum RequestRateLimiterBackend {
|
||||
Local {
|
||||
states: Arc<Mutex<HashMap<String, LocalBucketState>>>,
|
||||
},
|
||||
Shared {
|
||||
store: Arc<dyn RateLimitStateStore>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
struct LocalBucketState {
|
||||
tokens: f64,
|
||||
last_refill: Instant,
|
||||
last_seen: Instant,
|
||||
}
|
||||
|
||||
impl RequestRateLimiter {
|
||||
pub fn new(config: RequestRateLimitConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
backend: RequestRateLimiterBackend::Local {
|
||||
states: Arc::new(Mutex::new(HashMap::new())),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_shared(config: RequestRateLimitConfig, store: Arc<dyn RateLimitStateStore>) -> Self {
|
||||
Self {
|
||||
config,
|
||||
backend: RequestRateLimiterBackend::Shared { store },
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn check(&self, key: &str) -> Result<(), RateLimitRejection> {
|
||||
match &self.backend {
|
||||
RequestRateLimiterBackend::Local { .. } => self.check_local_at(key, Instant::now()),
|
||||
RequestRateLimiterBackend::Shared { store } => {
|
||||
self.check_shared_at(store.as_ref(), key, now_unix_ms())
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn check_local_at(&self, key: &str, now: Instant) -> Result<(), RateLimitRejection> {
|
||||
let RequestRateLimiterBackend::Local { states } = &self.backend else {
|
||||
panic!("check_local_at called for non-local limiter");
|
||||
};
|
||||
|
||||
let mut states = states
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
states.retain(|_, state| now.duration_since(state.last_seen) <= STALE_KEY_TTL);
|
||||
|
||||
let burst = self.config.burst as f64;
|
||||
let requests_per_second = self.config.requests_per_second as f64;
|
||||
let state = states.entry(key.to_owned()).or_insert(LocalBucketState {
|
||||
tokens: burst,
|
||||
last_refill: now,
|
||||
last_seen: now,
|
||||
});
|
||||
|
||||
let elapsed = now.duration_since(state.last_refill).as_secs_f64();
|
||||
state.tokens = (state.tokens + elapsed * requests_per_second).min(burst);
|
||||
state.last_refill = now;
|
||||
state.last_seen = now;
|
||||
|
||||
if state.tokens >= 1.0 {
|
||||
state.tokens -= 1.0;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let missing_tokens = (1.0 - state.tokens).max(0.0);
|
||||
let retry_after_ms = ((missing_tokens / requests_per_second) * 1000.0)
|
||||
.ceil()
|
||||
.max(1.0) as u64;
|
||||
Err(RateLimitRejection { retry_after_ms })
|
||||
}
|
||||
|
||||
async fn check_shared_at(
|
||||
&self,
|
||||
store: &dyn RateLimitStateStore,
|
||||
key: &str,
|
||||
now_unix_ms: i64,
|
||||
) -> Result<(), RateLimitRejection> {
|
||||
let burst_tokens = u64::from(self.config.burst) * TOKEN_SCALE;
|
||||
let refill_per_second = u64::from(self.config.requests_per_second) * TOKEN_SCALE;
|
||||
let mut state =
|
||||
store
|
||||
.get_bucket(key)
|
||||
.await
|
||||
.unwrap_or(None)
|
||||
.unwrap_or(RateLimitBucketState {
|
||||
tokens_micros: burst_tokens,
|
||||
last_refill_unix_ms: now_unix_ms,
|
||||
});
|
||||
|
||||
let elapsed_ms = (now_unix_ms - state.last_refill_unix_ms).max(0) as u64;
|
||||
let replenished =
|
||||
state.tokens_micros + (elapsed_ms.saturating_mul(refill_per_second) / 1000);
|
||||
state.tokens_micros = replenished.min(burst_tokens);
|
||||
state.last_refill_unix_ms = now_unix_ms;
|
||||
|
||||
if state.tokens_micros >= TOKEN_SCALE {
|
||||
state.tokens_micros -= TOKEN_SCALE;
|
||||
let _ = store.put_bucket(key, state, STALE_KEY_TTL).await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let missing_tokens = TOKEN_SCALE.saturating_sub(state.tokens_micros);
|
||||
let retry_after_ms = missing_tokens.div_ceil(refill_per_second).max(1);
|
||||
let _ = store.put_bucket(key, state, STALE_KEY_TTL).await;
|
||||
Err(RateLimitRejection { retry_after_ms })
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_ms() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_else(|_| Duration::from_secs(0))
|
||||
.as_millis()
|
||||
.try_into()
|
||||
.unwrap_or(i64::MAX)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use crate::InMemoryRateLimitStateStore;
|
||||
|
||||
use super::{RequestRateLimitConfig, RequestRateLimitConfigError, RequestRateLimiter};
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_requests_per_second() {
|
||||
assert_eq!(
|
||||
RequestRateLimitConfig::new(0, 1),
|
||||
Err(RequestRateLimitConfigError::InvalidRequestsPerSecond(0))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_zero_burst() {
|
||||
assert_eq!(
|
||||
RequestRateLimitConfig::new(1, 0),
|
||||
Err(RequestRateLimitConfigError::InvalidBurst(0))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_burst_then_rejects_until_refilled_for_local_limiter() {
|
||||
let limiter = RequestRateLimiter::new(RequestRateLimitConfig::new(2, 2).unwrap());
|
||||
let start = Instant::now();
|
||||
|
||||
assert!(limiter.check_local_at("key", start).is_ok());
|
||||
assert!(limiter.check_local_at("key", start).is_ok());
|
||||
|
||||
let rejection = limiter.check_local_at("key", start).unwrap_err();
|
||||
assert_eq!(rejection.retry_after_ms, 500);
|
||||
|
||||
assert!(
|
||||
limiter
|
||||
.check_local_at("key", start + Duration::from_millis(500))
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn allows_burst_then_rejects_until_refilled_for_shared_limiter() {
|
||||
let limiter = RequestRateLimiter::new_shared(
|
||||
RequestRateLimitConfig::new(2, 2).unwrap(),
|
||||
Arc::new(InMemoryRateLimitStateStore::default()),
|
||||
);
|
||||
|
||||
assert!(limiter.check_shared_at_store("key", 0).await.is_ok());
|
||||
assert!(limiter.check_shared_at_store("key", 0).await.is_ok());
|
||||
|
||||
let rejection = limiter.check_shared_at_store("key", 0).await.unwrap_err();
|
||||
assert_eq!(rejection.retry_after_ms, 1);
|
||||
|
||||
assert!(limiter.check_shared_at_store("key", 500).await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shared_limiter_state_is_visible_across_instances() {
|
||||
let store = Arc::new(InMemoryRateLimitStateStore::default());
|
||||
let first = RequestRateLimiter::new_shared(
|
||||
RequestRateLimitConfig::new(1, 1).unwrap(),
|
||||
store.clone(),
|
||||
);
|
||||
let second =
|
||||
RequestRateLimiter::new_shared(RequestRateLimitConfig::new(1, 1).unwrap(), store);
|
||||
|
||||
assert!(first.check_shared_at_store("shared", 0).await.is_ok());
|
||||
assert!(second.check_shared_at_store("shared", 0).await.is_err());
|
||||
assert!(second.check_shared_at_store("shared", 1000).await.is_ok());
|
||||
}
|
||||
|
||||
impl RequestRateLimiter {
|
||||
async fn check_shared_at_store(
|
||||
&self,
|
||||
key: &str,
|
||||
now_unix_ms: i64,
|
||||
) -> Result<(), super::RateLimitRejection> {
|
||||
let super::RequestRateLimiterBackend::Shared { store } = &self.backend else {
|
||||
panic!("check_shared_at_store called for non-shared limiter");
|
||||
};
|
||||
self.check_shared_at(store.as_ref(), key, now_unix_ms).await
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
use serde_json::Value;
|
||||
|
||||
pub fn redact_paths(value: &mut Value, paths: &[String]) {
|
||||
for path in paths {
|
||||
let segments = parse_path(path);
|
||||
if segments.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
redact_path(value, &segments);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn truncate_item_fields(value: &mut Value, max_len: usize) {
|
||||
match value {
|
||||
Value::String(text) => {
|
||||
if text != "[REDACTED]" && text.chars().count() > max_len {
|
||||
let truncated = text.chars().take(max_len).collect::<String>();
|
||||
*text = format!("{truncated}...");
|
||||
}
|
||||
}
|
||||
Value::Array(items) => {
|
||||
for item in items {
|
||||
truncate_item_fields(item, max_len);
|
||||
}
|
||||
}
|
||||
Value::Object(map) => {
|
||||
for value in map.values_mut() {
|
||||
truncate_item_fields(value, max_len);
|
||||
}
|
||||
}
|
||||
Value::Null | Value::Bool(_) | Value::Number(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
enum PathSegment {
|
||||
Field(String),
|
||||
Wildcard,
|
||||
}
|
||||
|
||||
fn parse_path(path: &str) -> Vec<PathSegment> {
|
||||
let path = path.strip_prefix("$.").or_else(|| path.strip_prefix('$'));
|
||||
let Some(path) = path else {
|
||||
return Vec::new();
|
||||
};
|
||||
|
||||
path.split('.')
|
||||
.flat_map(|segment| {
|
||||
if let Some(field) = segment.strip_suffix("[*]") {
|
||||
vec![PathSegment::Field(field.to_owned()), PathSegment::Wildcard]
|
||||
} else if segment == "*" {
|
||||
vec![PathSegment::Wildcard]
|
||||
} else {
|
||||
vec![PathSegment::Field(segment.to_owned())]
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn redact_path(value: &mut Value, segments: &[PathSegment]) {
|
||||
let Some((current, rest)) = segments.split_first() else {
|
||||
*value = Value::String("[REDACTED]".to_owned());
|
||||
return;
|
||||
};
|
||||
|
||||
match current {
|
||||
PathSegment::Field(field) => {
|
||||
if let Some(next) = value
|
||||
.as_object_mut()
|
||||
.and_then(|object| object.get_mut(field))
|
||||
{
|
||||
redact_path(next, rest);
|
||||
}
|
||||
}
|
||||
PathSegment::Wildcard => {
|
||||
if let Some(items) = value.as_array_mut() {
|
||||
for item in items {
|
||||
redact_path(item, rest);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{redact_paths, truncate_item_fields};
|
||||
|
||||
#[test]
|
||||
fn redacts_simple_and_wildcard_paths() {
|
||||
let mut value = json!({
|
||||
"summary": { "token": "secret" },
|
||||
"items": [
|
||||
{ "token": "a", "message": "keep" },
|
||||
{ "token": "b", "message": "keep" }
|
||||
]
|
||||
});
|
||||
|
||||
redact_paths(
|
||||
&mut value,
|
||||
&["$.summary.token".to_owned(), "$.items[*].token".to_owned()],
|
||||
);
|
||||
|
||||
assert_eq!(value["summary"]["token"], json!("[REDACTED]"));
|
||||
assert_eq!(value["items"][0]["token"], json!("[REDACTED]"));
|
||||
assert_eq!(value["items"][1]["token"], json!("[REDACTED]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_long_strings_recursively() {
|
||||
let mut value = json!({
|
||||
"message": "abcdefghijklmnop",
|
||||
"nested": [{ "message": "qrstuvwxyz" }]
|
||||
});
|
||||
|
||||
truncate_item_fields(&mut value, 4);
|
||||
|
||||
assert_eq!(value["message"], json!("abcd..."));
|
||||
assert_eq!(value["nested"][0]["message"], json!("qrst..."));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crank_core::{AgentId, InvocationSource, WorkspaceId};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ResponseCacheScope {
|
||||
pub workspace_key: String,
|
||||
pub agent_key: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RuntimeRequestContext {
|
||||
pub request_id: String,
|
||||
pub correlation_id: String,
|
||||
pub response_cache_scope: Option<ResponseCacheScope>,
|
||||
pub metering_context: Option<MeteringContext>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct MeteringContext {
|
||||
pub workspace_id: WorkspaceId,
|
||||
pub agent_id: Option<AgentId>,
|
||||
pub source: InvocationSource,
|
||||
}
|
||||
|
||||
impl RuntimeRequestContext {
|
||||
pub fn new(request_id: impl Into<String>, correlation_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
request_id: request_id.into(),
|
||||
correlation_id: correlation_id.into(),
|
||||
response_cache_scope: None,
|
||||
metering_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_request_id(request_id: impl Into<String>) -> Self {
|
||||
let request_id = request_id.into();
|
||||
Self::new(request_id.clone(), request_id)
|
||||
}
|
||||
|
||||
pub fn outbound_headers(&self) -> BTreeMap<String, String> {
|
||||
BTreeMap::from([
|
||||
("x-request-id".to_owned(), self.request_id.clone()),
|
||||
("x-correlation-id".to_owned(), self.correlation_id.clone()),
|
||||
])
|
||||
}
|
||||
|
||||
pub fn with_response_cache_scope(
|
||||
mut self,
|
||||
workspace_key: impl Into<String>,
|
||||
agent_key: impl Into<String>,
|
||||
) -> Self {
|
||||
self.response_cache_scope = Some(ResponseCacheScope {
|
||||
workspace_key: workspace_key.into(),
|
||||
agent_key: agent_key.into(),
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
pub fn response_cache_scope(&self) -> Option<&ResponseCacheScope> {
|
||||
self.response_cache_scope.as_ref()
|
||||
}
|
||||
|
||||
pub fn with_metering_context(
|
||||
mut self,
|
||||
workspace_id: WorkspaceId,
|
||||
agent_id: Option<AgentId>,
|
||||
source: InvocationSource,
|
||||
) -> Self {
|
||||
self.metering_context = Some(MeteringContext {
|
||||
workspace_id,
|
||||
agent_id,
|
||||
source,
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
pub fn metering_context(&self) -> Option<&MeteringContext> {
|
||||
self.metering_context.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ResponseCacheScope> for crank_core::ResponseCacheScope {
|
||||
fn from(value: ResponseCacheScope) -> Self {
|
||||
Self {
|
||||
workspace_key: value.workspace_key,
|
||||
agent_key: value.agent_key,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&ResponseCacheScope> for crank_core::ResponseCacheScope {
|
||||
fn from(value: &ResponseCacheScope) -> Self {
|
||||
Self {
|
||||
workspace_key: value.workspace_key.clone(),
|
||||
agent_key: value.agent_key.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&RuntimeRequestContext> for crank_core::RuntimeRequestContext {
|
||||
fn from(value: &RuntimeRequestContext) -> Self {
|
||||
Self {
|
||||
request_id: value.request_id.clone(),
|
||||
correlation_id: value.correlation_id.clone(),
|
||||
response_cache_scope: value.response_cache_scope.as_ref().map(Into::into),
|
||||
metering_context: value.metering_context.as_ref().map(Into::into),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<MeteringContext> for crank_core::MeteringContext {
|
||||
fn from(value: MeteringContext) -> Self {
|
||||
Self {
|
||||
workspace_id: value.workspace_id,
|
||||
agent_id: value.agent_id,
|
||||
source: value.source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&MeteringContext> for crank_core::MeteringContext {
|
||||
fn from(value: &MeteringContext) -> Self {
|
||||
Self {
|
||||
workspace_id: value.workspace_id.clone(),
|
||||
agent_id: value.agent_id.clone(),
|
||||
source: value.source,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crank_core::{InvocationSource, WorkspaceId};
|
||||
|
||||
use super::RuntimeRequestContext;
|
||||
|
||||
#[test]
|
||||
fn uses_request_id_for_default_correlation_id() {
|
||||
let context = RuntimeRequestContext::from_request_id("req_123");
|
||||
|
||||
assert_eq!(context.request_id, "req_123");
|
||||
assert_eq!(context.correlation_id, "req_123");
|
||||
assert!(context.response_cache_scope.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attaches_response_cache_scope() {
|
||||
let context = RuntimeRequestContext::from_request_id("req_123")
|
||||
.with_response_cache_scope("ws_01", "agent_01");
|
||||
|
||||
let scope = context.response_cache_scope().unwrap();
|
||||
assert_eq!(scope.workspace_key, "ws_01");
|
||||
assert_eq!(scope.agent_key, "agent_01");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attaches_metering_context() {
|
||||
let context = RuntimeRequestContext::from_request_id("req_123").with_metering_context(
|
||||
WorkspaceId::new("ws_01"),
|
||||
None,
|
||||
InvocationSource::AdminTestRun,
|
||||
);
|
||||
|
||||
let metering = context.metering_context().unwrap();
|
||||
assert_eq!(metering.workspace_id, WorkspaceId::new("ws_01"));
|
||||
assert_eq!(metering.agent_id, None);
|
||||
assert_eq!(metering.source, InvocationSource::AdminTestRun);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,274 @@
|
||||
use aes_gcm::{
|
||||
Aes256Gcm, KeyInit, Nonce,
|
||||
aead::{Aead, OsRng, rand_core::RngCore},
|
||||
};
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use hkdf::Hkdf;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::RuntimeError;
|
||||
|
||||
const LEGACY_KEY_VERSION: &str = "v1";
|
||||
const CURRENT_KEY_VERSION: &str = "v2";
|
||||
const SECRET_ENVELOPE_INFO: &[u8] = b"crank.secret-envelope.v2";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct SecretCrypto {
|
||||
current_cipher: Aes256Gcm,
|
||||
legacy_cipher: Aes256Gcm,
|
||||
key_version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct CipherEnvelope {
|
||||
nonce_b64: String,
|
||||
ciphertext_b64: String,
|
||||
}
|
||||
|
||||
impl SecretCrypto {
|
||||
pub fn new(master_key: &str) -> Result<Self, RuntimeError> {
|
||||
let trimmed = master_key.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err(RuntimeError::SecretCrypto {
|
||||
operation: "initialize secret crypto",
|
||||
details: "CRANK_MASTER_KEY must not be empty".to_owned(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
current_cipher: derive_hkdf_cipher(trimmed)?,
|
||||
legacy_cipher: derive_legacy_cipher(trimmed)?,
|
||||
key_version: CURRENT_KEY_VERSION.to_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn key_version(&self) -> &str {
|
||||
&self.key_version
|
||||
}
|
||||
|
||||
pub fn encrypt(&self, value: &Value) -> Result<String, RuntimeError> {
|
||||
encrypt_value_with_cipher(&self.current_cipher, value, "encrypt secret value")
|
||||
}
|
||||
|
||||
pub fn decrypt(&self, key_version: &str, ciphertext: &str) -> Result<Value, RuntimeError> {
|
||||
let envelope: CipherEnvelope =
|
||||
serde_json::from_str(ciphertext).map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "decode secret envelope",
|
||||
details: format!("failed to decode secret envelope: {error}"),
|
||||
})?;
|
||||
let nonce_bytes =
|
||||
STANDARD
|
||||
.decode(envelope.nonce_b64)
|
||||
.map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "decode secret nonce",
|
||||
details: format!("failed to decode secret nonce: {error}"),
|
||||
})?;
|
||||
let ciphertext_bytes = STANDARD.decode(envelope.ciphertext_b64).map_err(|error| {
|
||||
RuntimeError::SecretCrypto {
|
||||
operation: "decode secret payload",
|
||||
details: format!("failed to decode secret payload: {error}"),
|
||||
}
|
||||
})?;
|
||||
let cipher = self.cipher_for_version(key_version)?;
|
||||
let plaintext = cipher
|
||||
.decrypt(Nonce::from_slice(&nonce_bytes), ciphertext_bytes.as_ref())
|
||||
.map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "decrypt secret value",
|
||||
details: format!("failed to decrypt secret value: {error}"),
|
||||
})?;
|
||||
|
||||
serde_json::from_slice(&plaintext).map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "deserialize secret value",
|
||||
details: format!("failed to deserialize secret value: {error}"),
|
||||
})
|
||||
}
|
||||
|
||||
fn cipher_for_version(&self, key_version: &str) -> Result<&Aes256Gcm, RuntimeError> {
|
||||
match key_version {
|
||||
LEGACY_KEY_VERSION => Ok(&self.legacy_cipher),
|
||||
CURRENT_KEY_VERSION => Ok(&self.current_cipher),
|
||||
other => Err(RuntimeError::SecretCrypto {
|
||||
operation: "select secret key version",
|
||||
details: format!("unsupported secret key version: {other}"),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn derive_legacy_cipher(master_key: &str) -> Result<Aes256Gcm, RuntimeError> {
|
||||
let digest = Sha256::digest(master_key.as_bytes());
|
||||
Aes256Gcm::new_from_slice(digest.as_slice()).map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "initialize legacy secret crypto",
|
||||
details: format!("failed to initialize legacy secret crypto: {error}"),
|
||||
})
|
||||
}
|
||||
|
||||
fn derive_hkdf_cipher(master_key: &str) -> Result<Aes256Gcm, RuntimeError> {
|
||||
let hkdf = Hkdf::<Sha256>::new(None, master_key.as_bytes());
|
||||
let mut key_bytes = [0_u8; 32];
|
||||
hkdf.expand(SECRET_ENVELOPE_INFO, &mut key_bytes)
|
||||
.map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "derive secret key with hkdf",
|
||||
details: format!("failed to derive secret key with HKDF: {error}"),
|
||||
})?;
|
||||
Aes256Gcm::new_from_slice(&key_bytes).map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "initialize secret crypto",
|
||||
details: format!("failed to initialize secret crypto: {error}"),
|
||||
})
|
||||
}
|
||||
|
||||
fn encrypt_value_with_cipher(
|
||||
cipher: &Aes256Gcm,
|
||||
value: &Value,
|
||||
action: &str,
|
||||
) -> Result<String, RuntimeError> {
|
||||
let plaintext = serde_json::to_vec(value).map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "serialize secret value",
|
||||
details: format!("failed to serialize secret value: {error}"),
|
||||
})?;
|
||||
let mut nonce_bytes = [0_u8; 12];
|
||||
OsRng.fill_bytes(&mut nonce_bytes);
|
||||
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||
let ciphertext =
|
||||
cipher
|
||||
.encrypt(nonce, plaintext.as_ref())
|
||||
.map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "encrypt secret value",
|
||||
details: format!("failed to {action}: {error}"),
|
||||
})?;
|
||||
let envelope = CipherEnvelope {
|
||||
nonce_b64: STANDARD.encode(nonce_bytes),
|
||||
ciphertext_b64: STANDARD.encode(ciphertext),
|
||||
};
|
||||
|
||||
serde_json::to_string(&envelope).map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "encode secret ciphertext",
|
||||
details: format!("failed to encode secret ciphertext: {error}"),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn encrypt_with_legacy_scheme(master_key: &str, value: &Value) -> Result<String, RuntimeError> {
|
||||
let cipher = derive_legacy_cipher(master_key)?;
|
||||
encrypt_value_with_cipher(&cipher, value, "encrypt secret value with legacy scheme")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn current_key_bytes(master_key: &str) -> Result<[u8; 32], RuntimeError> {
|
||||
let hkdf = Hkdf::<Sha256>::new(None, master_key.as_bytes());
|
||||
let mut key_bytes = [0_u8; 32];
|
||||
hkdf.expand(SECRET_ENVELOPE_INFO, &mut key_bytes)
|
||||
.map_err(|error| RuntimeError::SecretCrypto {
|
||||
operation: "derive secret key with hkdf",
|
||||
details: format!("failed to derive secret key with HKDF: {error}"),
|
||||
})?;
|
||||
Ok(key_bytes)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn legacy_key_bytes(master_key: &str) -> [u8; 32] {
|
||||
let digest = Sha256::digest(master_key.as_bytes());
|
||||
let mut key_bytes = [0_u8; 32];
|
||||
key_bytes.copy_from_slice(digest.as_slice());
|
||||
key_bytes
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::RuntimeError;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
CURRENT_KEY_VERSION, LEGACY_KEY_VERSION, SecretCrypto, current_key_bytes,
|
||||
encrypt_with_legacy_scheme, legacy_key_bytes,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn roundtrips_secret_payload_with_current_scheme() {
|
||||
let crypto = SecretCrypto::new("test-master-key").unwrap();
|
||||
let plaintext = json!({
|
||||
"token": "top-secret",
|
||||
"username": "demo"
|
||||
});
|
||||
|
||||
let ciphertext = crypto.encrypt(&plaintext).unwrap();
|
||||
let decrypted = crypto.decrypt(CURRENT_KEY_VERSION, &ciphertext).unwrap();
|
||||
|
||||
assert_eq!(decrypted, plaintext);
|
||||
assert_eq!(crypto.key_version(), CURRENT_KEY_VERSION);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decrypts_legacy_v1_payloads() {
|
||||
let plaintext = json!({
|
||||
"token": "top-secret",
|
||||
"username": "demo"
|
||||
});
|
||||
let ciphertext = encrypt_with_legacy_scheme("test-master-key", &plaintext).unwrap();
|
||||
let crypto = SecretCrypto::new("test-master-key").unwrap();
|
||||
|
||||
let decrypted = crypto.decrypt(LEGACY_KEY_VERSION, &ciphertext).unwrap();
|
||||
|
||||
assert_eq!(decrypted, plaintext);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_master_key() {
|
||||
let error = SecretCrypto::new(" ").err().unwrap();
|
||||
match error {
|
||||
RuntimeError::SecretCrypto { operation, details } => {
|
||||
assert_eq!(operation, "initialize secret crypto");
|
||||
assert_eq!(details, "CRANK_MASTER_KEY must not be empty");
|
||||
}
|
||||
other => panic!("unexpected error: {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_master_key_derives_stable_hkdf_key() {
|
||||
let lhs = current_key_bytes("test-master-key").unwrap();
|
||||
let rhs = current_key_bytes("test-master-key").unwrap();
|
||||
|
||||
assert_eq!(lhs, rhs);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn current_scheme_key_differs_from_legacy_scheme() {
|
||||
let current = current_key_bytes("test-master-key").unwrap();
|
||||
let legacy = legacy_key_bytes("test-master-key");
|
||||
|
||||
assert_ne!(current, legacy);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_master_keys_produce_different_ciphertexts() {
|
||||
let plaintext = json!({
|
||||
"token": "top-secret",
|
||||
"username": "demo"
|
||||
});
|
||||
let left = SecretCrypto::new("test-master-key-a").unwrap();
|
||||
let right = SecretCrypto::new("test-master-key-b").unwrap();
|
||||
|
||||
let left_ciphertext = left.encrypt(&plaintext).unwrap();
|
||||
let right_ciphertext = right.encrypt(&plaintext).unwrap();
|
||||
|
||||
assert_ne!(left_ciphertext, right_ciphertext);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unknown_key_version() {
|
||||
let crypto = SecretCrypto::new("test-master-key").unwrap();
|
||||
let ciphertext = crypto.encrypt(&json!({"token": "top-secret"})).unwrap();
|
||||
let error = crypto.decrypt("v999", &ciphertext).unwrap_err();
|
||||
|
||||
match error {
|
||||
RuntimeError::SecretCrypto { operation, details } => {
|
||||
assert_eq!(operation, "select secret key version");
|
||||
assert!(details.contains("unsupported secret key version"));
|
||||
}
|
||||
other => panic!("unexpected error: {other}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct WindowExecutionResult {
|
||||
pub summary: Value,
|
||||
pub items: Vec<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cursor: Option<Value>,
|
||||
pub window_complete: bool,
|
||||
pub truncated: bool,
|
||||
pub has_more: bool,
|
||||
}
|
||||
|
||||
impl From<crank_core::WindowExecutionResult> for WindowExecutionResult {
|
||||
fn from(value: crank_core::WindowExecutionResult) -> Self {
|
||||
Self {
|
||||
summary: value.summary,
|
||||
items: value.items,
|
||||
cursor: value.cursor,
|
||||
window_complete: value.window_complete,
|
||||
truncated: value.truncated,
|
||||
has_more: value.has_more,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user