rig_core/client/
signed_in.rs

1//! A catalog model served with a signed-in credential that each request
2//! reads afresh, so a long session outlives its access token. ChatGPT's and
3//! Copilot's authenticators both sign in this way.
4
5use std::future::Future;
6use std::sync::{Arc, Mutex};
7
8use crate::catalog::{Catalog, ModelSpec};
9use crate::completion::UnsupportedOption;
10use crate::effect::{EffectKind, HandlerDescriptor, family};
11use crate::error::{ErrorKind, ErrorReport};
12use crate::http_client::DynHttpClient;
13use crate::operation::Completion;
14use crate::providers::internal::auth::{AuthContext, AuthError};
15use crate::providers::registry::{ConnectError, ConnectOptions};
16use crate::providers::{chatgpt, copilot};
17use crate::serve::adapters::ModelAdapter;
18use crate::serve::{Dispatch, Reply, Serve};
19use crate::wasm_compat::{WasmCompatSend, WasmCompatSync};
20
21/// Where a [`SignedInModel`] reads its credential on every request.
22pub trait CredentialSource: WasmCompatSend + WasmCompatSync + 'static {
23    /// The credential to send with now, refreshed through `http` when it
24    /// has expired.
25    fn auth_context(
26        &self,
27        http: &DynHttpClient,
28    ) -> impl Future<Output = Result<AuthContext, AuthError>> + WasmCompatSend;
29}
30
31impl CredentialSource for chatgpt::auth::Authenticator {
32    fn auth_context(
33        &self,
34        http: &DynHttpClient,
35    ) -> impl Future<Output = Result<AuthContext, AuthError>> + WasmCompatSend {
36        chatgpt::auth::Authenticator::auth_context(self, http)
37    }
38}
39
40impl CredentialSource for copilot::auth::Authenticator {
41    fn auth_context(
42        &self,
43        http: &DynHttpClient,
44    ) -> impl Future<Output = Result<AuthContext, AuthError>> + WasmCompatSend {
45        copilot::auth::Authenticator::auth_context(self, http)
46    }
47}
48
49/// A catalog model as an effect handler whose every request carries the
50/// credential its [`CredentialSource`] resolves at that moment, refreshed
51/// when it has expired: the key, the account it belongs to, and the API root the
52/// sign-in names. The model connected with a credential serves every
53/// request until the credential changes, so its memos last across turns.
54/// Described as [`ModelAdapter`] describes the model, under its
55/// [`reference`](ModelSpec::reference).
56///
57/// ```no_run
58/// use rig_core::catalog::Catalog;
59/// use rig_core::providers::chatgpt::auth::{AuthSource, Authenticator, DeviceCodeHandler, FileTokenStore};
60/// use rig_core::client::signed_in::SignedInModel;
61/// use rig_core::serve::ErasedHandler;
62///
63/// # fn run(http: rig_core::http_client::DynHttpClient) -> Result<(), Box<dyn std::error::Error>> {
64/// let spec = Catalog::builtin().resolve("chatgpt/gpt-5.5")?.spec.clone();
65/// let auth = Authenticator::new(AuthSource::OAuth, FileTokenStore::new("/path/to/credentials"), DeviceCodeHandler::default(), false);
66/// let handler = ErasedHandler::new(SignedInModel::new(spec, auth, http)?);
67/// # let _ = handler;
68/// # Ok(())
69/// # }
70/// ```
71pub struct SignedInModel<A> {
72    spec: ModelSpec,
73    credentials: A,
74    http: DynHttpClient,
75    /// The model connected without a credential: what it is and what it
76    /// would leave out, neither of which sends anything.
77    unsigned: ModelAdapter<Completion>,
78    /// The last credential and the model connected with it.
79    connected: Mutex<Option<(AuthContext, Arc<ModelAdapter<Completion>>)>>,
80}
81
82impl<A: CredentialSource> SignedInModel<A> {
83    /// `spec` served with what `credentials` resolves, through `http`.
84    ///
85    /// # Errors
86    ///
87    /// What [`Catalog::connect_with`] reports for `spec`.
88    pub fn new(spec: ModelSpec, credentials: A, http: DynHttpClient) -> Result<Self, ConnectError> {
89        let unsigned = Catalog::builtin()
90            .connect_with(&spec, ConnectOptions::new().api_key("").http(http.clone()))?;
91        Ok(Self {
92            unsigned: ModelAdapter::new(spec.reference(), unsigned),
93            spec,
94            credentials,
95            http,
96            connected: Mutex::default(),
97        })
98    }
99
100    /// The model connected with the current credential: the one already
101    /// connected while the credential is the same.
102    async fn model(&self) -> Result<Arc<ModelAdapter<Completion>>, ErrorReport> {
103        let context = self
104            .credentials
105            .auth_context(&self.http)
106            .await
107            .map_err(|error| {
108                ErrorReport::new(
109                    ErrorKind::Provider,
110                    format!("the sign-in gave no credential ({error}); sign in again"),
111                )
112                .with_retryable(false)
113            })?;
114        if let Ok(connected) = self.connected.lock()
115            && let Some((held, model)) = connected.as_ref()
116            && *held == context
117        {
118            return Ok(Arc::clone(model));
119        }
120        let mut options = ConnectOptions::new()
121            .api_key(context.api_key.clone())
122            .http(self.http.clone());
123        if let Some(account_id) = &context.account_id {
124            options = options.account_id(account_id);
125        }
126        if let Some(api_base) = &context.api_base {
127            options = options.base_url(api_base);
128        }
129        let model = Catalog::builtin()
130            .connect_with(&self.spec, options)
131            .map_err(|error| ErrorReport::new(ErrorKind::Internal, error.to_string()))?;
132        let model = Arc::new(ModelAdapter::new(self.spec.reference(), model));
133        if let Ok(mut connected) = self.connected.lock() {
134            *connected = Some((context, Arc::clone(&model)));
135        }
136        Ok(model)
137    }
138}
139
140impl<A: CredentialSource> Serve for SignedInModel<A> {
141    type Family = family::Completion;
142
143    fn descriptor(&self) -> HandlerDescriptor {
144        self.unsigned.descriptor()
145    }
146
147    /// What the model leaves out, from the unsigned model: a check is a
148    /// dry run, which needs no credential.
149    fn check(&self, kind: &EffectKind) -> Vec<UnsupportedOption> {
150        self.unsigned.check(kind)
151    }
152
153    async fn serve(&self, kind: EffectKind, dispatch: Dispatch) -> Reply {
154        match self.model().await {
155            Ok(model) => model.serve(kind, dispatch).await,
156            Err(report) => Reply::Outcome(Err(report)),
157        }
158    }
159}
rig_core/client/signed_in.rs

rig_core/client/
signed_in.rs
