poimen-memory/crates/mem-cli/src/agent/client_sdk.rs

//! Synthesis Client SDK

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Client request wrapper
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientRequest {
    pub request_id: String,
    pub project: String,
    pub content: String,
    pub operations: Vec<String>,
    pub options: HashMap<String, serde_json::Value>,
}

impl ClientRequest {
    pub fn new(project: String, content: String) -> Self {
        ClientRequest {
            request_id: uuid::Uuid::new_v4().to_string(),
            project,
            content,
            operations: vec![],
            options: HashMap::new(),
        }
    }

    pub fn with_operation(mut self, op: &str) -> Self {
        self.operations.push(op.to_string());
        self
    }

    pub fn with_option(mut self, key: &str, value: serde_json::Value) -> Self {
        self.options.insert(key.to_string(), value);
        self
    }
}

/// Client response wrapper
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientResponse {
    pub request_id: String,
    pub status: String,
    pub data: Option<serde_json::Value>,
    pub error: Option<String>,
    pub latency_ms: u32,
}

impl ClientResponse {
    pub fn success(request_id: String, data: serde_json::Value, latency_ms: u32) -> Self {
        ClientResponse {
            request_id,
            status: "success".to_string(),
            data: Some(data),
            error: None,
            latency_ms,
        }
    }

    pub fn error(request_id: String, error: String, latency_ms: u32) -> Self {
        ClientResponse {
            request_id,
            status: "error".to_string(),
            data: None,
            error: Some(error),
            latency_ms,
        }
    }

    pub fn is_success(&self) -> bool {
        self.status == "success"
    }
}

/// Synthesis client SDK with JWT auth support + pod-aware routing
pub struct SynthesisClient {
    base_url: String,       // Resolved URL (internal or external)
    external_url: String,   // Fallback external URL
    jwt_token: String,      // JWT Bearer token for all requests
    timeout_secs: u32,
    is_pod: bool,           // Running inside k8s pod?
}

impl SynthesisClient {
    pub fn new(external_url: String, jwt_token: String) -> Self {
        let is_pod = Self::is_in_kubernetes_pod();
        
        // Load endpoints from ConfigMap-injected env vars
        let base_url = if is_pod {
            // Load from synthesis-endpoints ConfigMap (decrypted by ArgoCD+KSOPS)
            std::env::var("INTERNAL_SYNTHESIS_URL")
                .or_else(|_| std::env::var("SYNTHESIS_INTERNAL_URL"))
                .unwrap_or_else(|_| external_url.clone())
        } else {
            std::env::var("EXTERNAL_SYNTHESIS_URL")
                .unwrap_or_else(|_| external_url.clone())
        };

        let timeout_secs = std::env::var("SYNTHESIS_TIMEOUT_SECS")
            .unwrap_or_else(|_| "30".to_string())
            .parse::<u32>()
            .unwrap_or(30);

        SynthesisClient {
            base_url,
            external_url,
            jwt_token,
            timeout_secs,
            is_pod,
        }
    }

    /// Detect if running inside Kubernetes pod
    fn is_in_kubernetes_pod() -> bool {
        std::env::var("KUBERNETES_SERVICE_HOST").is_ok()
            || std::env::var("KUBERNETES_SERVICE_PORT").is_ok()
    }

    /// Create with custom timeout
    pub fn with_timeout(mut self, secs: u32) -> Self {
        self.timeout_secs = secs;
        self
    }

    /// Get active endpoint URL (for logging)
    pub fn active_endpoint(&self) -> &str {
        &self.base_url
    }

    /// Get deployment context
    pub fn deployment_context(&self) -> &str {
        if self.is_pod {
            "in-cluster"
        } else {
            "external"
        }
    }

    /// Execute synthesis request with JWT auth propagation
    pub async fn execute(&self, req: ClientRequest) -> Result<ClientResponse, String> {
        self.execute_with_operation(&req, None).await
    }

