Kigi-CLI/crates/codegen/kigi-shell/src/extensions/session_admin.rs

//! Session-administration extension handlers.
//!
//! Methods grouped here are operational/admin endpoints that mutate
//! persistent or shared agent state but are not part of the per-turn prompt
//! lifecycle:
//!
//! - `x.ai/session/rename`                  rename a session locally
//! - `x.ai/session/delete`                  delete a session locally
//! - `x.ai/session/update_mcp_servers`      mid-session MCP server swap
//! - `x.ai/session/fork`                    fork a session into a new one
//! - `x.ai/internal/reload_all_mcp_servers` config hot-reload, all sessions
//! - `x.ai/internal/reload_project_mcp_servers` config hot-reload, cwd-scoped
//! - `x.ai/internal/reload_skills`          skills file watcher fan-out
//! - `x.ai/internal/reload_models`          model list hot-reload from config.toml
//! - `x.ai/internal/reload_models_cache`    model catalog hot-reload from disk cache
//! - `x.ai/internal/auth_cleared`           auth hot-clear cleanup
//! - `x.ai/plugins/reload`                  rebuild shared plugin registry
//! - `x.ai/commands/list`                   list slash commands

use std::path::Path;
use std::sync::Arc;

use agent_client_protocol as acp;
use agent_client_protocol::Client as _;
use serde::Deserialize;

use super::{ExtResult, parse_params, to_raw_response};
use crate::agent::MvpAgent;
use crate::session::persistence::list_summaries;
use crate::session::storage::StorageAdapter;
use crate::session::storage::jsonl::JsonlStorageAdapter;
use crate::session::unified_list::SessionKind;
use crate::session::{ExtMethodResult, SessionCommand};

#[tracing::instrument(skip_all, fields(method = %args.method))]
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    match args.method.as_ref() {
        "x.ai/session/rename" => handle_session_rename(agent, args).await,
        "x.ai/session/delete" => handle_session_delete(agent, args).await,
        "x.ai/session/update_mcp_servers" => handle_update_mcp_servers(agent, args).await,
        "x.ai/session/fork" => handle_session_fork(agent, args).await,
        "x.ai/internal/reload_all_mcp_servers" => handle_reload_all_mcp_servers(agent).await,
        "x.ai/internal/reload_project_mcp_servers" => {
            handle_reload_project_mcp_servers(agent, args).await
        }
        "x.ai/internal/reload_skills" => handle_reload_skills(agent),
        "x.ai/internal/reload_models" => handle_reload_models(agent),
        "x.ai/internal/reload_models_cache" => handle_reload_models_cache(agent),
        "x.ai/internal/auth_cleared" => handle_auth_cleared(agent),
        "x.ai/plugins/reload" => handle_plugins_reload(agent).await,
        "x.ai/commands/list" => handle_commands_list(agent, args).await,
        _ => Err(acp::Error::method_not_found()),
    }
}

// session/rename

