aether-acp-utils 0.4.1

//! Typed wire-format types for Aether's custom ACP extension requests and
//! notifications.
use std::path::PathBuf;

use agent_client_protocol::schema::v1::{AuthMethod, Meta};
use agent_client_protocol::{JsonRpcNotification, JsonRpcRequest, JsonRpcResponse};
pub use mcp_utils::display_meta::{ToolDisplayMeta, ToolResultMeta};
use serde::{Deserialize, Serialize, de::DeserializeOwned};

pub use mcp_utils::status::{McpServerAuthCapability, McpServerStatus, McpServerStatusEntry};

pub const AETHER_META_NAMESPACE: &str = "contextbridge/aether";

/// Parameters for `_aether/session_usage` notifications.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonRpcNotification)]
#[notification(method = "_aether/session_usage")]
pub struct SessionUsageParams {
    pub usage: llm::SessionUsageEvent,
}

/// Parameters for `_aether/context_compaction` notifications.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
#[notification(method = "_aether/context_compaction")]
pub struct ContextCompactionParams {
    pub active: bool,
}

/// Parameters for `_aether/context_cleared` notifications.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, JsonRpcNotification)]
#[notification(method = "_aether/context_cleared")]
pub struct ContextClearedParams {}

/// Parameters for `_aether/auth_methods_updated` notifications.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
#[notification(method = "_aether/auth_methods_updated")]
pub struct AuthMethodsUpdatedParams {
    pub auth_methods: Vec<AuthMethod>,
}

/// Parameters for the `_aether/prompt_search` request.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
#[request(method = "_aether/prompt_search", response = PromptSearchResponse)]
#[serde(rename_all = "camelCase")]
pub struct PromptSearchParams {
    pub query: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub limit: Option<usize>,
}

/// Response for the `_aether/prompt_search` request.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
#[serde(rename_all = "camelCase")]
pub struct PromptSearchResponse {
    pub query: String,
    pub results: Vec<PromptSearchResult>,
    pub truncated: bool,
}

/// A single prompt-history search hit.
///
/// `match_start` and `match_end` are UTF-8 byte offsets into `prompt` and are
/// guaranteed to fall on char boundaries.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct PromptSearchResult {
    pub session_id: String,
    pub cwd: PathBuf,
    pub session_created_at: String,
    pub prompt: String,
    pub match_start: usize,
    pub match_end: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
#[request(method = "_aether/session_preview", response = SessionPreviewResponse)]
#[serde(rename_all = "camelCase")]
pub struct SessionPreviewParams {
    pub session_id: String,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
#[serde(rename_all = "camelCase")]
pub struct SessionPreviewResponse {
    pub session_id: String,
    pub cwd: PathBuf,
    pub created_at: String,
    pub model: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub selected_mode: Option<String>,
    pub transcript: Vec<SessionPreviewTurn>,
    pub tool_call_count: usize,
    pub truncated: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct SessionPreviewTurn {
    pub role: SessionPreviewRole,
    pub text: String,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub enum SessionPreviewRole {
    User,
    Assistant,
}

/// Parameters for the `_aether/workspace_list` request.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
#[request(method = "_aether/workspace_list", response = WorkspaceListResponse)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceListParams {
    pub session_id: String,
}

/// Response for the `_aether/workspace_list` request: every managed workspace
/// originating from the same git repository as the session's working directory.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceListResponse {
    pub workspaces: Vec<WorkspaceEntry>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceEntry {
    pub path: PathBuf,
    pub is_current: bool,
}

/// Parameters for the `_aether/workspace_move` request.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
#[request(method = "_aether/workspace_move", response = WorkspaceMoveResponse)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceMoveParams {
    pub session_id: String,
    pub target: WorkspaceMoveTarget,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "kind", rename_all = "camelCase")]
pub enum WorkspaceMoveTarget {
    Existing { path: PathBuf },
    New { name: String },
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceMoveResponse {
    pub new_cwd: PathBuf,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "camelCase")]
pub struct SessionDisplayMeta {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub selected_mode: Option<String>,
}

impl SessionDisplayMeta {
    #[must_use]
    pub fn new(model: impl Into<String>, selected_mode: Option<String>) -> Self {
        Self { model: Some(model.into()), selected_mode }
    }

    #[must_use]
    pub fn to_meta(&self) -> Meta {
        to_aether_meta(self)
    }

    #[must_use]
    pub fn from_meta(meta: Option<&Meta>) -> Self {
        from_aether_meta(meta)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "camelCase")]
pub struct AetherCapabilities {
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub prompt_search: bool,
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub session_preview: bool,
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub workspace_move: bool,
}

impl AetherCapabilities {
    #[must_use]
    pub fn to_meta(self) -> Meta {
        to_aether_meta(&self)
    }

