systemprompt-agent 0.55.1

//! Context insert/update/delete mutations.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use chrono::Utc;

use super::ContextRepository;
use crate::models::context::ContextKind;
use systemprompt_identifiers::{ContextId, SessionId, UserId};
use systemprompt_traits::RepositoryError;

impl ContextRepository {
    pub async fn create_context(
        &self,
        user_id: &UserId,
        session_id: Option<&SessionId>,
        name: &str,
        kind: ContextKind,
    ) -> Result<ContextId, RepositoryError> {
        let context_id = ContextId::generate();
        let now = Utc::now();
        let session_id_str = session_id.map(SessionId::as_str);

        sqlx::query!(
            "INSERT INTO user_contexts (context_id, user_id, session_id, name, kind, created_at, \
             updated_at)
             VALUES ($1, $2, $3, $4, $5, $6, $6)",
            context_id.as_str(),
            user_id.as_str(),
            session_id_str,
            name,
            kind.as_str(),
            now
        )
        .execute(&*self.write_pool)
        .await
        .map_err(RepositoryError::database)?;

        Ok(context_id)
    }

    // Why: the context id derives from caller-supplied metadata, so the
    // conflict update is scoped to the owning user — another user's call on
    // the same id must be a no-op rather than a write into their row.
    pub async fn ensure_context(
        &self,
        params: &systemprompt_traits::EnsureContextParams<'_>,
        kind: ContextKind,
    ) -> Result<(), RepositoryError> {
        let context_id = params.context_id;
        let user_id = params.user_id;
        let session_id = params.session_id;
        let name = params.name;
        let now = Utc::now();

        let result = sqlx::query!(
            "INSERT INTO user_contexts (context_id, user_id, session_id, name, kind, created_at, \
             updated_at)
             VALUES ($1, $2, $3, $4, $5, $6, $6)
             ON CONFLICT (context_id) DO UPDATE
             SET updated_at = EXCLUDED.updated_at,
                 session_id = COALESCE(user_contexts.session_id, EXCLUDED.session_id)
             WHERE user_contexts.user_id = EXCLUDED.user_id",
            context_id.as_str(),
            user_id.as_str(),
            session_id.map(SessionId::as_str),
            name,
            kind.as_str(),
            now
        )
        .execute(&*self.write_pool)
        .await
        .map_err(RepositoryError::database)?;

        if result.rows_affected() != 1 {
            return Err(RepositoryError::NotFound(format!(
                "Context {} not found for user {}",
                context_id, user_id
            )));
        }
        Ok(())
    }

    // Why: the legacy context is the one row whose owner may be rebound — it is
    // system-owned and adopted by whichever admin the current profile names.
    // The id is fixed here so no caller can use this path to reassign another
    // user's context.
    pub async fn ensure_legacy_context(
        &self,
        system_admin: &UserId,
    ) -> Result<(), RepositoryError> {
        let now = Utc::now();
        let legacy = ContextId::legacy();
        sqlx::query!(
            "INSERT INTO user_contexts (context_id, user_id, session_id, name, kind, created_at, \
             updated_at)
             VALUES ($1, $2, NULL, $3, $4, $5, $5)
             ON CONFLICT (context_id) DO UPDATE
             SET user_id = EXCLUDED.user_id,
                 updated_at = EXCLUDED.updated_at",
            legacy.as_str(),
            system_admin.as_str(),
            "Legacy (pre-context)",
            ContextKind::Legacy.as_str(),
            now
        )
        .execute(&*self.write_pool)
        .await
        .map_err(RepositoryError::database)?;
        Ok(())
    }

    pub async fn get_or_create_cli_context(
        &self,
        user_id: &UserId,
        session_id: &SessionId,
        name: &str,
    ) -> Result<ContextId, RepositoryError> {
        let now = Utc::now();

        let adopted = sqlx::query_scalar!(
            r#"UPDATE user_contexts SET session_id = $1, updated_at = $2
             WHERE context_id = (
                 SELECT context_id FROM user_contexts
                 WHERE user_id = $3 AND kind = $4 AND name = $5
                 ORDER BY updated_at DESC LIMIT 1
             )
             RETURNING context_id"#,
            session_id.as_str(),
            now,
            user_id.as_str(),
            ContextKind::CliSession.as_str(),
            name
        )
        .fetch_optional(&*self.write_pool)
        .await
        .map_err(RepositoryError::database)?;

        match adopted {
            Some(context_id) => ContextId::try_new(context_id)
                .map_err(|e| RepositoryError::InvalidData(e.to_string())),
            None => {
                self.create_context(user_id, Some(session_id), name, ContextKind::CliSession)
                    .await
            },
        }
    }

    pub async fn validate_context_ownership(
        &self,
        context_id: &ContextId,
        user_id: &UserId,
    ) -> Result<(), RepositoryError> {
        let result = sqlx::query_scalar!(
            "SELECT context_id FROM user_contexts WHERE context_id = $1 AND user_id = $2",
            context_id.as_str(),
            user_id.as_str()
        )
        .fetch_optional(&*self.pool)
        .await
        .map_err(RepositoryError::database)?;

        match result {
            Some(_) => Ok(()),
            None => Err(RepositoryError::NotFound(format!(
                "Context {} not found or user {} does not have access",
                context_id, user_id
            ))),
        }
    }

    pub async fn update_context_name(
        &self,
        context_id: &ContextId,
        user_id: &UserId,
        name: &str,
    ) -> Result<(), RepositoryError> {
        let now = Utc::now();

        let result = sqlx::query!(
            "UPDATE user_contexts SET name = $1, updated_at = $2
             WHERE context_id = $3 AND user_id = $4",
            name,
            now,
            context_id.as_str(),
            user_id.as_str()
        )
        .execute(&*self.write_pool)
        .await
        .map_err(RepositoryError::database)?;

        if result.rows_affected() == 0 {
            return Err(RepositoryError::NotFound(format!(
                "Context {} not found for user {}",
                context_id, user_id
            )));
        }

        Ok(())
    }

    pub async fn delete_context(
        &self,
        context_id: &ContextId,
        user_id: &UserId,
    ) -> Result<(), RepositoryError> {
        let result = sqlx::query!(
            "DELETE FROM user_contexts WHERE context_id = $1 AND user_id = $2",
            context_id.as_str(),
            user_id.as_str()
        )
        .execute(&*self.write_pool)
        .await
        .map_err(RepositoryError::database)?;

        if result.rows_affected() == 0 {
            return Err(RepositoryError::NotFound(format!(
                "Context {} not found for user {}",
                context_id, user_id
            )));
        }

        Ok(())
    }
}