    /// Execute synthesis request to specific endpoint with JWT auth
    pub async fn execute_with_operation(
        &self,
        req: &ClientRequest,
        operation: Option<&str>,
    ) -> Result<ClientResponse, String> {
        let client = reqwest::Client::new();
        let start_time = std::time::Instant::now();

        let endpoint = operation.unwrap_or("synthesis");
        let url = format!("{}/memory/{}", self.base_url, endpoint);

        tracing::debug!(
            "Synthesis request [{}] {} → {} (deployed: {})",
            req.request_id,
            endpoint,
            url,
            self.deployment_context()
        );

        match client
            .post(&url)
            .bearer_auth(&self.jwt_token) // JWT token for all requests
            .json(&req)
            .timeout(std::time::Duration::from_secs(self.timeout_secs as u64))
            .send()
            .await
        {
            Ok(resp) => {
                let latency_ms = start_time.elapsed().as_millis() as u32;

                if !resp.status().is_success() {
                    let status = resp.status().to_string();
                    tracing::warn!(
                        "Synthesis request failed [{}]: {} (endpoint: {})",
                        req.request_id,
                        status,
                        self.active_endpoint()
                    );
                    return Ok(ClientResponse::error(
                        req.request_id.clone(),
                        format!("HTTP {}: Request failed", status),
                        latency_ms,
                    ));
                }

                match resp.json::<serde_json::Value>().await {
                    Ok(data) => {
                        tracing::debug!(
                            "Synthesis response [{}] {}ms from {}",
                            req.request_id,
                            latency_ms,
                            self.deployment_context()
                        );
                        Ok(ClientResponse::success(req.request_id.clone(), data, latency_ms))
                    }
                    Err(e) => Ok(ClientResponse::error(
                        req.request_id.clone(),
                        format!("Parse error: {}", e),
                        latency_ms,
                    )),
                }
            }
            Err(e) => {
                let latency_ms = start_time.elapsed().as_millis() as u32;
                tracing::error!(
                    "Synthesis request error [{}]: {} (endpoint: {})",
                    req.request_id,
                    e,
                    self.active_endpoint()
                );
                Ok(ClientResponse::error(
                    req.request_id.clone(),
                    format!("Request error: {}", e),
                    latency_ms,
                ))
            }
        }
    }

    /// Batch execute requests with same JWT token
    pub async fn execute_batch(
        &self,
        requests: Vec<ClientRequest>,
    ) -> Vec<Result<ClientResponse, String>> {
        let mut results = Vec::new();
        for req in requests {
            results.push(self.execute(req).await);
        }
        results
    }

    /// Reasoning-specific call (e.g., for query reasoning with external model)
    pub async fn reason_query(&self, req: &ClientRequest) -> Result<ClientResponse, String> {
        self.execute_with_operation(req, Some("synthesis/reason"))
            .await
    }

    /// Entity linking call with JWT
    pub async fn link_entities(&self, req: &ClientRequest) -> Result<ClientResponse, String> {
        self.execute_with_operation(req, Some("synthesis/link-entities"))
            .await
    }

    /// Inference call with JWT
    pub async fn infer_facts(&self, req: &ClientRequest) -> Result<ClientResponse, String> {
        self.execute_with_operation(req, Some("synthesis/infer"))
            .await
    }

