huskies/server/src/service/ws/dispatch.rs

//! Pure request dispatch logic — no side effects.
//!
//! Contains the branching logic for resolving permission responses and
//! classifying incoming requests. All functions are pure data transformations;
//! I/O (socket reads, spawning tasks) lives in `io.rs`.

use crate::http::context::PermissionDecision;
use std::collections::HashMap;
use tokio::sync::oneshot;

use super::message::{WsRequest, WsResponse};

/// The result of dispatching a single `WsRequest` outside an active chat session.
pub enum DispatchResult {
    /// Start a chat session with the given messages and config.
    StartChat {
        messages: Vec<crate::llm::types::Message>,
        config: crate::llm::chat::ProviderConfig,
    },
    /// Cancel the current chat session.
    CancelChat,
    /// Respond with a pong.
    Pong,
    /// Permission response outside an active chat — silently ignored.
    IgnoredPermission,
    /// Start a side question.
    StartSideQuestion {
        question: String,
        context_messages: Vec<crate::llm::types::Message>,
        config: crate::llm::chat::ProviderConfig,
    },
    /// The request could not be parsed.
    ParseError(String),
}

/// Parse a raw JSON text into a [`DispatchResult`].
///
/// This is the outer-loop dispatch: determines what action to take for a
/// message received when no chat session is active.
pub fn dispatch_outer(text: &str) -> DispatchResult {
    match serde_json::from_str::<WsRequest>(text) {
        Ok(WsRequest::Chat { messages, config }) => DispatchResult::StartChat { messages, config },
        Ok(WsRequest::Cancel) => DispatchResult::CancelChat,
        Ok(WsRequest::Ping) => DispatchResult::Pong,
        Ok(WsRequest::PermissionResponse { .. }) => DispatchResult::IgnoredPermission,
        Ok(WsRequest::SideQuestion {
            question,
            context_messages,
            config,
        }) => DispatchResult::StartSideQuestion {
            question,
            context_messages,
            config,
        },
        Err(err) => DispatchResult::ParseError(format!("Invalid request: {err}")),
    }
}

/// The result of dispatching a message during an active chat session.
pub enum InnerDispatchResult {
    /// A permission response was successfully resolved.
    PermissionResolved,
    /// Cancel the current chat session.
    CancelChat,
    /// Respond with a pong.
    Pong,
    /// Start a side question (can run concurrently with the chat).
    StartSideQuestion {
        question: String,
        context_messages: Vec<crate::llm::types::Message>,
        config: crate::llm::chat::ProviderConfig,
    },
    /// The message was not actionable during a chat (unknown type, etc.).
    Ignored,
}

/// Parse a raw JSON text and dispatch it within an active chat session.
///
/// Permission responses are resolved against the `pending_perms` map.
/// Returns what action the caller should take.
pub fn dispatch_inner(
    text: &str,
    pending_perms: &mut HashMap<String, oneshot::Sender<PermissionDecision>>,
) -> InnerDispatchResult {
    match serde_json::from_str::<WsRequest>(text) {
        Ok(WsRequest::PermissionResponse {
            request_id,
            approved,
            always_allow,
        }) => {
            if let Some(resp_tx) = pending_perms.remove(&request_id) {
                let decision = resolve_permission(approved, always_allow);
                let _ = resp_tx.send(decision);
            }
            InnerDispatchResult::PermissionResolved
        }
        Ok(WsRequest::Cancel) => InnerDispatchResult::CancelChat,
        Ok(WsRequest::Ping) => InnerDispatchResult::Pong,
        Ok(WsRequest::SideQuestion {
            question,
            context_messages,
            config,
        }) => InnerDispatchResult::StartSideQuestion {
            question,
            context_messages,
            config,
        },
        _ => InnerDispatchResult::Ignored,
    }
}

/// Map the `approved` and `always_allow` flags to a [`PermissionDecision`].
pub fn resolve_permission(approved: bool, always_allow: bool) -> PermissionDecision {
    if always_allow {
        PermissionDecision::AlwaysAllow
    } else if approved {
        PermissionDecision::Approve
    } else {
        PermissionDecision::Deny
    }
}

