use std::collections::BTreeMap;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::path::{Component, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use anyhow::{Context, Result as AnyResult};
use axum::body::{Body, Bytes, to_bytes};
use axum::extract::{DefaultBodyLimit, Path, Query, Request, State};
use axum::http::header::{
CACHE_CONTROL, CONTENT_SECURITY_POLICY as CONTENT_SECURITY_POLICY_HEADER, CONTENT_TYPE, COOKIE,
HeaderValue, LOCATION, REFERRER_POLICY, SET_COOKIE, X_CONTENT_TYPE_OPTIONS,
};
use axum::http::{HeaderMap, Response, StatusCode};
use axum::middleware::Next;
use axum::response::IntoResponse;
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::routing::{get, post, put};
use axum::{Json, Router};
use base64::Engine as _;
use hmac::{Hmac, KeyInit, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use tokio::sync::Semaphore;
use tokio::sync::{mpsc, watch};
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use mj_core::attachment::{AttachmentRef, AttachmentStore, MAX_IMAGE_BYTES, MAX_IMAGES};
use mj_core::config::{Config, TargetTemplate, project_history_host, validate_id};
use mj_core::elicitation::{ElicitationRequest, ElicitationResponse, MAX_ELICITATION_BYTES};
use mj_core::state::{
MoveOperation, MovePhase, MovePreparation, MoveSelection, MoveSessionRequest,
ProjectSourceIdentity, SessionResourceAllocation, SessionState, SessionTransitionKind,
State as AppState,
};
use crate::targets::AdditionalMount;
use crate::dictation::{
DictationError, DictationOperation, DictationRequest, DictationResponse, MAX_AUDIO_BYTES,
validate_wav,
};
use crate::image::optimize_image;
pub mod api;
pub use api::{
ApiFailure, ApiSession, PromptRequest, PromptResponse, SessionListResponse,
StartSessionRequest, StartSessionResponse, SubagentBackend, WaitOutcome, WaitRequest,
WaitResponse, api_token_path, load_or_create_api_token, map_stop_reason, resolve_wait,
};
pub use mj_client::web::{
BrowserDiffStat, BrowserTranscript, BrowserTranscriptEntry, WebListenerProcess,
WebViewerAccess, WebViewerRecovery,
};
pub use mj_core::state::ResumeQueueDisposition;
pub fn install_rustls_crypto_provider() {
let _ = rustls::crypto::ring::default_provider().install_default();
}
pub const COOKIE_NAME: &str = "hel_viewer_session";
const DEFAULT_SESSION_TTL: Duration = Duration::from_secs(30 * 24 * 60 * 60);
const EPHEMERAL_SESSION_TTL: Duration = Duration::from_secs(24 * 60 * 60);
const MAX_BODY_BYTES: usize = 128 * 1024;
const MAX_CODE_FAILURES: u32 = 5;
const CODE_LOCKOUT_BASE: Duration = Duration::from_secs(30);
const CODE_LOCKOUT_CAP: Duration = Duration::from_secs(60 * 60);
const MAX_TITLE_CHARS: usize = 120;
const MAX_PROMPT_CHARS: usize = 64 * 1024;
const MAX_DIRTY_ACKNOWLEDGEMENTS: usize = 32;
const MAX_DRAFT_BYTES: usize = 64 * 1024;
pub const MAX_HISTORY_MATCHES: usize = 40;
const MAX_PROMPT_BODY_BYTES: usize = 32 * 1024 * 1024;
const MAX_ATTACHMENT_UPLOAD_BYTES: usize = 64 * 1024 * 1024;
const MAX_CONCURRENT_DICTATIONS: usize = 2;
pub const MAX_PROMPT_IMAGES: usize = MAX_IMAGES;
const COOKIE_KEY_BYTES: usize = 32;
const COOKIE_KEY_FILE: &str = "phone-cookie-key";
pub const fn default_session_ttl() -> Duration {
DEFAULT_SESSION_TTL
}
pub fn cookie_key_path() -> PathBuf {
mj_core::config::data_dir().join(COOKIE_KEY_FILE)
}
pub fn load_or_create_cookie_key(path: &std::path::Path) -> AnyResult<Vec<u8>> {
match std::fs::read(path) {
Ok(key) if key.len() >= COOKIE_KEY_BYTES => return Ok(key),
Ok(key) => tracing::warn!(
path = %path.display(),
bytes = key.len(),
"phone cookie key is shorter than {COOKIE_KEY_BYTES} bytes; generating a new key signs every phone out"
),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
Err(error) => tracing::warn!(
path = %path.display(),
"could not read the phone cookie key ({error}); generating a new key signs every phone out"
),
}
let key = generate_cookie_key()?;
mj_core::config::atomic_write(path, &key)
.with_context(|| format!("persist Mjolnir phone cookie key {}", path.display()))?;
Ok(key.to_vec())
}
#[derive(Clone)]
pub struct ServerOptions {
pub bind: SocketAddr,
pub snapshot_rx: watch::Receiver<ViewerSnapshot>,
pub conversation_rx: watch::Receiver<BTreeMap<String, BrowserTranscript>>,
pub action_tx: mpsc::Sender<ControllerRequest>,
pub bundle_tx: mpsc::Sender<BundleRequest>,
pub receipt_tx: mpsc::Sender<ReadReceiptRequest>,
pub preflight_tx: mpsc::Sender<PreflightRequest>,
pub move_preparation_tx: mpsc::Sender<MovePreparationRequest>,
pub client_state_tx: mpsc::Sender<ClientStateRequest>,
pub dictation_tx: mpsc::Sender<DictationRequest>,
background_task_stop_tx: mpsc::Sender<BackgroundTaskStopRequest>,
pub shutdown: CancellationToken,
pub session_ttl: Duration,
pub secure_cookie: bool,
tls_config: Option<axum_server::tls_rustls::RustlsConfig>,
viewer_code: String,
login_token: String,
cookie_key: Vec<u8>,
api_token: String,
subagent: Option<Arc<dyn api::SubagentBackend>>,
}
pub struct ServerRequests {
pub action_tx: mpsc::Sender<ControllerRequest>,
pub bundle_tx: mpsc::Sender<BundleRequest>,
pub receipt_tx: mpsc::Sender<ReadReceiptRequest>,
pub preflight_tx: mpsc::Sender<PreflightRequest>,
pub move_preparation_tx: mpsc::Sender<MovePreparationRequest>,
pub client_state_tx: mpsc::Sender<ClientStateRequest>,
pub dictation_tx: mpsc::Sender<DictationRequest>,
}
impl ServerOptions {
pub fn new(
bind: SocketAddr,
snapshot_rx: watch::Receiver<ViewerSnapshot>,
conversation_rx: watch::Receiver<BTreeMap<String, BrowserTranscript>>,
requests: ServerRequests,
) -> AnyResult<Self> {
Ok(Self {
bind,
snapshot_rx,
conversation_rx,
action_tx: requests.action_tx,
bundle_tx: requests.bundle_tx,
receipt_tx: requests.receipt_tx,
preflight_tx: requests.preflight_tx,
move_preparation_tx: requests.move_preparation_tx,
client_state_tx: requests.client_state_tx,
dictation_tx: requests.dictation_tx,
background_task_stop_tx: mpsc::channel(1).0,
shutdown: CancellationToken::new(),
session_ttl: DEFAULT_SESSION_TTL,
secure_cookie: true,
tls_config: None,
viewer_code: generate_viewer_code()?,
login_token: generate_login_token()?,
cookie_key: generate_cookie_key()?.to_vec(),
api_token: String::new(),
subagent: None,
})
}
pub fn viewer_code(&self) -> &str {
&self.viewer_code
}
pub fn login_token(&self) -> &str {
&self.login_token
}
pub fn set_tls_config(&mut self, config: axum_server::tls_rustls::RustlsConfig) {
self.tls_config = Some(config);
self.secure_cookie = true;
}
pub fn set_cookie_key(&mut self, key: Vec<u8>) -> AnyResult<()> {
anyhow::ensure!(
key.len() >= COOKIE_KEY_BYTES,
"cookie signing key must be at least {COOKIE_KEY_BYTES} bytes"
);
self.cookie_key = key;
Ok(())
}
pub fn set_background_task_stop_tx(&mut self, tx: mpsc::Sender<BackgroundTaskStopRequest>) {
self.background_task_stop_tx = tx;
}
pub fn set_api_token(&mut self, token: String) {
self.api_token = token;
}
pub fn set_subagent_backend(&mut self, backend: Arc<dyn api::SubagentBackend>) {
self.subagent = Some(backend);
}
#[cfg(test)]
fn with_test_credentials(mut self, code: &str, key: &[u8]) -> Self {
self.viewer_code = code.to_string();
self.login_token = "test-login-token".into();
self.cookie_key = key.to_vec();
self.secure_cookie = false;
self.api_token = "test-api-token".into();
self
}
}
pub async fn run_server(options: ServerOptions) -> AnyResult<()> {
let listener = tokio::net::TcpListener::bind(options.bind)
.await
.with_context(|| format!("bind web viewer to {}", options.bind))?;
run_server_on_listener(options, listener).await
}
pub async fn run_server_on_listener(
options: ServerOptions,
listener: tokio::net::TcpListener,
) -> AnyResult<()> {
let mut options = options;
let bind = listener.local_addr().context("read web viewer address")?;
let shutdown = options.shutdown.clone();
let viewer_code = options.viewer_code.clone();
let tls_config = options.tls_config.take();
let app = router(options);
println!("Mjolnir viewer code: {viewer_code}");
let listener = listener.into_std().context("prepare web viewer listener")?;
let handle = axum_server::Handle::new();
let shutdown_handle = handle.clone();
let serve = async move {
if let Some(tls_config) = tls_config {
axum_server::from_tcp_rustls(listener, tls_config)
.handle(handle)
.serve(app.into_make_service())
.await
} else {
axum_server::from_tcp(listener)
.handle(handle)
.serve(app.into_make_service())
.await
}
};
tokio::pin!(serve);
tokio::select! {
result = &mut serve => result,
_ = shutdown.cancelled() => {
shutdown_handle.graceful_shutdown(Some(Duration::from_secs(2)));
serve.await
}
}
.with_context(|| format!("serve web viewer on {bind}"))
}
mod viewer_types;
pub use viewer_types::*;
mod actions;
pub use actions::*;
mod routes;
use routes::*;
mod handlers;
use handlers::*;
mod validation;
use validation::*;
mod errors;
use errors::*;
mod auth;
pub use auth::*;
mod assets;
use assets::*;
mod config_view;
pub use config_view::*;
#[cfg(test)]
mod tests;