searchi/apps/desktop/src-tauri/src/ai_commands.rs

//! Explicit optional requests, grounded in a single prepared local source.
//! Credentials are request/session memory only; this module never saves config.
use crate::AppState;
use searchi_ai::{AiAnswer, AiClient, AiConfig, AiMode, AiRequest, AiSource};
use searchi_core::{SearchIndex, SourceLocation};
use serde::{Deserialize, Serialize};
use std::sync::{
    atomic::{AtomicBool, AtomicU64, Ordering},
    Arc, Mutex,
};
use tauri::State;

#[derive(Default)]
pub(crate) struct AiSession {
    prepared: Mutex<Option<Prepared>>,
    running: AtomicBool,
    latest_revision: AtomicU64,
}
#[derive(Clone)]
struct Prepared {
    id: String,
    path: String,
    source: SourceLocation,
    text: String,
    revision: u64,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct SourceIdentity {
    kind: String,
    number: u32,
}
#[derive(Serialize)]
pub(crate) struct SourcePreview {
    id: String,
    text: String,
    location: String,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct GenerateInput {
    enabled: bool,
    config: AiConfig,
    source_id: String,
    mode: AiMode,
    question: Option<String>,
}

impl AiSession {
    fn prepare(
        &self,
        index: &SearchIndex,
        path: &str,
        source: SourceLocation,
        revision: u64,
    ) -> Result<SourcePreview, String> {
        let passage = index
            .source_passage(path, &source, 4096)
            .map_err(|_| "Selected passage is unavailable.".to_owned())?
            .ok_or("Select an unchanged indexed text passage before using optional AI.")?;
        let preview = SourcePreview {
            id: uuid::Uuid::new_v4().to_string(),
            text: passage.text,
            location: format!("{} {}", source.kind, source.number),
        };
        if preview.text.trim().is_empty() || preview.text.len() > searchi_ai::MAX_SOURCE_TEXT_BYTES
        {
            return Err("The selected passage is empty or exceeds the AI text limit.".into());
        }
        let mut catalog = self
            .prepared
            .lock()
            .map_err(|_| "Selected AI source is unavailable.")?;
        if revision != self.latest_revision.load(Ordering::Acquire) {
            return Err("The selected passage changed while preparing it.".into());
        }
        *catalog = Some(Prepared {
            id: preview.id.clone(),
            path: passage.path,
            source,
            text: preview.text.clone(),
            revision,
        });
        Ok(preview)
    }
    fn grounded_request(
        &self,
        index: &SearchIndex,
        input: &GenerateInput,
    ) -> Result<AiRequest, String> {
        if !input.enabled {
            return Err(
                "Enable optional AI and explicitly request an answer before any text is sent."
                    .into(),
            );
        }
        input.config.validate().map_err(|error| error.to_string())?;
        if input.source_id.len() > 64
            || input
                .question
                .as_ref()
                .is_some_and(|question| question.len() > searchi_ai::MAX_QUESTION_BYTES)
        {
            return Err("AI request exceeds its source or question limit.".into());
        }
        let prepared = self
            .prepared
            .lock()
            .map_err(|_| "Selected AI source is unavailable.")?
            .clone()
            .filter(|value| value.id == input.source_id)
            .ok_or("Prepare and review the current selected passage before requesting AI.")?;
        if prepared.revision != self.latest_revision.load(Ordering::Acquire) {
            return Err("Review the current selected passage before requesting AI.".into());
        }
        let actual = index
            .source_passage(&prepared.path, &prepared.source, 4096)
            .map_err(|_| "Selected passage is unavailable.".to_owned())?
            .ok_or(
                "The file changed or is unavailable. Reindex it and review the passage again.",
            )?;
        if actual.text != prepared.text {
            return Err("The indexed passage changed. Review it again before sending.".into());
        }
        if matches!(input.mode, AiMode::Question)
            && input
                .question
                .as_ref()
                .is_none_or(|value| value.trim().is_empty())
        {
            return Err("Enter a question about the selected passage.".into());
        }
        Ok(AiRequest {
            mode: input.mode.clone(),
            question: if matches!(input.mode, AiMode::Question) {
                input.question.clone()
            } else {
                None
            },
            sources: vec![AiSource {
                id: "selected".into(),
                title: "Selected passage".into(),
                text: actual.text,
                location: Some(format!("{} {}", actual.source.kind, actual.source.number)),
            }],
        })
    }
}

struct RunningGuard(Arc<AiSession>);
impl Drop for RunningGuard {
    fn drop(&mut self) {
        self.0.running.store(false, Ordering::Release);
    }
}

#[tauri::command]
pub(crate) async fn prepare_ai_source(
    enabled: bool,
    path: String,
    source: SourceIdentity,
    revision: u64,
    state: State<'_, AppState>,
) -> Result<SourcePreview, String> {
    if !enabled {
        return Err("Optional AI is disabled.".into());
    }
    if path.len() > 32768 || source.kind.len() > 16 {
        return Err("Selected source identity exceeds its limit.".into());
    }
    let database = state.database.clone();
    let session = state.ai.clone();
    session
        .latest_revision
        .fetch_max(revision, Ordering::AcqRel);
    tauri::async_runtime::spawn_blocking(move || {
        let index = SearchIndex::open(database)
            .map_err(|_| "Local source index is unavailable.".to_owned())?;
        session.prepare(
            &index,
            &path,
            SourceLocation {
                kind: source.kind,
                number: source.number,
            },
            revision,
        )
    })
    .await
    .map_err(|_| "Selected passage could not be prepared.".to_owned())?
}

#[tauri::command]
pub(crate) async fn generate_ai(
    input: GenerateInput,
    state: State<'_, AppState>,
) -> Result<AiAnswer, String> {
    if !input.enabled {
        return Err("Optional AI is disabled. No text was sent.".into());
    }
    input.config.validate().map_err(|error| error.to_string())?;
    if state
        .ai
        .running
        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
        .is_err()
    {
        return Err("An earlier AI request is still finishing. Wait before trying again.".into());
    }
    let guard = RunningGuard(state.ai.clone());
    let database = state.database.clone();
    let session = state.ai.clone();
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = guard;
        let index = SearchIndex::open(database)
            .map_err(|_| "Local source index is unavailable.".to_owned())?;
        let request = session.grounded_request(&index, &input)?;
        // This explicit generate action is the sole endpoint call. Paths and
        // filenames never enter the source records sent to the client.
        let key = input.config.api_key.clone();
        let mut answer = AiClient::new(input.config)
            .and_then(|client| client.generate(&request))
            .map_err(|error| error.to_string())?;
        if let Some(key) = key.filter(|key| !key.is_empty()) {
            answer.text = answer.text.replace(&key, "[redacted credential]");
            answer.model = answer.model.replace(&key, "[redacted credential]");
        }
        Ok(answer)
    })
    .await
    .map_err(|_| "AI request could not finish.".to_owned())?
}