/// Handles renaming a session.
async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    #[derive(Deserialize)]
    #[serde(rename_all = "camelCase")]
    struct RenameRequest {
        session_id: String,
        title: String,
        #[serde(default)]
        cwd: Option<String>,
        #[serde(default)]
        kind: SessionKind,
    }

    let mut req: RenameRequest = parse_params(args)?;
    // Manual titles must be non-blank: `Summary.title_is_manual` binds to a
    // real `generated_title`, so reject whitespace-only input at the boundary.
    req.title = req.title.trim().to_string();
    if req.title.is_empty() {
        return Err(acp::Error::invalid_request().data("title must not be blank"));
    }

    if req.kind == SessionKind::Chat {
        return Err(acp::Error::invalid_request()
            .data("chat conversations are not available in kigi (local sessions only)"));
    }

    let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str()));

    // Find the session info, scoping to cwd if provided
    let summaries = list_summaries(req.cwd.as_deref())
        .await
        .map_err(|e| acp::Error::internal_error().data(format!("failed to list sessions: {e}")))?;

    let summary = summaries
        .iter()
        .find(|s| s.info.id == session_id)
        .ok_or_else(|| {
            acp::Error::invalid_request().data(format!("session not found: {}", req.session_id))
        })?;

    let info = summary.info.clone();

    // Update the session title in local storage
    let storage = JsonlStorageAdapter::default();
    storage
        .update_session_title(&info, req.title.clone())
        .await
        .map_err(|e| {
            acp::Error::internal_error().data(format!("failed to update session title: {e}"))
        })?;

    // Update session search index with new title
    crate::session::storage::search::notify_session_updated(&info.id.to_string(), &info.cwd);

    // Send a SessionSummaryGenerated notification so the TUI updates its title
    notify_session_title(agent, session_id, &req.title).await;

    // Hook 2: update session replica with summary (fire-and-forget)
    if let Some(client) = agent.session_registry_client() {
        let sid = req.session_id.to_string();
        let title = Some(req.title.clone());
        tokio::spawn(async move {
            let update = crate::agent::session_registry_client::UpdateRequest {
                summary: title,
                first_prompt: None,
                last_turn_number: None,
                repo_head_at_end: None,
                restorable_turn_number: None,
            };
            if let Err(e) = client.update(&sid, &update).await {
                tracing::warn!(error = %e, "session registry summary update failed (non-fatal)");
            }
        });
    }

    tracing::info!(session_id = %req.session_id, title = %req.title, "Session renamed");

    to_raw_response(&serde_json::json!({ "success": true }))
}

/// Notify connected clients of a session's new title via
/// `SessionSummaryGenerated`.
async fn notify_session_title(agent: &MvpAgent, session_id: acp::SessionId, title: &str) {
    use crate::extensions::notification::{SessionNotification, SessionUpdate};

    let notification = SessionNotification {
        session_id,
        update: SessionUpdate::SessionSummaryGenerated {
            session_summary: title.to_owned(),
        },
        meta: None,
    };
    if let Ok(params) = serde_json::value::to_raw_value(&notification) {
        let ext_notification =
            acp::ExtNotification::new("x.ai/session_notification", params.into());
        let _ = agent.gateway.ext_notification(ext_notification).await;
    }
}

// session/delete

/// Delete a session from history.
async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    #[derive(Deserialize)]
    #[serde(rename_all = "camelCase")]
    struct DeleteRequest {
        session_id: String,
        #[serde(default)]
        cwd: Option<String>,
        #[serde(default)]
        kind: SessionKind,
    }

    let req: DeleteRequest = parse_params(args)?;

    if req.kind == SessionKind::Chat {
        return Err(acp::Error::invalid_request()
            .data("chat conversations are not available in kigi (local sessions only)"));
    }

    let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str()));

    // Local disk + FTS eviction. Mirrored by the `kigi sessions delete <id>`
    // CLI path.
    crate::session::persistence::delete_session_history(&req.session_id, req.cwd.as_deref())
        .await
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))?;

    // If an in-memory live session exists for this id (e.g. the user
    // deleted history for a session that is still open in another agent
    // or the current one), shut it down and drop the MvpAgent bookkeeping
    // so we don't leave a live actor whose on-disk/FTS state is gone.
    if agent.sessions.borrow().contains_key(&session_id) {
        agent.request_session_shutdown(&session_id);
        agent.remove_session(&session_id);
    }

    tracing::info!(session_id = %req.session_id, "Session deleted");

    to_raw_response(&serde_json::json!({ "success": true }))
}

// session/update_mcp_servers

