magi-code 0.77.1

use super::{
    dispatcher::{RequestIdGuard, ServiceOutbound},
    protocol::{
        RequestIdentity, ResponsePayload, ServiceErrorCode, ServiceEvent, ServiceMessage,
        ServiceRequest, ServiceResponse,
    },
    runtime::ServiceRuntime,
};
use crate::{
    config,
    providers::{ANTHROPIC_PROVIDER, OPENAI_CODEX_PROVIDER},
};
use crossbeam_channel::{Receiver, Sender, bounded};
use serde::Deserialize;
use serde_json::{Value, json};
use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
    mpsc::{SyncSender, sync_channel},
};
use std::thread::JoinHandle;

pub(crate) struct AuthWorkerMessage {
    state: &'static str,
    url: Option<String>,
    terminal: bool,
}

struct ActiveLogin {
    id: String,
    request_id: String,
    guard: Option<RequestIdGuard>,
    cancel_request: Sender<bool>,
    cancellation_requested: Arc<AtomicBool>,
    cancellation_finished: Arc<AtomicBool>,
    callback: Option<SyncSender<String>>,
    worker: JoinHandle<()>,
}

impl ActiveLogin {
    fn request_cancellation(&self) {
        // Stop deferred generation capture immediately. Only the cancellation worker
        // changes the separate commit flag under the protected credential lock.
        self.cancellation_requested.store(true, Ordering::Release);
        let _ = self.cancel_request.try_send(true);
    }
}

pub(super) struct LoginJob {
    paths: config::McPaths,
    generation: u64,
    pub(super) cancel: Arc<AtomicBool>,
    callback: std::sync::mpsc::Receiver<String>,
    sender: Sender<AuthWorkerMessage>,
}

pub(super) type LoginWorker = Arc<dyn Fn(LoginJob) -> &'static str + Send + Sync>;

fn run_login(job: LoginJob) -> &'static str {
    crate::login::service_codex_login(
        &job.paths,
        job.generation,
        &job.cancel,
        &job.callback,
        |state, url| {
            let _ = job.sender.send(AuthWorkerMessage {
                state,
                url,
                terminal: false,
            });
        },
    )
}

