rig-core 0.45.0

//! A catalog model served with a signed-in credential that each request
//! reads afresh, so a long session outlives its access token. ChatGPT's and
//! Copilot's authenticators both sign in this way.

use std::future::Future;
use std::sync::{Arc, Mutex};

use crate::catalog::{Catalog, ModelSpec};
use crate::completion::UnsupportedOption;
use crate::effect::{EffectKind, HandlerDescriptor, family};
use crate::error::{ErrorKind, ErrorReport};
use crate::http_client::DynHttpClient;
use crate::operation::Completion;
use crate::providers::internal::auth::{AuthContext, AuthError};
use crate::providers::registry::{ConnectError, ConnectOptions};
use crate::providers::{chatgpt, copilot};
use crate::serve::adapters::ModelAdapter;
use crate::serve::{Dispatch, Reply, Serve};
use crate::wasm_compat::{WasmCompatSend, WasmCompatSync};

/// Where a [`SignedInModel`] reads its credential on every request.
pub trait CredentialSource: WasmCompatSend + WasmCompatSync + 'static {
    /// The credential to send with now, refreshed through `http` when it
    /// has expired.
    fn auth_context(
        &self,
        http: &DynHttpClient,
    ) -> impl Future<Output = Result<AuthContext, AuthError>> + WasmCompatSend;
}

impl CredentialSource for chatgpt::auth::Authenticator {
    fn auth_context(
        &self,
        http: &DynHttpClient,
    ) -> impl Future<Output = Result<AuthContext, AuthError>> + WasmCompatSend {
        chatgpt::auth::Authenticator::auth_context(self, http)
    }
}

impl CredentialSource for copilot::auth::Authenticator {
    fn auth_context(
        &self,
        http: &DynHttpClient,
    ) -> impl Future<Output = Result<AuthContext, AuthError>> + WasmCompatSend {
        copilot::auth::Authenticator::auth_context(self, http)
    }
}

/// A catalog model as an effect handler whose every request carries the
/// credential its [`CredentialSource`] resolves at that moment, refreshed
/// when it has expired: the key, the account it belongs to, and the API root the
/// sign-in names. The model connected with a credential serves every
/// request until the credential changes, so its memos last across turns.
/// Described as [`ModelAdapter`] describes the model, under its
/// [`reference`](ModelSpec::reference).
///
/// ```no_run
/// use rig_core::catalog::Catalog;
/// use rig_core::providers::chatgpt::auth::{AuthSource, Authenticator, DeviceCodeHandler, FileTokenStore};
/// use rig_core::client::signed_in::SignedInModel;
/// use rig_core::serve::ErasedHandler;
///
/// # fn run(http: rig_core::http_client::DynHttpClient) -> Result<(), Box<dyn std::error::Error>> {
/// let spec = Catalog::builtin().resolve("chatgpt/gpt-5.5")?.spec.clone();
/// let auth = Authenticator::new(AuthSource::OAuth, FileTokenStore::new("/path/to/credentials"), DeviceCodeHandler::default(), false);
/// let handler = ErasedHandler::new(SignedInModel::new(spec, auth, http)?);
/// # let _ = handler;
/// # Ok(())
/// # }
/// ```
pub struct SignedInModel<A> {
    spec: ModelSpec,
    credentials: A,
    http: DynHttpClient,
    /// The model connected without a credential: what it is and what it
    /// would leave out, neither of which sends anything.
    unsigned: ModelAdapter<Completion>,
    /// The last credential and the model connected with it.
    connected: Mutex<Option<(AuthContext, Arc<ModelAdapter<Completion>>)>>,
}

impl<A: CredentialSource> SignedInModel<A> {
    /// `spec` served with what `credentials` resolves, through `http`.
    ///
    /// # Errors
    ///
    /// What [`Catalog::connect_with`] reports for `spec`.
    pub fn new(spec: ModelSpec, credentials: A, http: DynHttpClient) -> Result<Self, ConnectError> {
        let unsigned = Catalog::builtin()
            .connect_with(&spec, ConnectOptions::new().api_key("").http(http.clone()))?;
        Ok(Self {
            unsigned: ModelAdapter::new(spec.reference(), unsigned),
            spec,
            credentials,
            http,
            connected: Mutex::default(),
        })
    }

    /// The model connected with the current credential: the one already
    /// connected while the credential is the same.
    async fn model(&self) -> Result<Arc<ModelAdapter<Completion>>, ErrorReport> {
        let context = self
            .credentials
            .auth_context(&self.http)
            .await
            .map_err(|error| {
                ErrorReport::new(
                    ErrorKind::Provider,
                    format!("the sign-in gave no credential ({error}); sign in again"),
                )
                .with_retryable(false)
            })?;
        if let Ok(connected) = self.connected.lock()
            && let Some((held, model)) = connected.as_ref()
            && *held == context
        {
            return Ok(Arc::clone(model));
        }
        let mut options = ConnectOptions::new()
            .api_key(context.api_key.clone())
            .http(self.http.clone());
        if let Some(account_id) = &context.account_id {
            options = options.account_id(account_id);
        }
        if let Some(api_base) = &context.api_base {
            options = options.base_url(api_base);
        }
        let model = Catalog::builtin()
            .connect_with(&self.spec, options)
            .map_err(|error| ErrorReport::new(ErrorKind::Internal, error.to_string()))?;
        let model = Arc::new(ModelAdapter::new(self.spec.reference(), model));
        if let Ok(mut connected) = self.connected.lock() {
            *connected = Some((context, Arc::clone(&model)));
        }
        Ok(model)
    }
}

impl<A: CredentialSource> Serve for SignedInModel<A> {
    type Family = family::Completion;

    fn descriptor(&self) -> HandlerDescriptor {
        self.unsigned.descriptor()
    }

    /// What the model leaves out, from the unsigned model: a check is a
    /// dry run, which needs no credential.
    fn check(&self, kind: &EffectKind) -> Vec<UnsupportedOption> {
        self.unsigned.check(kind)
    }

    async fn serve(&self, kind: EffectKind, dispatch: Dispatch) -> Reply {
        match self.model().await {
            Ok(model) => model.serve(kind, dispatch).await,
            Err(report) => Reply::Outcome(Err(report)),
        }
    }
}