aether-agent-cli 0.9.2

use super::agent::acp_agent_builder;
use super::protocol::notify;
use acp_utils::notifications::{
    AetherCapabilities, AuthMethodsUpdatedParams, McpRequest, PromptSearchParams, PromptSearchResponse,
    RemoteServerInfo, SessionDisplayMeta, SessionPreviewParams, SessionPreviewResponse, WorkspaceListParams,
    WorkspaceListResponse, WorkspaceMoveParams, WorkspaceMoveResponse,
};
use aether_auth::OAuthCredentialStorage;
use aether_telemetry::TelemetryRuntime;
use agent_client_protocol::schema::ProtocolVersion;
use agent_client_protocol::schema::v2::{
    self as acp, AgentCapabilities, AuthMethod, CancelSessionNotification, CloseSessionRequest, CloseSessionResponse,
    Implementation, InitializeRequest, InitializeResponse, ListSessionsRequest, ListSessionsResponse, LoginAuthRequest,
    LoginAuthResponse, LogoutAuthRequest, LogoutAuthResponse, McpCapabilities, NewSessionRequest, NewSessionResponse,
    PromptAudioCapabilities, PromptCapabilities, PromptEmbeddedContextCapabilities, PromptImageCapabilities,
    PromptRequest, PromptResponse, ResumeSessionRequest, ResumeSessionResponse, SetSessionConfigOptionRequest,
    SetSessionConfigOptionResponse,
};
use agent_client_protocol::util::internal_error;
use agent_client_protocol::{Agent, Client, ConnectTo, ConnectionTo, Error, Responder};
use llm::catalog::{LlmModel, ModelSpec};
use llm::{ContentBlock, ProviderConnectionOverrides};
use mcp_utils::client::{client_capabilities, client_capabilities_for};
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::sync::{Mutex, oneshot, watch};
use tokio::task::spawn_blocking;
use tokio_util::sync::CancellationToken;
use tracing::{error, info};

use super::protocol::content::map_acp_to_content_blocks;
use super::session::actor::SessionCommand;
#[cfg(any(test, feature = "testing"))]
use super::session::actor::{SessionActor, SessionActorInit};
use super::session::config_setting::ConfigSetting;
use super::session::factory::SessionFactory;
use super::session::model::supports_prompt_audio;
use super::session::{SessionRegistry, paginate_summaries};
use crate::resolve::InitialSessionSelection;
use crate::settings_args::SettingsSourceArgs;
use crate::workspace::{WorkspaceError, WorkspaceManager};
use aether_sessions::{SessionStore, SessionStoreError};

#[async_trait::async_trait]
pub(crate) trait ProviderLogin: Send + Sync {
    async fn login(&self, store: &dyn OAuthCredentialStorage) -> Result<(), llm::LlmError>;
}

/// Host control plane owning one active session actor independently of its client.
pub(crate) struct AcpState {
    client_slot: ClientSlot,
    login: Arc<dyn ProviderLogin>,
    registry: SessionRegistry,
    stop: CancellationToken,
    session_store: Arc<SessionStore>,
    workspace_manager: Arc<WorkspaceManager>,
    oauth_credential_store: Arc<dyn OAuthCredentialStorage>,
    factory: SessionFactory,
    telemetry: Option<Arc<TelemetryRuntime>>,
    mcp_capabilities: Mutex<rmcp::model::ClientCapabilities>,
    cwd: PathBuf,
}

pub(crate) struct AcpStateConfig {
    pub(crate) session_store: Arc<SessionStore>,
    pub(crate) workspace_manager: Arc<WorkspaceManager>,
    pub(crate) oauth_credential_store: Arc<dyn OAuthCredentialStorage>,
    pub(crate) initial_selection: InitialSessionSelection,
    pub(crate) settings_source: SettingsSourceArgs,
    pub(crate) provider_connections: ProviderConnectionOverrides,
    pub(crate) telemetry: Option<Arc<TelemetryRuntime>>,
    pub(crate) runtime_factory: Option<Arc<dyn super::session::runtime::RuntimeFactory>>,
    pub(crate) cwd: PathBuf,
}

