crank/crates/crank-core/src/streaming.rs

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());
    }
}