Skip to main content

systemprompt_users/sessions/
ai_provider.rs

1//! AI session lifecycle and usage accounting owned by users.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use async_trait::async_trait;
7use systemprompt_identifiers::SessionId;
8use systemprompt_traits::{
9    ActiveSession, AiProviderError, AiProviderResult, AiSessionProvider, CreateAiSessionParams,
10};
11
12use crate::repository::SessionRepository;
13use systemprompt_traits::session_store::{
14    CreateSessionParams, SessionAttribution, SessionIdentity, SessionNetwork,
15};
16
17#[derive(Debug)]
18pub struct UsersAiSessionProvider {
19    session_repo: SessionRepository,
20}
21
22impl UsersAiSessionProvider {
23    pub const fn from_repository(session_repo: SessionRepository) -> Self {
24        Self { session_repo }
25    }
26}
27
28#[async_trait]
29impl AiSessionProvider for UsersAiSessionProvider {
30    async fn create_session(&self, params: CreateAiSessionParams<'_>) -> AiProviderResult<()> {
31        let full_params = CreateSessionParams {
32            identity: SessionIdentity {
33                session_id: params.session_id,
34                user_id: params.user_id,
35                session_source: params.session_source,
36                fingerprint_hash: None,
37            },
38            network: SessionNetwork::default(),
39            attribution: SessionAttribution::default(),
40            is_bot: false,
41            is_ai_crawler: false,
42            expires_at: params.expires_at,
43        };
44
45        self.session_repo
46            .create_session(&full_params)
47            .await
48            .map_err(|e| AiProviderError::Internal(e.into()))
49    }
50
51    async fn increment_ai_usage(
52        &self,
53        session_id: &SessionId,
54        tokens: i32,
55        cost_microdollars: i64,
56    ) -> AiProviderResult<()> {
57        self.session_repo
58            .increment_ai_usage(session_id, tokens, cost_microdollars)
59            .await
60            .map_err(|e| AiProviderError::Internal(e.into()))
61    }
62
63    async fn find_live_session(
64        &self,
65        session_id: &SessionId,
66    ) -> AiProviderResult<Option<ActiveSession>> {
67        let session = self
68            .session_repo
69            .find_active_by_id(session_id)
70            .await
71            .map_err(|e| AiProviderError::Internal(e.into()))?;
72        Ok(session.map(|row| ActiveSession {
73            user_id: row.user_id,
74        }))
75    }
76}