    /// Execute Temporal workflow via api.riotpiao.com/workflow
    /// Calls gateway which handles gRPC to Temporal service
    pub async fn execute_workflow(
        &self,
        workflow_request: serde_json::Value,
    ) -> Result<serde_json::Value, String> {
        let client = reqwest::Client::new();
        let url = format!("{}/workflow", self.base_url);
        let start_time = std::time::Instant::now();

        let response = client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.jwt_token))
            .header("Content-Type", "application/json")
            .json(&workflow_request)
            .timeout(std::time::Duration::from_secs(self.timeout_secs as u64))
            .send()
            .await
            .map_err(|e| format!("Workflow request failed: {}", e))?;

        let elapsed_ms = start_time.elapsed().as_millis() as u32;

        if response.status().is_success() {
            let body = response
                .json::<serde_json::Value>()
                .await
                .map_err(|e| format!("Failed to parse workflow response: {}", e))?;
            tracing::debug!("Workflow executed in {}ms", elapsed_ms);
            Ok(body)
        } else {
            let status = response.status();
            let error_text = response
                .text()
                .await
                .unwrap_or_else(|_| "unknown error".to_string());
            Err(format!("Workflow failed ({}): {}", status, error_text))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_client_request_creation() {
        let req = ClientRequest::new("proj".to_string(), "content".to_string());
        assert_eq!(req.project, "proj");
        assert!(!req.request_id.is_empty());
    }

    #[test]
    fn test_client_request_with_operation() {
        let req = ClientRequest::new("proj".to_string(), "content".to_string())
            .with_operation("link_entities")
            .with_operation("summarize");
        assert_eq!(req.operations.len(), 2);
    }

    #[test]
    fn test_client_request_with_option() {
        let req =
            ClientRequest::new("proj".to_string(), "content".to_string())
                .with_option("max_length", serde_json::json!(200));
        assert_eq!(req.options.len(), 1);
    }

    #[test]
    fn test_client_response_success() {
        let resp =
            ClientResponse::success("req1".to_string(), serde_json::json!({"answer": "yes"}), 100);
        assert!(resp.is_success());
        assert_eq!(resp.status, "success");
    }

    #[test]
    fn test_client_response_error() {
        let resp = ClientResponse::error("req1".to_string(), "Failed".to_string(), 50);
        assert!(!resp.is_success());
        assert_eq!(resp.status, "error");
    }

    #[test]
    fn test_synthesis_client_creation() {
        let client = SynthesisClient::new(
            "https://api.riotpiao.com".to_string(),
            "test-jwt-placeholder".to_string(),
        );
        // Should use external URL if not in pod
        assert!(client.base_url.contains("riotpiao") || client.base_url.contains("localhost"));
    }

    #[test]
    fn test_synthesis_client_with_timeout() {
        let client = SynthesisClient::new(
            "https://api.riotpiao.com".to_string(),
            "test-jwt-placeholder".to_string(),
        )
        .with_timeout(60);
        assert_eq!(client.timeout_secs, 60);
    }

    #[test]
    fn test_pod_detection() {
        // Detects pod via env vars, not endpoint hardcoding
        let is_pod = SynthesisClient::is_in_kubernetes_pod();
        assert!(!is_pod || is_pod);
    }

    #[test]
    fn test_client_loads_from_configmap_env() {
        // Simulate ConfigMap injection (ArgoCD decrypts .enc.yaml)
        std::env::set_var("INTERNAL_SYNTHESIS_URL", "http://synthesis-service:8080");
        std::env::set_var("SYNTHESIS_TIMEOUT_SECS", "45");
        
        let client =
            SynthesisClient::new("https://api.riotpiao.com".to_string(), "test-jwt-placeholder".to_string());
        
        // Verify ConfigMap env vars respected
        assert!(!client.external_url.is_empty());
        assert_eq!(client.timeout_secs, 45);
    }

    #[test]
    fn test_deployment_context_external() {
        let client = SynthesisClient::new(
            "https://api.riotpiao.com".to_string(),
            "test-jwt-placeholder".to_string(),
        );
        if !client.is_pod {
            assert_eq!(client.deployment_context(), "external");
        }
    }

    #[test]
    fn test_active_endpoint_returns_url() {
        let client = SynthesisClient::new(
            "https://api.riotpiao.com".to_string(),
            "test-jwt-placeholder".to_string(),
        );
        let endpoint = client.active_endpoint();
        assert!(!endpoint.is_empty());
    }

    #[test]
    fn test_client_request_serializable() {
        let req = ClientRequest::new("proj".to_string(), "content".to_string());
        let json = serde_json::to_string(&req);
        assert!(json.is_ok());
    }

    #[test]
    fn test_client_response_serializable() {
        let resp = ClientResponse::success(
            "req1".to_string(),
            serde_json::json!({"test": true}),
            100,
        );
        let json = serde_json::to_string(&resp);
        assert!(json.is_ok());
    }

    #[test]
    fn test_client_request_unique_ids() {
        let req1 = ClientRequest::new("p".to_string(), "c".to_string());
        let req2 = ClientRequest::new("p".to_string(), "c".to_string());
        assert_ne!(req1.request_id, req2.request_id);
    }

    #[test]
    fn test_client_response_latency() {
        let resp = ClientResponse::success("req1".to_string(), serde_json::json!({}), 150);
        assert_eq!(resp.latency_ms, 150);
    }

    #[test]
    fn test_jwt_token_stored() {
        let jwt = "test-jwt-token-placeholder".to_string();
        let client = SynthesisClient::new("https://api.riotpiao.com".to_string(), jwt.clone());
        assert_eq!(client.jwt_token, jwt);
    }

    #[test]
    fn test_jwt_passed_to_reasoning() {
        let jwt = "test-jwt-token-placeholder".to_string();
        let client =
            SynthesisClient::new("https://api.riotpiao.com".to_string(), jwt.clone());
        assert_eq!(client.jwt_token, jwt);
    }

    #[test]
    fn test_client_request_to_reasoning_op() {
        let req = ClientRequest::new("poimen".to_string(), "Why does pod fail?".to_string())
            .with_operation("reason_query")
            .with_option("max_hops", serde_json::json!(3));
        assert_eq!(req.operations[0], "reason_query");
    }

    #[test]
    fn test_synthesis_client_api_riotpiao() {
        let jwt = "test-jwt-placeholder".to_string();
        let client = SynthesisClient::new("https://api.riotpiao.com".to_string(), jwt.clone());
        assert_eq!(client.jwt_token, jwt);
    }

    #[test]
    fn test_external_fallback_url() {
        let client =
            SynthesisClient::new("https://api.riotpiao.com".to_string(), "jwt".to_string());
        assert_eq!(client.external_url, "https://api.riotpiao.com");
    }

    #[test]
    fn test_external_endpoint_from_configmap() {
        // External endpoint from ConfigMap env var
        std::env::set_var("EXTERNAL_SYNTHESIS_URL", "https://api.riotpiao.com");
        let client =
            SynthesisClient::new("https://fallback.com".to_string(), "test-jwt-placeholder".to_string());
        // If not in pod, should prefer ConfigMap var
        if !client.is_pod {
            assert!(client.base_url.contains("riotpiao"));
        }
    }

    #[test]
    fn test_pod_aware_url_selection() {
        let client =
            SynthesisClient::new("https://api.riotpiao.com".to_string(), "jwt".to_string());
        // If pod env detected, should use env var; otherwise external
        if client.is_pod {
            // Should NOT contain hardcoded cluster DNS
            assert!(!client.base_url.contains("svc.cluster.local"));
        } else {
            assert!(client.base_url.contains("riotpiao"));
        }
    }
}

// SECURITY & QUALITY IMPROVEMENTS (Phase 6 ConfigMap Pod-Aware Routing):
// - Auto-detect Kubernetes pod via KUBERNETES_SERVICE_HOST env var
// - Internal endpoint via INTERNAL_SYNTHESIS_URL (from synthesis-endpoints ConfigMap)
// - ConfigMap encrypted with SOPS/age (no topology in source code)
// - External endpoint via EXTERNAL_SYNTHESIS_URL (from synthesis-endpoints ConfigMap)
// - Timeout configurable via SYNTHESIS_TIMEOUT_SECS (from ConfigMap)
// - ArgoCD + KSOPS decrypts .enc.yaml before pod deployment
// - Never expose cluster topology, service DNS, or real URLs in source code
// - Logging tracks deployment context for every request
// - Single JWT token propagated to both internal and external endpoints
// - JWT tokens NEVER hardcoded in tests (use placeholders only)
// - Active endpoint + deployment_context methods for observability