mnote/rust/crates/mnote-web/src/acp_session_manager.rs

/// ACP Session Manager — session lifecycle management.
///
/// Wraps an [`AcpClient`] to provide strongly-typed session operations:
/// create, prompt, cancel, and receive typed events from the agent.
///
/// Reference:
/// - `reference-code/hermes-vscode-main/src/sessionManager.ts` (primary)
/// - `reference-code/hermes-vscode-main/src/protocol.ts` (dedup logic)

use crate::acp_client::AcpClient;
use crate::acp_types::{
    ContentBlock, SessionNewParams, SessionNewResult, SessionPromptParams,
    SessionPromptResult, SessionUpdate, ToolCallStatus,
};
use serde_json::{json, Value};
use std::sync::{Arc, Mutex};
use tracing::{debug, info, warn};

#[cfg(test)]
use tokio::time::{sleep, Duration};

// ── Events ───────────────────────────────────────────

/// Strongly-typed event emitted by the session manager when a `session/update` arrives.
#[derive(Debug, Clone)]
pub enum AcpSessionEvent {
    /// Streaming text from the agent's response message.
    TextDelta { text: String },
    /// Streaming reasoning/thinking text.
    ThoughtDelta { text: String },
    /// Tool call started.
    ToolCall {
        tool_call_id: String,
        title: String,
        kind: String,
        status: ToolCallStatus,
    },
    /// Tool call status update (with optional result content).
    ToolCallUpdate {
        tool_call_id: String,
        status: ToolCallStatus,
    },
    /// Context usage update.
    UsageUpdate { used: u64, size: u64 },
    /// Session metadata update (e.g. auto-title).
    SessionInfoUpdate { title: String },
    /// Plan entries update.
    PlanUpdate { entries: Vec<String> },
    /// Connection closed/error.
    Disconnected { reason: String },
}

/// Handler for session events.
pub type SessionEventHandler = Arc<dyn Fn(AcpSessionEvent) + Send + Sync + 'static>;

// ── Session state ────────────────────────────────────

#[derive(Debug, Clone, PartialEq)]
pub enum SessionState {
    Idle,
    Running,
    Cancelling,
    Closed,
}

// ── AcpSessionManager ────────────────────────────────

/// Manages ACP sessions — create, prompt, cancel, and event dispatch.
///
/// Currently supports one active session at a time. The internal [`AcpClient`]
/// handles the JSON-RPC wire protocol; this layer adds session semantics and
/// typed event dispatching.
pub struct AcpSessionManager {
    client: Arc<AcpClient>,
    session_id: Arc<Mutex<Option<String>>>,
    state: Arc<Mutex<SessionState>>,
    event_handler: Arc<Mutex<Option<SessionEventHandler>>>,
    /// Accumulated text for deduplication (per-turn).
    accumulated: Arc<Mutex<String>>,
    /// Whether we're currently inside a prompt (for dedup gating).
    in_prompt: Arc<Mutex<bool>>,
}

