dirigent/crates/dirigent_acp_api/src/jsonrpc.rs

//! JSON-RPC 2.0 types for the ACP Server
//!
//! This module implements JSON-RPC 2.0 request/response types according to
//! the specification at https://www.jsonrpc.org/specification.
//!
//! Key features:
//! - Support for both numeric and string IDs
//! - Batch request/response handling
//! - Proper serialization of null vs missing fields

use serde::{Deserialize, Serialize};

use crate::error::JsonRpcErrorObject;

/// JSON-RPC protocol version constant
pub const JSONRPC_VERSION: &str = "2.0";

/// JSON-RPC request/response identifier
///
/// According to the JSON-RPC 2.0 spec, an id can be a String, Number,
/// or Null. This type uses an untagged enum to handle both string and
/// number identifiers.
///
/// Note: The spec recommends not using Null as an id for requests.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum JsonRpcId {
    /// Numeric identifier (integer)
    Number(i64),

    /// String identifier
    String(String),

    /// Null identifier (typically used in error responses for invalid requests)
    Null,
}

impl From<i64> for JsonRpcId {
    fn from(n: i64) -> Self {
        JsonRpcId::Number(n)
    }
}

impl From<String> for JsonRpcId {
    fn from(s: String) -> Self {
        JsonRpcId::String(s)
    }
}

impl From<&str> for JsonRpcId {
    fn from(s: &str) -> Self {
        JsonRpcId::String(s.to_string())
    }
}

/// A JSON-RPC 2.0 request object
///
/// Represents a remote procedure call with optional parameters.
/// The `id` field determines whether this is a request (with id) or
/// notification (without id).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcRequest {
    /// JSON-RPC protocol version (must be "2.0")
    pub jsonrpc: String,

    /// A String containing the name of the method to be invoked
    pub method: String,

    /// Optional structured value that holds the parameter values
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<serde_json::Value>,

    /// Optional identifier established by the client
    ///
    /// If absent, the request is a notification (no response expected)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<JsonRpcId>,
}

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

    /// Create a new JSON-RPC notification (request without id)
    pub fn notification(method: impl Into<String>, params: Option<serde_json::Value>) -> Self {
        Self {
            jsonrpc: JSONRPC_VERSION.to_string(),
            method: method.into(),
            params,
            id: None,
        }
    }

    /// Check if this request is a notification (no id)
    pub fn is_notification(&self) -> bool {
        self.id.is_none()
    }

    /// Validate the request format
    pub fn validate(&self) -> Result<(), String> {
        if self.jsonrpc != JSONRPC_VERSION {
            return Err(format!(
                "Invalid JSON-RPC version: expected '{}', got '{}'",
                JSONRPC_VERSION, self.jsonrpc
            ));
        }

        if self.method.is_empty() {
            return Err("Method name cannot be empty".to_string());
        }

        // Methods starting with "rpc." are reserved for internal use
        if self.method.starts_with("rpc.") {
            return Err(format!(
                "Method name '{}' is reserved (starts with 'rpc.')",
                self.method
            ));
        }

        Ok(())
    }
}

/// A JSON-RPC 2.0 response object
///
/// Contains either a result (success) or an error (failure), never both.
/// The id must match the corresponding request id.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcResponse {
    /// JSON-RPC protocol version (must be "2.0")
    pub jsonrpc: String,

    /// The result of the call (on success)
    ///
    /// This member is REQUIRED on success and MUST NOT exist on error.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<serde_json::Value>,

    /// The error object (on failure)
    ///
    /// This member is REQUIRED on error and MUST NOT exist on success.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<JsonRpcErrorObject>,

    /// The identifier matching the request
    ///
    /// If there was an error detecting the id in the Request object
    /// (e.g. Parse error/Invalid Request), it MUST be Null.
    pub id: JsonRpcId,
}

impl JsonRpcResponse {
    /// Create a successful response
    pub fn success(result: serde_json::Value, id: JsonRpcId) -> Self {
        Self {
            jsonrpc: JSONRPC_VERSION.to_string(),
            result: Some(result),
            error: None,
            id,
        }
    }

    /// Create an error response
    pub fn error(error: JsonRpcErrorObject, id: JsonRpcId) -> Self {
        Self {
            jsonrpc: JSONRPC_VERSION.to_string(),
            result: None,
            error: Some(error),
            id,
        }
    }

