use std::sync::Arc;
use axum::{
extract::{
ws::{Message, WebSocket, WebSocketUpgrade},
Path, State,
},
http::{header, StatusCode},
response::{IntoResponse, Json, Response},
};
use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;
use crate::state::SessionError;
use crate::ServerState;
fn with_rolling_cookie(mut response: Response, cookie: Option<String>) -> Response {
if let Some(cookie_str) = cookie {
if let Ok(value) = cookie_str.parse() {
response.headers_mut().insert(header::SET_COOKIE, value);
}
}
response
}
#[derive(Debug, Serialize)]
pub struct SessionResponse {
pub workflow_state: String,
pub output_lines: Vec<String>,
pub todo_lines: Vec<String>,
pub mcp_servers: Vec<McpServerJson>,
pub context_tokens: usize,
pub input: String,
pub resume_info_present: bool,
pub session_name: Option<String>,
pub project_name: Option<String>,
pub session_id: Option<String>,
pub compaction: abk::resume::CompactionStatus,
}
#[derive(Debug, Serialize)]
pub struct McpServerJson {
pub name: String,
pub connected: bool,
pub tool_count: usize,
pub error: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct CommandRequest {
pub command: String,
pub session_id: Option<String>,
#[serde(default)]
pub identity: Option<String>,
#[serde(default)]
pub model: Option<String>,
#[serde(default)]
pub attachments: Vec<AttachmentInput>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct AttachmentInput {
pub mime: String,
pub data: String,
#[serde(default)]
pub filename: Option<String>,
}
pub const MAX_ATTACHMENTS: usize = 4;
pub const MAX_IMAGE_BYTES: usize = 6 * 1024 * 1024; pub const MAX_TEXT_BYTES: usize = 256 * 1024; pub const ATTACH_BODY_LIMIT: usize = 40 * 1024 * 1024;
pub const ALLOWED_IMAGE_MIMES: [&str; 4] = ["image/jpeg", "image/png", "image/gif", "image/webp"];
pub const ALLOWED_TEXT_MIMES: [&str; 3] = ["text/plain", "text/markdown", "text/x-markdown"];
#[derive(Debug, Default)]
pub struct ValidatedAttachments {
pub images: Vec<umf::chatml::ImageAttachment>,
pub text_files: Vec<trustee_core::session::TextFileAttachment>,
}
fn decode_attachment_base64(
idx: usize,
name: &str,
input: &AttachmentInput,
) -> Result<Vec<u8>, String> {
use base64::Engine as _;
let raw = input
.data
.trim()
.strip_prefix("data:")
.and_then(|rest| rest.split_once(";base64,"))
.map(|(_, b64)| b64)
.unwrap_or_else(|| input.data.trim());
let cleaned: String = raw.chars().filter(|c| !c.is_whitespace()).collect();
base64::engine::general_purpose::STANDARD
.decode(cleaned.as_bytes())
.map_err(|e| format!("Attachment {} ({}): invalid base64: {}", idx + 1, name, e))
}
pub fn validate_attachments(inputs: &[AttachmentInput]) -> Result<ValidatedAttachments, String> {
if inputs.len() > MAX_ATTACHMENTS {
return Err(format!(
"Too many attachments: {} (max {})",
inputs.len(),
MAX_ATTACHMENTS
));
}
let mut out = ValidatedAttachments::default();
for (idx, input) in inputs.iter().enumerate() {
let name = input.filename.as_deref().unwrap_or("unnamed");
let mime = input.mime.trim().to_ascii_lowercase();
let decoded = decode_attachment_base64(idx, name, input)?;
if ALLOWED_IMAGE_MIMES.contains(&mime.as_str()) {
if decoded.len() > MAX_IMAGE_BYTES {
return Err(format!(
"Attachment {} ({}): {} MiB decoded exceeds the {} MiB limit",
idx + 1,
name,
decoded.len() / (1024 * 1024),
MAX_IMAGE_BYTES / (1024 * 1024)
));
}
use base64::Engine as _;
let cleaned_b64 = base64::engine::general_purpose::STANDARD.encode(&decoded);
let display_name = if name.is_empty() {
format!(
"attachment-{}.{}",
idx + 1,
mime.trim_start_matches("image/")
)
} else {
name.to_string()
};
out.images.push(
umf::chatml::ImageAttachment::new(mime.clone(), cleaned_b64)
.with_filename(display_name),
);
} else if ALLOWED_TEXT_MIMES.contains(&mime.as_str()) {
let mime = if mime == "text/x-markdown" {
"text/markdown".to_string()
} else {
mime
};
if decoded.len() > MAX_TEXT_BYTES {
return Err(format!(
"Attachment {} ({}): {} KiB of text exceeds the {} KiB limit",
idx + 1,
name,
decoded.len() / 1024,
MAX_TEXT_BYTES / 1024
));
}
let content = String::from_utf8(decoded).map_err(|_| {
format!(
"Attachment {} ({}): not valid UTF-8 — text attachments must decode to UTF-8 (is this a binary file?)",
idx + 1,
name
)
})?;
let filename = if name.is_empty() {
format!(
"attachment-{}.{}",
idx + 1,
if mime == "text/markdown" { "md" } else { "txt" }
)
} else {
name.to_string()
};
out.text_files
.push(trustee_core::session::TextFileAttachment {
mime,
filename,
content,
});
} else {
return Err(format!(
"Attachment {} ({}): unsupported type '{}' — expected one of {} (images) or {} (text/markdown files)",
idx + 1,
name,
input.mime,
ALLOWED_IMAGE_MIMES.join(", "),
["text/plain", "text/markdown"].join(", ")
));
}
}
Ok(out)
}
#[derive(Debug, Serialize)]
pub struct CommandResponse {
pub accepted: bool,
}
#[derive(Debug, Serialize)]
pub struct HealthResponse {
pub status: String,
pub version: String,
}
#[derive(Debug, Serialize)]
pub struct SessionListResponse {
pub sessions: Vec<trustee_core::sessions::SessionSummary>,
}
#[derive(Debug, Serialize)]
pub struct SessionDetailResponse {
pub session: trustee_core::sessions::SessionSummary,
pub checkpoints: Vec<trustee_core::sessions::CheckpointSummary>,
}
#[derive(Debug, Serialize)]
pub struct ResumeResponse {
pub accepted: bool,
pub session_id: String,
pub checkpoint_id: String,
pub iteration: u32,
pub tombstoned_at: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct ResumeRequestBody {
pub checkpoint_id: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct SessionHistoryResponse {
pub session_id: String,
pub checkpoint_id: String,
pub task_description: String,
pub iteration: u32,
pub total_messages: usize,
pub messages: Vec<trustee_core::sessions::HistoryMessage>,
pub tombstoned_at: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct CreateSessionRequest {
#[serde(default)]
pub session_name: Option<String>,
#[serde(default)]
pub resume_from: Option<String>,
#[serde(default)]
pub identity: Option<String>,
#[serde(default)]
pub model: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct CreateSessionResponse {
pub session_id: String,
pub accepted: bool,
}
#[derive(Debug, Deserialize)]
pub struct SetNameRequest {
pub name: String,
}
#[derive(Debug, Deserialize)]
pub struct NewSessionRequest {
pub session_name: Option<String>,
#[serde(default)]
pub identity: Option<String>,
#[serde(default)]
pub model: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct SetNameResponse {
pub accepted: bool,
pub name: String,
}
#[derive(Debug, Serialize)]
pub struct NewSessionResponse {
pub accepted: bool,
}
#[derive(Debug, Serialize)]
pub struct ModelInfo {
pub model: String,
}
#[derive(Debug, Serialize)]
pub struct ProviderInfo {
pub name: Option<String>,
pub base_url: Option<String>,
pub default_model: Option<String>,
pub models: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct ModelsResponse {
pub provider: ProviderInfo,
pub providers: Vec<ProviderInfo>,
}
fn model_catalog(resolved_config_toml: &str) -> Vec<String> {
fn push(catalog: &mut Vec<String>, m: Option<String>) {
if let Some(m) = m {
let m = m.trim().to_string();
if !m.is_empty() && !catalog.contains(&m) {
catalog.push(m);
}
}
}
fn push_models(catalog: &mut Vec<String>, v: &toml::Value) {
if let Some(arr) = v.get("models").and_then(|m| m.as_array()) {
for m in arr {
push(catalog, m.as_str().map(|s| s.to_string()));
}
}
}
let mut catalog: Vec<String> = Vec::new();
let Ok(v) = resolved_config_toml.parse::<toml::Value>() else {
return catalog;
};
let Some(llm) = v.get("llm").and_then(|l| l.as_table()) else {
return catalog;
};
if let Some(p) = llm.get("provider") {
push(
&mut catalog,
p.get("model").and_then(|m| m.as_str()).map(|s| s.to_string()),
);
push_models(&mut catalog, p);
}
if let Some(ps) = llm.get("providers").and_then(|p| p.as_table()) {
for (key, pv) in ps {
push(&mut catalog, Some(key.clone()));
push(
&mut catalog,
pv.get("model").and_then(|m| m.as_str()).map(|s| s.to_string()),
);
push_models(&mut catalog, pv);
}
}
catalog
}
fn validate_model_override(resolved: Option<&str>, model: &str) -> Result<String, Vec<String>> {
let trimmed = model.trim();
if trimmed.is_empty() {
return Ok(String::new());
}
let Some(config) = resolved else {
return Ok(trimmed.to_string());
};
let catalog = model_catalog(config);
if catalog.is_empty() {
return Ok(trimmed.to_string());
}
if catalog.iter().any(|m| m == trimmed) {
return Ok(trimmed.to_string());
}
Err(catalog)
}
fn resolved_command_model(resolved: Option<&str>, request_model: Option<&str>) -> Option<String> {
if let Some(m) = request_model.map(str::trim).filter(|s| !s.is_empty()) {
return Some(m.to_string());
}
if let Some(config) = resolved {
if let Ok(v) = config.parse::<toml::Value>() {
if let Some(m) = v
.get("llm")
.and_then(|l| l.get("provider"))
.and_then(|p| p.get("model"))
.and_then(|m| m.as_str())
.map(str::trim)
.filter(|s| !s.is_empty())
{
return Some(m.to_string());
}
}
}
std::env::var("OPENAI_DEFAULT_MODEL")
.ok()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
}
fn loud_tombstoned_chain_error(
state: &ServerState,
user_key: &str,
session_id: &str,
status: StatusCode,
verb: &str,
) -> (StatusCode, String) {
match state.chain_tombstone_stamp(user_key, session_id) {
Some(stamp) => (
status,
format!(
"chain {session_id} was DESTROYED {stamp} — {verb} refused (destroyed chains never auto-resurrect); transcript preserved on disk: GET /sessions/{session_id}/history reads it, POST /sessions/{session_id}/resume restores it explicitly"
),
),
None => (StatusCode::NOT_FOUND, "Session not found".to_string()),
}
}
fn session_to_response(session: &trustee_core::session::Session) -> SessionResponse {
let workflow_state = match session.workflow_state {
trustee_core::types::WorkflowState::Idle => "Idle",
trustee_core::types::WorkflowState::Running => "Running",
trustee_core::types::WorkflowState::Cancelling => "Cancelling",
};
let mcp_servers = session
.mcp_servers
.iter()
.map(|s| McpServerJson {
name: s.name.clone(),
connected: s.status == trustee_core::types::McpServerStatus::Connected,
tool_count: s.tool_count,
error: s.error.clone(),
})
.collect();
SessionResponse {
workflow_state: workflow_state.to_string(),
output_lines: session.output_lines.clone(),
todo_lines: session.todo_lines.clone(),
mcp_servers,
context_tokens: session.current_context_tokens,
input: session.input.clone(),
resume_info_present: session.resume_info.is_some(),
session_name: session.session_name.clone(),
project_name: session.project_name.clone(),
session_id: session.session_id.clone(),
compaction: session
.compaction_status
.lock()
.map(|s| s.clone())
.unwrap_or_default(),
}
}
pub async fn health() -> Json<HealthResponse> {
Json(HealthResponse {
status: "ok".to_string(),
version: crate::bin_version().to_string(),
})
}
pub async fn list_models(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::LIST_MODELS)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let empty = || ModelsResponse {
provider: ProviderInfo {
name: None,
base_url: None,
default_model: None,
models: vec![],
},
providers: vec![],
};
let Some(config_toml) = state.resolve_user_config(&user_key) else {
return Ok(with_rolling_cookie(Json(empty()).into_response(), cookie));
};
let Ok(table) = config_toml.parse::<toml::Value>() else {
return Ok(with_rolling_cookie(Json(empty()).into_response(), cookie));
};
let get_str = |v: &toml::Value, k: &str| -> Option<String> {
v.get(k).and_then(|m| m.as_str()).map(|s| s.to_string())
};
let get_models = |v: &toml::Value| -> Vec<String> {
v.get("models")
.and_then(|m| m.as_array())
.map(|arr| {
arr.iter()
.filter_map(|m| m.as_str().map(|s| s.to_string()))
.collect()
})
.unwrap_or_default()
};
let to_info = |v: &toml::Value| -> ProviderInfo {
let default_model = get_str(v, "model");
let mut models = get_models(v);
if models.is_empty() {
if let Some(ref m) = default_model {
models = vec![m.clone()];
}
}
ProviderInfo {
name: get_str(v, "name"),
base_url: get_str(v, "base_url"),
default_model,
models,
}
};
let provider = table
.get("llm")
.and_then(|l| l.get("provider"))
.map(|p| to_info(p))
.unwrap_or(ProviderInfo {
name: None,
base_url: None,
default_model: None,
models: vec![],
});
let providers = table
.get("llm")
.and_then(|l| l.get("providers"))
.and_then(|p| p.as_table())
.map(|ps| ps.iter().map(|(_, v)| to_info(v)).collect::<Vec<_>>())
.unwrap_or_default();
let resp = Json(ModelsResponse {
provider,
providers,
});
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
async fn handle_ws(
socket: WebSocket,
session_arc: std::sync::Arc<tokio::sync::Mutex<trustee_core::session::Session>>,
ws_tx: broadcast::Sender<String>,
route_session_id: String,
) {
use futures::{SinkExt, StreamExt};
let (mut sender, mut receiver) = socket.split();
let mut ws_rx = ws_tx.subscribe();
{
let session = session_arc.lock().await;
let mut snapshot = session_to_response(&session);
snapshot.session_id = Some(route_session_id);
if let Ok(json) = serde_json::to_string(&snapshot) {
let _ = sender.send(Message::Text(json.into())).await;
}
}
loop {
tokio::select! {
msg = ws_rx.recv() => {
match msg {
Ok(text) => {
if sender.send(Message::Text(text.into())).await.is_err() {
break;
}
}
Err(broadcast::error::RecvError::Lagged(n)) => {
let warn = serde_json::json!({"type":"Warning","message":format!("Lagged {} messages", n)});
let _ = sender.send(Message::Text(warn.to_string().into())).await;
}
Err(broadcast::error::RecvError::Closed) => break,
}
}
msg = receiver.next() => {
match msg {
Some(Ok(Message::Close(_))) | None => break,
_ => {}
}
}
}
}
}
pub async fn list_sessions(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::LIST_SESSIONS)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (config_toml, home_dir) = {
let (config, home) = state.get_user_config_and_home(&user_key);
match config {
Some(c) => (c, home),
None => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
"Configuration not loaded".to_string(),
))
}
}
};
let sessions = trustee_core::sessions::list_all_sessions(&config_toml, home_dir.as_deref())
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let resp = Json(SessionListResponse { sessions });
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn list_live_sessions(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::LIST_SESSIONS)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let items = state.list_sessions(&user_key).await;
let resp = Json(serde_json::json!({ "sessions": items }));
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn get_session_detail(
State(state): State<ServerState>,
Path(session_id): Path<String>,
headers: axum::http::HeaderMap,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::VIEW_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (config_toml, home_dir) = {
let (config, home) = state.get_user_config_and_home(&user_key);
match config {
Some(c) => (c, home),
None => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
"Configuration not loaded".to_string(),
))
}
}
};
let detail =
trustee_core::sessions::get_session_detail(&config_toml, &session_id, home_dir.as_deref())
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
match detail {
Some((session, checkpoints)) => {
let resp = Json(SessionDetailResponse {
session,
checkpoints,
});
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
None => Err((StatusCode::NOT_FOUND, "Session not found".to_string())),
}
}
pub async fn resume_session(
State(state): State<ServerState>,
Path(checkpoint_session_id): Path<String>,
headers: axum::http::HeaderMap,
_body: Option<Json<ResumeRequestBody>>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::RESUME_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (config_toml, home_dir) = {
let (config, home) = state.get_user_config_and_home(&user_key);
match config {
Some(c) => (c, home),
None => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
"Configuration not loaded".to_string(),
))
}
}
};
let resume_info = trustee_core::sessions::create_resume_info(
&config_toml,
&checkpoint_session_id,
home_dir.as_deref(),
)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let resume_info = match resume_info {
Some(info) => info,
None => {
return Err((
StatusCode::NOT_FOUND,
"Session or checkpoint not found".to_string(),
))
}
};
let checkpoint_id = resume_info.checkpoint_id.clone();
let iteration = resume_info.iteration;
let session_name = match trustee_core::sessions::get_session_detail(
&config_toml,
&checkpoint_session_id,
home_dir.as_deref(),
)
.await
{
Ok(Some((summary, _))) => summary.description.filter(|d| !d.trim().is_empty()),
_ => None,
}
.map(|d| format!("Resumed: {}", d))
.unwrap_or_else(|| format!("Resumed: {}", checkpoint_session_id));
let new_sid = state
.create_session(&user_key, Some(session_name), None)
.await
.map_err(|e| match e {
crate::state::SessionError::MaxSessionsReached(n) => (
StatusCode::TOO_MANY_REQUESTS,
format!("Max {} sessions reached", n),
),
crate::state::SessionError::NotFound(_) => {
(StatusCode::NOT_FOUND, "Session not found".to_string())
}
crate::state::SessionError::NotIdle(_) => {
(StatusCode::CONFLICT, "Session is busy".to_string())
}
})?;
trustee_core::sessions::write_resume_pointer(
&config_toml,
&checkpoint_session_id,
&new_sid,
home_dir.as_deref(),
)
.await
.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to record resume handle: {}", e),
)
})?;
{
let (session_arc, _ws_tx) = state.get_session(&user_key, &new_sid).await.ok_or((
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to get new session".to_string(),
))?;
let mut session = session_arc.lock().await;
session.resume_info = Some(resume_info);
session.session_id = Some(checkpoint_session_id.clone());
if let Some(history) = trustee_core::sessions::load_session_history(
&checkpoint_session_id,
home_dir.as_deref(),
)
.await
.ok()
.flatten()
{
session.output_lines.clear();
for msg in &history.messages {
if msg.role == "user" {
session.output_lines.push(format!("> {}", msg.content));
} else {
if let Some(ref reasoning) = msg.reasoning {
if !reasoning.trim().is_empty() {
session.output_lines.push(format!("\x01{}", reasoning));
}
}
if !msg.content.is_empty() {
session.output_lines.push(msg.content.clone());
}
}
}
}
}
let (_, ws_tx) = state.get_session(&user_key, &new_sid).await.ok_or((
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to get session WS".to_string(),
))?;
let msg = serde_json::json!({
"type": "SessionResumed",
"session_id": checkpoint_session_id,
"checkpoint_id": checkpoint_id,
});
let _ = ws_tx.send(msg.to_string());
let resp = Json(ResumeResponse {
accepted: true,
session_id: new_sid.clone(),
checkpoint_id,
iteration,
tombstoned_at: state.chain_tombstone_stamp(&user_key, &checkpoint_session_id),
});
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn fork_session(
State(state): State<ServerState>,
Path(source_session_id): Path<String>,
headers: axum::http::HeaderMap,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::CREATE_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (config_toml, home_dir) = {
let (config, home) = state.get_user_config_and_home(&user_key);
match config {
Some(c) => (c, home),
None => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
"Configuration not loaded".to_string(),
))
}
}
};
let source_info = trustee_core::sessions::create_resume_info(
&config_toml,
&source_session_id,
home_dir.as_deref(),
)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let source_info = source_info.ok_or((
StatusCode::NOT_FOUND,
"Session or checkpoint not found".to_string(),
))?;
let checkpoint_id = source_info.checkpoint_id.clone();
let iteration = source_info.iteration;
let session_name = match trustee_core::sessions::get_session_detail(
&config_toml,
&source_session_id,
home_dir.as_deref(),
)
.await
{
Ok(Some((summary, _))) => summary.description.filter(|d| !d.trim().is_empty()),
_ => None,
}
.map(|d| format!("Forked: {}", d))
.unwrap_or_else(|| format!("Forked: {}", source_session_id));
let new_sid = state
.create_session(&user_key, Some(session_name), None)
.await
.map_err(|e| match e {
crate::state::SessionError::MaxSessionsReached(n) => (
StatusCode::TOO_MANY_REQUESTS,
format!("Max {} sessions reached", n),
),
crate::state::SessionError::NotFound(_) => {
(StatusCode::NOT_FOUND, "Session not found".to_string())
}
crate::state::SessionError::NotIdle(_) => {
(StatusCode::CONFLICT, "Session is busy".to_string())
}
})?;
let fork_info = trustee_core::sessions::fork_session_chain(
&config_toml,
&source_session_id,
&new_sid,
home_dir.as_deref(),
)
.await
.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to fork session chain: {}", e),
)
})?
.ok_or((
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to fork session chain: source checkpoints vanished".to_string(),
))?;
{
let (session_arc, _ws_tx) = state.get_session(&user_key, &new_sid).await.ok_or((
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to get new session".to_string(),
))?;
let mut session = session_arc.lock().await;
session.resume_info = Some(fork_info);
}
let (_, ws_tx) = state.get_session(&user_key, &new_sid).await.ok_or((
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to get session WS".to_string(),
))?;
let msg = serde_json::json!({
"type": "SessionForked",
"session_id": new_sid,
"forked_from": source_session_id,
"checkpoint_id": checkpoint_id,
});
let _ = ws_tx.send(msg.to_string());
let resp = Json(ResumeResponse {
accepted: true,
session_id: new_sid,
checkpoint_id,
iteration,
tombstoned_at: None,
});
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn get_session_history(
State(state): State<ServerState>,
Path(session_id): Path<String>,
headers: axum::http::HeaderMap,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::VIEW_HISTORY)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let home_dir = state.get_user_home_dir(&user_key);
let history = trustee_core::sessions::load_session_history(&session_id, home_dir.as_deref())
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
let tombstoned_at = state.chain_tombstone_stamp(&user_key, &session_id);
match history {
Some(h) => {
let resp = Json(SessionHistoryResponse {
session_id: h.session_id,
checkpoint_id: h.checkpoint_id,
task_description: h.task_description,
iteration: h.iteration,
total_messages: h.total_messages,
messages: h.messages,
tombstoned_at,
});
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
None => Err((StatusCode::NOT_FOUND, "Session not found".to_string())),
}
}
pub async fn create_session(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Json(req): Json<CreateSessionRequest>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::CREATE_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let session_id = state
.create_session(&user_key, req.session_name.clone(), req.identity.clone())
.await
.map_err(|e| match e {
SessionError::MaxSessionsReached(n) => (
StatusCode::TOO_MANY_REQUESTS,
format!("Max {} sessions per user", n),
),
_ => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
})?;
if let Some(ref model) = req.model {
let session_config = match state.get_session(&user_key, &session_id).await {
Some((session_arc, _)) => {
let session = session_arc.lock().await;
session.config_toml.clone()
}
None => None,
};
let validated = validate_model_override(session_config.as_deref(), model).map_err(|catalog| {
(
StatusCode::UNPROCESSABLE_ENTITY,
format!(
"Unknown model '{}'. Valid models: {}. Source values from GET /api/v1/models or omit the field.",
model,
catalog.join(", ")
),
)
})?;
if !validated.is_empty() {
if let Some((session_arc, _)) = state.get_session(&user_key, &session_id).await {
let mut session = session_arc.lock().await;
session.model = Some(validated);
}
}
}
if let Some(resume_from) = req.resume_from {
if let Some((session_arc, _)) = state.get_session(&user_key, &session_id).await {
let mut session = session_arc.lock().await;
if let Some(ref config_toml) = session.config_toml.clone() {
let home_dir = session.home_dir.as_deref();
if let Ok(Some(info)) =
trustee_core::sessions::create_resume_info(config_toml, &resume_from, home_dir)
.await
{
session.resume_info = Some(info);
}
}
}
}
let resp = Json(CreateSessionResponse {
session_id,
accepted: true,
});
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn get_live_session(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Path(session_id): Path<String>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::VIEW_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (_live_key, session_arc, _ws_tx) = state
.get_or_restore_session_by_any_id(&user_key, &session_id)
.await
.ok_or_else(|| {
loud_tombstoned_chain_error(
&state,
&user_key,
&session_id,
StatusCode::NOT_FOUND,
"no live session",
)
})?;
let session = session_arc.lock().await;
let resp = Json(session_to_response(&session));
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn post_command_session(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Path(session_id): Path<String>,
Json(req): Json<CommandRequest>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::COMMAND_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (_live_key, session_arc, ws_tx) = state
.get_or_restore_session_by_any_id(&user_key, &session_id)
.await
.ok_or_else(|| {
loud_tombstoned_chain_error(
&state,
&user_key,
&session_id,
StatusCode::CONFLICT,
"command",
)
})?;
let token_store = {
let user_sessions = state.sessions.get(&user_key).unwrap();
user_sessions.token_store.clone()
};
execute_command_inner(
&state,
&headers,
&session_arc,
&ws_tx,
&token_store,
&user_key,
req,
)
.await?;
let resp = Json(CommandResponse { accepted: true });
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn post_cancel_session(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Path(session_id): Path<String>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::CANCEL_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (_live_key, session_arc, ws_tx) = state
.get_or_restore_session_by_any_id(&user_key, &session_id)
.await
.ok_or_else(|| {
loud_tombstoned_chain_error(
&state,
&user_key,
&session_id,
StatusCode::CONFLICT,
"cancel",
)
})?;
let cancelled = {
let session = session_arc.lock().await;
let c = session.workflow_state == trustee_core::types::WorkflowState::Running;
if c {
session.cancel_token.cancel();
}
c
};
if cancelled {
let state_msg = serde_json::json!({"type": "StateChanged", "state": "Cancelling"});
let _ = ws_tx.send(state_msg.to_string());
}
let resp = Json(CommandResponse { accepted: true });
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn post_handoff_session(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Path(session_id): Path<String>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::HANDOFF_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (_live_key, session_arc, _ws_tx) = state
.get_or_restore_session_by_any_id(&user_key, &session_id)
.await
.ok_or_else(|| {
loud_tombstoned_chain_error(
&state,
&user_key,
&session_id,
StatusCode::CONFLICT,
"handoff",
)
})?;
let mut session = session_arc.lock().await;
session.request_handoff(String::new());
let resp = Json(CommandResponse { accepted: true });
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn set_session_name_session(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Path(session_id): Path<String>,
Json(req): Json<SetNameRequest>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::UPDATE_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (_live_key, session_arc, _ws_tx) = state
.get_or_restore_session_by_any_id(&user_key, &session_id)
.await
.ok_or((StatusCode::NOT_FOUND, "Session not found".to_string()))?;
{
let mut session = session_arc.lock().await;
session.session_name = Some(req.name.clone());
}
let resp = Json(SetNameResponse {
accepted: true,
name: req.name,
});
Ok(with_rolling_cookie(resp.into_response(), cookie))
}
pub async fn destroy_session(
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Path(session_id): Path<String>,
) -> Result<Response, (StatusCode, String)> {
let (cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::DELETE_SESSION)
.await
.map_err(|s| (s, "Unauthorized".to_string()))?;
let (live_key, _session_arc, _ws_tx) = state
.get_or_restore_session_by_any_id(&user_key, &session_id)
.await
.ok_or((StatusCode::NOT_FOUND, "Session not found".to_string()))?;
state
.destroy_session(&user_key, &live_key)
.await
.map_err(|e| match e {
SessionError::NotFound(_) => (StatusCode::NOT_FOUND, e.to_string()),
SessionError::NotIdle(_) => (StatusCode::CONFLICT, e.to_string()),
_ => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
})?;
Ok(with_rolling_cookie(
StatusCode::NO_CONTENT.into_response(),
cookie,
))
}
pub async fn ws_session_handler(
ws: WebSocketUpgrade,
State(state): State<ServerState>,
headers: axum::http::HeaderMap,
Path(session_id): Path<String>,
) -> Result<Response, StatusCode> {
let (_cookie, user_key) =
crate::auth::check_auth(&state.auth, &headers, crate::auth::actions::VIEW_SESSION).await?;
let (session_arc, ws_tx) = state
.get_or_restore_session_by_any_id(&user_key, &session_id)
.await
.map(|(_, s, ws)| (s, ws))
.ok_or(StatusCode::NOT_FOUND)?;
Ok(ws.on_upgrade(move |socket| handle_ws(socket, session_arc, ws_tx, session_id)))
}
pub async fn serve_index() -> Response {
match trustee_web::Asset::get("index.html") {
Some(content) => {
let html = String::from_utf8_lossy(&content.data)
.replace("__TRUSTEE_VERSION__", crate::bin_version());
(
StatusCode::OK,
[(header::CONTENT_TYPE, "text/html; charset=utf-8")],
html.into_bytes(),
)
.into_response()
}
None => (
StatusCode::NOT_FOUND,
[(header::CONTENT_TYPE, "text/plain")],
"Not found".to_string().into_bytes(),
)
.into_response(),
}
}
pub async fn serve_static(Path(file): Path<String>) -> Response {
match trustee_web::Asset::get(&file) {
Some(content) => {
let mime = mime_guess::from_path(&file).first_or_octet_stream();
(
StatusCode::OK,
[(header::CONTENT_TYPE, mime.as_ref())],
content.data.to_vec(),
)
.into_response()
}
None => (
StatusCode::NOT_FOUND,
[(header::CONTENT_TYPE, "text/plain")],
"Not found".to_string().into_bytes(),
)
.into_response(),
}
}
#[allow(clippy::too_many_arguments)]
async fn execute_command_inner(
state: &ServerState,
headers: &axum::http::HeaderMap,
session_arc: &std::sync::Arc<tokio::sync::Mutex<trustee_core::session::Session>>,
ws_tx: &broadcast::Sender<String>,
token_store: &Arc<pep::MemoryTokenStore>,
user_key: &str,
req: CommandRequest,
) -> Result<(), (StatusCode, String)> {
let attachments =
validate_attachments(&req.attachments).map_err(|e| (StatusCode::BAD_REQUEST, e))?;
let agent_name = {
let session = session_arc.lock().await;
session.agent_name.clone()
};
inject_session_token(&state.auth, headers, &agent_name, token_store).await;
{
let mut session = session_arc.lock().await;
if session.workflow_state != trustee_core::types::WorkflowState::Idle {
return Err((
StatusCode::CONFLICT,
"Workflow is running or cancelling".to_string(),
));
}
if let Some(ref client_session_id) = req.session_id {
if session.session_id.is_none() {
session.session_id = Some(client_session_id.clone());
}
let id_matches_current_chain = session
.session_id
.as_deref()
.is_some_and(|sid| sid == client_session_id.as_str());
let rotation_pending_first_checkpoint =
session.handoff_count > 0 && session.resume_info.is_none();
let stale_during_rotation_window =
rotation_pending_first_checkpoint && !id_matches_current_chain;
if session.resume_info.is_none() && !stale_during_rotation_window {
if let Some(ref config_toml) = session.config_toml.clone() {
let home_dir = session.home_dir.as_deref();
if let Ok(Some(info)) = trustee_core::sessions::create_resume_info(
config_toml,
client_session_id,
home_dir,
)
.await
{
session.resume_info = Some(info);
}
}
}
}
if session.identity.is_none() {
session.identity = req.identity;
}
let validated = match req.model.as_deref() {
Some(m) => Some(
validate_model_override(session.config_toml.as_deref(), m).map_err(|catalog| {
(
StatusCode::UNPROCESSABLE_ENTITY,
format!(
"Unknown model '{}'. Valid models: {}. Source values from GET /api/v1/models or omit the field.",
m,
catalog.join(", ")
),
)
})?,
),
None => None,
};
if resolved_command_model(session.config_toml.as_deref(), req.model.as_deref()).is_none() {
return Err((
StatusCode::UNPROCESSABLE_ENTITY,
"no model configured: set [llm.provider].model or OPENAI_DEFAULT_MODEL (or send a valid `model` override)".to_string(),
));
}
session.model = validated.filter(|s| !s.is_empty());
let permit = state
.workflow_semaphore
.clone()
.acquire_owned()
.await
.map_err(|_| {
(
StatusCode::SERVICE_UNAVAILABLE,
"Server overloaded".to_string(),
)
})?;
let mcp_loader = state
.get_or_build_mcp_loader(user_key, token_store)
.await
.map_err(|e| {
(
StatusCode::BAD_GATEWAY,
format!("MCP loader unavailable: {}", e),
)
})?;
session.mcp_loader = mcp_loader;
session.token_store = Some(token_store.clone());
session.workflow_permit = Some(permit);
let command = req.command.clone();
session.input = req.command;
session.input_images = attachments.images;
session.input_text_files = attachments.text_files;
session.execute_command();
let echo = serde_json::json!({
"type": "OutputLine",
"line": format!("> {}", command),
});
let _ = ws_tx.send(echo.to_string());
}
let state_msg = serde_json::json!({"type": "StateChanged", "state": "Running"});
let _ = ws_tx.send(state_msg.to_string());
Ok(())
}
const WEB_SESSION_CRED_NAME: &str = "__web_session";
async fn inject_session_token(
auth: &Option<Arc<crate::auth::AuthState>>,
headers: &axum::http::HeaderMap,
_agent_name: &str,
token_store: &pep::MemoryTokenStore,
) {
use pep::{StoredToken, TokenStore};
let Some(auth_state) = auth.as_ref() else {
return;
};
let access_token = match resolve_access_token_for_mcp(auth_state, headers).await {
Ok(token) => token,
Err(e) => {
tracing::debug!("Skipping MCP session token injection: {}", e);
return;
}
};
let expires_at = jwt_expiry(&access_token).unwrap_or_else(|| {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
compute_rfc3339(now + 900)
});
let stored = StoredToken::new(&access_token, None, "Bearer", &expires_at, None);
if let Err(e) = token_store.save(WEB_SESSION_CRED_NAME, &stored) {
tracing::warn!("Failed to write session token to MemoryTokenStore: {}", e);
} else {
tracing::debug!(
"Injected session token for web-session MCP credentials (expires {})",
expires_at
);
}
}
async fn resolve_access_token_for_mcp(
auth: &crate::auth::AuthState,
headers: &axum::http::HeaderMap,
) -> Result<String, String> {
if let Some(token) = headers
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.map(|s| s.to_string())
{
if token.starts_with("dev:") {
return Err("dev tokens not supported for MCP".to_string());
}
return Ok(token);
}
let session_id = headers
.get(header::COOKIE)
.and_then(|v| v.to_str().ok())
.and_then(|cookies| {
cookies.split(';').map(|c| c.trim()).find_map(|c| {
c.strip_prefix(&format!("{}=", auth.config.cookie_name))
.map(|s| s.to_string())
})
})
.ok_or("no session cookie")?;
if session_id.starts_with("dev:") {
return Err("dev tokens not supported for MCP".to_string());
}
auth.session_manager
.get_token(&session_id)
.await
.map_err(|e| format!("session lookup: {e}"))
}
fn jwt_expiry(token: &str) -> Option<String> {
let parts: Vec<&str> = token.split('.').collect();
if parts.len() < 2 {
return None;
}
use base64::Engine;
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(parts[1])
.or_else(|_| base64::engine::general_purpose::STANDARD_NO_PAD.decode(parts[1]))
.ok()?;
let json: serde_json::Value = serde_json::from_slice(&payload).ok()?;
let exp = json.get("exp")?.as_u64()?;
Some(compute_rfc3339(exp))
}
fn compute_rfc3339(epoch_secs: u64) -> String {
let days = epoch_secs / 86400;
let rem = epoch_secs % 86400;
let h = rem / 3600;
let m = (rem % 3600) / 60;
let s = rem % 60;
let z = days as i64 + 719468;
let era = if z >= 0 { z } else { z - 146096 } / 146097;
let doe = (z - era * 146097) as u64;
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
let y = yoe as i64 + era * 400;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
let mp = (5 * doy + 2) / 153;
let d = doy - (153 * mp + 2) / 5 + 1;
let mon = if mp < 10 { mp + 3 } else { mp - 9 };
let yr = if mon <= 2 { y + 1 } else { y };
format!("{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z", yr, mon, d, h, m, s)
}
#[cfg(test)]
mod serve_index_tests {
use super::*;
#[tokio::test]
async fn index_serves_real_version_no_placeholder() {
let resp = serve_index().await;
assert_eq!(resp.status(), StatusCode::OK);
let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX)
.await
.expect("body");
let html = String::from_utf8(bytes.to_vec()).expect("utf8");
assert!(
!html.contains("__TRUSTEE_VERSION__"),
"placeholder must be replaced at serve time"
);
assert!(
html.contains(&format!("trustee v{}", crate::bin_version())),
"burger-menu footer must show the version label"
);
}
}
#[cfg(test)]
mod attachment_tests {
use super::*;
const TINY_PNG_B64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==";
fn input(mime: &str, data: &str) -> AttachmentInput {
AttachmentInput {
mime: mime.to_string(),
data: data.to_string(),
filename: Some("test.img".to_string()),
}
}
#[test]
fn test_validate_attachments_happy_path() {
let out = validate_attachments(&[input("image/png", TINY_PNG_B64)]).unwrap();
assert_eq!(out.images.len(), 1);
assert_eq!(out.images[0].mime, "image/png");
assert_eq!(out.images[0].data, TINY_PNG_B64);
assert_eq!(out.images[0].filename.as_deref(), Some("test.img"));
}
#[test]
fn test_validate_attachments_mime_normalized_and_whitelisted() {
assert!(validate_attachments(&[input("IMAGE/JPEG", TINY_PNG_B64)]).is_ok());
assert!(validate_attachments(&[input("image/gif", TINY_PNG_B64)]).is_ok());
assert!(validate_attachments(&[input("image/webp", TINY_PNG_B64)]).is_ok());
assert!(validate_attachments(&[input("application/pdf", TINY_PNG_B64)]).is_err());
assert!(validate_attachments(&[input("image/svg+xml", TINY_PNG_B64)]).is_err());
}
#[test]
fn test_validate_attachments_data_url_prefix_stripped() {
let data_url = format!("data:image/png;base64,{}", TINY_PNG_B64);
let out = validate_attachments(&[input("image/png", &data_url)]).unwrap();
assert_eq!(out.images[0].data, TINY_PNG_B64);
}
#[test]
fn test_validate_attachments_invalid_base64_rejected() {
let err = validate_attachments(&[input("image/png", "not@@base64!!!")]).unwrap_err();
assert!(err.contains("invalid base64"), "got: {}", err);
}
#[test]
fn test_validate_attachments_too_many_rejected() {
let inputs: Vec<_> = (0..=MAX_ATTACHMENTS)
.map(|_| input("image/png", TINY_PNG_B64))
.collect();
let err = validate_attachments(&inputs).unwrap_err();
assert!(err.contains("Too many attachments"), "got: {}", err);
}
#[test]
fn test_validate_attachments_oversize_rejected() {
let big = "QUJD".repeat(2 * 1024 * 1024 + 4); let err = validate_attachments(&[input("image/png", &big)]).unwrap_err();
assert!(err.contains("exceeds"), "got: {}", err);
}
fn text_input(mime: &str, content: &str) -> AttachmentInput {
use base64::Engine as _;
AttachmentInput {
mime: mime.to_string(),
data: base64::engine::general_purpose::STANDARD.encode(content.as_bytes()),
filename: Some("notes.md".to_string()),
}
}
#[test]
fn test_validate_attachments_text_happy_path() {
let out = validate_attachments(&[
text_input("text/plain", "hello world"),
text_input("text/markdown", "# Title\n\nbody"),
])
.unwrap();
assert!(out.images.is_empty());
assert_eq!(out.text_files.len(), 2);
assert_eq!(out.text_files[0].mime, "text/plain");
assert_eq!(out.text_files[0].content, "hello world");
assert_eq!(out.text_files[0].filename, "notes.md");
assert_eq!(out.text_files[1].mime, "text/markdown");
assert_eq!(out.text_files[1].content, "# Title\n\nbody");
}
#[test]
fn test_validate_attachments_x_markdown_normalized() {
let out = validate_attachments(&[text_input("text/x-markdown", "# x")]).unwrap();
assert_eq!(out.text_files[0].mime, "text/markdown");
}
#[test]
fn test_validate_attachments_text_data_url_prefix_stripped() {
use base64::Engine as _;
let b64 = base64::engine::general_purpose::STANDARD.encode(b"readme body");
let data_url = format!("data:text/plain;base64,{}", b64);
let out = validate_attachments(&[input("text/plain", &data_url)]).unwrap();
assert_eq!(out.text_files[0].content, "readme body");
}
#[test]
fn test_validate_attachments_text_invalid_utf8_rejected() {
let err = validate_attachments(&[input("text/plain", "/v8A")]).unwrap_err(); assert!(err.contains("UTF-8"), "got: {}", err);
}
#[test]
fn test_validate_attachments_text_oversize_rejected() {
let big = "x".repeat(MAX_TEXT_BYTES + 1024);
let err = validate_attachments(&[text_input("text/plain", &big)]).unwrap_err();
assert!(err.contains("exceeds"), "got: {}", err);
let at_cap = "x".repeat(MAX_TEXT_BYTES);
assert!(validate_attachments(&[text_input("text/plain", &at_cap)]).is_ok());
}
#[test]
fn test_validate_attachments_mixed_families_preserved() {
let out = validate_attachments(&[
input("image/png", TINY_PNG_B64),
text_input("text/markdown", "# plan"),
])
.unwrap();
assert_eq!(out.images.len(), 1);
assert_eq!(out.text_files.len(), 1);
assert_eq!(out.text_files[0].content, "# plan");
}
#[test]
fn test_validate_attachments_too_many_across_families() {
let mut inputs = vec![
input("image/png", TINY_PNG_B64),
input("image/png", TINY_PNG_B64),
input("image/png", TINY_PNG_B64),
text_input("text/plain", "one"),
];
assert!(validate_attachments(&inputs).is_ok()); inputs.push(text_input("text/markdown", "two")); let err = validate_attachments(&inputs).unwrap_err();
assert!(err.contains("Too many attachments"), "got: {}", err);
}
#[test]
fn test_validate_attachments_unsupported_names_both_families() {
let err = validate_attachments(&[input("application/pdf", TINY_PNG_B64)]).unwrap_err();
assert!(err.contains("text/plain"), "got: {}", err);
assert!(err.contains("image/jpeg"), "got: {}", err);
}
#[test]
fn test_command_request_attachments_default_empty() {
let req: CommandRequest = serde_json::from_str(r#"{"command": "hello"}"#).unwrap();
assert!(req.attachments.is_empty());
let req: CommandRequest = serde_json::from_str(
r#"{"command": "hi", "attachments": [{"mime": "image/png", "data": "AAAA"}]}"#,
)
.unwrap();
assert_eq!(req.attachments.len(), 1);
assert!(req.attachments[0].filename.is_none());
}
}
#[cfg(test)]
mod tombstone_disclosure_tests {
use super::*;
use crate::state::ServerState;
use trustee_core::session::Session;
struct TempUser {
key: String,
}
impl TempUser {
fn new(tag: &str) -> Self {
let key = format!("tomb-disclose-{tag}-{}@test.invalid", std::process::id());
let home = dirs::home_dir()
.expect("HOME available in test env")
.join(".trustee")
.join("users")
.join(trustee_core::user_hash(&key));
std::fs::create_dir_all(home).expect("create user home");
Self { key }
}
}
fn fresh_state() -> ServerState {
let (session, _rx) = Session::new();
let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
let mut state = ServerState::new(session, ws_tx, None);
state.config_toml = Some("[agent]\nname = \"trustee\"\n".to_string());
state
}
async fn destroyed_chain(state: &ServerState, user: &TempUser, chain: &str) {
let sid = state
.create_session(&user.key, None, None)
.await
.expect("create session");
let (session, _ws) = state.get_session(&user.key, &sid).await.unwrap();
{
let mut s = session.lock().await;
s.session_id = Some(chain.to_string());
}
state
.destroy_session(&user.key, &sid)
.await
.expect("destroy");
}
#[tokio::test]
async fn tombstoned_chain_error_names_the_destruction() {
let state = fresh_state();
let user = TempUser::new("loud");
let chain = "session_2026_10_04_07_24_feedface";
destroyed_chain(&state, &user, chain).await;
let (status, body) =
loud_tombstoned_chain_error(&state, &user.key, chain, StatusCode::CONFLICT, "command");
assert_eq!(status, StatusCode::CONFLICT);
assert!(
body.contains("DESTROYED"),
"body must name the destruction: {body}"
);
assert!(body.contains(chain), "body must name the chain: {body}");
assert!(
body.contains("/history"),
"body must point at the preserved transcript: {body}"
);
assert!(
body.contains("/resume"),
"body must point at the explicit restore lane: {body}"
);
let (status2, body2) = loud_tombstoned_chain_error(
&state,
&user.key,
"session_2099_01_01_00_00_unknown0",
StatusCode::CONFLICT,
"command",
);
assert_eq!(status2, StatusCode::NOT_FOUND);
assert_eq!(body2, "Session not found");
}
#[tokio::test]
async fn chain_tombstone_stamp_roundtrips() {
let state = fresh_state();
let user = TempUser::new("stamp");
let chain = "session_2026_10_04_07_24_c0ffee00";
destroyed_chain(&state, &user, chain).await;
let stamp = state
.chain_tombstone_stamp(&user.key, chain)
.expect("destroyed chain must carry a stamp");
assert!(
stamp.starts_with("20"),
"stamp must be a timestamp, got: {stamp}"
);
assert!(
state
.chain_tombstone_stamp(&user.key, "session_2099_01_01_00_00_nomarker")
.is_none(),
"unknown chain must carry no stamp"
);
}
}
#[cfg(test)]
mod model_override_tests {
use super::*;
use crate::state::ServerState;
use trustee_core::session::Session;
const CONFIG_TOML: &str = "\
[agent]\nname = \"trustee\"\n\n[llm.provider]\nname = \"glm-zai\"\nprovider_type = \"openai\"\nbase_url = \"https://api.z.ai/api/coding/paas/v4\"\napi_key = \"test-key\"\nmodel = \"GLM-5.3-Flash\"\nmodels = [\"GLM-5.3\", \"GLM-5.2\", \"GLM-5-Turbo\", \"GLM-4.7-Flash\", \"GLM-5.3-Flash\"]\n\n[llm.providers.ninfer]\nmodel = \"qwen3.8-27b\"\nbase_url = \"https://10.253.1.9/v1\"\napi_key = \"test-key\"\n";
fn fresh_state_with_config() -> ServerState {
let (session, _rx) = Session::new();
let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
let mut state = ServerState::new(session, ws_tx, None);
state.config_toml = Some(CONFIG_TOML.to_string());
state
}
struct TempUser {
key: String,
}
impl TempUser {
fn new(tag: &str) -> Self {
let key = format!("model-ovr-{tag}-{}@test.invalid", std::process::id());
let home = dirs::home_dir()
.expect("HOME available in test env")
.join(".trustee")
.join("users")
.join(trustee_core::user_hash(&key));
std::fs::create_dir_all(home).expect("create user home");
Self { key }
}
}
#[test]
fn catalog_carries_every_resolvable_string() {
let catalog = model_catalog(CONFIG_TOML);
for m in [
"GLM-5.3-Flash", "GLM-5.3", "GLM-5.2",
"GLM-5-Turbo",
"GLM-4.7-Flash",
"ninfer", "qwen3.8-27b", ] {
assert!(
catalog.iter().any(|c| c == m),
"{m} must be in the catalog: {catalog:?}"
);
}
assert!(
!catalog.iter().any(|c| c.contains('@')),
"the inert materialized selector form must NOT be in the catalog: {catalog:?}"
);
}
#[test]
fn unknown_model_rejected_with_options() {
let err = validate_model_override(Some(CONFIG_TOML), "gpt-4o-mini")
.expect_err("the fictional fallback must be rejected");
assert!(
err.contains(&"GLM-5.3-Flash".to_string()),
"options carry the catalog: {err:?}"
);
}
#[test]
fn catalog_members_accepted() {
for m in ["GLM-5.3-Flash", "GLM-5.2", "ninfer", "qwen3.8-27b"] {
assert_eq!(
validate_model_override(Some(CONFIG_TOML), m).expect("catalog member"),
m
);
}
}
#[test]
fn empty_and_whitespace_handled() {
assert_eq!(validate_model_override(Some(CONFIG_TOML), "").expect("empty"), "");
assert_eq!(validate_model_override(Some(CONFIG_TOML), " ").expect("ws"), "");
assert_eq!(
validate_model_override(Some(CONFIG_TOML), " GLM-5.2 ").expect("trim"),
"GLM-5.2"
);
}
#[test]
fn no_config_is_not_a_rejection() {
assert_eq!(
validate_model_override(None, "anything").expect("no config"),
"anything"
);
}
#[test]
fn resolved_command_model_chain() {
std::env::remove_var("OPENAI_DEFAULT_MODEL");
assert_eq!(
resolved_command_model(Some(CONFIG_TOML), Some("GLM-5.2")).as_deref(),
Some("GLM-5.2")
);
assert_eq!(
resolved_command_model(Some(CONFIG_TOML), None).as_deref(),
Some("GLM-5.3-Flash")
);
let no_llm = "[agent]\nname = \"trustee\"\n";
assert_eq!(resolved_command_model(Some(no_llm), None), None);
}
#[tokio::test]
async fn command_unknown_model_422_with_options() {
let state = fresh_state_with_config();
let user = TempUser::new("model422");
let sid = state
.create_session(&user.key, None, None)
.await
.expect("create session");
let (session_arc, _ws) = state.get_session(&user.key, &sid).await.unwrap();
let (ws_tx, _rx) = tokio::sync::broadcast::channel::<String>(16);
let token_store = Arc::new(pep::MemoryTokenStore::new());
let req = CommandRequest {
command: "hello".to_string(),
session_id: Some(sid.clone()),
identity: None,
model: Some("gpt-4o-mini".to_string()),
attachments: vec![],
};
let result = execute_command_inner(
&state,
&axum::http::HeaderMap::new(),
&session_arc,
&ws_tx,
&token_store,
&user.key,
req,
)
.await;
let (status, body) = result.expect_err("unknown model must 422");
assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY);
assert!(body.contains("Valid models"), "options in the body: {body}");
assert!(body.contains("GLM-5.3-Flash"), "catalog in the body: {body}");
}
#[tokio::test]
async fn command_no_model_configured_422_honest() {
std::env::remove_var("OPENAI_DEFAULT_MODEL");
let (session, _rx) = Session::new();
let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
let mut state = ServerState::new(session, ws_tx, None);
state.config_toml = Some("[agent]\nname = \"trustee\"\n".to_string());
let user = TempUser::new("nomodel");
let sid = state
.create_session(&user.key, None, None)
.await
.expect("create session");
let (session_arc, _ws) = state.get_session(&user.key, &sid).await.unwrap();
let (ws_tx, _rx) = tokio::sync::broadcast::channel::<String>(16);
let token_store = Arc::new(pep::MemoryTokenStore::new());
let req = CommandRequest {
command: "hello".to_string(),
session_id: Some(sid.clone()),
identity: None,
model: None,
attachments: vec![],
};
let result = execute_command_inner(
&state,
&axum::http::HeaderMap::new(),
&session_arc,
&ws_tx,
&token_store,
&user.key,
req,
)
.await;
let (status, body) = result.expect_err("no model must 422");
assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY);
assert!(
body.contains("no model configured"),
"honest message: {body}"
);
assert!(body.contains("[llm.provider].model"), "names the key: {body}");
}
}