async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    #[derive(Deserialize)]
    #[serde(rename_all = "camelCase")]
    struct Params {
        session_id: acp::SessionId,
        mcp_servers: Vec<acp::McpServer>,
    }

    let params: Params = parse_params(args)?;

    let (handle, cwd) = {
        let sessions = agent.sessions.borrow();
        let h = sessions
            .get(&params.session_id)
            .cloned()
            .ok_or_else(|| acp::Error::invalid_params().data("unknown session id"))?;
        let cwd = std::path::PathBuf::from(&h.info.cwd);
        (h, cwd)
    };

    let managed = agent.get_managed_mcp_configs().await;
    let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
        params.mcp_servers.clone(),
        &cwd,
        &managed,
        agent.plugin_registry_handle().snapshot().as_deref(),
        &agent.cfg.borrow().compat_resolved,
    );

    let (tx, rx) = tokio::sync::oneshot::channel();
    handle
        .cmd_tx
        .send(SessionCommand::UpdateMcpServers {
            mcp_servers: merged,
            respond_to: tx,
        })
        .map_err(|_| acp::Error::internal_error().data("session closed"))?;

    // Wait for the session actor to finish MCP re-initialization.
    rx.await
        .map_err(|_| acp::Error::internal_error().data("session closed"))?
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))?;

    // Persist the new client set on the handle so config hot-reloads
    // (`reload_all_mcp_servers` / `reload_project_mcp_servers`) re-merge from
    // the client's latest intent rather than the `session/new` snapshot —
    // otherwise a reload would resurrect servers the client just removed
    // (or drop ones it just added).
    if let Some(h) = agent.sessions.borrow_mut().get_mut(&params.session_id) {
        h.initial_client_mcp_servers = params.mcp_servers;
    }

    ExtMethodResult::success(serde_json::json!({ "ok": true }))
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

// internal/reload_skills

