huskies/server/src/chat/transport/matrix/bot/context.rs

//! Matrix bot context — shared state for the Matrix bot (rooms, history, permissions).
use crate::chat::ChatTransport;
use crate::service::timer::TimerStore;
use crate::services::Services;
use matrix_sdk::ruma::{OwnedEventId, OwnedRoomId, OwnedUserId};
use std::collections::{BTreeMap, HashSet};
use std::sync::Arc;
use tokio::sync::Mutex as TokioMutex;
use tokio::sync::RwLock;

use super::history::ConversationHistory;

/// Shared context injected into Matrix event handlers.
#[derive(Clone)]
pub struct BotContext {
    /// Shared services bundle (project root, agent pool, bot identity, permissions).
    pub services: Arc<Services>,
    /// Matrix-specific parsed user ID (e.g. `@timmy:homeserver.local`).
    /// Transport-specific — kept separate from `services.bot_user_id` (String)
    /// because Matrix SDK APIs require `OwnedUserId` for comparisons and
    /// `.localpart()` extraction.
    pub matrix_user_id: OwnedUserId,
    /// All room IDs the bot listens in.
    pub target_room_ids: Vec<OwnedRoomId>,
    pub allowed_users: Vec<String>,
    /// Shared, per-room rolling conversation history.
    pub history: ConversationHistory,
    /// Maximum number of entries to keep per room before trimming the oldest.
    pub history_size: usize,
    /// Event IDs of messages the bot has sent.  Used to detect replies to the
    /// bot so it can continue a conversation thread without requiring an
    /// explicit `@mention` on every follow-up.
    pub bot_sent_event_ids: Arc<TokioMutex<HashSet<OwnedEventId>>>,
    /// Per-room htop monitoring sessions.  Keyed by room ID; each entry holds
    /// a stop-signal sender that the background task watches.
    pub htop_sessions: super::super::htop::HtopSessions,
    /// Chat transport used for sending and editing messages.
    ///
    /// All message I/O goes through this abstraction so the bot logic works
    /// with any platform, not just Matrix.
    pub transport: Arc<dyn ChatTransport>,
    /// Persistent store for pending deferred-start timers.
    pub timer_store: Arc<TimerStore>,
    /// In gateway mode: the currently active project (shared with the gateway HTTP handler).
    /// `None` in standalone single-project mode.
    pub gateway_active_project: Option<Arc<RwLock<String>>>,
    /// In gateway mode: valid project names accepted by the `switch` command.
    /// Empty in standalone mode.
    pub gateway_projects: Vec<String>,
    /// In gateway mode: mapping of project name → base URL (e.g. `"http://localhost:3001"`).
    /// Used to proxy bot commands to the active project over WebSocket (`/ws`).
    /// Empty in standalone mode.
    pub gateway_project_urls: BTreeMap<String, String>,
    /// Pipeline transition events buffered since the last LLM turn.
    ///
    /// A background task appends one compact audit line per real stage
    /// transition.  `handle_message` drains this buffer and injects it as a
    /// `<system-reminder>` block at the head of the next user prompt so Timmy
    /// sees pipeline activity without requiring a separate message.
    pub pending_pipeline_events: Arc<TokioMutex<Vec<String>>>,
}

impl BotContext {
    /// Resolve the effective project root for command dispatch.
    ///
    /// In gateway mode the bot's `project_root` is the gateway config directory.
    /// Each project lives in a subdirectory named after the project, so the
    /// effective root for commands is `project_root / active_project_name`.
    /// In standalone (single-project) mode this returns `project_root` unchanged.
    pub async fn effective_project_root(&self) -> std::path::PathBuf {
        if let Some(ref ap) = self.gateway_active_project {
            let name = ap.read().await.clone();
            self.services.project_root.join(&name)
        } else {
            self.services.project_root.clone()
        }
    }

    /// Returns `true` if the bot is running in gateway mode.
    pub fn is_gateway(&self) -> bool {
        self.gateway_active_project.is_some()
    }

    /// Return the base URL for the currently active project, if in gateway mode.
    pub async fn active_project_url(&self) -> Option<String> {
        let ap = self.gateway_active_project.as_ref()?;
        let name = ap.read().await.clone();
        self.gateway_project_urls.get(&name).cloned()
    }

