yorishiro-server 0.12.1

use std::sync::Arc;

use axum::extract::FromRef;
use sqlx::PgPool;
use tokio::sync::Semaphore;
use tokio_util::task::TaskTracker;
use uuid::Uuid;
use yorishiro_core::ResultExt;
use yorishiro_core::db::TenantDb;
use yorishiro_core::repositories::entities::EntityRecord;
use yorishiro_core::services::embedding::EmbeddingProvider;
use yorishiro_core::services::embedding::sync as embedding_sync;

/// Cap on concurrent background embedding syncs. Each sync task holds a pool connection for
/// the duration of the embedding API call (up to tens of seconds), so spawning without limit
/// would exhaust the connections needed for request handling (20 total in the pool) during a
/// write burst. Tasks beyond the cap aren't dropped — they wait on the semaphore without
/// holding a connection.
const EMBEDDING_SYNC_MAX_CONCURRENCY: usize = 4;

/// Application state shared by both the REST and MCP handlers. Using this struct as axum's
/// `State` — rather than `TenantDb` alone — lets search handlers also reach the
/// `EmbeddingProvider`.
#[derive(Clone)]
pub struct AppState {
    pub tenant_db: TenantDb,
    /// A connection pool using the admin/migration role (not `yorishiro_app`), reserved for
    /// the handful of control-plane endpoints (signup, login, invite redemption) that must
    /// read/write `identity.users`/`identity.tenant_memberships`/`identity.invites` before any
    /// tenant or workspace context can be established -- the same role `admin.rs`'s CLI
    /// commands already use, for the same reason (see the role-separation migration's comment
    /// on why `yorishiro_app` has no grant on those tables at all). Every other handler must
    /// keep using `tenant_db` instead: this pool bypasses RLS entirely.
    pub identity_pool: PgPool,
    pub embedding_provider: Arc<dyn EmbeddingProvider>,
    embedding_sync_permits: Arc<Semaphore>,
    embedding_tasks: TaskTracker,
}

impl AppState {
    pub fn new(
        tenant_db: TenantDb,
        identity_pool: PgPool,
        embedding_provider: Arc<dyn EmbeddingProvider>,
    ) -> Self {
        Self {
            tenant_db,
            identity_pool,
            embedding_provider,
            embedding_sync_permits: Arc::new(Semaphore::new(EMBEDDING_SYNC_MAX_CONCURRENCY)),
            embedding_tasks: TaskTracker::new(),
        }
    }

    /// Tracker used to wait for in-flight embedding syncs during graceful shutdown.
    /// `main` calls `close()` + `wait()` on it after the HTTP server stops.
    pub fn embedding_tasks(&self) -> &TaskTracker {
        &self.embedding_tasks
    }

    /// Syncs the `embedding` column in the background after an entity create/update
    /// succeeds. The embedding API call can take up to tens of seconds, so the request isn't
    /// made to wait for it, and a fresh connection is acquired from the pool instead of
    /// reusing the request's own connection (satisfying the no-same-transaction constraint
    /// documented on `sync_embedding`). Failures are only logged: embedding is an auxiliary
    /// feature and must not affect whether the entity write itself succeeds.
    pub fn spawn_embedding_sync(
        &self,
        tenant_id: Uuid,
        workspace_id: Uuid,
        record: EntityRecord,
    ) -> tokio::task::JoinHandle<()> {
        let db = self.tenant_db.clone();
        let provider = Arc::clone(&self.embedding_provider);
        let permits = Arc::clone(&self.embedding_sync_permits);
        // Spawning through the TaskTracker lets graceful shutdown wait for the embedding
        // sync of an already-written entity to finish (an immediate SIGTERM exit would lose
        // the sync, leaving that entity permanently missing from search).
        self.embedding_tasks.spawn(async move {
            // The order matters: acquire the permit before the connection. Reversing it
            // would let every waiting task hold a connection, defeating the point of the cap.
            let Ok(_permit) = permits.acquire_owned().await else {
                // Unreachable in practice: the semaphore is never closed.
                return;
            };

            let result = async {
                let mut conn = db
                    .acquire_for_workspace(tenant_id, workspace_id)
                    .await
                    .internal()?;
                embedding_sync::sync_embedding_for_record(
                    &mut conn,
                    workspace_id,
                    &record,
                    provider.as_ref(),
                )
                .await
            }
            .await;

            if let Err(err) = result {
                tracing::warn!(entity_id = %record.id, error = %err, "embedding sync failed");
            }
        })
    }
}

impl FromRef<AppState> for TenantDb {
    fn from_ref(state: &AppState) -> Self {
        state.tenant_db.clone()
    }
}