Files
ai-app/server/src/routes.rs
T

1718 lines
57 KiB
Rust

use std::convert::Infallible;
use std::path::{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::{get, post};
use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, mpsc};
use tokio_stream::StreamExt;
use tokio_stream::wrappers::{BroadcastStream, ReceiverStream};
use crate::session::driver::{SessionCommand, Unqueued};
use crate::session::pending::Operation;
use crate::session::subagent::{Subagent, SubagentInfo};
use crate::session::transcript::{CATCH_UP_LIMIT, CatchUp, SeqEvent, catch_up};
use crate::session::{LiveSession, SessionInfo, SessionManager, SpawnSpec};
pub fn router(manager: Arc<SessionManager>) -> Router {
Router::new()
.route("/setups", get(list_setups).post(add_setup))
.route("/setups/probe", post(probe_setup))
.route("/setups/{id}/importable", get(list_importable))
// A batch at a time, never a session at a time -- see
// [`delete_importable`].
.route("/setups/{id}/importable/delete", post(delete_importable))
.route("/setups/{id}/importable/import", post(start_import))
.route("/setups/{id}/importable/events", get(importable_events))
.route(
"/setups/{id}",
get(read_setup).put(update_setup).delete(delete_setup),
)
.route("/setups/{id}/models", get(setup_models))
// The filesystem of the machine a setup names. Under the setup
// rather than under a session because a filesystem is a property of
// a machine; a session only says where to start looking.
.route("/setups/{id}/dir", get(list_dir).post(create_dir))
.route(
"/setups/{id}/file",
get(read_file).put(write_file).post(create_file),
)
.route("/sessions", get(list_sessions).post(spawn_session))
.route("/sessions/{id}", get(read_session).delete(delete_session))
.route("/sessions/{id}/events", get(events))
.route("/sessions/{id}/transcript", get(transcript))
.route("/sessions/{id}/subagents", get(list_subagents))
.route(
"/sessions/{id}/subagents/{sub}/transcript",
get(subagent_transcript),
)
.route(
"/sessions/{id}/subagents/{sub}/events",
get(subagent_events),
)
.route("/sessions/{id}/message", post(message))
.route("/sessions/{id}/unqueue", post(unqueue))
.route("/sessions/{id}/answer", post(answer))
.route("/sessions/{id}/interrupt", post(interrupt))
.route("/sessions/{id}/stop", post(stop))
.route("/sessions/{id}/start", post(start))
.route("/sessions/{id}/title", post(rename))
.route("/sessions/{id}/cwd", post(set_cwd))
.route("/sessions/{id}/model", post(set_model))
.route("/sessions/{id}/permission-mode", post(set_permission_mode))
.route("/sessions/{id}/effort", post(set_effort))
.route("/defaults", get(defaults).post(set_defaults))
.route("/sessions/{id}/notify", post(set_notify))
.route("/sessions/{id}/auto-resume", post(set_auto_resume))
.route("/notifications", get(notifications))
.route("/sessions/{id}/compact", post(compact))
.route("/sessions/{id}/command", post(command))
.route(
"/sessions/{id}/attachments",
// A trace or a log is bigger than a photo; the cap below is for
// everything else, and the innermost limit is the one axum
// applies.
post(upload_attachment).layer(axum::extract::DefaultBodyLimit::max(ATTACHMENT_LIMIT)),
)
.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 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("{0}")]
Conflict(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::Conflict(_) => StatusCode::CONFLICT,
Self::Internal(err) => {
tracing::error!("{err:#}");
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}
};
(status, self.to_string()).into_response()
}
}
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}")))
}
/// A session's subagent by id -- the second half of the lookup every
/// `/sessions/{id}/subagents/{sub}/...` route needs. `Arc` because reopening
/// one from disk (a subagent this process has not touched yet) inserts it
/// into the registry, and a route holding a borrow across that would be
/// holding the registry's lock the whole request.
fn lookup_subagent(session: &LiveSession, sub: &str) -> Result<Arc<Subagent>, ApiError> {
session
.subagents()
.get(sub)
.ok_or_else(|| ApiError::NotFound(format!("no subagent {sub}")))
}
async fn list_sessions(State(manager): State<Arc<SessionManager>>) -> axum::Json<Vec<SessionInfo>> {
axum::Json(manager.sessions())
}
async fn read_session(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<axum::Json<SessionInfo>, ApiError> {
manager
.sessions()
.into_iter()
.find(|session| session.id == id)
.map(axum::Json)
.ok_or_else(|| ApiError::NotFound(format!("no session {id}")))
}
/// One list rather than two, because the halves are not independent. A
/// provider only exists on a machine that has it installed, so listing them
/// separately offered the whole cross-product -- including "the Claude CLI on
/// the box that hasn't got it".
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct SetupInfo {
id: String,
name: String,
#[serde(skip_serializing_if = "Option::is_none")]
address: Option<String>,
providers: Vec<ProviderInfo>,
}
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct ProviderInfo {
name: String,
kind: crate::config::DriverKind,
models: Vec<String>,
}
async fn list_setups(State(manager): State<Arc<SessionManager>>) -> axum::Json<Vec<SetupInfo>> {
axum::Json(manager.setups().into_iter().map(info_for).collect())
}
fn info_for(setup: crate::config::SetupConfig) -> SetupInfo {
SetupInfo {
id: setup.id,
name: setup.name,
address: setup.ssh.map(|ssh| ssh.address),
providers: setup
.providers
.into_iter()
.map(|provider| ProviderInfo {
name: provider.name,
kind: provider.kind,
models: provider.models,
})
.collect(),
}
}
/// Note what is absent: nothing here names a program. Providers are found by
/// asking the machine (`crate::setups`), never sent, so the enrolled token
/// cannot introduce something to run.
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct SshRequest {
address: String,
#[serde(default)]
port: Option<u16>,
/// A path on the *backend*, not a key itself: private keys do not travel,
/// so this names one that must already be there.
#[serde(default)]
identity_file: Option<String>,
#[serde(default)]
options: Vec<String>,
#[serde(default)]
attachments_dir: Option<String>,
#[serde(default)]
models_dir: Option<String>,
}
impl SshRequest {
fn into_config(self) -> Result<crate::config::SshConfig, ApiError> {
let address = crate::setups::tidy(&self.address)
.ok_or_else(|| ApiError::BadRequest("a machine needs an address".to_string()))?;
Ok(crate::config::SshConfig {
address,
port: self.port,
identity_file: self
.identity_file
.as_deref()
.and_then(crate::setups::tidy)
.map(std::path::PathBuf::from),
options: self
.options
.iter()
.filter_map(|o| crate::setups::tidy(o))
.collect(),
attachments_dir: self
.attachments_dir
.as_deref()
.map(str::trim)
.filter(|dir| !dir.is_empty())
.map(std::path::PathBuf::from),
models_dir: self
.models_dir
.as_deref()
.map(str::trim)
.filter(|dir| !dir.is_empty())
.map(std::path::PathBuf::from),
})
}
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct AddSetupRequest {
name: String,
#[serde(default)]
ssh: Option<SshRequest>,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct ProbeRequest {
#[serde(default)]
ssh: Option<SshRequest>,
}
async fn probe_setup(
axum::Json(body): axum::Json<ProbeRequest>,
) -> Result<axum::Json<Vec<ProviderInfo>>, ApiError> {
let ssh = body.ssh.map(SshRequest::into_config).transpose()?;
let providers = probe(ssh, "this setup").await?;
Ok(axum::Json(
providers
.into_iter()
.map(|provider| ProviderInfo {
name: provider.name,
kind: provider.kind,
models: provider.models,
})
.collect(),
))
}
async fn probe(
ssh: Option<crate::config::SshConfig>,
label: &str,
) -> Result<Vec<crate::config::ProviderConfig>, ApiError> {
let transport = match ssh {
Some(ssh) => crate::session::transport::Transport::Ssh {
name: label.to_string(),
ssh,
},
None => crate::session::transport::Transport::Here,
};
crate::setups::discover(&transport)
.await
.map_err(bad_request)
}
async fn add_setup(
State(manager): State<Arc<SessionManager>>,
axum::Json(body): axum::Json<AddSetupRequest>,
) -> Result<axum::Json<SetupInfo>, ApiError> {
let ssh = body.ssh.map(SshRequest::into_config).transpose()?;
let providers = probe(ssh.clone(), &body.name).await?;
let setup = manager
.add_setup(&body.name, ssh, providers)
.map_err(bad_request)?;
Ok(axum::Json(info_for(setup)))
}
fn setup_by_id(
manager: &Arc<SessionManager>,
id: &str,
) -> Result<crate::config::SetupConfig, ApiError> {
manager
.setups()
.into_iter()
.find(|setup| setup.id == id)
.ok_or_else(|| ApiError::NotFound(format!("no setup {id}")))
}
async fn read_setup(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<axum::Json<SetupInfo>, ApiError> {
setup_by_id(&manager, &id).map(|setup| axum::Json(info_for(setup)))
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct UpdateSetupRequest {
#[serde(default)]
name: Option<String>,
#[serde(default)]
rediscover: bool,
}
async fn update_setup(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<UpdateSetupRequest>,
) -> Result<axum::Json<SetupInfo>, ApiError> {
let providers = if body.rediscover {
let existing = manager
.setups()
.into_iter()
.find(|setup| setup.id == id)
.ok_or_else(|| ApiError::NotFound(format!("no setup {id}")))?;
let transport = crate::session::transport::Transport::for_setup(&existing);
Some(
crate::setups::discover(&transport)
.await
.map_err(bad_request)?,
)
} else {
None
};
let setup = manager
.update_setup(&id, body.name.as_deref(), providers)
.map_err(bad_request)?;
Ok(axum::Json(info_for(setup)))
}
async fn delete_setup(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<StatusCode, ApiError> {
manager.delete_setup(&id).map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
fn files_on(
manager: &Arc<SessionManager>,
id: &str,
path: &str,
) -> Result<(crate::session::transport::Transport, String), ApiError> {
let setup = setup_by_id(manager, id)?;
let path = crate::files::check_path(path).map_err(bad_request)?;
Ok((
crate::session::transport::Transport::for_setup(&setup),
path,
))
}
fn from_machine(err: anyhow::Error) -> ApiError {
ApiError::BadRequest(format!("{err:#}"))
}
#[derive(Deserialize)]
struct PathQuery {
path: String,
}
async fn setup_models(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<axum::Json<Vec<crate::models::LocalModel>>, ApiError> {
let setup = setup_by_id(&manager, &id)?;
let transport = crate::session::transport::Transport::for_setup(&setup);
let dir = crate::models::dir_on(&transport, manager.models_dir());
crate::models::on_machine(&transport, &dir)
.await
.map(axum::Json)
.map_err(from_machine)
}
async fn list_dir(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
Query(query): Query<PathQuery>,
) -> Result<axum::Json<crate::files::Listing>, ApiError> {
let (transport, path) = files_on(&manager, &id, &query.path)?;
crate::files::list(&transport, &path)
.await
.map(axum::Json)
.map_err(from_machine)
}
#[derive(Serialize)]
struct FileResponse {
path: String,
#[serde(flatten)]
read: crate::files::FileRead,
}
async fn read_file(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
Query(query): Query<PathQuery>,
) -> Result<axum::Json<FileResponse>, ApiError> {
let (transport, path) = files_on(&manager, &id, &query.path)?;
let read = crate::files::read(&transport, &path)
.await
.map_err(from_machine)?;
Ok(axum::Json(FileResponse { path, read }))
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct WriteFileRequest {
path: String,
content: String,
if_sha256: String,
}
async fn write_file(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<WriteFileRequest>,
) -> Result<axum::Json<crate::files::Written>, ApiError> {
let (transport, path) = files_on(&manager, &id, &body.path)?;
crate::files::write(
&transport,
&path,
&body.if_sha256,
body.content.into_bytes(),
)
.await
.map_err(from_machine)?
.map(axum::Json)
.map_err(|crate::files::Stale| {
ApiError::Conflict("this file changed on the machine since you opened it".to_string())
})
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct CreateRequest {
path: String,
}
async fn create_file(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<CreateRequest>,
) -> Result<StatusCode, ApiError> {
let (transport, path) = files_on(&manager, &id, &body.path)?;
crate::files::create_file(&transport, &path)
.await
.map_err(from_machine)?;
Ok(StatusCode::NO_CONTENT)
}
async fn create_dir(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<CreateRequest>,
) -> Result<StatusCode, ApiError> {
let (transport, path) = files_on(&manager, &id, &body.path)?;
crate::files::create_dir(&transport, &path)
.await
.map_err(from_machine)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct SpawnRequest {
setup: String,
provider: String,
#[serde(default)]
title: Option<String>,
#[serde(default)]
model: Option<String>,
#[serde(default)]
cwd: Option<PathBuf>,
#[serde(default)]
permission_mode: Option<String>,
#[serde(default)]
effort: Option<String>,
#[serde(default)]
params: std::collections::BTreeMap<String, String>,
#[serde(default)]
import: Option<String>,
}
async fn list_importable(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<axum::Json<Vec<ImportableRow>>, ApiError> {
let setup = setup_by_id(&manager, &id)?;
let transport = crate::session::transport::Transport::for_setup(&setup);
let mut found = crate::session::import::list(&transport)
.await
.map_err(bad_request)?;
// Anything this app is already continuing is not offered again. Left out
// rather than shown-and-disabled, because it has not disappeared: it is in
// the session list, which is where it now belongs.
found.retain(|candidate| {
manager.pending().running(&id, &candidate.id).is_some()
|| manager.session_driving(&candidate.id).is_none()
});
// What the server is doing to each of them, joined on here because a phone
// that was asleep or freshly opened never heard the events -- see
// `pending`. A row being imported has to stay visible, marked, or the list
// would say the work never started.
let present: Vec<String> = found.iter().map(|row| row.id.clone()).collect();
manager.pending().prune(&id, &present);
let rows: Vec<ImportableRow> = found
.into_iter()
.map(|importable| ImportableRow {
pending: manager
.pending()
.running(&id, &importable.id)
.map(|operation| operation.label()),
error: manager.pending().failure(&id, &importable.id),
importable,
})
.collect();
Ok(axum::Json(rows))
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ImportableRow {
#[serde(flatten)]
importable: crate::session::import::Importable,
#[serde(skip_serializing_if = "Option::is_none")]
pending: Option<&'static str>,
/// How the last attempt on this row failed, if it did. Kept until
/// something replaces it, because the phone that needs to see it may not
/// have been connected when it happened.
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
/// A batch and never a single session, which is why this is a POST with a body
/// rather than a `DELETE` on each id. One request per row made a handover only
/// as atomic as the network: leave the screen, lose signal, or have the fourth
/// of six requests fail, and some rows are being deleted while the rest are
/// untouched, with nothing anywhere that knows the difference. Here every id
/// is registered as in flight before the 202 goes back.
///
/// Registering is what has to be atomic; the work is not. Each row settles on
/// its own event from its own outcome, because six deletes that must all
/// succeed or all roll back is not something a filesystem offers.
async fn delete_importable(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<DeleteBatch>,
) -> Result<StatusCode, ApiError> {
let setup = setup_by_id(&manager, &id)?;
let transport = crate::session::transport::Transport::for_setup(&setup);
let running: Vec<(String, crate::session::pending::InFlight)> = body
.sessions
.iter()
.map(|session| {
(
session.clone(),
manager.pending().begin(&id, session, Operation::Deleting),
)
})
.collect();
let sessions = body.sessions;
tokio::spawn(async move {
let outcomes = match crate::session::import::delete(&transport, &sessions).await {
Ok(outcomes) => outcomes,
Err(err) => {
let message = format!("{err:#}");
tracing::warn!(
"deleting {} sessions on {id} failed: {message}",
running.len()
);
for (_, flight) in running {
flight.failed(message.clone());
}
return;
}
};
for (session, flight) in running {
match outcomes.get(&session) {
Some(Ok(())) => {
tracing::info!("deleting {session} on {id}: done");
flight.succeeded();
}
Some(Err(message)) => {
tracing::warn!("deleting {session} on {id} failed: {message}");
flight.failed(message.clone());
}
None => flight.failed(format!("nothing was reported about {session}")),
}
}
});
Ok(StatusCode::ACCEPTED)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct DeleteBatch {
sessions: Vec<String>,
}
/// Separate from `POST /sessions` because the two are asked different
/// questions. That one means "start this and take me to it", so it waits and
/// answers with the session. This one is the import screen's batch: several at
/// once, nobody waiting on any particular one, and the answer arrives as a row
/// changing rather than as a reply -- the screen it was started from may well
/// be gone by then.
async fn start_import(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<ImportRequest>,
) -> Result<StatusCode, ApiError> {
// Checked before accepting, so an unknown machine is an error the caller
// sees rather than one it has to go and read off a row.
setup_by_id(&manager, &id)?;
for session in body.sessions {
let request = SpawnRequest {
setup: id.clone(),
provider: body.provider.clone(),
title: None,
model: body.model.clone(),
cwd: None,
permission_mode: body.permission_mode.clone(),
effort: body.effort.clone(),
params: std::collections::BTreeMap::new(),
import: Some(session.clone()),
};
let inner = Arc::clone(&manager);
in_background(
&manager,
id.clone(),
session,
Operation::Importing,
async move {
spawn(&inner, request)
.await
.map(|_| ())
.map_err(|err| anyhow::anyhow!("{err}"))
},
);
}
Ok(StatusCode::ACCEPTED)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct ImportRequest {
sessions: Vec<String>,
provider: String,
#[serde(default)]
model: Option<String>,
#[serde(default)]
permission_mode: Option<String>,
#[serde(default)]
effort: Option<String>,
}
/// Runs `work` on the server, marked as in flight for as long as it takes.
/// Spawned rather than awaited, which is the whole difference: the phone asked
/// for it, but the phone leaving must not cancel it. What replaces the reply
/// is the pending registry.
fn in_background<F>(
manager: &Arc<SessionManager>,
setup: String,
session: String,
operation: Operation,
work: F,
) where
F: std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
{
let running = manager.pending().begin(&setup, &session, operation);
tokio::spawn(async move {
match work.await {
Ok(()) => {
tracing::info!("{} {session} on {setup}: done", operation.label());
running.succeeded();
}
Err(err) => {
tracing::warn!("{} {session} on {setup} failed: {err:#}", operation.label());
running.failed(format!("{err:#}"));
}
}
});
}
async fn importable_events(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Sse<impl tokio_stream::Stream<Item = Result<SseEvent, Infallible>>> {
let live = manager.pending().subscribe();
let stream = BroadcastStream::new(live).filter_map(move |item| {
let change = item.ok()?;
if change.setup() != id {
return None;
}
Some(Ok(SseEvent::default().json_data(&change).ok()?))
});
Sse::new(stream).keep_alive(KeepAlive::default())
}
async fn spawn_session(
State(manager): State<Arc<SessionManager>>,
axum::Json(body): axum::Json<SpawnRequest>,
) -> Result<axum::Json<SessionInfo>, ApiError> {
spawn(&manager, body).await.map(axum::Json)
}
/// A function rather than only a handler because the import screen's batch runs
/// this from a background task. Spawning has to mean exactly the same thing
/// either way: the same refusal when something else already has the
/// conversation open, the same title, the same working directory.
async fn spawn(manager: &Arc<SessionManager>, body: SpawnRequest) -> Result<SessionInfo, ApiError> {
// Resolved before the spawn because both halves are the machine's answer,
// not the phone's: which file that id names, and what is in it.
let seed = match &body.import {
Some(want) => {
let setup = setup_by_id(manager, &body.setup)?;
let transport = crate::session::transport::Transport::for_setup(&setup);
let chosen = crate::session::import::find(&transport, want)
.await
.map_err(bad_request)?
.ok_or_else(|| {
ApiError::NotFound(format!(
"setup \"{}\" has no Claude Code session {want} to import",
body.setup
))
})?;
if let Some(existing) = manager.session_driving(want) {
return Err(ApiError::BadRequest(format!(
"session {existing} is already continuing that one -- delete it first if you \
want a fresh copy. Deleting it here does not touch the conversation itself, \
only this app's view of it."
)));
}
if chosen.in_use == crate::session::import::InUse::Yes {
return Err(ApiError::BadRequest(format!(
"{want} is open in a terminal right now. Importing it would put a second \
Claude Code on the same conversation, which duplicates it and re-reads the \
whole thing. Close it there first, then import it here."
)));
}
let records = crate::session::import::read_tail(&transport, &chosen.path)
.await
.map_err(bad_request)?;
let mut chosen = chosen;
if !crate::session::import::directory_exists(&transport, &chosen.cwd).await {
tracing::warn!(
"imported session {} recorded {} as its directory, which is not there any \
more -- starting in the default one instead",
chosen.id,
chosen.cwd,
);
chosen.cwd = String::new();
}
Some((chosen, records))
}
None => None,
};
let spec = SpawnSpec {
setup: body.setup,
provider: body.provider,
title: body
.title
.filter(|title| !title.trim().is_empty())
.or_else(|| {
seed.as_ref()
.map(|(chosen, _)| chosen.title.clone())
.filter(|title| !title.trim().is_empty())
}),
model: body.model,
cwd: body.cwd.or_else(|| {
seed.as_ref()
.map(|(chosen, _)| PathBuf::from(&chosen.cwd))
.filter(|cwd| cwd.as_os_str() != "")
}),
permission_mode: body.permission_mode,
effort: body.effort,
params: body.params,
};
let info = match seed {
Some((chosen, records)) => {
tracing::info!(
"importing Claude Code session {} ({} lines replayed)",
chosen.id,
records.lines().count()
);
manager.spawn_imported(
spec,
crate::session::Seed {
resume: chosen.id.clone(),
cursor: crate::session::import::Cursor {
path: chosen.path,
lines: chosen.lines,
},
records,
},
)
}
None => manager.spawn_session(spec),
}
.map_err(bad_request)?;
tracing::info!(
"spawned {} session {} ({})",
info.provider,
info.id,
info.title
);
Ok(info)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct DeleteSessionQuery {
/// Also remove the machine's own transcript of this conversation -- the
/// file Claude Code keeps under `~/.claude/projects`, which this server's
/// delete does not otherwise touch.
///
/// Off by default, because the two deletes differ in what they cost:
/// leaving the machine's copy is recoverable and removing it is not, and a
/// default is the one choice nobody is shown.
#[serde(default)]
delete_foreign: bool,
}
async fn delete_session(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
Query(query): Query<DeleteSessionQuery>,
) -> Result<StatusCode, ApiError> {
// Before the session goes, because only the session record says which file
// on which machine this conversation is.
let foreign = query
.delete_foreign
.then(|| manager.foreign_transcript(&id))
.flatten();
if let Some((setup, session)) = &foreign {
let setup = setup_by_id(&manager, setup)?;
let transport = crate::session::transport::Transport::for_setup(&setup);
crate::session::import::delete(&transport, std::slice::from_ref(session))
.await
.map_err(bad_request)?
.remove(session)
.unwrap_or_else(|| Err(format!("nothing was reported about {session}")))
.map_err(|message| bad_request(anyhow::anyhow!("{message}")))?;
tracing::info!("deleted Claude Code session {session} with ai-app session {id}");
}
manager.delete_session(&id).map_err(bad_request)?;
tracing::info!("deleted session {id}");
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct MessageRequest {
text: String,
#[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> {
lookup(&manager, &id)?;
if body.text.trim().is_empty() && body.attachment_ids.is_empty() {
return Err(ApiError::BadRequest("message is empty".to_string()));
}
manager
.send_message(&id, body.text, body.attachment_ids)
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct UnqueueRequest {
message_id: String,
}
async fn unqueue(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<UnqueueRequest>,
) -> Result<StatusCode, ApiError> {
match lookup(&manager, &id)?.unqueue(&body.message_id) {
Unqueued::Dropped => Ok(StatusCode::NO_CONTENT),
Unqueued::AlreadySent => Err(ApiError::Conflict(
"the session has already been given this message".to_string(),
)),
Unqueued::Unknown => Err(ApiError::NotFound(
"this message is not waiting to be read".to_string(),
)),
}
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct AnswerRequest {
question_id: String,
answers: Vec<String>,
}
async fn answer(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<AnswerRequest>,
) -> Result<StatusCode, ApiError> {
if body.answers.is_empty() {
return Err(bad_request(anyhow::anyhow!(
"an answer needs at least one choice"
)));
}
lookup(&manager, &id)?.answer_question(&body.question_id, &body.answers);
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)
}
async fn stop(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<StatusCode, ApiError> {
manager.stop_session(&id).map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
async fn start(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<StatusCode, ApiError> {
manager.start_session(&id).map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Clone)]
pub struct UsageState {
monitor: Arc<crate::usage::UsageMonitor>,
manager: Arc<SessionManager>,
}
/// 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>,
manager: Arc<SessionManager>,
) -> Router {
Router::new()
.route("/usage", get(usage))
.with_state(UsageState { monitor, manager })
}
async fn usage(
State(state): State<UsageState>,
) -> Result<axum::Json<Vec<crate::usage::UsageSnapshot>>, ApiError> {
let setups = state.manager.setups();
let snapshots = tokio::task::spawn_blocking(move || state.monitor.snapshots(&setups))
.await
.context("usage fetch panicked")?;
Ok(axum::Json(snapshots))
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct TitleRequest {
title: String,
}
async fn rename(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<TitleRequest>,
) -> Result<StatusCode, ApiError> {
manager
.rename_session(&id, &body.title)
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct CwdRequest {
cwd: PathBuf,
}
/// The directory is checked here rather than in the manager because checking
/// it is an ssh round trip on a remote setup, and the manager is not async.
///
/// Checked rather than trusted, and refused rather than corrected: a mistyped
/// path that was accepted would leave a session recorded somewhere its process
/// cannot start, and the failure would arrive later with nothing pointing at
/// the typo. The spawn path corrects instead because it is resuming a
/// directory the *machine* recorded, which can be gone through nobody's fault.
async fn set_cwd(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<CwdRequest>,
) -> Result<StatusCode, ApiError> {
let session = manager
.sessions()
.into_iter()
.find(|session| session.id == id)
.ok_or_else(|| ApiError::NotFound(format!("no session {id}")))?;
// Absolute, because the alternative is relative to whatever the CLI is
// launched from, which is not something the person typing it can see. The
// same question the explorer asks of every path, asked in one place.
let cwd = crate::files::check_path(&body.cwd.to_string_lossy()).map_err(bad_request)?;
let setup = setup_by_id(&manager, &session.setup)?;
let transport = crate::session::transport::Transport::for_setup(&setup);
if !crate::session::import::directory_exists(&transport, &cwd).await {
return Err(ApiError::BadRequest(format!(
"{} has no directory {cwd}",
setup.name
)));
}
let stored = if setup.ssh.is_none() {
crate::setups::shorten_home(&cwd)
} else {
cwd.clone()
};
manager
.set_session_cwd(&id, PathBuf::from(&stored))
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
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)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct PermissionModeRequest {
mode: String,
}
/// What new sessions start at. One field today; a struct rather than a bare
/// value because "the defaults" is the thing a phone asks for, and the next
/// one to move here -- the permission mode, which the spawn screen still
/// hardcodes -- must not need a second route.
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct Defaults {
#[serde(default, skip_serializing_if = "Option::is_none")]
effort: Option<String>,
}
async fn defaults(State(manager): State<Arc<SessionManager>>) -> axum::Json<Defaults> {
axum::Json(Defaults {
effort: manager.default_effort(),
})
}
async fn set_defaults(
State(manager): State<Arc<SessionManager>>,
axum::Json(body): axum::Json<Defaults>,
) -> Result<StatusCode, ApiError> {
manager
.set_default_effort(body.effort.as_deref())
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
struct EffortRequest {
#[serde(default)]
effort: Option<String>,
}
async fn set_effort(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<EffortRequest>,
) -> Result<StatusCode, ApiError> {
manager
.set_session_effort(&id, body.effort.as_deref())
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
async fn set_permission_mode(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<PermissionModeRequest>,
) -> Result<StatusCode, ApiError> {
manager
.set_session_permission_mode(&id, &body.mode)
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct NotifyRequest {
notify: bool,
}
async fn set_notify(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<NotifyRequest>,
) -> Result<StatusCode, ApiError> {
manager
.set_session_notify(&id, body.notify)
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct AutoResumeRequest {
auto_resume: bool,
#[serde(default)]
message: Option<String>,
}
/// One request for both, because they are one decision: switching it on
/// without saying what to send is the ordinary case, and changing the words
/// while it is off is how somebody sets it up before it is needed.
async fn set_auto_resume(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<AutoResumeRequest>,
) -> Result<StatusCode, ApiError> {
manager
.set_session_auto_resume(&id, body.auto_resume, body.message.as_deref())
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct CommandRequest {
text: String,
}
/// The two this server understands are turned into the operations it has, and
/// everything else is passed to the session verbatim, because a dialect's
/// vocabulary is its own and grows without this file.
async fn command(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
axum::Json(body): axum::Json<CommandRequest>,
) -> Result<StatusCode, ApiError> {
let text = body.text.trim();
let (name, rest) = match text.split_once(char::is_whitespace) {
Some((name, rest)) => (name, rest.trim()),
None => (text, ""),
};
// A rename still goes through `rename_session` rather than being a command
// like the rest, because the name is persisted and listed as well as
// forwarded, and that is one operation.
let command = match (name, rest) {
("/compact", _) => SessionCommand::Compact,
("/clear", _) => SessionCommand::Clear,
("/rename", "") => return Err(bad_request(anyhow::anyhow!("a session needs a name"))),
("/rename", title) => {
manager.rename_session(&id, title).map_err(bad_request)?;
return Ok(StatusCode::NO_CONTENT);
}
_ => SessionCommand::Raw(text.to_string()),
};
lookup(&manager, &id)?;
manager.run_command(&id, command).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)?;
manager
.run_command(&id, SessionCommand::Compact)
.map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
const ATTACHMENT_LIMIT: usize = 4 * 1024 * 1024 * 1024;
/// Written to disk as it arrives rather than collected first: a trace is bigger
/// than this process should hold. Under a `.part` name until it is whole, so a
/// tunnel that drops mid-upload leaves nothing a message could reference.
///
/// A file for a session on another machine is copied there too, because the
/// path the session is told has to exist where the session runs. The copy is
/// part of the upload: if it fails, the upload fails and says so, rather than
/// a message later naming a file that is not there.
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 mut 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("application/octet-stream")
.to_string();
let file_name = field.file_name().map(str::to_string);
let (name, path) = session
.new_attachment(&content_type, file_name.as_deref())
.map_err(bad_request)?;
let part = path.with_file_name(format!("{name}.part"));
let received: Result<(), ApiError> = async {
use tokio::io::AsyncWriteExt;
let mut file = tokio::fs::File::create(&part)
.await
.with_context(|| format!("create {}", part.display()))?;
while let Some(chunk) = field
.chunk()
.await
.map_err(|err| ApiError::BadRequest(format!("upload read failed: {err}")))?
{
file.write_all(&chunk)
.await
.with_context(|| format!("write {}", part.display()))?;
}
file.flush()
.await
.with_context(|| format!("finish {}", part.display()))?;
Ok(())
}
.await;
if let Err(err) = received {
let _ = tokio::fs::remove_file(&part).await;
return Err(err);
}
tokio::fs::rename(&part, &path)
.await
.with_context(|| format!("name {}", path.display()))
.map_err(ApiError::Internal)?;
if crate::media::media_type_for(&name).is_none()
&& let Some((ssh, cwd)) = manager.remote_of(&id)
{
match ship_attachment(&ssh, cwd.as_deref(), &path, &name).await {
Ok(remote) => {
tokio::fs::write(remote_marker(&path), remote)
.await
.with_context(|| format!("record where {name} went"))
.map_err(ApiError::Internal)?;
}
Err(err) => {
let _ = tokio::fs::remove_file(&path).await;
return Err(ApiError::BadRequest(format!(
"{name} reached the server but couldn't be copied to {}: {err:#}",
ssh.address
)));
}
}
}
Ok(axum::Json(serde_json::json!({ "id": name })))
}
async fn ship_attachment(
ssh: &crate::config::SshConfig,
cwd: Option<&Path>,
local: &Path,
name: &str,
) -> anyhow::Result<String> {
let dir = ssh.attachments_dir.as_deref().or(cwd);
let mut script = String::new();
if let Some(dir) = dir {
let dir = crate::ssh::quote_path(&dir.to_string_lossy());
script.push_str(&format!("mkdir -p {dir} && cd {dir} && "));
}
script.push_str(&format!("cat > {} && pwd -P", crate::ssh::quote(name)));
let source = std::fs::File::open(local).with_context(|| format!("open {}", local.display()))?;
let transport = crate::session::transport::Transport::Ssh {
name: ssh.address.clone(),
ssh: ssh.clone(),
};
let launch = crate::session::transport::Launch::new("sh", vec!["-c".to_string(), script], None);
let stdout = transport
.capture_with_input(&launch, crate::session::transport::Input::File(source))
.await?
.ok()?;
let dir = String::from_utf8_lossy(&stdout).trim().to_string();
if dir.is_empty() {
anyhow::bail!("the remote shell did not say where it put the file");
}
Ok(format!("{dir}/{name}"))
}
fn remote_marker(local: &Path) -> std::path::PathBuf {
let name = local.file_name().unwrap_or_default().to_string_lossy();
local.with_file_name(format!("{name}.remote"))
}
async fn serve_file(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, name)): UrlPath<(String, String)>,
) -> Result<Response, ApiError> {
if !name
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '.' || c == '-' || 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}"
)));
};
let bytes = std::fs::read(path)
.with_context(|| format!("read {}", path.display()))
.map_err(ApiError::Internal)?;
let content_type = crate::media::media_type_for(&name).unwrap_or("application/octet-stream");
Ok(([(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
}
#[derive(Deserialize)]
struct EventsQuery {
#[serde(default)]
after: u64,
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct TranscriptQuery {
#[serde(default)]
before: Option<u64>,
#[serde(default = "default_window")]
limit: usize,
/// Join each reply's streamed deltas into one event, so a page counts rows
/// rather than tokens. The scroll-back pager asks for this; the
/// anchor-restore path does not, because it counts events to reach a known
/// seq. See `read_window`.
#[serde(default)]
coalesce: bool,
/// Return nothing at or below this seq; the page stops here instead of at
/// `limit`. The phone passes the end of what it already holds, so a page
/// never overlaps it. Exclusive, like the SSE route's `after`.
#[serde(default)]
after: Option<u64>,
}
fn default_window() -> usize {
80
}
async fn transcript(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
Query(query): Query<TranscriptQuery>,
) -> Result<axum::Json<Vec<crate::session::transcript::SeqEvent>>, ApiError> {
let session = lookup(&manager, &id)?;
transcript_page(session.transcript_path(), &id, query)
}
async fn subagent_transcript(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, sub)): UrlPath<(String, String)>,
Query(query): Query<TranscriptQuery>,
) -> Result<axum::Json<Vec<crate::session::transcript::SeqEvent>>, ApiError> {
let session = lookup(&manager, &id)?;
let subagent = lookup_subagent(&session, &sub)?;
transcript_page(
&subagent.transcript_path(),
&format!("{id}/subagents/{sub}"),
query,
)
}
fn transcript_page(
path: &Path,
label: &str,
query: TranscriptQuery,
) -> Result<axum::Json<Vec<crate::session::transcript::SeqEvent>>, ApiError> {
let events = crate::session::transcript::read_window(
path,
query.before,
query.after,
query.limit,
query.coalesce,
)
.map_err(bad_request)?;
tracing::debug!(
session = %label,
before = ?query.before,
after = ?query.after,
limit = query.limit,
got = events.len(),
oldest = ?events.first().map(|entry| entry.seq),
"transcript page"
);
Ok(axum::Json(events))
}
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 = cursor_of(&headers, query.after);
let live = session.subscribe();
Ok(sse_stream(
session.transcript_path().to_path_buf(),
cursor,
live,
))
}
async fn subagent_events(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, sub)): UrlPath<(String, String)>,
Query(query): Query<EventsQuery>,
headers: HeaderMap,
) -> Result<Sse<impl tokio_stream::Stream<Item = Result<SseEvent, Infallible>>>, ApiError> {
let session = lookup(&manager, &id)?;
let subagent = lookup_subagent(&session, &sub)?;
let cursor = cursor_of(&headers, query.after);
let live = subagent.subscribe();
Ok(sse_stream(subagent.transcript_path(), cursor, live))
}
fn cursor_of(headers: &HeaderMap, query_after: u64) -> u64 {
headers
.get("last-event-id")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse().ok())
.unwrap_or(query_after)
}
fn sse_stream(
transcript: PathBuf,
cursor: u64,
live: broadcast::Receiver<SeqEvent>,
) -> Sse<impl tokio_stream::Stream<Item = Result<SseEvent, Infallible>>> {
let (tx, stream) = mpsc::channel(64);
tokio::spawn(stream_session(transcript, cursor, live, tx));
Sse::new(ReceiverStream::new(stream).map(Ok)).keep_alive(KeepAlive::default())
}
async fn list_subagents(
State(manager): State<Arc<SessionManager>>,
UrlPath(id): UrlPath<String>,
) -> Result<axum::Json<Vec<SubagentInfo>>, ApiError> {
let session = lookup(&manager, &id)?;
let running = !matches!(
session.status(),
crate::session::driver::SessionStatus::Exited
| crate::session::driver::SessionStatus::Unknown
);
Ok(axum::Json(session.subagents().list(running)))
}
async fn notifications(
State(manager): State<Arc<SessionManager>>,
) -> Sse<impl tokio_stream::Stream<Item = Result<SseEvent, Infallible>>> {
let live = manager.subscribe_notifications();
let stream = BroadcastStream::new(live).filter_map(|item| {
let notification = item.ok()?;
Some(Ok(SseEvent::default().json_data(&notification).ok()?))
});
Sse::new(stream).keep_alive(KeepAlive::default())
}
async fn stream_session(
transcript: PathBuf,
mut last: u64,
mut live: broadcast::Receiver<SeqEvent>,
tx: mpsc::Sender<SseEvent>,
) {
if !send_backlog(&transcript, &mut last, &tx).await {
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(_)) => {
if !send_backlog(&transcript, &mut last, &tx).await {
return;
}
}
Err(broadcast::error::RecvError::Closed) => return,
}
}
}
/// Both ways into a backlog come through here -- the first replay and the
/// recovery from a lapped broadcast -- because either can be arbitrarily far
/// behind and owes the client the same answer.
async fn send_backlog(transcript: &Path, last: &mut u64, tx: &mpsc::Sender<SseEvent>) -> bool {
let cursor = *last;
let entries = match catch_up(transcript, *last, CATCH_UP_LIMIT) {
Ok(CatchUp::Continue(entries)) => {
// The pair at debug, because "was this subscriber reset, and how
// far behind was it" is a question about a phone that nothing else
// here can answer: the app sees a window arrive and cannot tell how
// far it had fallen.
tracing::debug!(cursor, sent = entries.len(), "stream backlog: continue");
entries
}
Ok(CatchUp::Restart(entries)) => {
tracing::debug!(cursor, sent = entries.len(), "stream backlog: reset");
if tx.send(SseEvent::default().event("reset")).await.is_err() {
return false;
}
entries
}
Err(err) => {
tracing::error!("transcript replay failed: {err:#}");
return false;
}
};
for entry in entries {
*last = entry.seq;
if send_event(tx, &entry).await.is_err() {
return false;
}
}
true
}
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
}
/// Separate router because its state is the model store, like `usage`'s.
pub fn models_router(store: Arc<crate::models::ModelStore>) -> Router {
Router::new()
.route("/models", get(list_models))
.route("/models/search", get(search_models))
.route("/models/files", get(repo_files))
.route("/models/download", post(start_download))
.route("/models/cancel", post(cancel_download))
.route("/models/delete", post(delete_model))
.with_state(store)
}
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct ModelsResponse {
local: Vec<crate::models::LocalModel>,
downloads: Vec<crate::models::DownloadStatus>,
}
async fn list_models(
State(store): State<Arc<crate::models::ModelStore>>,
) -> Result<axum::Json<ModelsResponse>, ApiError> {
let listing = tokio::task::spawn_blocking(move || ModelsResponse {
local: store.list(),
downloads: store.downloads(),
})
.await
.context("listing models panicked")?;
Ok(axum::Json(listing))
}
#[derive(Deserialize)]
struct SearchQuery {
q: String,
}
async fn search_models(
Query(query): Query<SearchQuery>,
) -> Result<axum::Json<Vec<crate::models::RemoteRepo>>, ApiError> {
let found = tokio::task::spawn_blocking(move || crate::models::search(&query.q))
.await
.context("model search panicked")?
.map_err(bad_request)?;
Ok(axum::Json(found))
}
#[derive(Deserialize)]
struct RepoQuery {
repo: String,
}
async fn repo_files(
State(store): State<Arc<crate::models::ModelStore>>,
Query(query): Query<RepoQuery>,
) -> Result<axum::Json<Vec<crate::models::RemoteFile>>, ApiError> {
let files = tokio::task::spawn_blocking(move || crate::models::files(&query.repo, &store))
.await
.context("listing repository files panicked")?
.map_err(bad_request)?;
Ok(axum::Json(files))
}
#[derive(Deserialize)]
struct DownloadRequest {
repo: String,
file: String,
}
async fn start_download(
State(store): State<Arc<crate::models::ModelStore>>,
axum::Json(body): axum::Json<DownloadRequest>,
) -> Result<axum::Json<crate::models::DownloadStatus>, ApiError> {
let status = store.start(&body.repo, &body.file).map_err(bad_request)?;
Ok(axum::Json(status))
}
#[derive(Deserialize)]
struct KeyRequest {
key: String,
}
async fn cancel_download(
State(store): State<Arc<crate::models::ModelStore>>,
axum::Json(body): axum::Json<KeyRequest>,
) -> Result<axum::Json<crate::models::DownloadStatus>, ApiError> {
let status = store.cancel(&body.key).map_err(bad_request)?;
Ok(axum::Json(status))
}
async fn delete_model(
State(store): State<Arc<crate::models::ModelStore>>,
axum::Json(body): axum::Json<KeyRequest>,
) -> Result<StatusCode, ApiError> {
store.delete(&body.key).map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}