tidaldb/applications/forage/server/src/main.rs

use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;

use axum::Router;
use axum::routing::{get, post};
use clap::Parser;
use forage_engine::ForageEngine;
use tokio::sync::broadcast;
use tower_http::cors::CorsLayer;
use tower_http::services::ServeDir;

mod handlers;

/// Tracks autonomous discovery agent state.
///
/// Handlers read/write this to report agent liveness and run history.
/// All fields are async Mutexes so SSE handlers (which are async) can
/// update them without blocking the Tokio runtime.
pub struct DiscoveryState {
    /// Wall-clock timestamp of the last successful `POST /capture` from the discovery agent.
    /// Stored as `SystemTime` so it can be formatted as ISO 8601 in `/discovery/status`.
    pub last_discovery_at: tokio::sync::Mutex<Option<std::time::SystemTime>>,
    /// Timestamp of the last `POST /discovery/heartbeat` from the agent.
    pub agent_last_seen: tokio::sync::Mutex<Option<std::time::Instant>>,
    /// Number of items captured in the most recently completed discovery cycle.
    pub items_last_run: tokio::sync::Mutex<u32>,
}

impl Default for DiscoveryState {
    fn default() -> Self {
        Self {
            last_discovery_at: tokio::sync::Mutex::new(None),
            agent_last_seen: tokio::sync::Mutex::new(None),
            items_last_run: tokio::sync::Mutex::new(0),
        }
    }
}

impl DiscoveryState {
    pub fn new() -> Self {
        Self::default()
    }
}

/// Shared application state passed to every handler via Axum's `State` extractor.
pub struct AppState {
    /// The Forage engine (tidalDB wrapper + MAB logic).
    pub engine: Arc<ForageEngine>,
    /// Static bearer token for single-user auth.  Empty string means auth is
    /// disabled (the default, for backwards-compatible local dev).
    pub token: String,
    /// Broadcast channel for SSE push to connected feed pages.
    /// Sent on every successful `/capture`. Capacity of 64 is generous for
    /// a single local user — lagged receivers simply drop old events.
    pub events: broadcast::Sender<handlers::CaptureEvent>,
    /// State for the autonomous discovery loop (heartbeat, run history).
    pub discovery: Arc<DiscoveryState>,
}

#[derive(Parser)]
#[command(name = "forage-server", about = "Forage personalized feed server")]
struct Args {
    /// Use ephemeral (in-memory) storage.
    #[arg(long)]
    ephemeral: bool,

    /// Directory for persistent storage.
    #[arg(long)]
    data_dir: Option<PathBuf>,

    /// URL of the forage-embedder sidecar (e.g. http://localhost:4243).
    /// When set, add_item calls the sidecar for 1536-dim semantic vectors
    /// instead of 8-dim category-axis vectors.
    /// Requires forage-embedder to be running before the server starts.
    #[arg(long)]
    embedder: Option<String>,

    /// Directory containing static web assets (index.html, etc.).
    /// Defaults to `applications/forage/server/static` relative to CWD.
    #[arg(long)]
    static_dir: Option<PathBuf>,

    /// Port to listen on.
    #[arg(long, default_value = "4242")]
    port: u16,

    /// Static bearer token for single-user auth.
    /// When set, all API requests must include `Authorization: Bearer <token>`.
    /// Omit to disable auth (the default for local dev).
    #[arg(long)]
    token: Option<String>,
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "forage_server=debug,tower_http=info".parse().unwrap()),
        )
        .init();

    let args = Args::parse();

    let mut builder = ForageEngine::builder();
    if args.ephemeral {
        builder = builder.ephemeral();
    } else {
        let dir = args.data_dir.unwrap_or_else(|| {
            dirs_next::home_dir()
                .expect("cannot determine home directory; use --data-dir to specify storage path")
                .join(".forage/data")
        });
        builder = builder.data_dir(dir);
    }
    if let Some(url) = args.embedder {
        eprintln!("[forage-server] semantic embeddings via {url}");
        builder = builder.with_embedder(url);
    } else {
        eprintln!(
            "[forage-server] using category-axis embeddings \
             (pass --embedder <url> to enable semantic mode)"
        );
    }
    let engine = builder.open().expect("failed to open engine");

    let (events_tx, _) = broadcast::channel::<handlers::CaptureEvent>(64);
    let token = args.token.unwrap_or_default();
    if !token.is_empty() {
        eprintln!("[forage-server] auth: token required");
    }

    let state = Arc::new(AppState {
        engine: Arc::new(engine),
        token,
        events: events_tx,
        discovery: Arc::new(DiscoveryState::new()),
    });

    // Resolve static file directory.
    let static_dir = args.static_dir.unwrap_or_else(|| {
        std::env::current_dir()
            .unwrap_or_default()
            .join("applications/forage/server/static")
    });
    eprintln!("[forage-server] static files: {}", static_dir.display());
    if !static_dir.exists() {
        eprintln!(
            "[forage-server] WARNING: static dir does not exist — UI will return 404. \
             Run from the repo root or pass --static-dir <path>."
        );
    }

    let app = Router::new()
        .route("/signal", post(handlers::post_signal))
        .route("/capture", post(handlers::post_capture))
        .route("/feed", get(handlers::get_feed))
        .route("/prefs", get(handlers::get_prefs))
        .route("/items", get(handlers::get_items))
        .route("/events", get(handlers::get_events))
        .route("/browse-tasks", get(handlers::get_browse_tasks))
        .route("/discovery/heartbeat", post(handlers::post_heartbeat))
        .route("/discovery/status", get(handlers::get_discovery_status))
        .nest_service("/", ServeDir::new(static_dir))
        .layer(CorsLayer::permissive())
        .with_state(state);

    let addr = SocketAddr::from(([127, 0, 0, 1], args.port));
    eprintln!("[forage-server] listening on http://{addr}");

    let listener = tokio::net::TcpListener::bind(addr)
        .await
        .expect("failed to bind");
    axum::serve(listener, app).await.expect("server error");
}