chore: publish clean community baseline
Deploy / deploy (push) Successful in 37s
CI / Rust Checks (push) Successful in 5m33s
CI / UI Checks (push) Successful in 5s
CI / Deployment Manifests (push) Successful in 3s
CI / Frontend E2E (push) Successful in 4m24s

This commit is contained in:
github-ops
2026-06-17 06:15:46 +00:00
commit 338bb4d74a
309 changed files with 71223 additions and 0 deletions
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "crank-adapter-rest"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
async-trait = "0.1"
crank-core = { path = "../crank-core" }
futures-util = "0.3"
reqwest = { workspace = true, features = ["stream"] }
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]
axum.workspace = true
tokio.workspace = true
+463
View File
@@ -0,0 +1,463 @@
use std::{collections::BTreeMap, time::Duration};
use crank_core::{HttpMethod, RestTarget};
use reqwest::{
Client,
header::{HeaderMap, HeaderName, HeaderValue},
};
use serde_json::Value;
use crate::{RestAdapterError, RestRequest, RestResponse, RestWindowRequest, RestWindowResponse};
#[derive(Clone, Debug)]
pub struct RestAdapter {
client: Client,
}
impl Default for RestAdapter {
fn default() -> Self {
Self::new()
}
}
impl RestAdapter {
pub fn new() -> Self {
Self {
client: Client::new(),
}
}
pub async fn execute(
&self,
target: &RestTarget,
request: &RestRequest,
) -> Result<RestResponse, RestAdapterError> {
let url = build_url(target, request)?;
let headers = build_headers(target, request)?;
let mut builder = self
.client
.request(to_reqwest_method(target.method), url)
.headers(headers)
.timeout(Duration::from_millis(request.timeout_ms));
if let Some(body) = &request.body {
builder = builder.json(body);
}
let response = builder.send().await?;
let status = response.status();
let headers = normalize_headers(response.headers());
let body = decode_body(response).await?;
if !status.is_success() {
return Err(RestAdapterError::UnexpectedStatus {
status: status.as_u16(),
body,
});
}
Ok(RestResponse {
status_code: status.as_u16(),
headers,
body,
})
}
pub async fn execute_window(
&self,
target: &RestTarget,
request: &RestWindowRequest,
) -> Result<RestWindowResponse, RestAdapterError> {
let url = build_url(target, &request.request)?;
let mut headers = build_headers(target, &request.request)?;
headers.insert(
reqwest::header::ACCEPT,
HeaderValue::from_static("text/event-stream"),
);
let mut builder = self
.client
.request(to_reqwest_method(target.method), url)
.headers(headers)
.timeout(Duration::from_millis(request.request.timeout_ms));
if let Some(body) = &request.request.body {
builder = builder.json(body);
}
let response = builder.send().await?;
let status = response.status();
if !status.is_success() {
let headers = normalize_headers(response.headers());
let body = decode_body(response).await?;
return Err(RestAdapterError::UnexpectedStatus {
status: status.as_u16(),
body: Value::Object(
[
(
"headers".to_owned(),
serde_json::to_value(headers).unwrap_or(Value::Null),
),
("body".to_owned(), body),
]
.into_iter()
.collect(),
),
});
}
let (status_code, headers, body) =
crate::sse::collect_sse_window(response, request.window_duration_ms, request.max_items)
.await?;
Ok(RestWindowResponse {
status_code,
headers,
body,
})
}
}
fn build_url(target: &RestTarget, request: &RestRequest) -> Result<reqwest::Url, RestAdapterError> {
let base_url =
reqwest::Url::parse(&target.base_url).map_err(|_| RestAdapterError::InvalidBaseUrl {
url: target.base_url.clone(),
})?;
let path = substitute_path_params(&target.path_template, &request.path_params);
let mut url = base_url.join(path.trim_start_matches('/')).map_err(|_| {
RestAdapterError::InvalidBaseUrl {
url: target.base_url.clone(),
}
})?;
{
let mut query = url.query_pairs_mut();
for (key, value) in &request.query_params {
query.append_pair(key, value);
}
}
Ok(url)
}
fn substitute_path_params(path_template: &str, path_params: &BTreeMap<String, String>) -> String {
let mut rendered = path_template.to_owned();
for (key, value) in path_params {
rendered = rendered.replace(&format!("{{{key}}}"), value);
}
rendered
}
fn build_headers(
target: &RestTarget,
request: &RestRequest,
) -> Result<HeaderMap, RestAdapterError> {
let mut headers = HeaderMap::new();
for (name, value) in &target.static_headers {
insert_header(&mut headers, name, value)?;
}
for (name, value) in &request.headers {
insert_header(&mut headers, name, value)?;
}
Ok(headers)
}
fn insert_header(headers: &mut HeaderMap, name: &str, value: &str) -> Result<(), RestAdapterError> {
let header_name =
HeaderName::try_from(name).map_err(|_| RestAdapterError::InvalidHeaderName {
header: name.to_owned(),
})?;
let header_value =
HeaderValue::try_from(value).map_err(|_| RestAdapterError::InvalidHeaderValue {
header: name.to_owned(),
})?;
headers.insert(header_name, header_value);
Ok(())
}
async fn decode_body(response: reqwest::Response) -> Result<Value, RestAdapterError> {
let bytes = response.bytes().await?;
if bytes.is_empty() {
return Ok(Value::Null);
}
match serde_json::from_slice::<Value>(&bytes) {
Ok(value) => Ok(value),
Err(_) => Ok(Value::String(
String::from_utf8_lossy(&bytes).trim().to_owned(),
)),
}
}
fn normalize_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
headers
.iter()
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_owned(), value.to_owned()))
})
.collect()
}
fn to_reqwest_method(method: HttpMethod) -> reqwest::Method {
match method {
HttpMethod::Get => reqwest::Method::GET,
HttpMethod::Post => reqwest::Method::POST,
HttpMethod::Put => reqwest::Method::PUT,
HttpMethod::Patch => reqwest::Method::PATCH,
HttpMethod::Delete => reqwest::Method::DELETE,
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use axum::{
Json, Router,
extract::{Path, Query},
http::HeaderMap,
response::sse::{Event, KeepAlive, Sse},
routing::{get, post},
};
use crank_core::{HttpMethod, RestTarget};
use futures_util::stream;
use serde_json::{Value, json};
use tokio::net::TcpListener;
use crate::{RestAdapter, RestAdapterError, RestRequest, RestWindowRequest};
#[tokio::test]
async fn executes_rest_request_and_normalizes_json_response() {
let base_url = spawn_test_server().await;
let adapter = RestAdapter::new();
let target = RestTarget {
base_url,
method: HttpMethod::Post,
path_template: "/users/{user_id}".to_owned(),
static_headers: BTreeMap::from([("x-static".to_owned(), "static".to_owned())]),
};
let request = RestRequest {
path_params: BTreeMap::from([("user_id".to_owned(), "42".to_owned())]),
query_params: BTreeMap::from([("expand".to_owned(), "true".to_owned())]),
headers: BTreeMap::from([("x-trace-id".to_owned(), "trace-123".to_owned())]),
body: Some(json!({ "name": "Ada" })),
timeout_ms: 1_000,
};
let response = adapter.execute(&target, &request).await.unwrap();
assert_eq!(response.status_code, 200);
assert_eq!(
response.body,
json!({
"id": "42",
"query": "true",
"trace": "trace-123",
"static": "static",
"payload": { "name": "Ada" }
})
);
}
#[tokio::test]
async fn returns_unexpected_status_with_normalized_body() {
let base_url = spawn_test_server().await;
let adapter = RestAdapter::new();
let target = RestTarget {
base_url,
method: HttpMethod::Get,
path_template: "/fail".to_owned(),
static_headers: BTreeMap::new(),
};
let request = RestRequest {
path_params: BTreeMap::new(),
query_params: BTreeMap::new(),
headers: BTreeMap::new(),
body: None,
timeout_ms: 1_000,
};
let error = adapter.execute(&target, &request).await.unwrap_err();
assert!(matches!(
error,
RestAdapterError::UnexpectedStatus {
status: 502,
body: Value::Object(_)
}
));
}
#[tokio::test]
async fn collects_sse_events_with_window_bounds() {
let base_url = spawn_test_server().await;
let adapter = RestAdapter::new();
let target = RestTarget {
base_url,
method: HttpMethod::Get,
path_template: "/events".to_owned(),
static_headers: BTreeMap::new(),
};
let request = RestWindowRequest {
request: RestRequest {
path_params: BTreeMap::new(),
query_params: BTreeMap::new(),
headers: BTreeMap::new(),
body: None,
timeout_ms: 1_000,
},
window_duration_ms: 1_000,
max_items: Some(2),
};
let response = adapter.execute_window(&target, &request).await.unwrap();
assert_eq!(response.status_code, 200);
assert_eq!(
response.body,
json!({
"items": [
{ "message": "one" },
{ "message": "two" }
],
"done": false
})
);
}
#[tokio::test]
async fn returns_timeout_window_when_no_events_arrive_before_deadline() {
let base_url = spawn_test_server().await;
let adapter = RestAdapter::new();
let target = RestTarget {
base_url,
method: HttpMethod::Get,
path_template: "/events-idle".to_owned(),
static_headers: BTreeMap::new(),
};
let request = RestWindowRequest {
request: RestRequest {
path_params: BTreeMap::new(),
query_params: BTreeMap::new(),
headers: BTreeMap::new(),
body: None,
timeout_ms: 1_000,
},
window_duration_ms: 50,
max_items: Some(10),
};
let response = adapter.execute_window(&target, &request).await.unwrap();
assert_eq!(response.body, json!({ "items": [], "done": true }));
}
#[tokio::test]
async fn rejects_malformed_sse_payloads() {
let base_url = spawn_test_server().await;
let adapter = RestAdapter::new();
let target = RestTarget {
base_url,
method: HttpMethod::Get,
path_template: "/events-broken".to_owned(),
static_headers: BTreeMap::new(),
};
let request = RestWindowRequest {
request: RestRequest {
path_params: BTreeMap::new(),
query_params: BTreeMap::new(),
headers: BTreeMap::new(),
body: None,
timeout_ms: 1_000,
},
window_duration_ms: 1_000,
max_items: Some(10),
};
let error = adapter.execute_window(&target, &request).await.unwrap_err();
assert!(matches!(error, RestAdapterError::InvalidSseEvent));
}
async fn spawn_test_server() -> String {
let app = Router::new()
.route("/users/{user_id}", post(create_user))
.route("/fail", get(fail))
.route("/events", get(sse_events))
.route("/events-idle", get(sse_idle))
.route("/events-broken", get(sse_broken));
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{}", address)
}
async fn create_user(
Path(user_id): Path<String>,
Query(query): Query<BTreeMap<String, String>>,
headers: HeaderMap,
Json(payload): Json<Value>,
) -> Json<Value> {
let trace = headers
.get("x-trace-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
let static_header = headers
.get("x-static")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
Json(json!({
"id": user_id,
"query": query.get("expand").cloned().unwrap_or_default(),
"trace": trace,
"static": static_header,
"payload": payload
}))
}
async fn fail() -> (axum::http::StatusCode, Json<Value>) {
(
axum::http::StatusCode::BAD_GATEWAY,
Json(json!({ "error": "upstream failed" })),
)
}
async fn sse_events()
-> Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>> {
let events = vec![
Ok(Event::default().data("{\"message\":\"one\"}")),
Ok(Event::default().data("{\"message\":\"two\"}")),
Ok(Event::default().data("{\"message\":\"three\"}")),
];
Sse::new(stream::iter(events)).keep_alive(KeepAlive::default())
}
async fn sse_idle()
-> Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>> {
Sse::new(stream::pending()).keep_alive(KeepAlive::default())
}
async fn sse_broken()
-> Sse<impl futures_util::Stream<Item = Result<Event, std::convert::Infallible>>> {
let events = vec![Ok(Event::default().data("{broken-json}"))];
Sse::new(stream::iter(events)).keep_alive(KeepAlive::default())
}
}
+24
View File
@@ -0,0 +1,24 @@
use serde_json::Value;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum RestAdapterError {
#[error("invalid base url: {url}")]
InvalidBaseUrl { url: String },
#[error("invalid path parameter {key}")]
InvalidPathParameter { key: String },
#[error("invalid query parameter {key}")]
InvalidQueryParameter { key: String },
#[error("invalid header name {header}")]
InvalidHeaderName { header: String },
#[error("invalid header value for {header}")]
InvalidHeaderValue { header: String },
#[error("request failed")]
Transport(#[from] reqwest::Error),
#[error("sse collection window expired before stream completed")]
WindowExpired,
#[error("rest endpoint returned status {status}")]
UnexpectedStatus { status: u16, body: Value },
#[error("sse stream produced malformed event payload")]
InvalidSseEvent,
}
+123
View File
@@ -0,0 +1,123 @@
mod client;
mod error;
mod model;
mod sse;
use async_trait::async_trait;
use crank_core::{
AdapterResponse, ExecutionMode, PreparedRequest, Protocol, ProtocolAdapter,
ProtocolAdapterError, RestTarget, RuntimeRequestContext, Target, WindowExecutionResult,
};
use serde_json::{Map, Value};
pub use client::RestAdapter;
pub use error::RestAdapterError;
pub use model::{RestRequest, RestResponse, RestWindowRequest, RestWindowResponse};
#[async_trait]
impl ProtocolAdapter for RestAdapter {
fn protocol(&self) -> Protocol {
Protocol::Rest
}
fn supports_mode(&self, mode: ExecutionMode) -> bool {
matches!(mode, ExecutionMode::Unary | ExecutionMode::Window)
}
async fn invoke_unary(
&self,
target: &Target,
prepared: &PreparedRequest,
_context: &RuntimeRequestContext,
) -> Result<AdapterResponse, ProtocolAdapterError> {
let target = rest_target(target)?;
let request = RestRequest {
path_params: prepared.path_params.clone(),
query_params: prepared.query_params.clone(),
headers: prepared.headers.clone(),
body: prepared.body.clone(),
timeout_ms: prepared.timeout_ms,
};
let response = self.execute(target, &request).await?;
Ok(AdapterResponse {
status_code: response.status_code,
headers: response.headers,
data: response.body.clone(),
body: response.body,
})
}
async fn invoke_window(
&self,
target: &Target,
prepared: &PreparedRequest,
window_duration_ms: u64,
max_items: Option<u32>,
_context: &RuntimeRequestContext,
) -> Result<WindowExecutionResult, ProtocolAdapterError> {
let target = rest_target(target)?;
let request = RestWindowRequest {
request: RestRequest {
path_params: prepared.path_params.clone(),
query_params: prepared.query_params.clone(),
headers: prepared.headers.clone(),
body: prepared.body.clone(),
timeout_ms: prepared.timeout_ms,
},
window_duration_ms,
max_items: max_items.map(|value| value.saturating_add(1)),
};
let response = self.execute_window(target, &request).await?;
Ok(window_result_from_response(response.body, max_items))
}
}
fn rest_target(target: &Target) -> Result<&RestTarget, ProtocolAdapterError> {
match target {
Target::Rest(target) => Ok(target),
other => Err(ProtocolAdapterError::Message(format!(
"rest adapter cannot handle target {other:?}"
))),
}
}
fn window_result_from_response(body: Value, max_items: Option<u32>) -> WindowExecutionResult {
let done = body.get("done").and_then(Value::as_bool).unwrap_or(true);
let mut items = body
.get("items")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let mut truncated = false;
let mut has_more = !done;
if let Some(max_items) = max_items.map(|value| value as usize) {
if items.len() > max_items {
items.truncate(max_items);
truncated = true;
has_more = true;
}
}
let summary = body
.get("summary")
.cloned()
.unwrap_or_else(|| Value::Object(Map::new()));
let cursor = body.get("cursor").cloned();
WindowExecutionResult {
summary,
items,
cursor,
window_complete: !has_more,
truncated,
has_more,
}
}
impl From<RestAdapterError> for ProtocolAdapterError {
fn from(value: RestAdapterError) -> Self {
ProtocolAdapterError::Message(value.to_string())
}
}
+42
View File
@@ -0,0 +1,42 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
pub struct RestRequest {
#[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 body: Option<Value>,
pub timeout_ms: u64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RestResponse {
pub status_code: u16,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub headers: BTreeMap<String, String>,
pub body: Value,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
pub struct RestWindowRequest {
#[serde(flatten)]
pub request: RestRequest,
pub window_duration_ms: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_items: Option<u32>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RestWindowResponse {
pub status_code: u16,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub headers: BTreeMap<String, String>,
pub body: Value,
}
+160
View File
@@ -0,0 +1,160 @@
use std::collections::BTreeMap;
use futures_util::StreamExt;
use reqwest::header::HeaderMap;
use serde_json::{Value, json};
use tokio::time::{Duration, Instant, timeout_at};
use crate::RestAdapterError;
pub async fn collect_sse_window(
response: reqwest::Response,
window_duration_ms: u64,
max_items: Option<u32>,
) -> Result<(u16, BTreeMap<String, String>, Value), RestAdapterError> {
let status = response.status();
let headers = normalize_headers(response.headers());
let deadline = Instant::now() + Duration::from_millis(window_duration_ms);
let mut stream = response.bytes_stream();
let mut buffer = String::new();
let mut items = Vec::new();
let mut done = true;
loop {
if max_items.is_some_and(|limit| items.len() >= limit as usize) {
done = false;
break;
}
let next_chunk = match timeout_at(deadline, stream.next()).await {
Ok(next_chunk) => next_chunk,
Err(_) => break,
};
let Some(next_chunk) = next_chunk else {
break;
};
let chunk = next_chunk?;
buffer.push_str(&String::from_utf8_lossy(&chunk));
while let Some(event_end) = find_event_boundary(&buffer) {
let event = buffer[..event_end].to_owned();
let boundary_len = boundary_length(&buffer[event_end..]);
buffer = buffer[event_end + boundary_len..].to_owned();
if let Some(item) = parse_sse_event(&event)? {
items.push(item);
if max_items.is_some_and(|limit| items.len() >= limit as usize) {
done = false;
break;
}
}
}
if !done {
break;
}
}
Ok((
status.as_u16(),
headers,
json!({
"items": items,
"done": done
}),
))
}
fn parse_sse_event(raw: &str) -> Result<Option<Value>, RestAdapterError> {
let mut data_lines = Vec::new();
for line in raw.lines() {
if line.is_empty() || line.starts_with(':') {
continue;
}
if let Some(value) = line.strip_prefix("data:") {
data_lines.push(value.trim_start().to_owned());
}
}
if data_lines.is_empty() {
return Ok(None);
}
let payload = data_lines.join("\n");
if payload.is_empty() {
return Ok(None);
}
serde_json::from_str::<Value>(&payload)
.map(Some)
.or_else(|_| {
if payload.starts_with('{') || payload.starts_with('[') {
Err(RestAdapterError::InvalidSseEvent)
} else {
Ok(Some(Value::String(payload)))
}
})
}
fn find_event_boundary(buffer: &str) -> Option<usize> {
buffer
.find("\r\n\r\n")
.or_else(|| buffer.find("\n\n"))
.or_else(|| buffer.find("\r\r"))
}
fn boundary_length(boundary: &str) -> usize {
if boundary.starts_with("\r\n\r\n") {
4
} else {
2
}
}
fn normalize_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
headers
.iter()
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_owned(), value.to_owned()))
})
.collect()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::parse_sse_event;
use crate::RestAdapterError;
#[test]
fn parses_json_event_payload() {
let event = "event: message\ndata: {\"message\":\"ok\"}\n\n";
let parsed = parse_sse_event(event).unwrap();
assert_eq!(parsed, Some(json!({ "message": "ok" })));
}
#[test]
fn ignores_comment_only_events() {
let parsed = parse_sse_event(": keepalive\n\n").unwrap();
assert_eq!(parsed, None);
}
#[test]
fn rejects_malformed_json_like_payload() {
let error = parse_sse_event("data: {broken-json}\n\n").unwrap_err();
assert!(matches!(error, RestAdapterError::InvalidSseEvent));
}
}
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "crank-community-auth"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
argon2.workspace = true
async-trait = "0.1"
axum-extra.workspace = true
base64.workspace = true
crank-core = { path = "../crank-core" }
crank-registry = { path = "../crank-registry" }
rand.workspace = true
sha2.workspace = true
thiserror.workspace = true
time.workspace = true
tracing.workspace = true
uuid.workspace = true
@@ -0,0 +1,61 @@
#[cfg(debug_assertions)]
use argon2::{Algorithm, Params, Version};
use argon2::{
Argon2,
password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString, rand_core::OsRng},
};
#[derive(Debug, thiserror::Error)]
pub enum HashPasswordError {
#[error("failed to initialize password hasher: {0}")]
Hasher(String),
#[error("failed to hash password: {0}")]
Hash(String),
}
pub fn hash_password(password: &str, pepper: &str) -> Result<String, HashPasswordError> {
let salt = SaltString::generate(&mut OsRng);
let password = format!("{password}{pepper}");
password_hasher()?
.hash_password(password.as_bytes(), &salt)
.map(|hash| hash.to_string())
.map_err(|error| HashPasswordError::Hash(error.to_string()))
}
pub fn verify_password(password: &str, pepper: &str, password_hash: &str) -> bool {
let Ok(parsed) = PasswordHash::new(password_hash) else {
return false;
};
let password = format!("{password}{pepper}");
match password_hasher() {
Ok(argon2) => argon2.verify_password(password.as_bytes(), &parsed).is_ok(),
Err(_) => false,
}
}
fn password_hasher() -> Result<Argon2<'static>, HashPasswordError> {
#[cfg(debug_assertions)]
{
let params = Params::new(8 * 1024, 1, 1, None)
.map_err(|error| HashPasswordError::Hasher(error.to_string()))?;
Ok(Argon2::new(Algorithm::Argon2id, Version::V0x13, params))
}
#[cfg(not(debug_assertions))]
{
Ok(Argon2::default())
}
}
#[cfg(test)]
mod tests {
use super::{hash_password, verify_password};
#[test]
fn hashes_and_verifies_password() {
let password_hash = hash_password("test-password", "pepper").unwrap();
assert!(verify_password("test-password", "pepper", &password_hash));
assert!(!verify_password("wrong-password", "pepper", &password_hash));
}
}
+10
View File
@@ -0,0 +1,10 @@
pub mod hashing;
pub mod password_provider;
pub mod session_cookie;
pub use hashing::{HashPasswordError, hash_password, verify_password};
pub use password_provider::PasswordIdentityProvider;
pub use session_cookie::{
SESSION_COOKIE_NAME, SessionCookie, SessionCookieError, cleared_session_cookie,
create_session_cookie, extract_session_token, hash_session_secret, session_cookie,
};
@@ -0,0 +1,73 @@
use async_trait::async_trait;
use crank_core::{
AuthenticatedIdentity, IdentityError, IdentityProvider, IdentityProviderKind, LoginOutcome,
};
use crank_registry::PostgresRegistry;
use tracing::debug;
use crate::hashing::verify_password;
#[derive(Clone)]
pub struct PasswordIdentityProvider {
registry: PostgresRegistry,
password_pepper: String,
}
impl PasswordIdentityProvider {
pub fn new(registry: PostgresRegistry, password_pepper: impl Into<String>) -> Self {
Self {
registry,
password_pepper: password_pepper.into(),
}
}
}
#[async_trait]
impl IdentityProvider for PasswordIdentityProvider {
fn id(&self) -> &str {
"community-password"
}
fn kind(&self) -> IdentityProviderKind {
IdentityProviderKind::Password
}
async fn login_password(
&self,
payload: crank_core::LoginPayload,
) -> Result<LoginOutcome, IdentityError> {
let user = self
.registry
.get_auth_user_by_email(&payload.email)
.await
.map_err(|error| IdentityError::Internal(error.to_string()))?
.ok_or(IdentityError::BadCredentials)?;
if user.user.status != crank_core::UserStatus::Active {
return Err(IdentityError::AccountDisabled);
}
if !verify_password(
&payload.password,
&self.password_pepper,
&user.password_hash,
) {
debug!(email = %payload.email, "password identity provider rejected credentials");
return Err(IdentityError::BadCredentials);
}
let current_workspace_id = self
.registry
.list_workspaces_for_user(&user.user.id)
.await
.map_err(|error| IdentityError::Internal(error.to_string()))?
.first()
.map(|membership| membership.workspace.id.clone());
Ok(LoginOutcome::Authenticated(AuthenticatedIdentity {
user: user.user,
memberships: vec![],
current_workspace_id,
}))
}
}
@@ -0,0 +1,116 @@
use axum_extra::extract::cookie::{Cookie, CookieJar, SameSite};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use crank_core::UserSessionId;
use rand::RngCore;
use sha2::{Digest, Sha256};
use time::{Duration, OffsetDateTime};
pub const SESSION_COOKIE_NAME: &str = "crank_session";
#[derive(Clone)]
pub struct SessionCookie {
pub session_id: UserSessionId,
pub value: String,
pub expires_at: OffsetDateTime,
}
#[derive(Debug, thiserror::Error)]
pub enum SessionCookieError {
#[error("failed to compute session expiration")]
InvalidExpiration,
}
pub fn create_session_cookie(session_ttl_hours: i64) -> Result<SessionCookie, SessionCookieError> {
let session_id = UserSessionId::new(format!("sess_{}", uuid::Uuid::now_v7().simple()));
let mut secret_bytes = [0_u8; 32];
rand::thread_rng().fill_bytes(&mut secret_bytes);
let secret = URL_SAFE_NO_PAD.encode(secret_bytes);
let expires_at = OffsetDateTime::now_utc()
.checked_add(Duration::hours(session_ttl_hours))
.ok_or(SessionCookieError::InvalidExpiration)?;
Ok(SessionCookie {
session_id,
value: format!("{secret}.{}", expires_at.unix_timestamp_nanos()),
expires_at,
})
}
pub fn hash_session_secret(
session_id: &UserSessionId,
session_value: &str,
session_secret: &str,
) -> String {
let mut digest = Sha256::new();
digest.update(session_id.as_str().as_bytes());
digest.update(b":");
digest.update(session_value.as_bytes());
digest.update(b":");
digest.update(session_secret.as_bytes());
URL_SAFE_NO_PAD.encode(digest.finalize())
}
pub fn session_cookie(token: &str, cookie_secure: bool, session_ttl_hours: i64) -> Cookie<'static> {
Cookie::build((SESSION_COOKIE_NAME, token.to_owned()))
.http_only(true)
.same_site(SameSite::Lax)
.secure(cookie_secure)
.path("/")
.max_age(Duration::hours(session_ttl_hours))
.build()
}
pub fn cleared_session_cookie(cookie_secure: bool) -> Cookie<'static> {
Cookie::build((SESSION_COOKIE_NAME, String::new()))
.http_only(true)
.same_site(SameSite::Lax)
.secure(cookie_secure)
.path("/")
.max_age(Duration::seconds(0))
.build()
}
pub fn extract_session_token(jar: &CookieJar) -> Option<(UserSessionId, String)> {
let cookie = jar.get(SESSION_COOKIE_NAME)?;
let value = cookie.value();
let (session_id, session_value) = value.split_once('.')?;
Some((UserSessionId::new(session_id), session_value.to_owned()))
}
#[cfg(test)]
mod tests {
use axum_extra::extract::cookie::CookieJar;
use super::{
SESSION_COOKIE_NAME, cleared_session_cookie, create_session_cookie, extract_session_token,
session_cookie,
};
#[test]
fn extracts_session_token_from_cookie_jar() {
let token = "sess_01.secret-value";
let jar = CookieJar::new().add(session_cookie(token, false, 24));
let (session_id, secret) = extract_session_token(&jar).unwrap();
assert_eq!(session_id.as_str(), "sess_01");
assert_eq!(secret, "secret-value");
}
#[test]
fn clears_session_cookie() {
let cookie = cleared_session_cookie(false);
assert_eq!(cookie.name(), SESSION_COOKIE_NAME);
assert_eq!(cookie.value(), "");
}
#[test]
fn creates_session_cookie_with_generated_values() {
let session = create_session_cookie(24).unwrap();
assert!(session.session_id.as_str().starts_with("sess_"));
assert!(session.value.contains('.'));
}
}
+25
View File
@@ -0,0 +1,25 @@
[package]
name = "crank-community-mcp"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
async-trait = "0.1"
axum.workspace = true
base64.workspace = true
crank-core = { path = "../crank-core" }
crank-registry = { path = "../crank-registry" }
crank-runtime = { path = "../crank-runtime" }
crank-schema = { path = "../crank-schema" }
futures-util = "0.3"
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
sqlx.workspace = true
thiserror.workspace = true
time.workspace = true
tokio = { workspace = true, features = ["sync"] }
tracing.workspace = true
uuid.workspace = true
File diff suppressed because it is too large Load Diff
+22
View File
@@ -0,0 +1,22 @@
use async_trait::async_trait;
pub use crank_core::{
MachineCredentialVerifier, MachineCredentialVerifierError, SharedMachineCredentialVerifier,
VerifiedMachineCredential,
};
#[derive(Clone, Default)]
pub struct StaticAgentKeyVerifier;
#[async_trait]
impl MachineCredentialVerifier for StaticAgentKeyVerifier {
async fn verify_bearer_token(
&self,
_workspace_slug: &str,
_agent_slug: &str,
_token: &str,
) -> Result<Option<VerifiedMachineCredential>, MachineCredentialVerifierError> {
Ok(None)
}
}
pub use StaticAgentKeyVerifier as CommunityMachineCredentialVerifier;
+206
View File
@@ -0,0 +1,206 @@
use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::info;
#[derive(Clone)]
pub struct PublishedToolCatalog {
registry: PostgresRegistry,
refresh_interval: Duration,
coordination_store: Arc<dyn CoordinationStateStore>,
cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct CatalogKey {
workspace_slug: String,
agent_slug: String,
}
#[derive(Default)]
struct CachedCatalog {
loaded_at: Option<Instant>,
tools: Vec<PublishedAgentTool>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
struct CatalogSnapshot {
tools: Vec<PublishedAgentTool>,
}
impl PublishedToolCatalog {
pub fn new(
registry: PostgresRegistry,
refresh_interval: Duration,
coordination_store: Arc<dyn CoordinationStateStore>,
) -> Self {
Self {
registry,
refresh_interval,
coordination_store,
cached: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn list_tools(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<Vec<PublishedAgentTool>, RegistryError> {
self.refresh_if_stale(workspace_slug, agent_slug).await?;
let guard = self.cached.read().await;
Ok(guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.clone())
.unwrap_or_default())
}
async fn refresh_if_stale(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<(), RegistryError> {
let key = CatalogKey::new(workspace_slug, agent_slug);
let should_refresh = {
let guard = self.cached.read().await;
match guard.get(&key).and_then(|entry| entry.loaded_at) {
Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval,
None => true,
}
};
if !should_refresh {
return Ok(());
}
if let Some(tools) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
let mut guard = self.cached.write().await;
guard.insert(
key,
CachedCatalog {
loaded_at: Some(Instant::now()),
tools,
},
);
return Ok(());
}
let tools = match self
.registry
.get_published_agent_tools_by_slug(workspace_slug, agent_slug)
.await
{
Ok(tools) => tools,
Err(RegistryError::PublishedAgentNotFound { .. }) => Vec::new(),
Err(error) => return Err(error),
};
self.store_shared_snapshot(workspace_slug, agent_slug, &tools)
.await;
let mut guard = self.cached.write().await;
let previous_count = guard
.get(&key)
.map(|entry| entry.tools.len())
.unwrap_or_default();
guard.insert(
key,
CachedCatalog {
loaded_at: Some(Instant::now()),
tools,
},
);
info!(
workspace_slug,
agent_slug,
published_tool_count = guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.len())
.unwrap_or_default(),
previous_published_tool_count = previous_count,
"published agent catalog refreshed"
);
Ok(())
}
async fn load_shared_snapshot(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Option<Vec<PublishedAgentTool>> {
if self.refresh_interval.is_zero() {
return None;
}
let key = catalog_snapshot_key(workspace_slug, agent_slug);
let value = match self
.coordination_store
.get_value(CacheScope::Coordination, &key)
.await
{
Ok(value) => value?,
Err(_) => return None,
};
serde_json::from_value::<CatalogSnapshot>(value.payload)
.ok()
.map(|snapshot| snapshot.tools)
}
async fn store_shared_snapshot(
&self,
workspace_slug: &str,
agent_slug: &str,
tools: &[PublishedAgentTool],
) {
let Some(ttl) = catalog_snapshot_ttl(self.refresh_interval) else {
return;
};
let payload = match serde_json::to_value(CatalogSnapshot {
tools: tools.to_vec(),
}) {
Ok(payload) => payload,
Err(_) => return,
};
let key = catalog_snapshot_key(workspace_slug, agent_slug);
let _ = self
.coordination_store
.put_value(
CacheScope::Coordination,
&key,
CoordinationStateValue { payload },
ttl,
)
.await;
}
}
impl CatalogKey {
fn new(workspace_slug: &str, agent_slug: &str) -> Self {
Self {
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
}
}
}
fn catalog_snapshot_key(workspace_slug: &str, agent_slug: &str) -> String {
format!("published_catalog:{workspace_slug}:{agent_slug}")
}
fn catalog_snapshot_ttl(refresh_interval: Duration) -> Option<Duration> {
if refresh_interval.is_zero() {
return None;
}
refresh_interval.checked_mul(2).or(Some(refresh_interval))
}
+63
View File
@@ -0,0 +1,63 @@
use serde_json::{Value, json};
pub const CURRENT_PROTOCOL_VERSION: &str = "2025-11-25";
pub const DEFAULT_PROTOCOL_VERSION: &str = "2025-03-26";
pub const SUPPORTED_PROTOCOL_VERSIONS: &[&str] = &[
DEFAULT_PROTOCOL_VERSION,
"2025-06-18",
CURRENT_PROTOCOL_VERSION,
];
pub fn request_id(message: &Value) -> Value {
message.get("id").cloned().unwrap_or(Value::Null)
}
pub fn method_name(message: &Value) -> Option<&str> {
message.get("method").and_then(Value::as_str)
}
pub fn params(message: &Value) -> Value {
message
.get("params")
.cloned()
.unwrap_or_else(|| Value::Object(Default::default()))
}
pub fn is_notification(message: &Value) -> bool {
method_name(message).is_some() && message.get("id").is_none()
}
pub fn is_request(message: &Value) -> bool {
method_name(message).is_some() && message.get("id").is_some()
}
pub fn is_response(message: &Value) -> bool {
method_name(message).is_none()
&& (message.get("result").is_some() || message.get("error").is_some())
}
pub fn jsonrpc_result(id: Value, result: Value) -> Value {
json!({
"jsonrpc": "2.0",
"id": id,
"result": result
})
}
pub fn jsonrpc_error(id: Value, code: i64, message: impl Into<String>) -> Value {
json!({
"jsonrpc": "2.0",
"id": id,
"error": {
"code": code,
"message": message.into()
}
})
}
pub fn negotiated_protocol_version(requested: &str) -> Option<&'static str> {
SUPPORTED_PROTOCOL_VERSIONS
.iter()
.copied()
.find(|version| *version == requested)
}
+7
View File
@@ -0,0 +1,7 @@
mod app;
pub mod auth;
pub mod catalog;
pub mod jsonrpc;
pub mod session;
pub use app::build_app;
+550
View File
@@ -0,0 +1,550 @@
use std::{collections::HashMap, sync::Arc};
use async_trait::async_trait;
use crank_registry::PostgresPoolConfig;
use sqlx::{
PgPool,
postgres::{PgConnectOptions, PgPoolOptions},
query,
};
use thiserror::Error;
use time::OffsetDateTime;
use tokio::sync::RwLock;
use uuid::Uuid;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionState {
pub id: String,
pub protocol_version: String,
pub initialized: bool,
pub workspace_slug: String,
pub agent_slug: String,
pub created_at: OffsetDateTime,
pub updated_at: OffsetDateTime,
pub expires_at: Option<OffsetDateTime>,
}
#[derive(Debug, Error)]
#[error("transport session store is unavailable: {details}")]
pub struct SessionStoreError {
pub details: String,
}
#[async_trait]
pub trait TransportSessionStore: Send + Sync {
async fn create(
&self,
protocol_version: &str,
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError>;
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError>;
async fn mark_initialized(
&self,
session_id: &str,
now: OffsetDateTime,
) -> Result<bool, SessionStoreError>;
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError>;
}
pub type SharedSessionStore = Arc<dyn TransportSessionStore>;
#[derive(Clone, Debug)]
pub struct PostgresTransportSessionStore {
pool: PgPool,
}
impl PostgresTransportSessionStore {
pub async fn connect_with_options_and_pool_config(
connect_options: PgConnectOptions,
pool_config: PostgresPoolConfig,
) -> Result<Self, SessionStoreError> {
let pool = PgPoolOptions::new()
.max_connections(pool_config.max_connections)
.min_connections(pool_config.min_connections)
.acquire_timeout(std::time::Duration::from_millis(
pool_config.acquire_timeout_ms,
))
.idle_timeout(Some(std::time::Duration::from_millis(
pool_config.idle_timeout_ms,
)))
.max_lifetime(Some(std::time::Duration::from_millis(
pool_config.max_lifetime_ms,
)))
.connect_with(connect_options)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
apply_postgres_migrations(&pool).await?;
Ok(Self { pool })
}
}
#[derive(Clone, Default)]
pub struct InMemorySessionStore {
inner: Arc<RwLock<HashMap<String, SessionState>>>,
}
#[async_trait]
impl TransportSessionStore for InMemorySessionStore {
async fn create(
&self,
protocol_version: &str,
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError> {
let session_id = Uuid::now_v7().to_string();
let mut guard = self.inner.write().await;
guard.insert(
session_id.clone(),
SessionState {
id: session_id.clone(),
protocol_version: protocol_version.to_owned(),
initialized: false,
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
created_at: now,
updated_at: now,
expires_at,
},
);
Ok(session_id)
}
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
{
let guard = self.inner.read().await;
if let Some(session) = guard.get(session_id) {
if !is_expired(session, OffsetDateTime::now_utc()) {
return Ok(Some(session.clone()));
}
} else {
return Ok(None);
}
}
let mut guard = self.inner.write().await;
guard.remove(session_id);
Ok(None)
}
async fn mark_initialized(
&self,
session_id: &str,
now: OffsetDateTime,
) -> Result<bool, SessionStoreError> {
let mut guard = self.inner.write().await;
if let Some(session) = guard.get_mut(session_id) {
session.initialized = true;
session.updated_at = now;
return Ok(true);
}
Ok(false)
}
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError> {
let mut guard = self.inner.write().await;
Ok(guard.remove(session_id).is_some())
}
}
#[async_trait]
impl TransportSessionStore for PostgresTransportSessionStore {
async fn create(
&self,
protocol_version: &str,
workspace_slug: &str,
agent_slug: &str,
now: OffsetDateTime,
expires_at: Option<OffsetDateTime>,
) -> Result<String, SessionStoreError> {
let session_id = Uuid::now_v7().to_string();
query(
"insert into mcp_transport_sessions (
id,
protocol_version,
initialized,
workspace_slug,
agent_slug,
created_at,
updated_at,
expires_at
) values (
$1, $2, false, $3, $4, $5::timestamptz, $5::timestamptz, $6::timestamptz
)",
)
.bind(&session_id)
.bind(protocol_version)
.bind(workspace_slug)
.bind(agent_slug)
.bind(now)
.bind(expires_at)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(session_id)
}
async fn get(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
let row = sqlx::query!(
"select
id,
protocol_version,
initialized,
workspace_slug,
agent_slug,
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
expires_at as \"expires_at: OffsetDateTime\"
from mcp_transport_sessions
where id = $1",
session_id,
)
.fetch_optional(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
let Some(session) = row.map(|row| SessionState {
id: row.id,
protocol_version: row.protocol_version,
initialized: row.initialized,
workspace_slug: row.workspace_slug,
agent_slug: row.agent_slug,
created_at: row.created_at,
updated_at: row.updated_at,
expires_at: row.expires_at,
}) else {
return Ok(None);
};
if is_expired(&session, OffsetDateTime::now_utc()) {
query("delete from mcp_transport_sessions where id = $1")
.bind(session_id)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
return Ok(None);
}
Ok(Some(session))
}
async fn mark_initialized(
&self,
session_id: &str,
now: OffsetDateTime,
) -> Result<bool, SessionStoreError> {
let result = query(
"update mcp_transport_sessions
set initialized = true,
updated_at = $2::timestamptz
where id = $1",
)
.bind(session_id)
.bind(now)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(result.rows_affected() > 0)
}
async fn delete(&self, session_id: &str) -> Result<bool, SessionStoreError> {
let result = query("delete from mcp_transport_sessions where id = $1")
.bind(session_id)
.execute(&self.pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(result.rows_affected() > 0)
}
}
async fn apply_postgres_migrations(pool: &PgPool) -> Result<(), SessionStoreError> {
query(
"create table if not exists mcp_transport_sessions (
id text primary key,
protocol_version text not null,
initialized boolean not null default false,
workspace_slug text not null,
agent_slug text not null,
created_at timestamptz not null,
updated_at timestamptz not null,
expires_at timestamptz null
)",
)
.execute(pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
query(
"alter table mcp_transport_sessions add column if not exists expires_at timestamptz null",
)
.execute(pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
query(
"create index if not exists mcp_transport_sessions_workspace_agent_idx
on mcp_transport_sessions(workspace_slug, agent_slug, updated_at desc)",
)
.execute(pool)
.await
.map_err(|error| SessionStoreError {
details: error.to_string(),
})?;
Ok(())
}
fn is_expired(session: &SessionState, now: OffsetDateTime) -> bool {
session
.expires_at
.is_some_and(|expires_at| expires_at <= now)
}
#[cfg(test)]
mod tests {
use std::env;
use sqlx::{
Executor,
postgres::{PgConnectOptions, PgPoolOptions},
};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use uuid::Uuid;
use super::{
InMemorySessionStore, PostgresTransportSessionStore, SessionStoreError,
TransportSessionStore,
};
use crank_registry::PostgresPoolConfig;
fn timestamp(value: &str) -> time::OffsetDateTime {
time::OffsetDateTime::parse(value, &Rfc3339).unwrap()
}
fn truncate_to_micros(value: OffsetDateTime) -> OffsetDateTime {
value
.replace_nanosecond((value.nanosecond() / 1_000) * 1_000)
.unwrap()
}
#[tokio::test]
async fn creates_and_reads_transport_sessions() {
let store = InMemorySessionStore::default();
let created_at = timestamp("2026-05-01T10:00:00Z");
let session_id = store
.create("2025-11-25", "default", "sales", created_at, None)
.await
.unwrap();
let session = store.get(&session_id).await.unwrap().unwrap();
assert_eq!(session.id, session_id);
assert_eq!(session.protocol_version, "2025-11-25");
assert_eq!(session.workspace_slug, "default");
assert_eq!(session.agent_slug, "sales");
assert!(!session.initialized);
assert_eq!(session.created_at, created_at);
assert_eq!(session.updated_at, created_at);
}
#[tokio::test]
async fn marks_transport_sessions_initialized() {
let store = InMemorySessionStore::default();
let created_at = timestamp("2026-05-01T10:00:00Z");
let initialized_at = timestamp("2026-05-01T10:00:05Z");
let session_id = store
.create("2025-11-25", "default", "sales", created_at, None)
.await
.unwrap();
assert!(
store
.mark_initialized(&session_id, initialized_at)
.await
.unwrap()
);
let session = store.get(&session_id).await.unwrap().unwrap();
assert!(session.initialized);
assert_eq!(session.updated_at, initialized_at);
}
#[tokio::test]
async fn drops_expired_in_memory_transport_sessions_on_read() {
let store = InMemorySessionStore::default();
let created_at = timestamp("2026-05-01T10:00:00Z");
let expires_at = timestamp("2026-05-01T10:00:01Z");
let session_id = store
.create(
"2025-11-25",
"default",
"sales",
created_at,
Some(expires_at),
)
.await
.unwrap();
assert!(store.get(&session_id).await.unwrap().is_none());
}
#[test]
fn formats_transport_session_store_error() {
let error = SessionStoreError {
details: "postgres is unavailable".to_owned(),
};
assert_eq!(
error.to_string(),
"transport session store is unavailable: postgres is unavailable"
);
}
#[tokio::test]
async fn postgres_transport_sessions_survive_store_reconnect() {
let database_url = env::var("TEST_DATABASE_URL")
.unwrap_or_else(|_| "postgres://crank:crank@127.0.0.1:15432/crank".to_owned());
let admin_pool = PgPoolOptions::new()
.max_connections(1)
.connect(&database_url)
.await
.unwrap();
let schema = format!("test_mcp_transport_{}", Uuid::now_v7().simple());
admin_pool
.execute(sqlx::query(&format!("create schema {schema}")))
.await
.unwrap();
let connect_options = format!("{database_url}?options=-csearch_path%3D{schema}")
.parse::<PgConnectOptions>()
.unwrap();
let pool_config = PostgresPoolConfig::default();
let store_a = PostgresTransportSessionStore::connect_with_options_and_pool_config(
connect_options.clone(),
pool_config,
)
.await
.unwrap();
let created_at = OffsetDateTime::now_utc() - time::Duration::hours(1);
let initialized_at = created_at + time::Duration::seconds(5);
let session_id = store_a
.create(
"2025-11-25",
"default",
"sales",
created_at,
Some(created_at + time::Duration::days(30)),
)
.await
.unwrap();
store_a
.mark_initialized(&session_id, initialized_at)
.await
.unwrap();
let store_b = PostgresTransportSessionStore::connect_with_options_and_pool_config(
connect_options,
pool_config,
)
.await
.unwrap();
let session = store_b.get(&session_id).await.unwrap().unwrap();
assert_eq!(session.id, session_id);
assert!(session.initialized);
assert_eq!(session.updated_at, truncate_to_micros(initialized_at));
assert_eq!(session.workspace_slug, "default");
assert_eq!(session.agent_slug, "sales");
}
#[tokio::test]
async fn postgres_transport_sessions_evict_expired_rows_on_read() {
let database_url = env::var("TEST_DATABASE_URL")
.unwrap_or_else(|_| "postgres://crank:crank@127.0.0.1:15432/crank".to_owned());
let admin_pool = PgPoolOptions::new()
.max_connections(1)
.connect(&database_url)
.await
.unwrap();
let schema = format!("test_mcp_transport_{}", Uuid::now_v7().simple());
admin_pool
.execute(sqlx::query(&format!("create schema {schema}")))
.await
.unwrap();
let connect_options = format!("{database_url}?options=-csearch_path%3D{schema}")
.parse::<PgConnectOptions>()
.unwrap();
let store = PostgresTransportSessionStore::connect_with_options_and_pool_config(
connect_options.clone(),
PostgresPoolConfig::default(),
)
.await
.unwrap();
let session_id = store
.create(
"2025-11-25",
"default",
"sales",
timestamp("2026-05-01T10:00:00Z"),
Some(timestamp("2026-05-01T10:00:01Z")),
)
.await
.unwrap();
assert!(store.get(&session_id).await.unwrap().is_none());
let pool = PgPoolOptions::new()
.max_connections(1)
.connect_with(connect_options)
.await
.unwrap();
let remaining = sqlx::query_scalar::<_, i64>(
"select count(*) from mcp_transport_sessions where id = $1",
)
.bind(&session_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(remaining, 0);
}
}
+17
View File
@@ -0,0 +1,17 @@
[package]
name = "crank-core"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
async-trait = "0.1"
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
time.workspace = true
[dev-dependencies]
serde_yaml.workspace = true
tokio.workspace = true
+154
View File
@@ -0,0 +1,154 @@
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use crate::ids::{AgentId, InvitationId, PlatformApiKeyId, UserId, WorkspaceId};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UserStatus {
Active,
Disabled,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MembershipRole {
Owner,
Admin,
Operator,
Viewer,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InvitationStatus {
Pending,
Accepted,
Revoked,
Expired,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PlatformApiKeyStatus {
Active,
Revoked,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PlatformApiKeyScope {
Read,
Write,
Deploy,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct User {
pub id: UserId,
pub email: String,
pub display_name: String,
pub status: UserStatus,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Membership {
pub workspace_id: WorkspaceId,
pub user_id: UserId,
pub role: MembershipRole,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct InvitationToken {
pub id: InvitationId,
pub workspace_id: WorkspaceId,
pub email: String,
pub role: MembershipRole,
pub status: InvitationStatus,
pub token_hash: String,
#[serde(with = "time::serde::rfc3339")]
pub expires_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PlatformApiKey {
pub id: PlatformApiKeyId,
pub workspace_id: WorkspaceId,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub agent_id: Option<AgentId>,
pub name: String,
pub prefix: String,
pub scopes: Vec<PlatformApiKeyScope>,
pub status: PlatformApiKeyStatus,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339::option")]
pub last_used_at: Option<OffsetDateTime>,
}
#[cfg(test)]
mod tests {
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{PlatformApiKey, PlatformApiKeyScope, PlatformApiKeyStatus, User, UserStatus};
use crate::ids::{AgentId, PlatformApiKeyId, UserId, WorkspaceId};
#[test]
fn user_serializes_created_at_as_rfc3339() {
let user = User {
id: UserId::new("user_01"),
email: "owner@example.com".to_owned(),
display_name: "Owner".to_owned(),
status: UserStatus::Active,
created_at: OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap(),
};
let value = serde_json::to_value(&user).unwrap();
assert_eq!(value["created_at"], json!("2026-03-25T12:00:00Z"));
}
#[test]
fn user_deserializes_created_at_from_rfc3339() {
let user: User = serde_json::from_value(json!({
"id": "user_01",
"email": "owner@example.com",
"display_name": "Owner",
"status": "active",
"created_at": "2026-03-25T12:00:00Z"
}))
.unwrap();
assert_eq!(
user.created_at,
OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap()
);
}
#[test]
fn platform_api_key_serializes_timestamps_as_rfc3339() {
let api_key = PlatformApiKey {
id: PlatformApiKeyId::new("pk_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: Some(AgentId::new("agent_01")),
name: "Primary".to_owned(),
prefix: "crk_live".to_owned(),
scopes: vec![PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
status: PlatformApiKeyStatus::Active,
created_at: OffsetDateTime::parse("2026-03-25T12:01:00Z", &Rfc3339).unwrap(),
last_used_at: Some(OffsetDateTime::parse("2026-03-25T12:05:00Z", &Rfc3339).unwrap()),
};
let value = serde_json::to_value(&api_key).unwrap();
assert_eq!(value["created_at"], json!("2026-03-25T12:01:00Z"));
assert_eq!(value["last_used_at"], json!("2026-03-25T12:05:00Z"));
}
}
+102
View File
@@ -0,0 +1,102 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use crate::ids::{AgentId, OperationId, WorkspaceId};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentStatus {
Draft,
Published,
Archived,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Agent {
pub id: AgentId,
pub workspace_id: WorkspaceId,
pub slug: String,
pub display_name: String,
pub description: String,
pub status: AgentStatus,
pub current_draft_version: u32,
pub latest_published_version: Option<u32>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
#[serde(default, with = "time::serde::rfc3339::option")]
pub published_at: Option<OffsetDateTime>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentVersion {
pub agent_id: AgentId,
pub version: u32,
pub status: AgentStatus,
pub instructions: Value,
pub tool_selection_policy: Value,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentOperationBinding {
pub operation_id: OperationId,
pub operation_version: u32,
pub tool_name: String,
pub tool_title: String,
pub tool_description_override: Option<String>,
pub enabled: bool,
}
#[cfg(test)]
mod tests {
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{Agent, AgentStatus, AgentVersion};
use crate::ids::{AgentId, WorkspaceId};
#[test]
fn agent_serializes_timestamps_as_rfc3339() {
let agent = Agent {
id: AgentId::new("agent_01"),
workspace_id: WorkspaceId::new("ws_01"),
slug: "triage".to_owned(),
display_name: "Triage".to_owned(),
description: "Triage agent".to_owned(),
status: AgentStatus::Published,
current_draft_version: 2,
latest_published_version: Some(2),
created_at: OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap(),
updated_at: OffsetDateTime::parse("2026-03-25T12:05:00Z", &Rfc3339).unwrap(),
published_at: Some(OffsetDateTime::parse("2026-03-25T12:10:00Z", &Rfc3339).unwrap()),
};
let value = serde_json::to_value(&agent).unwrap();
assert_eq!(value["created_at"], json!("2026-03-25T12:00:00Z"));
assert_eq!(value["updated_at"], json!("2026-03-25T12:05:00Z"));
assert_eq!(value["published_at"], json!("2026-03-25T12:10:00Z"));
}
#[test]
fn agent_version_deserializes_created_at_from_rfc3339() {
let version: AgentVersion = serde_json::from_value(json!({
"agent_id": "agent_01",
"version": 1,
"status": "draft",
"instructions": {"system": "triage"},
"tool_selection_policy": {"mode": "allow_list"},
"created_at": "2026-03-25T12:00:00Z"
}))
.unwrap();
assert_eq!(
version.created_at,
OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap()
);
}
}
+195
View File
@@ -0,0 +1,195 @@
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use crate::{
edition::{MachineAccessMode, OperationSecurityLevel},
ids::{AgentId, AuthProfileId, OperationId, SecretId, WorkspaceId},
protocol::AuthKind,
};
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BearerAuthConfig {
pub header_name: String,
pub secret_id: SecretId,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BasicAuthConfig {
pub username_secret_id: SecretId,
pub password_secret_id: SecretId,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ApiKeyHeaderAuthConfig {
pub header_name: String,
pub secret_id: SecretId,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ApiKeyQueryAuthConfig {
pub param_name: String,
pub secret_id: SecretId,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuthConfig {
Bearer(BearerAuthConfig),
Basic(BasicAuthConfig),
ApiKeyHeader(ApiKeyHeaderAuthConfig),
ApiKeyQuery(ApiKeyQueryAuthConfig),
}
impl AuthConfig {
pub fn secret_ids(&self) -> Vec<&SecretId> {
match self {
Self::Bearer(config) => vec![&config.secret_id],
Self::Basic(config) => vec![&config.username_secret_id, &config.password_secret_id],
Self::ApiKeyHeader(config) => vec![&config.secret_id],
Self::ApiKeyQuery(config) => vec![&config.secret_id],
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AuthProfile {
pub id: AuthProfileId,
pub workspace_id: WorkspaceId,
pub name: String,
pub kind: AuthKind,
pub config: AuthConfig,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AgentTokenGrantType {
AgentKey,
RefreshToken,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct IssueAgentTokenRequest {
pub grant_type: AgentTokenGrantType,
#[serde(skip_serializing_if = "Option::is_none")]
pub agent_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub refresh_token: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub scope: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct IssueOneTimeAgentTokenRequest {
pub agent_key: String,
pub operation_id: OperationId,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub scope: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct IssuedAgentTokenResponse {
pub access_token: String,
pub token_type: String,
pub expires_in: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub refresh_token: Option<String>,
pub machine_access_mode: MachineAccessMode,
pub security_level: OperationSecurityLevel,
pub agent_id: AgentId,
}
#[cfg(test)]
mod tests {
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{
AgentTokenGrantType, ApiKeyHeaderAuthConfig, AuthConfig, AuthProfile,
IssueAgentTokenRequest, IssueOneTimeAgentTokenRequest, IssuedAgentTokenResponse,
};
use crate::{
edition::{MachineAccessMode, OperationSecurityLevel},
ids::{AuthProfileId, SecretId, WorkspaceId},
protocol::AuthKind,
};
#[test]
fn auth_profile_serializes_timestamps_as_rfc3339() {
let profile = AuthProfile {
id: AuthProfileId::new("auth_01"),
workspace_id: WorkspaceId::new("ws_default"),
name: "header".to_owned(),
kind: AuthKind::ApiKeyHeader,
config: AuthConfig::ApiKeyHeader(ApiKeyHeaderAuthConfig {
header_name: "X-Api-Key".to_owned(),
secret_id: SecretId::new("secret_01"),
}),
created_at: OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap(),
updated_at: OffsetDateTime::parse("2026-03-25T12:05:00Z", &Rfc3339).unwrap(),
};
let value = serde_json::to_value(&profile).unwrap();
assert_eq!(value["created_at"], json!("2026-03-25T12:00:00Z"));
assert_eq!(value["updated_at"], json!("2026-03-25T12:05:00Z"));
}
#[test]
fn issue_agent_token_request_serializes_stable_contract() {
let request = IssueAgentTokenRequest {
grant_type: AgentTokenGrantType::AgentKey,
agent_key: Some("crk_agent_secret".to_owned()),
refresh_token: None,
scope: vec!["tools:call".to_owned()],
};
let value = serde_json::to_value(&request).unwrap();
assert_eq!(value["grant_type"], json!("agent_key"));
assert_eq!(value["agent_key"], json!("crk_agent_secret"));
assert_eq!(value["scope"], json!(["tools:call"]));
assert_eq!(value.get("refresh_token"), None);
}
#[test]
fn issue_one_time_agent_token_request_serializes_stable_contract() {
let request = IssueOneTimeAgentTokenRequest {
agent_key: "crk_agent_secret".to_owned(),
operation_id: "op_01".into(),
scope: vec!["tools:call".to_owned()],
};
let value = serde_json::to_value(&request).unwrap();
assert_eq!(value["agent_key"], json!("crk_agent_secret"));
assert_eq!(value["operation_id"], json!("op_01"));
assert_eq!(value["scope"], json!(["tools:call"]));
}
#[test]
fn issued_agent_token_response_serializes_machine_access_metadata() {
let response = IssuedAgentTokenResponse {
access_token: "tok_01".to_owned(),
token_type: "Bearer".to_owned(),
expires_in: 300,
refresh_token: Some("rfr_01".to_owned()),
machine_access_mode: MachineAccessMode::ShortLivedToken,
security_level: OperationSecurityLevel::Elevated,
agent_id: "agt_01".into(),
};
let value = serde_json::to_value(&response).unwrap();
assert_eq!(value["access_token"], json!("tok_01"));
assert_eq!(value["token_type"], json!("Bearer"));
assert_eq!(value["expires_in"], json!(300));
assert_eq!(value["refresh_token"], json!("rfr_01"));
assert_eq!(value["machine_access_mode"], json!("short_lived_token"));
assert_eq!(value["security_level"], json!("elevated"));
assert_eq!(value["agent_id"], json!("agt_01"));
}
}
+201
View File
@@ -0,0 +1,201 @@
use std::{fmt, str::FromStr, time::Duration};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CacheBackend {
Memory,
Valkey,
Redis,
}
impl CacheBackend {
pub fn is_external(self) -> bool {
matches!(self, Self::Valkey | Self::Redis)
}
}
impl fmt::Display for CacheBackend {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let value = match self {
Self::Memory => "memory",
Self::Valkey => "valkey",
Self::Redis => "redis",
};
f.write_str(value)
}
}
impl FromStr for CacheBackend {
type Err = ParseCacheBackendError;
fn from_str(value: &str) -> Result<Self, Self::Err> {
match value {
"memory" => Ok(Self::Memory),
"valkey" => Ok(Self::Valkey),
"redis" => Ok(Self::Redis),
_ => Err(ParseCacheBackendError {
value: value.to_owned(),
}),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CacheScope {
Response,
RateLimit,
ReplayGuard,
Coordination,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CachedHeader {
pub name: String,
pub value: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CachedResponse {
pub status: u16,
pub headers: Vec<CachedHeader>,
pub body: Vec<u8>,
pub data: Vec<u8>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RateLimitBucketState {
pub tokens_micros: u64,
pub last_refill_unix_ms: i64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReplayGuardStatus {
Fresh,
AlreadySeen,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CoordinationStateValue {
pub payload: Value,
}
#[async_trait]
pub trait ResponseCacheStore: Send + Sync {
async fn get(&self, key: &str) -> Result<Option<CachedResponse>, CacheStoreError>;
async fn put(
&self,
key: &str,
value: CachedResponse,
ttl: Duration,
) -> Result<(), CacheStoreError>;
async fn delete(&self, key: &str) -> Result<(), CacheStoreError>;
}
#[async_trait]
pub trait RateLimitStateStore: Send + Sync {
async fn get_bucket(&self, key: &str) -> Result<Option<RateLimitBucketState>, CacheStoreError>;
async fn put_bucket(
&self,
key: &str,
value: RateLimitBucketState,
ttl: Duration,
) -> Result<(), CacheStoreError>;
async fn delete_bucket(&self, key: &str) -> Result<(), CacheStoreError>;
}
#[async_trait]
pub trait ReplayGuardStore: Send + Sync {
async fn mark_seen(
&self,
key: &str,
ttl: Duration,
) -> Result<ReplayGuardStatus, CacheStoreError>;
async fn clear(&self, key: &str) -> Result<(), CacheStoreError>;
}
#[async_trait]
pub trait CoordinationStateStore: Send + Sync {
async fn get_value(
&self,
scope: CacheScope,
key: &str,
) -> Result<Option<CoordinationStateValue>, CacheStoreError>;
async fn put_value(
&self,
scope: CacheScope,
key: &str,
value: CoordinationStateValue,
ttl: Duration,
) -> Result<(), CacheStoreError>;
async fn delete_value(&self, scope: CacheScope, key: &str) -> Result<(), CacheStoreError>;
}
#[derive(Debug, Error, PartialEq, Eq)]
pub enum CacheStoreError {
#[error("cache backend is unavailable: {message}")]
Unavailable { message: String },
#[error("cache key is invalid: {message}")]
InvalidKey { message: String },
#[error("cache value serialization failed: {message}")]
Serialization { message: String },
}
#[derive(Debug, Error, PartialEq, Eq)]
#[error("unsupported cache backend {value}")]
pub struct ParseCacheBackendError {
pub value: String,
}
#[cfg(test)]
mod tests {
use super::{CacheBackend, CacheScope, ParseCacheBackendError, ReplayGuardStatus};
#[test]
fn cache_backend_roundtrip_is_stable() {
for backend in [
CacheBackend::Memory,
CacheBackend::Valkey,
CacheBackend::Redis,
] {
let encoded = backend.to_string();
let decoded = encoded.parse::<CacheBackend>().unwrap();
assert_eq!(decoded, backend);
}
}
#[test]
fn rejects_unknown_cache_backend() {
assert_eq!(
"memcached".parse::<CacheBackend>(),
Err(ParseCacheBackendError {
value: "memcached".to_owned()
})
);
}
#[test]
fn distinguishes_external_cache_backends() {
assert!(!CacheBackend::Memory.is_external());
assert!(CacheBackend::Valkey.is_external());
assert!(CacheBackend::Redis.is_external());
}
#[test]
fn serialized_contracts_use_snake_case_values() {
assert_eq!(
serde_json::to_string(&CacheScope::ReplayGuard).unwrap(),
"\"replay_guard\""
);
assert_eq!(
serde_json::to_string(&ReplayGuardStatus::AlreadySeen).unwrap(),
"\"already_seen\""
);
}
}
+49
View File
@@ -0,0 +1,49 @@
use serde::{Deserialize, Serialize};
use crate::Protocol;
#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ProductEdition {
Community,
Enterprise,
Cloud,
}
#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum MachineAccessMode {
StaticAgentKey,
ShortLivedToken,
OneTimeToken,
}
#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum OperationSecurityLevel {
Standard,
Elevated,
Strict,
}
impl Default for OperationSecurityLevel {
fn default() -> Self {
Self::Standard
}
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct EditionLimits {
pub max_workspaces: Option<u32>,
pub max_users_per_workspace: Option<u32>,
pub max_agents_per_workspace: Option<u32>,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct EditionCapabilities {
pub edition: ProductEdition,
pub supported_protocols: Vec<Protocol>,
pub supported_security_levels: Vec<OperationSecurityLevel>,
pub machine_access_modes: Vec<MachineAccessMode>,
pub limits: EditionLimits,
}
+142
View File
@@ -0,0 +1,142 @@
use crate::{MembershipRole, UserId, WorkspaceId};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SessionActor {
pub user_id: UserId,
pub workspace_id: WorkspaceId,
pub role: MembershipRole,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[allow(clippy::enum_variant_names)]
pub enum PolicyAction {
ReadWorkspace,
WriteWorkspace,
ReadWorkspaceAccess,
WriteWorkspaceAccess,
ReadOperation,
WriteOperation,
ReadAgent,
WriteAgent,
ReadPlatformApiKey,
WritePlatformApiKey,
ReadSecret,
WriteSecret,
ReadAuthProfile,
WriteAuthProfile,
ReadObservability,
ReadCapability,
}
impl PolicyAction {
pub fn is_read(self) -> bool {
matches!(
self,
Self::ReadWorkspace
| Self::ReadWorkspaceAccess
| Self::ReadOperation
| Self::ReadAgent
| Self::ReadPlatformApiKey
| Self::ReadSecret
| Self::ReadAuthProfile
| Self::ReadObservability
| Self::ReadCapability
)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PolicyScope {
Workspace(WorkspaceId),
Global,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PolicyDecision {
Allow,
Deny { reason: &'static str },
}
pub trait PolicyEngine: Send + Sync {
fn check(
&self,
actor: &SessionActor,
action: PolicyAction,
scope: PolicyScope,
) -> PolicyDecision;
}
#[derive(Clone, Copy, Debug, Default)]
pub struct OwnerOnlyPolicyEngine;
impl PolicyEngine for OwnerOnlyPolicyEngine {
fn check(
&self,
actor: &SessionActor,
action: PolicyAction,
_scope: PolicyScope,
) -> PolicyDecision {
if matches!(actor.role, MembershipRole::Owner) || action.is_read() {
PolicyDecision::Allow
} else {
PolicyDecision::Deny {
reason: "write actions require owner role in community",
}
}
}
}
#[cfg(test)]
mod tests {
use super::{
OwnerOnlyPolicyEngine, PolicyAction, PolicyDecision, PolicyEngine, PolicyScope,
SessionActor,
};
use crate::{MembershipRole, UserId, WorkspaceId};
fn actor(role: MembershipRole) -> SessionActor {
SessionActor {
user_id: UserId::new("user_01"),
workspace_id: WorkspaceId::new("ws_01"),
role,
}
}
#[test]
fn owner_is_allowed_for_write_actions() {
let decision = OwnerOnlyPolicyEngine.check(
&actor(MembershipRole::Owner),
PolicyAction::WriteWorkspace,
PolicyScope::Workspace(WorkspaceId::new("ws_01")),
);
assert_eq!(decision, PolicyDecision::Allow);
}
#[test]
fn non_owner_is_allowed_for_read_actions() {
let decision = OwnerOnlyPolicyEngine.check(
&actor(MembershipRole::Viewer),
PolicyAction::ReadWorkspace,
PolicyScope::Workspace(WorkspaceId::new("ws_01")),
);
assert_eq!(decision, PolicyDecision::Allow);
}
#[test]
fn non_owner_is_denied_for_write_actions() {
let decision = OwnerOnlyPolicyEngine.check(
&actor(MembershipRole::Admin),
PolicyAction::WriteWorkspace,
PolicyScope::Workspace(WorkspaceId::new("ws_01")),
);
assert_eq!(
decision,
PolicyDecision::Deny {
reason: "write actions require owner role in community",
}
);
}
}
+94
View File
@@ -0,0 +1,94 @@
use std::sync::Arc;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use crate::{AuditEventId, UserId, UserSessionId, WorkspaceId};
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AuditActor {
pub user_id: UserId,
pub email: String,
pub session_id: Option<UserSessionId>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuditTargetKind {
Workspace,
Membership,
Invitation,
Operation,
Agent,
PlatformApiKey,
Secret,
AuthProfile,
Other,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AuditTarget {
pub workspace_id: WorkspaceId,
pub kind: AuditTargetKind,
pub id: String,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AuditEvent {
pub id: AuditEventId,
#[serde(with = "time::serde::rfc3339")]
pub occurred_at: OffsetDateTime,
pub actor: AuditActor,
pub action: String,
pub target: AuditTarget,
pub payload: Value,
pub source_ip: Option<String>,
pub user_agent: Option<String>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AuditQuery {
pub workspace_id: Option<WorkspaceId>,
pub action: Option<String>,
pub cursor: Option<String>,
pub limit: Option<u32>,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct AuditEventPage {
pub items: Vec<AuditEvent>,
pub next_cursor: Option<String>,
}
#[derive(Debug, thiserror::Error)]
pub enum AuditError {
#[error("audit sink failed: {0}")]
Backend(String),
}
#[async_trait]
pub trait AuditSink: Send + Sync {
async fn record(&self, event: AuditEvent) -> Result<(), AuditError>;
async fn list(&self, _query: AuditQuery) -> Result<AuditEventPage, AuditError> {
Ok(AuditEventPage::default())
}
async fn get(&self, _id: &AuditEventId) -> Result<Option<AuditEvent>, AuditError> {
Ok(None)
}
}
pub type SharedAuditSink = Arc<dyn AuditSink>;
#[derive(Clone, Copy, Debug, Default)]
pub struct NoopAuditSink;
#[async_trait]
impl AuditSink for NoopAuditSink {
async fn record(&self, _event: AuditEvent) -> Result<(), AuditError> {
Ok(())
}
}
+301
View File
@@ -0,0 +1,301 @@
use std::sync::Arc;
use async_trait::async_trait;
use time::OffsetDateTime;
use crate::{
IssueAgentTokenRequest, IssueOneTimeAgentTokenRequest, IssuedAgentTokenResponse,
MachineAccessMode, Membership, MembershipRole, OperationSecurityLevel, PlatformApiKeyScope,
User, UserId, WorkspaceId,
};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VerifiedMachineCredential {
pub machine_access_mode: MachineAccessMode,
pub max_security_level: OperationSecurityLevel,
pub scopes: Vec<PlatformApiKeyScope>,
}
#[derive(Debug, thiserror::Error)]
#[error("machine credential verifier failed")]
pub struct MachineCredentialVerifierError;
#[async_trait]
pub trait MachineCredentialVerifier: Send + Sync {
async fn verify_bearer_token(
&self,
workspace_slug: &str,
agent_slug: &str,
token: &str,
) -> Result<Option<VerifiedMachineCredential>, MachineCredentialVerifierError>;
}
pub type SharedMachineCredentialVerifier = Arc<dyn MachineCredentialVerifier>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TokenIssuerActor {
pub user_id: UserId,
pub workspace_id: WorkspaceId,
pub role: MembershipRole,
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum TokenIssuerError {
#[error("machine tokens are not available in this edition")]
NotSupportedInEdition,
#[error("invalid grant: {0}")]
InvalidGrant(String),
#[error("agent key is unknown")]
AgentKeyUnknown,
#[error("operation is not strict")]
OperationNotStrict,
#[error("operation is not published for agent")]
OperationNotPublishedForAgent,
#[error("registry failure: {0}")]
RegistryFailure(String),
#[error("replay guard failure: {0}")]
ReplayGuardFailure(String),
}
#[async_trait]
pub trait MachineTokenIssuer: Send + Sync {
async fn issue_short_lived(
&self,
request: IssueAgentTokenRequest,
actor: &TokenIssuerActor,
) -> Result<IssuedAgentTokenResponse, TokenIssuerError>;
async fn issue_one_time(
&self,
request: IssueOneTimeAgentTokenRequest,
actor: &TokenIssuerActor,
) -> Result<IssuedAgentTokenResponse, TokenIssuerError>;
}
pub type SharedMachineTokenIssuer = Arc<dyn MachineTokenIssuer>;
#[derive(Clone, Copy, Debug, Default)]
pub struct NoMachineTokenIssuer;
#[async_trait]
impl MachineTokenIssuer for NoMachineTokenIssuer {
async fn issue_short_lived(
&self,
_request: IssueAgentTokenRequest,
_actor: &TokenIssuerActor,
) -> Result<IssuedAgentTokenResponse, TokenIssuerError> {
Err(TokenIssuerError::NotSupportedInEdition)
}
async fn issue_one_time(
&self,
_request: IssueOneTimeAgentTokenRequest,
_actor: &TokenIssuerActor,
) -> Result<IssuedAgentTokenResponse, TokenIssuerError> {
Err(TokenIssuerError::NotSupportedInEdition)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum IdentityProviderKind {
Password,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AuthenticatedIdentity {
pub user: User,
pub memberships: Vec<Membership>,
pub current_workspace_id: Option<WorkspaceId>,
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum IdentityError {
#[error("not supported for this provider")]
NotSupportedForProvider,
#[error("bad credentials")]
BadCredentials,
#[error("account disabled")]
AccountDisabled,
#[error("internal: {0}")]
Internal(String),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LoginPayload {
pub email: String,
pub password: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SsoAuthorizeRequest {
pub workspace_id: WorkspaceId,
pub provider_id: String,
pub base_origin: String,
pub return_to: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SsoAuthorizeRedirect {
pub authorization_url: String,
pub state: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SsoCallbackRequest {
pub workspace_id: WorkspaceId,
pub provider_id: String,
pub code: String,
pub base_origin: String,
pub return_to: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorPendingIdentity {
pub user_id: UserId,
pub workspace_id: Option<WorkspaceId>,
pub return_to: Option<String>,
pub expires_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorStatus {
pub enabled: bool,
pub pending_setup: bool,
pub recovery_codes_remaining: usize,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorSetup {
pub secret: String,
pub otpauth_url: String,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TwoFactorActivation {
pub recovery_codes: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct TwoFactorChallenge {
pub code: Option<String>,
pub recovery_code: Option<String>,
}
pub enum LoginOutcome {
Authenticated(AuthenticatedIdentity),
TwoFactorRequired(TwoFactorPendingIdentity),
}
#[async_trait]
pub trait IdentityProvider: Send + Sync {
fn id(&self) -> &str;
fn kind(&self) -> IdentityProviderKind;
async fn login_password(&self, payload: LoginPayload) -> Result<LoginOutcome, IdentityError>;
async fn begin_sso(
&self,
_request: SsoAuthorizeRequest,
) -> Result<SsoAuthorizeRedirect, IdentityError> {
Err(IdentityError::NotSupportedForProvider)
}
async fn complete_sso(
&self,
_request: SsoCallbackRequest,
) -> Result<LoginOutcome, IdentityError> {
Err(IdentityError::NotSupportedForProvider)
}
async fn get_two_factor_status(
&self,
_user_id: &UserId,
) -> Result<TwoFactorStatus, IdentityError> {
Err(IdentityError::NotSupportedForProvider)
}
async fn begin_two_factor_setup(
&self,
_user_id: &UserId,
_email: &str,
) -> Result<TwoFactorSetup, IdentityError> {
Err(IdentityError::NotSupportedForProvider)
}
async fn activate_two_factor(
&self,
_user_id: &UserId,
_code: &str,
) -> Result<TwoFactorActivation, IdentityError> {
Err(IdentityError::NotSupportedForProvider)
}
async fn disable_two_factor(
&self,
_user_id: &UserId,
_challenge: TwoFactorChallenge,
) -> Result<(), IdentityError> {
Err(IdentityError::NotSupportedForProvider)
}
async fn verify_two_factor_login(
&self,
_pending: &TwoFactorPendingIdentity,
_challenge: TwoFactorChallenge,
) -> Result<AuthenticatedIdentity, IdentityError> {
Err(IdentityError::NotSupportedForProvider)
}
}
pub type SharedIdentityProvider = Arc<dyn IdentityProvider>;
#[cfg(test)]
mod tests {
use super::{MachineTokenIssuer, NoMachineTokenIssuer, TokenIssuerActor, TokenIssuerError};
use crate::{
AgentTokenGrantType, IssueAgentTokenRequest, IssueOneTimeAgentTokenRequest, MembershipRole,
UserId, WorkspaceId,
};
fn actor() -> TokenIssuerActor {
TokenIssuerActor {
user_id: UserId::new("user_01"),
workspace_id: WorkspaceId::new("ws_01"),
role: MembershipRole::Owner,
}
}
#[tokio::test]
async fn no_machine_token_issuer_rejects_short_lived_issue() {
let result = NoMachineTokenIssuer
.issue_short_lived(
IssueAgentTokenRequest {
grant_type: AgentTokenGrantType::AgentKey,
agent_key: Some("crk_agent_secret".to_owned()),
refresh_token: None,
scope: vec![],
},
&actor(),
)
.await;
assert_eq!(result, Err(TokenIssuerError::NotSupportedInEdition));
}
#[tokio::test]
async fn no_machine_token_issuer_rejects_one_time_issue() {
let result = NoMachineTokenIssuer
.issue_one_time(
IssueOneTimeAgentTokenRequest {
agent_key: "crk_agent_secret".to_owned(),
operation_id: "op_01".into(),
scope: vec![],
},
&actor(),
)
.await;
assert_eq!(result, Err(TokenIssuerError::NotSupportedInEdition));
}
}
+75
View File
@@ -0,0 +1,75 @@
use std::sync::Arc;
use async_trait::async_trait;
use crate::{TenantId, UsageRollup};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum BillingGate {
Allow,
Deny { reason: String },
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum BillingError {
#[error("billing integration is not available in this edition")]
NotSupportedInEdition,
#[error("billing provider failure: {0}")]
Provider(String),
}
#[async_trait]
pub trait BillingHook: Send + Sync {
async fn on_rollup(&self, rollup: UsageRollup) -> Result<(), BillingError>;
async fn check_billing_gate(&self, tenant: &TenantId) -> Result<BillingGate, BillingError>;
}
pub type SharedBillingHook = Arc<dyn BillingHook>;
#[derive(Clone, Copy, Debug, Default)]
pub struct NoopBillingHook;
#[async_trait]
impl BillingHook for NoopBillingHook {
async fn on_rollup(&self, _rollup: UsageRollup) -> Result<(), BillingError> {
Ok(())
}
async fn check_billing_gate(&self, _tenant: &TenantId) -> Result<BillingGate, BillingError> {
Ok(BillingGate::Allow)
}
}
#[cfg(test)]
mod tests {
use super::{BillingGate, BillingHook, NoopBillingHook};
use crate::{TenantId, UsagePeriod, UsageRollup, WorkspaceId};
#[tokio::test]
async fn noop_billing_hook_allows_default_gate_and_rollups() {
let hook = NoopBillingHook;
hook.on_rollup(UsageRollup {
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: None,
period: UsagePeriod::Last24Hours,
calls_total: 10,
calls_ok: 9,
calls_error: 1,
p50_ms: 12,
p95_ms: 42,
p99_ms: 77,
})
.await
.unwrap();
let gate = hook
.check_billing_gate(&TenantId::new("tenant_01"))
.await
.unwrap();
assert_eq!(gate, BillingGate::Allow);
}
}
+27
View File
@@ -0,0 +1,27 @@
use crate::{
EditionCapabilities, EditionLimits, MachineAccessMode, OperationSecurityLevel, ProductEdition,
Protocol,
};
pub trait CapabilityProfile: Send + Sync {
fn capabilities(&self) -> EditionCapabilities;
}
#[derive(Clone, Copy, Debug, Default)]
pub struct CommunityCapabilityProfile;
impl CapabilityProfile for CommunityCapabilityProfile {
fn capabilities(&self) -> EditionCapabilities {
EditionCapabilities {
edition: ProductEdition::Community,
supported_protocols: vec![Protocol::Rest],
supported_security_levels: vec![OperationSecurityLevel::Standard],
machine_access_modes: vec![MachineAccessMode::StaticAgentKey],
limits: EditionLimits {
max_workspaces: Some(1),
max_users_per_workspace: Some(1),
max_agents_per_workspace: Some(1),
},
}
}
}
+55
View File
@@ -0,0 +1,55 @@
use std::sync::Arc;
use async_trait::async_trait;
use time::OffsetDateTime;
use crate::{AgentId, InvocationSource, InvocationStatus, OperationId, WorkspaceId};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MeteringEvent {
pub workspace_id: WorkspaceId,
pub agent_id: Option<AgentId>,
pub operation_id: OperationId,
pub source: InvocationSource,
pub status: InvocationStatus,
pub duration_ms: u64,
pub created_at: OffsetDateTime,
}
#[async_trait]
pub trait MeteringSink: Send + Sync {
async fn record(&self, event: MeteringEvent);
}
pub type SharedMeteringSink = Arc<dyn MeteringSink>;
#[derive(Clone, Copy, Debug, Default)]
pub struct NoopMeteringSink;
#[async_trait]
impl MeteringSink for NoopMeteringSink {
async fn record(&self, _event: MeteringEvent) {}
}
#[cfg(test)]
mod tests {
use time::OffsetDateTime;
use super::{MeteringEvent, MeteringSink, NoopMeteringSink};
use crate::{InvocationSource, InvocationStatus, OperationId, WorkspaceId};
#[tokio::test]
async fn noop_metering_sink_accepts_event() {
NoopMeteringSink
.record(MeteringEvent {
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
source: InvocationSource::AgentToolCall,
status: InvocationStatus::Ok,
duration_ms: 42,
created_at: OffsetDateTime::now_utc(),
})
.await;
}
}
+8
View File
@@ -0,0 +1,8 @@
pub mod access;
pub mod audit;
pub mod auth;
pub mod billing;
pub mod capability;
pub mod metering;
pub mod protocol;
pub mod tenancy;
+322
View File
@@ -0,0 +1,322 @@
use std::{collections::BTreeMap, sync::Arc};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
AgentId, ExecutionMode, InvocationSource, Protocol, StreamSession, Target, WorkspaceId,
};
#[derive(Clone, Debug, PartialEq, Eq, Default)]
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()
}
}
#[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,
}
#[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,
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum ProtocolAdapterError {
#[error("protocol {protocol:?} does not support execution mode {mode:?}")]
UnsupportedMode {
protocol: Protocol,
mode: ExecutionMode,
},
#[error("{0}")]
Message(String),
}
#[async_trait]
pub trait ProtocolAdapter: Send + Sync {
fn protocol(&self) -> Protocol;
fn supports_mode(&self, mode: ExecutionMode) -> bool;
async fn invoke_unary(
&self,
target: &Target,
prepared: &PreparedRequest,
context: &RuntimeRequestContext,
) -> Result<AdapterResponse, ProtocolAdapterError>;
async fn invoke_window(
&self,
_target: &Target,
_prepared: &PreparedRequest,
_window_duration_ms: u64,
_max_items: Option<u32>,
_context: &RuntimeRequestContext,
) -> Result<WindowExecutionResult, ProtocolAdapterError> {
Err(ProtocolAdapterError::UnsupportedMode {
protocol: self.protocol(),
mode: ExecutionMode::Window,
})
}
async fn open_session(
&self,
_target: &Target,
_prepared: &PreparedRequest,
_context: &RuntimeRequestContext,
) -> Result<StreamSession, ProtocolAdapterError> {
Err(ProtocolAdapterError::UnsupportedMode {
protocol: self.protocol(),
mode: ExecutionMode::Session,
})
}
}
pub type SharedProtocolAdapter = Arc<dyn ProtocolAdapter>;
#[derive(Clone, Default)]
pub struct AdapterRegistry {
adapters: Vec<SharedProtocolAdapter>,
}
impl AdapterRegistry {
pub fn empty() -> Self {
Self::default()
}
pub fn register(mut self, adapter: SharedProtocolAdapter) -> Self {
self.adapters
.retain(|existing| existing.protocol() != adapter.protocol());
self.adapters.push(adapter);
self
}
pub fn get(&self, protocol: Protocol) -> Option<SharedProtocolAdapter> {
self.adapters
.iter()
.find(|adapter| adapter.protocol() == protocol)
.cloned()
}
pub fn protocols(&self) -> Vec<Protocol> {
self.adapters
.iter()
.map(|adapter| adapter.protocol())
.collect()
}
}
#[cfg(test)]
mod tests {
use std::{collections::BTreeMap, sync::Arc};
use async_trait::async_trait;
use serde_json::json;
use super::{
AdapterRegistry, AdapterResponse, PreparedRequest, ProtocolAdapter, ProtocolAdapterError,
RuntimeRequestContext,
};
use crate::{
ExecutionMode, HttpMethod, InvocationSource, Protocol, RestTarget, Target, WorkspaceId,
};
#[derive(Clone)]
struct StubAdapter(Protocol);
#[async_trait]
impl ProtocolAdapter for StubAdapter {
fn protocol(&self) -> Protocol {
self.0
}
fn supports_mode(&self, mode: ExecutionMode) -> bool {
matches!(mode, ExecutionMode::Unary)
}
async fn invoke_unary(
&self,
_target: &Target,
_prepared: &PreparedRequest,
_context: &RuntimeRequestContext,
) -> Result<AdapterResponse, ProtocolAdapterError> {
Ok(AdapterResponse {
status_code: 200,
headers: BTreeMap::new(),
body: json!({"ok": true}),
data: json!({"ok": true}),
})
}
}
#[test]
fn registry_returns_registered_protocols() {
let registry = AdapterRegistry::empty()
.register(Arc::new(StubAdapter(Protocol::Rest)))
.register(Arc::new(StubAdapter(Protocol::Graphql)));
assert_eq!(
registry.protocols(),
vec![Protocol::Rest, Protocol::Graphql]
);
}
#[test]
fn registry_replaces_adapter_for_same_protocol() {
let registry = AdapterRegistry::empty()
.register(Arc::new(StubAdapter(Protocol::Rest)))
.register(Arc::new(StubAdapter(Protocol::Rest)));
assert_eq!(registry.protocols(), vec![Protocol::Rest]);
assert!(registry.get(Protocol::Rest).is_some());
}
#[test]
fn request_context_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 request_context_attaches_metering_context() {
let context = RuntimeRequestContext::from_request_id("req_123").with_metering_context(
WorkspaceId::new("ws_01"),
None,
InvocationSource::AgentToolCall,
);
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::AgentToolCall);
}
#[tokio::test]
async fn registry_returns_registered_adapter() {
let registry = AdapterRegistry::empty().register(Arc::new(StubAdapter(Protocol::Rest)));
let adapter = registry.get(Protocol::Rest).unwrap();
let response = adapter
.invoke_unary(
&Target::Rest(RestTarget {
base_url: "https://example.test".to_owned(),
method: HttpMethod::Get,
path_template: "/health".to_owned(),
static_headers: BTreeMap::new(),
}),
&PreparedRequest::default(),
&RuntimeRequestContext::from_request_id("req_1"),
)
.await
.unwrap();
assert_eq!(response.status_code, 200);
}
}
+81
View File
@@ -0,0 +1,81 @@
use std::sync::Arc;
use async_trait::async_trait;
use crate::{TenantId, WorkspaceId};
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TenantResolutionContext {
pub workspace_id: Option<WorkspaceId>,
pub workspace_slug: Option<String>,
pub host: Option<String>,
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum TenancyError {
#[error("tenant lookup is not supported in this edition")]
UnsupportedLookup,
#[error("internal tenancy failure: {0}")]
Internal(String),
}
#[async_trait]
pub trait TenantController: Send + Sync {
fn resolve_tenant(&self, request: &TenantResolutionContext) -> Result<TenantId, TenancyError>;
async fn list_workspaces_for_tenant(
&self,
tenant: &TenantId,
) -> Result<Vec<WorkspaceId>, TenancyError>;
}
pub type SharedTenantController = Arc<dyn TenantController>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SingleTenantController {
tenant_id: TenantId,
}
impl Default for SingleTenantController {
fn default() -> Self {
Self {
tenant_id: TenantId::new("default"),
}
}
}
#[async_trait]
impl TenantController for SingleTenantController {
fn resolve_tenant(&self, _request: &TenantResolutionContext) -> Result<TenantId, TenancyError> {
Ok(self.tenant_id.clone())
}
async fn list_workspaces_for_tenant(
&self,
_tenant: &TenantId,
) -> Result<Vec<WorkspaceId>, TenancyError> {
Ok(Vec::new())
}
}
#[cfg(test)]
mod tests {
use super::{SingleTenantController, TenantController, TenantResolutionContext};
use crate::TenantId;
#[tokio::test]
async fn single_tenant_controller_returns_default_tenant() {
let controller = SingleTenantController::default();
let tenant = controller
.resolve_tenant(&TenantResolutionContext::default())
.unwrap();
let workspaces = controller
.list_workspaces_for_tenant(&tenant)
.await
.unwrap();
assert_eq!(tenant, TenantId::new("default"));
assert!(workspaces.is_empty());
}
}
+74
View File
@@ -0,0 +1,74 @@
use serde::{Deserialize, Serialize};
use std::fmt;
macro_rules! define_id {
($name:ident) => {
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct $name(String);
impl $name {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<String> for $name {
fn from(value: String) -> Self {
Self(value)
}
}
impl From<&str> for $name {
fn from(value: &str) -> Self {
Self(value.to_owned())
}
}
impl AsRef<str> for $name {
fn as_ref(&self) -> &str {
self.as_str()
}
}
impl fmt::Display for $name {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(self.as_str())
}
}
};
}
define_id!(OperationId);
define_id!(DescriptorId);
define_id!(ToolId);
define_id!(SampleId);
define_id!(AuthProfileId);
define_id!(WorkspaceId);
define_id!(TenantId);
define_id!(UserId);
define_id!(UserSessionId);
define_id!(AgentId);
define_id!(InvitationId);
define_id!(PlatformApiKeyId);
define_id!(InvocationLogId);
define_id!(AuditEventId);
define_id!(SecretId);
define_id!(StreamSessionId);
define_id!(AsyncJobId);
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ids_implement_display() {
let workspace_id = WorkspaceId::new("ws_default");
assert_eq!(workspace_id.to_string(), "ws_default");
assert_eq!(format!("{workspace_id}"), "ws_default");
}
}
+88
View File
@@ -0,0 +1,88 @@
pub mod access;
pub mod agent;
pub mod auth;
pub mod cache;
pub mod edition;
pub mod ext;
pub mod ids;
pub mod observability;
pub mod operation;
pub mod protocol;
pub mod secret;
pub mod soap;
pub mod stream_session;
pub mod streaming;
pub mod workspace;
pub use access::{
InvitationStatus, InvitationToken, Membership, MembershipRole, PlatformApiKey,
PlatformApiKeyScope, PlatformApiKeyStatus, User, UserStatus,
};
pub use agent::{Agent, AgentOperationBinding, AgentStatus, AgentVersion};
pub use auth::{
AgentTokenGrantType, ApiKeyHeaderAuthConfig, ApiKeyQueryAuthConfig, AuthConfig, AuthProfile,
BasicAuthConfig, BearerAuthConfig, IssueAgentTokenRequest, IssueOneTimeAgentTokenRequest,
IssuedAgentTokenResponse,
};
pub use cache::{
CacheBackend, CacheScope, CacheStoreError, CachedHeader, CachedResponse,
CoordinationStateStore, CoordinationStateValue, ParseCacheBackendError, RateLimitBucketState,
RateLimitStateStore, ReplayGuardStatus, ReplayGuardStore, ResponseCacheStore,
};
pub use edition::{
EditionCapabilities, EditionLimits, MachineAccessMode, OperationSecurityLevel, ProductEdition,
};
pub use ext::access::{
OwnerOnlyPolicyEngine, PolicyAction, PolicyDecision, PolicyEngine, PolicyScope, SessionActor,
};
pub use ext::audit::{
AuditActor, AuditError, AuditEvent, AuditEventPage, AuditQuery, AuditSink, AuditTarget,
AuditTargetKind, NoopAuditSink, SharedAuditSink,
};
pub use ext::auth::{
AuthenticatedIdentity, IdentityError, IdentityProvider, IdentityProviderKind, LoginOutcome,
LoginPayload, MachineCredentialVerifier, MachineCredentialVerifierError, MachineTokenIssuer,
NoMachineTokenIssuer, SharedIdentityProvider, SharedMachineCredentialVerifier,
SharedMachineTokenIssuer, SsoAuthorizeRedirect, SsoAuthorizeRequest, SsoCallbackRequest,
TokenIssuerActor, TokenIssuerError, TwoFactorActivation, TwoFactorChallenge,
TwoFactorPendingIdentity, TwoFactorSetup, TwoFactorStatus, VerifiedMachineCredential,
};
pub use ext::billing::{
BillingError, BillingGate, BillingHook, NoopBillingHook, SharedBillingHook,
};
pub use ext::capability::{CapabilityProfile, CommunityCapabilityProfile};
pub use ext::metering::{MeteringEvent, MeteringSink, NoopMeteringSink, SharedMeteringSink};
pub use ext::protocol::{
AdapterRegistry, AdapterResponse, MeteringContext, PreparedRequest, ProtocolAdapter,
ProtocolAdapterError, ResponseCacheScope, RuntimeRequestContext, SharedProtocolAdapter,
WindowExecutionResult,
};
pub use ext::tenancy::{
SharedTenantController, SingleTenantController, TenancyError, TenantController,
TenantResolutionContext,
};
pub use ids::{
AgentId, AsyncJobId, AuditEventId, AuthProfileId, DescriptorId, InvitationId, InvocationLogId,
OperationId, PlatformApiKeyId, SampleId, SecretId, StreamSessionId, TenantId, ToolId, UserId,
UserSessionId, WorkspaceId,
};
pub use observability::{
InvocationLevel, InvocationLog, InvocationSource, InvocationStatus, UsagePeriod, UsageRollup,
};
pub use operation::{
ConfigExport, ExecutionConfig, GeneratedDraft, GeneratedDraftStatus, GraphqlTarget,
GrpcProtocolOptions, GrpcTarget, Operation, OperationStatus, ProtocolOptions,
ResponseCachePolicy, RestTarget, RetryPolicy, Samples, SoapProtocolOptions, SoapTarget, Target,
ToolDescription, ToolExample, WebsocketProtocolOptions, WebsocketTarget,
};
pub use protocol::{AuthKind, ExportMode, GraphqlOperationType, HttpMethod, Protocol};
pub use secret::{Secret, SecretKind, SecretStatus, SecretVersion};
pub use soap::{
SoapBindingStyle, SoapFaultContract, SoapHeaderConfig, SoapOperationMetadata, SoapVersion,
};
pub use stream_session::{AsyncJobHandle, JobStatus, StreamSession, StreamStatus};
pub use streaming::{
AggregationMode, ExecutionMode, StreamingConfig, StreamingConfigError, ToolFamilyConfig,
TransportBehavior,
};
pub use workspace::{Workspace, WorkspaceStatus};
+132
View File
@@ -0,0 +1,132 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use crate::{AgentId, OperationId, WorkspaceId};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InvocationSource {
AdminTestRun,
AgentToolCall,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InvocationLevel {
Debug,
Info,
Warn,
Error,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InvocationStatus {
Ok,
Error,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum UsagePeriod {
#[serde(rename = "30m")]
Last30Minutes,
#[serde(rename = "1h")]
LastHour,
#[serde(rename = "6h")]
Last6Hours,
#[serde(rename = "24h")]
Last24Hours,
#[serde(rename = "7d")]
Last7Days,
#[serde(rename = "30d")]
Last30Days,
#[serde(rename = "90d")]
Last90Days,
#[serde(rename = "this_month")]
ThisMonth,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct InvocationLog {
pub id: crate::ids::InvocationLogId,
pub workspace_id: WorkspaceId,
pub agent_id: Option<AgentId>,
pub operation_id: OperationId,
pub source: InvocationSource,
pub level: InvocationLevel,
pub status: InvocationStatus,
pub tool_name: String,
pub message: String,
pub request_id: Option<String>,
pub status_code: Option<u16>,
pub duration_ms: u64,
pub error_kind: Option<String>,
pub request_preview: Value,
pub response_preview: Value,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UsageRollup {
pub workspace_id: WorkspaceId,
pub agent_id: Option<AgentId>,
pub operation_id: Option<OperationId>,
pub period: UsagePeriod,
pub calls_total: u64,
pub calls_ok: u64,
pub calls_error: u64,
pub p50_ms: u64,
pub p95_ms: u64,
pub p99_ms: u64,
}
#[cfg(test)]
mod tests {
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{InvocationLevel, InvocationLog, InvocationSource, InvocationStatus, UsagePeriod};
use crate::{AgentId, OperationId, WorkspaceId, ids::InvocationLogId};
fn timestamp(value: &str) -> OffsetDateTime {
OffsetDateTime::parse(value, &Rfc3339).unwrap()
}
#[test]
fn invocation_log_serializes_timestamps_as_rfc3339() {
let log = InvocationLog {
id: InvocationLogId::new("log_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: Some(AgentId::new("agent_01")),
operation_id: OperationId::new("op_01"),
source: InvocationSource::AdminTestRun,
level: InvocationLevel::Info,
status: InvocationStatus::Ok,
tool_name: "create_lead".to_owned(),
message: "ok".to_owned(),
request_id: Some("req_01".to_owned()),
status_code: Some(200),
duration_ms: 123,
error_kind: None,
request_preview: json!({"input": "value"}),
response_preview: json!({"ok": true}),
created_at: timestamp("2026-04-19T12:34:56Z"),
};
let value = serde_json::to_value(&log).unwrap();
assert_eq!(value["created_at"], "2026-04-19T12:34:56Z");
assert_eq!(value["source"], "admin_test_run");
assert_eq!(value["level"], "info");
assert_eq!(value["status"], "ok");
}
#[test]
fn usage_period_serializes_compact_labels() {
let value = serde_json::to_value(UsagePeriod::Last7Days).unwrap();
assert_eq!(value, "7d");
}
}
+586
View File
@@ -0,0 +1,586 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use crate::{
edition::OperationSecurityLevel,
ids::{AuthProfileId, DescriptorId, OperationId, SampleId},
protocol::{ExportMode, GraphqlOperationType, HttpMethod, Protocol},
soap::{
SoapBindingStyle, SoapFaultContract, SoapHeaderConfig, SoapOperationMetadata, SoapVersion,
},
streaming::StreamingConfig,
};
fn default_operation_category() -> String {
"general".to_owned()
}
fn default_operation_security_level() -> OperationSecurityLevel {
OperationSecurityLevel::Standard
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OperationStatus {
Draft,
Testing,
Published,
Archived,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RestTarget {
pub base_url: String,
pub method: HttpMethod,
pub path_template: String,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub static_headers: BTreeMap<String, String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct GraphqlTarget {
pub endpoint: String,
pub operation_type: GraphqlOperationType,
pub operation_name: String,
pub query_template: String,
pub response_path: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct GrpcTarget {
pub server_addr: String,
pub package: String,
pub service: String,
pub method: String,
#[serde(default)]
pub read_only: bool,
pub descriptor_ref: DescriptorId,
pub descriptor_set_b64: String,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WebsocketTarget {
pub url: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub subprotocols: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub subscribe_message_template: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub unsubscribe_message_template: Option<Value>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub static_headers: BTreeMap<String, String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SoapTarget {
pub wsdl_ref: SampleId,
pub service_name: String,
pub port_name: String,
pub operation_name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub endpoint_override: Option<String>,
pub soap_version: SoapVersion,
#[serde(skip_serializing_if = "Option::is_none")]
pub soap_action: Option<String>,
pub binding_style: SoapBindingStyle,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub headers: Vec<SoapHeaderConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub fault_contract: Option<SoapFaultContract>,
#[serde(default)]
pub metadata: SoapOperationMetadata,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Target {
Rest(RestTarget),
Graphql(GraphqlTarget),
Grpc(GrpcTarget),
Websocket(WebsocketTarget),
Soap(SoapTarget),
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RetryPolicy {
pub max_attempts: u32,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResponseCachePolicy {
pub ttl_ms: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct GrpcProtocolOptions {
pub use_tls: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct WebsocketProtocolOptions {
#[serde(skip_serializing_if = "Option::is_none")]
pub heartbeat_interval_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reconnect_max_attempts: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reconnect_backoff_ms: Option<u64>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct SoapProtocolOptions {
#[serde(default)]
pub use_tls: bool,
#[serde(default)]
pub validate_certificate: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub ws_security_profile: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct ProtocolOptions {
#[serde(skip_serializing_if = "Option::is_none")]
pub grpc: Option<GrpcProtocolOptions>,
#[serde(skip_serializing_if = "Option::is_none")]
pub websocket: Option<WebsocketProtocolOptions>,
#[serde(skip_serializing_if = "Option::is_none")]
pub soap: Option<SoapProtocolOptions>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ExecutionConfig {
pub timeout_ms: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub retry_policy: Option<RetryPolicy>,
#[serde(skip_serializing_if = "Option::is_none")]
pub response_cache: Option<ResponseCachePolicy>,
#[serde(skip_serializing_if = "Option::is_none")]
pub auth_profile_ref: Option<AuthProfileId>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub headers: BTreeMap<String, String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub protocol_options: Option<ProtocolOptions>,
#[serde(skip_serializing_if = "Option::is_none")]
pub streaming: Option<StreamingConfig>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolExample {
pub input: Value,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolDescription {
pub title: String,
pub description: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub tags: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub examples: Vec<ToolExample>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct Samples {
#[serde(skip_serializing_if = "Option::is_none")]
pub input_json_sample_ref: Option<SampleId>,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_json_sample_ref: Option<SampleId>,
#[serde(skip_serializing_if = "Option::is_none")]
pub proto_file_ref: Option<SampleId>,
#[serde(skip_serializing_if = "Option::is_none")]
pub descriptor_ref: Option<DescriptorId>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum GeneratedDraftStatus {
None,
Available,
Stale,
Failed,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct GeneratedDraft {
pub status: GeneratedDraftStatus,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub source_types: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub generated_at: Option<String>,
pub input_schema_generated: bool,
pub output_schema_generated: bool,
pub input_mapping_generated: bool,
pub output_mapping_generated: bool,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub warnings: Vec<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ConfigExport {
pub format_version: String,
pub export_mode: ExportMode,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Operation<TSchema, TMapping> {
pub id: OperationId,
pub name: String,
pub display_name: String,
#[serde(default = "default_operation_category")]
pub category: String,
pub protocol: Protocol,
#[serde(default = "default_operation_security_level")]
pub security_level: OperationSecurityLevel,
pub status: OperationStatus,
pub version: u32,
pub target: Target,
pub input_schema: TSchema,
pub output_schema: TSchema,
pub input_mapping: TMapping,
pub output_mapping: TMapping,
pub execution_config: ExecutionConfig,
pub tool_description: ToolDescription,
#[serde(skip_serializing_if = "Option::is_none")]
pub samples: Option<Samples>,
#[serde(skip_serializing_if = "Option::is_none")]
pub generated_draft: Option<GeneratedDraft>,
#[serde(skip_serializing_if = "Option::is_none")]
pub config_export: Option<ConfigExport>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub published_at: Option<OffsetDateTime>,
}
impl<TSchema, TMapping> Operation<TSchema, TMapping> {
pub fn tool_name(&self) -> &str {
&self.name
}
pub fn is_draft(&self) -> bool {
self.status == OperationStatus::Draft
}
pub fn is_published(&self) -> bool {
self.status == OperationStatus::Published
}
pub fn protocol(&self) -> Protocol {
self.protocol
}
pub fn auth_profile_ref(&self) -> Option<&AuthProfileId> {
self.execution_config.auth_profile_ref.as_ref()
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use crate::{
auth::{AuthConfig, AuthProfile, BearerAuthConfig},
edition::OperationSecurityLevel,
ids::{AuthProfileId, OperationId, SampleId},
operation::{
ConfigExport, ExecutionConfig, GraphqlTarget, Operation, OperationStatus,
ProtocolOptions, RestTarget, Samples, SoapProtocolOptions, SoapTarget, Target,
ToolDescription, ToolExample, WebsocketProtocolOptions, WebsocketTarget,
},
protocol::{AuthKind, ExportMode, GraphqlOperationType, HttpMethod, Protocol},
soap::{
SoapBindingStyle, SoapFaultContract, SoapHeaderConfig, SoapOperationMetadata,
SoapVersion,
},
};
fn timestamp(value: &str) -> OffsetDateTime {
OffsetDateTime::parse(value, &Rfc3339).unwrap()
}
#[test]
fn rest_target_serializes_with_kind_tag() {
let target = Target::Rest(RestTarget {
base_url: "https://api.example.com".to_owned(),
method: HttpMethod::Post,
path_template: "/v1/leads".to_owned(),
static_headers: BTreeMap::new(),
});
let value = serde_json::to_value(target).unwrap();
assert_eq!(value["kind"], "rest");
assert_eq!(value["method"], "POST");
}
#[test]
fn graphql_target_serializes_response_path() {
let target = Target::Graphql(GraphqlTarget {
endpoint: "https://api.example.com/graphql".to_owned(),
operation_type: GraphqlOperationType::Mutation,
operation_name: "CreateLead".to_owned(),
query_template: "mutation {}".to_owned(),
response_path: "$.response.body.data.createLead".to_owned(),
});
let value = serde_json::to_value(target).unwrap();
assert_eq!(value["kind"], "graphql");
assert_eq!(value["operation_type"], "mutation");
}
#[test]
fn websocket_target_serializes_templates() {
let target = Target::Websocket(WebsocketTarget {
url: "wss://events.example.com/stream".to_owned(),
subprotocols: vec!["graphql-transport-ws".to_owned()],
subscribe_message_template: Some(json!({ "type": "subscribe" })),
unsubscribe_message_template: Some(json!({ "type": "unsubscribe" })),
static_headers: BTreeMap::from([("x-env".to_owned(), "test".to_owned())]),
});
let value = serde_json::to_value(target).unwrap();
assert_eq!(value["kind"], "websocket");
assert_eq!(value["subprotocols"][0], "graphql-transport-ws");
assert_eq!(value["subscribe_message_template"]["type"], "subscribe");
}
#[test]
fn soap_target_serializes_binding_metadata() {
let target = Target::Soap(SoapTarget {
wsdl_ref: SampleId::new("sample_wsdl_01"),
service_name: "LeadService".to_owned(),
port_name: "LeadPort".to_owned(),
operation_name: "CreateLead".to_owned(),
endpoint_override: Some("https://soap.example.com/lead".to_owned()),
soap_version: SoapVersion::Soap12,
soap_action: Some("urn:createLead".to_owned()),
binding_style: SoapBindingStyle::DocumentLiteral,
headers: vec![SoapHeaderConfig {
name: "CorrelationId".to_owned(),
namespace_uri: Some("urn:crm".to_owned()),
required: true,
value_path: Some("$.mcp.correlation_id".to_owned()),
}],
fault_contract: Some(SoapFaultContract {
code_path: Some("$.Envelope.Body.Fault.Code.Value".to_owned()),
message_path: Some("$.Envelope.Body.Fault.Reason.Text".to_owned()),
detail_path: Some("$.Envelope.Body.Fault.Detail".to_owned()),
}),
metadata: SoapOperationMetadata {
input_part_names: vec!["LeadRequest".to_owned()],
output_part_names: vec!["LeadResponse".to_owned()],
namespaces: vec!["urn:crm".to_owned()],
},
});
let value = serde_json::to_value(target).unwrap();
assert_eq!(value["kind"], "soap");
assert_eq!(value["soap_version"], "soap_12");
assert_eq!(value["binding_style"], "document_literal");
assert_eq!(value["metadata"]["input_part_names"][0], "LeadRequest");
}
#[test]
fn operation_exposes_local_domain_helpers() {
let operation = Operation {
id: OperationId::new("op_01"),
name: "crm_create_lead".to_owned(),
display_name: "Create Lead".to_owned(),
category: "sales".to_owned(),
protocol: Protocol::Rest,
security_level: OperationSecurityLevel::Standard,
status: OperationStatus::Draft,
version: 1,
target: Target::Rest(RestTarget {
base_url: "https://api.example.com".to_owned(),
method: HttpMethod::Post,
path_template: "/v1/leads".to_owned(),
static_headers: BTreeMap::new(),
}),
input_schema: json!({"type":"object"}),
output_schema: json!({"type":"object"}),
input_mapping: json!({"rules":[]}),
output_mapping: json!({"rules":[]}),
execution_config: ExecutionConfig {
timeout_ms: 10_000,
retry_policy: None,
response_cache: None,
auth_profile_ref: Some(AuthProfileId::new("auth_01")),
headers: BTreeMap::new(),
protocol_options: Some(ProtocolOptions {
grpc: None,
websocket: Some(WebsocketProtocolOptions {
heartbeat_interval_ms: Some(5_000),
reconnect_max_attempts: Some(3),
reconnect_backoff_ms: Some(250),
}),
soap: Some(SoapProtocolOptions {
use_tls: true,
validate_certificate: true,
ws_security_profile: Some("username_token".to_owned()),
}),
}),
streaming: None,
},
tool_description: ToolDescription {
title: "Create CRM lead".to_owned(),
description: "Creates a new lead.".to_owned(),
tags: Vec::new(),
examples: Vec::new(),
},
samples: Some(Samples::default()),
generated_draft: None,
config_export: None,
created_at: timestamp("2026-03-25T08:00:00Z"),
updated_at: timestamp("2026-03-25T08:00:00Z"),
published_at: None,
};
assert_eq!(operation.tool_name(), "crm_create_lead");
assert!(operation.is_draft());
assert!(!operation.is_published());
assert_eq!(operation.protocol(), Protocol::Rest);
assert_eq!(
operation.auth_profile_ref().map(|value| value.as_str()),
Some("auth_01")
);
}
#[test]
fn auth_profile_serializes_secret_ids_without_secret_values() {
let profile = AuthProfile {
id: AuthProfileId::new("auth_01"),
workspace_id: crate::ids::WorkspaceId::new("ws_01"),
name: "crm-prod-bearer".to_owned(),
kind: AuthKind::Bearer,
config: AuthConfig::Bearer(BearerAuthConfig {
header_name: "Authorization".to_owned(),
secret_id: crate::ids::SecretId::new("secret_crm_prod_token"),
}),
created_at: timestamp("2026-03-25T08:00:00Z"),
updated_at: timestamp("2026-03-25T08:00:00Z"),
};
let value = serde_json::to_value(profile).unwrap();
assert_eq!(value["kind"], "bearer");
assert_eq!(
value["config"]["bearer"]["secret_id"],
"secret_crm_prod_token"
);
}
#[test]
fn operation_roundtrips_through_yaml() {
let operation = Operation {
id: OperationId::new("op_01"),
name: "crm_create_lead".to_owned(),
display_name: "Create Lead".to_owned(),
category: "sales".to_owned(),
protocol: Protocol::Rest,
security_level: OperationSecurityLevel::Standard,
status: OperationStatus::Published,
version: 3,
target: Target::Rest(RestTarget {
base_url: "https://api.example.com".to_owned(),
method: HttpMethod::Post,
path_template: "/v1/leads".to_owned(),
static_headers: BTreeMap::from([("X-App-Source".to_owned(), "crank".to_owned())]),
}),
input_schema: json!({"type":"object","fields":{"email":{"type":"string","required":true}}}),
output_schema: json!({"type":"object","fields":{"id":{"type":"string","required":true}}}),
input_mapping: json!({"rules":[{"source":"$.mcp.email","target":"$.request.body.email"}]}),
output_mapping: json!({"rules":[{"source":"$.response.body.id","target":"$.output.id"}]}),
execution_config: ExecutionConfig {
timeout_ms: 10_000,
retry_policy: None,
response_cache: None,
auth_profile_ref: Some(AuthProfileId::new("auth_01")),
headers: BTreeMap::new(),
protocol_options: Some(ProtocolOptions::default()),
streaming: None,
},
tool_description: ToolDescription {
title: "Create CRM lead".to_owned(),
description: "Creates a new lead.".to_owned(),
tags: vec!["crm".to_owned()],
examples: vec![ToolExample {
input: json!({"email":"user@example.com"}),
}],
},
samples: Some(Samples::default()),
generated_draft: None,
config_export: Some(ConfigExport {
format_version: "1".to_owned(),
export_mode: ExportMode::Portable,
}),
created_at: timestamp("2026-03-25T08:00:00Z"),
updated_at: timestamp("2026-03-25T08:10:00Z"),
published_at: Some(timestamp("2026-03-25T08:15:00Z")),
};
let yaml = serde_yaml::to_string(&operation).unwrap();
let restored: Operation<serde_json::Value, serde_json::Value> =
serde_yaml::from_str(&yaml).unwrap();
assert!(yaml.contains("protocol: rest"));
assert!(yaml.contains("security_level: standard"));
assert!(yaml.contains("export_mode: portable"));
assert_eq!(restored, operation);
}
#[test]
fn operation_yaml_deserializes_without_optional_published_at() {
let yaml = r#"
id: op_01
name: crm_create_lead
display_name: Create Lead
category: sales
protocol: rest
security_level: standard
status: draft
version: 1
target:
kind: rest
base_url: https://api.example.com
method: POST
path_template: /v1/leads
static_headers: {}
input_schema:
type: object
output_schema:
type: object
input_mapping:
rules: []
output_mapping:
rules: []
execution_config:
timeout_ms: 10000
headers: {}
tool_description:
title: Create CRM lead
description: Creates a new lead.
tags: []
examples: []
created_at: 2026-03-25T08:00:00Z
updated_at: 2026-03-25T08:10:00Z
"#;
let restored: Operation<serde_json::Value, serde_json::Value> =
serde_yaml::from_str(yaml).unwrap();
assert_eq!(restored.security_level, OperationSecurityLevel::Standard);
assert_eq!(restored.published_at, None);
}
}
+116
View File
@@ -0,0 +1,116 @@
use serde::{Deserialize, Serialize};
use crate::streaming::{ExecutionMode, TransportBehavior};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Protocol {
Rest,
Graphql,
Grpc,
Websocket,
Soap,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
pub enum HttpMethod {
Get,
Post,
Put,
Patch,
Delete,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum GraphqlOperationType {
Query,
Mutation,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AuthKind {
Bearer,
Basic,
ApiKeyHeader,
ApiKeyQuery,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExportMode {
Portable,
Bundle,
}
impl Protocol {
pub fn supports_execution_mode(self, mode: ExecutionMode) -> bool {
match self {
Self::Rest => true,
Self::Graphql => matches!(mode, ExecutionMode::Unary),
Self::Grpc => true,
Self::Websocket => !matches!(mode, ExecutionMode::Unary),
Self::Soap => matches!(mode, ExecutionMode::Unary | ExecutionMode::AsyncJob),
}
}
pub fn supports_transport_behavior(self, behavior: TransportBehavior) -> bool {
match self {
Self::Rest => true,
Self::Graphql => matches!(behavior, TransportBehavior::RequestResponse),
Self::Grpc => !matches!(behavior, TransportBehavior::DeferredResult),
Self::Websocket => !matches!(behavior, TransportBehavior::RequestResponse),
Self::Soap => !matches!(behavior, TransportBehavior::StatefulSession),
}
}
}
#[cfg(test)]
mod tests {
use super::Protocol;
use crate::streaming::{ExecutionMode, TransportBehavior};
#[test]
fn graphql_support_matrix_is_restricted() {
assert!(Protocol::Graphql.supports_execution_mode(ExecutionMode::Unary));
assert!(!Protocol::Graphql.supports_execution_mode(ExecutionMode::Window));
assert!(!Protocol::Graphql.supports_execution_mode(ExecutionMode::Session));
assert!(!Protocol::Graphql.supports_transport_behavior(TransportBehavior::ServerStream));
}
#[test]
fn grpc_supports_session_but_not_deferred_result() {
assert!(Protocol::Grpc.supports_execution_mode(ExecutionMode::Session));
assert!(!Protocol::Grpc.supports_transport_behavior(TransportBehavior::DeferredResult));
}
#[test]
fn websocket_requires_streaming_modes() {
assert!(!Protocol::Websocket.supports_execution_mode(ExecutionMode::Unary));
assert!(Protocol::Websocket.supports_execution_mode(ExecutionMode::Window));
assert!(Protocol::Websocket.supports_execution_mode(ExecutionMode::Session));
assert!(Protocol::Websocket.supports_execution_mode(ExecutionMode::AsyncJob));
assert!(
!Protocol::Websocket.supports_transport_behavior(TransportBehavior::RequestResponse)
);
assert!(Protocol::Websocket.supports_transport_behavior(TransportBehavior::ServerStream));
assert!(
Protocol::Websocket.supports_transport_behavior(TransportBehavior::StatefulSession)
);
assert!(Protocol::Websocket.supports_transport_behavior(TransportBehavior::DeferredResult));
}
#[test]
fn soap_is_request_response_first_with_async_job_support() {
assert!(Protocol::Soap.supports_execution_mode(ExecutionMode::Unary));
assert!(!Protocol::Soap.supports_execution_mode(ExecutionMode::Window));
assert!(!Protocol::Soap.supports_execution_mode(ExecutionMode::Session));
assert!(Protocol::Soap.supports_execution_mode(ExecutionMode::AsyncJob));
assert!(Protocol::Soap.supports_transport_behavior(TransportBehavior::RequestResponse));
assert!(Protocol::Soap.supports_transport_behavior(TransportBehavior::ServerStream));
assert!(!Protocol::Soap.supports_transport_behavior(TransportBehavior::StatefulSession));
assert!(Protocol::Soap.supports_transport_behavior(TransportBehavior::DeferredResult));
}
}
+96
View File
@@ -0,0 +1,96 @@
use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use crate::ids::{SecretId, UserId, WorkspaceId};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SecretKind {
Token,
UsernamePassword,
Header,
Generic,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SecretStatus {
Active,
Disabled,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Secret {
pub id: SecretId,
pub workspace_id: WorkspaceId,
pub name: String,
pub kind: SecretKind,
pub status: SecretStatus,
pub current_version: u32,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339::option")]
pub last_used_at: Option<OffsetDateTime>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SecretVersion {
pub secret_id: SecretId,
pub version: u32,
pub ciphertext: String,
pub key_version: String,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
pub created_by: Option<UserId>,
}
#[cfg(test)]
mod tests {
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{Secret, SecretKind, SecretStatus, SecretVersion};
use crate::ids::{SecretId, UserId, WorkspaceId};
#[test]
fn secret_serializes_timestamps_as_rfc3339() {
let secret = Secret {
id: SecretId::new("secret_01"),
workspace_id: WorkspaceId::new("ws_01"),
name: "primary".to_owned(),
kind: SecretKind::Token,
status: SecretStatus::Active,
current_version: 2,
created_at: OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap(),
updated_at: OffsetDateTime::parse("2026-03-25T12:05:00Z", &Rfc3339).unwrap(),
last_used_at: Some(OffsetDateTime::parse("2026-03-25T12:07:00Z", &Rfc3339).unwrap()),
};
let value = serde_json::to_value(&secret).unwrap();
assert_eq!(value["created_at"], json!("2026-03-25T12:00:00Z"));
assert_eq!(value["updated_at"], json!("2026-03-25T12:05:00Z"));
assert_eq!(value["last_used_at"], json!("2026-03-25T12:07:00Z"));
}
#[test]
fn secret_version_deserializes_created_at_from_rfc3339() {
let version: SecretVersion = serde_json::from_value(json!({
"secret_id": "secret_01",
"version": 2,
"ciphertext": "opaque",
"key_version": "v2",
"created_at": "2026-03-25T12:07:00Z",
"created_by": "user_01"
}))
.unwrap();
assert_eq!(
version.created_at,
OffsetDateTime::parse("2026-03-25T12:07:00Z", &Rfc3339).unwrap()
);
assert_eq!(version.created_by, Some(UserId::new("user_01")));
}
}
+101
View File
@@ -0,0 +1,101 @@
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum SoapVersion {
#[serde(rename = "soap_11")]
Soap11,
#[serde(rename = "soap_12")]
Soap12,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SoapBindingStyle {
DocumentLiteral,
RpcLiteral,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SoapHeaderConfig {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub namespace_uri: Option<String>,
pub required: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub value_path: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SoapFaultContract {
#[serde(skip_serializing_if = "Option::is_none")]
pub code_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub message_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub detail_path: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct SoapOperationMetadata {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub input_part_names: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub output_part_names: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub namespaces: Vec<String>,
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
SoapBindingStyle, SoapFaultContract, SoapHeaderConfig, SoapOperationMetadata, SoapVersion,
};
#[test]
fn soap_metadata_serializes_compactly() {
let value = serde_json::to_value(SoapOperationMetadata {
input_part_names: vec!["LeadRequest".to_owned()],
output_part_names: vec!["LeadResponse".to_owned()],
namespaces: vec!["urn:crm".to_owned()],
})
.unwrap();
assert_eq!(value["input_part_names"][0], "LeadRequest");
assert_eq!(value["output_part_names"][0], "LeadResponse");
assert_eq!(value["namespaces"][0], "urn:crm");
}
#[test]
fn soap_supporting_types_roundtrip() {
let value = json!({
"version": "soap_12",
"style": "document_literal",
"header": {
"name": "CorrelationId",
"namespace_uri": "urn:crm",
"required": true,
"value_path": "$.mcp.correlation_id"
},
"fault": {
"code_path": "$.Envelope.Body.Fault.Code.Value",
"message_path": "$.Envelope.Body.Fault.Reason.Text",
"detail_path": "$.Envelope.Body.Fault.Detail"
}
});
let version: SoapVersion = serde_json::from_value(value["version"].clone()).unwrap();
let style: SoapBindingStyle = serde_json::from_value(value["style"].clone()).unwrap();
let header: SoapHeaderConfig = serde_json::from_value(value["header"].clone()).unwrap();
let fault: SoapFaultContract = serde_json::from_value(value["fault"].clone()).unwrap();
assert_eq!(version, SoapVersion::Soap12);
assert_eq!(style, SoapBindingStyle::DocumentLiteral);
assert_eq!(header.name, "CorrelationId");
assert_eq!(
fault.detail_path.as_deref(),
Some("$.Envelope.Body.Fault.Detail")
);
}
}
+318
View File
@@ -0,0 +1,318 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use crate::{
ids::{AgentId, AsyncJobId, OperationId, StreamSessionId, WorkspaceId},
protocol::Protocol,
streaming::ExecutionMode,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StreamStatus {
Created,
Running,
Stopped,
Failed,
Expired,
}
impl StreamStatus {
pub fn is_terminal(self) -> bool {
matches!(self, Self::Stopped | Self::Failed | Self::Expired)
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StreamSession {
pub id: StreamSessionId,
pub workspace_id: WorkspaceId,
#[serde(skip_serializing_if = "Option::is_none")]
pub agent_id: Option<AgentId>,
pub operation_id: OperationId,
pub protocol: Protocol,
pub mode: ExecutionMode,
pub status: StreamStatus,
#[serde(skip_serializing_if = "Option::is_none")]
pub cursor: Option<Value>,
pub state: Value,
#[serde(with = "time::serde::rfc3339")]
pub expires_at: OffsetDateTime,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub last_poll_at: Option<OffsetDateTime>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub closed_at: Option<OffsetDateTime>,
}
impl StreamSession {
pub fn is_expired(&self, now: OffsetDateTime) -> bool {
self.expires_at <= now
}
pub fn can_poll(&self, now: OffsetDateTime) -> bool {
!self.status.is_terminal() && !self.is_expired(now)
}
pub fn remaining_poll_delay_ms(&self, now: OffsetDateTime, poll_interval_ms: u64) -> u64 {
let Some(last_poll_at) = self.last_poll_at else {
return 0;
};
let elapsed_ms = (now - last_poll_at).whole_milliseconds().max(0) as u64;
poll_interval_ms.saturating_sub(elapsed_ms)
}
pub fn mark_polled(&mut self, now: OffsetDateTime) {
self.last_poll_at = Some(now);
}
pub fn mark_closed(&mut self, now: OffsetDateTime) {
self.status = StreamStatus::Stopped;
self.last_poll_at = Some(now);
self.closed_at = Some(now);
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum JobStatus {
Created,
Running,
Completed,
Failed,
Cancelled,
Expired,
}
impl JobStatus {
pub fn is_terminal(self) -> bool {
matches!(
self,
Self::Completed | Self::Failed | Self::Cancelled | Self::Expired
)
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AsyncJobHandle {
pub id: AsyncJobId,
pub workspace_id: WorkspaceId,
#[serde(skip_serializing_if = "Option::is_none")]
pub agent_id: Option<AgentId>,
pub operation_id: OperationId,
pub status: JobStatus,
pub progress: Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub expires_at: Option<OffsetDateTime>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub last_poll_at: Option<OffsetDateTime>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(with = "time::serde::rfc3339::option")]
pub finished_at: Option<OffsetDateTime>,
}
impl AsyncJobHandle {
pub fn is_finished(&self) -> bool {
self.status.is_terminal()
}
pub fn can_cancel(&self) -> bool {
matches!(self.status, JobStatus::Created | JobStatus::Running)
}
pub fn remaining_poll_delay_ms(&self, now: OffsetDateTime, poll_interval_ms: u64) -> u64 {
let Some(last_poll_at) = self.last_poll_at else {
return 0;
};
let elapsed_ms = (now - last_poll_at).whole_milliseconds().max(0) as u64;
poll_interval_ms.saturating_sub(elapsed_ms)
}
pub fn mark_finished(&mut self, now: OffsetDateTime, result: Value) {
self.status = JobStatus::Completed;
self.result = Some(result);
self.error = None;
self.updated_at = now;
self.finished_at = Some(now);
}
pub fn mark_failed(&mut self, now: OffsetDateTime, error: Value) {
self.status = JobStatus::Failed;
self.error = Some(error);
self.updated_at = now;
self.finished_at = Some(now);
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{AsyncJobHandle, JobStatus, StreamSession, StreamStatus};
use crate::{
ids::{AsyncJobId, OperationId, StreamSessionId, WorkspaceId},
protocol::Protocol,
streaming::ExecutionMode,
};
fn timestamp(value: &str) -> OffsetDateTime {
OffsetDateTime::parse(value, &Rfc3339).unwrap()
}
#[test]
fn stream_session_tracks_poll_and_close_transitions() {
let mut session = StreamSession {
id: StreamSessionId::new("stream_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
protocol: Protocol::Rest,
mode: ExecutionMode::Session,
status: StreamStatus::Running,
cursor: None,
state: json!({"cursor":"abc"}),
expires_at: timestamp("2026-04-06T12:05:00Z"),
last_poll_at: None,
created_at: timestamp("2026-04-06T12:00:00Z"),
closed_at: None,
};
assert!(session.can_poll(timestamp("2026-04-06T12:01:00Z")));
session.mark_polled(timestamp("2026-04-06T12:01:00Z"));
session.mark_closed(timestamp("2026-04-06T12:02:00Z"));
assert_eq!(session.status, StreamStatus::Stopped);
assert_eq!(session.closed_at, Some(timestamp("2026-04-06T12:02:00Z")));
assert!(!session.can_poll(timestamp("2026-04-06T12:03:00Z")));
}
#[test]
fn first_poll_has_no_required_delay() {
let session = StreamSession {
id: StreamSessionId::new("stream_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
protocol: Protocol::Rest,
mode: ExecutionMode::Session,
status: StreamStatus::Running,
cursor: None,
state: json!({"cursor":"abc"}),
expires_at: timestamp("2026-04-06T12:05:00Z"),
last_poll_at: None,
created_at: timestamp("2026-04-06T12:00:00Z"),
closed_at: None,
};
assert_eq!(
session.remaining_poll_delay_ms(timestamp("2026-04-06T12:00:00.100Z"), 250),
0
);
}
#[test]
fn rapid_repeat_poll_reports_remaining_delay() {
let session = StreamSession {
id: StreamSessionId::new("stream_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
protocol: Protocol::Rest,
mode: ExecutionMode::Session,
status: StreamStatus::Running,
cursor: None,
state: json!({"cursor":"abc"}),
expires_at: timestamp("2026-04-06T12:05:00Z"),
last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")),
created_at: timestamp("2026-04-06T12:00:00Z"),
closed_at: None,
};
assert_eq!(
session.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250),
150
);
}
#[test]
fn async_job_tracks_finish_and_failure() {
let mut job = AsyncJobHandle {
id: AsyncJobId::new("job_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
status: JobStatus::Running,
progress: json!({"percent": 60}),
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: None,
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
};
assert!(job.can_cancel());
job.mark_finished(timestamp("2026-04-06T12:01:00Z"), json!({"ok": true}));
assert!(job.is_finished());
assert_eq!(job.status, JobStatus::Completed);
let mut failed_job = job.clone();
failed_job.status = JobStatus::Running;
failed_job.result = None;
failed_job.finished_at = None;
failed_job.mark_failed(
timestamp("2026-04-06T12:02:00Z"),
json!({"message": "boom"}),
);
assert_eq!(failed_job.status, JobStatus::Failed);
assert_eq!(
failed_job.finished_at,
Some(timestamp("2026-04-06T12:02:00Z"))
);
}
#[test]
fn async_job_reports_remaining_poll_delay() {
let job = AsyncJobHandle {
id: AsyncJobId::new("job_01"),
workspace_id: WorkspaceId::new("ws_01"),
agent_id: None,
operation_id: OperationId::new("op_01"),
status: JobStatus::Running,
progress: json!({"percent": 60}),
result: None,
error: None,
expires_at: Some(timestamp("2026-04-06T12:05:00Z")),
last_poll_at: Some(timestamp("2026-04-06T12:01:00Z")),
created_at: timestamp("2026-04-06T12:00:00Z"),
updated_at: timestamp("2026-04-06T12:00:00Z"),
finished_at: None,
};
assert_eq!(
job.remaining_poll_delay_ms(timestamp("2026-04-06T12:01:00.100Z"), 250),
150
);
}
}
+364
View File
@@ -0,0 +1,364 @@
use serde::{Deserialize, Serialize};
use thiserror::Error;
use crate::protocol::Protocol;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExecutionMode {
Unary,
Window,
Session,
AsyncJob,
}
impl ExecutionMode {
pub fn is_stateful(self) -> bool {
matches!(self, Self::Session | Self::AsyncJob)
}
pub fn requires_tool_family(self) -> bool {
self.is_stateful()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AggregationMode {
RawItems,
SummaryOnly,
SummaryPlusSamples,
Stats,
LatestState,
}
impl AggregationMode {
pub fn needs_items(self) -> bool {
matches!(self, Self::RawItems | Self::SummaryPlusSamples)
}
pub fn needs_summary(self) -> bool {
!matches!(self, Self::RawItems)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TransportBehavior {
RequestResponse,
ServerStream,
StatefulSession,
DeferredResult,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct ToolFamilyConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub start_tool_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub poll_tool_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop_tool_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub status_tool_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub result_tool_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cancel_tool_name: Option<String>,
}
impl ToolFamilyConfig {
pub fn validate_for_mode(&self, mode: ExecutionMode) -> Result<(), StreamingConfigError> {
match mode {
ExecutionMode::Unary | ExecutionMode::Window => {
if self.start_tool_name.is_some()
|| self.poll_tool_name.is_some()
|| self.stop_tool_name.is_some()
|| self.status_tool_name.is_some()
|| self.result_tool_name.is_some()
|| self.cancel_tool_name.is_some()
{
return Err(StreamingConfigError::UnexpectedToolFamily(mode));
}
}
ExecutionMode::Session => {
if self.start_tool_name.is_none()
|| self.poll_tool_name.is_none()
|| self.stop_tool_name.is_none()
{
return Err(StreamingConfigError::MissingSessionToolNames);
}
}
ExecutionMode::AsyncJob => {
if self.start_tool_name.is_none()
|| self.status_tool_name.is_none()
|| self.result_tool_name.is_none()
|| self.cancel_tool_name.is_none()
{
return Err(StreamingConfigError::MissingAsyncJobToolNames);
}
}
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StreamingConfig {
pub mode: ExecutionMode,
pub transport_behavior: TransportBehavior,
#[serde(skip_serializing_if = "Option::is_none")]
pub window_duration_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub poll_interval_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_timeout_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub idle_timeout_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_session_lifetime_ms: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_items: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_bytes: Option<u32>,
pub aggregation_mode: AggregationMode,
#[serde(skip_serializing_if = "Option::is_none")]
pub summary_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub items_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cursor_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub status_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub done_path: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub redacted_paths: Vec<String>,
#[serde(default)]
pub truncate_item_fields: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_field_length: Option<u32>,
#[serde(default)]
pub drop_duplicates: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub sampling_rate: Option<f64>,
#[serde(default)]
pub tool_family: ToolFamilyConfig,
}
impl StreamingConfig {
pub fn validate_common(&self) -> Result<(), StreamingConfigError> {
self.tool_family.validate_for_mode(self.mode)?;
if self.max_items.is_some_and(|value| value == 0) {
return Err(StreamingConfigError::InvalidMaxItems);
}
if self.max_bytes.is_some_and(|value| value == 0) {
return Err(StreamingConfigError::InvalidMaxBytes);
}
if self.max_field_length.is_some_and(|value| value == 0) {
return Err(StreamingConfigError::InvalidMaxFieldLength);
}
if self
.sampling_rate
.is_some_and(|value| value <= 0.0 || value > 1.0)
{
return Err(StreamingConfigError::InvalidSamplingRate);
}
match self.mode {
ExecutionMode::Unary => {
if self.window_duration_ms.is_some()
|| self.poll_interval_ms.is_some()
|| self.idle_timeout_ms.is_some()
|| self.max_session_lifetime_ms.is_some()
{
return Err(StreamingConfigError::UnexpectedStatefulLimits(
ExecutionMode::Unary,
));
}
}
ExecutionMode::Window => {
if self.window_duration_ms.is_none() {
return Err(StreamingConfigError::MissingWindowDuration);
}
}
ExecutionMode::Session => {
if self.poll_interval_ms.is_none() || self.max_session_lifetime_ms.is_none() {
return Err(StreamingConfigError::MissingSessionLimits);
}
}
ExecutionMode::AsyncJob => {}
}
Ok(())
}
pub fn validate_for_protocol(&self, protocol: Protocol) -> Result<(), StreamingConfigError> {
self.validate_common()?;
if !protocol.supports_execution_mode(self.mode) {
return Err(StreamingConfigError::UnsupportedExecutionMode {
protocol,
mode: self.mode,
});
}
if !protocol.supports_transport_behavior(self.transport_behavior) {
return Err(StreamingConfigError::UnsupportedTransportBehavior {
protocol,
behavior: self.transport_behavior,
});
}
Ok(())
}
}
#[derive(Clone, Debug, Error, PartialEq)]
pub enum StreamingConfigError {
#[error("window mode requires window_duration_ms")]
MissingWindowDuration,
#[error("session mode requires poll_interval_ms and max_session_lifetime_ms")]
MissingSessionLimits,
#[error("max_items must be greater than zero")]
InvalidMaxItems,
#[error("max_bytes must be greater than zero")]
InvalidMaxBytes,
#[error("max_field_length must be greater than zero")]
InvalidMaxFieldLength,
#[error("sampling_rate must be in range (0, 1]")]
InvalidSamplingRate,
#[error("{0:?} mode cannot use session/window-only limits")]
UnexpectedStatefulLimits(ExecutionMode),
#[error("{mode:?} is not supported for protocol {protocol:?}")]
UnsupportedExecutionMode {
protocol: Protocol,
mode: ExecutionMode,
},
#[error("{behavior:?} is not supported for protocol {protocol:?}")]
UnsupportedTransportBehavior {
protocol: Protocol,
behavior: TransportBehavior,
},
#[error("session mode requires start, poll and stop tool names")]
MissingSessionToolNames,
#[error("async_job mode requires start, status, result and cancel tool names")]
MissingAsyncJobToolNames,
#[error("{0:?} mode cannot define tool-family names")]
UnexpectedToolFamily(ExecutionMode),
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
AggregationMode, ExecutionMode, StreamingConfig, StreamingConfigError, ToolFamilyConfig,
TransportBehavior,
};
use crate::protocol::Protocol;
fn session_config() -> StreamingConfig {
StreamingConfig {
mode: ExecutionMode::Session,
transport_behavior: TransportBehavior::StatefulSession,
window_duration_ms: None,
poll_interval_ms: Some(1_000),
upstream_timeout_ms: Some(5_000),
idle_timeout_ms: Some(15_000),
max_session_lifetime_ms: Some(60_000),
max_items: Some(100),
max_bytes: Some(65_536),
aggregation_mode: AggregationMode::SummaryPlusSamples,
summary_path: Some("$.summary".to_owned()),
items_path: Some("$.items".to_owned()),
cursor_path: Some("$.cursor".to_owned()),
status_path: Some("$.status".to_owned()),
done_path: Some("$.done".to_owned()),
redacted_paths: vec!["$.items[*].token".to_owned()],
truncate_item_fields: true,
max_field_length: Some(256),
drop_duplicates: true,
sampling_rate: Some(0.5),
tool_family: ToolFamilyConfig {
start_tool_name: Some("logs_start".to_owned()),
poll_tool_name: Some("logs_poll".to_owned()),
stop_tool_name: Some("logs_stop".to_owned()),
status_tool_name: None,
result_tool_name: None,
cancel_tool_name: None,
},
}
}
#[test]
fn streaming_config_roundtrips_through_json() {
let config = session_config();
let value = serde_json::to_value(&config).unwrap();
let decoded: StreamingConfig = serde_json::from_value(value.clone()).unwrap();
assert_eq!(decoded, config);
assert_eq!(value["mode"], json!("session"));
assert_eq!(value["tool_family"]["start_tool_name"], json!("logs_start"));
}
#[test]
fn unary_mode_rejects_session_specific_fields() {
let config = StreamingConfig {
mode: ExecutionMode::Unary,
transport_behavior: TransportBehavior::RequestResponse,
window_duration_ms: None,
poll_interval_ms: Some(1_000),
upstream_timeout_ms: None,
idle_timeout_ms: None,
max_session_lifetime_ms: None,
max_items: None,
max_bytes: None,
aggregation_mode: AggregationMode::SummaryOnly,
summary_path: None,
items_path: None,
cursor_path: None,
status_path: None,
done_path: None,
redacted_paths: Vec::new(),
truncate_item_fields: false,
max_field_length: None,
drop_duplicates: false,
sampling_rate: None,
tool_family: ToolFamilyConfig::default(),
};
assert_eq!(
config.validate_common(),
Err(StreamingConfigError::UnexpectedStatefulLimits(
ExecutionMode::Unary
))
);
}
#[test]
fn protocol_validation_rejects_graphql_session() {
let config = session_config();
assert_eq!(
config.validate_for_protocol(Protocol::Graphql),
Err(StreamingConfigError::UnsupportedExecutionMode {
protocol: Protocol::Graphql,
mode: ExecutionMode::Session,
})
);
}
#[test]
fn protocol_validation_accepts_rest_session() {
let config = session_config();
assert!(config.validate_for_protocol(Protocol::Rest).is_ok());
}
}
+75
View File
@@ -0,0 +1,75 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use crate::ids::WorkspaceId;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkspaceStatus {
Active,
Archived,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Workspace {
pub id: WorkspaceId,
pub slug: String,
pub display_name: String,
pub status: WorkspaceStatus,
pub settings: Value,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339")]
pub updated_at: OffsetDateTime,
}
#[cfg(test)]
mod tests {
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use super::{Workspace, WorkspaceStatus};
use crate::ids::WorkspaceId;
#[test]
fn workspace_serializes_timestamps_as_rfc3339() {
let workspace = Workspace {
id: WorkspaceId::new("ws_01"),
slug: "primary".to_owned(),
display_name: "Primary".to_owned(),
status: WorkspaceStatus::Active,
settings: json!({"region":"eu"}),
created_at: OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap(),
updated_at: OffsetDateTime::parse("2026-03-25T12:05:00Z", &Rfc3339).unwrap(),
};
let value = serde_json::to_value(&workspace).unwrap();
assert_eq!(value["created_at"], json!("2026-03-25T12:00:00Z"));
assert_eq!(value["updated_at"], json!("2026-03-25T12:05:00Z"));
}
#[test]
fn workspace_deserializes_timestamps_from_rfc3339() {
let workspace: Workspace = serde_json::from_value(json!({
"id": "ws_01",
"slug": "primary",
"display_name": "Primary",
"status": "active",
"settings": {"region":"eu"},
"created_at": "2026-03-25T12:00:00Z",
"updated_at": "2026-03-25T12:05:00Z"
}))
.unwrap();
assert_eq!(
workspace.created_at,
OffsetDateTime::parse("2026-03-25T12:00:00Z", &Rfc3339).unwrap()
);
assert_eq!(
workspace.updated_at,
OffsetDateTime::parse("2026-03-25T12:05:00Z", &Rfc3339).unwrap()
);
}
}
+15
View File
@@ -0,0 +1,15 @@
[package]
name = "crank-mapping"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
crank-core = { path = "../crank-core" }
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
[dev-dependencies]
serde_yaml.workspace = true
+20
View File
@@ -0,0 +1,20 @@
use thiserror::Error;
#[derive(Clone, Debug, PartialEq, Eq, Error)]
pub enum MappingError {
#[error("invalid JSONPath: {path}")]
InvalidJsonPath { path: String },
#[error("unsupported source root {root} for {context} mapping")]
UnsupportedSourceRoot { root: String, context: &'static str },
#[error("unsupported target root {root} for {context} mapping")]
UnsupportedTargetRoot { root: String, context: &'static str },
#[error("mixed mapping target contexts are not allowed")]
MixedTargetContext,
#[error("missing required value at {path}")]
MissingRequiredValue { path: String },
#[error("transform {transform} cannot be applied to {actual}")]
InvalidTransformInput {
transform: &'static str,
actual: &'static str,
},
}
+372
View File
@@ -0,0 +1,372 @@
use serde_json::{Map, Number, Value};
use crate::{
JsonPath, JsonPathSegment, MappingError, MappingSet, MappingTargetContext, TransformKind,
};
impl MappingSet {
pub fn apply(&self, source: &Value) -> Result<Value, MappingError> {
self.validate_paths()?;
let mut target = empty_target_context(self.target_context());
for rule in &self.rules {
if let Some(condition) = &rule.condition {
let condition_path = JsonPath::parse(&condition.source)?;
let condition_value = read_path(source, &condition_path);
if condition_value != Some(&condition.equals) {
continue;
}
}
let source_path = JsonPath::parse(&rule.source)?;
let target_path = JsonPath::parse(&rule.target)?;
let Some(source_value) = read_path(source, &source_path)
.cloned()
.or(rule.default_value.clone())
else {
if rule.required {
return Err(MappingError::MissingRequiredValue {
path: source_path.to_string_path(),
});
}
continue;
};
let value = apply_transform(
source_value,
rule.transform.as_ref().map(|value| &value.kind),
)?;
write_path(&mut target, &target_path, value)?;
}
Ok(target)
}
}
fn empty_target_context(context: MappingTargetContext) -> Value {
match context {
MappingTargetContext::Input => Value::Object(Map::from_iter([(
"request".to_owned(),
Value::Object(Map::from_iter([
("path".to_owned(), Value::Object(Map::new())),
("query".to_owned(), Value::Object(Map::new())),
("headers".to_owned(), Value::Object(Map::new())),
("body".to_owned(), Value::Object(Map::new())),
("variables".to_owned(), Value::Object(Map::new())),
("grpc".to_owned(), Value::Object(Map::new())),
])),
)])),
MappingTargetContext::Output => Value::Object(Map::from_iter([(
"output".to_owned(),
Value::Object(Map::new()),
)])),
_ => Value::Object(Map::new()),
}
}
fn read_path<'a>(value: &'a Value, path: &JsonPath) -> Option<&'a Value> {
path.read(value)
}
fn write_path(target: &mut Value, path: &JsonPath, value: Value) -> Result<(), MappingError> {
let mut segments = path
.root
.root_fields()
.iter()
.map(|field| JsonPathSegment::Field((*field).to_owned()))
.collect::<Vec<_>>();
segments.extend(path.segments.clone());
write_segments(target, &segments, value)
}
fn write_segments(
target: &mut Value,
segments: &[JsonPathSegment],
value: Value,
) -> Result<(), MappingError> {
if segments.is_empty() {
*target = value;
return Ok(());
}
match &segments[0] {
JsonPathSegment::Field(field) => {
if !target.is_object() {
*target = Value::Object(Map::new());
}
let Some(object) = target.as_object_mut() else {
return Err(MappingError::InvalidJsonPath {
path: field.clone(),
});
};
let entry = object.entry(field.clone()).or_insert(Value::Null);
write_segments(entry, &segments[1..], value)
}
JsonPathSegment::Index(index) => {
if !target.is_array() {
*target = Value::Array(Vec::new());
}
let Some(array) = target.as_array_mut() else {
return Err(MappingError::InvalidJsonPath {
path: index.to_string(),
});
};
while array.len() <= *index {
array.push(Value::Null);
}
write_segments(&mut array[*index], &segments[1..], value)
}
}
}
fn apply_transform(value: Value, transform: Option<&TransformKind>) -> Result<Value, MappingError> {
let Some(transform) = transform else {
return Ok(value);
};
match transform {
TransformKind::Identity => Ok(value),
TransformKind::ToString => Ok(Value::String(to_string_value(&value))),
TransformKind::ToNumber => to_number(value),
TransformKind::ToBoolean => to_boolean(value),
TransformKind::Join => join_values(value),
TransformKind::Split => split_value(value),
TransformKind::WrapArray => Ok(match value {
Value::Array(values) => Value::Array(values),
other => Value::Array(vec![other]),
}),
TransformKind::UnwrapSingleton => unwrap_singleton(value),
}
}
fn to_string_value(value: &Value) -> String {
match value {
Value::Null => "null".to_owned(),
Value::Bool(boolean) => boolean.to_string(),
Value::Number(number) => number.to_string(),
Value::String(text) => text.clone(),
Value::Array(_) | Value::Object(_) => serde_json::to_string(value).unwrap_or_default(),
}
}
fn to_number(value: Value) -> Result<Value, MappingError> {
let number = match value {
Value::Number(number) => return Ok(Value::Number(number)),
Value::String(text) => {
if let Ok(integer) = text.parse::<i64>() {
Number::from(integer)
} else {
Number::from_f64(text.parse::<f64>().map_err(|_| {
MappingError::InvalidTransformInput {
transform: "to_number",
actual: "string",
}
})?)
.ok_or(MappingError::InvalidTransformInput {
transform: "to_number",
actual: "string",
})?
}
}
Value::Bool(boolean) => Number::from(if boolean { 1 } else { 0 }),
other => {
return Err(MappingError::InvalidTransformInput {
transform: "to_number",
actual: type_name(&other),
});
}
};
Ok(Value::Number(number))
}
fn to_boolean(value: Value) -> Result<Value, MappingError> {
let boolean = match value {
Value::Bool(boolean) => boolean,
Value::Number(number) => {
number.as_i64().unwrap_or_default() != 0 || number.as_f64().unwrap_or_default() != 0.0
}
Value::String(text) => match text.to_ascii_lowercase().as_str() {
"true" | "1" | "yes" => true,
"false" | "0" | "no" => false,
_ => {
return Err(MappingError::InvalidTransformInput {
transform: "to_boolean",
actual: "string",
});
}
},
other => {
return Err(MappingError::InvalidTransformInput {
transform: "to_boolean",
actual: type_name(&other),
});
}
};
Ok(Value::Bool(boolean))
}
fn join_values(value: Value) -> Result<Value, MappingError> {
let Value::Array(values) = value else {
return Err(MappingError::InvalidTransformInput {
transform: "join",
actual: type_name(&value),
});
};
Ok(Value::String(
values
.iter()
.map(to_string_value)
.collect::<Vec<_>>()
.join(","),
))
}
fn split_value(value: Value) -> Result<Value, MappingError> {
let Value::String(text) = value else {
return Err(MappingError::InvalidTransformInput {
transform: "split",
actual: type_name(&value),
});
};
Ok(Value::Array(
text.split(',')
.map(|part| Value::String(part.trim().to_owned()))
.collect(),
))
}
fn unwrap_singleton(value: Value) -> Result<Value, MappingError> {
match value {
Value::Array(mut values) if values.len() == 1 => Ok(values.remove(0)),
Value::Array(_) => Err(MappingError::InvalidTransformInput {
transform: "unwrap_singleton",
actual: "array",
}),
other => Err(MappingError::InvalidTransformInput {
transform: "unwrap_singleton",
actual: type_name(&other),
}),
}
}
fn type_name(value: &Value) -> &'static str {
match value {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use crate::{MappingRule, MappingSet, Transform, TransformKind};
#[test]
fn applies_input_mapping_into_request_context() {
let mapping = MappingSet {
rules: vec![
MappingRule {
source: "$.mcp.email".to_owned(),
target: "$.request.body.contact.email".to_owned(),
required: true,
default_value: None,
transform: None,
condition: None,
notes: None,
},
MappingRule {
source: "$.mcp.tags".to_owned(),
target: "$.request.query.tags".to_owned(),
required: false,
default_value: None,
transform: Some(Transform {
kind: TransformKind::Join,
}),
condition: None,
notes: None,
},
],
};
let result = mapping
.apply(&json!({
"mcp": {
"email": "user@example.com",
"tags": ["warm", "priority"]
}
}))
.unwrap();
assert_eq!(
result["request"]["body"]["contact"]["email"],
"user@example.com"
);
assert_eq!(result["request"]["query"]["tags"], "warm,priority");
}
#[test]
fn applies_output_mapping_into_output_context() {
let mapping = MappingSet {
rules: vec![MappingRule {
source: "$.response.body.lead_id".to_owned(),
target: "$.output.id".to_owned(),
required: true,
default_value: None,
transform: Some(Transform {
kind: TransformKind::ToString,
}),
condition: None,
notes: None,
}],
};
let result = mapping
.apply(&json!({
"response": {
"body": {
"lead_id": 42
}
}
}))
.unwrap();
assert_eq!(result["output"]["id"], "42");
}
#[test]
fn uses_default_value_for_missing_optional_source() {
let mapping = MappingSet {
rules: vec![MappingRule {
source: "$.mcp.region".to_owned(),
target: "$.request.headers.X-Region".to_owned(),
required: false,
default_value: Some(json!("eu-central")),
transform: None,
condition: None,
notes: None,
}],
};
let result = mapping.apply(&json!({ "mcp": {} })).unwrap();
assert_eq!(result["request"]["headers"]["X-Region"], "eu-central");
}
}
+116
View File
@@ -0,0 +1,116 @@
use std::collections::BTreeMap;
use serde_json::Value;
use crate::{JsonPathRoot, MappingRule, MappingSet};
pub fn infer_mapping_from_samples(
source_sample: &Value,
source_root: JsonPathRoot,
target_sample: &Value,
target_root: JsonPathRoot,
) -> MappingSet {
let source_paths = collect_leaf_paths(source_sample, source_root);
let target_paths = collect_leaf_paths(target_sample, target_root);
let mut rules = Vec::new();
for (leaf_name, source_path) in source_paths {
let Some(target_path) = target_paths.get(&leaf_name) else {
continue;
};
rules.push(MappingRule {
source: source_path,
target: target_path.clone(),
required: false,
default_value: None,
transform: None,
condition: None,
notes: None,
});
}
MappingSet { rules }
}
fn collect_leaf_paths(sample: &Value, root: JsonPathRoot) -> BTreeMap<String, String> {
let mut collected = BTreeMap::new();
let mut segments = Vec::new();
collect_leaf_paths_recursive(sample, &mut segments, &mut collected);
collected
.into_iter()
.map(|(leaf, path_segments)| {
let mut path = format!("$.{}", root.as_str());
for segment in path_segments {
path.push('.');
path.push_str(&segment);
}
(leaf, path)
})
.collect()
}
fn collect_leaf_paths_recursive(
sample: &Value,
segments: &mut Vec<String>,
collected: &mut BTreeMap<String, Vec<String>>,
) {
match sample {
Value::Object(object) => {
for (field, value) in object {
segments.push(field.clone());
collect_leaf_paths_recursive(value, segments, collected);
segments.pop();
}
}
Value::Array(values) => {
if let Some(first) = values.first() {
collect_leaf_paths_recursive(first, segments, collected);
}
}
Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {
if let Some(leaf) = segments.last() {
collected
.entry(leaf.clone())
.or_insert_with(|| segments.clone());
}
}
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use crate::{JsonPathRoot, infer_mapping_from_samples};
#[test]
fn infers_rules_from_matching_leaf_names() {
let mapping = infer_mapping_from_samples(
&json!({
"lead": {
"email": "user@example.com",
"name": "Ada"
}
}),
JsonPathRoot::Mcp,
&json!({
"contact": {
"email": "user@example.com"
},
"name": "Ada"
}),
JsonPathRoot::RequestBody,
);
assert_eq!(mapping.rules.len(), 2);
assert_eq!(mapping.rules[0].source, "$.mcp.lead.email");
assert_eq!(mapping.rules[0].target, "$.request.body.contact.email");
assert_eq!(mapping.rules[1].source, "$.mcp.lead.name");
assert_eq!(mapping.rules[1].target, "$.request.body.name");
}
}
+233
View File
@@ -0,0 +1,233 @@
use serde_json::Value;
use crate::MappingError;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum JsonPathRoot {
Mcp,
RequestPath,
RequestQuery,
RequestHeaders,
RequestBody,
RequestVariables,
RequestGrpc,
ResponseBody,
ResponseData,
ResponseGrpc,
Output,
}
impl JsonPathRoot {
pub fn as_str(&self) -> &'static str {
match self {
Self::Mcp => "mcp",
Self::RequestPath => "request.path",
Self::RequestQuery => "request.query",
Self::RequestHeaders => "request.headers",
Self::RequestBody => "request.body",
Self::RequestVariables => "request.variables",
Self::RequestGrpc => "request.grpc",
Self::ResponseBody => "response.body",
Self::ResponseData => "response.data",
Self::ResponseGrpc => "response.grpc",
Self::Output => "output",
}
}
pub fn root_fields(&self) -> &'static [&'static str] {
match self {
Self::Mcp => &["mcp"],
Self::RequestPath => &["request", "path"],
Self::RequestQuery => &["request", "query"],
Self::RequestHeaders => &["request", "headers"],
Self::RequestBody => &["request", "body"],
Self::RequestVariables => &["request", "variables"],
Self::RequestGrpc => &["request", "grpc"],
Self::ResponseBody => &["response", "body"],
Self::ResponseData => &["response", "data"],
Self::ResponseGrpc => &["response", "grpc"],
Self::Output => &["output"],
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum JsonPathSegment {
Field(String),
Index(usize),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct JsonPath {
pub root: JsonPathRoot,
pub segments: Vec<JsonPathSegment>,
}
impl JsonPath {
pub fn parse(path: &str) -> Result<Self, MappingError> {
let body = path
.strip_prefix("$.")
.ok_or_else(|| MappingError::InvalidJsonPath {
path: path.to_owned(),
})?;
let (root, rest) = match_root(body).ok_or_else(|| MappingError::InvalidJsonPath {
path: path.to_owned(),
})?;
if !(rest.is_empty() || rest.starts_with('.') || rest.starts_with('[')) {
return Err(MappingError::InvalidJsonPath {
path: path.to_owned(),
});
}
let segments = parse_segments(rest, path)?;
Ok(Self { root, segments })
}
pub fn to_string_path(&self) -> String {
let mut path = format!("$.{}", self.root.as_str());
for segment in &self.segments {
match segment {
JsonPathSegment::Field(field) => {
path.push('.');
path.push_str(field);
}
JsonPathSegment::Index(index) => {
path.push('[');
path.push_str(&index.to_string());
path.push(']');
}
}
}
path
}
pub fn read<'a>(&self, value: &'a Value) -> Option<&'a Value> {
let mut current = value;
for field in self.root.root_fields() {
current = current.get(*field)?;
}
for segment in &self.segments {
current = match segment {
JsonPathSegment::Field(field) => current.get(field)?,
JsonPathSegment::Index(index) => current.get(*index)?,
};
}
Some(current)
}
}
fn match_root(input: &str) -> Option<(JsonPathRoot, &str)> {
const ROOTS: [(&str, JsonPathRoot); 11] = [
("request.variables", JsonPathRoot::RequestVariables),
("request.headers", JsonPathRoot::RequestHeaders),
("response.body", JsonPathRoot::ResponseBody),
("response.data", JsonPathRoot::ResponseData),
("response.grpc", JsonPathRoot::ResponseGrpc),
("request.path", JsonPathRoot::RequestPath),
("request.query", JsonPathRoot::RequestQuery),
("request.body", JsonPathRoot::RequestBody),
("request.grpc", JsonPathRoot::RequestGrpc),
("output", JsonPathRoot::Output),
("mcp", JsonPathRoot::Mcp),
];
for (prefix, root) in ROOTS {
if let Some(rest) = input.strip_prefix(prefix) {
return Some((root, rest));
}
}
None
}
fn parse_segments(mut input: &str, original: &str) -> Result<Vec<JsonPathSegment>, MappingError> {
let mut segments = Vec::new();
while !input.is_empty() {
if let Some(rest) = input.strip_prefix('.') {
let end = rest.find(['.', '[']).unwrap_or(rest.len());
let field = &rest[..end];
if field.is_empty() || !is_valid_field(field) {
return Err(MappingError::InvalidJsonPath {
path: original.to_owned(),
});
}
segments.push(JsonPathSegment::Field(field.to_owned()));
input = &rest[end..];
continue;
}
if let Some(rest) = input.strip_prefix('[') {
let Some(end) = rest.find(']') else {
return Err(MappingError::InvalidJsonPath {
path: original.to_owned(),
});
};
let index =
rest[..end]
.parse::<usize>()
.map_err(|_| MappingError::InvalidJsonPath {
path: original.to_owned(),
})?;
segments.push(JsonPathSegment::Index(index));
input = &rest[end + 1..];
continue;
}
return Err(MappingError::InvalidJsonPath {
path: original.to_owned(),
});
}
Ok(segments)
}
fn is_valid_field(value: &str) -> bool {
value
.chars()
.all(|char| char.is_ascii_alphanumeric() || char == '_' || char == '-')
}
#[cfg(test)]
mod tests {
use super::{JsonPath, JsonPathRoot, JsonPathSegment};
#[test]
fn parses_nested_jsonpath_with_array_index() {
let path = JsonPath::parse("$.request.body.contacts[0].email").unwrap();
assert_eq!(path.root, JsonPathRoot::RequestBody);
assert_eq!(
path.segments,
vec![
JsonPathSegment::Field("contacts".to_owned()),
JsonPathSegment::Index(0),
JsonPathSegment::Field("email".to_owned())
]
);
}
#[test]
fn rejects_unknown_root() {
let error = JsonPath::parse("$.unknown.field").unwrap_err();
assert_eq!(
error,
crate::MappingError::InvalidJsonPath {
path: "$.unknown.field".to_owned(),
}
);
}
}
+12
View File
@@ -0,0 +1,12 @@
mod error;
mod execute;
mod infer;
mod jsonpath;
mod model;
pub use error::MappingError;
pub use infer::infer_mapping_from_samples;
pub use jsonpath::{JsonPath, JsonPathRoot, JsonPathSegment};
pub use model::{
MappingCondition, MappingRule, MappingSet, MappingTargetContext, Transform, TransformKind,
};
+311
View File
@@ -0,0 +1,311 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{JsonPath, JsonPathRoot, MappingError};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MappingTargetContext {
Unknown,
Input,
Output,
Mixed,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct MappingCondition {
pub source: String,
pub equals: Value,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TransformKind {
Identity,
ToString,
ToNumber,
ToBoolean,
Join,
Split,
WrapArray,
UnwrapSingleton,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Transform {
pub kind: TransformKind,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct MappingRule {
pub source: String,
pub target: String,
#[serde(default)]
pub required: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub default_value: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub transform: Option<Transform>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub condition: Option<MappingCondition>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub notes: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Default)]
pub struct MappingSet {
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub rules: Vec<MappingRule>,
}
impl MappingSet {
pub fn is_empty(&self) -> bool {
self.rules.is_empty()
}
pub fn len(&self) -> usize {
self.rules.len()
}
pub fn target_context(&self) -> MappingTargetContext {
let mut detected = MappingTargetContext::Unknown;
for rule in &self.rules {
let Ok(path) = JsonPath::parse(&rule.target) else {
return MappingTargetContext::Mixed;
};
let current = match target_root_context(path.root) {
Some(context) => context,
None => return MappingTargetContext::Mixed,
};
detected = match (detected, current) {
(MappingTargetContext::Unknown, context) => context,
(existing, current) if existing == current => existing,
_ => MappingTargetContext::Mixed,
};
}
detected
}
pub fn validate_paths(&self) -> Result<(), MappingError> {
let context = self.target_context();
if context == MappingTargetContext::Mixed {
return Err(MappingError::MixedTargetContext);
}
for rule in &self.rules {
let source = JsonPath::parse(&rule.source)?;
let target = JsonPath::parse(&rule.target)?;
validate_source_root(source.root, context)?;
validate_target_root(target.root, context)?;
if let Some(condition) = &rule.condition {
let condition_path = JsonPath::parse(&condition.source)?;
validate_source_root(condition_path.root, context)?;
}
}
Ok(())
}
}
fn target_root_context(root: JsonPathRoot) -> Option<MappingTargetContext> {
match root {
JsonPathRoot::RequestPath
| JsonPathRoot::RequestQuery
| JsonPathRoot::RequestHeaders
| JsonPathRoot::RequestBody
| JsonPathRoot::RequestVariables
| JsonPathRoot::RequestGrpc => Some(MappingTargetContext::Input),
JsonPathRoot::Output => Some(MappingTargetContext::Output),
_ => None,
}
}
fn validate_source_root(
root: JsonPathRoot,
context: MappingTargetContext,
) -> Result<(), MappingError> {
let valid = match context {
MappingTargetContext::Unknown => true,
MappingTargetContext::Input => root == JsonPathRoot::Mcp,
MappingTargetContext::Output => {
matches!(
root,
JsonPathRoot::ResponseBody
| JsonPathRoot::ResponseData
| JsonPathRoot::ResponseGrpc
)
}
MappingTargetContext::Mixed => false,
};
if valid {
return Ok(());
}
Err(MappingError::UnsupportedSourceRoot {
root: root.as_str().to_owned(),
context: context_name(context),
})
}
fn validate_target_root(
root: JsonPathRoot,
context: MappingTargetContext,
) -> Result<(), MappingError> {
let valid = match context {
MappingTargetContext::Unknown => true,
MappingTargetContext::Input => matches!(
root,
JsonPathRoot::RequestPath
| JsonPathRoot::RequestQuery
| JsonPathRoot::RequestHeaders
| JsonPathRoot::RequestBody
| JsonPathRoot::RequestVariables
| JsonPathRoot::RequestGrpc
),
MappingTargetContext::Output => root == JsonPathRoot::Output,
MappingTargetContext::Mixed => false,
};
if valid {
return Ok(());
}
Err(MappingError::UnsupportedTargetRoot {
root: root.as_str().to_owned(),
context: context_name(context),
})
}
fn context_name(context: MappingTargetContext) -> &'static str {
match context {
MappingTargetContext::Unknown => "unknown",
MappingTargetContext::Input => "input",
MappingTargetContext::Output => "output",
MappingTargetContext::Mixed => "mixed",
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{MappingRule, MappingSet, MappingTargetContext, Transform, TransformKind};
#[test]
fn mapping_set_reports_non_empty_rules() {
let mapping = MappingSet {
rules: vec![MappingRule {
source: "$.mcp.email".to_owned(),
target: "$.request.body.contact.email".to_owned(),
required: true,
default_value: None,
transform: Some(Transform {
kind: TransformKind::Identity,
}),
condition: None,
notes: None,
}],
};
assert!(!mapping.is_empty());
assert_eq!(mapping.len(), 1);
}
#[test]
fn mapping_rule_serializes_jsonpath_fields() {
let mapping = MappingSet {
rules: vec![MappingRule {
source: "$.response.body.id".to_owned(),
target: "$.output.id".to_owned(),
required: false,
default_value: Some(json!("lead_123")),
transform: None,
condition: None,
notes: Some("map identifier".to_owned()),
}],
};
let value = serde_json::to_value(mapping).unwrap();
assert_eq!(value["rules"][0]["source"], "$.response.body.id");
assert_eq!(value["rules"][0]["target"], "$.output.id");
assert_eq!(value["rules"][0]["default_value"], "lead_123");
}
#[test]
fn mapping_set_roundtrips_through_yaml() {
let mapping = MappingSet {
rules: vec![MappingRule {
source: "$.mcp.tags".to_owned(),
target: "$.request.body.tags".to_owned(),
required: false,
default_value: None,
transform: Some(Transform {
kind: TransformKind::WrapArray,
}),
condition: None,
notes: Some("normalize tags".to_owned()),
}],
};
let yaml = serde_yaml::to_string(&mapping).unwrap();
let restored: MappingSet = serde_yaml::from_str(&yaml).unwrap();
assert!(yaml.contains("source: $.mcp.tags"));
assert_eq!(restored, mapping);
}
#[test]
fn detects_input_context_from_target_roots() {
let mapping = MappingSet {
rules: vec![MappingRule {
source: "$.mcp.email".to_owned(),
target: "$.request.body.email".to_owned(),
required: true,
default_value: None,
transform: None,
condition: None,
notes: None,
}],
};
assert_eq!(mapping.target_context(), MappingTargetContext::Input);
}
#[test]
fn rejects_mixed_target_contexts() {
let mapping = MappingSet {
rules: vec![
MappingRule {
source: "$.mcp.email".to_owned(),
target: "$.request.body.email".to_owned(),
required: true,
default_value: None,
transform: None,
condition: None,
notes: None,
},
MappingRule {
source: "$.response.body.id".to_owned(),
target: "$.output.id".to_owned(),
required: true,
default_value: None,
transform: None,
condition: None,
notes: None,
},
],
};
let error = mapping.validate_paths().unwrap_err();
assert_eq!(error, crate::MappingError::MixedTargetContext);
}
}
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "crank-registry"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
crank-core = { path = "../crank-core" }
crank-mapping = { path = "../crank-mapping" }
crank-schema = { path = "../crank-schema" }
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
thiserror.workspace = true
time.workspace = true
uuid.workspace = true
[dev-dependencies]
tokio.workspace = true
+89
View File
@@ -0,0 +1,89 @@
use thiserror::Error;
#[derive(Debug, Error)]
pub enum RegistryError {
#[error(transparent)]
Storage(#[from] sqlx::Error),
#[error(transparent)]
Serialization(#[from] serde_json::Error),
#[error("workspace {workspace_id} was not found")]
WorkspaceNotFound { workspace_id: String },
#[error("workspace with slug {slug} already exists")]
WorkspaceSlugAlreadyExists { slug: String },
#[error("user {user_id} was not found")]
UserNotFound { user_id: String },
#[error("user with email {email} already exists")]
UserEmailAlreadyExists { email: String },
#[error("membership for user {user_id} in workspace {workspace_id} was not found")]
MembershipNotFound {
workspace_id: String,
user_id: String,
},
#[error("invitation {invitation_id} was not found")]
InvitationNotFound { invitation_id: String },
#[error("platform api key {key_id} was not found")]
PlatformApiKeyNotFound { key_id: String },
#[error("secret {secret_id} was not found")]
SecretNotFound { secret_id: String },
#[error("stream session {session_id} was not found")]
StreamSessionNotFound { session_id: String },
#[error("async job {job_id} was not found")]
AsyncJobNotFound { job_id: String },
#[error("invalid stream session transition for {session_id}: {from} -> {to}")]
InvalidStreamSessionTransition {
session_id: String,
from: String,
to: String,
},
#[error("invalid async job transition for {job_id}: {from} -> {to}")]
InvalidAsyncJobTransition {
job_id: String,
from: String,
to: String,
},
#[error("secret with name {name} already exists in workspace {workspace_id}")]
SecretNameAlreadyExists { workspace_id: String, name: String },
#[error("secret {secret_id} is referenced by auth profile {auth_profile_id}")]
SecretReferencedByAuthProfile {
secret_id: String,
auth_profile_id: String,
},
#[error("invocation log {log_id} was not found")]
InvocationLogNotFound { log_id: String },
#[error("agent {agent_id} was not found")]
AgentNotFound { agent_id: String },
#[error("published agent {workspace_slug}/{agent_slug} was not found")]
PublishedAgentNotFound {
workspace_slug: String,
agent_slug: String,
},
#[error("operation {operation_id} already exists")]
OperationAlreadyExists { operation_id: String },
#[error("operation {operation_id} was not found")]
OperationNotFound { operation_id: String },
#[error("operation version {version} for {operation_id} was not found")]
OperationVersionNotFound { operation_id: String, version: u32 },
#[error("operation {operation_id} cannot be deleted while it is bound to a published agent")]
OperationHasPublishedAgentBindings { operation_id: String },
#[error("operation {operation_id} must start with version 1, got {version}")]
InvalidInitialVersion { operation_id: String, version: u32 },
#[error("operation {operation_id} expected next version {expected}, got {actual}")]
InvalidVersionSequence {
operation_id: String,
expected: u32,
actual: u32,
},
#[error("operation {operation_id} changed immutable field {field}")]
ImmutableOperationFieldChanged {
operation_id: String,
field: &'static str,
},
#[error("auth profile {auth_profile_id} was not found")]
AuthProfileNotFound { auth_profile_id: String },
#[error("yaml import job {job_id} was not found")]
YamlImportJobNotFound { job_id: String },
#[error("unsupported enum representation for field {field}")]
InvalidEnumRepresentation { field: &'static str },
#[error("invalid numeric value for field {field}: {value}")]
InvalidNumericValue { field: &'static str, value: i64 },
}
+77
View File
@@ -0,0 +1,77 @@
use std::sync::Arc;
use sqlx::{PgPool, query};
use crate::RegistryError;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ExtensionMigration {
pub version: u32,
pub sql: &'static str,
}
pub trait RegistryExtension: Send + Sync {
fn name(&self) -> &str;
fn migrations(&self) -> &[ExtensionMigration];
}
pub async fn apply_extension_migrations(
pool: &PgPool,
extensions: &[Arc<dyn RegistryExtension>],
) -> Result<(), RegistryError> {
query(
"create table if not exists __crank_ext_migrations (
extension_name text not null,
version integer not null,
applied_at timestamptz not null default now(),
primary key (extension_name, version)
)",
)
.execute(pool)
.await?;
for extension in extensions {
for migration in extension.migrations() {
let already_applied = query(
"select 1
from __crank_ext_migrations
where extension_name = $1 and version = $2",
)
.bind(extension.name())
.bind(i32::try_from(migration.version).map_err(|_| {
RegistryError::InvalidNumericValue {
field: "extension_migration.version",
value: migration.version as i64,
}
})?)
.fetch_optional(pool)
.await?
.is_some();
if already_applied {
continue;
}
let mut tx = pool.begin().await?;
query(migration.sql).execute(&mut *tx).await?;
query(
"insert into __crank_ext_migrations (extension_name, version)
values ($1, $2)",
)
.bind(extension.name())
.bind(i32::try_from(migration.version).map_err(|_| {
RegistryError::InvalidNumericValue {
field: "extension_migration.version",
value: migration.version as i64,
}
})?)
.execute(&mut *tx)
.await?;
tx.commit().await?;
}
}
Ok(())
}
+26
View File
@@ -0,0 +1,26 @@
mod error;
mod ext;
mod migrations;
mod model;
mod postgres;
pub use error::RegistryError;
pub use ext::{ExtensionMigration, RegistryExtension, apply_extension_migrations};
pub use model::{
AgentSummary, AgentVersionRecord, AsyncJobFilter, AsyncJobRecord, AuthUserRecord,
CreateAgentRequest, CreateAsyncJobRequest, CreateInvitationRequest, CreateInvocationLogRequest,
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateStreamSessionRequest,
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind,
DescriptorMetadata, InvitationRecord, InvocationLogRecord, ListInvocationLogsQuery,
MembershipRecord, OperationAgentRef, OperationSampleMetadata, OperationSummary,
OperationUsageSummary, OperationVersionRecord, Page, PlatformApiKeyRecord, PublishAgentRequest,
PublishRequest, PublishedAgentTool, RegistryOperation, RotateSecretRequest, SampleKind,
SaveAgentBindingsRequest, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
SaveSampleMetadataRequest, SecretRecord, SecretVersionRecord, SessionRecord,
StreamSessionFilter, StreamSessionRecord, UpdateAsyncJobStatusRequest,
UpdateStreamSessionStateRequest, UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket,
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
WorkspaceMembershipRecord, WorkspaceRecord, YamlImportJob, YamlImportJobCompletion,
YamlImportJobId, YamlImportJobStatus,
};
pub use postgres::{PostgresPoolConfig, PostgresPoolConfigError, PostgresRegistry};
+599
View File
@@ -0,0 +1,599 @@
use sqlx::{PgPool, query};
pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
query(
"create table if not exists workspaces (
id text primary key,
slug text not null unique,
display_name text not null,
status text not null,
settings_json jsonb not null default '{}'::jsonb,
created_at timestamptz not null,
updated_at timestamptz not null
)",
)
.execute(pool)
.await?;
query(
"create table if not exists users (
id text primary key,
email text not null unique,
display_name text not null,
password_hash text null,
status text not null,
created_at timestamptz not null
)",
)
.execute(pool)
.await?;
query("alter table users add column if not exists password_hash text null")
.execute(pool)
.await?;
query(
"insert into users (
id,
email,
display_name,
status,
created_at
) values (
'user_default_owner',
'owner@crank.local',
'Workspace Owner',
'active',
now()
)
on conflict (id) do nothing",
)
.execute(pool)
.await?;
query(
"create table if not exists memberships (
workspace_id text not null references workspaces(id) on delete cascade,
user_id text not null references users(id) on delete cascade,
role text not null,
created_at timestamptz not null,
primary key (workspace_id, user_id)
)",
)
.execute(pool)
.await?;
query(
"create table if not exists user_sessions (
id text primary key,
user_id text not null references users(id) on delete cascade,
current_workspace_id text null references workspaces(id) on delete set null,
secret_hash text not null,
status text not null,
expires_at timestamptz not null,
last_seen_at timestamptz null,
created_at timestamptz not null
)",
)
.execute(pool)
.await?;
query(
"alter table user_sessions
add column if not exists current_workspace_id text null references workspaces(id) on delete set null",
)
.execute(pool)
.await?;
query(
"insert into workspaces (
id,
slug,
display_name,
status,
settings_json,
created_at,
updated_at
) values (
'ws_default',
'default',
'Default Workspace',
'active',
'{}'::jsonb,
now(),
now()
)
on conflict (id) do nothing",
)
.execute(pool)
.await?;
query(
"insert into memberships (
workspace_id,
user_id,
role,
created_at
) values (
'ws_default',
'user_default_owner',
'owner',
now()
)
on conflict (workspace_id, user_id) do nothing",
)
.execute(pool)
.await?;
query(
"create table if not exists invitation_tokens (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
email text not null,
role text not null,
status text not null,
token_hash text not null,
expires_at timestamptz not null,
created_at timestamptz not null
)",
)
.execute(pool)
.await?;
query(
"create table if not exists platform_api_keys (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
agent_id text null,
name text not null,
prefix text not null,
secret_hash text not null,
scopes_json jsonb not null,
status text not null,
created_at timestamptz not null,
last_used_at timestamptz null,
revoked_at timestamptz null
)",
)
.execute(pool)
.await?;
query(
"create unique index if not exists platform_api_keys_workspace_name_idx on platform_api_keys(workspace_id, name)",
)
.execute(pool)
.await?;
query("alter table platform_api_keys add column if not exists agent_id text null")
.execute(pool)
.await?;
query(
"insert into workspaces (
id,
slug,
display_name,
status,
settings_json,
created_at,
updated_at
) values (
'ws_default',
'default',
'Default Workspace',
'active',
'{}'::jsonb,
now(),
now()
)
on conflict (id) do nothing",
)
.execute(pool)
.await?;
query(
"create table if not exists operations (
id text primary key,
workspace_id text null references workspaces(id) on delete cascade,
name text not null,
display_name text not null,
category text not null default 'general',
protocol text not null,
security_level text not null default 'standard',
status text not null,
current_draft_version integer not null default 1,
latest_published_version integer null,
created_at timestamptz not null,
updated_at timestamptz not null,
published_at timestamptz null
)",
)
.execute(pool)
.await?;
query("alter table operations add column if not exists workspace_id text null references workspaces(id) on delete cascade")
.execute(pool)
.await?;
query(
"alter table operations add column if not exists category text not null default 'general'",
)
.execute(pool)
.await?;
query(
"alter table operations add column if not exists security_level text not null default 'standard'",
)
.execute(pool)
.await?;
query("update operations set workspace_id = 'ws_default' where workspace_id is null")
.execute(pool)
.await?;
query("alter table operations alter column workspace_id set not null")
.execute(pool)
.await?;
query("alter table operations drop constraint if exists operations_name_key")
.execute(pool)
.await?;
query(
"create unique index if not exists operations_workspace_name_idx on operations(workspace_id, name)",
)
.execute(pool)
.await?;
query(
"create table if not exists operation_versions (
operation_id text not null references operations(id) on delete cascade,
version integer not null,
status text not null,
target_json jsonb not null,
input_schema_json jsonb not null,
output_schema_json jsonb not null,
input_mapping_json jsonb not null,
output_mapping_json jsonb not null,
execution_config_json jsonb not null,
tool_description_json jsonb not null,
samples_json jsonb null,
generated_draft_json jsonb null,
config_export_json jsonb null,
change_note text null,
created_at timestamptz not null,
created_by text null,
primary key (operation_id, version)
)",
)
.execute(pool)
.await?;
query(
"create table if not exists published_operations (
operation_id text primary key references operations(id) on delete cascade,
version integer not null,
published_at timestamptz not null,
published_by text null,
foreign key (operation_id, version) references operation_versions(operation_id, version) on delete cascade
)",
)
.execute(pool)
.await?;
query(
"create table if not exists operation_samples (
id text primary key,
operation_id text not null references operations(id) on delete cascade,
version integer not null,
sample_kind text not null,
storage_ref text not null,
content_type text not null,
file_name text null,
created_at timestamptz not null,
foreign key (operation_id, version) references operation_versions(operation_id, version) on delete cascade
)",
)
.execute(pool)
.await?;
query(
"create table if not exists descriptors (
id text primary key,
operation_id text null references operations(id) on delete cascade,
version integer null,
descriptor_kind text not null,
storage_ref text not null,
source_name text null,
package_index_json jsonb null,
created_at timestamptz not null,
foreign key (operation_id, version) references operation_versions(operation_id, version) on delete cascade
)",
)
.execute(pool)
.await?;
query(
"create table if not exists auth_profiles (
id text primary key,
workspace_id text null references workspaces(id) on delete cascade,
name text not null,
kind text not null,
config_json jsonb not null,
created_at timestamptz not null,
updated_at timestamptz not null
)",
)
.execute(pool)
.await?;
query("alter table auth_profiles add column if not exists workspace_id text null references workspaces(id) on delete cascade")
.execute(pool)
.await?;
query("update auth_profiles set workspace_id = 'ws_default' where workspace_id is null")
.execute(pool)
.await?;
query("alter table auth_profiles alter column workspace_id set not null")
.execute(pool)
.await?;
query("alter table auth_profiles drop constraint if exists auth_profiles_name_key")
.execute(pool)
.await?;
query(
"create unique index if not exists auth_profiles_workspace_name_idx on auth_profiles(workspace_id, name)",
)
.execute(pool)
.await?;
query(
"create table if not exists secrets (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
name text not null,
kind text not null,
status text not null,
current_version integer not null,
last_used_at timestamptz null,
created_at timestamptz not null,
updated_at timestamptz not null
)",
)
.execute(pool)
.await?;
query(
"create unique index if not exists secrets_workspace_name_idx on secrets(workspace_id, name)",
)
.execute(pool)
.await?;
query(
"create table if not exists secret_versions (
secret_id text not null references secrets(id) on delete cascade,
version integer not null,
ciphertext text not null,
key_version text not null,
created_at timestamptz not null,
created_by text null references users(id) on delete set null,
primary key (secret_id, version)
)",
)
.execute(pool)
.await?;
query(
"create table if not exists yaml_import_jobs (
id text primary key,
source_sample_id text null references operation_samples(id) on delete set null,
status text not null,
format_version text not null,
mode text not null,
result_operation_id text null references operations(id) on delete set null,
result_version integer null,
error_text text null,
created_at timestamptz not null,
finished_at timestamptz null
)",
)
.execute(pool)
.await?;
query(
"create table if not exists agents (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
slug text not null,
display_name text not null,
description text not null,
status text not null,
current_draft_version integer not null default 1,
latest_published_version integer null,
created_at timestamptz not null,
updated_at timestamptz not null,
published_at timestamptz null
)",
)
.execute(pool)
.await?;
query(
"create unique index if not exists agents_workspace_slug_idx on agents(workspace_id, slug)",
)
.execute(pool)
.await?;
query(
"create table if not exists agent_versions (
agent_id text not null references agents(id) on delete cascade,
version integer not null,
status text not null,
instructions_json jsonb not null,
tool_selection_policy_json jsonb not null,
created_at timestamptz not null,
primary key (agent_id, version)
)",
)
.execute(pool)
.await?;
query(
"create table if not exists agent_operation_bindings (
agent_id text not null references agents(id) on delete cascade,
agent_version integer not null,
operation_id text not null references operations(id) on delete cascade,
operation_version integer not null,
tool_name text not null,
tool_title text not null,
tool_description_override text null,
enabled boolean not null default true,
foreign key (agent_id, agent_version) references agent_versions(agent_id, version) on delete cascade,
foreign key (operation_id, operation_version) references operation_versions(operation_id, version) on delete cascade
)",
)
.execute(pool)
.await?;
query(
"create unique index if not exists agent_bindings_tool_name_idx on agent_operation_bindings(agent_id, agent_version, tool_name)",
)
.execute(pool)
.await?;
query(
"create table if not exists published_agents (
agent_id text primary key references agents(id) on delete cascade,
version integer not null,
published_at timestamptz not null,
published_by text null,
foreign key (agent_id, version) references agent_versions(agent_id, version) on delete cascade
)",
)
.execute(pool)
.await?;
query(
"create table if not exists invocation_logs (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
agent_id text null references agents(id) on delete set null,
operation_id text not null references operations(id) on delete cascade,
source text not null,
level text not null,
status text not null,
tool_name text not null,
message text not null,
request_id text null,
status_code integer null,
duration_ms bigint not null,
error_kind text null,
request_preview_json jsonb not null,
response_preview_json jsonb not null,
created_at timestamptz not null
)",
)
.execute(pool)
.await?;
query(
"create index if not exists invocation_logs_workspace_created_idx on invocation_logs(workspace_id, created_at desc)",
)
.execute(pool)
.await?;
query(
"create index if not exists invocation_logs_workspace_operation_created_idx on invocation_logs(workspace_id, operation_id, created_at desc)",
)
.execute(pool)
.await?;
query(
"create index if not exists invocation_logs_workspace_agent_created_idx on invocation_logs(workspace_id, agent_id, created_at desc)",
)
.execute(pool)
.await?;
query(
"create table if not exists usage_rollups (
workspace_id text not null references workspaces(id) on delete cascade,
agent_id text null references agents(id) on delete cascade,
operation_id text null references operations(id) on delete cascade,
period text not null,
calls_total bigint not null,
calls_ok bigint not null,
calls_error bigint not null,
p50_ms bigint not null,
p95_ms bigint not null,
p99_ms bigint not null,
updated_at timestamptz not null
)",
)
.execute(pool)
.await?;
query(
"create table if not exists stream_sessions (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
agent_id text null references agents(id) on delete set null,
operation_id text not null references operations(id) on delete cascade,
protocol text not null,
mode text not null,
status text not null,
cursor_json jsonb null,
state_json jsonb not null,
expires_at timestamptz not null,
last_poll_at timestamptz null,
created_at timestamptz not null,
closed_at timestamptz null
)",
)
.execute(pool)
.await?;
query(
"create index if not exists stream_sessions_workspace_status_idx
on stream_sessions(workspace_id, status, created_at desc)",
)
.execute(pool)
.await?;
query(
"create index if not exists stream_sessions_expires_at_idx
on stream_sessions(expires_at)",
)
.execute(pool)
.await?;
query(
"create table if not exists async_jobs (
id text primary key,
workspace_id text not null references workspaces(id) on delete cascade,
agent_id text null references agents(id) on delete set null,
operation_id text not null references operations(id) on delete cascade,
status text not null,
progress_json jsonb not null,
result_json jsonb null,
error_json jsonb null,
expires_at timestamptz null,
last_poll_at timestamptz null,
created_at timestamptz not null,
updated_at timestamptz not null,
finished_at timestamptz null
)",
)
.execute(pool)
.await?;
query("alter table async_jobs add column if not exists last_poll_at timestamptz null")
.execute(pool)
.await?;
query(
"create index if not exists async_jobs_workspace_status_idx
on async_jobs(workspace_id, status, updated_at desc)",
)
.execute(pool)
.await?;
query(
"create index if not exists async_jobs_expires_at_idx
on async_jobs(expires_at)",
)
.execute(pool)
.await?;
Ok(())
}
+525
View File
@@ -0,0 +1,525 @@
use crank_core::{
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
AuthProfile, DescriptorId, ExecutionMode, ExportMode, InvitationToken, InvocationLevel,
InvocationLog, InvocationSource, JobStatus, MembershipRole, Operation, OperationId,
OperationSecurityLevel, OperationStatus, PlatformApiKey, Protocol, SampleId, Secret, SecretId,
SecretVersion, StreamSession, StreamSessionId, StreamStatus, UsagePeriod, UsageRollup, User,
UserSessionId, Workspace, WorkspaceId,
};
use crank_mapping::MappingSet;
use crank_schema::Schema;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
macro_rules! define_registry_id {
($name:ident) => {
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct $name(String);
impl $name {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl From<String> for $name {
fn from(value: String) -> Self {
Self(value)
}
}
impl From<&str> for $name {
fn from(value: &str) -> Self {
Self(value.to_owned())
}
}
};
}
define_registry_id!(YamlImportJobId);
pub type RegistryOperation = Operation<Schema, MappingSet>;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Page<T> {
pub items: Vec<T>,
pub total: u64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WorkspaceRecord {
pub workspace: Workspace,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct MembershipRecord {
pub workspace_id: WorkspaceId,
pub user: User,
pub role: MembershipRole,
pub created_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WorkspaceMembershipRecord {
pub workspace: Workspace,
pub role: MembershipRole,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AuthUserRecord {
pub user: User,
pub password_hash: String,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SessionRecord {
pub session_id: UserSessionId,
pub user: User,
pub memberships: Vec<WorkspaceMembershipRecord>,
pub current_workspace_id: Option<WorkspaceId>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct InvitationRecord {
pub invitation: InvitationToken,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct PlatformApiKeyRecord {
pub api_key: PlatformApiKey,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SecretRecord {
pub secret: Secret,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct SecretVersionRecord {
pub secret_version: SecretVersion,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StreamSessionRecord {
pub session: StreamSession,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AsyncJobRecord {
pub job: AsyncJobHandle,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct InvocationLogRecord {
pub log: InvocationLog,
pub operation_name: String,
pub operation_display_name: String,
pub agent_slug: Option<String>,
pub agent_display_name: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UsageRollupRecord {
pub rollup: UsageRollup,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UsageTimelinePoint {
#[serde(with = "time::serde::rfc3339")]
pub bucket_start: OffsetDateTime,
pub calls_ok: u64,
pub calls_error: u64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UsageOperationBreakdown {
pub operation_id: OperationId,
pub operation_name: String,
pub operation_display_name: String,
pub protocol: Protocol,
pub calls_total: u64,
pub calls_error: u64,
pub p50_ms: u64,
pub p95_ms: u64,
pub p99_ms: u64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UsageAgentBreakdown {
pub agent_id: AgentId,
pub agent_slug: String,
pub agent_display_name: String,
pub calls_total: u64,
pub calls_error: u64,
pub p50_ms: u64,
pub p95_ms: u64,
pub p99_ms: u64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AgentSummary {
pub id: AgentId,
pub workspace_id: WorkspaceId,
pub slug: String,
pub display_name: String,
pub description: String,
pub status: AgentStatus,
pub current_draft_version: u32,
pub latest_published_version: Option<u32>,
pub created_at: OffsetDateTime,
pub updated_at: OffsetDateTime,
pub published_at: Option<OffsetDateTime>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AgentVersionRecord {
pub agent_id: AgentId,
pub workspace_id: WorkspaceId,
pub version: u32,
pub status: AgentStatus,
pub created_at: OffsetDateTime,
pub snapshot: AgentVersion,
pub bindings: Vec<AgentOperationBinding>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct PublishedAgentTool {
pub workspace_id: WorkspaceId,
pub workspace_slug: String,
pub agent_id: AgentId,
pub agent_slug: String,
pub operation: RegistryOperation,
pub tool_name: String,
pub tool_title: String,
pub tool_description: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct OperationSummary {
pub id: OperationId,
pub workspace_id: WorkspaceId,
pub name: String,
pub display_name: String,
pub category: String,
pub protocol: Protocol,
pub security_level: OperationSecurityLevel,
pub target_url: String,
pub target_action: String,
pub status: OperationStatus,
pub current_draft_version: u32,
pub latest_published_version: Option<u32>,
pub created_at: OffsetDateTime,
pub updated_at: OffsetDateTime,
pub published_at: Option<OffsetDateTime>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OperationUsageSummary {
pub operation_id: OperationId,
pub calls_today: u64,
pub error_rate_pct: f64,
pub avg_latency_ms: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct OperationAgentRef {
pub operation_id: OperationId,
pub agent_id: AgentId,
pub agent_slug: String,
pub display_name: String,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OperationVersionRecord {
pub operation_id: OperationId,
pub workspace_id: WorkspaceId,
pub version: u32,
pub status: OperationStatus,
pub change_note: Option<String>,
pub created_at: OffsetDateTime,
pub created_by: Option<String>,
pub snapshot: RegistryOperation,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SampleKind {
InputJson,
OutputJson,
YamlImportSource,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct OperationSampleMetadata {
pub id: SampleId,
pub operation_id: OperationId,
pub version: u32,
pub sample_kind: SampleKind,
pub storage_ref: String,
pub content_type: String,
pub file_name: Option<String>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DescriptorKind {
ProtoUpload,
DescriptorSet,
ReflectionSnapshot,
WsdlUpload,
XsdUpload,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct DescriptorMetadata {
pub id: DescriptorId,
pub operation_id: Option<OperationId>,
pub version: Option<u32>,
pub descriptor_kind: DescriptorKind,
pub storage_ref: String,
pub source_name: Option<String>,
pub package_index: Option<Value>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum YamlImportJobStatus {
Pending,
Completed,
Failed,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct YamlImportJob {
pub id: YamlImportJobId,
pub source_sample_id: Option<SampleId>,
pub status: YamlImportJobStatus,
pub format_version: String,
pub mode: ExportMode,
pub result_operation_id: Option<OperationId>,
pub result_version: Option<u32>,
pub error_text: Option<String>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339::option")]
pub finished_at: Option<OffsetDateTime>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct YamlImportJobCompletion {
pub status: YamlImportJobStatus,
pub result_operation_id: Option<OperationId>,
pub result_version: Option<u32>,
pub error_text: Option<String>,
#[serde(with = "time::serde::rfc3339")]
pub finished_at: OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateInvocationLogRequest<'a> {
pub log: &'a InvocationLog,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ListInvocationLogsQuery<'a> {
pub workspace_id: &'a WorkspaceId,
pub level: Option<InvocationLevel>,
pub search_text: Option<&'a str>,
pub source: Option<InvocationSource>,
pub operation_id: Option<&'a OperationId>,
pub agent_id: Option<&'a AgentId>,
pub created_after: Option<&'a str>,
pub limit: u32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum UsageBucket {
Hour,
Day,
Week,
Month,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UsageQuery<'a> {
pub workspace_id: &'a WorkspaceId,
pub period: UsagePeriod,
pub source: Option<InvocationSource>,
pub created_after: &'a str,
pub bucket: UsageBucket,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct UsageSummary {
pub rollup: UsageRollup,
pub success_rate: f64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateVersionRequest<'a> {
pub workspace_id: &'a WorkspaceId,
pub snapshot: &'a RegistryOperation,
pub change_note: Option<&'a str>,
pub created_by: Option<&'a str>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PublishRequest<'a> {
pub workspace_id: &'a WorkspaceId,
pub operation_id: &'a OperationId,
pub version: u32,
pub published_at: &'a OffsetDateTime,
pub published_by: Option<&'a str>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CreateYamlImportJobRequest<'a> {
pub id: &'a YamlImportJobId,
pub source_sample_id: Option<&'a SampleId>,
pub format_version: &'a str,
pub mode: ExportMode,
pub created_at: &'a OffsetDateTime,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SaveAuthProfileRequest<'a> {
pub workspace_id: &'a WorkspaceId,
pub profile: &'a AuthProfile,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateWorkspaceRequest<'a> {
pub workspace: &'a Workspace,
}
#[derive(Clone, Debug, PartialEq)]
pub struct UpdateWorkspaceRequest<'a> {
pub workspace: &'a Workspace,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateAgentRequest<'a> {
pub agent: &'a Agent,
pub version: &'a AgentVersion,
pub bindings: &'a [AgentOperationBinding],
}
#[derive(Clone, Debug, PartialEq)]
pub struct SaveAgentBindingsRequest<'a> {
pub workspace_id: &'a WorkspaceId,
pub agent_id: &'a AgentId,
pub agent_version: u32,
pub bindings: &'a [AgentOperationBinding],
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PublishAgentRequest<'a> {
pub workspace_id: &'a WorkspaceId,
pub agent_id: &'a AgentId,
pub version: u32,
pub published_at: &'a OffsetDateTime,
pub published_by: Option<&'a str>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CreateInvitationRequest<'a> {
pub invitation: &'a InvitationToken,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreatePlatformApiKeyRequest<'a> {
pub api_key: &'a PlatformApiKey,
pub secret_hash: &'a str,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateSecretRequest<'a> {
pub secret: &'a Secret,
pub ciphertext: &'a str,
pub key_version: &'a str,
pub created_by: Option<&'a crank_core::UserId>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RotateSecretRequest<'a> {
pub workspace_id: &'a WorkspaceId,
pub secret_id: &'a SecretId,
pub ciphertext: &'a str,
pub key_version: &'a str,
pub created_at: &'a OffsetDateTime,
pub updated_at: &'a OffsetDateTime,
pub created_by: Option<&'a crank_core::UserId>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateStreamSessionRequest<'a> {
pub session: &'a StreamSession,
}
#[derive(Clone, Debug, PartialEq)]
pub struct UpdateStreamSessionStateRequest<'a> {
pub session_id: &'a StreamSessionId,
pub current_status: StreamStatus,
pub next_status: StreamStatus,
pub cursor: Option<&'a Value>,
pub state: &'a Value,
pub expires_at: Option<&'a OffsetDateTime>,
pub last_poll_at: Option<&'a OffsetDateTime>,
pub closed_at: Option<&'a OffsetDateTime>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct StreamSessionFilter<'a> {
pub workspace_id: &'a WorkspaceId,
pub agent_id: Option<&'a AgentId>,
pub operation_id: Option<&'a OperationId>,
pub status: Option<StreamStatus>,
pub mode: Option<ExecutionMode>,
pub limit: u32,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CreateAsyncJobRequest<'a> {
pub job: &'a AsyncJobHandle,
}
#[derive(Clone, Debug, PartialEq)]
pub struct UpdateAsyncJobStatusRequest<'a> {
pub job_id: &'a AsyncJobId,
pub current_status: JobStatus,
pub next_status: JobStatus,
pub progress: &'a Value,
pub result: Option<&'a Value>,
pub error: Option<&'a Value>,
pub expires_at: Option<&'a OffsetDateTime>,
pub updated_at: &'a OffsetDateTime,
pub finished_at: Option<&'a OffsetDateTime>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AsyncJobFilter<'a> {
pub workspace_id: &'a WorkspaceId,
pub agent_id: Option<&'a AgentId>,
pub operation_id: Option<&'a OperationId>,
pub status: Option<JobStatus>,
pub limit: u32,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SaveSampleMetadataRequest<'a> {
pub sample: &'a OperationSampleMetadata,
}
#[derive(Clone, Debug, PartialEq)]
pub struct SaveDescriptorMetadataRequest<'a> {
pub descriptor: &'a DescriptorMetadata,
}
+552
View File
@@ -0,0 +1,552 @@
use super::*;
impl PostgresRegistry {
pub async fn list_agents(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<AgentSummary>, RegistryError> {
let rows = sqlx::query!(
"select
id,
workspace_id,
slug,
display_name,
description,
status,
current_draft_version,
latest_published_version,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\",
published_at as \"published_at: time::OffsetDateTime\"
from agents
where workspace_id = $1
order by slug asc",
workspace_id.as_str(),
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
build_agent_summary(
row.id,
row.workspace_id,
row.slug,
row.display_name,
row.description,
row.status,
row.current_draft_version,
row.latest_published_version,
row.created_at,
row.updated_at,
row.published_at,
)
})
.collect()
}
pub async fn get_agent_summary(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
) -> Result<Option<AgentSummary>, RegistryError> {
let row = sqlx::query!(
"select
id,
workspace_id,
slug,
display_name,
description,
status,
current_draft_version,
latest_published_version,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\",
published_at as \"published_at: time::OffsetDateTime\"
from agents
where workspace_id = $1 and id = $2",
workspace_id.as_str(),
agent_id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
build_agent_summary(
row.id,
row.workspace_id,
row.slug,
row.display_name,
row.description,
row.status,
row.current_draft_version,
row.latest_published_version,
row.created_at,
row.updated_at,
row.published_at,
)
})
.transpose()
}
pub async fn create_agent(&self, request: CreateAgentRequest<'_>) -> Result<(), RegistryError> {
let mut tx = self.pool.begin().await?;
sqlx::query(
"insert into agents (
id,
workspace_id,
slug,
display_name,
description,
status,
current_draft_version,
latest_published_version,
created_at,
updated_at,
published_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9::timestamptz, $10::timestamptz, $11::timestamptz
)",
)
.bind(request.agent.id.as_str())
.bind(request.agent.workspace_id.as_str())
.bind(&request.agent.slug)
.bind(&request.agent.display_name)
.bind(&request.agent.description)
.bind(serialize_enum_text(&request.agent.status, "status")?)
.bind(to_db_version(request.agent.current_draft_version))
.bind(request.agent.latest_published_version.map(to_db_version))
.bind(request.agent.created_at)
.bind(request.agent.updated_at)
.bind(request.agent.published_at)
.execute(&mut *tx)
.await?;
insert_agent_version_row(&mut tx, request.version).await?;
replace_agent_bindings_rows(
&mut tx,
&request.agent.id,
request.version.version,
request.bindings,
)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn get_agent_version(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
version: u32,
) -> Result<Option<AgentVersionRecord>, RegistryError> {
let summary = match self.get_agent_summary(workspace_id, agent_id).await? {
Some(value) => value,
None => return Ok(None),
};
let row = sqlx::query!(
"select
agent_id,
version,
status,
instructions_json,
tool_selection_policy_json,
created_at as \"created_at!: time::OffsetDateTime\"
from agent_versions
where agent_id = $1 and version = $2",
agent_id.as_str(),
to_db_version(version),
)
.fetch_optional(&self.pool)
.await?;
let Some(row) = row else {
return Ok(None);
};
let bindings = self.list_agent_bindings(agent_id, version).await?;
Ok(Some(build_agent_version_record(
&summary,
row.version,
row.status,
row.instructions_json,
row.tool_selection_policy_json,
row.created_at,
bindings,
)?))
}
pub async fn save_agent_bindings(
&self,
request: SaveAgentBindingsRequest<'_>,
) -> Result<(), RegistryError> {
if self
.get_agent_summary(request.workspace_id, request.agent_id)
.await?
.is_none()
{
return Err(RegistryError::AgentNotFound {
agent_id: request.agent_id.as_str().to_owned(),
});
}
let mut tx = self.pool.begin().await?;
replace_agent_bindings_rows(
&mut tx,
request.agent_id,
request.agent_version,
request.bindings,
)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn update_agent_summary(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
slug: &str,
display_name: &str,
description: &str,
updated_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update agents
set slug = $3,
display_name = $4,
description = $5,
updated_at = $6::timestamptz
where workspace_id = $1 and id = $2",
)
.bind(workspace_id.as_str())
.bind(agent_id.as_str())
.bind(slug)
.bind(display_name)
.bind(description)
.bind(updated_at)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn delete_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
) -> Result<(), RegistryError> {
let result = sqlx::query("delete from agents where workspace_id = $1 and id = $2")
.bind(workspace_id.as_str())
.bind(agent_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn publish_agent(
&self,
request: PublishAgentRequest<'_>,
) -> Result<(), RegistryError> {
if self
.get_agent_version(request.workspace_id, request.agent_id, request.version)
.await?
.is_none()
{
return Err(RegistryError::AgentNotFound {
agent_id: request.agent_id.as_str().to_owned(),
});
}
let mut tx = self.pool.begin().await?;
sqlx::query(
"insert into published_agents (
agent_id,
version,
published_at,
published_by
) values ($1, $2, $3::timestamptz, $4)
on conflict(agent_id) do update set
version = excluded.version,
published_at = excluded.published_at,
published_by = excluded.published_by",
)
.bind(request.agent_id.as_str())
.bind(to_db_version(request.version))
.bind(request.published_at)
.bind(request.published_by)
.execute(&mut *tx)
.await?;
sqlx::query(
"update agent_versions
set status = $1
where agent_id = $2 and version = $3",
)
.bind(serialize_enum_text(&AgentStatus::Published, "status")?)
.bind(request.agent_id.as_str())
.bind(to_db_version(request.version))
.execute(&mut *tx)
.await?;
sqlx::query(
"update agents
set status = $1,
latest_published_version = $2,
published_at = $3::timestamptz,
updated_at = $4::timestamptz
where id = $5 and workspace_id = $6",
)
.bind(serialize_enum_text(&AgentStatus::Published, "status")?)
.bind(to_db_version(request.version))
.bind(request.published_at)
.bind(request.published_at)
.bind(request.agent_id.as_str())
.bind(request.workspace_id.as_str())
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
pub async fn unpublish_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
updated_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let mut tx = self.pool.begin().await?;
sqlx::query("delete from published_agents where agent_id = $1")
.bind(agent_id.as_str())
.execute(&mut *tx)
.await?;
sqlx::query(
"update agent_versions
set status = $1
where agent_id = $2
and version = (
select current_draft_version
from agents
where id = $2 and workspace_id = $3
)",
)
.bind(serialize_enum_text(&AgentStatus::Draft, "status")?)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
let result = sqlx::query(
"update agents
set status = $1,
published_at = null,
updated_at = $2::timestamptz
where id = $3 and workspace_id = $4",
)
.bind(serialize_enum_text(&AgentStatus::Draft, "status")?)
.bind(updated_at)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
tx.commit().await?;
Ok(())
}
pub async fn archive_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
updated_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let mut tx = self.pool.begin().await?;
sqlx::query("delete from published_agents where agent_id = $1")
.bind(agent_id.as_str())
.execute(&mut *tx)
.await?;
sqlx::query(
"update agent_versions
set status = $1
where agent_id = $2
and version = (
select current_draft_version
from agents
where id = $2 and workspace_id = $3
)",
)
.bind(serialize_enum_text(&AgentStatus::Archived, "status")?)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
let result = sqlx::query(
"update agents
set status = $1,
published_at = null,
updated_at = $2::timestamptz
where id = $3 and workspace_id = $4",
)
.bind(serialize_enum_text(&AgentStatus::Archived, "status")?)
.bind(updated_at)
.bind(agent_id.as_str())
.bind(workspace_id.as_str())
.execute(&mut *tx)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::AgentNotFound {
agent_id: agent_id.as_str().to_owned(),
});
}
tx.commit().await?;
Ok(())
}
pub async fn get_published_agent_tools_by_slug(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<Vec<PublishedAgentTool>, RegistryError> {
let rows = sqlx::query!(
"select
w.id as workspace_id,
w.slug as workspace_slug,
a.id as agent_id,
a.slug as agent_slug,
b.tool_name,
b.tool_title,
coalesce(b.tool_description_override, ov.tool_description_json->>'description') as \"tool_description!\",
o.id,
o.name,
o.display_name,
o.category,
o.protocol,
o.security_level,
o.created_at as \"operation_created_at!: time::OffsetDateTime\",
o.updated_at as \"operation_updated_at!: time::OffsetDateTime\",
o.published_at as \"operation_published_at: time::OffsetDateTime\",
ov.version,
ov.status,
ov.target_json,
ov.input_schema_json,
ov.output_schema_json,
ov.input_mapping_json,
ov.output_mapping_json,
ov.execution_config_json,
ov.tool_description_json,
ov.samples_json,
ov.generated_draft_json,
ov.config_export_json,
ov.change_note,
ov.created_at as \"created_at!: time::OffsetDateTime\",
ov.created_by
from workspaces w
join agents a on a.workspace_id = w.id
join published_agents pa on pa.agent_id = a.id
join agent_operation_bindings b on b.agent_id = a.id and b.agent_version = pa.version
join operation_versions ov on ov.operation_id = b.operation_id and ov.version = b.operation_version
join operations o on o.id = ov.operation_id and o.workspace_id = w.id
where w.slug = $1 and a.slug = $2 and b.enabled = true
order by b.tool_name asc",
workspace_slug,
agent_slug,
)
.fetch_all(&self.pool)
.await?;
if rows.is_empty() {
let exists = sqlx::query!(
"select 1 as \"present!\"
from workspaces w
join agents a on a.workspace_id = w.id
join published_agents pa on pa.agent_id = a.id
where w.slug = $1 and a.slug = $2",
workspace_slug,
agent_slug,
)
.fetch_optional(&self.pool)
.await?;
if exists.is_none() {
return Err(RegistryError::PublishedAgentNotFound {
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
});
}
}
rows.into_iter()
.map(|row| {
let workspace_id = row.workspace_id.clone();
build_published_agent_tool(
row.workspace_id,
row.workspace_slug,
row.agent_id,
row.agent_slug,
row.tool_name,
row.tool_title,
row.tool_description,
build_operation_version_record(
row.id,
workspace_id,
row.name,
row.display_name,
row.category,
row.protocol,
row.security_level,
row.operation_created_at,
row.operation_updated_at,
row.operation_published_at,
row.version,
row.status,
row.target_json,
row.input_schema_json,
row.output_schema_json,
row.input_mapping_json,
row.output_mapping_json,
row.execution_config_json,
row.tool_description_json,
row.samples_json,
row.generated_draft_json,
row.config_export_json,
row.change_note,
row.created_at,
row.created_by,
)?
.snapshot,
)
})
.collect()
}
}
@@ -0,0 +1,321 @@
use super::*;
impl PostgresRegistry {
pub async fn list_platform_api_keys_for_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
) -> Result<Vec<PlatformApiKeyRecord>, RegistryError> {
let rows = sqlx::query(
"select
id,
workspace_id,
agent_id,
name,
prefix,
scopes_json,
status,
created_at,
last_used_at
from platform_api_keys
where workspace_id = $1
and agent_id = $2
order by created_at desc",
)
.bind(workspace_id.as_str())
.bind(agent_id.as_str())
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok(PlatformApiKeyRecord {
api_key: PlatformApiKey {
id: PlatformApiKeyId::new(row.get::<String, _>("id")),
workspace_id: WorkspaceId::new(row.get::<String, _>("workspace_id")),
agent_id: row.get::<Option<String>, _>("agent_id").map(AgentId::new),
name: row.get("name"),
prefix: row.get("prefix"),
scopes: deserialize_json_value(row.get::<Value, _>("scopes_json"))?,
status: deserialize_enum_text(&row.get::<String, _>("status"), "status")?,
created_at: row.get("created_at"),
last_used_at: row.get("last_used_at"),
},
})
})
.collect()
}
pub async fn list_platform_api_keys(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<PlatformApiKeyRecord>, RegistryError> {
let rows = sqlx::query(
"select
id,
workspace_id,
agent_id,
name,
prefix,
scopes_json,
status,
created_at,
last_used_at
from platform_api_keys
where workspace_id = $1
order by created_at desc",
)
.bind(workspace_id.as_str())
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok(PlatformApiKeyRecord {
api_key: PlatformApiKey {
id: PlatformApiKeyId::new(row.get::<String, _>("id")),
workspace_id: WorkspaceId::new(row.get::<String, _>("workspace_id")),
agent_id: row.get::<Option<String>, _>("agent_id").map(AgentId::new),
name: row.get("name"),
prefix: row.get("prefix"),
scopes: deserialize_json_value(row.get::<Value, _>("scopes_json"))?,
status: deserialize_enum_text(&row.get::<String, _>("status"), "status")?,
created_at: row.get("created_at"),
last_used_at: row.get("last_used_at"),
},
})
})
.collect()
}
pub async fn get_platform_api_key_by_secret_for_agent_slug(
&self,
workspace_slug: &str,
agent_slug: &str,
secret_hash: &str,
) -> Result<Option<PlatformApiKeyRecord>, RegistryError> {
let row = sqlx::query(
"select
k.id,
k.workspace_id,
k.agent_id,
k.name,
k.prefix,
k.scopes_json,
k.status,
k.created_at,
k.last_used_at
from platform_api_keys k
join workspaces w on w.id = k.workspace_id
join agents a on a.id = k.agent_id
where w.slug = $1
and a.slug = $2
and k.secret_hash = $3
and k.status = 'active'
limit 1",
)
.bind(workspace_slug)
.bind(agent_slug)
.bind(secret_hash)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(PlatformApiKeyRecord {
api_key: PlatformApiKey {
id: PlatformApiKeyId::new(row.get::<String, _>("id")),
workspace_id: WorkspaceId::new(row.get::<String, _>("workspace_id")),
agent_id: row.get::<Option<String>, _>("agent_id").map(AgentId::new),
name: row.get("name"),
prefix: row.get("prefix"),
scopes: deserialize_json_value(row.get::<Value, _>("scopes_json"))?,
status: deserialize_enum_text(&row.get::<String, _>("status"), "status")?,
created_at: row.get("created_at"),
last_used_at: row.get("last_used_at"),
},
})
})
.transpose()
}
pub async fn create_platform_api_key(
&self,
request: CreatePlatformApiKeyRequest<'_>,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"insert into platform_api_keys (
id,
workspace_id,
agent_id,
name,
prefix,
secret_hash,
scopes_json,
status,
created_at,
last_used_at,
revoked_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9::timestamptz, $10::timestamptz, $11::timestamptz
)",
)
.bind(request.api_key.id.as_str())
.bind(request.api_key.workspace_id.as_str())
.bind(request.api_key.agent_id.as_ref().map(AgentId::as_str))
.bind(&request.api_key.name)
.bind(&request.api_key.prefix)
.bind(request.secret_hash)
.bind(Json(serialize_json_value(&request.api_key.scopes)?))
.bind(serialize_enum_text(&request.api_key.status, "status")?)
.bind(request.api_key.created_at)
.bind(request.api_key.last_used_at)
.bind(Option::<&str>::None)
.execute(&self.pool)
.await;
match result {
Ok(_) => Ok(()),
Err(sqlx::Error::Database(error))
if error.constraint() == Some("platform_api_keys_workspace_name_idx") =>
{
Err(RegistryError::Storage(sqlx::Error::Database(error)))
}
Err(error) => Err(RegistryError::Storage(error)),
}
}
pub async fn revoke_platform_api_key(
&self,
workspace_id: &WorkspaceId,
key_id: &PlatformApiKeyId,
revoked_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update platform_api_keys
set status = $1,
revoked_at = $2::timestamptz
where workspace_id = $3 and id = $4",
)
.bind(serialize_enum_text(
&PlatformApiKeyStatus::Revoked,
"status",
)?)
.bind(revoked_at)
.bind(workspace_id.as_str())
.bind(key_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::PlatformApiKeyNotFound {
key_id: key_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn revoke_platform_api_key_for_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
key_id: &PlatformApiKeyId,
revoked_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update platform_api_keys
set status = $1,
revoked_at = $2::timestamptz
where workspace_id = $3 and agent_id = $4 and id = $5",
)
.bind(serialize_enum_text(
&PlatformApiKeyStatus::Revoked,
"status",
)?)
.bind(revoked_at)
.bind(workspace_id.as_str())
.bind(agent_id.as_str())
.bind(key_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::PlatformApiKeyNotFound {
key_id: key_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn delete_platform_api_key(
&self,
workspace_id: &WorkspaceId,
key_id: &PlatformApiKeyId,
) -> Result<(), RegistryError> {
let result =
sqlx::query("delete from platform_api_keys where workspace_id = $1 and id = $2")
.bind(workspace_id.as_str())
.bind(key_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::PlatformApiKeyNotFound {
key_id: key_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn delete_platform_api_key_for_agent(
&self,
workspace_id: &WorkspaceId,
agent_id: &AgentId,
key_id: &PlatformApiKeyId,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"delete from platform_api_keys where workspace_id = $1 and agent_id = $2 and id = $3",
)
.bind(workspace_id.as_str())
.bind(agent_id.as_str())
.bind(key_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::PlatformApiKeyNotFound {
key_id: key_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn touch_platform_api_key(
&self,
workspace_id: &WorkspaceId,
key_id: &PlatformApiKeyId,
used_at: &time::OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update platform_api_keys
set last_used_at = $3::timestamptz
where workspace_id = $1 and id = $2",
)
.bind(workspace_id.as_str())
.bind(key_id.as_str())
.bind(used_at)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::PlatformApiKeyNotFound {
key_id: key_id.as_str().to_owned(),
});
}
Ok(())
}
}
+490
View File
@@ -0,0 +1,490 @@
use super::*;
use time::OffsetDateTime;
impl PostgresRegistry {
pub async fn upsert_bootstrap_user(
&self,
email: &str,
display_name: &str,
password_hash: &str,
) -> Result<UserId, RegistryError> {
let user_id = format!("user_{}", uuid::Uuid::now_v7().simple());
let row = sqlx::query!(
"insert into users (
id,
email,
display_name,
password_hash,
status,
created_at
) values (
$1, $2, $3, $4, 'active', now()
)
on conflict (email) do update
set display_name = excluded.display_name,
password_hash = excluded.password_hash,
status = 'active'
returning id",
user_id,
email,
display_name,
password_hash,
)
.fetch_one(&self.pool)
.await?;
Ok(UserId::new(row.id))
}
pub async fn ensure_membership(
&self,
workspace_id: &WorkspaceId,
user_id: &UserId,
role: MembershipRole,
) -> Result<(), RegistryError> {
sqlx::query(
"insert into memberships (
workspace_id,
user_id,
role,
created_at
) values (
$1, $2, $3, now()
)
on conflict (workspace_id, user_id) do update
set role = excluded.role",
)
.bind(workspace_id.as_str())
.bind(user_id.as_str())
.bind(serialize_enum_text(&role, "role")?)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn get_auth_user_by_email(
&self,
email: &str,
) -> Result<Option<AuthUserRecord>, RegistryError> {
let row = sqlx::query!(
"select
id,
email,
display_name,
password_hash as \"password_hash!\",
status,
created_at as \"created_at!: OffsetDateTime\"
from users
where email = $1
limit 1",
email,
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(AuthUserRecord {
user: User {
id: UserId::new(row.id),
email: row.email,
display_name: row.display_name,
status: deserialize_enum_text(&row.status, "status")?,
created_at: row.created_at,
},
password_hash: row.password_hash,
})
})
.transpose()
}
pub async fn get_auth_user_by_id(
&self,
user_id: &UserId,
) -> Result<Option<AuthUserRecord>, RegistryError> {
let row = sqlx::query!(
"select
id,
email,
display_name,
password_hash as \"password_hash!\",
status,
created_at as \"created_at!: OffsetDateTime\"
from users
where id = $1
limit 1",
user_id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(AuthUserRecord {
user: User {
id: UserId::new(row.id),
email: row.email,
display_name: row.display_name,
status: deserialize_enum_text(&row.status, "status")?,
created_at: row.created_at,
},
password_hash: row.password_hash,
})
})
.transpose()
}
pub async fn update_user_profile(
&self,
user_id: &UserId,
email: &str,
display_name: &str,
) -> Result<User, RegistryError> {
let row = sqlx::query!(
"update users
set email = $2,
display_name = $3
where id = $1
returning
id,
email,
display_name,
status,
created_at as \"created_at!: OffsetDateTime\"",
user_id.as_str(),
email,
display_name,
)
.fetch_one(&self.pool)
.await
.map_err(|error| map_user_update_error(error, user_id, email))?;
build_user(
row.id,
row.email,
row.display_name,
row.status,
row.created_at,
)
}
pub async fn update_user_password(
&self,
user_id: &UserId,
password_hash: &str,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update users
set password_hash = $2
where id = $1",
)
.bind(user_id.as_str())
.bind(password_hash)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::UserNotFound {
user_id: user_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn create_user_session(
&self,
session_id: &UserSessionId,
user_id: &UserId,
current_workspace_id: Option<&WorkspaceId>,
secret_hash: &str,
expires_at: &OffsetDateTime,
) -> Result<(), RegistryError> {
sqlx::query(
"insert into user_sessions (
id,
user_id,
current_workspace_id,
secret_hash,
status,
expires_at,
last_seen_at,
created_at
) values (
$1,
$2,
$3,
$4,
'active',
$5::timestamptz,
now(),
now()
)",
)
.bind(session_id.as_str())
.bind(user_id.as_str())
.bind(current_workspace_id.map(|id| id.as_str()))
.bind(secret_hash)
.bind(*expires_at)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn get_user_session(
&self,
session_id: &UserSessionId,
secret_hash: &str,
) -> Result<Option<SessionRecord>, RegistryError> {
let row = sqlx::query!(
"select
s.id,
s.user_id,
s.current_workspace_id,
u.email,
u.display_name,
u.status,
u.created_at as \"created_at!: OffsetDateTime\"
from user_sessions s
join users u on u.id = s.user_id
where s.id = $1
and s.secret_hash = $2
and s.status = 'active'
and s.expires_at > now()
limit 1",
session_id.as_str(),
secret_hash,
)
.fetch_optional(&self.pool)
.await?;
let Some(row) = row else {
return Ok(None);
};
let user_id = UserId::new(row.user_id);
let user = User {
id: user_id.clone(),
email: row.email,
display_name: row.display_name,
status: deserialize_enum_text(&row.status, "status")?,
created_at: row.created_at,
};
let memberships = self.list_workspaces_for_user(&user_id).await?;
let stored_workspace_id = row.current_workspace_id.map(WorkspaceId::new);
let current_workspace_id = stored_workspace_id
.filter(|workspace_id| {
memberships
.iter()
.any(|membership| membership.workspace.id == *workspace_id)
})
.or_else(|| {
memberships
.first()
.map(|membership| membership.workspace.id.clone())
});
Ok(Some(SessionRecord {
session_id: UserSessionId::new(row.id),
user,
memberships,
current_workspace_id,
}))
}
pub async fn touch_user_session(
&self,
session_id: &UserSessionId,
) -> Result<(), RegistryError> {
sqlx::query(
"update user_sessions
set last_seen_at = now()
where id = $1",
)
.bind(session_id.as_str())
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn revoke_user_session(
&self,
session_id: &UserSessionId,
) -> Result<(), RegistryError> {
sqlx::query(
"update user_sessions
set status = 'revoked'
where id = $1",
)
.bind(session_id.as_str())
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn set_user_session_current_workspace(
&self,
session_id: &UserSessionId,
workspace_id: &WorkspaceId,
) -> Result<(), RegistryError> {
sqlx::query(
"update user_sessions
set current_workspace_id = $2
where id = $1",
)
.bind(session_id.as_str())
.bind(workspace_id.as_str())
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn user_has_workspace_access(
&self,
user_id: &UserId,
workspace_id: &WorkspaceId,
) -> Result<bool, RegistryError> {
let row = sqlx::query!(
"select exists(
select 1
from memberships
where user_id = $1
and workspace_id = $2
) as \"allowed!\"",
user_id.as_str(),
workspace_id.as_str(),
)
.fetch_one(&self.pool)
.await?;
Ok(row.allowed)
}
pub async fn save_auth_profile(
&self,
request: SaveAuthProfileRequest<'_>,
) -> Result<(), RegistryError> {
sqlx::query(
"insert into auth_profiles (
id,
workspace_id,
name,
kind,
config_json,
created_at,
updated_at
) values ($1, $2, $3, $4, $5, $6::timestamptz, $7::timestamptz)
on conflict(id) do update set
workspace_id = excluded.workspace_id,
name = excluded.name,
kind = excluded.kind,
config_json = excluded.config_json,
updated_at = excluded.updated_at",
)
.bind(request.profile.id.as_str())
.bind(request.workspace_id.as_str())
.bind(&request.profile.name)
.bind(serialize_enum_text(&request.profile.kind, "kind")?)
.bind(Json(serialize_json_value(&request.profile.config)?))
.bind(request.profile.created_at)
.bind(request.profile.updated_at)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn get_auth_profile(
&self,
workspace_id: &WorkspaceId,
auth_profile_id: &crank_core::AuthProfileId,
) -> Result<Option<AuthProfile>, RegistryError> {
let row = sqlx::query!(
"select
id,
workspace_id,
name,
kind,
config_json,
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\"
from auth_profiles
where workspace_id = $1 and id = $2",
workspace_id.as_str(),
auth_profile_id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
build_auth_profile(
row.id,
row.workspace_id,
row.name,
row.kind,
row.config_json,
row.created_at,
row.updated_at,
)
})
.transpose()
}
pub async fn list_auth_profiles(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<AuthProfile>, RegistryError> {
let rows = sqlx::query!(
"select
id,
workspace_id,
name,
kind,
config_json,
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\"
from auth_profiles
where workspace_id = $1
order by name asc",
workspace_id.as_str(),
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
build_auth_profile(
row.id,
row.workspace_id,
row.name,
row.kind,
row.config_json,
row.created_at,
row.updated_at,
)
})
.collect()
}
pub async fn list_auth_profiles_referencing_secret(
&self,
workspace_id: &WorkspaceId,
secret_id: &SecretId,
) -> Result<Vec<AuthProfile>, RegistryError> {
let profiles = self.list_auth_profiles(workspace_id).await?;
Ok(profiles
.into_iter()
.filter(|profile| {
profile
.config
.secret_ids()
.into_iter()
.any(|candidate| candidate == secret_id)
})
.collect())
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,446 @@
use super::*;
impl PostgresRegistry {
pub async fn create_invocation_log(
&self,
request: CreateInvocationLogRequest<'_>,
) -> Result<(), RegistryError> {
sqlx::query(
"insert into invocation_logs (
id,
workspace_id,
agent_id,
operation_id,
source,
level,
status,
tool_name,
message,
request_id,
status_code,
duration_ms,
error_kind,
request_preview_json,
response_preview_json,
created_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16::timestamptz
)",
)
.bind(request.log.id.as_str())
.bind(request.log.workspace_id.as_str())
.bind(request.log.agent_id.as_ref().map(|value| value.as_str()))
.bind(request.log.operation_id.as_str())
.bind(serialize_enum_text(&request.log.source, "source")?)
.bind(serialize_enum_text(&request.log.level, "level")?)
.bind(serialize_enum_text(&request.log.status, "status")?)
.bind(&request.log.tool_name)
.bind(&request.log.message)
.bind(&request.log.request_id)
.bind(request.log.status_code.map(i32::from))
.bind(i64::try_from(request.log.duration_ms).map_err(|_| {
RegistryError::InvalidNumericValue {
field: "duration_ms",
value: i64::MAX,
}
})?)
.bind(&request.log.error_kind)
.bind(Json(request.log.request_preview.clone()))
.bind(Json(request.log.response_preview.clone()))
.bind(request.log.created_at)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn list_invocation_logs(
&self,
query: ListInvocationLogsQuery<'_>,
) -> Result<Vec<InvocationLogRecord>, RegistryError> {
let rows = sqlx::query(
"select
l.id,
l.workspace_id,
l.agent_id,
l.operation_id,
l.source,
l.level,
l.status,
l.tool_name,
l.message,
l.request_id,
l.status_code,
l.duration_ms,
l.error_kind,
l.request_preview_json,
l.response_preview_json,
l.created_at as created_at,
o.name as operation_name,
o.display_name as operation_display_name,
a.slug as agent_slug,
a.display_name as agent_display_name
from invocation_logs l
join operations o on o.id = l.operation_id
left join agents a on a.id = l.agent_id
where l.workspace_id = $1
and ($2::text is null or l.level = $2)
and ($3::text is null or l.source = $3)
and ($4::text is null or l.operation_id = $4)
and ($5::text is null or l.agent_id = $5)
and ($6::timestamptz is null or l.created_at >= $6::timestamptz)
and (
$7::text is null
or l.tool_name ilike '%' || $7 || '%'
or l.message ilike '%' || $7 || '%'
or o.name ilike '%' || $7 || '%'
or o.display_name ilike '%' || $7 || '%'
)
order by l.created_at desc
limit $8",
)
.bind(query.workspace_id.as_str())
.bind(
query
.level
.as_ref()
.map(|value| serialize_enum_text(value, "level"))
.transpose()?,
)
.bind(
query
.source
.as_ref()
.map(|value| serialize_enum_text(value, "source"))
.transpose()?,
)
.bind(query.operation_id.map(|value| value.as_str()))
.bind(query.agent_id.map(|value| value.as_str()))
.bind(query.created_after)
.bind(query.search_text)
.bind(i64::from(query.limit))
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_invocation_log_record).collect()
}
pub async fn get_invocation_log(
&self,
workspace_id: &WorkspaceId,
log_id: &InvocationLogId,
) -> Result<Option<InvocationLogRecord>, RegistryError> {
let row = sqlx::query(
"select
l.id,
l.workspace_id,
l.agent_id,
l.operation_id,
l.source,
l.level,
l.status,
l.tool_name,
l.message,
l.request_id,
l.status_code,
l.duration_ms,
l.error_kind,
l.request_preview_json,
l.response_preview_json,
l.created_at as created_at,
o.name as operation_name,
o.display_name as operation_display_name,
a.slug as agent_slug,
a.display_name as agent_display_name
from invocation_logs l
join operations o on o.id = l.operation_id
left join agents a on a.id = l.agent_id
where l.workspace_id = $1 and l.id = $2",
)
.bind(workspace_id.as_str())
.bind(log_id.as_str())
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_invocation_log_record).transpose()
}
pub async fn summarize_usage(
&self,
query: UsageQuery<'_>,
) -> Result<UsageSummary, RegistryError> {
let row = sqlx::query(
"select
count(*)::bigint as calls_total,
count(*) filter (where status = 'ok')::bigint as calls_ok,
count(*) filter (where status = 'error')::bigint as calls_error,
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
from invocation_logs
where workspace_id = $1
and created_at >= $2::timestamptz
and ($3::text is null or source = $3)",
)
.bind(query.workspace_id.as_str())
.bind(query.created_after)
.bind(
query
.source
.as_ref()
.map(|value| serialize_enum_text(value, "source"))
.transpose()?,
)
.fetch_one(&self.pool)
.await?;
let calls_total = row.try_get::<i64, _>("calls_total")?;
let calls_ok = row.try_get::<i64, _>("calls_ok")?;
let calls_error = row.try_get::<i64, _>("calls_error")?;
let rollup = UsageRollup {
workspace_id: query.workspace_id.clone(),
agent_id: None,
operation_id: None,
period: query.period,
calls_total: to_u64(calls_total, "calls_total")?,
calls_ok: to_u64(calls_ok, "calls_ok")?,
calls_error: to_u64(calls_error, "calls_error")?,
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
};
let success_rate = if rollup.calls_total == 0 {
0.0
} else {
(rollup.calls_ok as f64 / rollup.calls_total as f64) * 100.0
};
Ok(UsageSummary {
rollup,
success_rate,
})
}
pub async fn list_usage_timeline(
&self,
query: UsageQuery<'_>,
) -> Result<Vec<UsageTimelinePoint>, RegistryError> {
let bucket = usage_bucket_sql(query.bucket);
let sql = format!(
"select
(date_trunc('{bucket}', created_at at time zone 'UTC') at time zone 'UTC') as bucket_start,
count(*) filter (where status = 'ok')::bigint as calls_ok,
count(*) filter (where status = 'error')::bigint as calls_error
from invocation_logs
where workspace_id = $1
and created_at >= $2::timestamptz
and ($3::text is null or source = $3)
group by 1
order by 1 asc"
);
let rows = sqlx::query(&sql)
.bind(query.workspace_id.as_str())
.bind(query.created_after)
.bind(
query
.source
.as_ref()
.map(|value| serialize_enum_text(value, "source"))
.transpose()?,
)
.fetch_all(&self.pool)
.await?;
rows.iter()
.map(|row| {
Ok(UsageTimelinePoint {
bucket_start: row.try_get("bucket_start")?,
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
})
})
.collect()
}
pub async fn list_usage_by_operation(
&self,
query: UsageQuery<'_>,
) -> Result<Vec<UsageOperationBreakdown>, RegistryError> {
let rows = sqlx::query(
"select
o.id as operation_id,
o.name as operation_name,
o.display_name as operation_display_name,
o.protocol,
count(*)::bigint as calls_total,
count(*) filter (where l.status = 'error')::bigint as calls_error,
coalesce(percentile_cont(0.5) within group (order by l.duration_ms), 0)::bigint as p50_ms,
coalesce(percentile_cont(0.95) within group (order by l.duration_ms), 0)::bigint as p95_ms,
coalesce(percentile_cont(0.99) within group (order by l.duration_ms), 0)::bigint as p99_ms
from invocation_logs l
join operations o on o.id = l.operation_id
where l.workspace_id = $1
and l.created_at >= $2::timestamptz
and ($3::text is null or l.source = $3)
group by o.id, o.name, o.display_name, o.protocol
order by calls_total desc, o.name asc",
)
.bind(query.workspace_id.as_str())
.bind(query.created_after)
.bind(
query
.source
.as_ref()
.map(|value| serialize_enum_text(value, "source"))
.transpose()?,
)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_usage_operation_breakdown).collect()
}
pub async fn get_usage_for_operation(
&self,
query: UsageQuery<'_>,
operation_id: &OperationId,
) -> Result<Option<UsageRollupRecord>, RegistryError> {
let row = sqlx::query(
"select
count(*)::bigint as calls_total,
count(*) filter (where status = 'ok')::bigint as calls_ok,
count(*) filter (where status = 'error')::bigint as calls_error,
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
from invocation_logs
where workspace_id = $1
and operation_id = $2
and created_at >= $3::timestamptz
and ($4::text is null or source = $4)",
)
.bind(query.workspace_id.as_str())
.bind(operation_id.as_str())
.bind(query.created_after)
.bind(
query
.source
.as_ref()
.map(|value| serialize_enum_text(value, "source"))
.transpose()?,
)
.fetch_one(&self.pool)
.await?;
let calls_total = to_u64(row.try_get::<i64, _>("calls_total")?, "calls_total")?;
if calls_total == 0 {
return Ok(None);
}
Ok(Some(UsageRollupRecord {
rollup: UsageRollup {
workspace_id: query.workspace_id.clone(),
agent_id: None,
operation_id: Some(operation_id.clone()),
period: query.period,
calls_total,
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
},
}))
}
pub async fn list_usage_by_agent(
&self,
query: UsageQuery<'_>,
) -> Result<Vec<UsageAgentBreakdown>, RegistryError> {
let rows = sqlx::query(
"select
a.id as agent_id,
a.slug as agent_slug,
a.display_name as agent_display_name,
count(*)::bigint as calls_total,
count(*) filter (where l.status = 'error')::bigint as calls_error,
coalesce(percentile_cont(0.5) within group (order by l.duration_ms), 0)::bigint as p50_ms,
coalesce(percentile_cont(0.95) within group (order by l.duration_ms), 0)::bigint as p95_ms,
coalesce(percentile_cont(0.99) within group (order by l.duration_ms), 0)::bigint as p99_ms
from invocation_logs l
join agents a on a.id = l.agent_id
where l.workspace_id = $1
and l.created_at >= $2::timestamptz
and ($3::text is null or l.source = $3)
group by a.id, a.slug, a.display_name
order by calls_total desc, a.slug asc",
)
.bind(query.workspace_id.as_str())
.bind(query.created_after)
.bind(
query
.source
.as_ref()
.map(|value| serialize_enum_text(value, "source"))
.transpose()?,
)
.fetch_all(&self.pool)
.await?;
rows.iter().map(map_usage_agent_breakdown).collect()
}
pub async fn get_usage_for_agent(
&self,
query: UsageQuery<'_>,
agent_id: &AgentId,
) -> Result<Option<UsageRollupRecord>, RegistryError> {
let row = sqlx::query(
"select
count(*)::bigint as calls_total,
count(*) filter (where status = 'ok')::bigint as calls_ok,
count(*) filter (where status = 'error')::bigint as calls_error,
coalesce(percentile_cont(0.5) within group (order by duration_ms), 0)::bigint as p50_ms,
coalesce(percentile_cont(0.95) within group (order by duration_ms), 0)::bigint as p95_ms,
coalesce(percentile_cont(0.99) within group (order by duration_ms), 0)::bigint as p99_ms
from invocation_logs
where workspace_id = $1
and agent_id = $2
and created_at >= $3::timestamptz
and ($4::text is null or source = $4)",
)
.bind(query.workspace_id.as_str())
.bind(agent_id.as_str())
.bind(query.created_after)
.bind(
query
.source
.as_ref()
.map(|value| serialize_enum_text(value, "source"))
.transpose()?,
)
.fetch_one(&self.pool)
.await?;
let calls_total = to_u64(row.try_get::<i64, _>("calls_total")?, "calls_total")?;
if calls_total == 0 {
return Ok(None);
}
Ok(Some(UsageRollupRecord {
rollup: UsageRollup {
workspace_id: query.workspace_id.clone(),
agent_id: Some(agent_id.clone()),
operation_id: None,
period: query.period,
calls_total,
calls_ok: to_u64(row.try_get::<i64, _>("calls_ok")?, "calls_ok")?,
calls_error: to_u64(row.try_get::<i64, _>("calls_error")?, "calls_error")?,
p50_ms: to_u64(row.try_get::<i64, _>("p50_ms")?, "p50_ms")?,
p95_ms: to_u64(row.try_get::<i64, _>("p95_ms")?, "p95_ms")?,
p99_ms: to_u64(row.try_get::<i64, _>("p99_ms")?, "p99_ms")?,
},
}))
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,304 @@
use super::*;
impl PostgresRegistry {
pub async fn list_secrets(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<SecretRecord>, RegistryError> {
let rows = sqlx::query!(
"select
id,
workspace_id,
name,
kind,
status,
current_version,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\",
last_used_at as \"last_used_at: time::OffsetDateTime\"
from secrets
where workspace_id = $1
order by name asc",
workspace_id.as_str(),
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok(SecretRecord {
secret: Secret {
id: SecretId::new(row.id),
workspace_id: WorkspaceId::new(row.workspace_id),
name: row.name,
kind: deserialize_enum_text(&row.kind, "kind")?,
status: deserialize_enum_text(&row.status, "status")?,
current_version: from_db_version(row.current_version, "current_version")?,
created_at: row.created_at,
updated_at: row.updated_at,
last_used_at: row.last_used_at,
},
})
})
.collect()
}
pub async fn get_secret(
&self,
workspace_id: &WorkspaceId,
secret_id: &SecretId,
) -> Result<Option<SecretRecord>, RegistryError> {
let row = sqlx::query!(
"select
id,
workspace_id,
name,
kind,
status,
current_version,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\",
last_used_at as \"last_used_at: time::OffsetDateTime\"
from secrets
where workspace_id = $1 and id = $2",
workspace_id.as_str(),
secret_id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(SecretRecord {
secret: Secret {
id: SecretId::new(row.id),
workspace_id: WorkspaceId::new(row.workspace_id),
name: row.name,
kind: deserialize_enum_text(&row.kind, "kind")?,
status: deserialize_enum_text(&row.status, "status")?,
current_version: from_db_version(row.current_version, "current_version")?,
created_at: row.created_at,
updated_at: row.updated_at,
last_used_at: row.last_used_at,
},
})
})
.transpose()
}
pub async fn get_current_secret_version(
&self,
workspace_id: &WorkspaceId,
secret_id: &SecretId,
) -> Result<Option<SecretVersionRecord>, RegistryError> {
let row = sqlx::query!(
"select
sv.secret_id,
sv.version,
sv.ciphertext,
sv.key_version,
sv.created_at as \"created_at!: time::OffsetDateTime\",
sv.created_by
from secrets s
join secret_versions sv
on sv.secret_id = s.id and sv.version = s.current_version
where s.workspace_id = $1 and s.id = $2",
workspace_id.as_str(),
secret_id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(SecretVersionRecord {
secret_version: SecretVersion {
secret_id: SecretId::new(row.secret_id),
version: from_db_version(row.version, "version")?,
ciphertext: row.ciphertext,
key_version: row.key_version,
created_at: row.created_at,
created_by: row.created_by.map(UserId::new),
},
})
})
.transpose()
}
pub async fn create_secret(
&self,
request: CreateSecretRequest<'_>,
) -> Result<(), RegistryError> {
let mut tx = self.pool.begin().await?;
let result = sqlx::query(
"insert into secrets (
id,
workspace_id,
name,
kind,
status,
current_version,
last_used_at,
created_at,
updated_at
) values (
$1, $2, $3, $4, $5, $6, $7::timestamptz, $8::timestamptz, $9::timestamptz
)",
)
.bind(request.secret.id.as_str())
.bind(request.secret.workspace_id.as_str())
.bind(&request.secret.name)
.bind(serialize_enum_text(&request.secret.kind, "kind")?)
.bind(serialize_enum_text(&request.secret.status, "status")?)
.bind(to_db_version(request.secret.current_version))
.bind(request.secret.last_used_at)
.bind(request.secret.created_at)
.bind(request.secret.updated_at)
.execute(&mut *tx)
.await;
match result {
Ok(_) => {
sqlx::query(
"insert into secret_versions (
secret_id,
version,
ciphertext,
key_version,
created_at,
created_by
) values (
$1, $2, $3, $4, $5::timestamptz, $6
)",
)
.bind(request.secret.id.as_str())
.bind(to_db_version(request.secret.current_version))
.bind(request.ciphertext)
.bind(request.key_version)
.bind(request.secret.created_at)
.bind(request.created_by.map(|value| value.as_str()))
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
Err(sqlx::Error::Database(error))
if error.constraint() == Some("secrets_workspace_name_idx") =>
{
Err(RegistryError::SecretNameAlreadyExists {
workspace_id: request.secret.workspace_id.as_str().to_owned(),
name: request.secret.name.clone(),
})
}
Err(error) => Err(RegistryError::Storage(error)),
}
}
pub async fn rotate_secret(
&self,
request: RotateSecretRequest<'_>,
) -> Result<SecretVersionRecord, RegistryError> {
let existing = self
.get_secret(request.workspace_id, request.secret_id)
.await?
.ok_or_else(|| RegistryError::SecretNotFound {
secret_id: request.secret_id.as_str().to_owned(),
})?;
let next_version = existing.secret.current_version + 1;
let mut tx = self.pool.begin().await?;
sqlx::query(
"insert into secret_versions (
secret_id,
version,
ciphertext,
key_version,
created_at,
created_by
) values (
$1, $2, $3, $4, $5::timestamptz, $6
)",
)
.bind(request.secret_id.as_str())
.bind(to_db_version(next_version))
.bind(request.ciphertext)
.bind(request.key_version)
.bind(request.created_at)
.bind(request.created_by.map(|value| value.as_str()))
.execute(&mut *tx)
.await?;
sqlx::query(
"update secrets
set current_version = $3,
updated_at = $4::timestamptz
where workspace_id = $1 and id = $2",
)
.bind(request.workspace_id.as_str())
.bind(request.secret_id.as_str())
.bind(to_db_version(next_version))
.bind(request.updated_at)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(SecretVersionRecord {
secret_version: SecretVersion {
secret_id: request.secret_id.clone(),
version: next_version,
ciphertext: request.ciphertext.to_owned(),
key_version: request.key_version.to_owned(),
created_at: *request.created_at,
created_by: request.created_by.cloned(),
},
})
}
pub async fn delete_secret(
&self,
workspace_id: &WorkspaceId,
secret_id: &SecretId,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"delete from secrets
where workspace_id = $1 and id = $2",
)
.bind(workspace_id.as_str())
.bind(secret_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::SecretNotFound {
secret_id: secret_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn touch_secret(
&self,
workspace_id: &WorkspaceId,
secret_id: &SecretId,
used_at: &OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update secrets
set last_used_at = $3::timestamptz
where workspace_id = $1 and id = $2",
)
.bind(workspace_id.as_str())
.bind(secret_id.as_str())
.bind(used_at)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::SecretNotFound {
secret_id: secret_id.as_str().to_owned(),
});
}
Ok(())
}
}
@@ -0,0 +1,634 @@
use super::*;
use time::OffsetDateTime;
impl PostgresRegistry {
pub async fn create_stream_session(
&self,
request: CreateStreamSessionRequest<'_>,
) -> Result<StreamSession, RegistryError> {
sqlx::query(
"insert into stream_sessions (
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
expires_at,
last_poll_at,
created_at,
closed_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9::jsonb, $10::timestamptz,
$11::timestamptz, $12::timestamptz, $13::timestamptz
)",
)
.bind(request.session.id.as_str())
.bind(request.session.workspace_id.as_str())
.bind(
request
.session
.agent_id
.as_ref()
.map(|value| value.as_str()),
)
.bind(request.session.operation_id.as_str())
.bind(serialize_enum_text(&request.session.protocol, "protocol")?)
.bind(serialize_enum_text(&request.session.mode, "mode")?)
.bind(serialize_enum_text(&request.session.status, "status")?)
.bind(request.session.cursor.clone().map(Json))
.bind(Json(request.session.state.clone()))
.bind(request.session.expires_at)
.bind(request.session.last_poll_at)
.bind(request.session.created_at)
.bind(request.session.closed_at)
.execute(&self.pool)
.await?;
Ok(request.session.clone())
}
pub async fn get_stream_session(
&self,
id: &StreamSessionId,
) -> Result<Option<StreamSession>, RegistryError> {
let row = sqlx::query!(
"select
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
expires_at as \"expires_at!: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
closed_at as \"closed_at: OffsetDateTime\"
from stream_sessions
where id = $1",
id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(StreamSession {
id: StreamSessionId::new(row.id),
workspace_id: WorkspaceId::new(row.workspace_id),
agent_id: row.agent_id.map(AgentId::new),
operation_id: OperationId::new(row.operation_id),
protocol: deserialize_enum_text(&row.protocol, "protocol")?,
mode: deserialize_enum_text(&row.mode, "mode")?,
status: deserialize_enum_text(&row.status, "status")?,
cursor: row.cursor_json,
state: row.state_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
closed_at: row.closed_at,
})
})
.transpose()
}
pub async fn update_stream_session_state(
&self,
request: UpdateStreamSessionStateRequest<'_>,
) -> Result<StreamSession, RegistryError> {
validate_stream_session_transition(
request.session_id,
request.current_status,
request.next_status,
)?;
let row = sqlx::query(
"update stream_sessions
set status = $3,
cursor_json = $4::jsonb,
state_json = $5::jsonb,
expires_at = coalesce($6::timestamptz, expires_at),
last_poll_at = coalesce($7::timestamptz, last_poll_at),
closed_at = case
when $8::timestamptz is null then closed_at
else $8::timestamptz
end
where id = $1
and status = $2
returning
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
expires_at,
last_poll_at,
created_at,
closed_at",
)
.bind(request.session_id.as_str())
.bind(serialize_enum_text(&request.current_status, "status")?)
.bind(serialize_enum_text(&request.next_status, "status")?)
.bind(request.cursor.cloned().map(Json))
.bind(Json(request.state.clone()))
.bind(request.expires_at)
.bind(request.last_poll_at)
.bind(request.closed_at)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_stream_session(row),
None => {
if self.get_stream_session(request.session_id).await?.is_some() {
Err(RegistryError::InvalidStreamSessionTransition {
session_id: request.session_id.as_str().to_owned(),
from: serialize_enum_text(&request.current_status, "status")?,
to: serialize_enum_text(&request.next_status, "status")?,
})
} else {
Err(RegistryError::StreamSessionNotFound {
session_id: request.session_id.as_str().to_owned(),
})
}
}
}
}
pub async fn close_stream_session(
&self,
id: &StreamSessionId,
now: &OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update stream_sessions
set status = 'stopped',
last_poll_at = $2::timestamptz,
closed_at = $2::timestamptz
where id = $1
and status in ('created', 'running')",
)
.bind(id.as_str())
.bind(now)
.execute(&self.pool)
.await?;
if result.rows_affected() > 0 {
return Ok(());
}
match self.get_stream_session(id).await? {
Some(existing) => Err(RegistryError::InvalidStreamSessionTransition {
session_id: id.as_str().to_owned(),
from: serialize_enum_text(&existing.status, "status")?,
to: "stopped".to_owned(),
}),
None => Err(RegistryError::StreamSessionNotFound {
session_id: id.as_str().to_owned(),
}),
}
}
pub async fn touch_stream_session_poll(
&self,
id: &StreamSessionId,
now: &OffsetDateTime,
) -> Result<StreamSession, RegistryError> {
let row = sqlx::query(
"update stream_sessions
set last_poll_at = $2::timestamptz
where id = $1
returning
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
expires_at,
last_poll_at,
created_at,
closed_at",
)
.bind(id.as_str())
.bind(now)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_stream_session(row),
None => Err(RegistryError::StreamSessionNotFound {
session_id: id.as_str().to_owned(),
}),
}
}
pub async fn list_stream_sessions(
&self,
filter: StreamSessionFilter<'_>,
) -> Result<Page<StreamSession>, RegistryError> {
let status = filter
.status
.as_ref()
.map(|value| serialize_enum_text(value, "status"))
.transpose()?;
let mode = filter
.mode
.as_ref()
.map(|value| serialize_enum_text(value, "mode"))
.transpose()?;
let limit = i64::from(filter.limit);
let rows = sqlx::query!(
"select
id,
workspace_id,
agent_id,
operation_id,
protocol,
mode,
status,
cursor_json,
state_json,
expires_at as \"expires_at!: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
closed_at as \"closed_at: OffsetDateTime\"
from stream_sessions
where workspace_id = $1
and ($2::text is null or agent_id = $2)
and ($3::text is null or operation_id = $3)
and ($4::text is null or status = $4)
and ($5::text is null or mode = $5)
order by created_at desc
limit $6",
filter.workspace_id.as_str(),
filter.agent_id.map(|value| value.as_str()),
filter.operation_id.map(|value| value.as_str()),
status,
mode,
limit,
)
.fetch_all(&self.pool)
.await?;
let items = rows
.into_iter()
.map(|row| {
Ok(StreamSession {
id: StreamSessionId::new(row.id),
workspace_id: WorkspaceId::new(row.workspace_id),
agent_id: row.agent_id.map(AgentId::new),
operation_id: OperationId::new(row.operation_id),
protocol: deserialize_enum_text(&row.protocol, "protocol")?,
mode: deserialize_enum_text(&row.mode, "mode")?,
status: deserialize_enum_text(&row.status, "status")?,
cursor: row.cursor_json,
state: row.state_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
closed_at: row.closed_at,
})
})
.collect::<Result<Vec<_>, RegistryError>>()?;
let total = items.len() as u64;
Ok(Page { items, total })
}
pub async fn delete_expired_stream_sessions(
&self,
now: &OffsetDateTime,
) -> Result<u64, RegistryError> {
let result = sqlx::query(
"delete from stream_sessions
where expires_at <= $1::timestamptz",
)
.bind(now)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
pub async fn create_async_job(
&self,
request: CreateAsyncJobRequest<'_>,
) -> Result<AsyncJobHandle, RegistryError> {
sqlx::query(
"insert into async_jobs (
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at
) values (
$1, $2, $3, $4, $5, $6::jsonb, $7::jsonb, $8::jsonb, $9::timestamptz,
$10::timestamptz, $11::timestamptz, $12::timestamptz, $13::timestamptz
)",
)
.bind(request.job.id.as_str())
.bind(request.job.workspace_id.as_str())
.bind(request.job.agent_id.as_ref().map(|value| value.as_str()))
.bind(request.job.operation_id.as_str())
.bind(serialize_enum_text(&request.job.status, "status")?)
.bind(Json(request.job.progress.clone()))
.bind(request.job.result.clone().map(Json))
.bind(request.job.error.clone().map(Json))
.bind(request.job.expires_at)
.bind(request.job.last_poll_at)
.bind(request.job.created_at)
.bind(request.job.updated_at)
.bind(request.job.finished_at)
.execute(&self.pool)
.await?;
Ok(request.job.clone())
}
pub async fn get_async_job(
&self,
id: &AsyncJobId,
) -> Result<Option<AsyncJobHandle>, RegistryError> {
let row = sqlx::query!(
"select
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at as \"expires_at: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
finished_at as \"finished_at: OffsetDateTime\"
from async_jobs
where id = $1",
id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(AsyncJobHandle {
id: AsyncJobId::new(row.id),
workspace_id: WorkspaceId::new(row.workspace_id),
agent_id: row.agent_id.map(AgentId::new),
operation_id: OperationId::new(row.operation_id),
status: deserialize_enum_text(&row.status, "status")?,
progress: row.progress_json,
result: row.result_json,
error: row.error_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
updated_at: row.updated_at,
finished_at: row.finished_at,
})
})
.transpose()
}
pub async fn update_async_job_status(
&self,
request: UpdateAsyncJobStatusRequest<'_>,
) -> Result<AsyncJobHandle, RegistryError> {
validate_async_job_transition(request.job_id, request.current_status, request.next_status)?;
let row = sqlx::query(
"update async_jobs
set status = $3,
progress_json = $4::jsonb,
result_json = $5::jsonb,
error_json = $6::jsonb,
expires_at = coalesce($7::timestamptz, expires_at),
updated_at = $8::timestamptz,
finished_at = case
when $9::timestamptz is null then finished_at
else $9::timestamptz
end
where id = $1
and status = $2
returning
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at",
)
.bind(request.job_id.as_str())
.bind(serialize_enum_text(&request.current_status, "status")?)
.bind(serialize_enum_text(&request.next_status, "status")?)
.bind(Json(request.progress.clone()))
.bind(request.result.cloned().map(Json))
.bind(request.error.cloned().map(Json))
.bind(request.expires_at)
.bind(request.updated_at)
.bind(request.finished_at)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_async_job(row),
None => {
if self.get_async_job(request.job_id).await?.is_some() {
Err(RegistryError::InvalidAsyncJobTransition {
job_id: request.job_id.as_str().to_owned(),
from: serialize_enum_text(&request.current_status, "status")?,
to: serialize_enum_text(&request.next_status, "status")?,
})
} else {
Err(RegistryError::AsyncJobNotFound {
job_id: request.job_id.as_str().to_owned(),
})
}
}
}
}
pub async fn touch_async_job_poll(
&self,
id: &AsyncJobId,
now: &OffsetDateTime,
) -> Result<AsyncJobHandle, RegistryError> {
let row = sqlx::query(
"update async_jobs
set last_poll_at = $2::timestamptz
where id = $1
returning
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at,
last_poll_at,
created_at,
updated_at,
finished_at",
)
.bind(id.as_str())
.bind(now)
.fetch_optional(&self.pool)
.await?;
match row.as_ref() {
Some(row) => map_async_job(row),
None => Err(RegistryError::AsyncJobNotFound {
job_id: id.as_str().to_owned(),
}),
}
}
pub async fn cancel_async_job(
&self,
id: &AsyncJobId,
now: &OffsetDateTime,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update async_jobs
set status = 'cancelled',
updated_at = $2::timestamptz,
finished_at = $2::timestamptz
where id = $1
and status in ('created', 'running')",
)
.bind(id.as_str())
.bind(now)
.execute(&self.pool)
.await?;
if result.rows_affected() > 0 {
return Ok(());
}
match self.get_async_job(id).await? {
Some(existing) => Err(RegistryError::InvalidAsyncJobTransition {
job_id: id.as_str().to_owned(),
from: serialize_enum_text(&existing.status, "status")?,
to: "cancelled".to_owned(),
}),
None => Err(RegistryError::AsyncJobNotFound {
job_id: id.as_str().to_owned(),
}),
}
}
pub async fn list_async_jobs(
&self,
filter: AsyncJobFilter<'_>,
) -> Result<Page<AsyncJobHandle>, RegistryError> {
let status = filter
.status
.as_ref()
.map(|value| serialize_enum_text(value, "status"))
.transpose()?;
let limit = i64::from(filter.limit);
let rows = sqlx::query!(
"select
id,
workspace_id,
agent_id,
operation_id,
status,
progress_json,
result_json,
error_json,
expires_at as \"expires_at: OffsetDateTime\",
last_poll_at as \"last_poll_at: OffsetDateTime\",
created_at as \"created_at!: OffsetDateTime\",
updated_at as \"updated_at!: OffsetDateTime\",
finished_at as \"finished_at: OffsetDateTime\"
from async_jobs
where workspace_id = $1
and ($2::text is null or agent_id = $2)
and ($3::text is null or operation_id = $3)
and ($4::text is null or status = $4)
order by updated_at desc
limit $5",
filter.workspace_id.as_str(),
filter.agent_id.map(|value| value.as_str()),
filter.operation_id.map(|value| value.as_str()),
status,
limit,
)
.fetch_all(&self.pool)
.await?;
let items = rows
.into_iter()
.map(|row| {
Ok(AsyncJobHandle {
id: AsyncJobId::new(row.id),
workspace_id: WorkspaceId::new(row.workspace_id),
agent_id: row.agent_id.map(AgentId::new),
operation_id: OperationId::new(row.operation_id),
status: deserialize_enum_text(&row.status, "status")?,
progress: row.progress_json,
result: row.result_json,
error: row.error_json,
expires_at: row.expires_at,
last_poll_at: row.last_poll_at,
created_at: row.created_at,
updated_at: row.updated_at,
finished_at: row.finished_at,
})
})
.collect::<Result<Vec<_>, RegistryError>>()?;
let total = items.len() as u64;
Ok(Page { items, total })
}
pub async fn delete_expired_async_jobs(
&self,
now: &OffsetDateTime,
) -> Result<u64, RegistryError> {
let result = sqlx::query(
"delete from async_jobs
where expires_at is not null
and expires_at <= $1::timestamptz",
)
.bind(now)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
}
@@ -0,0 +1,391 @@
use super::*;
impl PostgresRegistry {
pub async fn list_workspaces(&self) -> Result<Vec<WorkspaceRecord>, RegistryError> {
let rows = sqlx::query!(
"select
id,
slug,
display_name,
status,
settings_json,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\"
from workspaces
order by slug asc",
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok(WorkspaceRecord {
workspace: Workspace {
id: WorkspaceId::new(row.id),
slug: row.slug,
display_name: row.display_name,
status: deserialize_enum_text(&row.status, "status")?,
settings: row.settings_json,
created_at: row.created_at,
updated_at: row.updated_at,
},
})
})
.collect()
}
pub async fn list_workspaces_for_user(
&self,
user_id: &UserId,
) -> Result<Vec<WorkspaceMembershipRecord>, RegistryError> {
let rows = sqlx::query!(
"select
w.id,
w.slug,
w.display_name,
w.status,
w.settings_json,
w.created_at as \"created_at!: time::OffsetDateTime\",
w.updated_at as \"updated_at!: time::OffsetDateTime\",
m.role
from memberships m
join workspaces w on w.id = m.workspace_id
where m.user_id = $1
order by w.slug asc",
user_id.as_str(),
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok(WorkspaceMembershipRecord {
workspace: Workspace {
id: WorkspaceId::new(row.id),
slug: row.slug,
display_name: row.display_name,
status: deserialize_enum_text(&row.status, "status")?,
settings: row.settings_json,
created_at: row.created_at,
updated_at: row.updated_at,
},
role: deserialize_enum_text(&row.role, "role")?,
})
})
.collect()
}
pub async fn list_memberships(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<MembershipRecord>, RegistryError> {
let rows = sqlx::query!(
"select
m.workspace_id,
m.user_id,
m.role,
m.created_at as \"created_at!: time::OffsetDateTime\",
u.email,
u.display_name,
u.status,
u.created_at as \"user_created_at!: time::OffsetDateTime\"
from memberships m
join users u on u.id = m.user_id
where m.workspace_id = $1
order by u.email asc",
workspace_id.as_str(),
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok(MembershipRecord {
workspace_id: WorkspaceId::new(row.workspace_id),
user: User {
id: UserId::new(row.user_id),
email: row.email,
display_name: row.display_name,
status: deserialize_enum_text(&row.status, "status")?,
created_at: row.user_created_at,
},
role: deserialize_enum_text(&row.role, "role")?,
created_at: row.created_at,
})
})
.collect()
}
pub async fn update_membership_role(
&self,
workspace_id: &WorkspaceId,
user_id: &UserId,
role: MembershipRole,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update memberships
set role = $3
where workspace_id = $1 and user_id = $2",
)
.bind(workspace_id.as_str())
.bind(user_id.as_str())
.bind(serialize_enum_text(&role, "role")?)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::MembershipNotFound {
workspace_id: workspace_id.as_str().to_owned(),
user_id: user_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn delete_membership(
&self,
workspace_id: &WorkspaceId,
user_id: &UserId,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"delete from memberships
where workspace_id = $1 and user_id = $2",
)
.bind(workspace_id.as_str())
.bind(user_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::MembershipNotFound {
workspace_id: workspace_id.as_str().to_owned(),
user_id: user_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn list_invitations(
&self,
workspace_id: &WorkspaceId,
) -> Result<Vec<InvitationRecord>, RegistryError> {
let rows = sqlx::query!(
"select
id,
workspace_id,
email,
role,
status,
token_hash,
expires_at as \"expires_at!: time::OffsetDateTime\",
created_at as \"created_at!: time::OffsetDateTime\"
from invitation_tokens
where workspace_id = $1
order by created_at desc",
workspace_id.as_str(),
)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok(InvitationRecord {
invitation: InvitationToken {
id: InvitationId::new(row.id),
workspace_id: WorkspaceId::new(row.workspace_id),
email: row.email,
role: deserialize_enum_text(&row.role, "role")?,
status: deserialize_enum_text(&row.status, "status")?,
token_hash: row.token_hash,
expires_at: row.expires_at,
created_at: row.created_at,
},
})
})
.collect()
}
pub async fn create_invitation(
&self,
request: CreateInvitationRequest<'_>,
) -> Result<(), RegistryError> {
sqlx::query(
"insert into invitation_tokens (
id,
workspace_id,
email,
role,
status,
token_hash,
expires_at,
created_at
) values (
$1, $2, $3, $4, $5, $6, $7::timestamptz, $8::timestamptz
)",
)
.bind(request.invitation.id.as_str())
.bind(request.invitation.workspace_id.as_str())
.bind(&request.invitation.email)
.bind(serialize_enum_text(&request.invitation.role, "role")?)
.bind(serialize_enum_text(&request.invitation.status, "status")?)
.bind(&request.invitation.token_hash)
.bind(request.invitation.expires_at)
.bind(request.invitation.created_at)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn delete_invitation(
&self,
workspace_id: &WorkspaceId,
invitation_id: &InvitationId,
) -> Result<(), RegistryError> {
let result =
sqlx::query("delete from invitation_tokens where workspace_id = $1 and id = $2")
.bind(workspace_id.as_str())
.bind(invitation_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::InvitationNotFound {
invitation_id: invitation_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn delete_workspace(&self, workspace_id: &WorkspaceId) -> Result<(), RegistryError> {
let result = sqlx::query("delete from workspaces where id = $1")
.bind(workspace_id.as_str())
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::WorkspaceNotFound {
workspace_id: workspace_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn get_workspace(
&self,
workspace_id: &WorkspaceId,
) -> Result<Option<WorkspaceRecord>, RegistryError> {
let row = sqlx::query!(
"select
id,
slug,
display_name,
status,
settings_json,
created_at as \"created_at!: time::OffsetDateTime\",
updated_at as \"updated_at!: time::OffsetDateTime\"
from workspaces
where id = $1",
workspace_id.as_str(),
)
.fetch_optional(&self.pool)
.await?;
row.map(|row| {
Ok(WorkspaceRecord {
workspace: Workspace {
id: WorkspaceId::new(row.id),
slug: row.slug,
display_name: row.display_name,
status: deserialize_enum_text(&row.status, "status")?,
settings: row.settings_json,
created_at: row.created_at,
updated_at: row.updated_at,
},
})
})
.transpose()
}
pub async fn create_workspace(
&self,
request: CreateWorkspaceRequest<'_>,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"insert into workspaces (
id,
slug,
display_name,
status,
settings_json,
created_at,
updated_at
) values (
$1, $2, $3, $4, $5, $6::timestamptz, $7::timestamptz
)",
)
.bind(request.workspace.id.as_str())
.bind(&request.workspace.slug)
.bind(&request.workspace.display_name)
.bind(serialize_enum_text(&request.workspace.status, "status")?)
.bind(Json(request.workspace.settings.clone()))
.bind(request.workspace.created_at)
.bind(request.workspace.updated_at)
.execute(&self.pool)
.await;
match result {
Ok(_) => Ok(()),
Err(sqlx::Error::Database(error))
if error.constraint() == Some("workspaces_slug_key") =>
{
Err(RegistryError::WorkspaceSlugAlreadyExists {
slug: request.workspace.slug.clone(),
})
}
Err(error) => Err(RegistryError::Storage(error)),
}
}
pub async fn update_workspace(
&self,
request: UpdateWorkspaceRequest<'_>,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update workspaces
set slug = $1,
display_name = $2,
status = $3,
settings_json = $4,
updated_at = $5::timestamptz
where id = $6",
)
.bind(&request.workspace.slug)
.bind(&request.workspace.display_name)
.bind(serialize_enum_text(&request.workspace.status, "status")?)
.bind(Json(request.workspace.settings.clone()))
.bind(request.workspace.updated_at)
.bind(request.workspace.id.as_str())
.execute(&self.pool)
.await;
match result {
Ok(result) if result.rows_affected() == 0 => Err(RegistryError::WorkspaceNotFound {
workspace_id: request.workspace.id.as_str().to_owned(),
}),
Ok(_) => Ok(()),
Err(sqlx::Error::Database(error))
if error.constraint() == Some("workspaces_slug_key") =>
{
Err(RegistryError::WorkspaceSlugAlreadyExists {
slug: request.workspace.slug.clone(),
})
}
Err(error) => Err(RegistryError::Storage(error)),
}
}
}
@@ -0,0 +1,101 @@
use super::*;
impl PostgresRegistry {
pub async fn create_yaml_import_job(
&self,
request: CreateYamlImportJobRequest<'_>,
) -> Result<(), RegistryError> {
sqlx::query(
"insert into yaml_import_jobs (
id,
source_sample_id,
status,
format_version,
mode,
result_operation_id,
result_version,
error_text,
created_at,
finished_at
) values (
$1, $2, $3, $4, $5, null, null, null, $6::timestamptz, null
)",
)
.bind(request.id.as_str())
.bind(request.source_sample_id.map(|value| value.as_str()))
.bind(serialize_enum_text(
&YamlImportJobStatus::Pending,
"status",
)?)
.bind(request.format_version)
.bind(serialize_enum_text(&request.mode, "mode")?)
.bind(request.created_at)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn finish_yaml_import_job(
&self,
job_id: &YamlImportJobId,
completion: &YamlImportJobCompletion,
) -> Result<(), RegistryError> {
let result = sqlx::query(
"update yaml_import_jobs
set status = $2,
result_operation_id = $3,
result_version = $4,
error_text = $5,
finished_at = $6::timestamptz
where id = $1",
)
.bind(job_id.as_str())
.bind(serialize_enum_text(&completion.status, "status")?)
.bind(
completion
.result_operation_id
.as_ref()
.map(|value| value.as_str()),
)
.bind(completion.result_version.map(to_db_version))
.bind(completion.error_text.as_deref())
.bind(completion.finished_at)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(RegistryError::YamlImportJobNotFound {
job_id: job_id.as_str().to_owned(),
});
}
Ok(())
}
pub async fn get_yaml_import_job(
&self,
job_id: &YamlImportJobId,
) -> Result<Option<YamlImportJob>, RegistryError> {
let row = sqlx::query(
"select
id,
source_sample_id,
status,
format_version,
mode,
result_operation_id,
result_version,
error_text,
created_at,
finished_at
from yaml_import_jobs
where id = $1",
)
.bind(job_id.as_str())
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(map_yaml_import_job).transpose()
}
}
+36
View File
@@ -0,0 +1,36 @@
[package]
name = "crank-runtime"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[lib]
path = "src/lib.rs"
[features]
default = []
[dependencies]
aes-gcm.workspace = true
async-trait = "0.1"
base64.workspace = true
crank-adapter-rest = { path = "../crank-adapter-rest" }
crank-core = { path = "../crank-core" }
crank-mapping = { path = "../crank-mapping" }
crank-schema = { path = "../crank-schema" }
hkdf.workspace = true
redis = { version = "0.29", features = ["tokio-comp", "connection-manager"] }
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
time.workspace = true
tokio = { workspace = true, features = ["sync"] }
tracing.workspace = true
[dev-dependencies]
axum.workspace = true
futures-util = "0.3"
time.workspace = true
tracing-subscriber.workspace = true
+292
View File
@@ -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..."));
}
}
+245
View File
@@ -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")
);
}
}
+903
View File
@@ -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
}
);
}
}
+45
View File
@@ -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);
}
}
+53
View File
@@ -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)
}
+35
View File
@@ -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;
+119
View File
@@ -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");
}
}
}
+106
View File
@@ -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,
}
}
}
+269
View File
@@ -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
}
}
}
+124
View File
@@ -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..."));
}
}
+170
View File
@@ -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);
}
}
+274
View File
@@ -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}"),
}
}
}
+26
View File
@@ -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,
}
}
}
+15
View File
@@ -0,0 +1,15 @@
[package]
name = "crank-schema"
edition.workspace = true
license.workspace = true
rust-version.workspace = true
version.workspace = true
[dependencies]
crank-core = { path = "../crank-core" }
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
[dev-dependencies]
serde_yaml.workspace = true
+55
View File
@@ -0,0 +1,55 @@
use thiserror::Error;
#[derive(Clone, Debug, PartialEq, Eq, Error)]
pub enum SchemaError {
#[error("missing required field at {path}")]
MissingField { path: String },
#[error("invalid value at {path}: expected {expected}, got {actual}")]
InvalidType {
path: String,
expected: &'static str,
actual: &'static str,
},
#[error("invalid enum value at {path}: expected one of {expected:?}, got {actual}")]
InvalidEnumValue {
path: String,
expected: Vec<String>,
actual: String,
},
#[error("no matching oneof variant at {path}")]
OneofMismatch { path: String },
}
impl SchemaError {
pub fn missing_field(path: impl Into<String>) -> Self {
Self::MissingField { path: path.into() }
}
pub fn invalid_type(
path: impl Into<String>,
expected: &'static str,
actual: &'static str,
) -> Self {
Self::InvalidType {
path: path.into(),
expected,
actual,
}
}
pub fn invalid_enum_value(
path: impl Into<String>,
expected: Vec<String>,
actual: impl Into<String>,
) -> Self {
Self::InvalidEnumValue {
path: path.into(),
expected,
actual: actual.into(),
}
}
pub fn oneof_mismatch(path: impl Into<String>) -> Self {
Self::OneofMismatch { path: path.into() }
}
}
+7
View File
@@ -0,0 +1,7 @@
mod error;
mod model;
mod normalize;
mod validate;
pub use error::SchemaError;
pub use model::{Schema, SchemaKind, XmlNodeKind, XmlQualifiedName, XmlSchemaBinding};
+226
View File
@@ -0,0 +1,226 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum XmlNodeKind {
Element,
Attribute,
Text,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct XmlQualifiedName {
pub local_name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub namespace_uri: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub prefix: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct XmlSchemaBinding {
pub node_kind: XmlNodeKind,
pub name: XmlQualifiedName,
#[serde(default)]
pub repeated: bool,
#[serde(default)]
pub nillable: bool,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SchemaKind {
Object,
Array,
String,
Integer,
Number,
Boolean,
Enum,
#[serde(rename = "null")]
Null,
Oneof,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Schema {
#[serde(rename = "type")]
pub kind: SchemaKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(default)]
pub required: bool,
#[serde(default)]
pub nullable: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub default_value: Option<Value>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub fields: BTreeMap<String, Schema>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub items: Option<Box<Schema>>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub enum_values: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub variants: Vec<Schema>,
}
impl Schema {
pub fn is_object(&self) -> bool {
self.kind == SchemaKind::Object
}
pub fn field(&self, name: &str) -> Option<&Schema> {
self.fields.get(name)
}
pub fn has_required_fields(&self) -> bool {
self.fields.values().any(|field| field.required)
}
pub fn expected_type_name(&self) -> &'static str {
match self.kind {
SchemaKind::Object => "object",
SchemaKind::Array => "array",
SchemaKind::String => "string",
SchemaKind::Integer => "integer",
SchemaKind::Number => "number",
SchemaKind::Boolean => "boolean",
SchemaKind::Enum => "enum",
SchemaKind::Null => "null",
SchemaKind::Oneof => "oneof",
}
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{Schema, SchemaKind, XmlNodeKind, XmlQualifiedName, XmlSchemaBinding};
#[test]
fn object_schema_exposes_fields() {
let mut fields = BTreeMap::new();
fields.insert(
"email".to_owned(),
Schema {
kind: SchemaKind::String,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
},
);
let schema = Schema {
kind: SchemaKind::Object,
description: Some("User input".to_owned()),
required: true,
nullable: false,
default_value: None,
fields,
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
};
assert!(schema.is_object());
assert!(schema.has_required_fields());
assert_eq!(
schema.field("email").map(|field| field.required),
Some(true)
);
}
#[test]
fn schema_serializes_type_field() {
let schema = Schema {
kind: SchemaKind::Array,
description: None,
required: false,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: Some(Box::new(Schema {
kind: SchemaKind::String,
description: None,
required: false,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
})),
enum_values: Vec::new(),
variants: Vec::new(),
};
let value = serde_json::to_value(schema).unwrap();
assert_eq!(value["type"], "array");
assert_eq!(value["items"]["type"], "string");
}
#[test]
fn schema_roundtrips_through_yaml() {
let schema = Schema {
kind: SchemaKind::Object,
description: Some("Lead output".to_owned()),
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::from([(
"id".to_owned(),
Schema {
kind: SchemaKind::String,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
},
)]),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
};
let yaml = serde_yaml::to_string(&schema).unwrap();
let restored: Schema = serde_yaml::from_str(&yaml).unwrap();
assert!(yaml.contains("type: object"));
assert_eq!(restored, schema);
}
#[test]
fn xml_schema_binding_roundtrips() {
let binding = XmlSchemaBinding {
node_kind: XmlNodeKind::Element,
name: XmlQualifiedName {
local_name: "CreateLeadRequest".to_owned(),
namespace_uri: Some("urn:crm".to_owned()),
prefix: Some("crm".to_owned()),
},
repeated: false,
nillable: true,
};
let value = serde_json::to_value(&binding).unwrap();
let restored: XmlSchemaBinding = serde_json::from_value(value.clone()).unwrap();
assert_eq!(value["node_kind"], "element");
assert_eq!(value["name"]["local_name"], "CreateLeadRequest");
assert_eq!(restored, binding);
}
}
+163
View File
@@ -0,0 +1,163 @@
use std::collections::BTreeMap;
use serde_json::{Map, Value};
use crate::{Schema, SchemaKind};
impl Schema {
pub fn from_json_sample(value: &Value) -> Self {
match value {
Value::Null => Self {
kind: SchemaKind::Null,
description: None,
required: true,
nullable: true,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
},
Value::Bool(_) => scalar_schema(SchemaKind::Boolean),
Value::String(_) => scalar_schema(SchemaKind::String),
Value::Number(number) if number.is_i64() || number.is_u64() => {
scalar_schema(SchemaKind::Integer)
}
Value::Number(_) => scalar_schema(SchemaKind::Number),
Value::Array(items) => array_schema(items),
Value::Object(fields) => object_schema(fields),
}
}
}
fn scalar_schema(kind: SchemaKind) -> Schema {
Schema {
kind,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
}
}
fn array_schema(items: &[Value]) -> Schema {
Schema {
kind: SchemaKind::Array,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: items
.first()
.map(|item| Box::new(Schema::from_json_sample(item))),
enum_values: Vec::new(),
variants: Vec::new(),
}
}
fn object_schema(fields: &Map<String, Value>) -> Schema {
Schema {
kind: SchemaKind::Object,
description: None,
required: true,
nullable: false,
default_value: None,
fields: fields
.iter()
.map(|(name, value)| (name.clone(), Schema::from_json_sample(value)))
.collect(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use crate::{Schema, SchemaKind};
#[test]
fn infers_object_schema_from_json_sample() {
let schema = Schema::from_json_sample(&json!({
"lead": {
"email": "user@example.com",
"active": true
},
"score": 42
}));
assert_eq!(schema.kind, SchemaKind::Object);
assert_eq!(
schema.field("lead").map(|field| field.kind.clone()),
Some(SchemaKind::Object)
);
assert_eq!(
schema
.field("lead")
.and_then(|field| field.field("email"))
.map(|field| field.kind.clone()),
Some(SchemaKind::String)
);
assert_eq!(
schema.field("score").map(|field| field.kind.clone()),
Some(SchemaKind::Integer)
);
}
#[test]
fn infers_array_item_schema_from_first_sample_item() {
let schema = Schema::from_json_sample(&json!([
{
"id": "lead_01",
"tags": ["warm", "priority"]
}
]));
assert_eq!(schema.kind, SchemaKind::Array);
assert_eq!(
schema
.items
.as_deref()
.and_then(|item| item.field("id"))
.map(|field| field.kind.clone()),
Some(SchemaKind::String)
);
assert_eq!(
schema
.items
.as_deref()
.and_then(|item| item.field("tags"))
.map(|field| field.kind.clone()),
Some(SchemaKind::Array)
);
}
#[test]
fn infers_null_schema_for_null_sample() {
let schema = Schema::from_json_sample(&serde_json::Value::Null);
assert_eq!(
schema,
Schema {
kind: SchemaKind::Null,
description: None,
required: true,
nullable: true,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
}
);
}
}
+260
View File
@@ -0,0 +1,260 @@
use serde_json::Value;
use crate::{Schema, SchemaError, SchemaKind};
impl Schema {
pub fn validate_shape(&self, value: &Value) -> Result<(), SchemaError> {
validate_value(self, value, "$")
}
}
fn validate_value(schema: &Schema, value: &Value, path: &str) -> Result<(), SchemaError> {
if value.is_null() {
if schema.nullable || schema.kind == SchemaKind::Null {
return Ok(());
}
return Err(SchemaError::invalid_type(
path,
schema.expected_type_name(),
"null",
));
}
match schema.kind {
SchemaKind::Object => validate_object(schema, value, path),
SchemaKind::Array => validate_array(schema, value, path),
SchemaKind::String => validate_scalar(schema, value, path, Value::is_string),
SchemaKind::Integer => validate_scalar(schema, value, path, Value::is_i64_or_u64),
SchemaKind::Number => validate_scalar(schema, value, path, Value::is_number),
SchemaKind::Boolean => validate_scalar(schema, value, path, Value::is_boolean),
SchemaKind::Enum => validate_enum(schema, value, path),
SchemaKind::Null => Err(SchemaError::invalid_type(
path,
"null",
actual_type_name(value),
)),
SchemaKind::Oneof => validate_oneof(schema, value, path),
}
}
fn validate_object(schema: &Schema, value: &Value, path: &str) -> Result<(), SchemaError> {
let Some(object) = value.as_object() else {
return Err(SchemaError::invalid_type(
path,
"object",
actual_type_name(value),
));
};
for (field_name, field_schema) in &schema.fields {
let field_path = format!("{path}.{field_name}");
match object.get(field_name) {
Some(field_value) => validate_value(field_schema, field_value, &field_path)?,
None if field_schema.required => return Err(SchemaError::missing_field(field_path)),
None => {}
}
}
Ok(())
}
fn validate_array(schema: &Schema, value: &Value, path: &str) -> Result<(), SchemaError> {
let Some(items) = value.as_array() else {
return Err(SchemaError::invalid_type(
path,
"array",
actual_type_name(value),
));
};
if let Some(item_schema) = schema.items.as_deref() {
for (index, item) in items.iter().enumerate() {
let item_path = format!("{path}[{index}]");
validate_value(item_schema, item, &item_path)?;
}
}
Ok(())
}
fn validate_scalar(
schema: &Schema,
value: &Value,
path: &str,
is_valid: impl Fn(&Value) -> bool,
) -> Result<(), SchemaError> {
if is_valid(value) {
return Ok(());
}
Err(SchemaError::invalid_type(
path,
schema.expected_type_name(),
actual_type_name(value),
))
}
fn validate_enum(schema: &Schema, value: &Value, path: &str) -> Result<(), SchemaError> {
let Some(actual) = value.as_str() else {
return Err(SchemaError::invalid_type(
path,
"enum",
actual_type_name(value),
));
};
if schema.enum_values.iter().any(|expected| expected == actual) {
return Ok(());
}
Err(SchemaError::invalid_enum_value(
path,
schema.enum_values.clone(),
actual,
))
}
fn validate_oneof(schema: &Schema, value: &Value, path: &str) -> Result<(), SchemaError> {
if schema
.variants
.iter()
.any(|variant| validate_value(variant, value, path).is_ok())
{
return Ok(());
}
Err(SchemaError::oneof_mismatch(path))
}
fn actual_type_name(value: &Value) -> &'static str {
match value {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(number) if number.is_i64() || number.is_u64() => "integer",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
trait JsonValueExt {
fn is_i64_or_u64(&self) -> bool;
}
impl JsonValueExt for Value {
fn is_i64_or_u64(&self) -> bool {
self.as_i64().is_some() || self.as_u64().is_some()
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use crate::{Schema, SchemaError, SchemaKind};
#[test]
fn validates_nested_object_shape() {
let schema = Schema {
kind: SchemaKind::Object,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::from([(
"lead".to_owned(),
Schema {
kind: SchemaKind::Object,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::from([(
"email".to_owned(),
Schema {
kind: SchemaKind::String,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
},
)]),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
},
)]),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
};
let result = schema.validate_shape(&json!({
"lead": {
"email": "user@example.com"
}
}));
assert!(result.is_ok());
}
#[test]
fn rejects_missing_required_field() {
let schema = Schema {
kind: SchemaKind::Object,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::from([(
"email".to_owned(),
Schema {
kind: SchemaKind::String,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
},
)]),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
};
let error = schema.validate_shape(&json!({})).unwrap_err();
assert_eq!(error, SchemaError::missing_field("$.email"));
}
#[test]
fn rejects_invalid_scalar_type() {
let schema = Schema {
kind: SchemaKind::Boolean,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
};
let error = schema.validate_shape(&json!("true")).unwrap_err();
assert_eq!(error, SchemaError::invalid_type("$", "boolean", "string"));
}
}