#[cfg(test)]
mod tests {
    use super::*;
    fn fixture() -> (tempfile::TempDir, SearchIndex, String) {
        let root = tempfile::tempdir().unwrap();
        let path = root.path().join("private-title.txt");
        std::fs::write(&path, "Supplier cancellation requires sixty days notice.").unwrap();
        let mut index = SearchIndex::open(root.path().join("index.sqlite")).unwrap();
        index.index_folder(root.path()).unwrap();
        let selected_path = index.search("private-title", 1).unwrap()[0].path.clone();
        (root, index, selected_path)
    }
    fn config() -> AiConfig {
        AiConfig {
            base_url: "http://127.0.0.1:1/v1".into(),
            model: "configured-model".into(),
            api_key: Some("session-secret-key".into()),
        }
    }
    #[test]
    fn disabled_unknown_and_changed_sources_cannot_form_network_requests() {
        let (_root, index, path) = fixture();
        let session = AiSession::default();
        let preview = session
            .prepare(
                &index,
                &path,
                SourceLocation {
                    kind: "line".into(),
                    number: 1,
                },
                0,
            )
            .unwrap();
        let mut input = GenerateInput {
            enabled: false,
            config: config(),
            source_id: preview.id.clone(),
            mode: AiMode::Summary,
            question: None,
        };
        assert!(session.grounded_request(&index, &input).is_err());
        input.enabled = true;
        input.source_id = "unknown".into();
        assert!(session.grounded_request(&index, &input).is_err());
        input.source_id = preview.id;
        std::fs::write(path, "changed").unwrap();
        assert!(session.grounded_request(&index, &input).is_err());
    }
    #[test]
    fn prepared_preview_equals_outgoing_text_without_filename_path_or_credential() {
        let (_root, index, path) = fixture();
        let session = AiSession::default();
        let preview = session
            .prepare(
                &index,
                &path,
                SourceLocation {
                    kind: "line".into(),
                    number: 1,
                },
                0,
            )
            .unwrap();
        let input = GenerateInput {
            enabled: true,
            config: config(),
            source_id: preview.id.clone(),
            mode: AiMode::Summary,
            question: None,
        };
        let request = session.grounded_request(&index, &input).unwrap();
        assert_eq!(request.sources[0].text, preview.text);
        let json = serde_json::to_string(&request).unwrap();
        assert!(!json.contains("private-title"));
        assert!(!json.contains(&path));
        assert!(!json.contains("session-secret"));
        assert!(
            !serde_json::to_string(&crate::settings::Settings::default())
                .unwrap()
                .contains("api_key")
        );
    }
    #[test]
    fn source_replacement_and_question_limits_are_enforced_offline() {
        let (_root, index, path) = fixture();
        let session = AiSession::default();
        let first = session
            .prepare(
                &index,
                &path,
                SourceLocation {
                    kind: "line".into(),
                    number: 1,
                },
                0,
            )
            .unwrap();
        session
            .prepare(
                &index,
                &path,
                SourceLocation {
                    kind: "line".into(),
                    number: 1,
                },
                0,
            )
            .unwrap();
        let input = GenerateInput {
            enabled: true,
            config: config(),
            source_id: first.id,
            mode: AiMode::Question,
            question: Some("x".repeat(2049)),
        };
        assert!(session.grounded_request(&index, &input).is_err());
    }

    #[test]
    fn older_preparation_cannot_replace_current_catalog_or_send_stale_source() {
        let (_root, index, path) = fixture();
        let session = AiSession::default();
        session.latest_revision.store(2, Ordering::Release);
        let source = SourceLocation {
            kind: "line".into(),
            number: 1,
        };
        let current = session.prepare(&index, &path, source.clone(), 2).unwrap();
        assert!(session.prepare(&index, &path, source, 1).is_err());
        let input = GenerateInput {
            enabled: true,
            config: config(),
            source_id: current.id,
            mode: AiMode::Summary,
            question: None,
        };
        assert!(session.grounded_request(&index, &input).is_ok());
        session.latest_revision.store(3, Ordering::Release);
        assert!(session.grounded_request(&index, &input).is_err());
    }
}