/// Build a [`WsResponse::Error`] from an error message.
pub fn error_response(message: String) -> WsResponse {
    WsResponse::Error { message }
}

/// Build the permission request forward message for the client.
pub fn permission_request_response(
    request_id: &str,
    tool_name: &str,
    tool_input: &serde_json::Value,
) -> WsResponse {
    WsResponse::PermissionRequest {
        request_id: request_id.to_string(),
        tool_name: tool_name.to_string(),
        tool_input: tool_input.clone(),
    }
}

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

    // ── dispatch_outer ──────────────────────────────────────────────

    #[test]
    fn dispatch_outer_chat() {
        let json = r#"{"type":"chat","messages":[{"role":"user","content":"hi"}],"config":{"provider":"ollama","model":"m"}}"#;
        let result = dispatch_outer(json);
        assert!(matches!(result, DispatchResult::StartChat { .. }));
    }

    #[test]
    fn dispatch_outer_cancel() {
        let result = dispatch_outer(r#"{"type":"cancel"}"#);
        assert!(matches!(result, DispatchResult::CancelChat));
    }

    #[test]
    fn dispatch_outer_ping() {
        let result = dispatch_outer(r#"{"type":"ping"}"#);
        assert!(matches!(result, DispatchResult::Pong));
    }

    #[test]
    fn dispatch_outer_permission_response_ignored() {
        let json = r#"{"type":"permission_response","request_id":"x","approved":true}"#;
        let result = dispatch_outer(json);
        assert!(matches!(result, DispatchResult::IgnoredPermission));
    }

    #[test]
    fn dispatch_outer_side_question() {
        let json = r#"{"type":"side_question","question":"what?","context_messages":[],"config":{"provider":"ollama","model":"m"}}"#;
        let result = dispatch_outer(json);
        assert!(matches!(result, DispatchResult::StartSideQuestion { .. }));
    }

    #[test]
    fn dispatch_outer_invalid_json() {
        let result = dispatch_outer("not json");
        match result {
            DispatchResult::ParseError(msg) => {
                assert!(msg.contains("Invalid request"));
            }
            _ => panic!("expected ParseError"),
        }
    }

    #[test]
    fn dispatch_outer_unknown_type() {
        let result = dispatch_outer(r#"{"type":"bogus"}"#);
        assert!(matches!(result, DispatchResult::ParseError(_)));
    }

    #[test]
    fn dispatch_outer_missing_type() {
        let result = dispatch_outer(r#"{"messages":[]}"#);
        assert!(matches!(result, DispatchResult::ParseError(_)));
    }

    // ── dispatch_inner ──────────────────────────────────────────────

    #[test]
    fn dispatch_inner_permission_response_resolves() {
        let (tx, rx) = oneshot::channel();
        let mut perms = HashMap::new();
        perms.insert("req-1".to_string(), tx);

        let json = r#"{"type":"permission_response","request_id":"req-1","approved":true,"always_allow":false}"#;
        let result = dispatch_inner(json, &mut perms);
        assert!(matches!(result, InnerDispatchResult::PermissionResolved));
        assert!(perms.is_empty());
        assert_eq!(rx.blocking_recv().unwrap(), PermissionDecision::Approve);
    }

    #[test]
    fn dispatch_inner_permission_response_always_allow() {
        let (tx, rx) = oneshot::channel();
        let mut perms = HashMap::new();
        perms.insert("req-2".to_string(), tx);

        let json = r#"{"type":"permission_response","request_id":"req-2","approved":true,"always_allow":true}"#;
        let result = dispatch_inner(json, &mut perms);
        assert!(matches!(result, InnerDispatchResult::PermissionResolved));
        assert_eq!(rx.blocking_recv().unwrap(), PermissionDecision::AlwaysAllow);
    }

    #[test]
    fn dispatch_inner_permission_response_deny() {
        let (tx, rx) = oneshot::channel();
        let mut perms = HashMap::new();
        perms.insert("req-3".to_string(), tx);

        let json = r#"{"type":"permission_response","request_id":"req-3","approved":false}"#;
        let result = dispatch_inner(json, &mut perms);
        assert!(matches!(result, InnerDispatchResult::PermissionResolved));
        assert_eq!(rx.blocking_recv().unwrap(), PermissionDecision::Deny);
    }

    #[test]
    fn dispatch_inner_permission_unknown_request_id() {
        let mut perms: HashMap<String, oneshot::Sender<PermissionDecision>> = HashMap::new();
        let json = r#"{"type":"permission_response","request_id":"unknown","approved":true}"#;
        let result = dispatch_inner(json, &mut perms);
        // Still returns PermissionResolved — the unknown ID is silently ignored.
        assert!(matches!(result, InnerDispatchResult::PermissionResolved));
    }

    #[test]
    fn dispatch_inner_cancel() {
        let mut perms = HashMap::new();
        let result = dispatch_inner(r#"{"type":"cancel"}"#, &mut perms);
        assert!(matches!(result, InnerDispatchResult::CancelChat));
    }

    #[test]
    fn dispatch_inner_ping() {
        let mut perms = HashMap::new();
        let result = dispatch_inner(r#"{"type":"ping"}"#, &mut perms);
        assert!(matches!(result, InnerDispatchResult::Pong));
    }

    #[test]
    fn dispatch_inner_side_question() {
        let mut perms = HashMap::new();
        let json = r#"{"type":"side_question","question":"what?","context_messages":[],"config":{"provider":"ollama","model":"m"}}"#;
        let result = dispatch_inner(json, &mut perms);
        assert!(matches!(
            result,
            InnerDispatchResult::StartSideQuestion { .. }
        ));
    }

    #[test]
    fn dispatch_inner_chat_during_chat_ignored() {
        let mut perms = HashMap::new();
        let json = r#"{"type":"chat","messages":[],"config":{"provider":"ollama","model":"m"}}"#;
        let result = dispatch_inner(json, &mut perms);
        assert!(matches!(result, InnerDispatchResult::Ignored));
    }

    #[test]
    fn dispatch_inner_invalid_json_ignored() {
        let mut perms = HashMap::new();
        let result = dispatch_inner("not json", &mut perms);
        assert!(matches!(result, InnerDispatchResult::Ignored));
    }

    // ── resolve_permission ──────────────────────────────────────────

    #[test]
    fn resolve_approve() {
        assert_eq!(resolve_permission(true, false), PermissionDecision::Approve);
    }

    #[test]
    fn resolve_deny() {
        assert_eq!(resolve_permission(false, false), PermissionDecision::Deny);
    }

    #[test]
    fn resolve_always_allow() {
        assert_eq!(
            resolve_permission(true, true),
            PermissionDecision::AlwaysAllow
        );
    }

    #[test]
    fn resolve_always_allow_overrides_denied() {
        // always_allow=true should win even if approved=false
        assert_eq!(
            resolve_permission(false, true),
            PermissionDecision::AlwaysAllow
        );
    }

    // ── error_response ──────────────────────────────────────────────

    #[test]
    fn error_response_creates_error_variant() {
        let resp = error_response("oops".to_string());
        let json = serde_json::to_value(&resp).unwrap();
        assert_eq!(json["type"], "error");
        assert_eq!(json["message"], "oops");
    }

    // ── permission_request_response ─────────────────────────────────

    #[test]
    fn permission_request_response_creates_correct_variant() {
        let input = serde_json::json!({"command": "rm -rf /"});
        let resp = permission_request_response("req-1", "Bash", &input);
        let json = serde_json::to_value(&resp).unwrap();
        assert_eq!(json["type"], "permission_request");
        assert_eq!(json["request_id"], "req-1");
        assert_eq!(json["tool_name"], "Bash");
        assert_eq!(json["tool_input"]["command"], "rm -rf /");
    }
}