owlen/crates/integration/jsonrpc/src/lib.rs

//! JSON-RPC 2.0 client for external tool communication
//!
//! This crate provides a JSON-RPC 2.0 client for invoking external tools
//! via stdio (spawning a process) or HTTP endpoints.

use color_eyre::eyre::{eyre, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::process::Stdio;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
use tokio::time::timeout;
use uuid::Uuid;

// ============================================================================
// JSON-RPC 2.0 Protocol Types
// ============================================================================

/// JSON-RPC 2.0 request
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcRequest {
    /// JSON-RPC version (must be "2.0")
    pub jsonrpc: String,
    /// Request ID for matching responses
    pub id: serde_json::Value,
    /// Method name to invoke
    pub method: String,
    /// Method parameters
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<serde_json::Value>,
}

impl JsonRpcRequest {
    /// Create a new JSON-RPC request
    pub fn new(method: impl Into<String>, params: Option<serde_json::Value>) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: serde_json::Value::String(Uuid::new_v4().to_string()),
            method: method.into(),
            params,
        }
    }

    /// Create a request with a specific ID
    pub fn with_id(id: impl Into<serde_json::Value>, method: impl Into<String>, params: Option<serde_json::Value>) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id: id.into(),
            method: method.into(),
            params,
        }
    }

    /// Serialize to JSON string with newline
    pub fn to_json_line(&self) -> Result<String> {
        let mut json = serde_json::to_string(self)?;
        json.push('\n');
        Ok(json)
    }
}

/// JSON-RPC 2.0 response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcResponse {
    /// JSON-RPC version (must be "2.0")
    pub jsonrpc: String,
    /// Request ID this is responding to
    pub id: serde_json::Value,
    /// Result (on success)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,
    /// Error (on failure)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<JsonRpcError>,
}

impl JsonRpcResponse {
    /// Check if the response is an error
    pub fn is_error(&self) -> bool {
        self.error.is_some()
    }

    /// Get the result or error as a Result type
    pub fn into_result(self) -> Result<serde_json::Value> {
        if let Some(err) = self.error {
            Err(eyre!("JSON-RPC error {}: {}", err.code, err.message))
        } else {
            self.result.ok_or_else(|| eyre!("No result in response"))
        }
    }
}

/// JSON-RPC 2.0 error
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcError {
    /// Error code (standard or application-specific)
    pub code: i64,
    /// Short error message
    pub message: String,
    /// Additional error data
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<serde_json::Value>,
}

/// Standard JSON-RPC 2.0 error codes
pub mod error_codes {
    /// Invalid JSON was received
    pub const PARSE_ERROR: i64 = -32700;
    /// The JSON sent is not a valid Request object
    pub const INVALID_REQUEST: i64 = -32600;
    /// The method does not exist or is not available
    pub const METHOD_NOT_FOUND: i64 = -32601;
    /// Invalid method parameter(s)
    pub const INVALID_PARAMS: i64 = -32602;
    /// Internal JSON-RPC error
    pub const INTERNAL_ERROR: i64 = -32603;
}

// ============================================================================
// Stdio Client
// ============================================================================

/// JSON-RPC client over stdio (spawned process)
pub struct StdioClient {
    /// Spawned child process
    child: Child,
    /// Request timeout
    timeout: Duration,
}

impl StdioClient {
    /// Spawn a new stdio client
    pub async fn spawn(
        command: impl AsRef<str>,
        args: &[String],
        env: &HashMap<String, String>,
        timeout_ms: u64,
    ) -> Result<Self> {
        let mut cmd = Command::new(command.as_ref());
        cmd.args(args)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit());

        // Add environment variables
        for (key, value) in env {
            cmd.env(key, value);
        }

        let child = cmd.spawn()?;

        Ok(Self {
            child,
            timeout: Duration::from_millis(timeout_ms),
        })
    }

    /// Spawn from a path (resolving relative paths)
    pub async fn spawn_from_path(
        command: &PathBuf,
        args: &[String],
        env: &HashMap<String, String>,
        working_dir: Option<&PathBuf>,
        timeout_ms: u64,
    ) -> Result<Self> {
        let mut cmd = Command::new(command);
        cmd.args(args)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit());

        if let Some(dir) = working_dir {
            cmd.current_dir(dir);
        }

        for (key, value) in env {
            cmd.env(key, value);
        }

        let child = cmd.spawn()?;

        Ok(Self {
            child,
            timeout: Duration::from_millis(timeout_ms),
        })
    }

    /// Send a request and wait for response
    pub async fn call(&mut self, method: &str, params: Option<serde_json::Value>) -> Result<serde_json::Value> {
        let request = JsonRpcRequest::new(method, params);
        let request_id = request.id.clone();

        // Get handles to stdin/stdout
        let stdin = self.child.stdin.as_mut()
            .ok_or_else(|| eyre!("Failed to get stdin handle"))?;
        let stdout = self.child.stdout.take()
            .ok_or_else(|| eyre!("Failed to get stdout handle"))?;

        // Write request
        let request_json = request.to_json_line()?;
        stdin.write_all(request_json.as_bytes()).await?;
        stdin.flush().await?;

        // Read response with timeout
        let mut reader = BufReader::new(stdout);
        let mut line = String::new();

        let result = timeout(self.timeout, reader.read_line(&mut line)).await;

        // Restore stdout for future calls
        self.child.stdout = Some(reader.into_inner());

        match result {
            Ok(Ok(0)) => Err(eyre!("Process closed stdout")),
            Ok(Ok(_)) => {
                let response: JsonRpcResponse = serde_json::from_str(&line)?;

                // Verify ID matches
                if response.id != request_id {
                    return Err(eyre!(
                        "Response ID mismatch: expected {:?}, got {:?}",
                        request_id,
                        response.id
                    ));
                }

                response.into_result()
            }
            Ok(Err(e)) => Err(e.into()),
            Err(_) => Err(eyre!("Request timed out after {:?}", self.timeout)),
        }
    }

    /// Check if the process is still running
    pub fn is_alive(&mut self) -> bool {
        matches!(self.child.try_wait(), Ok(None))
    }

    /// Kill the process
    pub async fn kill(&mut self) -> Result<()> {
        self.child.kill().await?;
        Ok(())
    }
}

