aether-agent-cli 0.9.3

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, SessionId,
    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, mpsc, 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::server::DetachedArgs;
use super::session::actor::{ClientConnection, 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,
    pub(crate) detached: DetachedArgs,
}

struct SpawnedSession {
    session_id: SessionId,
    config_options: Vec<acp::SessionConfigOption>,
    commands: mpsc::Sender<SessionCommand>,
}

#[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,
            config.detached,
        );
        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 created = self.spawn_session(req, Some(cx)).await?;
        Ok(NewSessionResponse::new(created.session_id).config_options(created.config_options))
    }

    pub(super) async fn start_session(&self, prompt: String) -> Result<SessionId, Error> {
        let request = NewSessionRequest::new(acp::AbsolutePath::new(self.cwd.clone()));
        let created = self.spawn_session(request, None).await?;
        let content = vec![ContentBlock::text(prompt)];
        created
            .commands
            .send(SessionCommand::Prompt {
                content: content.clone(),
                display_content: content,
                client_connection: ClientConnection::Detached,
            })
            .await
            .map_err(|_| Error::internal_error())?;
        Ok(created.session_id)
    }

    async fn spawn_session(
        &self,
        request: NewSessionRequest,
        connection: Option<&ConnectionTo<Client>>,
    ) -> Result<SpawnedSession, Error> {
        let mcp_capabilities = self.mcp_capabilities.lock().await.clone();
        let prepared = self.factory.prepare_new(request, connection, mcp_capabilities).await?;
        self.registry.stop().await;
        let created = prepared.start().await?;
        let spawned = SpawnedSession {
            session_id: created.session_id.clone(),
            config_options: created.config_options,
            commands: created.handle.command_sender(),
        };
        self.registry.register(&created.session_id, created.handle).await;
        Ok(spawned)
    }

    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 { client_connection: responder, .. }) = sender
            .send(SessionCommand::Prompt {
                content,
                display_content,
                client_connection: ClientConnection::attached(responder),
            })
            .await
            .map_err(|e| e.0)
        {
            error!("Session actor channel closed for prompt: {session_id}");
            responder.respond_with_error(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"),
            detached: DetachedArgs::default(),
        })
    }

    #[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());
    }
}