/// Reload skills for ALL active sessions. Called by the skills file watcher
/// (via ACP injection from `app.rs`) when `SKILL.md` files change.
fn handle_reload_skills(agent: &MvpAgent) -> ExtResult {
    let session_ids: Vec<acp::SessionId> = agent.sessions.borrow().keys().cloned().collect();
    for sid in &session_ids {
        if let Some(handle) = agent.sessions.borrow().get(sid).cloned() {
            let _ = handle.cmd_tx.send(SessionCommand::ReloadSkills);
        }
    }
    ExtMethodResult::success(serde_json::json!({ "reloaded": session_ids.len() }))
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

// internal/reload_all_mcp_servers

/// Reload MCP servers for ALL active sessions. Called by the config
/// hot-reload watcher when `[mcp_servers]` changes in config.toml.
async fn handle_reload_all_mcp_servers(agent: &MvpAgent) -> ExtResult {
    let session_ids: Vec<acp::SessionId> = agent.sessions.borrow().keys().cloned().collect();

    if session_ids.is_empty() {
        return ExtMethodResult::success(serde_json::json!({ "updated": 0 }))
            .to_ext_response()
            .map_err(|e| acp::Error::internal_error().data(e.to_string()));
    }

    let managed = agent.get_managed_mcp_configs().await;
    let mut updated = 0u32;
    for session_id in &session_ids {
        let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else {
            continue;
        };
        let cwd = std::path::PathBuf::from(&handle.info.cwd);
        let compat = agent.cfg.borrow().compat_resolved;
        // Re-seed the merge with the session's original client-provided MCP
        // servers (e.g. a managed connector injected at `session/new` by a
        // client session binding). `merge_managed_mcp_servers` already
        // re-reads every disk source (config.toml, plugins, ~/.claude.json,
        // ~/.cursor/mcp.json, .mcp.json) internally, so passing
        // `load_mcp_servers()` output here was redundant — and silently
        // dropped client servers that exist in no on-disk config, tearing
        // them down on every config hot-reload.
        let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
            handle.initial_client_mcp_servers.clone(),
            &cwd,
            &managed,
            agent.plugin_registry_handle().snapshot().as_deref(),
            &compat,
        );

        let (tx, _rx) = tokio::sync::oneshot::channel();
        if handle
            .cmd_tx
            .send(SessionCommand::UpdateMcpServers {
                mcp_servers: merged,
                respond_to: tx,
            })
            .is_ok()
        {
            updated += 1;
        }
    }

    tracing::info!(
        updated,
        total = session_ids.len(),
        "reloaded MCP servers for active sessions"
    );
    ExtMethodResult::success(serde_json::json!({ "updated": updated }))
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

// internal/reload_project_mcp_servers

/// Reload MCP servers for sessions whose `cwd` matches (or sits beneath)
/// the project root passed in `params.cwd`. Called by the config
/// hot-reload watcher when `<cwd>/.kigi/config.toml`,
/// `<cwd>/.mcp.json`, or `<cwd>/.claude.json` changes.
///
/// Sessions in unrelated cwds are intentionally NOT touched — that is
/// the whole point of [`crate::config::reloader::ConfigUpdate::
/// ProjectMcpServersChanged`] being a per-cwd variant. The legacy
/// [`handle_reload_all_mcp_servers`] is still the fan-out for global
/// `~/.kigi/config.toml` edits.
async fn handle_reload_project_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    #[derive(Deserialize)]
    struct Params {
        cwd: String,
    }

    let params: Params = parse_params(args)?;
    let target_cwd = std::path::PathBuf::from(&params.cwd);

    // Collect (session_id, cwd) pairs once so we don't hold the
    // `sessions` RefCell borrow across `.await` points.
    let session_ids: Vec<(acp::SessionId, std::path::PathBuf)> = agent
        .sessions
        .borrow()
        .iter()
        .map(|(sid, h)| (sid.clone(), std::path::PathBuf::from(&h.info.cwd)))
        .filter(|(_, cwd)| cwd_matches(cwd, &target_cwd))
        .collect();

    if session_ids.is_empty() {
        return ExtMethodResult::success(serde_json::json!({ "updated": 0 }))
            .to_ext_response()
            .map_err(|e| acp::Error::internal_error().data(e.to_string()));
    }

    let managed = agent.get_managed_mcp_configs().await;
    let mut updated = 0u32;
    for (session_id, cwd) in &session_ids {
        let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else {
            continue;
        };
        // See `handle_reload_all_mcp_servers`: seed with the session's
        // client-provided servers, not `load_mcp_servers()` — the merge
        // re-reads all disk sources itself, and client-provided servers
        // (session bindings) must survive config hot-reloads.
        let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
            handle.initial_client_mcp_servers.clone(),
            cwd,
            &managed,
            agent.plugin_registry_handle().snapshot().as_deref(),
            &agent.cfg.borrow().compat_resolved,
        );

        let (tx, _rx) = tokio::sync::oneshot::channel();
        if handle
            .cmd_tx
            .send(SessionCommand::UpdateMcpServers {
                mcp_servers: merged,
                respond_to: tx,
            })
            .is_ok()
        {
            updated += 1;
        }
    }

    tracing::info!(
        updated,
        total = session_ids.len(),
        cwd = %target_cwd.display(),
        "reloaded project MCP servers for matching sessions"
    );
    ExtMethodResult::success(serde_json::json!({ "updated": updated }))
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

/// Returns `true` iff `session_cwd` equals `target_cwd` or sits
/// beneath it (so a `<repo>/` edit reloads `<repo>/subdir/` sessions
/// too).
///
/// This uses `Path::starts_with`, which is
/// **component-aware** — `/repo-test` does NOT match `/repo` even
/// though the byte prefix matches. That is the desired behavior
/// (component-aware avoids the `/foo-bar` ⊂ `/foo` foot-gun). Paths
/// come from `SessionInfo::cwd` (always absolute) and the watcher's
/// emitted path (also absolute), so no canonicalization is needed
/// here. The `==` short-circuit is redundant (`Path::starts_with` is
/// reflexive) but kept for an explicit zero-allocation fast path.
fn cwd_matches(session_cwd: &std::path::Path, target_cwd: &std::path::Path) -> bool {
    session_cwd == target_cwd || session_cwd.starts_with(target_cwd)
}

// internal/reload_models