impl Drop for StdioClient {
    fn drop(&mut self) {
        // Try to kill the process on drop
        let _ = self.child.start_kill();
    }
}

// ============================================================================
// Tool Executor
// ============================================================================

/// Executor for external tools via JSON-RPC
pub struct ToolExecutor {
    /// Default timeout for tool calls
    default_timeout: Duration,
}

impl Default for ToolExecutor {
    fn default() -> Self {
        Self {
            default_timeout: Duration::from_secs(30),
        }
    }
}

impl ToolExecutor {
    /// Create a new tool executor with custom timeout
    pub fn with_timeout(timeout_ms: u64) -> Self {
        Self {
            default_timeout: Duration::from_millis(timeout_ms),
        }
    }

    /// Execute a tool via stdio
    pub async fn execute_stdio(
        &self,
        command: &str,
        args: &[String],
        env: &HashMap<String, String>,
        tool_params: serde_json::Value,
        timeout_ms: Option<u64>,
    ) -> Result<serde_json::Value> {
        let timeout = timeout_ms.unwrap_or(self.default_timeout.as_millis() as u64);

        let mut client = StdioClient::spawn(command, args, env, timeout).await?;

        // The standard method name for tool execution
        let result = client.call("execute", Some(tool_params)).await?;

        // Kill the process after we're done
        let _ = client.kill().await;

        Ok(result)
    }

    /// Execute a tool via HTTP (not implemented yet)
    pub async fn execute_http(
        &self,
        _url: &str,
        _tool_params: serde_json::Value,
        _timeout_ms: Option<u64>,
    ) -> Result<serde_json::Value> {
        Err(eyre!("HTTP transport not yet implemented"))
    }
}

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

    #[test]
    fn test_jsonrpc_request_serialization() {
        let request = JsonRpcRequest::new("test_method", Some(serde_json::json!({"arg": "value"})));

        assert_eq!(request.jsonrpc, "2.0");
        assert_eq!(request.method, "test_method");
        assert!(request.params.is_some());

        let json = serde_json::to_string(&request).unwrap();
        assert!(json.contains("\"jsonrpc\":\"2.0\""));
        assert!(json.contains("\"method\":\"test_method\""));
    }

    #[test]
    fn test_jsonrpc_response_success() {
        let json = r#"{"jsonrpc":"2.0","id":"123","result":{"data":"hello"}}"#;
        let response: JsonRpcResponse = serde_json::from_str(json).unwrap();

        assert!(!response.is_error());
        assert_eq!(response.id, serde_json::json!("123"));
        assert_eq!(response.result.unwrap()["data"], "hello");
    }

    #[test]
    fn test_jsonrpc_response_error() {
        let json = r#"{"jsonrpc":"2.0","id":"123","error":{"code":-32600,"message":"Invalid Request"}}"#;
        let response: JsonRpcResponse = serde_json::from_str(json).unwrap();

        assert!(response.is_error());
        let err = response.error.unwrap();
        assert_eq!(err.code, error_codes::INVALID_REQUEST);
        assert_eq!(err.message, "Invalid Request");
    }

    #[test]
    fn test_jsonrpc_request_line_format() {
        let request = JsonRpcRequest::with_id("test-id", "method", None);
        let line = request.to_json_line().unwrap();

        assert!(line.ends_with('\n'));
        assert!(!line.ends_with("\n\n"));
    }

    #[test]
    fn test_response_into_result_success() {
        let response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!("1"),
            result: Some(serde_json::json!({"success": true})),
            error: None,
        };

        let result = response.into_result().unwrap();
        assert_eq!(result["success"], true);
    }

    #[test]
    fn test_response_into_result_error() {
        let response = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: serde_json::json!("1"),
            result: None,
            error: Some(JsonRpcError {
                code: -32600,
                message: "Invalid".to_string(),
                data: None,
            }),
        };

        let result = response.into_result();
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("-32600"));
    }
}