    /// Create an error response with a null id (for parse errors)
    pub fn error_with_null_id(error: JsonRpcErrorObject) -> Self {
        Self::error(error, JsonRpcId::Null)
    }

    /// Check if this response represents success
    pub fn is_success(&self) -> bool {
        self.result.is_some() && self.error.is_none()
    }

    /// Check if this response represents an error
    pub fn is_error(&self) -> bool {
        self.error.is_some()
    }
}

/// Represents either a single request or a batch of requests
///
/// The JSON-RPC 2.0 spec allows sending multiple requests in a single
/// JSON array for batch processing.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum JsonRpcRequestBatch {
    /// A single request
    Single(JsonRpcRequest),

    /// A batch of requests
    Batch(Vec<JsonRpcRequest>),
}

impl JsonRpcRequestBatch {
    /// Check if this is an empty batch
    pub fn is_empty(&self) -> bool {
        match self {
            JsonRpcRequestBatch::Single(_) => false,
            JsonRpcRequestBatch::Batch(batch) => batch.is_empty(),
        }
    }

    /// Get the number of requests
    pub fn len(&self) -> usize {
        match self {
            JsonRpcRequestBatch::Single(_) => 1,
            JsonRpcRequestBatch::Batch(batch) => batch.len(),
        }
    }

    /// Convert to a vector of requests
    pub fn into_vec(self) -> Vec<JsonRpcRequest> {
        match self {
            JsonRpcRequestBatch::Single(req) => vec![req],
            JsonRpcRequestBatch::Batch(batch) => batch,
        }
    }
}

/// Represents either a single response or a batch of responses
///
/// The response format must match the request format: single request
/// gets single response, batch request gets batch response.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum JsonRpcResponseBatch {
    /// A single response
    Single(JsonRpcResponse),

    /// A batch of responses
    Batch(Vec<JsonRpcResponse>),
}

