bsodfather/crank
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add websocket upstream adapter

Browse Source
This commit is contained in:
a.tolmachev
2026-04-06 13:23:45 +03:00
parent b5f80c5d2f
commit 45ea011b7f
27 changed files with 978 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/crank-runtime/src/error.rs
+3
View File
@@ -1,6 +1,7 @@
use crank_adapter_graphql::GraphqlAdapterError;
use crank_adapter_grpc::GrpcAdapterError;
use crank_adapter_rest::RestAdapterError;
use crank_adapter_websocket::WebsocketAdapterError;
use crank_core::{ExecutionMode, Protocol};
use crank_mapping::MappingError;
use crank_schema::SchemaError;
@@ -18,6 +19,8 @@ pub enum RuntimeError {
GrpcAdapter(#[from] GrpcAdapterError),
#[error(transparent)]
RestAdapter(#[from] RestAdapterError),
#[error(transparent)]
WebsocketAdapter(#[from] WebsocketAdapterError),
#[error("protocol {protocol:?} is not supported by runtime")]
UnsupportedProtocol { protocol: Protocol },
#[error("operation {operation_id} does not define streaming config")]
crates/crank-runtime/src/executor.rs
+211 -5
View File
@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
use crank_adapter_graphql::{GraphqlAdapter, GraphqlRequest};
use crank_adapter_grpc::{GrpcAdapter, GrpcRequest, GrpcWindowRequest};
use crank_adapter_rest::{RestAdapter, RestRequest, RestWindowRequest};
use crank_adapter_websocket::{WebsocketAdapter, WebsocketWindowRequest};
use crank_core::{ExecutionMode, Target, TransportBehavior};
use serde_json::{Map, Value, json};
@@ -15,6 +16,7 @@ pub struct RuntimeExecutor {
graphql_adapter: GraphqlAdapter,
grpc_adapter: GrpcAdapter,
rest_adapter: RestAdapter,
websocket_adapter: WebsocketAdapter,
}
impl Default for RuntimeExecutor {
@@ -29,6 +31,7 @@ impl RuntimeExecutor {
graphql_adapter: GraphqlAdapter::new(),
grpc_adapter: GrpcAdapter::new(),
rest_adapter: RestAdapter::new(),
websocket_adapter: WebsocketAdapter::new(),
}
}
@@ -63,8 +66,10 @@ impl RuntimeExecutor {
let adapter_response = if matches!(
streaming.transport_behavior,
TransportBehavior::ServerStream
) && matches!(operation.target, Target::Rest(_) | Target::Grpc(_))
{
) && matches!(
operation.target,
Target::Rest(_) | Target::Grpc(_) | Target::Websocket(_)
) {
self.execute_window_adapter(operation, prepared_request)
.await?
} else {
@@ -200,6 +205,10 @@ impl RuntimeExecutor {
data: Value::Null,
})
}
Target::Websocket(_) => Err(RuntimeError::UnsupportedExecutionMode {
operation_id: operation.operation_id.as_str().to_owned(),
mode: ExecutionMode::Unary,
}),
}
}
@@ -274,6 +283,46 @@ impl RuntimeExecutor {
data: Value::Null,
})
}
Target::Websocket(target) => {
let Some(streaming) = operation.execution_config.streaming.as_ref() else {
return Err(RuntimeError::MissingStreamingConfig {
operation_id: operation.operation_id.as_str().to_owned(),
});
};
let websocket_options = operation
.execution_config
.protocol_options
.as_ref()
.and_then(|options| options.websocket.as_ref());
let request = WebsocketWindowRequest {
headers: merge_headers(
&BTreeMap::new(),
&operation.execution_config.headers,
&prepared_request.headers,
),
window_duration_ms: streaming.window_duration_ms.unwrap_or_default(),
max_items: streaming.max_items.map(|value| value.saturating_add(1)),
heartbeat_interval_ms: websocket_options
.and_then(|options| options.heartbeat_interval_ms),
reconnect_max_attempts: websocket_options
.and_then(|options| options.reconnect_max_attempts)
.unwrap_or_default(),
reconnect_backoff_ms: websocket_options
.and_then(|options| options.reconnect_backoff_ms)
.unwrap_or_default(),
};
let response = self
.websocket_adapter
.execute_window(target, &request)
.await?;
Ok(AdapterResponse {
status_code: response.status_code,
headers: response.headers,
body: response.body,
data: Value::Null,
})
}
_ => self.execute_adapter(operation, prepared_request).await,
}
}
@@ -391,14 +440,16 @@ mod tests {
use crank_core::{
AggregationMode, DescriptorId, ExecutionConfig, ExecutionMode, GeneratedDraft,
GeneratedDraftStatus, GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod,
Operation, OperationId, OperationStatus, Protocol, RestTarget, Samples, StreamingConfig,
Target, ToolDescription, ToolExample, ToolFamilyConfig, TransportBehavior,
Operation, OperationId, OperationStatus, Protocol, ProtocolOptions, RestTarget, Samples,
StreamingConfig, Target, ToolDescription, ToolExample, ToolFamilyConfig, TransportBehavior,
WebsocketProtocolOptions, WebsocketTarget,
};
use crank_mapping::{MappingRule, MappingSet};
use crank_schema::{Schema, SchemaKind};
use futures_util::stream;
use futures_util::{SinkExt, StreamExt, stream};
use serde_json::{Value, json};
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
use crate::{RuntimeError, RuntimeExecutor, RuntimeOperation};
@@ -464,6 +515,26 @@ mod tests {
assert!(!result.window_complete);
}
#[tokio::test]
async fn executes_websocket_window_mode_end_to_end() {
let target_url = spawn_websocket_server().await;
let executor = RuntimeExecutor::new();
let operation =
test_websocket_window_operation(&target_url, AggregationMode::RawItems, Some(3));
let result = executor
.execute_window(&operation, &json!({}))
.await
.unwrap();
assert_eq!(result.items.len(), 3);
assert_eq!(result.items[0], json!({ "seq": 1, "value": 101 }));
assert_eq!(result.items[2], json!({ "seq": 3, "value": 103 }));
assert!(result.window_complete);
assert!(!result.truncated);
assert!(!result.has_more);
}
#[tokio::test]
async fn rejects_invalid_input_shape() {
let base_url = spawn_runtime_server().await;
@@ -652,6 +723,43 @@ mod tests {
format!("http://{}", address)
}
async fn spawn_websocket_server() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
tokio::spawn(async move {
loop {
let (stream, _) = listener.accept().await.unwrap();
tokio::spawn(async move {
let websocket = accept_async(stream).await.unwrap();
let (mut sink, mut source) = websocket.split();
if let Some(message) = source.next().await {
let message = message.unwrap();
if !matches!(message, tokio_tungstenite::tungstenite::Message::Text(_)) {
return;
}
}
for payload in [
json!({ "seq": 1, "value": 101 }),
json!({ "seq": 2, "value": 102 }),
json!({ "seq": 3, "value": 103 }),
] {
sink.send(tokio_tungstenite::tungstenite::Message::Text(
payload.to_string().into(),
))
.await
.unwrap();
}
let _ = sink.close().await;
});
}
});
format!("ws://{}", address)
}
async fn create_lead(Json(payload): Json<Value>) -> (axum::http::StatusCode, Json<Value>) {
let should_fail = payload
.get("fail")
@@ -1168,6 +1276,104 @@ mod tests {
operation
}
fn test_websocket_window_operation(
target_url: &str,
aggregation_mode: AggregationMode,
max_items: Option<u32>,
) -> RuntimeOperation {
RuntimeOperation::from(Operation {
id: OperationId::new("op_websocket_window_runtime"),
name: "telemetry_window_ws".to_owned(),
display_name: "Telemetry Window WebSocket".to_owned(),
category: "ops".to_owned(),
protocol: Protocol::Websocket,
status: OperationStatus::Published,
version: 1,
target: Target::Websocket(WebsocketTarget {
url: target_url.to_owned(),
subprotocols: Vec::new(),
subscribe_message_template: Some(json!({"type":"subscribe","topic":"telemetry"})),
unsubscribe_message_template: Some(
json!({"type":"unsubscribe","topic":"telemetry"}),
),
static_headers: BTreeMap::new(),
}),
input_schema: Schema {
kind: SchemaKind::Object,
description: None,
required: true,
nullable: false,
default_value: None,
fields: BTreeMap::new(),
items: None,
enum_values: Vec::new(),
variants: Vec::new(),
},
output_schema: object_schema("ignored", SchemaKind::String),
input_mapping: MappingSet {
rules: vec![MappingRule {
source: "$.mcp".to_owned(),
target: "$.request.body".to_owned(),
required: false,
default_value: None,
transform: None,
condition: None,
notes: None,
}],
},
output_mapping: MappingSet { rules: Vec::new() },
execution_config: ExecutionConfig {
timeout_ms: 1_000,
retry_policy: None,
auth_profile_ref: None,
headers: BTreeMap::new(),
protocol_options: Some(ProtocolOptions {
grpc: None,
websocket: Some(WebsocketProtocolOptions {
heartbeat_interval_ms: Some(100),
reconnect_max_attempts: Some(1),
reconnect_backoff_ms: Some(10),
}),
}),
streaming: Some(StreamingConfig {
mode: ExecutionMode::Window,
transport_behavior: TransportBehavior::ServerStream,
window_duration_ms: Some(2_000),
poll_interval_ms: None,
upstream_timeout_ms: Some(1_000),
idle_timeout_ms: None,
max_session_lifetime_ms: None,
max_items,
max_bytes: None,
aggregation_mode,
summary_path: None,
items_path: Some("$.items".to_owned()),
cursor_path: None,
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(),
}),
},
tool_description: ToolDescription {
title: "Telemetry Window WebSocket".to_owned(),
description: "Collects bounded WebSocket telemetry.".to_owned(),
tags: vec!["websocket".to_owned(), "stream".to_owned()],
examples: Vec::new(),
},
samples: None,
generated_draft: None,
config_export: None,
created_at: "2026-04-06T12:00:00Z".to_owned(),
updated_at: "2026-04-06T12:00:00Z".to_owned(),
published_at: Some("2026-04-06T12:00:00Z".to_owned()),
})
}
fn test_window_sse_operation(
base_url: &str,
aggregation_mode: AggregationMode,