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