Skip to main content

systemprompt_cli/session/creation/
helpers.rs

1//! User resolution for CLI session creation.
2//!
3//! Local administrator lookup uses the configured username. Credentialed
4//! paths resolve email addresses from session hints or cloud credentials.
5//!
6//! Copyright (c) systemprompt.io — Business Source License 1.1.
7//! See <https://systemprompt.io> for licensing details.
8
9use std::sync::Arc;
10
11use anyhow::{Context, Result};
12use chrono::Duration as ChronoDuration;
13use systemprompt_agent::repository::context::ContextRepository;
14use systemprompt_cloud::{
15    CliSession, CloudCredentials, CredentialsBootstrap, SessionBinding, SessionIdentity, SessionKey,
16};
17use systemprompt_config::SecretsBootstrap;
18use systemprompt_database::{Database, DbPool, PoolConfig};
19use systemprompt_identifiers::{ContextId, Email, ProfileName, SessionId, SessionToken};
20use systemprompt_models::auth::{Permission, RateLimitTier, UserType};
21use systemprompt_security::{SessionGenerator, SessionParams};
22use systemprompt_users::{UserRepository, UserService};
23
24use crate::session::resolution::ProfileContext;
25
26pub(super) struct ResolvedSecrets {
27    pub database_url: String,
28    pub database_write_url: Option<String>,
29}
30
31pub(super) fn load_secrets() -> Result<ResolvedSecrets> {
32    let secrets = SecretsBootstrap::get().map_err(|e| {
33        anyhow::anyhow!(
34            "Secrets not initialized: {}\n\nEnsure your profile has a valid secrets \
35             configuration.\nCheck that secrets.json exists or environment variables are set.",
36            e
37        )
38    })?;
39
40    Ok(ResolvedSecrets {
41        database_url: secrets.database_url.clone(),
42        database_write_url: secrets.database_write_url.clone(),
43    })
44}
45
46pub(super) async fn connect_database(secrets: &ResolvedSecrets) -> Result<DbPool> {
47    let db = Database::from_config_with_write(
48        "postgres",
49        &secrets.database_url,
50        secrets.database_write_url.as_deref(),
51        &PoolConfig::default(),
52    )
53    .await
54    .context("Failed to connect to database")?;
55    Ok(DbPool::from(Arc::new(db)))
56}
57
58pub async fn get_or_create_admin(
59    db_pool: &DbPool,
60    email: &str,
61    context_type: &str,
62) -> Result<systemprompt_users::User> {
63    let email = Email::try_new(email).map_err(|e| {
64        anyhow::anyhow!("refusing to provision an admin for an invalid address: {e}")
65    })?;
66    let email = email.as_str();
67
68    let user_service = UserService::new(Arc::new(UserRepository::new(db_pool)?));
69
70    if let Some(user) = user_service
71        .find_by_email(email)
72        .await
73        .context("Failed to query user by email")?
74    {
75        if user.is_admin() {
76            return Ok(user);
77        }
78
79        tracing::info!(email = %email, context = %context_type, "Promoting existing user to admin");
80
81        return user_service
82            .assign_roles(&user.id, &["admin".to_owned()])
83            .await
84            .context("Failed to assign admin role to existing user");
85    }
86
87    let name = email.split('@').next().unwrap_or("admin").to_owned();
88
89    tracing::info!(email = %email, name = %name, context = %context_type, "Auto-provisioning user");
90
91    let user = match user_service
92        .create_if_absent(&name, email, None, None)
93        .await
94        .with_context(|| format!("Failed to create user in {context_type} database"))?
95    {
96        Some(user) => user,
97        None => user_service
98            .find_by_email(email)
99            .await
100            .context("Failed to query user by email")?
101            .with_context(|| format!("User {email} vanished between provisioning and lookup"))?,
102    };
103
104    user_service
105        .assign_roles(&user.id, &["admin".to_owned()])
106        .await
107        .context("Failed to assign admin role to new user")
108}
109
110pub fn generate_admin_token(
111    issuer: &str,
112    user: &systemprompt_users::User,
113    session_id: &SessionId,
114) -> Result<SessionToken> {
115    let generator = SessionGenerator::new(issuer);
116    generator
117        .generate(&SessionParams {
118            user_id: &user.id,
119            session_id,
120            email: &user.email,
121            duration: ChronoDuration::hours(crate::session::api::DEFAULT_CLI_SESSION_HOURS),
122            user_type: UserType::Admin,
123            permissions: vec![Permission::Admin],
124            roles: vec!["admin".to_owned()],
125            attributes: std::collections::BTreeMap::new(),
126            rate_limit_tier: RateLimitTier::Admin,
127        })
128        .context("Failed to generate session token")
129}
130
131pub async fn create_cli_context(
132    db_pool: DbPool,
133    user: &systemprompt_users::User,
134    session_id: &SessionId,
135    profile_name: &str,
136) -> Result<ContextId> {
137    let context_repo = ContextRepository::new(&db_pool)?;
138    context_repo
139        .get_or_create_cli_context(
140            &user.id,
141            session_id,
142            &format!("CLI Session - {}", profile_name),
143        )
144        .await
145        .context("Failed to create CLI context")
146}
147
148pub(super) struct SessionComponents {
149    pub session_token: SessionToken,
150    pub session_id: SessionId,
151    pub context_id: ContextId,
152}
153
154pub(super) fn build_cli_session(
155    profile_ctx: &ProfileContext<'_>,
156    session_key: &SessionKey,
157    components: SessionComponents,
158    admin_user: &systemprompt_users::User,
159    issuer: &str,
160) -> Result<CliSession> {
161    let profile_name = ProfileName::try_new(profile_ctx.name)
162        .map_err(|e| anyhow::anyhow!("Invalid profile name: {}", e))?;
163    let email =
164        Email::try_new(&admin_user.email).map_err(|e| anyhow::anyhow!("Invalid email: {}", e))?;
165
166    Ok(CliSession::builder(
167        SessionBinding::new(profile_name, issuer.to_owned()),
168        components.session_token,
169        components.session_id,
170        components.context_id,
171        SessionIdentity::new(admin_user.id.clone(), email, UserType::Admin),
172    )
173    .with_session_key(session_key)
174    .with_profile_path(profile_ctx.path.clone())
175    .build())
176}
177
178pub async fn resolve_local_admin(
179    db_pool: &DbPool,
180    admin_name: &str,
181) -> Result<systemprompt_users::User> {
182    let user_service = UserService::new(Arc::new(UserRepository::new(db_pool)?));
183
184    let user = user_service
185        .find_by_name(admin_name)
186        .await
187        .context("Failed to query the local admin user by name")?
188        .with_context(|| {
189            format!(
190                "Local admin user '{admin_name}' not found.\n\nRun 'systemprompt admin bootstrap \
191                 --email <your email>' to create it with a real address."
192            )
193        })?;
194
195    if !user.is_active() {
196        anyhow::bail!("Local admin user '{admin_name}' exists but is not active.");
197    }
198    if !user.is_admin() {
199        anyhow::bail!(
200            "User '{admin_name}' exists but does not hold the admin role. Run 'systemprompt admin \
201             bootstrap' to repair it."
202        );
203    }
204
205    Ok(user)
206}
207
208#[doc(hidden)]
209pub async fn resolve_credentialed_user_email(session_email_hint: Option<&str>) -> Result<Email> {
210    if let Some(email) = session_email_hint {
211        return Email::try_new(email).context("session email hint is not a valid email address");
212    }
213
214    CredentialsBootstrap::try_init()
215        .await
216        .context("Failed to initialize credentials. Run 'systemprompt cloud auth login'.")?;
217
218    let creds = CredentialsBootstrap::require().map_err(|_e| {
219        anyhow::anyhow!(
220            "Cloud authentication required for new sessions.\n\nRun 'systemprompt cloud auth \
221             login' to authenticate."
222        )
223    })?;
224    Ok(creds.user_email.clone())
225}
226
227#[doc(hidden)]
228pub async fn resolve_admin_with_fallback(
229    db_pool: &DbPool,
230    user_email: &str,
231    session_email_hint: Option<&str>,
232    context_type: &str,
233) -> Result<systemprompt_users::User> {
234    match get_or_create_admin(db_pool, user_email, context_type).await {
235        Ok(user) => Ok(user),
236        Err(e) if session_email_hint.is_some() => {
237            tracing::warn!(
238                email = %user_email,
239                error = %e,
240                "Session user lookup failed, falling back to cloud credentials"
241            );
242            if let Err(init_err) = CredentialsBootstrap::try_init().await {
243                tracing::debug!(error = %init_err, "Credentials init failed during fallback");
244            }
245            if let Ok(creds) = CredentialsBootstrap::require()
246                && creds.user_email.as_str() != user_email
247            {
248                return get_or_create_admin(db_pool, creds.user_email.as_str(), context_type).await;
249            }
250            Err(e)
251        },
252        Err(e) => Err(e),
253    }
254}
255
256pub async fn resolve_tenant_admin_with_fallback(
257    db_pool: &DbPool,
258    creds: &CloudCredentials,
259    user_email: &str,
260    session_email_hint: Option<&str>,
261) -> Result<systemprompt_users::User> {
262    match get_or_create_admin(db_pool, user_email, "tenant").await {
263        Ok(user) => Ok(user),
264        Err(e) if session_email_hint.is_some() && creds.user_email.as_str() != user_email => {
265            tracing::warn!(
266                email = %user_email,
267                error = %e,
268                "Session user lookup failed, falling back to cloud credentials"
269            );
270            get_or_create_admin(db_pool, creds.user_email.as_str(), "tenant").await
271        },
272        Err(e) => Err(e),
273    }
274}