impl JsonRpcResponseBatch {
    /// Create a batch response from a vector
    ///
    /// Returns Single if there's exactly one response, otherwise Batch.
    pub fn from_vec(responses: Vec<JsonRpcResponse>) -> Self {
        if responses.len() == 1 {
            JsonRpcResponseBatch::Single(responses.into_iter().next().unwrap())
        } else {
            JsonRpcResponseBatch::Batch(responses)
        }
    }
}

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

    #[test]
    fn test_jsonrpc_id_number() {
        let id = JsonRpcId::Number(42);
        let json = serde_json::to_string(&id).unwrap();
        assert_eq!(json, "42");

        let parsed: JsonRpcId = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, id);
    }

    #[test]
    fn test_jsonrpc_id_string() {
        let id = JsonRpcId::String("abc-123".to_string());
        let json = serde_json::to_string(&id).unwrap();
        assert_eq!(json, "\"abc-123\"");

        let parsed: JsonRpcId = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, id);
    }

    #[test]
    fn test_jsonrpc_id_null() {
        let id = JsonRpcId::Null;
        let json = serde_json::to_string(&id).unwrap();
        assert_eq!(json, "null");

        let parsed: JsonRpcId = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, id);
    }

    #[test]
    fn test_jsonrpc_id_from_conversions() {
        let id: JsonRpcId = 42i64.into();
        assert_eq!(id, JsonRpcId::Number(42));

        let id: JsonRpcId = "test".into();
        assert_eq!(id, JsonRpcId::String("test".to_string()));

        let id: JsonRpcId = String::from("owned").into();
        assert_eq!(id, JsonRpcId::String("owned".to_string()));
    }

    #[test]
    fn test_request_creation() {
        let req = JsonRpcRequest::new("session.new", Some(json!({"title": "Test"})), 1.into());
        assert_eq!(req.jsonrpc, "2.0");
        assert_eq!(req.method, "session.new");
        assert!(req.params.is_some());
        assert_eq!(req.id, Some(JsonRpcId::Number(1)));
        assert!(!req.is_notification());
    }

    #[test]
    fn test_notification_creation() {
        let notif = JsonRpcRequest::notification("event.ping", None);
        assert!(notif.is_notification());
        assert_eq!(notif.id, None);
    }

    #[test]
    fn test_request_validation() {
        let valid = JsonRpcRequest::new("test.method", None, 1.into());
        assert!(valid.validate().is_ok());

        // Invalid version
        let mut invalid_version = valid.clone();
        invalid_version.jsonrpc = "1.0".to_string();
        assert!(invalid_version.validate().is_err());

        // Empty method
        let mut empty_method = valid.clone();
        empty_method.method = String::new();
        assert!(empty_method.validate().is_err());

        // Reserved method
        let mut reserved = valid.clone();
        reserved.method = "rpc.internal".to_string();
        assert!(reserved.validate().is_err());
    }

    #[test]
    fn test_request_serialization() {
        let req = JsonRpcRequest::new(
            "session.prompt",
            Some(json!({"session_id": "abc", "content": "Hello"})),
            "req-123".into(),
        );

        let json = serde_json::to_string(&req).unwrap();
        assert!(json.contains("\"jsonrpc\":\"2.0\""));
        assert!(json.contains("\"method\":\"session.prompt\""));
        assert!(json.contains("\"id\":\"req-123\""));

        // Deserialize back
        let parsed: JsonRpcRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.method, "session.prompt");
    }

    #[test]
    fn test_response_success() {
        let resp = JsonRpcResponse::success(json!({"session_id": "new-123"}), 1.into());
        assert!(resp.is_success());
        assert!(!resp.is_error());
        assert_eq!(resp.result, Some(json!({"session_id": "new-123"})));
        assert!(resp.error.is_none());
    }

    #[test]
    fn test_response_error() {
        let error = JsonRpcErrorObject::method_not_found("unknown.method");
        let resp = JsonRpcResponse::error(error, 1.into());
        assert!(!resp.is_success());
        assert!(resp.is_error());
        assert!(resp.result.is_none());
        assert!(resp.error.is_some());
    }

    #[test]
    fn test_response_serialization() {
        // Success response
        let success = JsonRpcResponse::success(json!({"ok": true}), 42.into());
        let json = serde_json::to_string(&success).unwrap();
        assert!(json.contains("\"result\""));
        assert!(!json.contains("\"error\""));

        // Error response
        let error = JsonRpcResponse::error(
            JsonRpcErrorObject::internal_error("Something broke"),
            42.into(),
        );
        let json = serde_json::to_string(&error).unwrap();
        assert!(!json.contains("\"result\""));
        assert!(json.contains("\"error\""));
    }

    #[test]
    fn test_batch_request_single() {
        let req = JsonRpcRequest::new("test", None, 1.into());
        let batch = JsonRpcRequestBatch::Single(req);
        assert_eq!(batch.len(), 1);
        assert!(!batch.is_empty());

        let vec = batch.into_vec();
        assert_eq!(vec.len(), 1);
    }

    #[test]
    fn test_batch_request_multiple() {
        let req1 = JsonRpcRequest::new("test1", None, 1.into());
        let req2 = JsonRpcRequest::new("test2", None, 2.into());
        let batch = JsonRpcRequestBatch::Batch(vec![req1, req2]);
        assert_eq!(batch.len(), 2);

        let vec = batch.into_vec();
        assert_eq!(vec.len(), 2);
    }

    #[test]
    fn test_batch_request_empty() {
        let batch = JsonRpcRequestBatch::Batch(vec![]);
        assert!(batch.is_empty());
        assert_eq!(batch.len(), 0);
    }

    #[test]
    fn test_batch_request_deserialization() {
        // Single request
        let single_json = r#"{"jsonrpc":"2.0","method":"test","id":1}"#;
        let single: JsonRpcRequestBatch = serde_json::from_str(single_json).unwrap();
        assert_eq!(single.len(), 1);

        // Batch request
        let batch_json = r#"[{"jsonrpc":"2.0","method":"test1","id":1},{"jsonrpc":"2.0","method":"test2","id":2}]"#;
        let batch: JsonRpcRequestBatch = serde_json::from_str(batch_json).unwrap();
        assert_eq!(batch.len(), 2);
    }

    #[test]
    fn test_batch_response_from_vec() {
        // Single response becomes Single variant
        let responses = vec![JsonRpcResponse::success(json!(null), 1.into())];
        let batch = JsonRpcResponseBatch::from_vec(responses);
        matches!(batch, JsonRpcResponseBatch::Single(_));

        // Multiple responses become Batch variant
        let responses = vec![
            JsonRpcResponse::success(json!(null), 1.into()),
            JsonRpcResponse::success(json!(null), 2.into()),
        ];
        let batch = JsonRpcResponseBatch::from_vec(responses);
        matches!(batch, JsonRpcResponseBatch::Batch(_));
    }
}