impl AcpSessionManager {
    /// Create a new session manager wrapping an existing [`AcpClient`].
    ///
    /// Registers an internal notification handler that dispatches
    /// `session/update` notifications as typed [`AcpSessionEvent`]s.
    pub fn new(client: Arc<AcpClient>) -> Self {
        let session_id: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
        let state: Arc<Mutex<SessionState>> = Arc::new(Mutex::new(SessionState::Idle));
        let event_handler: Arc<Mutex<Option<SessionEventHandler>>> =
            Arc::new(Mutex::new(None));
        let accumulated: Arc<Mutex<String>> = Arc::new(Mutex::new(String::new()));
        let in_prompt: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));

        // Wire up the ACP notification handler
        let session_id_clone = session_id.clone();
        let event_handler_clone = event_handler.clone();
        let accumulated_clone = accumulated.clone();
        let in_prompt_clone = in_prompt.clone();

        client.on_notification(move |method, params| {
            if method != "session/update" {
                return;
            }
            let sid = session_id_clone.lock().unwrap().clone();
            if let Some(ref session_id) = sid {
                if let Some(msg_sid) = params.get("sessionId").and_then(|v| v.as_str()) {
                    if msg_sid != session_id {
                        return; // not our session
                    }
                }
            }

            // Parse the update
            let update: SessionUpdate = match serde_json::from_value(
                params.get("update").cloned().unwrap_or(Value::Null),
            ) {
                Ok(u) => u,
                Err(e) => {
                    warn!("ACP session/update parse error: {e}");
                    return;
                }
            };

            let is_in_prompt = *in_prompt_clone.lock().unwrap();
            let event = Self::session_update_to_event(&update, &accumulated_clone, is_in_prompt);

            if let Some(ev) = event {
                let handler = event_handler_clone.lock().unwrap();
                if let Some(ref h) = *handler {
                    h(ev);
                }
            }
        });

        Self {
            client,
            session_id,
            state,
            event_handler,
            accumulated,
            in_prompt,
        }
    }

    /// Register an event handler for session events.
    /// Only one handler at a time — subsequent calls replace the previous.
    pub fn on_event<F>(&self, handler: F)
    where
        F: Fn(AcpSessionEvent) + Send + Sync + 'static,
    {
        let mut guard = self.event_handler.lock().unwrap();
        *guard = Some(Arc::new(handler));
    }

    /// Create a new ACP session.
    ///
    /// Sends `session/new` to the agent and stores the returned `sessionId`.
    /// The `page_context` is optional metadata about the current document.
    pub async fn create_session(
        &self,
        cwd: Option<&str>,
        _page_context: Option<Value>,
    ) -> Result<String, crate::acp_client::AcpError> {
        let params = SessionNewParams {
            cwd: cwd.map(|s| s.to_string()),
            mcp_servers: None,
        };

        let result: SessionNewResult = self.client.request("session/new", params).await?;
        let mut sid_guard = self.session_id.lock().unwrap();
        *sid_guard = Some(result.session_id.clone());
        info!("ACP session created: {}", result.session_id);

        Ok(result.session_id)
    }

    /// Send a prompt to the agent and stream events.
    ///
    /// The `prompt` is a list of content blocks (text + optional page context).
    /// Returns once the agent finishes (stopReason received) or on error.
    ///
    /// Sets `in_prompt = true` during the call to enable deduplication,
    /// then resets to `false` and clears accumulated text on completion.
    pub async fn run_prompt(
        &self,
        prompt: Vec<ContentBlock>,
    ) -> Result<SessionPromptResult, crate::acp_client::AcpError> {
        {
            let mut in_prompt = self.in_prompt.lock().unwrap();
            *in_prompt = true;
        }
        {
            let mut acc = self.accumulated.lock().unwrap();
            acc.clear();
        }
        {
            let mut state = self.state.lock().unwrap();
            *state = SessionState::Running;
        }

        let sid = self.session_id.lock().unwrap().clone();
        let session_id = sid.ok_or_else(|| {
            crate::acp_client::AcpError::Internal(
                "no session created yet — call create_session first".into(),
            )
        })?;

        let params = SessionPromptParams {
            session_id: session_id.clone(),
            prompt,
        };

        debug!("ACP session/prompt (session={})", session_id);
        let result: SessionPromptResult = self.client.request("session/prompt", params).await?;
        debug!(
            "ACP session/prompt done (session={}, stop_reason={:?})",
            session_id, result.stop_reason
        );

        {
            let mut state = self.state.lock().unwrap();
            *state = SessionState::Idle;
        }
        {
            let mut in_prompt = self.in_prompt.lock().unwrap();
            *in_prompt = false;
        }
        {
            let mut acc = self.accumulated.lock().unwrap();
            acc.clear();
        }

        Ok(result)
    }

    /// Cancel the current prompt.
    ///
    /// Sends `session/cancel` notification to the agent.
    pub async fn cancel(&self) -> Result<(), crate::acp_client::AcpError> {
        let sid = self.session_id.lock().unwrap().clone();
        let session_id = sid.ok_or_else(|| {
            crate::acp_client::AcpError::Internal("no active session".into())
        })?;

        {
            let mut state = self.state.lock().unwrap();
            *state = SessionState::Cancelling;
        }

        self.client
            .notification(
                "session/cancel",
                json!({ "sessionId": session_id }),
            )
            .await?;

        info!("ACP session cancelled: {}", session_id);
        Ok(())
    }

    /// Close the session and the underlying ACP client.
    pub async fn close(&self) {
        {
            let mut state = self.state.lock().unwrap();
            *state = SessionState::Closed;
        }

        if let Some(handler) = self.event_handler.lock().unwrap().take() {
            handler(AcpSessionEvent::Disconnected {
                reason: "session closed".into(),
            });
        }
    }

    /// Get the current session ID, if any.
    pub async fn session_id(&self) -> Option<String> {
        self.session_id.lock().unwrap().clone()
    }

    /// Get the current session state.
    pub async fn state(&self) -> SessionState {
        self.state.lock().unwrap().clone()
    }

    // ── Internal: session/update → event mapping ─────

    /// Convert a parsed [`SessionUpdate`] into an [`AcpSessionEvent`],
    /// applying text deduplication for streaming text.
    ///
    /// Reference: `hermes-vscode-main/src/protocol.ts` `extractTextContent()`,
    /// `deduplicateChunk()`, `parseToolCall()`, `parseToolCallUpdate()`,
    /// `parseUsageUpdate()`, `parseSessionInfoUpdate()`
    fn session_update_to_event(
        update: &SessionUpdate,
        accumulated: &Arc<Mutex<String>>,
        is_in_prompt: bool,
    ) -> Option<AcpSessionEvent> {
        match update {
            SessionUpdate::AgentMessageChunk { content, .. }
            | SessionUpdate::AgentThoughtChunk { content, .. } => {
                let discrim = match update {
                    SessionUpdate::AgentMessageChunk { .. } => "msg",
                    SessionUpdate::AgentThoughtChunk { .. } => "thought",
                    _ => unreachable!(),
                };

                // Deduplication (reference: protocol.ts deduplicateChunk)
                if is_in_prompt {
                    let mut acc = accumulated.lock().unwrap();
                    let text = &content.text;
                    let event = if text == acc.as_str() {
                        // Exact full resend → drop
                        None
                    } else if text.len() > 10 && text.starts_with(acc.as_str()) {
                        // Superset resend → emit only the tail
                        let new_part = text[acc.len()..].to_string();
                        if new_part.is_empty() {
                            None
                        } else {
                            *acc = text.clone();
                            Some(new_part)
                        }
                    } else if text.len() > 10 && acc.ends_with(text) {
                        // Partial resend → drop
                        None
                    } else {
                        // Normal delta
                        let new_acc = format!("{}{}", acc, text);
                        *acc = new_acc;
                        Some(text.clone())
                    };

                    return event.map(|t| match discrim {
                        "msg" => AcpSessionEvent::TextDelta { text: t },
                        "thought" => AcpSessionEvent::ThoughtDelta { text: t },
                        _ => unreachable!(),
                    });
                }

                // Not in prompt — emit directly (historical playback)
                let text = content.text.clone();
                Some(match discrim {
                    "msg" => AcpSessionEvent::TextDelta { text },
                    "thought" => AcpSessionEvent::ThoughtDelta { text },
                    _ => unreachable!(),
                })
            }

            SessionUpdate::ToolCall {
                tool_call_id,
                title,
                kind,
                status,
                ..
            } => {
                let title = title.clone().unwrap_or_else(|| "tool".into());
                let kind_str = match kind {
                    Some(k) => format!("{:?}", k).to_lowercase(),
                    None => "other".into(),
                };
                let status = status.clone().unwrap_or(ToolCallStatus::Pending);
                Some(AcpSessionEvent::ToolCall {
                    tool_call_id: tool_call_id.clone(),
                    title,
                    kind: kind_str,
                    status,
                })
            }

            SessionUpdate::ToolCallUpdate {
                tool_call_id,
                status,
                ..
            } => {
                let status = status.clone().unwrap_or(ToolCallStatus::Completed);
                Some(AcpSessionEvent::ToolCallUpdate {
                    tool_call_id: tool_call_id.clone(),
                    status,
                })
            }

            SessionUpdate::UsageUpdate { used, size, .. } => {
                Some(AcpSessionEvent::UsageUpdate {
                    used: *used,
                    size: *size,
                })
            }

            SessionUpdate::SessionInfoUpdate { title, .. } => {
                Some(AcpSessionEvent::SessionInfoUpdate {
                    title: title.clone(),
                })
            }

            SessionUpdate::Plan { entries, .. } => {
                let summaries: Vec<String> =
                    entries.iter().map(|e| e.content.clone()).collect();
                Some(AcpSessionEvent::PlanUpdate { entries: summaries })
            }

            SessionUpdate::Unknown { .. } => {
                warn!("ACP unknown session/update variant");
                None
            }
        }
    }
}