#[derive(Default)]
pub(crate) struct ClientSlot {
    occupied: AtomicBool,
    generation: watch::Sender<u64>,
}

impl ClientSlot {
    #[cfg(any(test, feature = "testing"))]
    pub(crate) fn subscribe(&self) -> watch::Receiver<u64> {
        self.generation.subscribe()
    }
}

#[must_use = "the client must be released after serving or a failed handshake"]
pub(crate) struct ClientGuard {
    state: Option<Arc<AcpState>>,
}

impl ClientGuard {
    pub(crate) async fn serve(
        self,
        transport: impl ConnectTo<Agent> + 'static,
        stop: CancellationToken,
    ) -> Result<(), Error> {
        let state = self.state.as_ref().expect("attached client").clone();
        state.serve_attached(transport, stop, self, agent_client_protocol::NullRun).await
    }

    pub(crate) async fn release(mut self) {
        let state = self.state.as_ref().expect("attached client");
        state.release_client().await;
        self.state = None;
    }
}

impl Drop for ClientGuard {
    fn drop(&mut self) {
        if let Some(state) = self.state.take() {
            tokio::spawn(async move { state.release_client().await });
        }
    }
}

impl AcpState {
    pub(crate) fn try_attach(self: &Arc<Self>) -> Option<ClientGuard> {
        if self.stop.is_cancelled() {
            return None;
        }
        self.client_slot.occupied.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed).ok()?;
        Some(ClientGuard { state: Some(self.clone()) })
    }

    #[cfg(any(test, feature = "testing"))]
    pub(crate) fn client_slot(&self) -> &ClientSlot {
        &self.client_slot
    }

    pub(crate) async fn serve(
        self: &Arc<Self>,
        transport: impl ConnectTo<Agent> + 'static,
        stop: CancellationToken,
    ) -> Result<(), Error> {
        let guard = self.try_attach().ok_or_else(|| internal_error("client already attached"))?;
        guard.serve(transport, stop).await
    }

    #[cfg(any(test, feature = "testing"))]
    pub(crate) async fn serve_with(
        self: &Arc<Self>,
        transport: impl ConnectTo<Agent> + 'static,
        stop: CancellationToken,
        on_connected: oneshot::Sender<ConnectionTo<Client>>,
    ) -> Result<(), Error> {
        let guard = self.try_attach().ok_or_else(|| internal_error("client already attached"))?;
        self.serve_attached(transport, stop, guard, acp_utils::testing::CaptureConnection(on_connected)).await
    }

    pub(crate) fn new(config: AcpStateConfig) -> Self {
        Self::with_login(config, Arc::new(CodexLogin))
    }

    pub(crate) fn with_login(config: AcpStateConfig, login: Arc<dyn ProviderLogin>) -> Self {
        let factory = SessionFactory::new(
            config.settings_source,
            config.provider_connections,
            Arc::clone(&config.oauth_credential_store),
            Arc::clone(&config.session_store),
            config.initial_selection,
            config.telemetry.as_ref().map(|runtime| runtime.observer_factory()),
            config.runtime_factory,
        );
        Self {
            client_slot: ClientSlot::default(),
            login,
            registry: SessionRegistry::new(),
            stop: CancellationToken::new(),
            session_store: config.session_store,
            workspace_manager: config.workspace_manager,
            oauth_credential_store: config.oauth_credential_store,
            factory,
            telemetry: config.telemetry,
            mcp_capabilities: Mutex::new(client_capabilities()),
            cwd: config.cwd,
        }
    }

    pub(crate) async fn initialize(&self, args: InitializeRequest) -> Result<InitializeResponse, Error> {
        info!("Received initialize request: {:?}", args);
        *self.mcp_capabilities.lock().await = mcp_client_capabilities(&args.capabilities);
        let auth_methods = build_auth_methods(self.oauth_credential_store.as_ref());
        let available = self.factory.available_models().await.to_vec();
        let prompt_capabilities = prompt_capabilities_for_models(&available);
        let aether_capabilities =
            AetherCapabilities { prompt_search: true, session_preview: true, workspace_move: true };
        let session_capabilities = acp::SessionCapabilities::new()
            .prompt(prompt_capabilities)
            .mcp(McpCapabilities::new().stdio(acp::McpStdioCapabilities::new()).http(acp::McpHttpCapabilities::new()))
            .meta(Some(aether_capabilities.to_meta()));

        let session_id = self.registry.session_id().await;
        let cwd = session_id
            .as_ref()
            .and_then(|id| self.session_store.session_cwd(id.0.as_ref()))
            .unwrap_or_else(|| self.cwd.clone());
        let remote_meta = Some(RemoteServerInfo { cwd, session_id }.to_meta());
        Ok(InitializeResponse::new(ProtocolVersion::V2, Implementation::new("Aether", "0.1.0"))
            .meta(remote_meta)
            .capabilities(AgentCapabilities::new().session(session_capabilities))
            .auth_methods(auth_methods))
    }

    pub(crate) async fn login(
        &self,
        args: LoginAuthRequest,
        cx: &ConnectionTo<Client>,
    ) -> Result<LoginAuthResponse, Error> {
        info!("Received login request: {:?}", args);
        let method_id = args.method_id.0.as_ref();
        match method_id {
            "codex" => {
                self.login.login(self.oauth_credential_store.as_ref()).await.map_err(|e| {
                    error!("OAuth flow failed for {method_id}: {e}");
                    Error::internal_error()
                })?;
            }
            _ => return Err(Error::invalid_params()),
        }
        self.broadcast_auth_state(cx).await;
        Ok(LoginAuthResponse::new())
    }

    pub(crate) async fn new_session(
        &self,
        req: NewSessionRequest,
        cx: &ConnectionTo<Client>,
    ) -> Result<NewSessionResponse, Error> {
        let mcp_capabilities = self.mcp_capabilities.lock().await.clone();
        let prepared = self.factory.prepare_new(req, cx, mcp_capabilities).await?;
        self.registry.stop().await;
        let created = prepared.start().await?;
        let response = NewSessionResponse::new(created.session_id.clone()).config_options(created.config_options);
        self.registry.register(&created.session_id, created.handle).await;
        Ok(response)
    }

    pub(crate) async fn logout(
        &self,
        _req: LogoutAuthRequest,
        cx: &ConnectionTo<Client>,
    ) -> Result<LogoutAuthResponse, Error> {
        self.oauth_credential_store.delete("codex").await.map_err(|e| internal_error(e.to_string()))?;
        self.broadcast_auth_state(cx).await;
        Ok(LogoutAuthResponse::new())
    }

    pub(crate) async fn resume_session(
        &self,
        req: ResumeSessionRequest,
        cx: &ConnectionTo<Client>,
    ) -> Result<ResumeSessionResponse, Error> {
        let replay = match &req.replay_from {
            Some(acp::ReplayFrom::Start(_)) => true,
            None => false,
            Some(_) => return Err(Error::invalid_params()),
        };

        if let Some(sender) = self.registry.lookup(Some(req.session_id.0.as_ref())).await {
            let available = self.factory.available_models().await.to_vec();
            let (config_tx, config_rx) = oneshot::channel();

            sender
                .send(SessionCommand::Attach {
                    connection: cx.clone(),
                    cwd: req.cwd.into_inner(),
                    mcp_servers: req.mcp_servers,
                    replay,
                    available,
                    reply: config_tx,
                })
                .await
                .map_err(|_| Error::internal_error())?;

            let options = config_rx.await.map_err(|_| Error::internal_error())??;
            return Ok(ResumeSessionResponse::new().config_options(options));
        }

        let mcp_capabilities = self.mcp_capabilities.lock().await.clone();
        let prepared = self.factory.prepare_resume(req, cx, mcp_capabilities, replay).await?;
        self.registry.stop().await;
        let created = prepared.start().await?;
        let response = ResumeSessionResponse::new().config_options(created.config_options);
        self.registry.register(&created.session_id, created.handle).await;
        Ok(response)
    }

    pub(crate) async fn close_session(&self, req: CloseSessionRequest) -> Result<CloseSessionResponse, Error> {
        let session_id = req.session_id.0.to_string();
        if self.registry.lookup(Some(&session_id)).await.is_none() {
            error!("Session not found for close: {session_id}");
            return Err(Error::invalid_params().data(format!("unknown session: {session_id}")));
        }
        self.registry.stop().await;
        Ok(CloseSessionResponse::new())
    }

    pub(crate) fn list_sessions(&self, args: &ListSessionsRequest) -> Result<ListSessionsResponse, Error> {
        info!("Listing sessions, cwd filter: {:?}, cursor: {:?}", args.cwd, args.cursor);
        let mut summaries = self.session_store.list();

        if let Some(cwd) = args.cwd.as_ref() {
            summaries.retain(|s| s.meta.cwd == cwd.0);
        }

        let (summaries, next_cursor) =
            paginate_summaries(summaries, args.cursor.as_ref().map(|cursor| cursor.0.as_ref()))?;
        let sessions: Vec<acp::SessionInfo> = summaries
            .into_iter()
            .map(|s| {
                let cwd = acp::AbsolutePath::new(s.meta.cwd);
                acp::SessionInfo::new(s.meta.session_id, cwd)
                    .updated_at(s.meta.created_at)
                    .title(s.title)
                    .meta(SessionDisplayMeta::new(s.meta.model, s.meta.selected_mode).to_meta())
            })
            .collect();

        info!("Found {} sessions", sessions.len());
        Ok(ListSessionsResponse::new(sessions).next_cursor(next_cursor.map(acp::SessionListCursor::new)))
    }

    pub(crate) fn search_prompts(&self, params: &PromptSearchParams) -> Result<PromptSearchResponse, Error> {
        self.session_store.search_prompts(&params.query, params.limit).map_err(|e| {
            error!("Prompt search failed: {e}");
            Error::internal_error()
        })
    }

    pub(crate) fn session_preview(&self, params: &SessionPreviewParams) -> Result<SessionPreviewResponse, Error> {
        self.session_store.preview(&params.session_id).map_err(|e| {
            error!("Session preview failed: {e}");
            match e {
                SessionStoreError::Io(error) if error.kind() == std::io::ErrorKind::NotFound => Error::invalid_params(),
                _ => Error::internal_error(),
            }
        })
    }

    /// Lists managed workspaces sharing the session's repository.
    pub(crate) async fn workspace_list(&self, params: &WorkspaceListParams) -> Result<WorkspaceListResponse, Error> {
        self.run_workspace_request(params.session_id.clone(), |_session_store, manager, src_cwd| {
            manager
                .list(&src_cwd)
                .map(|workspaces| WorkspaceListResponse { workspaces })
                .map_err(WorkspaceRequestError::Workspace)
        })
        .await
    }

    /// Moves the session's uncommitted changes to the target workspace, then
    /// relocates the stored session to the target directory.
    pub(crate) async fn workspace_move(&self, params: &WorkspaceMoveParams) -> Result<WorkspaceMoveResponse, Error> {
        let target = params.target.clone();
        let session_id = params.session_id.clone();
        let response = self
            .run_workspace_request(session_id.clone(), move |session_store, manager, src_cwd| {
                let new_cwd = manager.move_to(&src_cwd, &target).map_err(WorkspaceRequestError::Workspace)?;
                session_store.relocate(&session_id, &new_cwd).map_err(WorkspaceRequestError::Relocate)?;
                Ok(WorkspaceMoveResponse { new_cwd })
            })
            .await?;

        self.registry.stop_matching(Some(&params.session_id)).await;
        info!("Moved session {} to workspace {}", params.session_id, response.new_cwd.display());
        Ok(response)
    }

    async fn run_workspace_request<T, F>(&self, session_id: String, op: F) -> Result<T, Error>
    where
        T: Send + 'static,
        F: FnOnce(&SessionStore, &WorkspaceManager, PathBuf) -> Result<T, WorkspaceRequestError> + Send + 'static,
    {
        let session_store = Arc::clone(&self.session_store);
        let manager = Arc::clone(&self.workspace_manager);
        spawn_blocking(move || {
            let src_cwd = session_store
                .session_cwd(&session_id)
                .ok_or_else(|| WorkspaceRequestError::UnknownSession(session_id.clone()))?;
            op(&session_store, &manager, src_cwd)
        })
        .await
        .map_err(|e| internal_error(format!("workspace task failed: {e}")))?
        .map_err(workspace_request_error)
    }

    /// Route a prompt to its session actor for validation and acceptance.
    pub(crate) async fn route_prompt(&self, args: PromptRequest, responder: Responder<PromptResponse>) {
        info!("Received prompt for session: {:?}", args.session_id);
        let session_id = args.session_id.0.to_string();
        let display_content = map_acp_to_content_blocks(acp_utils::content::display_content_blocks(&args.prompt));
        let content = map_acp_to_content_blocks(args.prompt);

        let Some(sender) = self.registry.lookup(Some(&session_id)).await else {
            error!("Session not found: {session_id}");
            respond_err(responder, Error::invalid_params());
            return;
        };

        if let Err(SessionCommand::Prompt { responder, .. }) =
            sender.send(SessionCommand::Prompt { content, display_content, responder }).await.map_err(|e| e.0)
        {
            error!("Session actor channel closed for prompt: {session_id}");
            respond_err(responder, Error::internal_error());
        }
    }

    pub(crate) async fn cancel(&self, args: CancelSessionNotification) -> Result<(), Error> {
        info!("Received cancel for session: {:?}", args.session_id);
        let session_id = args.session_id.0.to_string();
        let Some(sender) = self.registry.lookup(Some(&session_id)).await else {
            error!("Session not found for cancel: {session_id}");
            return Err(Error::invalid_params());
        };
        sender.send(SessionCommand::Cancel).await.map_err(|_| {
            error!("Session actor channel closed for cancel: {session_id}");
            Error::internal_error()
        })
    }

    /// Route a config change to its session actor. Discovers available models
    /// once, then the actor applies the change and answers with rebuilt options.
    pub(crate) async fn set_session_config_option(
        &self,
        args: SetSessionConfigOptionRequest,
        responder: Responder<SetSessionConfigOptionResponse>,
    ) {
        let session_id = args.session_id.0.to_string();
        let config_id = args.config_id.0.to_string();
        let value = match args.value {
            acp::SessionConfigOptionValue::Id { value } => value.0.to_string(),
            acp::SessionConfigOptionValue::Boolean { value } => value.to_string(),
            _ => {
                respond_err(responder, Error::invalid_params());
                return;
            }
        };
        info!("set_session_config_option: session={session_id}, config={config_id}, value={value}");

        let setting = match ConfigSetting::parse(&config_id, &value) {
            Ok(setting) => setting,
            Err(e) => {
                error!("{e}");
                respond_err(responder, Error::invalid_params());
                return;
            }
        };

        let Some(sender) = self.registry.lookup(Some(&session_id)).await else {
            error!("Session not found: {session_id}");
            respond_err(responder, Error::invalid_params());
            return;
        };

        let available = self.factory.available_models().await.to_vec();
        if let Err(SessionCommand::SetConfig { responder, .. }) =
            sender.send(SessionCommand::SetConfig { setting, available, responder }).await.map_err(|e| e.0)
        {
            error!("Session actor channel closed for set_config: {session_id}");
            respond_err(responder, Error::internal_error());
        }
    }

    pub(crate) async fn on_mcp_request(&self, request: McpRequest) -> Result<(), Error> {
        info!("Received MCP ext request: {:?}", request);
        match request {
            McpRequest::Authenticate { session_id, server_name } => {
                let Some(sender) = self.registry.lookup(Some(&session_id)).await else {
                    error!("Session not found for authenticate_mcp_server: {session_id}");
                    return Err(Error::invalid_params());
                };
                sender.send(SessionCommand::AuthenticateMcp { server_name }).await.map_err(|_| {
                    error!("Session actor channel closed for MCP auth: {session_id}");
                    Error::internal_error()
                })?;
            }
        }
        Ok(())
    }

    pub(super) fn stop_token(&self) -> CancellationToken {
        self.stop.clone()
    }

    /// Join the active actor when its host is shutting down.
    pub(super) async fn shutdown_all(&self) {
        self.stop.cancel();
        self.registry.stop().await;
        if let Some(telemetry) = &self.telemetry {
            telemetry.shutdown_or_log();
        }
    }

    #[cfg(any(test, feature = "testing"))]
    pub(crate) async fn register_session(&self, init: SessionActorInit) {
        self.registry.stop().await;
        let id = init.session_id.clone();
        let handle = SessionActor::spawn(init).await.expect("test session actor spawns");
        self.registry.register(&id, handle).await;
    }

    async fn serve_attached(
        self: &Arc<Self>,
        transport: impl ConnectTo<Agent> + 'static,
        stop: CancellationToken,
        guard: ClientGuard,
        runner: impl agent_client_protocol::RunWithConnectionTo<Client> + 'static,
    ) -> Result<(), Error> {
        let result = acp_agent_builder(self.clone())
            .with_runner(runner)
            .connect_with(transport, async move |cx| {
                tokio::select! {
                    biased;
                    () = stop.cancelled() => {},
                    () = cx.incoming_closed() => {},
                }
                Ok(())
            })
            .await;
        guard.release().await;
        result
    }

    async fn release_client(&self) {
        self.detach_client().await;
        self.client_slot.generation.send_modify(|generation| {
            *generation += 1;
            self.client_slot.occupied.store(false, Ordering::Release);
        });
    }

    /// Detach output without interrupting the active session's work.
    async fn detach_client(&self) {
        let Some(sender) = self.registry.lookup(None).await else { return };
        let (reply, response) = oneshot::channel();
        if sender.send(SessionCommand::Detach { reply }).await.is_ok() {
            let _ = response.await;
        }
    }

    async fn broadcast_auth_state(&self, cx: &ConnectionTo<Client>) {
        let auth_methods = build_auth_methods(self.oauth_credential_store.as_ref());
        notify(cx, AuthMethodsUpdatedParams { auth_methods });
        self.broadcast_config_options().await;
    }

    async fn broadcast_config_options(&self) {
        if let Some(sender) = self.registry.lookup(None).await {
            let available = self.factory.available_models().await.to_vec();
            let _ = sender.send(SessionCommand::RefreshConfigOptions { available }).await;
        }
    }
}