    /// Proxy a bot command to the active project over a WebSocket RPC call.
    ///
    /// Connects to `{base_url}/ws`, sends an `rpc_request` frame for the
    /// `bot.command` method, and returns the Markdown response from the
    /// `rpc_response` frame.  Returns an error message string if the
    /// connection or command fails.
    pub async fn proxy_bot_command(&self, command: &str, args: &str) -> Option<String> {
        use futures::{SinkExt, StreamExt};
        use tokio_tungstenite::tungstenite::Message as WsMsg;

        let base_url = self.active_project_url().await?;

        // Convert http(s):// → ws(s)://
        let ws_base = if let Some(rest) = base_url.strip_prefix("https://") {
            format!("wss://{rest}")
        } else if let Some(rest) = base_url.strip_prefix("http://") {
            format!("ws://{rest}")
        } else {
            base_url.clone()
        };
        let ws_url = format!("{ws_base}/ws");

        let correlation_id = uuid::Uuid::new_v4().to_string();
        let request = serde_json::json!({
            "kind": "rpc_request",
            "version": 1,
            "correlation_id": correlation_id,
            "ttl_ms": 30_000u64,
            "method": "bot.command",
            "params": { "command": command, "args": args },
        });
        let request_text = match serde_json::to_string(&request) {
            Ok(t) => t,
            Err(e) => return Some(format!("Failed to serialize RPC request: {e}")),
        };

        let ws_stream = match tokio_tungstenite::connect_async(&ws_url).await {
            Ok((stream, _)) => stream,
            Err(e) => {
                return Some(format!(
                    "Failed to connect to project server at {ws_url}: {e}"
                ));
            }
        };

        let (mut sink, mut stream) = ws_stream.split();

        if let Err(e) = sink.send(WsMsg::Text(request_text.into())).await {
            return Some(format!("Failed to send RPC request: {e}"));
        }

        while let Some(msg) = stream.next().await {
            match msg {
                Ok(WsMsg::Text(text)) => {
                    let Ok(frame) = serde_json::from_str::<serde_json::Value>(&text) else {
                        continue;
                    };
                    if frame.get("kind").and_then(|v| v.as_str()) != Some("rpc_response") {
                        continue;
                    }
                    if frame
                        .get("correlation_id")
                        .and_then(|v| v.as_str())
                        .map(|id| id != correlation_id)
                        .unwrap_or(true)
                    {
                        continue;
                    }
                    let ok = frame.get("ok").and_then(|v| v.as_bool()).unwrap_or(false);
                    if ok {
                        return frame
                            .get("result")
                            .and_then(|r| r.get("response"))
                            .and_then(|v| v.as_str())
                            .map(String::from)
                            .or_else(|| {
                                Some("Command succeeded with no response text".to_string())
                            });
                    } else {
                        let err = frame
                            .get("error")
                            .and_then(|v| v.as_str())
                            .unwrap_or("unknown error");
                        return Some(format!("Project server command failed: {err}"));
                    }
                }
                Ok(WsMsg::Close(_)) => break,
                Err(e) => return Some(format!("WebSocket error: {e}")),
                _ => continue,
            }
        }

        Some("Connection closed before receiving command response".to_string())
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use std::path::PathBuf;
    use tokio::sync::mpsc;

    fn make_user_id(s: &str) -> OwnedUserId {
        s.parse().unwrap()
    }

    /// Build a test `Services` bundle with the given project root.
    fn test_services(project_root: PathBuf) -> Arc<Services> {
        let (_perm_tx, perm_rx) = mpsc::unbounded_channel();
        Arc::new(Services {
            project_root,
            agents: Arc::new(crate::agents::AgentPool::new_test(3000)),
            bot_name: "Assistant".to_string(),
            bot_user_id: "@bot:example.com".to_string(),
            ambient_rooms: Arc::new(std::sync::Mutex::new(std::collections::HashSet::new())),
            perm_rx: Arc::new(TokioMutex::new(perm_rx)),
            pending_perm_replies: Arc::new(TokioMutex::new(HashMap::new())),
            permission_timeout_secs: 120,
            status: Arc::new(crate::service::status::StatusBroadcaster::new()),
        })
    }

    /// Build a minimal `BotContext` for testing with the given Services and
    /// optional gateway active project.
    fn test_bot_context(
        services: Arc<Services>,
        gateway_active_project: Option<Arc<RwLock<String>>>,
        gateway_projects: Vec<String>,
        gateway_project_urls: BTreeMap<String, String>,
    ) -> BotContext {
        BotContext {
            services,
            matrix_user_id: make_user_id("@bot:example.com"),
            target_room_ids: vec![],
            allowed_users: vec![],
            history: Arc::new(TokioMutex::new(HashMap::new())),
            history_size: 20,
            bot_sent_event_ids: Arc::new(TokioMutex::new(std::collections::HashSet::new())),
            htop_sessions: Arc::new(TokioMutex::new(HashMap::new())),
            transport: Arc::new(crate::chat::transport::whatsapp::WhatsAppTransport::new(
                "test-phone".to_string(),
                "test-token".to_string(),
                "pipeline_notification".to_string(),
            )),
            timer_store: Arc::new(crate::service::timer::TimerStore::load(
                std::path::PathBuf::from("/tmp/timers.json"),
            )),
            gateway_active_project,
            gateway_projects,
            gateway_project_urls,
            pending_pipeline_events: Arc::new(TokioMutex::new(Vec::new())),
        }
    }

    #[test]
    fn bot_context_is_clone() {
        // BotContext must be Clone for the Matrix event handler injection.
        fn assert_clone<T: Clone>() {}
        assert_clone::<BotContext>();
    }

    #[tokio::test]
    async fn effective_project_root_standalone_returns_project_root() {
        let services = test_services(PathBuf::from("/projects/myapp"));
        let ctx = test_bot_context(services, None, vec![], BTreeMap::new());
        assert_eq!(
            ctx.effective_project_root().await,
            PathBuf::from("/projects/myapp")
        );
    }

    #[tokio::test]
    async fn effective_project_root_gateway_uses_active_project_subdir() {
        let services = test_services(PathBuf::from("/gateway"));
        let active = Arc::new(RwLock::new("huskies".to_string()));
        let ctx = test_bot_context(
            services,
            Some(Arc::clone(&active)),
            vec!["huskies".into(), "robot-studio".into()],
            BTreeMap::from([
                ("huskies".into(), "http://localhost:3001".into()),
                ("robot-studio".into(), "http://localhost:3002".into()),
            ]),
        );
        assert_eq!(
            ctx.effective_project_root().await,
            PathBuf::from("/gateway/huskies")
        );
    }

    #[tokio::test]
    async fn effective_project_root_gateway_reflects_project_switch() {
        let services = test_services(PathBuf::from("/gateway"));
        let active = Arc::new(RwLock::new("huskies".to_string()));
        let ctx = test_bot_context(
            services,
            Some(Arc::clone(&active)),
            vec!["huskies".into(), "robot-studio".into()],
            BTreeMap::from([
                ("huskies".into(), "http://localhost:3001".into()),
                ("robot-studio".into(), "http://localhost:3002".into()),
            ]),
        );

        assert_eq!(
            ctx.effective_project_root().await,
            PathBuf::from("/gateway/huskies")
        );

        *active.write().await = "robot-studio".to_string();

        assert_eq!(
            ctx.effective_project_root().await,
            PathBuf::from("/gateway/robot-studio")
        );
    }

    #[test]
    fn bot_context_has_no_require_verified_devices_field() {
        let services = test_services(PathBuf::from("/tmp"));
        let ctx = test_bot_context(services, None, vec![], BTreeMap::new());
        let _cloned = ctx.clone();
    }

    /// A bot command issued in gateway mode must round-trip over WebSocket
    /// (using the `bot.command` RPC method) and must NOT use HTTP transport.
    #[tokio::test]
    async fn proxy_bot_command_uses_websocket_not_http() {
        use futures::{SinkExt, StreamExt};
        use tokio::net::TcpListener;
        use tokio_tungstenite::tungstenite::Message as WsMsg;

        // Bind an ephemeral port for our mock WebSocket server.
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();

        // Spawn a minimal WS server: accept one connection, verify the
        // request uses the `bot.command` RPC method (not HTTP), and reply.
        let server = tokio::spawn(async move {
            let (tcp, _addr) = listener.accept().await.unwrap();
            let mut ws = tokio_tungstenite::accept_async(tcp).await.unwrap();
            while let Some(Ok(msg)) = ws.next().await {
                if let WsMsg::Text(text) = msg {
                    let req: serde_json::Value =
                        serde_json::from_str(&text).expect("valid JSON from proxy");
                    assert_eq!(
                        req["kind"], "rpc_request",
                        "transport must use rpc_request, not HTTP"
                    );
                    assert_eq!(req["method"], "bot.command");
                    assert_eq!(req["params"]["command"], "status");
                    let correlation_id = req["correlation_id"].clone();
                    let resp = serde_json::json!({
                        "kind": "rpc_response",
                        "correlation_id": correlation_id,
                        "ok": true,
                        "result": { "response": "all systems go" },
                    });
                    ws.send(WsMsg::Text(resp.to_string().into())).await.unwrap();
                    break;
                }
            }
        });

        let base_url = format!("http://127.0.0.1:{port}");
        let services = test_services(PathBuf::from("/gateway"));
        let active = Arc::new(RwLock::new("huskies".to_string()));
        let ctx = test_bot_context(
            services,
            Some(Arc::clone(&active)),
            vec!["huskies".into()],
            BTreeMap::from([("huskies".into(), base_url)]),
        );

        let result = ctx.proxy_bot_command("status", "").await;
        assert_eq!(
            result.as_deref(),
            Some("all systems go"),
            "proxy must return the response text from the rpc_response frame"
        );

        server.await.unwrap();
    }
}