wsb-daemon 0.248.6265

//! Search, in two halves: the machine's sessions, and one worktree's files.
//!
//! Two routes rather than one, because the two are about different things:
//! every session on the machine — narrowed by where it ran — and the files of
//! the worktree a page is looking at. A page searching both opens both.

use wsb::prelude::*;
use axum::{
    extract::{Query, State},
    http::StatusCode,
    response::sse::{Event, KeepAlive, Sse},
    Json,
};
use std::convert::Infallible;
use std::sync::atomic::AtomicBool;

use super::sessions::ErrorResponse;
use crate::web::machine_web::{place_filter, MachineWeb};
use crate::web::search_sources::{CancelOnDrop, FileSource, SearchEvent, SessionSource};
use crate::web::state::WebState;

#[derive(Deserialize)]
pub struct SessionSearchQuery {
    pub q: String,
    pub sort: Option<String>,
    pub project: Option<String>,
    pub worktree: Option<String>,
    pub placed: Option<String>,
}

/// `GET /api/v1/search` — every session's messages.
pub async fn session_search(
    State(state): State<MachineWeb>,
    Query(params): Query<SessionSearchQuery>,
) -> Result<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>, (StatusCode, Json<ErrorResponse>)>
{
    let place = place_filter(params.project, params.worktree, params.placed)?;
    let cancelled = Arc::new(AtomicBool::new(false));
    let (tx, rx) = tokio::sync::mpsc::channel::<SearchEvent>(256);
    let source = SessionSource::new(state.db.clone(), params.sort, place);
    let c = cancelled.clone();
    tokio::spawn(async move { source.run(params.q, tx, c).await });
    Ok(stream(rx, CancelOnDrop(cancelled)))
}

#[derive(Deserialize)]
pub struct FileSearchQuery {
    pub q: String,
}

/// `GET /api/v1/worktrees/:id/search` — this worktree's files.
pub async fn file_search(
    State(state): State<WebState>,
    Query(params): Query<FileSearchQuery>,
) -> Sse<impl futures::Stream<Item = Result<Event, Infallible>>> {
    let cancelled = Arc::new(AtomicBool::new(false));
    let (tx, rx) = tokio::sync::mpsc::channel::<SearchEvent>(256);
    let source = FileSource::new(state.workspace_root.clone());
    let c = cancelled.clone();
    tokio::task::spawn_blocking(move || source.run(params.q, tx, c));
    stream(rx, CancelOnDrop(cancelled))
}

/// The results as they arrive, then `end`. Dropping the stream — the page
/// closing it — stops the search.
fn stream(
    mut rx: tokio::sync::mpsc::Receiver<SearchEvent>,
    cancel_guard: CancelOnDrop,
) -> Sse<impl futures::Stream<Item = Result<Event, Infallible>>> {
    let stream = async_stream::stream! {
        let _guard = cancel_guard;
        while let Some(event) = rx.recv().await {
            match serde_json::to_string(&event) {
                Ok(json) => yield Ok(Event::default().event("result").data(json)),
                Err(e) => yield Ok(Event::default().event("error").data(e.to_string())),
            }
        }
        yield Ok(Event::default().event("end").data("{}"));
    };
    Sse::new(stream).keep_alive(KeepAlive::default())
}