mi6_core/config/
otel.rs

1//! OpenTelemetry configuration utilities.
2//!
3//! This module handles OTel environment variable generation and management
4//! for automatic token tracking.
5
6use serde::{Deserialize, Serialize};
7
8/// Default OTel port (matches OpenTelemetry standard HTTP/JSON port).
9pub const DEFAULT_OTEL_PORT: u16 = 4318;
10
11/// OTel server processing mode.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
13#[serde(rename_all = "kebab-case")]
14pub enum OtelMode {
15    /// Shell out to `mi6 ingest otel` for processing (default).
16    /// Enables dynamic upgrades without server restart.
17    #[default]
18    RelayCli,
19    /// Use mi6-core library directly for processing.
20    /// Lower latency, but requires server restart for upgrades.
21    Library,
22}
23
24impl std::str::FromStr for OtelMode {
25    type Err = String;
26
27    fn from_str(s: &str) -> Result<Self, Self::Err> {
28        match s.to_lowercase().as_str() {
29            "relay-cli" | "relay_cli" | "relaycli" => Ok(Self::RelayCli),
30            "library" | "lib" => Ok(Self::Library),
31            _ => Err(format!(
32                "invalid mode '{}', expected 'relay-cli' or 'library'",
33                s
34            )),
35        }
36    }
37}
38
39impl std::fmt::Display for OtelMode {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        match self {
42            Self::RelayCli => write!(f, "relay-cli"),
43            Self::Library => write!(f, "library"),
44        }
45    }
46}
47
48/// OTel server configuration.
49#[derive(Debug, Clone, Default, Serialize, Deserialize)]
50pub struct OtelConfig {
51    /// Server processing mode.
52    #[serde(default)]
53    pub mode: OtelMode,
54}
55
56/// OTel-related environment variable keys.
57///
58/// These keys are used to configure OpenTelemetry for automatic token tracking.
59pub const OTEL_ENV_KEYS: &[&str] = &[
60    "CLAUDE_CODE_ENABLE_TELEMETRY",
61    "OTEL_LOGS_EXPORTER",
62    "OTEL_EXPORTER_OTLP_PROTOCOL",
63    "OTEL_EXPORTER_OTLP_ENDPOINT",
64];
65
66/// Generate OpenTelemetry environment variables configuration.
67///
68/// Returns a JSON object containing the environment variables needed
69/// to enable OpenTelemetry-based token tracking.
70///
71/// # Arguments
72/// * `port` - The port where the OTel server will listen.
73pub fn generate_otel_env(port: u16) -> serde_json::Value {
74    serde_json::json!({
75        "CLAUDE_CODE_ENABLE_TELEMETRY": "1",
76        "OTEL_LOGS_EXPORTER": "otlp",
77        "OTEL_EXPORTER_OTLP_PROTOCOL": "http/json",
78        "OTEL_EXPORTER_OTLP_ENDPOINT": format!("http://127.0.0.1:{}", port)
79    })
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn test_generate_otel_env() {
88        let env = generate_otel_env(4318);
89        assert_eq!(env["CLAUDE_CODE_ENABLE_TELEMETRY"], "1");
90        assert_eq!(env["OTEL_LOGS_EXPORTER"], "otlp");
91        assert_eq!(env["OTEL_EXPORTER_OTLP_PROTOCOL"], "http/json");
92        assert_eq!(env["OTEL_EXPORTER_OTLP_ENDPOINT"], "http://127.0.0.1:4318");
93    }
94
95    #[test]
96    fn test_generate_otel_env_custom_port() {
97        let env = generate_otel_env(9999);
98        assert_eq!(env["OTEL_EXPORTER_OTLP_ENDPOINT"], "http://127.0.0.1:9999");
99    }
100
101    #[test]
102    fn test_otel_env_keys() {
103        assert_eq!(OTEL_ENV_KEYS.len(), 4);
104        assert!(OTEL_ENV_KEYS.contains(&"CLAUDE_CODE_ENABLE_TELEMETRY"));
105        assert!(OTEL_ENV_KEYS.contains(&"OTEL_LOGS_EXPORTER"));
106        assert!(OTEL_ENV_KEYS.contains(&"OTEL_EXPORTER_OTLP_PROTOCOL"));
107        assert!(OTEL_ENV_KEYS.contains(&"OTEL_EXPORTER_OTLP_ENDPOINT"));
108    }
109
110    #[test]
111    fn test_otel_mode_default() {
112        let mode = OtelMode::default();
113        assert_eq!(mode, OtelMode::RelayCli);
114    }
115
116    #[test]
117    fn test_otel_mode_from_str() {
118        assert_eq!(
119            "relay-cli".parse::<OtelMode>().expect("parse"),
120            OtelMode::RelayCli
121        );
122        assert_eq!(
123            "relay_cli".parse::<OtelMode>().expect("parse"),
124            OtelMode::RelayCli
125        );
126        assert_eq!(
127            "relaycli".parse::<OtelMode>().expect("parse"),
128            OtelMode::RelayCli
129        );
130        assert_eq!(
131            "library".parse::<OtelMode>().expect("parse"),
132            OtelMode::Library
133        );
134        assert_eq!("lib".parse::<OtelMode>().expect("parse"), OtelMode::Library);
135        assert!("invalid".parse::<OtelMode>().is_err());
136    }
137
138    #[test]
139    fn test_otel_mode_display() {
140        assert_eq!(OtelMode::RelayCli.to_string(), "relay-cli");
141        assert_eq!(OtelMode::Library.to_string(), "library");
142    }
143
144    #[test]
145    fn test_otel_mode_serde() {
146        // Test serialization
147        let mode = OtelMode::RelayCli;
148        let json = serde_json::to_string(&mode).expect("serialize");
149        assert_eq!(json, "\"relay-cli\"");
150
151        let mode = OtelMode::Library;
152        let json = serde_json::to_string(&mode).expect("serialize");
153        assert_eq!(json, "\"library\"");
154
155        // Test deserialization
156        let mode: OtelMode = serde_json::from_str("\"relay-cli\"").expect("deserialize");
157        assert_eq!(mode, OtelMode::RelayCli);
158
159        let mode: OtelMode = serde_json::from_str("\"library\"").expect("deserialize");
160        assert_eq!(mode, OtelMode::Library);
161    }
162
163    #[test]
164    fn test_otel_config_default() {
165        let config = OtelConfig::default();
166        assert_eq!(config.mode, OtelMode::RelayCli);
167    }
168
169    #[test]
170    fn test_otel_config_from_toml() {
171        let toml_str = r#"
172            mode = "library"
173        "#;
174        let config: OtelConfig = toml::from_str(toml_str).expect("parse toml");
175        assert_eq!(config.mode, OtelMode::Library);
176    }
177
178    #[test]
179    fn test_otel_config_from_toml_default_mode() {
180        // Empty config should use default mode
181        let toml_str = "";
182        let config: OtelConfig = toml::from_str(toml_str).expect("parse toml");
183        assert_eq!(config.mode, OtelMode::RelayCli);
184    }
185}