/// Re-resolve the agent model list from config.toml. Called by the config
/// hot-reload watcher when `[model.*]` or `[models]` changes.
///
/// Re-reads config from disk, re-runs the same resolution logic as
/// `new_with_models()` for user TOML config entries, and swaps the model list
/// in-place. Prefetched (API) and default models are NOT re-fetched -- only
/// BYOK entries from config are updated.
fn handle_reload_models(agent: &MvpAgent) -> ExtResult {
    let disk_config = crate::config::load_effective_config()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))?;

    let toml_config = crate::agent::config::Config::new_from_toml_cfg(&disk_config)
        .map_err(|e| acp::Error::internal_error().data(e))?;

    // Merge TOML-derived model fields into the agent's in-memory config so
    // runtime-only fields (#[serde(skip)]: remote_settings, endpoints, CLI
    // flags) are preserved. Only model-related TOML fields are refreshed.
    {
        let agent_config = agent.cfg.borrow();
        let overrides = crate::config::ModelOverrideConfig::resolve(
            agent_config.web_search_model_override.as_deref(),
            agent_config.session_summary_model_override.as_deref(),
            &disk_config,
            agent_config.remote_settings.as_ref(),
        );
        drop(agent_config);
        let mut agent_config = agent.cfg.borrow_mut();
        agent_config.models = toml_config.models.clone();
        agent_config.config_models = toml_config.config_models.clone();
        agent_config.web_search_model = overrides.web_search;
        agent_config.session_summary_model = overrides.session_summary;
        agent_config.image_description_model = overrides.image_description;
        agent_config.prompt_suggest_model_pin = overrides.prompt_suggestion;
    }
    // Recompute the campaign overlay + `pre_campaign_default` (the catalog-miss
    // fallback) so reload matches spawn; `new_from_toml_cfg` reset it to None.
    {
        let mut agent_config = agent.cfg.borrow_mut();
        crate::util::config::sync_campaign_fields(&mut agent_config);
    }
    let merged_config = agent.cfg.borrow().clone();

    agent.models_manager.apply_config(merged_config);

    let count = agent.models_manager.models().len();
    tracing::info!(count, "model list reloaded from config.toml");
    ExtMethodResult::success(serde_json::json!({ "models": count }))
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

// internal/reload_models_cache

