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, #[serde(skip_serializing_if = "Option::is_none")] pub poll_tool_name: Option, #[serde(skip_serializing_if = "Option::is_none")] pub stop_tool_name: Option, #[serde(skip_serializing_if = "Option::is_none")] pub status_tool_name: Option, #[serde(skip_serializing_if = "Option::is_none")] pub result_tool_name: Option, #[serde(skip_serializing_if = "Option::is_none")] pub cancel_tool_name: Option, } 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, #[serde(skip_serializing_if = "Option::is_none")] pub poll_interval_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub upstream_timeout_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub idle_timeout_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub max_session_lifetime_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub max_items: Option, #[serde(skip_serializing_if = "Option::is_none")] pub max_bytes: Option, pub aggregation_mode: AggregationMode, #[serde(skip_serializing_if = "Option::is_none")] pub summary_path: Option, #[serde(skip_serializing_if = "Option::is_none")] pub items_path: Option, #[serde(skip_serializing_if = "Option::is_none")] pub cursor_path: Option, #[serde(skip_serializing_if = "Option::is_none")] pub status_path: Option, #[serde(skip_serializing_if = "Option::is_none")] pub done_path: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub redacted_paths: Vec, #[serde(default)] pub truncate_item_fields: bool, #[serde(skip_serializing_if = "Option::is_none")] pub max_field_length: Option, #[serde(default)] pub drop_duplicates: bool, #[serde(skip_serializing_if = "Option::is_none")] pub sampling_rate: Option, #[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()); } }