feat: add window execution to runtime
This commit is contained in:
a.tolmachev
@@ -0,0 +1,292 @@
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{RuntimeError, WindowExecutionResult};
|
||||
|
||||
pub fn collect_window_result(
|
||||
response: &Value,
|
||||
config: &crank_core::StreamingConfig,
|
||||
) -> Result<WindowExecutionResult, RuntimeError> {
|
||||
let mut items = extract_path_value(response, config.items_path.as_deref())
|
||||
.and_then(|value| value.as_array().cloned())
|
||||
.unwrap_or_default();
|
||||
let cursor = extract_path_value(response, config.cursor_path.as_deref()).cloned();
|
||||
let done = extract_path_value(response, config.done_path.as_deref())
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(true);
|
||||
|
||||
let mut truncated = false;
|
||||
let mut has_more = !done;
|
||||
|
||||
if let Some(max_items) = config.max_items.map(|value| value as usize) {
|
||||
if items.len() > max_items {
|
||||
items.truncate(max_items);
|
||||
truncated = true;
|
||||
has_more = true;
|
||||
}
|
||||
}
|
||||
|
||||
let mut summary = build_summary(response, &items, config);
|
||||
redact_and_truncate(&mut summary, &mut items, config);
|
||||
|
||||
if let Some(max_bytes) = config.max_bytes.map(|value| value as usize) {
|
||||
while payload_size_bytes(&summary, &items, cursor.as_ref()) > max_bytes && !items.is_empty()
|
||||
{
|
||||
items.pop();
|
||||
truncated = true;
|
||||
has_more = true;
|
||||
}
|
||||
|
||||
if payload_size_bytes(&summary, &items, cursor.as_ref()) > max_bytes {
|
||||
summary = json!({
|
||||
"notice": "window payload exceeded byte limit",
|
||||
"items_returned": items.len()
|
||||
});
|
||||
truncated = true;
|
||||
has_more = true;
|
||||
}
|
||||
}
|
||||
|
||||
if matches!(
|
||||
config.aggregation_mode,
|
||||
crank_core::AggregationMode::SummaryOnly
|
||||
) {
|
||||
items.clear();
|
||||
}
|
||||
|
||||
Ok(WindowExecutionResult {
|
||||
summary,
|
||||
items,
|
||||
cursor,
|
||||
window_complete: !has_more,
|
||||
truncated,
|
||||
has_more,
|
||||
})
|
||||
}
|
||||
|
||||
fn build_summary(response: &Value, items: &[Value], config: &crank_core::StreamingConfig) -> Value {
|
||||
match config.aggregation_mode {
|
||||
crank_core::AggregationMode::RawItems => Value::Null,
|
||||
crank_core::AggregationMode::SummaryOnly
|
||||
| crank_core::AggregationMode::SummaryPlusSamples => {
|
||||
extract_path_value(response, config.summary_path.as_deref())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| stats_summary(items))
|
||||
}
|
||||
crank_core::AggregationMode::Stats => stats_summary(items),
|
||||
crank_core::AggregationMode::LatestState => items.last().cloned().unwrap_or(Value::Null),
|
||||
}
|
||||
}
|
||||
|
||||
fn stats_summary(items: &[Value]) -> Value {
|
||||
let mut summary = Map::new();
|
||||
summary.insert("items_total".to_owned(), Value::from(items.len() as u64));
|
||||
summary.insert("items_sampled".to_owned(), Value::from(items.len() as u64));
|
||||
Value::Object(summary)
|
||||
}
|
||||
|
||||
fn redact_and_truncate(
|
||||
summary: &mut Value,
|
||||
items: &mut [Value],
|
||||
config: &crank_core::StreamingConfig,
|
||||
) {
|
||||
crate::redaction::redact_paths(summary, &config.redacted_paths);
|
||||
|
||||
for item in items.iter_mut() {
|
||||
crate::redaction::redact_paths(item, &config.redacted_paths);
|
||||
}
|
||||
|
||||
if config.truncate_item_fields {
|
||||
if let Some(max_len) = config.max_field_length.map(|value| value as usize) {
|
||||
for item in items.iter_mut() {
|
||||
crate::redaction::truncate_item_fields(item, max_len);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn payload_size_bytes(summary: &Value, items: &[Value], cursor: Option<&Value>) -> usize {
|
||||
serde_json::to_vec(&json!({
|
||||
"summary": summary,
|
||||
"items": items,
|
||||
"cursor": cursor
|
||||
}))
|
||||
.map(|value| value.len())
|
||||
.unwrap_or(usize::MAX)
|
||||
}
|
||||
|
||||
pub fn extract_path_value<'a>(value: &'a Value, path: Option<&str>) -> Option<&'a Value> {
|
||||
let path = path?;
|
||||
let path = path.strip_prefix("$.").or_else(|| path.strip_prefix('$'))?;
|
||||
|
||||
if path.is_empty() {
|
||||
return Some(value);
|
||||
}
|
||||
|
||||
let mut current = value;
|
||||
for segment in path.split('.') {
|
||||
current = current.get(segment)?;
|
||||
}
|
||||
|
||||
Some(current)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crank_core::{
|
||||
AggregationMode, ExecutionMode, StreamingConfig, ToolFamilyConfig, TransportBehavior,
|
||||
};
|
||||
|
||||
use super::collect_window_result;
|
||||
|
||||
fn window_config() -> StreamingConfig {
|
||||
StreamingConfig {
|
||||
mode: ExecutionMode::Window,
|
||||
transport_behavior: TransportBehavior::ServerStream,
|
||||
window_duration_ms: Some(3_000),
|
||||
poll_interval_ms: None,
|
||||
upstream_timeout_ms: Some(5_000),
|
||||
idle_timeout_ms: None,
|
||||
max_session_lifetime_ms: None,
|
||||
max_items: Some(10),
|
||||
max_bytes: Some(10_000),
|
||||
aggregation_mode: AggregationMode::RawItems,
|
||||
summary_path: Some("$.summary".to_owned()),
|
||||
items_path: Some("$.items".to_owned()),
|
||||
cursor_path: Some("$.cursor".to_owned()),
|
||||
status_path: None,
|
||||
done_path: Some("$.done".to_owned()),
|
||||
redacted_paths: Vec::new(),
|
||||
truncate_item_fields: false,
|
||||
max_field_length: None,
|
||||
drop_duplicates: false,
|
||||
sampling_rate: None,
|
||||
tool_family: ToolFamilyConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collects_raw_items_mode() {
|
||||
let config = window_config();
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }, { "message": "two" }],
|
||||
"summary": { "count": 2 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.summary, Value::Null);
|
||||
assert_eq!(result.items.len(), 2);
|
||||
assert!(result.window_complete);
|
||||
assert!(!result.truncated);
|
||||
assert!(!result.has_more);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_summary_only_mode() {
|
||||
let mut config = window_config();
|
||||
config.aggregation_mode = AggregationMode::SummaryOnly;
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }],
|
||||
"summary": { "count": 1 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.summary, json!({ "count": 1 }));
|
||||
assert!(result.items.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_summary_plus_samples() {
|
||||
let mut config = window_config();
|
||||
config.aggregation_mode = AggregationMode::SummaryPlusSamples;
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }, { "message": "two" }],
|
||||
"summary": { "count": 2 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.summary, json!({ "count": 2 }));
|
||||
assert_eq!(result.items.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_by_item_count() {
|
||||
let mut config = window_config();
|
||||
config.max_items = Some(1);
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [{ "message": "one" }, { "message": "two" }],
|
||||
"cursor": "next",
|
||||
"done": false
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.items.len(), 1);
|
||||
assert!(result.truncated);
|
||||
assert!(result.has_more);
|
||||
assert!(!result.window_complete);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_by_byte_limit() {
|
||||
let mut config = window_config();
|
||||
config.max_bytes = Some(120);
|
||||
config.aggregation_mode = AggregationMode::SummaryPlusSamples;
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [
|
||||
{ "message": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" },
|
||||
{ "message": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb" }
|
||||
],
|
||||
"summary": { "count": 2 },
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(result.truncated);
|
||||
assert!(result.items.len() < 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redacts_and_truncates_values() {
|
||||
let mut config = window_config();
|
||||
config.redacted_paths = vec!["$.secret".to_owned()];
|
||||
config.truncate_item_fields = true;
|
||||
config.max_field_length = Some(4);
|
||||
|
||||
let result = collect_window_result(
|
||||
&json!({
|
||||
"items": [
|
||||
{ "message": "abcdefghijklmnop", "secret": "token" }
|
||||
],
|
||||
"done": true
|
||||
}),
|
||||
&config,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.items[0]["secret"], json!("[REDACTED]"));
|
||||
assert_eq!(result.items[0]["message"], json!("abcd..."));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user