// ── Tests ────────────────────────────────────────────

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

    /// Creates a minimal ACP mock server for testing session operations.
    async fn spawn_mock_acp() -> Arc<AcpClient> {
        let script = r#"
import * as readline from 'node:readline';
import { stdin as input, stdout as output } from 'node:process';
const rl = readline.createInterface({ input, output, terminal: false });
rl.on('line', (line) => {
    const msg = JSON.parse(line);
    if (!msg.id) return;
    if (msg.method === 'session/new') {
        process.stdout.write(JSON.stringify({
            jsonrpc: '2.0', id: msg.id,
            result: { sessionId: 'test_session_1' }
        }) + '\n');
    } else if (msg.method === 'session/prompt') {
        const sessionId = msg.params?.sessionId || 'test';
        process.stdout.write(JSON.stringify({
            jsonrpc: '2.0',
            method: 'session/update',
            params: {
                sessionId,
                update: {
                    sessionUpdate: 'agent_message_chunk',
                    content: { type: 'text', text: 'Hello from mock ACP' }
                }
            }
        }) + '\n');
        process.stdout.write(JSON.stringify({
            jsonrpc: '2.0', id: msg.id,
            result: { stopReason: 'end_turn' }
        }) + '\n');
    } else if (msg.method === 'session/cancel') {
        // No response for notification
    } else if (msg.method === 'initialize') {
        process.stdout.write(JSON.stringify({
            jsonrpc: '2.0', id: msg.id,
            result: { protocolVersion: 1, agentCapabilities: {}, agentInfo: { name: 'mock', version: '1.0' }, authMethods: [] }
        }) + '\n');
    }
});
        "#;
        let dir = std::env::temp_dir();
        let script_path = dir.join("acp_session_test_mock.mjs");
        std::fs::write(&script_path, script).expect("write mock");
        let client = AcpClient::spawn("node", &[script_path.to_str().unwrap()])
            .await
            .expect("spawn");
        Arc::new(client)
    }

    #[tokio::test]
    async fn test_create_session() {
        let client = spawn_mock_acp().await;
        let mgr = AcpSessionManager::new(client);
        let sid = mgr
            .create_session(Some("/test"), None)
            .await
            .expect("create_session");
        assert_eq!(sid, "test_session_1");
        assert_eq!(mgr.session_id().await, Some("test_session_1".into()));
    }

    #[tokio::test]
    async fn test_run_prompt() {
        let client = spawn_mock_acp().await;
        let mgr = AcpSessionManager::new(client);
        mgr.create_session(Some("/test"), None)
            .await
            .expect("create_session");

        let prompt = vec![ContentBlock::Text {
            text: "Hello agent".into(),
        }];
        let result = mgr.run_prompt(prompt).await.expect("run_prompt");
        assert_eq!(
            format!("{:?}", result.stop_reason),
            "EndTurn".to_string()
        );
        // After prompt, state should be idle again
        assert_eq!(mgr.state().await, SessionState::Idle);
    }

    #[tokio::test]
    async fn test_cancel() {
        let client = spawn_mock_acp().await;
        let mgr = AcpSessionManager::new(client);
        mgr.create_session(Some("/test"), None)
            .await
            .expect("create_session");
        mgr.cancel().await.expect("cancel");
    }

    #[tokio::test]
    async fn test_event_handler() {
        use std::sync::atomic::{AtomicBool, Ordering};
        let client = spawn_mock_acp().await;
        let mgr = AcpSessionManager::new(client);
        let received = Arc::new(AtomicBool::new(false));
        let r = received.clone();

        mgr.on_event(move |ev| {
            if matches!(ev, AcpSessionEvent::TextDelta { .. }) {
                r.store(true, Ordering::SeqCst);
            }
        });

        mgr.create_session(Some("/test"), None)
            .await
            .expect("create_session");
        let prompt = vec![ContentBlock::Text {
            text: "Hello".into(),
        }];
        mgr.run_prompt(prompt).await.expect("run_prompt");

        // Give the notification handler time to process
        sleep(Duration::from_millis(200)).await;
        assert!(received.load(Ordering::SeqCst), "should have received TextDelta");
    }
}