/// Hot-reload the model catalog from `~/.kigi/models_cache.json` after an
/// external write detected by the config watcher.
///
/// Routed through the agent's ACP stream (injected by the
/// `ConfigUpdate::ModelsCacheChanged` arm in `agent/app.rs`) instead of being
/// applied directly on the manager from the config-update task: stream
/// requests are processed in order, so when `config.toml` and
/// `models_cache.json` change in the same watcher batch this runs strictly
/// after `reload_models`' `apply_config` accepted or rejected the new config,
/// rather than rebuilding the catalog and notifying clients mid-flight.
fn handle_reload_models_cache(agent: &MvpAgent) -> ExtResult {
    agent.models_manager.reload_from_disk_cache();
    ExtMethodResult::success(serde_json::json!({ "reloaded": true }))
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

fn handle_auth_cleared(agent: &MvpAgent) -> ExtResult {
    agent.disable_managed_gateway_tools_and_refresh_sessions();
    ExtMethodResult::success(serde_json::json!({ "ok": true }))
        .to_ext_response()
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
}

// plugins/reload

async fn handle_plugins_reload(agent: &MvpAgent) -> ExtResult {
    // Rebuild the shared registry so future/new sessions clone the latest.
    let session_cwd = agent
        .sessions
        .borrow()
        .values()
        .next()
        .map(|h| std::path::PathBuf::from(&h.info.cwd));
    let mut plugins = agent.cfg.borrow().plugins.clone();
    plugins.merge_claude_enabled_plugins(session_cwd.as_deref());
    let disk_cfg = plugins.to_discovery_config();
    // Folder-trust gates repo-local project plugins (hooks/MCP). Resolve and
    // record the verdict for this cwd (honoring the real remote), then gate
    // plugins on it.
    let project_trusted = session_cwd.as_deref().is_some_and(|c| {
        let remote_settings = agent.cfg.borrow().remote_settings.clone();
        crate::agent::folder_trust::resolve_and_record(c, remote_settings.as_ref(), false)
    });
    // Explicit desktop `x.ai/plugins/reload`: force a full local-install re-copy.
    agent
        .plugin_registry_handle()
        .reload(session_cwd.as_deref(), &disk_cfg, project_trusted, true);

    // Eagerly fan out the new registry to every live session: each adopts a
    // cwd-correct snapshot (hooks + MCP + skills + client slash-command
    // catalog), the same refresh the originating session of a reload gets.
    agent.broadcast_plugin_registry_to_sessions(None);

    super::to_ext_response(Ok(serde_json::json!({"ok": true})))
}

// commands/list

async fn handle_commands_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    let req: crate::session::slash_commands::ListCommandsRequest = parse_params(args)?;

    let skills_config = agent.cfg.borrow().skills.clone();
    let compat = agent.cfg.borrow().compat_resolved;
    let availability = agent.command_availability();

    // For a given cwd, compute the plugin registry the same way a session would
    // at spawn time (via build_for_cwd) and the same way reload_plugins_impl does
    // (ancestor project config walk + vendor compat merge). This is required so
    // that `x.ai/commands/list` (the pull used by grok-desktop after session
    // start) returns plugin-provided slash commands for the target cwd.
    //
    // The shared snapshot is only populated at agent boot (using process CWD)
    // and by explicit reloads. In desktop<->docker (and ssh) setups the agent's
    // launch CWD is unrelated to the user's chosen workspace dir, so relying on
    // snapshot() alone meant the post-start pull returned no project plugin
    // skills until the user manually reloaded.
    let plugin_reg = if let Some(cwd_str) = &req.cwd {
        let cwd = Path::new(cwd_str);

        // Folder-trust gates repo-local project plugins (hooks/MCP). Resolve and
        // record the verdict for this cwd (honoring the real remote) BEFORE the
        // plugins-config read below: that read gates its project-paths merge on
        // the recorded verdict, and a cold cwd (client-supplied, no session
        // resolve yet) must not first take the gate's remote-less backstop —
        // that would record a kill-switch-blind deny no later resolve can lift.
        let remote_settings = agent.cfg.borrow().remote_settings.clone();
        let project_trusted =
            crate::agent::folder_trust::resolve_and_record(cwd, remote_settings.as_ref(), false);

        // Effective [plugins] config (global + ancestor project configs +
        // vendor compat merge), shared with reload_plugins_impl and the eager
        // fan-out so the menu agrees with each session's registry for this cwd.
        let disk_cfg = crate::config::resolve_effective_plugins_config(cwd).to_discovery_config();

        // Fresh discovery for *this* cwd (includes .kigi/plugins under it, plus
        // the cli --plugin-dir dirs). Does not mutate the shared snapshot.
        agent
            .plugin_registry_handle()
            .build_for_cwd(cwd, &disk_cfg, &[], project_trusted)
    } else {
        // No cwd: global/user skills only (pre-session case). Use the boot snapshot.
        agent.plugin_registry_handle().snapshot()
    };

    let response = crate::session::slash_commands::list_commands(
        req.cwd.as_deref(),
        &skills_config,
        plugin_reg.as_deref(),
        availability,
        compat,
    )
    .await;
    Ok(acp::ExtResponse::new(Arc::from(
        serde_json::value::to_raw_value(&response)?,
    )))
}

// session/fork

async fn handle_session_fork(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
    use crate::session::fork::{ForkSessionRequest, fork_session};

    let request: ForkSessionRequest = parse_params(args)?;

    let response = fork_session(request)
        .await
        .map_err(|e| acp::Error::internal_error().data(e.to_string()))?;

    to_raw_response(&response)
}