struct CodexLogin;

#[async_trait::async_trait]
impl ProviderLogin for CodexLogin {
    async fn login(&self, store: &dyn OAuthCredentialStorage) -> Result<(), llm::LlmError> {
        llm::perform_codex_oauth_flow(store).await
    }
}

fn respond_err<T: agent_client_protocol::JsonRpcResponse>(responder: Responder<T>, error: Error) {
    if let Err(e) = responder.respond_with_error(error) {
        error!("failed to send error response: {e:?}");
    }
}

enum WorkspaceRequestError {
    UnknownSession(String),
    Workspace(WorkspaceError),
    Relocate(SessionStoreError),
}

fn workspace_request_error(e: WorkspaceRequestError) -> Error {
    match e {
        WorkspaceRequestError::UnknownSession(session_id) => {
            Error::invalid_params().data(format!("unknown session: {session_id}"))
        }
        WorkspaceRequestError::Workspace(e) => workspace_error(&e),
        WorkspaceRequestError::Relocate(e) => internal_error(format!("failed to relocate session: {e}")),
    }
}

fn workspace_error(e: &WorkspaceError) -> Error {
    if e.is_invalid_input() { Error::invalid_params().data(e.to_string()) } else { internal_error(e.to_string()) }
}

fn mcp_client_capabilities(client: &acp::ClientCapabilities) -> rmcp::model::ClientCapabilities {
    let elicitation = client.elicitation.as_ref();
    client_capabilities_for(
        elicitation.is_some_and(|capabilities| capabilities.form.is_some()),
        elicitation.is_some_and(|capabilities| capabilities.url.is_some()),
    )
}

