Polish community UI copy and cleanup
This commit is contained in:
github-ops
@@ -7,7 +7,7 @@ use reqwest::{
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{RestAdapterError, RestRequest, RestResponse, RestWindowRequest, RestWindowResponse};
|
||||
use crate::{RestAdapterError, RestRequest, RestResponse};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RestAdapter {
|
||||
@@ -62,61 +62,6 @@ impl RestAdapter {
|
||||
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> {
|
||||
@@ -227,15 +172,13 @@ mod tests {
|
||||
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};
|
||||
use crate::{RestAdapter, RestAdapterError, RestRequest};
|
||||
|
||||
#[tokio::test]
|
||||
async fn executes_rest_request_and_normalizes_json_response() {
|
||||
@@ -299,104 +242,10 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[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));
|
||||
.route("/fail", get(fail));
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
|
||||
@@ -437,27 +286,4 @@ mod tests {
|
||||
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())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,18 +1,16 @@
|
||||
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,
|
||||
ProtocolAdapterError, RestTarget, RuntimeRequestContext, Target,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub use client::RestAdapter;
|
||||
pub use error::RestAdapterError;
|
||||
pub use model::{RestRequest, RestResponse, RestWindowRequest, RestWindowResponse};
|
||||
pub use model::{RestRequest, RestResponse};
|
||||
|
||||
#[async_trait]
|
||||
impl ProtocolAdapter for RestAdapter {
|
||||
@@ -21,7 +19,7 @@ impl ProtocolAdapter for RestAdapter {
|
||||
}
|
||||
|
||||
fn supports_mode(&self, mode: ExecutionMode) -> bool {
|
||||
matches!(mode, ExecutionMode::Unary | ExecutionMode::Window)
|
||||
matches!(mode, ExecutionMode::Unary)
|
||||
}
|
||||
|
||||
async fn invoke_unary(
|
||||
@@ -47,72 +45,11 @@ impl ProtocolAdapter for RestAdapter {
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -23,20 +23,3 @@ pub struct RestResponse {
|
||||
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,
|
||||
}
|
||||
|
||||
@@ -1,160 +0,0 @@
|
||||
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));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,9 +4,13 @@ use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
AgentId, ExecutionMode, InvocationSource, Protocol, StreamSession, Target, WorkspaceId,
|
||||
};
|
||||
use crate::{AgentId, InvocationSource, Protocol, Target, WorkspaceId};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ExecutionMode {
|
||||
Unary,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Default)]
|
||||
pub struct ResponseCacheScope {
|
||||
@@ -95,10 +99,6 @@ pub struct PreparedRequest {
|
||||
#[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,
|
||||
@@ -113,17 +113,6 @@ pub struct AdapterResponse {
|
||||
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:?}")]
|
||||
@@ -147,32 +136,6 @@ pub trait ProtocolAdapter: Send + Sync {
|
||||
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>;
|
||||
@@ -211,18 +174,16 @@ impl AdapterRegistry {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
use std::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,
|
||||
AdapterRegistry, AdapterResponse, ExecutionMode, PreparedRequest, ProtocolAdapter,
|
||||
ProtocolAdapterError, RuntimeRequestContext,
|
||||
};
|
||||
use crate::{InvocationSource, Protocol, RestTarget, Target, WorkspaceId};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct StubAdapter(Protocol);
|
||||
@@ -245,25 +206,13 @@ mod tests {
|
||||
) -> Result<AdapterResponse, ProtocolAdapterError> {
|
||||
Ok(AdapterResponse {
|
||||
status_code: 200,
|
||||
headers: BTreeMap::new(),
|
||||
headers: Default::default(),
|
||||
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()
|
||||
@@ -271,46 +220,36 @@ mod tests {
|
||||
.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() {
|
||||
fn request_context_builds_headers_and_metering_context() {
|
||||
let context = RuntimeRequestContext::from_request_id("req_123").with_metering_context(
|
||||
WorkspaceId::new("ws_01"),
|
||||
None,
|
||||
InvocationSource::AgentToolCall,
|
||||
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::AgentToolCall);
|
||||
assert_eq!(
|
||||
context.outbound_headers().get("x-request-id"),
|
||||
Some(&"req_123".to_owned())
|
||||
);
|
||||
assert!(context.metering_context().is_some());
|
||||
}
|
||||
|
||||
#[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();
|
||||
async fn adapter_contract_invokes_unary() {
|
||||
let adapter = StubAdapter(Protocol::Rest);
|
||||
let target = Target::Rest(RestTarget {
|
||||
base_url: "https://api.example.com".to_owned(),
|
||||
method: crate::HttpMethod::Get,
|
||||
path_template: "/health".to_owned(),
|
||||
static_headers: Default::default(),
|
||||
});
|
||||
|
||||
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(),
|
||||
}),
|
||||
&target,
|
||||
&PreparedRequest::default(),
|
||||
&RuntimeRequestContext::from_request_id("req_1"),
|
||||
)
|
||||
|
||||
@@ -57,8 +57,6 @@ define_id!(PlatformApiKeyId);
|
||||
define_id!(InvocationLogId);
|
||||
define_id!(AuditEventId);
|
||||
define_id!(SecretId);
|
||||
define_id!(StreamSessionId);
|
||||
define_id!(AsyncJobId);
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
|
||||
@@ -9,9 +9,6 @@ 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::{
|
||||
@@ -53,36 +50,26 @@ pub use ext::billing::{
|
||||
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,
|
||||
AdapterRegistry, AdapterResponse, ExecutionMode, MeteringContext, PreparedRequest,
|
||||
ProtocolAdapter, ProtocolAdapterError, ResponseCacheScope, RuntimeRequestContext,
|
||||
SharedProtocolAdapter,
|
||||
};
|
||||
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,
|
||||
AgentId, AuditEventId, AuthProfileId, DescriptorId, InvitationId, InvocationLogId, OperationId,
|
||||
PlatformApiKeyId, SampleId, SecretId, 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, WizardState,
|
||||
ConfigExport, ExecutionConfig, GeneratedDraft, GeneratedDraftStatus, Operation,
|
||||
OperationStatus, ResponseCachePolicy, RestTarget, RetryPolicy, Samples, Target,
|
||||
ToolDescription, ToolExample, WizardState,
|
||||
};
|
||||
pub use protocol::{AuthKind, ExportMode, GraphqlOperationType, HttpMethod, Protocol};
|
||||
pub use protocol::{AuthKind, ExportMode, 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};
|
||||
|
||||
@@ -6,12 +6,8 @@ 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,
|
||||
ids::{AuthProfileId, OperationId, SampleId},
|
||||
protocol::{ExportMode, HttpMethod, Protocol},
|
||||
};
|
||||
|
||||
fn default_operation_category() -> String {
|
||||
@@ -40,68 +36,10 @@ pub struct RestTarget {
|
||||
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)]
|
||||
@@ -114,41 +52,6 @@ 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,
|
||||
@@ -160,10 +63,6 @@ pub struct ExecutionConfig {
|
||||
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)]
|
||||
@@ -187,10 +86,6 @@ pub struct Samples {
|
||||
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)]
|
||||
@@ -301,17 +196,12 @@ mod tests {
|
||||
use crate::{
|
||||
auth::{AuthConfig, AuthProfile, BearerAuthConfig},
|
||||
edition::OperationSecurityLevel,
|
||||
ids::{AuthProfileId, OperationId, SampleId},
|
||||
ids::{AuthProfileId, OperationId},
|
||||
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,
|
||||
ConfigExport, ExecutionConfig, Operation, OperationStatus, RestTarget, Samples, Target,
|
||||
ToolDescription, ToolExample,
|
||||
},
|
||||
protocol::{AuthKind, ExportMode, HttpMethod, Protocol},
|
||||
};
|
||||
|
||||
fn timestamp(value: &str) -> OffsetDateTime {
|
||||
@@ -333,132 +223,9 @@ mod tests {
|
||||
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,
|
||||
wizard_state: None,
|
||||
created_at: timestamp("2026-03-25T08:00:00Z"),
|
||||
updated_at: timestamp("2026-03-25T08:00:00Z"),
|
||||
published_at: None,
|
||||
};
|
||||
let operation = test_operation(OperationStatus::Draft);
|
||||
|
||||
assert_eq!(operation.tool_name(), "crm_create_lead");
|
||||
assert!(operation.is_draft());
|
||||
@@ -497,51 +264,14 @@ mod tests {
|
||||
#[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,
|
||||
}),
|
||||
wizard_state: None,
|
||||
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")),
|
||||
..test_operation(OperationStatus::Draft)
|
||||
};
|
||||
|
||||
let yaml = serde_yaml::to_string(&operation).unwrap();
|
||||
@@ -597,4 +327,49 @@ updated_at: 2026-03-25T08:10:00Z
|
||||
assert_eq!(restored.security_level, OperationSecurityLevel::Standard);
|
||||
assert_eq!(restored.published_at, None);
|
||||
}
|
||||
|
||||
fn test_operation(status: OperationStatus) -> Operation<serde_json::Value, serde_json::Value> {
|
||||
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,
|
||||
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","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(),
|
||||
},
|
||||
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: None,
|
||||
wizard_state: None,
|
||||
created_at: timestamp("2026-03-25T08:00:00Z"),
|
||||
updated_at: timestamp("2026-03-25T08:10:00Z"),
|
||||
published_at: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,15 +1,9 @@
|
||||
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)]
|
||||
@@ -22,13 +16,6 @@ pub enum HttpMethod {
|
||||
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 {
|
||||
@@ -44,73 +31,3 @@ 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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
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")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,318 +0,0 @@
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,364 +0,0 @@
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -7,10 +7,9 @@ mod postgres;
|
||||
pub use error::RegistryError;
|
||||
pub use ext::{ExtensionMigration, RegistryExtension, apply_extension_migrations};
|
||||
pub use model::{
|
||||
AgentSummary, AgentVersionRecord, AsyncJobFilter, AsyncJobRecord, AuthUserRecord,
|
||||
CreateAgentDraftVersionRequest, CreateAgentRequest, CreateAsyncJobRequest,
|
||||
CreateInvitationRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
|
||||
CreateSecretRequest, CreateStreamSessionRequest, CreateVersionRequest, CreateWorkspaceRequest,
|
||||
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentDraftVersionRequest,
|
||||
CreateAgentRequest, CreateInvitationRequest, CreateInvocationLogRequest,
|
||||
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateVersionRequest, CreateWorkspaceRequest,
|
||||
CreateYamlImportJobRequest, DescriptorKind, DescriptorMetadata, InvitationRecord,
|
||||
InvocationLogRecord, ListInvocationLogsQuery, MembershipRecord, OperationAgentRef,
|
||||
OperationSampleMetadata, OperationSummary, OperationUsageSummary, OperationVersionRecord, Page,
|
||||
@@ -18,10 +17,9 @@ pub use model::{
|
||||
RegistryOperation, RotateSecretRequest, SampleKind, SaveAgentBindingsRequest,
|
||||
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
|
||||
SaveWorkspaceUpstreamRequest, SecretRecord, SecretVersionRecord, SessionRecord,
|
||||
StreamSessionFilter, StreamSessionRecord, UpdateAsyncJobStatusRequest,
|
||||
UpdateStreamSessionStateRequest, UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket,
|
||||
UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary, UsageTimelinePoint,
|
||||
WorkspaceMembershipRecord, WorkspaceRecord, WorkspaceUpstream, WorkspaceUpstreamId,
|
||||
YamlImportJob, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
UpdateWorkspaceRequest, UsageAgentBreakdown, UsageBucket, UsageOperationBreakdown, UsageQuery,
|
||||
UsageRollupRecord, UsageSummary, UsageTimelinePoint, WorkspaceMembershipRecord,
|
||||
WorkspaceRecord, WorkspaceUpstream, WorkspaceUpstreamId, YamlImportJob,
|
||||
YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
};
|
||||
pub use postgres::{PostgresPoolConfig, PostgresPoolConfigError, PostgresRegistry};
|
||||
|
||||
@@ -583,77 +583,5 @@ pub async fn apply_postgres(pool: &PgPool) -> Result<(), sqlx::Error> {
|
||||
.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(())
|
||||
}
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
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,
|
||||
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, DescriptorId,
|
||||
ExportMode, InvitationToken, InvocationLevel, InvocationLog, InvocationSource, MembershipRole,
|
||||
Operation, OperationId, OperationSecurityLevel, OperationStatus, PlatformApiKey, Protocol,
|
||||
SampleId, Secret, SecretId, SecretVersion, UsagePeriod, UsageRollup, User, UserSessionId,
|
||||
Workspace, WorkspaceId,
|
||||
};
|
||||
use crank_mapping::MappingSet;
|
||||
use crank_schema::Schema;
|
||||
@@ -117,16 +116,6 @@ pub struct WorkspaceUpstream {
|
||||
pub updated_at: OffsetDateTime,
|
||||
}
|
||||
|
||||
#[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,
|
||||
@@ -487,60 +476,6 @@ pub struct RotateSecretRequest<'a> {
|
||||
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,
|
||||
|
||||
@@ -4,18 +4,15 @@ mod auth;
|
||||
mod observability;
|
||||
mod operation;
|
||||
mod secret;
|
||||
mod stream;
|
||||
mod upstream;
|
||||
mod workspace;
|
||||
mod yaml_import;
|
||||
|
||||
use crank_core::{
|
||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AsyncJobHandle, AsyncJobId,
|
||||
AuthProfile, GraphqlOperationType, HttpMethod, InvitationId, InvitationToken, InvocationLog,
|
||||
InvocationLogId, JobStatus, MembershipRole, OperationId, OperationStatus, PlatformApiKey,
|
||||
PlatformApiKeyId, PlatformApiKeyStatus, Secret, SecretId, SecretVersion, StreamSession,
|
||||
StreamSessionId, StreamStatus, Target, UsageRollup, User, UserId, UserSessionId, Workspace,
|
||||
WorkspaceId,
|
||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AuthProfile, HttpMethod,
|
||||
InvitationId, InvitationToken, InvocationLog, InvocationLogId, MembershipRole, OperationId,
|
||||
OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyStatus, Secret, SecretId,
|
||||
SecretVersion, Target, UsageRollup, User, UserId, UserSessionId, Workspace, WorkspaceId,
|
||||
};
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::Value;
|
||||
@@ -32,18 +29,16 @@ use crate::{
|
||||
error::RegistryError,
|
||||
migrations,
|
||||
model::{
|
||||
AgentSummary, AgentVersionRecord, AsyncJobFilter, AuthUserRecord,
|
||||
CreateAgentDraftVersionRequest, CreateAgentRequest, CreateAsyncJobRequest,
|
||||
CreateInvitationRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
|
||||
CreateSecretRequest, CreateStreamSessionRequest, CreateVersionRequest,
|
||||
AgentSummary, AgentVersionRecord, AuthUserRecord, CreateAgentDraftVersionRequest,
|
||||
CreateAgentRequest, CreateInvitationRequest, CreateInvocationLogRequest,
|
||||
CreatePlatformApiKeyRequest, CreateSecretRequest, CreateVersionRequest,
|
||||
CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorMetadata, InvitationRecord,
|
||||
InvocationLogRecord, ListInvocationLogsQuery, MembershipRecord, OperationAgentRef,
|
||||
OperationSampleMetadata, OperationSummary, OperationUsageSummary, OperationVersionRecord,
|
||||
Page, PlatformApiKeyRecord, PublishAgentRequest, PublishRequest, PublishedAgentTool,
|
||||
PlatformApiKeyRecord, PublishAgentRequest, PublishRequest, PublishedAgentTool,
|
||||
RegistryOperation, RotateSecretRequest, SaveAgentBindingsRequest, SaveAuthProfileRequest,
|
||||
SaveDescriptorMetadataRequest, SaveSampleMetadataRequest, SaveWorkspaceUpstreamRequest,
|
||||
SecretRecord, SecretVersionRecord, SessionRecord, StreamSessionFilter,
|
||||
UpdateAsyncJobStatusRequest, UpdateStreamSessionStateRequest, UpdateWorkspaceRequest,
|
||||
SecretRecord, SecretVersionRecord, SessionRecord, UpdateWorkspaceRequest,
|
||||
UsageAgentBreakdown, UsageOperationBreakdown, UsageQuery, UsageRollupRecord, UsageSummary,
|
||||
UsageTimelinePoint, WorkspaceMembershipRecord, WorkspaceRecord, WorkspaceUpstream,
|
||||
YamlImportJob, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
@@ -520,52 +515,6 @@ fn map_user_update_error(error: sqlx::Error, user_id: &UserId, email: &str) -> R
|
||||
}
|
||||
}
|
||||
|
||||
fn map_stream_session(row: &PgRow) -> Result<StreamSession, RegistryError> {
|
||||
Ok(StreamSession {
|
||||
id: StreamSessionId::new(row.try_get::<String, _>("id")?),
|
||||
workspace_id: WorkspaceId::new(row.try_get::<String, _>("workspace_id")?),
|
||||
agent_id: row
|
||||
.try_get::<Option<String>, _>("agent_id")?
|
||||
.map(AgentId::new),
|
||||
operation_id: OperationId::new(row.try_get::<String, _>("operation_id")?),
|
||||
protocol: deserialize_enum_text(&row.try_get::<String, _>("protocol")?, "protocol")?,
|
||||
mode: deserialize_enum_text(&row.try_get::<String, _>("mode")?, "mode")?,
|
||||
status: deserialize_enum_text(&row.try_get::<String, _>("status")?, "status")?,
|
||||
cursor: row
|
||||
.try_get::<Option<Json<Value>>, _>("cursor_json")?
|
||||
.map(|value| value.0),
|
||||
state: row.try_get::<Json<Value>, _>("state_json")?.0,
|
||||
expires_at: row.try_get("expires_at")?,
|
||||
last_poll_at: row.try_get("last_poll_at")?,
|
||||
created_at: row.try_get("created_at")?,
|
||||
closed_at: row.try_get("closed_at")?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_async_job(row: &PgRow) -> Result<AsyncJobHandle, RegistryError> {
|
||||
Ok(AsyncJobHandle {
|
||||
id: AsyncJobId::new(row.try_get::<String, _>("id")?),
|
||||
workspace_id: WorkspaceId::new(row.try_get::<String, _>("workspace_id")?),
|
||||
agent_id: row
|
||||
.try_get::<Option<String>, _>("agent_id")?
|
||||
.map(AgentId::new),
|
||||
operation_id: OperationId::new(row.try_get::<String, _>("operation_id")?),
|
||||
status: deserialize_enum_text(&row.try_get::<String, _>("status")?, "status")?,
|
||||
progress: row.try_get::<Json<Value>, _>("progress_json")?.0,
|
||||
result: row
|
||||
.try_get::<Option<Json<Value>>, _>("result_json")?
|
||||
.map(|value| value.0),
|
||||
error: row
|
||||
.try_get::<Option<Json<Value>>, _>("error_json")?
|
||||
.map(|value| value.0),
|
||||
expires_at: row.try_get("expires_at")?,
|
||||
last_poll_at: row.try_get("last_poll_at")?,
|
||||
created_at: row.try_get("created_at")?,
|
||||
updated_at: row.try_get("updated_at")?,
|
||||
finished_at: row.try_get("finished_at")?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_invocation_log_record(row: &PgRow) -> Result<InvocationLogRecord, RegistryError> {
|
||||
Ok(InvocationLogRecord {
|
||||
log: InvocationLog {
|
||||
@@ -627,22 +576,6 @@ fn target_summary(target: &Target) -> (String, String) {
|
||||
}
|
||||
.to_owned(),
|
||||
),
|
||||
Target::Graphql(graphql) => (
|
||||
graphql.endpoint.clone(),
|
||||
match graphql.operation_type {
|
||||
GraphqlOperationType::Query => "QUERY",
|
||||
GraphqlOperationType::Mutation => "MUTATION",
|
||||
}
|
||||
.to_owned(),
|
||||
),
|
||||
Target::Grpc(grpc) => (grpc.server_addr.clone(), grpc.method.clone()),
|
||||
Target::Websocket(websocket) => (websocket.url.clone(), "SUBSCRIBE".to_owned()),
|
||||
Target::Soap(soap) => (
|
||||
soap.endpoint_override
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("wsdl://{}", soap.service_name)),
|
||||
soap.operation_name.clone(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1061,59 +994,6 @@ fn from_db_version(value: i32, field: &'static str) -> Result<u32, RegistryError
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_stream_session_transition(
|
||||
session_id: &StreamSessionId,
|
||||
from: StreamStatus,
|
||||
to: StreamStatus,
|
||||
) -> Result<(), RegistryError> {
|
||||
let allowed = matches!(
|
||||
(from, to),
|
||||
(StreamStatus::Created, StreamStatus::Running)
|
||||
| (StreamStatus::Running, StreamStatus::Running)
|
||||
| (StreamStatus::Running, StreamStatus::Stopped)
|
||||
| (StreamStatus::Running, StreamStatus::Expired)
|
||||
| (StreamStatus::Running, StreamStatus::Failed)
|
||||
);
|
||||
|
||||
if allowed {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(RegistryError::InvalidStreamSessionTransition {
|
||||
session_id: session_id.as_str().to_owned(),
|
||||
from: serialize_enum_text(&from, "status").unwrap_or_else(|_| "unknown".to_owned()),
|
||||
to: serialize_enum_text(&to, "status").unwrap_or_else(|_| "unknown".to_owned()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_async_job_transition(
|
||||
job_id: &AsyncJobId,
|
||||
from: JobStatus,
|
||||
to: JobStatus,
|
||||
) -> Result<(), RegistryError> {
|
||||
let allowed = matches!(
|
||||
(from, to),
|
||||
(JobStatus::Created, JobStatus::Running)
|
||||
| (JobStatus::Running, JobStatus::Running)
|
||||
| (JobStatus::Running, JobStatus::Completed)
|
||||
| (JobStatus::Running, JobStatus::Failed)
|
||||
| (JobStatus::Running, JobStatus::Cancelled)
|
||||
| (JobStatus::Completed, JobStatus::Expired)
|
||||
| (JobStatus::Failed, JobStatus::Expired)
|
||||
| (JobStatus::Cancelled, JobStatus::Expired)
|
||||
);
|
||||
|
||||
if allowed {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(RegistryError::InvalidAsyncJobTransition {
|
||||
job_id: job_id.as_str().to_owned(),
|
||||
from: serialize_enum_text(&from, "status").unwrap_or_else(|_| "unknown".to_owned()),
|
||||
to: serialize_enum_text(&to, "status").unwrap_or_else(|_| "unknown".to_owned()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn to_u64(value: i64, field: &'static str) -> Result<u64, RegistryError> {
|
||||
u64::try_from(value).map_err(|_| RegistryError::InvalidNumericValue { field, value })
|
||||
}
|
||||
@@ -1137,13 +1017,12 @@ mod tests {
|
||||
|
||||
use crank_core::{
|
||||
AgentId, AgentOperationBinding, AgentStatus, AgentVersion, ApiKeyHeaderAuthConfig,
|
||||
AsyncJobHandle, AuthConfig, AuthKind, AuthProfile, ConfigExport, ExecutionConfig,
|
||||
ExportMode, GeneratedDraft, GeneratedDraftStatus, HttpMethod, InvocationLog, JobStatus,
|
||||
MembershipRole, OperationId, OperationSecurityLevel, OperationStatus, PlatformApiKey,
|
||||
PlatformApiKeyId, PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, RestTarget,
|
||||
RetryPolicy, Samples, SecretId, StreamSession, StreamSessionId, StreamStatus, Target,
|
||||
ToolDescription, ToolExample, User, UserId, UserSessionId, WizardState, Workspace,
|
||||
WorkspaceId,
|
||||
AuthConfig, AuthKind, AuthProfile, ConfigExport, ExecutionConfig, ExportMode,
|
||||
GeneratedDraft, GeneratedDraftStatus, HttpMethod, InvocationLog, MembershipRole,
|
||||
OperationId, OperationSecurityLevel, OperationStatus, PlatformApiKey, PlatformApiKeyId,
|
||||
PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, RestTarget, RetryPolicy, Samples,
|
||||
SecretId, Target, ToolDescription, ToolExample, User, UserId, UserSessionId, WizardState,
|
||||
Workspace, WorkspaceId,
|
||||
};
|
||||
use crank_mapping::{MappingRule, MappingSet};
|
||||
use crank_schema::{Schema, SchemaKind};
|
||||
@@ -1154,14 +1033,12 @@ mod tests {
|
||||
use crate::{
|
||||
PostgresRegistry, RegistryError,
|
||||
model::{
|
||||
AsyncJobFilter, CreateAgentRequest, CreateAsyncJobRequest, CreateInvocationLogRequest,
|
||||
CreatePlatformApiKeyRequest, CreateStreamSessionRequest, CreateVersionRequest,
|
||||
CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind, DescriptorMetadata,
|
||||
OperationSampleMetadata, PlatformApiKeyRecord, PublishAgentRequest, PublishRequest,
|
||||
RegistryOperation, SampleKind, SaveAuthProfileRequest, SaveDescriptorMetadataRequest,
|
||||
SaveSampleMetadataRequest, StreamSessionFilter, UpdateAsyncJobStatusRequest,
|
||||
UpdateStreamSessionStateRequest, WorkspaceRecord, YamlImportJobCompletion,
|
||||
YamlImportJobId, YamlImportJobStatus,
|
||||
CreateAgentRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
|
||||
CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest,
|
||||
DescriptorKind, DescriptorMetadata, OperationSampleMetadata, PlatformApiKeyRecord,
|
||||
PublishAgentRequest, PublishRequest, RegistryOperation, SampleKind,
|
||||
SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
|
||||
WorkspaceRecord, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
|
||||
},
|
||||
};
|
||||
|
||||
@@ -2032,284 +1909,6 @@ mod tests {
|
||||
|
||||
database.cleanup().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manages_stream_sessions_with_transitions_and_cleanup() {
|
||||
let database = TestDatabase::new().await;
|
||||
let registry = database.registry().await;
|
||||
let operation = test_operation("op_stream_01", 1, OperationStatus::Draft);
|
||||
let session = test_stream_session("stream_01", "op_stream_01", StreamStatus::Running);
|
||||
|
||||
registry
|
||||
.create_operation(&test_workspace_id(), &operation, None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
registry
|
||||
.create_stream_session(CreateStreamSessionRequest { session: &session })
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let loaded = registry
|
||||
.get_stream_session(&session.id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(loaded.status, StreamStatus::Running);
|
||||
|
||||
let closable_session =
|
||||
test_stream_session("stream_closable", "op_stream_01", StreamStatus::Running);
|
||||
registry
|
||||
.create_stream_session(CreateStreamSessionRequest {
|
||||
session: &closable_session,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.close_stream_session(&closable_session.id, ×tamp("2026-04-06T12:00:30Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let closed = registry
|
||||
.get_stream_session(&closable_session.id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(closed.status, StreamStatus::Stopped);
|
||||
|
||||
let touched = registry
|
||||
.touch_stream_session_poll(&session.id, ×tamp("2026-04-06T12:00:20Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
touched.last_poll_at,
|
||||
Some(timestamp("2026-04-06T12:00:20Z"))
|
||||
);
|
||||
|
||||
let updated = registry
|
||||
.update_stream_session_state(UpdateStreamSessionStateRequest {
|
||||
session_id: &session.id,
|
||||
current_status: StreamStatus::Running,
|
||||
next_status: StreamStatus::Failed,
|
||||
cursor: Some(&json!({"cursor":"next"})),
|
||||
state: &json!({"phase":"failed"}),
|
||||
expires_at: Some(×tamp("2026-04-06T12:10:00Z")),
|
||||
last_poll_at: Some(×tamp("2026-04-06T12:01:00Z")),
|
||||
closed_at: Some(×tamp("2026-04-06T12:01:00Z")),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(updated.status, StreamStatus::Failed);
|
||||
assert_eq!(updated.closed_at, Some(timestamp("2026-04-06T12:01:00Z")));
|
||||
|
||||
let page = registry
|
||||
.list_stream_sessions(StreamSessionFilter {
|
||||
workspace_id: &test_workspace_id(),
|
||||
agent_id: None,
|
||||
operation_id: None,
|
||||
status: Some(StreamStatus::Failed),
|
||||
mode: Some(crank_core::ExecutionMode::Session),
|
||||
limit: 10,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.items[0].id, session.id);
|
||||
|
||||
let expired_session =
|
||||
test_stream_session("stream_expired", "op_stream_01", StreamStatus::Running);
|
||||
registry
|
||||
.create_stream_session(CreateStreamSessionRequest {
|
||||
session: &StreamSession {
|
||||
id: expired_session.id.clone(),
|
||||
expires_at: timestamp("2026-04-06T11:00:00Z"),
|
||||
..expired_session
|
||||
},
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let deleted = registry
|
||||
.delete_expired_stream_sessions(×tamp("2026-04-06T12:00:00Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(deleted, 1);
|
||||
|
||||
database.cleanup().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_invalid_stream_session_transition() {
|
||||
let database = TestDatabase::new().await;
|
||||
let registry = database.registry().await;
|
||||
let operation = test_operation("op_stream_02", 1, OperationStatus::Draft);
|
||||
let session = test_stream_session("stream_invalid", "op_stream_02", StreamStatus::Stopped);
|
||||
|
||||
registry
|
||||
.create_operation(&test_workspace_id(), &operation, None)
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.create_stream_session(CreateStreamSessionRequest { session: &session })
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = registry
|
||||
.update_stream_session_state(UpdateStreamSessionStateRequest {
|
||||
session_id: &session.id,
|
||||
current_status: StreamStatus::Stopped,
|
||||
next_status: StreamStatus::Running,
|
||||
cursor: None,
|
||||
state: &json!({"phase":"resume"}),
|
||||
expires_at: None,
|
||||
last_poll_at: None,
|
||||
closed_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
RegistryError::InvalidStreamSessionTransition { .. }
|
||||
));
|
||||
|
||||
database.cleanup().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manages_async_jobs_with_transitions_and_cleanup() {
|
||||
let database = TestDatabase::new().await;
|
||||
let registry = database.registry().await;
|
||||
let operation = test_operation("op_job_01", 1, OperationStatus::Draft);
|
||||
let job = test_async_job("job_01", "op_job_01", JobStatus::Running);
|
||||
|
||||
registry
|
||||
.create_operation(&test_workspace_id(), &operation, None)
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.create_async_job(CreateAsyncJobRequest { job: &job })
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cancellable_job = test_async_job("job_cancel", "op_job_01", JobStatus::Running);
|
||||
registry
|
||||
.create_async_job(CreateAsyncJobRequest {
|
||||
job: &cancellable_job,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.cancel_async_job(&cancellable_job.id, ×tamp("2026-04-06T12:00:30Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let cancelled = registry
|
||||
.get_async_job(&cancellable_job.id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(cancelled.status, JobStatus::Cancelled);
|
||||
|
||||
let updated = registry
|
||||
.update_async_job_status(UpdateAsyncJobStatusRequest {
|
||||
job_id: &job.id,
|
||||
current_status: JobStatus::Running,
|
||||
next_status: JobStatus::Completed,
|
||||
progress: &json!({"percent":100}),
|
||||
result: Some(&json!({"ok":true})),
|
||||
error: None,
|
||||
expires_at: Some(×tamp("2026-04-06T12:15:00Z")),
|
||||
updated_at: ×tamp("2026-04-06T12:01:00Z"),
|
||||
finished_at: Some(×tamp("2026-04-06T12:01:00Z")),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(updated.status, JobStatus::Completed);
|
||||
assert_eq!(updated.result, Some(json!({"ok":true})));
|
||||
|
||||
let loaded = registry.get_async_job(&job.id).await.unwrap().unwrap();
|
||||
assert_eq!(loaded.status, JobStatus::Completed);
|
||||
|
||||
let touched = registry
|
||||
.touch_async_job_poll(&job.id, ×tamp("2026-04-06T12:02:00Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
touched.last_poll_at,
|
||||
Some(timestamp("2026-04-06T12:02:00Z"))
|
||||
);
|
||||
|
||||
let page = registry
|
||||
.list_async_jobs(AsyncJobFilter {
|
||||
workspace_id: &test_workspace_id(),
|
||||
agent_id: None,
|
||||
operation_id: None,
|
||||
status: Some(JobStatus::Completed),
|
||||
limit: 10,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.total, 1);
|
||||
|
||||
let expired_job = AsyncJobHandle {
|
||||
id: crank_core::AsyncJobId::new("job_expired"),
|
||||
expires_at: Some(timestamp("2026-04-06T11:00:00Z")),
|
||||
..test_async_job("job_expired", "op_job_01", JobStatus::Cancelled)
|
||||
};
|
||||
registry
|
||||
.create_async_job(CreateAsyncJobRequest { job: &expired_job })
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let deleted = registry
|
||||
.delete_expired_async_jobs(×tamp("2026-04-06T12:00:00Z"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(deleted, 1);
|
||||
|
||||
database.cleanup().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_invalid_async_job_transition() {
|
||||
let database = TestDatabase::new().await;
|
||||
let registry = database.registry().await;
|
||||
let operation = test_operation("op_job_02", 1, OperationStatus::Draft);
|
||||
let job = test_async_job("job_invalid", "op_job_02", JobStatus::Completed);
|
||||
|
||||
registry
|
||||
.create_operation(&test_workspace_id(), &operation, None)
|
||||
.await
|
||||
.unwrap();
|
||||
registry
|
||||
.create_async_job(CreateAsyncJobRequest { job: &job })
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = registry
|
||||
.update_async_job_status(UpdateAsyncJobStatusRequest {
|
||||
job_id: &job.id,
|
||||
current_status: JobStatus::Completed,
|
||||
next_status: JobStatus::Running,
|
||||
progress: &json!({"percent":50}),
|
||||
result: None,
|
||||
error: None,
|
||||
expires_at: None,
|
||||
updated_at: ×tamp("2026-04-06T12:01:00Z"),
|
||||
finished_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
RegistryError::InvalidAsyncJobTransition { .. }
|
||||
));
|
||||
|
||||
database.cleanup().await;
|
||||
}
|
||||
|
||||
struct TestDatabase {
|
||||
admin_pool: PgPool,
|
||||
database_url: String,
|
||||
@@ -2454,9 +2053,7 @@ mod tests {
|
||||
retry_policy: Some(RetryPolicy { max_attempts: 3 }),
|
||||
auth_profile_ref: Some("auth_crank".into()),
|
||||
headers: BTreeMap::from([("X-Request-Id".to_owned(), "static".to_owned())]),
|
||||
protocol_options: None,
|
||||
response_cache: None,
|
||||
streaming: None,
|
||||
},
|
||||
tool_description: ToolDescription {
|
||||
title: "Create lead".to_owned(),
|
||||
@@ -2469,8 +2066,6 @@ mod tests {
|
||||
samples: Some(Samples {
|
||||
input_json_sample_ref: Some("sample_input".into()),
|
||||
output_json_sample_ref: Some("sample_output".into()),
|
||||
proto_file_ref: None,
|
||||
descriptor_ref: None,
|
||||
}),
|
||||
generated_draft: Some(GeneratedDraft {
|
||||
status: GeneratedDraftStatus::Available,
|
||||
@@ -2496,25 +2091,6 @@ mod tests {
|
||||
published_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_stream_session(id: &str, operation_id: &str, status: StreamStatus) -> StreamSession {
|
||||
StreamSession {
|
||||
id: StreamSessionId::new(id),
|
||||
workspace_id: test_workspace_id(),
|
||||
agent_id: None,
|
||||
operation_id: OperationId::new(operation_id),
|
||||
protocol: Protocol::Rest,
|
||||
mode: crank_core::ExecutionMode::Session,
|
||||
status,
|
||||
cursor: Some(json!({"cursor":"initial"})),
|
||||
state: json!({"phase":"running"}),
|
||||
expires_at: timestamp("2026-04-06T12:05:00Z"),
|
||||
last_poll_at: None,
|
||||
created_at: timestamp("2026-04-06T12:00:00Z"),
|
||||
closed_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn test_agent(id: &str, status: AgentStatus) -> crank_core::Agent {
|
||||
crank_core::Agent {
|
||||
id: AgentId::new(id),
|
||||
@@ -2575,22 +2151,4 @@ mod tests {
|
||||
created_at: timestamp(created_at),
|
||||
}
|
||||
}
|
||||
|
||||
fn test_async_job(id: &str, operation_id: &str, status: JobStatus) -> AsyncJobHandle {
|
||||
AsyncJobHandle {
|
||||
id: crank_core::AsyncJobId::new(id),
|
||||
workspace_id: test_workspace_id(),
|
||||
agent_id: None,
|
||||
operation_id: OperationId::new(operation_id),
|
||||
status,
|
||||
progress: json!({"percent": 25}),
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,634 +0,0 @@
|
||||
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())
|
||||
}
|
||||
}
|
||||
@@ -1,292 +0,0 @@
|
||||
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..."));
|
||||
}
|
||||
}
|
||||
@@ -10,22 +10,12 @@ pub enum RuntimeError {
|
||||
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,
|
||||
@@ -35,8 +25,6 @@ pub enum RuntimeError {
|
||||
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")]
|
||||
|
||||
@@ -6,7 +6,6 @@ use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use crank_core::{
|
||||
AdapterRegistry, CachedHeader, CachedResponse, ExecutionMode, HttpMethod, InvocationStatus,
|
||||
MeteringEvent, ResponseCacheStore, SharedMeteringSink, SharedProtocolAdapter, Target,
|
||||
TransportBehavior,
|
||||
};
|
||||
use serde_json::{Map, Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
@@ -16,7 +15,7 @@ use tracing::debug;
|
||||
|
||||
use crate::{
|
||||
AdapterResponse, PreparedRequest, ResolvedAuth, RuntimeError, RuntimeLimits, RuntimeOperation,
|
||||
RuntimeRequestContext, WindowExecutionResult,
|
||||
RuntimeRequestContext,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -24,9 +23,6 @@ pub struct RuntimeExecutor {
|
||||
adapters: AdapterRegistry,
|
||||
limits: RuntimeLimits,
|
||||
unary_limiter: Arc<Semaphore>,
|
||||
window_limiter: Arc<Semaphore>,
|
||||
session_limiter: Arc<Semaphore>,
|
||||
job_limiter: Arc<Semaphore>,
|
||||
response_cache: Option<Arc<dyn ResponseCacheStore>>,
|
||||
metering_sink: SharedMeteringSink,
|
||||
}
|
||||
@@ -57,9 +53,6 @@ impl RuntimeExecutor {
|
||||
Self {
|
||||
adapters,
|
||||
unary_limiter: Arc::new(Semaphore::new(limits.max_concurrent_unary)),
|
||||
window_limiter: Arc::new(Semaphore::new(limits.max_concurrent_window)),
|
||||
session_limiter: Arc::new(Semaphore::new(limits.max_concurrent_sessions)),
|
||||
job_limiter: Arc::new(Semaphore::new(limits.max_concurrent_jobs)),
|
||||
limits,
|
||||
response_cache,
|
||||
metering_sink,
|
||||
@@ -123,156 +116,6 @@ impl RuntimeExecutor {
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn execute_window(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
self.execute_window_with_auth_and_context(operation, input, None, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_window_with_auth(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
resolved_auth: Option<&ResolvedAuth>,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
self.execute_window_with_auth_and_context(operation, input, resolved_auth, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_window_with_context(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
request_context: Option<&RuntimeRequestContext>,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
self.execute_window_with_auth_and_context(operation, input, None, request_context)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_window_with_auth_and_context(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
resolved_auth: Option<&ResolvedAuth>,
|
||||
request_context: Option<&RuntimeRequestContext>,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
log_runtime_event("window.execute", operation, request_context);
|
||||
let _permit = self.acquire_window_permit()?;
|
||||
let started_at = Instant::now();
|
||||
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
|
||||
let result = Err(RuntimeError::MissingStreamingConfig {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
});
|
||||
self.record_metering(operation, request_context, &result, started_at)
|
||||
.await;
|
||||
return result;
|
||||
};
|
||||
|
||||
if streaming.mode != ExecutionMode::Window {
|
||||
let result = Err(RuntimeError::UnsupportedExecutionMode {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
mode: streaming.mode,
|
||||
});
|
||||
self.record_metering(operation, request_context, &result, started_at)
|
||||
.await;
|
||||
return result;
|
||||
}
|
||||
|
||||
let prepared_request = self.prepare_request(operation, input)?;
|
||||
let prepared_request = apply_resolved_auth(prepared_request, resolved_auth);
|
||||
let adapter_response = if matches!(
|
||||
streaming.transport_behavior,
|
||||
TransportBehavior::ServerStream
|
||||
) {
|
||||
self.execute_window_adapter(operation, prepared_request, request_context)
|
||||
.await?
|
||||
} else {
|
||||
self.execute_adapter(operation, prepared_request, request_context)
|
||||
.await?
|
||||
};
|
||||
|
||||
let result = crate::aggregation::collect_window_result(&adapter_response.body, streaming);
|
||||
self.record_metering(operation, request_context, &result, started_at)
|
||||
.await;
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn execute_session_seed(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
self.execute_session_seed_with_auth_and_context(operation, input, None, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_session_seed_with_auth(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
resolved_auth: Option<&ResolvedAuth>,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
self.execute_session_seed_with_auth_and_context(operation, input, resolved_auth, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_session_seed_with_context(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
request_context: Option<&RuntimeRequestContext>,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
self.execute_session_seed_with_auth_and_context(operation, input, None, request_context)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn execute_session_seed_with_auth_and_context(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
input: &Value,
|
||||
resolved_auth: Option<&ResolvedAuth>,
|
||||
request_context: Option<&RuntimeRequestContext>,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
log_runtime_event("session.seed", operation, request_context);
|
||||
let _permit = self.acquire_session_permit()?;
|
||||
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
|
||||
return Err(RuntimeError::MissingStreamingConfig {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
});
|
||||
};
|
||||
|
||||
if streaming.mode != ExecutionMode::Session {
|
||||
return Err(RuntimeError::UnsupportedExecutionMode {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
mode: streaming.mode,
|
||||
});
|
||||
}
|
||||
|
||||
let batch_size = streaming.max_items.unwrap_or(10).max(1);
|
||||
let seed_limit = batch_size.saturating_mul(4);
|
||||
let mut seeded_operation = operation.clone();
|
||||
if let Some(config) = seeded_operation.execution_config.streaming.as_mut() {
|
||||
config.mode = ExecutionMode::Window;
|
||||
config.window_duration_ms = config
|
||||
.window_duration_ms
|
||||
.or(config.poll_interval_ms)
|
||||
.or(config.upstream_timeout_ms)
|
||||
.or(Some(operation.execution_config.timeout_ms));
|
||||
config.max_items = Some(seed_limit);
|
||||
}
|
||||
|
||||
self.execute_window_with_auth_and_context(
|
||||
&seeded_operation,
|
||||
input,
|
||||
resolved_auth,
|
||||
request_context,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn prepare_request(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
@@ -379,98 +222,14 @@ impl RuntimeExecutor {
|
||||
.map_err(|error| map_protocol_adapter_error(operation, error))
|
||||
}
|
||||
|
||||
async fn execute_window_adapter(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
prepared_request: PreparedRequest,
|
||||
request_context: Option<&RuntimeRequestContext>,
|
||||
) -> Result<AdapterResponse, RuntimeError> {
|
||||
log_runtime_event("window.adapter.dispatch", operation, request_context);
|
||||
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
|
||||
return Err(RuntimeError::MissingStreamingConfig {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
});
|
||||
};
|
||||
let adapter = self.adapter_for(operation)?;
|
||||
if !adapter.supports_mode(ExecutionMode::Window) {
|
||||
return Err(RuntimeError::UnsupportedExecutionMode {
|
||||
operation_id: operation.operation_id.as_str().to_owned(),
|
||||
mode: ExecutionMode::Window,
|
||||
});
|
||||
}
|
||||
|
||||
let prepared_request = adapter_prepared_request(
|
||||
operation,
|
||||
&prepared_request,
|
||||
request_context,
|
||||
streaming
|
||||
.upstream_timeout_ms
|
||||
.unwrap_or(operation.execution_config.timeout_ms),
|
||||
);
|
||||
let adapter_context = adapter_request_context(request_context);
|
||||
let result = adapter
|
||||
.invoke_window(
|
||||
&operation.target,
|
||||
&prepared_request.into(),
|
||||
streaming.window_duration_ms.unwrap_or_default(),
|
||||
streaming.max_items,
|
||||
&adapter_context,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| map_protocol_adapter_error(operation, error))?;
|
||||
|
||||
Ok(AdapterResponse {
|
||||
status_code: 200,
|
||||
headers: BTreeMap::new(),
|
||||
body: json!({
|
||||
"summary": result.summary,
|
||||
"items": result.items,
|
||||
"cursor": result.cursor,
|
||||
"done": result.window_complete,
|
||||
}),
|
||||
data: Value::Null,
|
||||
})
|
||||
}
|
||||
|
||||
fn acquire_unary_permit(
|
||||
&self,
|
||||
operation: &RuntimeOperation,
|
||||
_operation: &RuntimeOperation,
|
||||
) -> Result<OwnedSemaphorePermit, RuntimeError> {
|
||||
let (kind, limit, limiter) = if operation
|
||||
.execution_config
|
||||
.streaming
|
||||
.as_ref()
|
||||
.is_some_and(|streaming| streaming.mode == ExecutionMode::AsyncJob)
|
||||
{
|
||||
(
|
||||
"async_job",
|
||||
self.limits.max_concurrent_jobs,
|
||||
Arc::clone(&self.job_limiter),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
"unary",
|
||||
self.limits.max_concurrent_unary,
|
||||
Arc::clone(&self.unary_limiter),
|
||||
)
|
||||
};
|
||||
|
||||
try_acquire_limit(limiter, kind, limit)
|
||||
}
|
||||
|
||||
fn acquire_window_permit(&self) -> Result<OwnedSemaphorePermit, RuntimeError> {
|
||||
try_acquire_limit(
|
||||
Arc::clone(&self.window_limiter),
|
||||
"window",
|
||||
self.limits.max_concurrent_window,
|
||||
)
|
||||
}
|
||||
|
||||
fn acquire_session_permit(&self) -> Result<OwnedSemaphorePermit, RuntimeError> {
|
||||
try_acquire_limit(
|
||||
Arc::clone(&self.session_limiter),
|
||||
"session",
|
||||
self.limits.max_concurrent_sessions,
|
||||
Arc::clone(&self.unary_limiter),
|
||||
"unary",
|
||||
self.limits.max_concurrent_unary,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -564,8 +323,6 @@ impl PreparedRequest {
|
||||
path_params: read_string_map(request.get("path"), "request.path")?,
|
||||
query_params: read_string_map(request.get("query"), "request.query")?,
|
||||
headers: read_string_map(request.get("headers"), "request.headers")?,
|
||||
grpc: non_empty_payload(request.get("grpc").cloned()),
|
||||
variables: non_empty_payload(request.get("variables").cloned()),
|
||||
body: request.get("body").and_then(non_empty_body).cloned(),
|
||||
timeout_ms: 0,
|
||||
})
|
||||
@@ -586,10 +343,8 @@ fn adapter_prepared_request(
|
||||
request_context: Option<&RuntimeRequestContext>,
|
||||
timeout_ms: u64,
|
||||
) -> PreparedRequest {
|
||||
let empty_headers = BTreeMap::new();
|
||||
let static_headers = match &operation.target {
|
||||
Target::Rest(target) => &target.static_headers,
|
||||
_ => &empty_headers,
|
||||
};
|
||||
|
||||
let mut prepared_request = prepared_request.clone();
|
||||
@@ -717,10 +472,6 @@ fn log_runtime_event(
|
||||
}
|
||||
|
||||
fn response_cache_ttl(operation: &RuntimeOperation) -> Option<Duration> {
|
||||
if operation.execution_config.streaming.is_some() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let is_cacheable_protocol =
|
||||
matches!(&operation.target, Target::Rest(target) if target.method == HttpMethod::Get);
|
||||
if !is_cacheable_protocol {
|
||||
@@ -747,8 +498,6 @@ fn response_cache_key(
|
||||
"query_params": prepared_request.query_params,
|
||||
"headers": prepared_request.headers,
|
||||
"body": prepared_request.body,
|
||||
"variables": prepared_request.variables,
|
||||
"grpc": prepared_request.grpc,
|
||||
});
|
||||
let fingerprint_bytes = serde_json::to_vec(&fingerprint).ok()?;
|
||||
let fingerprint_hash = URL_SAFE_NO_PAD.encode(Sha256::digest(fingerprint_bytes));
|
||||
@@ -844,15 +593,7 @@ fn non_empty_body(value: &Value) -> Option<&Value> {
|
||||
}
|
||||
}
|
||||
|
||||
fn non_empty_payload(value: Option<Value>) -> Option<Value> {
|
||||
match value {
|
||||
None | Some(Value::Null) => None,
|
||||
Some(Value::Object(object)) if object.is_empty() => None,
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[cfg(any())]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::io;
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
mod aggregation;
|
||||
mod auth;
|
||||
mod cache;
|
||||
mod cache_factory;
|
||||
@@ -8,10 +7,8 @@ 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::{
|
||||
@@ -32,4 +29,3 @@ pub use rate_limit::{
|
||||
};
|
||||
pub use request_context::{MeteringContext, ResponseCacheScope, RuntimeRequestContext};
|
||||
pub use secret_crypto::SecretCrypto;
|
||||
pub use streaming::WindowExecutionResult;
|
||||
|
||||
@@ -3,25 +3,16 @@ 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -33,18 +24,6 @@ impl RuntimeLimits {
|
||||
"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,
|
||||
)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -92,9 +71,6 @@ mod tests {
|
||||
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]
|
||||
|
||||
@@ -48,11 +48,7 @@ pub struct PreparedRequest {
|
||||
#[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>,
|
||||
pub body: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub timeout_ms: u64,
|
||||
}
|
||||
@@ -72,8 +68,6 @@ impl From<PreparedRequest> for crank_core::PreparedRequest {
|
||||
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,
|
||||
}
|
||||
@@ -86,8 +80,6 @@ impl From<crank_core::PreparedRequest> for PreparedRequest {
|
||||
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,
|
||||
}
|
||||
|
||||
@@ -1,124 +0,0 @@
|
||||
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..."));
|
||||
}
|
||||
}
|
||||
@@ -1,26 +0,0 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct WindowExecutionResult {
|
||||
pub summary: Value,
|
||||
pub items: Vec<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cursor: Option<Value>,
|
||||
pub window_complete: bool,
|
||||
pub truncated: bool,
|
||||
pub has_more: bool,
|
||||
}
|
||||
|
||||
impl From<crank_core::WindowExecutionResult> for WindowExecutionResult {
|
||||
fn from(value: crank_core::WindowExecutionResult) -> Self {
|
||||
Self {
|
||||
summary: value.summary,
|
||||
items: value.items,
|
||||
cursor: value.cursor,
|
||||
window_complete: value.window_complete,
|
||||
truncated: value.truncated,
|
||||
has_more: value.has_more,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user