    #[must_use]
    pub fn from_meta(meta: Option<&Meta>) -> Self {
        from_aether_meta(meta)
    }
}

fn to_aether_meta<T: Serialize>(value: &T) -> Meta {
    let mut meta = Meta::new();
    meta.insert(AETHER_META_NAMESPACE.to_string(), serde_json::json!(value));
    meta
}

fn from_aether_meta<T: DeserializeOwned + Default>(meta: Option<&Meta>) -> T {
    meta.and_then(|m| m.get(AETHER_META_NAMESPACE))
        .cloned()
        .and_then(|value| serde_json::from_value(value).ok())
        .unwrap_or_default()
}

/// Server→client MCP extension notifications (relay → wisp).
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
#[notification(method = "_aether/mcp_event")]
pub enum McpNotification {
    ServerStatus { servers: Vec<McpServerStatusEntry> },
}

/// Client→server MCP extension requests (wisp → relay).
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
#[notification(method = "_aether/mcp_request")]
pub enum McpRequest {
    Authenticate { session_id: String, server_name: String },
}

/// Parameters for `_aether/sub_agent_progress` notifications.
///
/// This is the wire format sent from the ACP server (`aether-cli`) to clients like `wisp`.
#[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)]
#[notification(method = "_aether/sub_agent_progress")]
pub struct SubAgentProgressParams {
    pub parent_tool_id: String,
    pub task_id: String,
    pub agent_name: String,
    pub event: SubAgentEvent,
}

