crank/crates/crank-runtime/src/aggregation.rs

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..."));
    }
}