use crate::{capabilities, Event};
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;
pub const PROTOCOL_V0: u32 = 0;
pub const PROTOCOL_V1: u32 = 1;
pub const CURRENT_VERSION: u32 = PROTOCOL_V0;
pub const SUPPORTED_VERSIONS: &[u32] = &[PROTOCOL_V1, PROTOCOL_V0];
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NegotiationResult {
pub version: u32,
pub capabilities: Vec<String>,
pub fallback_reason: Option<String>,
pub service_info: HashMap<String, String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientInfo {
pub supported_versions: Vec<u32>,
pub requested_capabilities: Vec<String>,
pub client_metadata: HashMap<String, String>,
}
pub struct ProtocolNegotiator {
service_id: Uuid,
available_capabilities: Vec<String>,
service_metadata: HashMap<String, String>,
}
impl ProtocolNegotiator {
pub fn new(service_id: Uuid) -> Self {
let mut service_metadata = HashMap::new();
service_metadata.insert(
"service_name".to_string(),
"claude-code-rs-core".to_string(),
);
service_metadata.insert("version".to_string(), env!("CARGO_PKG_VERSION").to_string());
service_metadata.insert(
"build_timestamp".to_string(),
std::env::var("BUILD_TIMESTAMP").unwrap_or_else(|_| "unknown".to_string()),
);
service_metadata.insert(
"git_commit".to_string(),
std::env::var("GIT_COMMIT").unwrap_or_else(|_| "unknown".to_string()),
);
let available_capabilities = vec![
capabilities::SHELL_EXEC.to_string(),
capabilities::REPLAY.to_string(),
capabilities::TRACING.to_string(),
#[cfg(feature = "hooks-quickjs")]
capabilities::HOOKS_JS.to_string(),
#[cfg(feature = "hooks-rust")]
capabilities::HOOKS_RUST.to_string(),
#[cfg(feature = "nats")]
capabilities::NATS.to_string(),
];
#[cfg(feature = "protobuf")]
service_metadata.insert("protobuf_support".to_string(), "true".to_string());
Self {
service_id,
available_capabilities,
service_metadata,
}
}
pub fn negotiate(&self, client_info: ClientInfo) -> Result<NegotiationResult> {
let selected_version = self.select_version(&client_info.supported_versions)?;
let granted_capabilities = self.filter_capabilities(&client_info.requested_capabilities);
let fallback_reason =
self.check_fallback_conditions(selected_version, &granted_capabilities);
let mut service_info = self.service_metadata.clone();
service_info.insert(
"negotiated_version".to_string(),
selected_version.to_string(),
);
service_info.insert(
"granted_capabilities".to_string(),
granted_capabilities.len().to_string(),
);
Ok(NegotiationResult {
version: selected_version,
capabilities: granted_capabilities,
fallback_reason,
service_info,
})
}
fn select_version(&self, client_versions: &[u32]) -> Result<u32> {
for &server_version in SUPPORTED_VERSIONS {
if client_versions.contains(&server_version) {
return Ok(server_version);
}
}
Err(anyhow!(
"No compatible protocol version found. Server supports: {:?}, Client supports: {:?}",
SUPPORTED_VERSIONS,
client_versions
))
}
fn filter_capabilities(&self, requested: &[String]) -> Vec<String> {
requested
.iter()
.filter(|cap| self.available_capabilities.contains(cap))
.cloned()
.collect()
}
fn check_fallback_conditions(&self, version: u32, _capabilities: &[String]) -> Option<String> {
match version {
PROTOCOL_V1 => {
#[cfg(not(feature = "protobuf"))]
{
Some("Protobuf support not compiled in, falling back to JSONL".to_string())
}
#[cfg(feature = "protobuf")]
{
None
}
}
PROTOCOL_V0 => None, _ => Some(format!("Unsupported version {}, using v0", version)),
}
}
pub fn create_ready_event(&self, result: &NegotiationResult) -> Event {
Event::Ready {
version: result.version,
capabilities: result.capabilities.clone(),
service_id: self.service_id,
}
}
pub fn get_available_capabilities(&self) -> &[String] {
&self.available_capabilities
}
pub fn supports_capability(&self, capability: &str) -> bool {
self.available_capabilities
.contains(&capability.to_string())
}
pub fn get_service_metadata(&self) -> &HashMap<String, String> {
&self.service_metadata
}
}
pub struct CapabilityChecker;
impl CapabilityChecker {
pub fn check_compatibility(capabilities: &[String]) -> Result<Vec<String>> {
let mut warnings = Vec::new();
let has_js_hooks = capabilities.contains(&capabilities::HOOKS_JS.to_string());
let has_rust_hooks = capabilities.contains(&capabilities::HOOKS_RUST.to_string());
if has_js_hooks && has_rust_hooks {
warnings
.push("Both JS and Rust hooks enabled - performance may be impacted".to_string());
}
let has_nats = capabilities.contains(&capabilities::NATS.to_string());
let has_tracing = capabilities.contains(&capabilities::TRACING.to_string());
if has_nats && !has_tracing {
warnings.push("NATS enabled without tracing - reduced observability".to_string());
}
let has_replay = capabilities.contains(&capabilities::REPLAY.to_string());
if has_replay {
warnings.push(
"Replay enabled - performance overhead for recording all operations".to_string(),
);
}
Ok(warnings)
}
pub fn recommend_capabilities(use_case: &str) -> Vec<String> {
match use_case {
"development" => vec![
capabilities::SHELL_EXEC.to_string(),
capabilities::HOOKS_JS.to_string(),
capabilities::REPLAY.to_string(),
capabilities::TRACING.to_string(),
],
"production" => vec![
capabilities::SHELL_EXEC.to_string(),
capabilities::HOOKS_RUST.to_string(),
capabilities::TRACING.to_string(),
capabilities::NATS.to_string(),
],
"testing" => vec![
capabilities::SHELL_EXEC.to_string(),
capabilities::REPLAY.to_string(),
capabilities::TRACING.to_string(),
],
"minimal" => vec![capabilities::SHELL_EXEC.to_string()],
_ => vec![
capabilities::SHELL_EXEC.to_string(),
capabilities::TRACING.to_string(),
],
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_version_negotiation() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let client_info = ClientInfo {
supported_versions: vec![0, 1],
requested_capabilities: vec![
capabilities::SHELL_EXEC.to_string(),
capabilities::TRACING.to_string(),
],
client_metadata: HashMap::new(),
};
let result = negotiator.negotiate(client_info).unwrap();
assert!(result.version <= 1);
assert!(result
.capabilities
.contains(&capabilities::SHELL_EXEC.to_string()));
assert!(result
.capabilities
.contains(&capabilities::TRACING.to_string()));
}
#[test]
fn test_incompatible_versions() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let client_info = ClientInfo {
supported_versions: vec![999], requested_capabilities: vec![],
client_metadata: HashMap::new(),
};
let result = negotiator.negotiate(client_info);
assert!(result.is_err());
}
#[test]
fn test_capability_filtering() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let client_info = ClientInfo {
supported_versions: vec![0],
requested_capabilities: vec![
capabilities::SHELL_EXEC.to_string(),
"non_existent_capability".to_string(),
],
client_metadata: HashMap::new(),
};
let result = negotiator.negotiate(client_info).unwrap();
assert!(result
.capabilities
.contains(&capabilities::SHELL_EXEC.to_string()));
assert!(!result
.capabilities
.contains(&"non_existent_capability".to_string()));
}
#[test]
fn test_capability_compatibility() {
let warnings = CapabilityChecker::check_compatibility(&[
capabilities::HOOKS_JS.to_string(),
capabilities::HOOKS_RUST.to_string(),
])
.unwrap();
assert!(!warnings.is_empty());
assert!(warnings[0].contains("JS and Rust hooks"));
}
#[test]
fn test_use_case_recommendations() {
let dev_caps = CapabilityChecker::recommend_capabilities("development");
let prod_caps = CapabilityChecker::recommend_capabilities("production");
assert!(dev_caps.contains(&capabilities::REPLAY.to_string()));
assert!(dev_caps.contains(&capabilities::HOOKS_JS.to_string()));
assert!(!prod_caps.contains(&capabilities::REPLAY.to_string()));
assert!(prod_caps.contains(&capabilities::HOOKS_RUST.to_string()));
assert!(prod_caps.contains(&capabilities::NATS.to_string()));
}
#[test]
fn test_protocol_constants() {
assert_eq!(PROTOCOL_V0, 0);
assert_eq!(PROTOCOL_V1, 1);
assert_eq!(CURRENT_VERSION, PROTOCOL_V0);
assert!(SUPPORTED_VERSIONS.contains(&PROTOCOL_V0));
assert!(SUPPORTED_VERSIONS.contains(&PROTOCOL_V1));
assert!(SUPPORTED_VERSIONS.len() >= 2);
}
#[test]
fn test_negotiator_methods() {
let service_id = Uuid::new_v4();
let negotiator = ProtocolNegotiator::new(service_id);
assert!(negotiator.supports_capability(capabilities::SHELL_EXEC));
assert!(negotiator.supports_capability(capabilities::TRACING));
assert!(!negotiator.supports_capability("non_existent_capability"));
let caps = negotiator.get_available_capabilities();
assert!(!caps.is_empty());
assert!(caps.contains(&capabilities::SHELL_EXEC.to_string()));
let metadata = negotiator.get_service_metadata();
assert!(metadata.contains_key("service_name"));
assert!(metadata.contains_key("version"));
assert_eq!(metadata.get("service_name").unwrap(), "claude-code-rs-core");
}
#[test]
fn test_ready_event_creation() {
let service_id = Uuid::new_v4();
let negotiator = ProtocolNegotiator::new(service_id);
let negotiation_result = NegotiationResult {
version: PROTOCOL_V0,
capabilities: vec![capabilities::SHELL_EXEC.to_string()],
fallback_reason: None,
service_info: HashMap::new(),
};
let event = negotiator.create_ready_event(&negotiation_result);
match event {
Event::Ready {
version,
capabilities,
service_id: event_service_id,
} => {
assert_eq!(version, PROTOCOL_V0);
assert_eq!(capabilities, vec![capabilities::SHELL_EXEC.to_string()]);
assert_eq!(event_service_id, service_id);
}
_ => panic!("Expected Ready event"),
}
}
#[test]
fn test_fallback_conditions_v0() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let fallback = negotiator.check_fallback_conditions(PROTOCOL_V0, &[]);
assert!(fallback.is_none()); }
#[test]
fn test_fallback_conditions_unsupported_version() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let fallback = negotiator.check_fallback_conditions(999, &[]);
assert!(fallback.is_some());
assert!(fallback.unwrap().contains("Unsupported version 999"));
}
#[test]
fn test_capability_compatibility_nats_without_tracing() {
let warnings =
CapabilityChecker::check_compatibility(&[capabilities::NATS.to_string()]).unwrap();
assert!(!warnings.is_empty());
assert!(warnings
.iter()
.any(|w| w.contains("NATS enabled without tracing")));
}
#[test]
fn test_capability_compatibility_with_replay() {
let warnings =
CapabilityChecker::check_compatibility(&[capabilities::REPLAY.to_string()]).unwrap();
assert!(!warnings.is_empty());
assert!(warnings.iter().any(|w| w.contains("Replay enabled")));
}
#[test]
fn test_capability_compatibility_good_config() {
let warnings = CapabilityChecker::check_compatibility(&[
capabilities::SHELL_EXEC.to_string(),
capabilities::TRACING.to_string(),
])
.unwrap();
assert!(warnings.is_empty());
}
#[test]
fn test_use_case_recommendations_all_variants() {
let test_cases = vec![
(
"development",
vec![
capabilities::SHELL_EXEC,
capabilities::HOOKS_JS,
capabilities::REPLAY,
capabilities::TRACING,
],
),
(
"production",
vec![
capabilities::SHELL_EXEC,
capabilities::HOOKS_RUST,
capabilities::TRACING,
capabilities::NATS,
],
),
(
"testing",
vec![
capabilities::SHELL_EXEC,
capabilities::REPLAY,
capabilities::TRACING,
],
),
("minimal", vec![capabilities::SHELL_EXEC]),
(
"unknown_use_case",
vec![capabilities::SHELL_EXEC, capabilities::TRACING],
),
];
for (use_case, expected_caps) in test_cases {
let recommendations = CapabilityChecker::recommend_capabilities(use_case);
for cap in expected_caps {
assert!(
recommendations.contains(&cap.to_string()),
"Use case '{}' should include capability '{}'",
use_case,
cap
);
}
}
}
#[test]
fn test_client_info_and_negotiation_result_serialization() {
let mut metadata = HashMap::new();
metadata.insert("client_version".to_string(), "1.0.0".to_string());
let client_info = ClientInfo {
supported_versions: vec![0, 1],
requested_capabilities: vec![capabilities::SHELL_EXEC.to_string()],
client_metadata: metadata,
};
let json = serde_json::to_string(&client_info).unwrap();
let deserialized: ClientInfo = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.supported_versions, vec![0, 1]);
assert_eq!(
deserialized.requested_capabilities,
vec![capabilities::SHELL_EXEC.to_string()]
);
let mut service_info = HashMap::new();
service_info.insert("key".to_string(), "value".to_string());
let result = NegotiationResult {
version: PROTOCOL_V0,
capabilities: vec![capabilities::SHELL_EXEC.to_string()],
fallback_reason: Some("test fallback".to_string()),
service_info,
};
let json = serde_json::to_string(&result).unwrap();
let deserialized: NegotiationResult = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.version, PROTOCOL_V0);
assert_eq!(
deserialized.fallback_reason,
Some("test fallback".to_string())
);
}
#[test]
fn test_empty_client_versions_negotiation() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let client_info = ClientInfo {
supported_versions: vec![], requested_capabilities: vec![],
client_metadata: HashMap::new(),
};
let result = negotiator.negotiate(client_info);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("No compatible protocol version"));
}
#[test]
fn test_empty_requested_capabilities() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let client_info = ClientInfo {
supported_versions: vec![PROTOCOL_V0],
requested_capabilities: vec![], client_metadata: HashMap::new(),
};
let result = negotiator.negotiate(client_info).unwrap();
assert_eq!(result.version, PROTOCOL_V0);
assert!(result.capabilities.is_empty()); }
#[test]
fn test_version_selection_priority() {
let negotiator = ProtocolNegotiator::new(Uuid::new_v4());
let client_info = ClientInfo {
supported_versions: vec![PROTOCOL_V0, PROTOCOL_V1],
requested_capabilities: vec![],
client_metadata: HashMap::new(),
};
let result = negotiator.negotiate(client_info).unwrap();
assert!(SUPPORTED_VERSIONS.contains(&result.version));
}
}