/// Subset of agent message variants relevant for sub-agent status display.
///
/// The ACP server (`aether-cli`) converts `AgentEvent` to this type before
/// serializing, so the wire format only contains these known variants.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SubAgentEvent {
    ToolCall { request: SubAgentToolRequest },
    ToolCallUpdate { update: SubAgentToolCallUpdate },
    ToolResult { result: SubAgentToolResult },
    ToolError { error: SubAgentToolError },
    Done,
    Other,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubAgentToolRequest {
    pub id: String,
    pub name: String,
    pub arguments: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubAgentToolCallUpdate {
    pub id: String,
    pub chunk: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubAgentToolResult {
    pub id: String,
    pub name: String,
    pub result_meta: Option<ToolResultMeta>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SubAgentToolError {
    pub id: String,
    pub name: String,
}

#[cfg(test)]
mod tests {
    use agent_client_protocol::JsonRpcMessage;
    use agent_client_protocol::schema::v1::AuthMethodAgent;

    use super::*;

    #[test]
    fn wire_method_names_are_prefixed() {
        assert_eq!(ContextClearedParams::default().method(), "_aether/context_cleared");
        assert_eq!(AuthMethodsUpdatedParams { auth_methods: vec![] }.method(), "_aether/auth_methods_updated");
        assert_eq!(McpNotification::ServerStatus { servers: vec![] }.method(), "_aether/mcp_event");
        assert_eq!(
            McpRequest::Authenticate { session_id: String::new(), server_name: String::new() }.method(),
            "_aether/mcp_request"
        );
        assert_eq!(PromptSearchParams { query: String::new(), limit: None }.method(), "_aether/prompt_search");
        assert_eq!(SessionPreviewParams { session_id: String::new() }.method(), "_aether/session_preview");
        assert_eq!(WorkspaceListParams { session_id: String::new() }.method(), "_aether/workspace_list");
        let move_params =
            WorkspaceMoveParams { session_id: String::new(), target: WorkspaceMoveTarget::New { name: String::new() } };
        assert_eq!(move_params.method(), "_aether/workspace_move");
    }

    #[test]
    fn context_compaction_params_roundtrip() {
        for active in [true, false] {
            let params = ContextCompactionParams { active };
            let untyped = params.to_untyped_message().expect("serializable");
            assert_eq!(untyped.method(), "_aether/context_compaction");
            let parsed = ContextCompactionParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
            assert_eq!(parsed, params);
        }
    }

    #[test]
    fn context_cleared_params_roundtrip() {
        let params = ContextClearedParams::default();
        let untyped = params.to_untyped_message().expect("serializable");
        assert_eq!(untyped.method(), "_aether/context_cleared");
        let parsed = ContextClearedParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
        assert_eq!(parsed, params);
    }

    #[test]
    fn auth_methods_updated_roundtrip() {
        let params = AuthMethodsUpdatedParams {
            auth_methods: vec![
                AuthMethod::Agent(AuthMethodAgent::new("anthropic", "Anthropic").description("authenticated")),
                AuthMethod::Agent(AuthMethodAgent::new("openrouter", "OpenRouter")),
            ],
        };

        let untyped = params.to_untyped_message().expect("serializable");
        assert_eq!(untyped.method(), "_aether/auth_methods_updated");
        let parsed = AuthMethodsUpdatedParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
        assert_eq!(parsed, params);
    }

    #[test]
    fn mcp_request_authenticate_roundtrip() {
        let msg = McpRequest::Authenticate {
            session_id: "session-0".to_string(),
            server_name: "my oauth server".to_string(),
        };

        let untyped = msg.to_untyped_message().expect("serializable");
        assert_eq!(untyped.method(), "_aether/mcp_request");
        let parsed = McpRequest::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
        assert_eq!(parsed, msg);
    }

    #[test]
    fn mcp_notification_server_status_roundtrip() {
        let msg = McpNotification::ServerStatus {
            servers: vec![
                McpServerStatusEntry::new("github", McpServerStatus::Connected { tool_count: 5 }),
                McpServerStatusEntry::new("linear", McpServerStatus::NeedsOAuth)
                    .with_auth_capability(McpServerAuthCapability::OAuth),
                McpServerStatusEntry::new("slack", McpServerStatus::Failed { error: "connection timeout".to_string() }),
            ],
        };

        let untyped = msg.to_untyped_message().expect("serializable");
        assert_eq!(untyped.method(), "_aether/mcp_event");
        let parsed = McpNotification::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
        assert_eq!(parsed, msg);
    }

    #[test]
    fn sub_agent_progress_params_roundtrip() {
        let params = SubAgentProgressParams {
            parent_tool_id: "call_123".to_string(),
            task_id: "task_abc".to_string(),
            agent_name: "explorer".to_string(),
            event: SubAgentEvent::Done,
        };

        let untyped = params.to_untyped_message().expect("serializable");
        assert_eq!(untyped.method(), "_aether/sub_agent_progress");
    }

    #[test]
    fn mcp_server_status_entry_serde_roundtrip() {
        let entry = McpServerStatusEntry::new("test-server", McpServerStatus::Connected { tool_count: 3 })
            .with_auth_capability(McpServerAuthCapability::OAuth);

        let json = serde_json::to_string(&entry).unwrap();
        assert!(json.contains("\"auth_capability\":\"OAuth\""));
        assert!(json.contains("\"deferTools\":false"));
        let parsed: McpServerStatusEntry = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, entry);
        assert!(!parsed.deferred_tools);
        assert!(parsed.can_authenticate());
    }

    #[test]
    fn mcp_server_status_entry_deferred_tools_serde_roundtrip() {
        let entry = McpServerStatusEntry::new("math", McpServerStatus::NeedsOAuth)
            .with_auth_capability(McpServerAuthCapability::OAuth)
            .with_deferred_tools(true);

        let json = serde_json::to_string(&entry).unwrap();
        assert!(json.contains("\"deferTools\":true"));
        let parsed: McpServerStatusEntry = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, entry);
    }

    #[test]
    fn deserialize_tool_call_event() {
        let json = r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{\"pattern\":\"test\"}"},"model_name":"m"}}"#;
        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
        assert!(matches!(event, SubAgentEvent::ToolCall { .. }));
    }

    #[test]
    fn deserialize_tool_call_update_event() {
        let json = r#"{"ToolCallUpdate":{"update":{"id":"c1","chunk":"{\"pattern\":\"test\"}"},"model_name":"m"}}"#;
        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
        assert!(matches!(event, SubAgentEvent::ToolCallUpdate { .. }));
    }

    #[test]
    fn deserialize_tool_result_event() {
        let json = r#"{"ToolResult":{"result":{"id":"c1","name":"grep","result_meta":{"display":{"title":"Grep","value":"'test' in src (3 matches)"}}}}}"#;
        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
        match event {
            SubAgentEvent::ToolResult { result } => {
                let result_meta = result.result_meta.expect("expected result_meta");
                assert_eq!(result_meta.display.title, "Grep");
            }
            other => panic!("Expected ToolResult, got {other:?}"),
        }
    }

    #[test]
    fn deserialize_tool_error_event() {
        let json = r#"{"ToolError":{"error":{"id":"c1","name":"grep"}}}"#;
        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
        assert!(matches!(event, SubAgentEvent::ToolError { .. }));
    }

    #[test]
    fn deserialize_done_event() {
        let event: SubAgentEvent = serde_json::from_str(r#""Done""#).unwrap();
        assert!(matches!(event, SubAgentEvent::Done));
    }

    #[test]
    fn deserialize_other_variant() {
        let event: SubAgentEvent = serde_json::from_str(r#""Other""#).unwrap();
        assert!(matches!(event, SubAgentEvent::Other));
    }

    #[test]
    fn tool_result_meta_map_roundtrip() {
        let meta: ToolResultMeta = ToolDisplayMeta::new("Read file", "Cargo.toml, 156 lines").into();
        let map = meta.clone().into_map();
        let parsed = ToolResultMeta::from_map(&map).expect("should deserialize ToolResultMeta");
        assert_eq!(parsed, meta);
    }
}