pub(crate) struct ServiceAuthManager {
    runtime: Arc<ServiceRuntime>,
    active: Option<ActiveLogin>,
    sender: Sender<AuthWorkerMessage>,
    pub(crate) receiver: Receiver<AuthWorkerMessage>,
    worker: LoginWorker,
}

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ProviderParams {
    provider_id: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct CallbackParams {
    login_id: String,
    input: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct CancelParams {
    login_id: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct LogoutParams {
    provider_id: String,
    #[serde(default)]
    confirmed: bool,
}

impl ServiceAuthManager {
    pub(crate) fn new(runtime: Arc<ServiceRuntime>) -> Self {
        let (sender, receiver) = bounded(16);
        Self {
            runtime,
            active: None,
            sender,
            receiver,
            worker: Arc::new(run_login),
        }
    }

    #[cfg(test)]
    pub(super) fn set_worker(&mut self, worker: LoginWorker) {
        self.worker = worker;
    }

    pub(crate) fn dispatch(
        &mut self,
        request: ServiceRequest,
        guard: RequestIdGuard,
        initialized: bool,
    ) -> ServiceOutbound {
        let result = if !initialized {
            Err(ServiceErrorCode::NotInitialized)
        } else if request.session_id.is_some() {
            Err(ServiceErrorCode::InvalidPayload)
        } else {
            self.operation(&request)
        };
        match result {
            Ok(payload) => {
                let response = ServiceMessage::response(ServiceResponse::success(
                    &request,
                    ResponsePayload::Auth(payload),
                ));
                if request.method == "auth.login.start" {
                    // Start reserves its request ID until the terminal event has been written.
                    let active = self
                        .active
                        .as_mut()
                        .expect("successful start owns a worker");
                    active.guard = Some(guard);
                    ServiceOutbound::unguarded(vec![response])
                } else {
                    ServiceOutbound::guarded(vec![response], guard)
                }
            }
            Err(code) => ServiceOutbound::guarded(
                vec![ServiceMessage::response(ServiceResponse::error(
                    RequestIdentity::from_request(&request),
                    code,
                ))],
                guard,
            ),
        }
    }

    pub(super) fn operation(
        &mut self,
        request: &ServiceRequest,
    ) -> Result<Value, ServiceErrorCode> {
        let invalid = |_| ServiceErrorCode::InvalidPayload;
        match request.method.as_str() {
            "auth.status" => {
                let _: super::protocol::EmptyParams =
                    serde_json::from_value(request.payload.clone()).map_err(invalid)?;
                self.status()
            }
            "auth.login.start" => {
                let params: ProviderParams =
                    serde_json::from_value(request.payload.clone()).map_err(invalid)?;
                if params.provider_id != OPENAI_CODEX_PROVIDER {
                    return Err(ServiceErrorCode::ProviderUnavailable);
                }
                // The caller installs the request guard after starting the worker.
                self.start(request)
            }
            "auth.login.callback" => {
                let params: CallbackParams =
                    serde_json::from_value(request.payload.clone()).map_err(invalid)?;
                if params.input.is_empty() || params.input.len() > 4096 {
                    return Err(ServiceErrorCode::InvalidPayload);
                }
                let active = self
                    .active
                    .as_mut()
                    .filter(|a| a.id == params.login_id)
                    .ok_or(ServiceErrorCode::UnknownLogin)?;
                let sender = active.callback.as_ref().ok_or(ServiceErrorCode::AuthBusy)?;
                sender
                    .try_send(params.input)
                    .map_err(|_| ServiceErrorCode::AuthBusy)?;
                active.callback = None;
                Ok(json!({"login_id": active.id, "status": "submitted"}))
            }
            "auth.login.cancel" => {
                let params: CancelParams =
                    serde_json::from_value(request.payload.clone()).map_err(invalid)?;
                let active = self
                    .active
                    .as_ref()
                    .filter(|a| a.id == params.login_id)
                    .ok_or(ServiceErrorCode::UnknownLogin)?;
                active.request_cancellation();
                Ok(json!({"login_id": active.id, "status": "cancellation_requested"}))
            }
            "auth.logout" => {
                let provider_id = self.validate_logout(request)?;
                if provider_id == OPENAI_CODEX_PROVIDER {
                    self.cancel_all();
                    if let Some(barrier) = self.cancellation_barrier() {
                        while !barrier.load(Ordering::Acquire) {
                            std::thread::sleep(std::time::Duration::from_millis(1));
                        }
                    }
                }
                let removal =
                    config::remove_provider_auth(&self.runtime.config.paths, &provider_id)
                        .map_err(|_| ServiceErrorCode::InternalError)?;
                Ok(
                    json!({"provider_id": provider_id, "status": "logged_out", "removed": removal.removed, "environment_unchanged": true}),
                )
            }
            _ => Err(ServiceErrorCode::UnsupportedOperation),
        }
    }

    pub(super) fn validate_logout(
        &self,
        request: &ServiceRequest,
    ) -> Result<String, ServiceErrorCode> {
        let params: LogoutParams = serde_json::from_value(request.payload.clone())
            .map_err(|_| ServiceErrorCode::InvalidPayload)?;
        if !params.confirmed {
            return Err(ServiceErrorCode::ConfirmationRequired);
        }
        if params.provider_id != OPENAI_CODEX_PROVIDER
            && params.provider_id != ANTHROPIC_PROVIDER
            && !self
                .runtime
                .config
                .custom_providers
                .contains_key(&params.provider_id)
        {
            return Err(ServiceErrorCode::ProviderUnavailable);
        }
        Ok(params.provider_id)
    }
    fn status(&self) -> Result<Value, ServiceErrorCode> {
        let auth = config::read_auth(&self.runtime.config.paths)
            .map_err(|_| ServiceErrorCode::InternalError)?;
        let custom = &self.runtime.config.custom_providers;
        let mut providers = Vec::new();
        for id in [OPENAI_CODEX_PROVIDER, ANTHROPIC_PROVIDER]
            .into_iter()
            .chain(custom.keys().map(String::as_str))
        {
            let readiness = config::current_provider_auth_readiness(id, &auth, custom);
            let (method, source) = if id == OPENAI_CODEX_PROVIDER {
                ("oauth", "shared_store")
            } else if let Some(provider) = custom.get(id) {
                if provider.api_key_env_var.is_some() {
                    ("api_key", "environment")
                } else {
                    ("none", "none")
                }
            } else if std::env::var("ANTHROPIC_API_KEY")
                .ok()
                .is_some_and(|v| !v.is_empty())
            {
                ("api_key", "environment")
            } else {
                ("api_key", "shared_store")
            };
            let state = match readiness {
                config::CredentialReadiness::Ready | config::CredentialReadiness::Refreshable => {
                    "ready"
                }
                config::CredentialReadiness::Missing => "missing",
                config::CredentialReadiness::Invalid => "needs_login",
            };
            providers.push(json!({"provider_id": id, "auth_method": method, "source": source, "readiness": state, "ready": readiness.is_ready()}));
        }
        let payload = json!({"providers": providers});
        super::protocol::payload_is_bounded(&payload)
            .map_err(|_| ServiceErrorCode::LimitExceeded)?;
        Ok(payload)
    }

    fn start(&mut self, request: &ServiceRequest) -> Result<Value, ServiceErrorCode> {
        if self.active.is_some() {
            return Err(ServiceErrorCode::AuthBusy);
        }
        let paths = self.runtime.config.paths.clone();
        let id = uuid::Uuid::new_v4().to_string();
        let cancel = Arc::new(AtomicBool::new(false));
        let (callback, receiver) = sync_channel(1);
        let sender = self.sender.clone();
        let login_worker = Arc::clone(&self.worker);
        // Bound cancellation arbitration to one worker per login, never the coordinator.
        let (cancel_request, cancellation) = bounded(1);
        let cancellation_paths = paths.clone();
        let cancellation_flag = Arc::clone(&cancel);
        let cancellation_requested = Arc::new(AtomicBool::new(false));
        let worker_requested = Arc::clone(&cancellation_requested);
        let cancellation_finished = Arc::new(AtomicBool::new(false));
        let worker_finished = Arc::clone(&cancellation_finished);
        let cancellation_worker = std::thread::Builder::new()
            .name("service-auth-cancel".into())
            .spawn(move || {
                if cancellation.recv() == Ok(true)
                    && config::cancel_codex_login_before_commit(
                        &cancellation_paths,
                        &cancellation_flag,
                    )
                    .is_err()
                {
                    cancellation_flag.store(true, Ordering::Relaxed);
                }
                worker_finished.store(true, Ordering::Release);
            })
            .map_err(|_| ServiceErrorCode::InternalError)?;
        let finish_cancellation = cancel_request.clone();
        let worker = std::thread::Builder::new()
            .name("service-auth".into())
            .spawn(move || {
                let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                    let Ok(store) = config::read_auth_store(&paths) else {
                        return "internal_error";
                    };
                    if worker_requested.load(Ordering::Acquire) {
                        return "cancelled";
                    }
                    let generation = store.provider_generation(OPENAI_CODEX_PROVIDER);
                    login_worker(LoginJob {
                        paths,
                        generation,
                        cancel,
                        callback: receiver,
                        sender: sender.clone(),
                    })
                }))
                .unwrap_or("internal_error");
                let _ = finish_cancellation.try_send(false);
                let _ = cancellation_worker.join();
                let _ = sender.send(AuthWorkerMessage {
                    state: outcome,
                    url: None,
                    terminal: true,
                });
            })
            .map_err(|_| ServiceErrorCode::InternalError)?;
        self.active = Some(ActiveLogin {
            id: id.clone(),
            request_id: request.request_id.clone(),
            guard: None,
            cancel_request,
            cancellation_requested,
            cancellation_finished,
            callback: Some(callback),
            worker,
        });
        Ok(json!({"login_id": id, "provider_id": OPENAI_CODEX_PROVIDER, "status": "started"}))
    }

    pub(crate) fn output(&mut self, message: AuthWorkerMessage) -> Option<ServiceOutbound> {
        let active = self.active.as_ref()?;
        let mut payload = json!({"login_id": active.id, "provider_id": OPENAI_CODEX_PROVIDER, "state": message.state});
        if let Some(url) = message.url {
            payload["url"] = Value::String(url);
            payload["open_browser"] = json!(true);
        }
        if message.terminal {
            payload["cleanup_complete"] = json!(true);
        }
        let event = ServiceMessage::Event(ServiceEvent::new(
            active.request_id.clone(),
            None,
            if message.terminal {
                "auth.login.terminal"
            } else {
                "auth.login.progress"
            },
            payload,
        ));
        if message.terminal {
            let active = self.active.take()?;
            let _ = active.worker.join();
            Some(ServiceOutbound::guarded(vec![event], active.guard?))
        } else {
            Some(ServiceOutbound::unguarded(vec![event]))
        }
    }

    pub(crate) fn cancel_all(&self) {
        if let Some(active) = &self.active {
            active.request_cancellation();
        }
    }
    pub(super) fn cancellation_barrier(&self) -> Option<Arc<AtomicBool>> {
        self.active
            .as_ref()
            .map(|active| Arc::clone(&active.cancellation_finished))
    }
    pub(crate) fn is_active(&self) -> bool {
        self.active.is_some()
    }
}

impl Drop for ServiceAuthManager {
    fn drop(&mut self) {
        self.cancel_all();
        if let Some(active) = self.active.take() {
            // One flow sends at most five messages into the 16-slot queue, so joining
            // without an adapter draining output cannot block on event delivery.
            let _ = active.worker.join();
        }
    }
}