ai-app/server/src/routes.rs

//! The HTTP surface -- REST for actions, one SSE stream per open session
//! screen for events, all behind the bearer-token middleware `main.rs`
//! wraps the whole router in.
//!
//! ```text
//! GET    /providers                    what can be spawned
//! GET    /hosts                        machines a session can be run on
//! GET    /sessions                     list (id, provider, title, model, status, last activity)
//! POST   /sessions                     spawn {provider, title?, model?, cwd?, permissionMode?}
//! GET    /sessions/{id}/events?after=N SSE: transcript replay from N, then live
//! POST   /sessions/{id}/message        {text, attachmentIds?}
//! POST   /sessions/{id}/answer         {questionId, answer}   (questions and permissions)
//! POST   /sessions/{id}/interrupt
//! POST   /sessions/{id}/model          {model}
//! POST   /sessions/{id}/compact
//! POST   /sessions/{id}/attachments    multipart image upload -> {id}, referenced by /message
//! GET    /sessions/{id}/files/{name}   images the session produced or was sent
//! DELETE /sessions/{id}                kill process, delete transcript + files
//! GET    /usage                        cached usage windows per provider
//! ```
//!
//! Later phases add: `GET|PUT /hosts` and `/models` -- see PLAN.md's table.
//!
//! Everything here works purely in the common event model; nothing may
//! branch on the session kind (that's what drivers are for).

use std::convert::Infallible;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::Context;

use axum::Router;
use axum::extract::{Path as UrlPath, Query, State};
use axum::http::{HeaderMap, StatusCode};
use axum::response::sse::{Event as SseEvent, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use axum::routing::{delete, get, post};
use serde::Deserialize;
use tokio::sync::{broadcast, mpsc};
use tokio_stream::StreamExt;
use tokio_stream::wrappers::ReceiverStream;

use crate::session::transcript::{SeqEvent, read_after};
use crate::session::{LiveSession, SessionInfo, SessionManager, SpawnSpec};

pub fn router(manager: Arc<SessionManager>) -> Router {
    Router::new()
        .route("/providers", get(list_providers))
        .route("/hosts", get(list_hosts))
        .route("/sessions", get(list_sessions).post(spawn_session))
        .route("/sessions/{id}", delete(delete_session))
        .route("/sessions/{id}/events", get(events))
        .route("/sessions/{id}/message", post(message))
        .route("/sessions/{id}/answer", post(answer))
        .route("/sessions/{id}/interrupt", post(interrupt))
        .route("/sessions/{id}/model", post(set_model))
        .route("/sessions/{id}/compact", post(compact))
        .route("/sessions/{id}/attachments", post(upload_attachment))
        .route("/sessions/{id}/files/{name}", get(serve_file))
        // Phone photos overflow axum's 2 MB default body cap.
        .layer(axum::extract::DefaultBodyLimit::max(32 * 1024 * 1024))
        // An explicit fallback so the auth middleware (layered around the
        // whole router in main.rs) also covers unknown paths -- a scanner
        // gets the same 401 everywhere, never a route map.
        .fallback(|| async { ApiError::UnknownRoute })
        .with_state(manager)
}

#[derive(Debug, thiserror::Error)]
enum ApiError {
    #[error("{0}")]
    NotFound(String),
    #[error("no such route")]
    UnknownRoute,
    #[error("{0}")]
    BadRequest(String),
    #[error(transparent)]
    Internal(#[from] anyhow::Error),
}

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        let status = match self {
            Self::NotFound(_) | Self::UnknownRoute => StatusCode::NOT_FOUND,
            Self::BadRequest(_) => StatusCode::BAD_REQUEST,
            Self::Internal(err) => {
                // The only variant whose real cause isn't safe to hand
                // back verbatim, and the only one worth a log line.
                tracing::error!("{err:#}");
                return StatusCode::INTERNAL_SERVER_ERROR.into_response();
            }
        };
        (status, self.to_string()).into_response()
    }
}

/// An `anyhow` error from a session mutation is a message written *for*
/// the phone ("no session abc123") -- not an internal fault, so it comes
/// back as a 400 with that message rather than a 500 and a log line.
fn bad_request(err: anyhow::Error) -> ApiError {
    ApiError::BadRequest(format!("{err:#}"))
}

fn lookup(manager: &SessionManager, id: &str) -> Result<Arc<LiveSession>, ApiError> {
    manager.session(id).ok_or_else(|| ApiError::NotFound(format!("no session {id}")))
}

async fn list_sessions(State(manager): State<Arc<SessionManager>>) -> axum::Json<Vec<SessionInfo>> {
    axum::Json(manager.sessions())
}

/// What the spawn screen needs to render itself, so the phone holds no
/// hardcoded list: an entry added to `config.ron` shows up with no app
/// rebuild. Providers and hosts are listed separately because they are
/// independent choices -- any provider can be run on any host.
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct ProviderInfo {
    name: String,
    kind: crate::config::DriverKind,
    models: Vec<String>,
}