fn build_auth_methods(store: &dyn OAuthCredentialStorage) -> Vec<AuthMethod> {
    let mut seen = HashSet::new();
    LlmModel::all()
        .iter()
        .filter_map(LlmModel::oauth_provider_id)
        .filter(|id| seen.insert(*id))
        .map(|id| {
            let display = LlmModel::all()
                .iter()
                .find(|m| m.oauth_provider_id() == Some(id))
                .map_or(id, |m| m.provider_display_name());
            let mut method = acp::AuthMethodAgent::new(id, display);
            if store.contains(id) {
                method = method.description("authenticated");
            }
            AuthMethod::Agent(method)
        })
        .collect()
}

fn prompt_capabilities_for_models(models: &[LlmModel]) -> PromptCapabilities {
    PromptCapabilities::new()
        .embedded_context(PromptEmbeddedContextCapabilities::new())
        .image(models.iter().any(LlmModel::supports_image).then(PromptImageCapabilities::new))
        .audio(models.iter().any(supports_prompt_audio).then(PromptAudioCapabilities::new))
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
struct PromptModalities {
    image: bool,
    audio: bool,
}

impl PromptModalities {
    fn from_content(content: &[ContentBlock]) -> Self {
        Self {
            image: content.iter().any(ContentBlock::is_image),
            audio: content.iter().any(|block| matches!(block, ContentBlock::Audio { .. })),
        }
    }

    fn is_empty(self) -> bool {
        !self.image && !self.audio
    }
}

fn selected_models(model_value: &str) -> Result<Vec<LlmModel>, Error> {
    model_value.parse::<ModelSpec>().map(|spec| spec.models().to_vec()).map_err(|_| Error::invalid_params())
}

pub(crate) fn validate_prompt_support(model_value: &str, content: &[ContentBlock]) -> Result<(), Error> {
    let modalities = PromptModalities::from_content(content);
    if modalities.is_empty() {
        return Ok(());
    }

    let selected = selected_models(model_value)?;
    if modalities.image && selected.iter().any(|model| !model.supports_image()) {
        return Err(Error::invalid_params());
    }
    if modalities.audio && selected.iter().any(|model| !supports_prompt_audio(model)) {
        return Err(Error::invalid_params());
    }

    Ok(())
}

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

    const SONNET: &str = "anthropic:claude-sonnet-4-5";
    const AUDIO_ONLY: &str = "bedrock:mistral.voxtral-small-24b-2507";

    fn text_only_model() -> LlmModel {
        // Upstream model metadata changes weekly, so look the fixture up instead of pinning a model id.
        LlmModel::all()
            .iter()
            .find(|model| !model.supports_image() && !supports_prompt_audio(model))
            .cloned()
            .expect("catalog contains a text-only model")
    }

    fn fake_oauth_store() -> Arc<dyn OAuthCredentialStorage> {
        Arc::new(aether_auth::FakeOAuthCredentialStore::new())
    }

    fn test_state() -> AcpState {
        let session_store = Arc::new(SessionStore::from_path(PathBuf::from("/tmp/aether-test-sessions")));
        let workspace_manager =
            Arc::new(WorkspaceManager::from_registry_path(PathBuf::from("/tmp/aether-test-workspaces.json")));
        AcpState::new(AcpStateConfig {
            session_store,
            workspace_manager,
            oauth_credential_store: fake_oauth_store(),
            initial_selection: InitialSessionSelection::default(),
            settings_source: SettingsSourceArgs::default(),
            provider_connections: ProviderConnectionOverrides::default(),
            telemetry: None,
            runtime_factory: None,
            cwd: PathBuf::from("/tmp"),
        })
    }

    #[tokio::test]
    async fn initialize_advertises_session_lifecycle_support() {
        let state = test_state();
        let response = state
            .initialize(InitializeRequest::new(ProtocolVersion::V2, Implementation::new("test", "1")))
            .await
            .expect("initialize succeeds");
        let json = serde_json::to_string(&response).expect("response serializes");
        assert_eq!(response.protocol_version, ProtocolVersion::V2);
        let session = response.capabilities.session.unwrap();
        assert!(session.prompt.is_some());
        let mcp = session.mcp.unwrap();
        assert!(mcp.stdio.is_some());
        assert!(mcp.http.is_some());
        for obsolete in ["loadSession", "resume", "close", "list", "sse"] {
            assert!(!json.contains(&format!("\"{obsolete}\":")));
        }
    }

    #[tokio::test]
    async fn initialize_advertises_aether_capabilities_once() {
        let state = test_state();
        let response = state
            .initialize(InitializeRequest::new(ProtocolVersion::V2, Implementation::new("test", "1")))
            .await
            .expect("initialize succeeds");
        assert_eq!(
            AetherCapabilities::from_meta(
                response.capabilities.session.as_ref().unwrap().prompt.as_ref().unwrap().meta.as_ref()
            ),
            AetherCapabilities::default()
        );
        let capabilities = AetherCapabilities::from_meta(response.capabilities.session.as_ref().unwrap().meta.as_ref());
        assert!(capabilities.prompt_search);
        assert!(capabilities.session_preview);
        assert!(capabilities.workspace_move);
    }

    #[test]
    fn mcp_elicitation_capabilities_mirror_the_acp_client() {
        let none = mcp_client_capabilities(&acp::ClientCapabilities::new());
        assert!(none.elicitation.is_none());

        let form_only = acp::ClientCapabilities::new()
            .elicitation(acp::ElicitationCapabilities::new().form(acp::ElicitationFormCapabilities::new()));
        let form = mcp_client_capabilities(&form_only).elicitation.unwrap();
        assert!(form.form.is_some());
        assert!(form.url.is_none());

        let url_only = acp::ClientCapabilities::new()
            .elicitation(acp::ElicitationCapabilities::new().url(acp::ElicitationUrlCapabilities::new()));
        let url = mcp_client_capabilities(&url_only).elicitation.unwrap();
        assert!(url.form.is_none());
        assert!(url.url.is_some());
    }

    #[test]
    fn prompt_capabilities_reflect_available_modalities() {
        let image_only = prompt_capabilities_for_models(&[SONNET.parse().unwrap()]);
        assert!(image_only.image.is_some());
        assert!(image_only.audio.is_none());

        let audio_capable = prompt_capabilities_for_models(&[AUDIO_ONLY.parse().unwrap()]);
        assert!(audio_capable.image.is_none());
        assert!(audio_capable.audio.is_some());

        let text_only = prompt_capabilities_for_models(&[text_only_model()]);
        assert!(text_only.image.is_none());
        assert!(text_only.audio.is_none());
    }

    #[test]
    fn validate_prompt_support_requires_all_selected_models_to_support_media() {
        let image_content = vec![ContentBlock::Image { data: "aW1n".to_string(), mime_type: "image/png".to_string() }];
        let audio_content =
            vec![ContentBlock::Audio { data: "YXVkaW8=".to_string(), mime_type: "audio/wav".to_string() }];

        let text_only = text_only_model().to_string();
        assert!(validate_prompt_support(SONNET, &image_content).is_ok());
        assert!(validate_prompt_support(&text_only, &image_content).is_err());
        assert!(validate_prompt_support(AUDIO_ONLY, &audio_content).is_ok());
        assert!(validate_prompt_support(SONNET, &audio_content).is_err());
        assert!(validate_prompt_support(format!("{SONNET},{text_only}").as_str(), &image_content).is_err());
        assert!(validate_prompt_support(format!("{AUDIO_ONLY},{text_only}").as_str(), &audio_content).is_err());
    }
}