async fn list_providers(
    State(manager): State<Arc<SessionManager>>,
) -> axum::Json<Vec<ProviderInfo>> {
    axum::Json(
        manager
            .providers()
            .into_iter()
            .map(|provider| ProviderInfo {
                name: provider.name,
                kind: provider.kind,
                models: provider.models,
            })
            .collect(),
    )
}

#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct HostInfo {
    name: String,
    /// Shown under the name so a host can be told apart from its label.
    address: String,
}

/// Configured remote machines. Running on the backend itself is always
/// available and deliberately absent here -- it is the "no host" case, not
/// an entry that could be edited away.
async fn list_hosts(State(manager): State<Arc<SessionManager>>) -> axum::Json<Vec<HostInfo>> {
    axum::Json(
        manager
            .hosts()
            .into_iter()
            .map(|host| HostInfo { name: host.name, address: host.address })
            .collect(),
    )
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct SpawnRequest {
    provider: String,
    /// Name of a configured host; absent runs on the backend machine.
    #[serde(default)]
    host: Option<String>,
    #[serde(default)]
    title: Option<String>,
    #[serde(default)]
    model: Option<String>,
    #[serde(default)]
    cwd: Option<PathBuf>,
    #[serde(default)]
    permission_mode: Option<String>,
}

async fn spawn_session(
    State(manager): State<Arc<SessionManager>>,
    axum::Json(body): axum::Json<SpawnRequest>,
) -> Result<axum::Json<SessionInfo>, ApiError> {
    let info = manager
        .spawn_session(SpawnSpec {
            provider: body.provider,
            host: body.host,
            title: body.title,
            model: body.model,
            cwd: body.cwd,
            permission_mode: body.permission_mode,
        })
        .map_err(bad_request)?;
    tracing::info!("spawned {} session {} ({})", info.provider, info.id, info.title);
    Ok(axum::Json(info))
}

async fn delete_session(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
) -> Result<StatusCode, ApiError> {
    manager.delete_session(&id).map_err(bad_request)?;
    tracing::info!("deleted session {id}");
    Ok(StatusCode::NO_CONTENT)
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct MessageRequest {
    text: String,
    /// Ids from `POST /attachments`, uploaded before the message that
    /// references them.
    #[serde(default)]
    attachment_ids: Vec<String>,
}

async fn message(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
    axum::Json(body): axum::Json<MessageRequest>,
) -> Result<StatusCode, ApiError> {
    let session = lookup(&manager, &id)?;
    if body.text.trim().is_empty() && body.attachment_ids.is_empty() {
        return Err(ApiError::BadRequest("message is empty".to_string()));
    }
    session.send_message(body.text, body.attachment_ids);
    Ok(StatusCode::NO_CONTENT)
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct AnswerRequest {
    question_id: String,
    answer: String,
}

async fn answer(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
    axum::Json(body): axum::Json<AnswerRequest>,
) -> Result<StatusCode, ApiError> {
    lookup(&manager, &id)?.answer_question(&body.question_id, &body.answer);
    Ok(StatusCode::NO_CONTENT)
}

async fn interrupt(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
) -> Result<StatusCode, ApiError> {
    lookup(&manager, &id)?.interrupt();
    Ok(StatusCode::NO_CONTENT)
}

/// Separate router because its state is the usage monitor, not the
/// session manager; merged (and auth-wrapped) with the rest in `main`.
pub fn usage_router(monitor: Arc<crate::usage::UsageMonitor>) -> Router {
    Router::new().route("/usage", get(usage)).with_state(monitor)
}

async fn usage(
    State(monitor): State<Arc<crate::usage::UsageMonitor>>,
) -> Result<axum::Json<Vec<crate::usage::UsageSnapshot>>, ApiError> {
    // The fetch is blocking by design (see `usage`); off the workers.
    let snapshots = tokio::task::spawn_blocking(move || monitor.snapshots())
        .await
        .context("usage fetch panicked")?;
    Ok(axum::Json(snapshots))
}

#[derive(Deserialize)]
struct ModelRequest {
    model: String,
}

async fn set_model(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
    axum::Json(body): axum::Json<ModelRequest>,
) -> Result<StatusCode, ApiError> {
    manager.set_session_model(&id, &body.model).map_err(bad_request)?;
    Ok(StatusCode::NO_CONTENT)
}

async fn compact(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
) -> Result<StatusCode, ApiError> {
    lookup(&manager, &id)?.compact();
    Ok(StatusCode::NO_CONTENT)
}

/// Accepts one image (any multipart field) and stores it under the
/// session; the returned id goes into a later `/message`'s attachmentIds.
async fn upload_attachment(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
    mut multipart: axum::extract::Multipart,
) -> Result<axum::Json<serde_json::Value>, ApiError> {
    let session = lookup(&manager, &id)?;
    let field = multipart
        .next_field()
        .await
        .map_err(|err| ApiError::BadRequest(format!("bad upload: {err}")))?
        .ok_or_else(|| ApiError::BadRequest("no file in the upload".to_string()))?;
    let content_type = field.content_type().unwrap_or("image/jpeg").to_string();
    let bytes = field
        .bytes()
        .await
        .map_err(|err| ApiError::BadRequest(format!("upload read failed: {err}")))?;
    let name = session.save_attachment(&bytes, &content_type).map_err(bad_request)?;
    Ok(axum::Json(serde_json::json!({ "id": name })))
}

/// Serves a session's stored images -- both `files/` (produced by tools)
/// and `attachments/` (uploaded from the phone), by the id events and
/// uploads reference.
async fn serve_file(
    State(manager): State<Arc<SessionManager>>,
    UrlPath((id, name)): UrlPath<(String, String)>,
) -> Result<Response, ApiError> {
    // Ids are server-generated hex + extension; anything else (and any
    // path separator in particular) is refused, not resolved.
    if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '.') || name.contains("..") {
        return Err(ApiError::BadRequest("invalid file id".to_string()));
    }
    let session = lookup(&manager, &id)?;
    let candidates =
        [session.dir().join("files").join(&name), session.dir().join("attachments").join(&name)];
    let Some(path) = candidates.iter().find(|path| path.is_file()) else {
        return Err(ApiError::NotFound(format!("no file {name} in session {id}")));
    };
    // A file that is there but unreadable is this server's fault, not the
    // request's -- Internal logs it and says nothing more to the caller.
    let bytes = std::fs::read(path)
        .with_context(|| format!("read {}", path.display()))
        .map_err(ApiError::Internal)?;
    // Names are server-generated, so an unrecognized extension can only
    // mean a file this server didn't write.
    let content_type = crate::media::media_type_for(&name).unwrap_or("image/jpeg");
    Ok(([(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
}

#[derive(Deserialize)]
struct EventsQuery {
    #[serde(default)]
    after: u64,
}

/// The session screen's one data source: replay everything after the
/// cursor from the transcript, then live events as they happen. An SSE
/// auto-reconnect sends the last event id it saw as `Last-Event-ID`, which
/// takes precedence over `after` -- same cursor, native mechanism.
async fn events(
    State(manager): State<Arc<SessionManager>>,
    UrlPath(id): UrlPath<String>,
    Query(query): Query<EventsQuery>,
    headers: HeaderMap,
) -> Result<Sse<impl tokio_stream::Stream<Item = Result<SseEvent, Infallible>>>, ApiError> {
    let session = lookup(&manager, &id)?;
    let cursor = headers
        .get("last-event-id")
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse().ok())
        .unwrap_or(query.after);

    // Subscribe before reading the file so nothing can land in the gap
    // between replay and live; overlap is deduplicated by seq.
    let live = session.subscribe();
    let (tx, stream) = mpsc::channel(64);
    tokio::spawn(stream_session(
        session.transcript_path().to_path_buf(),
        cursor,
        live,
        tx,
    ));
    Ok(Sse::new(ReceiverStream::new(stream).map(Ok)).keep_alive(KeepAlive::default()))
}

/// Feeds one SSE subscriber: transcript replay after the cursor, then live
/// events, catching back up from the file whenever the broadcast channel
/// laps us. Ends when the client disconnects (send fails) or the session
/// is deleted (channel closed).
async fn stream_session(
    transcript: PathBuf,
    mut last: u64,
    mut live: broadcast::Receiver<SeqEvent>,
    tx: mpsc::Sender<SseEvent>,
) {
    // Synchronous file reads from an async task: transcript lines are
    // small and local; revisit if daily use produces transcripts where
    // this shows (phase 6 territory).
    let catch_up = |after: u64| match read_after(&transcript, after) {
        Ok(entries) => Some(entries),
        Err(err) => {
            tracing::error!("transcript replay failed: {err:#}");
            None
        }
    };

    let Some(replay) = catch_up(last) else { return };
    for entry in replay {
        last = entry.seq;
        if send_event(&tx, &entry).await.is_err() {
            return;
        }
    }
    loop {
        match live.recv().await {
            Ok(entry) => {
                if entry.seq <= last {
                    continue;
                }
                last = entry.seq;
                if send_event(&tx, &entry).await.is_err() {
                    return;
                }
            }
            Err(broadcast::error::RecvError::Lagged(_)) => {
                let Some(missed) = catch_up(last) else { return };
                for entry in missed {
                    last = entry.seq;
                    if send_event(&tx, &entry).await.is_err() {
                        return;
                    }
                }
            }
            Err(broadcast::error::RecvError::Closed) => return,
        }
    }
}

async fn send_event(
    tx: &mpsc::Sender<SseEvent>,
    entry: &SeqEvent,
) -> Result<(), mpsc::error::SendError<SseEvent>> {
    let data = serde_json::to_string(entry).expect("events always serialize");
    tx.send(SseEvent::default().id(entry.seq.to_string()).data(data)).await
}