Users, file manager, image generation, budgets and agent workflow
- User management: local accounts (Argon2) and OpenID Connect, admin/user roles, per-user folders plus admin-shared folders, read-only session sharing, event filtering per user - File manager with CodeMirror editor and git client (status, diff, stage, commit, branches, fetch/pull/push, log, clone/init) - Image generation via fal.ai: generate_image tool and image page - Budgets: monthly limits per provider, provider balances/costs (OpenRouter, Anthropic admin cost report, fal.ai), price discovery - Subagents on other models, background tasks, concurrency limits per host/model/subagent - MCP tools loaded on demand (mcp_search/mcp_call), read-only hints - Permission mode "auto", safe read-only commands always allowed - Persistent message queue with attachments, deletable - Live reasoning incl. <think> tags, activity animations, CSS themes, mobile layout, statistics per user, upload size handling - Docker: named volume for state, localhost-bound port Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Richard ZinkandClaude Opus 5.5 committed
1 parent
f3a2fe3816
commit
49262c612c
70 files changed
+8495
-714
No files matched your search
+405
-122
@@ -1,17 +1,20 @@
|
||||
//! HTTP + WebSocket API.
|
||||
//! HTTP + WebSocket API. Every route except the public auth endpoints requires a
|
||||
//! signed-in user; admin-only routes check the role.
|
||||
|
||||
use crate::app::{App, PermissionReply};
|
||||
use crate::app::{Access, App, PermissionReply};
|
||||
use crate::auth::{self, require_admin};
|
||||
use crate::config::{self, ConfigPatch};
|
||||
use crate::prompts;
|
||||
use crate::runner;
|
||||
use crate::store::User;
|
||||
use crate::types::{Event, ImageData, Part, Session, new_id, now_ms};
|
||||
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Path, Query, Request, State};
|
||||
use axum::http::{StatusCode, header};
|
||||
use axum::middleware::{self, Next};
|
||||
use axum::extract::{Path, Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::middleware;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use axum::routing::{delete, get, post, put};
|
||||
use axum::{Extension, Json, Router};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Map, Value, json};
|
||||
use std::path::PathBuf;
|
||||
@@ -19,7 +22,7 @@ use std::sync::Arc;
|
||||
|
||||
type AppState = Arc<App>;
|
||||
|
||||
pub struct ApiError(StatusCode, String);
|
||||
pub struct ApiError(pub StatusCode, pub String);
|
||||
|
||||
impl IntoResponse for ApiError {
|
||||
fn into_response(self) -> Response {
|
||||
@@ -33,15 +36,33 @@ impl<E: Into<anyhow::Error>> From<E> for ApiError {
|
||||
}
|
||||
}
|
||||
|
||||
type ApiResult<T = Json<Value>> = Result<T, ApiError>;
|
||||
pub type ApiResult<T = Json<Value>> = Result<T, ApiError>;
|
||||
|
||||
fn not_found(what: &str) -> ApiError {
|
||||
ApiError(StatusCode::NOT_FOUND, format!("{what} not found"))
|
||||
}
|
||||
|
||||
fn forbidden() -> ApiError {
|
||||
ApiError(StatusCode::FORBIDDEN, "no access to this session".into())
|
||||
}
|
||||
|
||||
pub fn router(app: AppState) -> Router {
|
||||
let api = Router::new()
|
||||
.route("/info", get(info))
|
||||
// public
|
||||
.route("/info", get(auth::info))
|
||||
.route("/auth/login", post(auth::login))
|
||||
.route("/auth/logout", post(auth::logout))
|
||||
.route("/auth/setup", post(auth::setup))
|
||||
.route("/auth/oidc/start", get(auth::oidc_start))
|
||||
.route("/auth/oidc/callback", get(auth::oidc_callback))
|
||||
// signed in
|
||||
.route("/auth/password", post(auth::change_password))
|
||||
.route("/users", get(auth::list_users).post(auth::create_user))
|
||||
.route("/users/directory", get(auth::directory))
|
||||
.route("/users/{id}", axum::routing::patch(auth::update_user).delete(auth::delete_user))
|
||||
.route("/folders", get(list_folders).post(create_folder))
|
||||
.route("/folders/{id}", delete(delete_folder))
|
||||
.route("/folders/{id}/access", put(set_folder_access))
|
||||
.route("/config", get(get_config).patch(patch_config))
|
||||
.route("/config/raw", get(get_config_raw).put(put_config_raw))
|
||||
.route("/providers/reload", post(reload_providers))
|
||||
@@ -55,15 +76,38 @@ pub fn router(app: AppState) -> Router {
|
||||
.route("/sessions/{id}/compact", post(compact))
|
||||
.route("/sessions/{id}/rewind", post(rewind))
|
||||
.route("/sessions/{id}/stats", get(session_stats))
|
||||
.route("/sessions/{id}/shares", get(get_shares).put(put_shares))
|
||||
.route("/sessions/{id}/queue", get(get_queue))
|
||||
.route("/sessions/{id}/queue/flush", post(flush_queue))
|
||||
.route("/sessions/{id}/queue/{item}", delete(delete_queue_item))
|
||||
.route("/permissions/{id}", post(reply_permission))
|
||||
.route("/questions/{id}", post(answer_question))
|
||||
.route("/pending", get(pending))
|
||||
.route("/mcp", get(mcp_status))
|
||||
.route("/mcp/reload", post(mcp_reload))
|
||||
.route("/fs/roots", get(list_roots))
|
||||
.route("/files/list", get(crate::files::list))
|
||||
.route("/files/read", get(crate::files::read))
|
||||
.route("/files/raw", get(crate::files::raw))
|
||||
.route("/files/write", put(crate::files::write))
|
||||
.route("/files/create", post(crate::files::create))
|
||||
.route("/files/move", post(crate::files::move_or_copy))
|
||||
.route("/files/delete", delete(crate::files::remove))
|
||||
.route("/files/upload", put(crate::files::upload))
|
||||
.route("/files/search", get(crate::files::search))
|
||||
.route("/git", post(crate::git::handle))
|
||||
.route("/fs/dirs", get(list_dirs))
|
||||
.route("/stats", get(stats))
|
||||
.route("/budget", get(budget))
|
||||
.route("/providers/{name}/prices", get(provider_prices))
|
||||
.route("/images", get(list_images).post(create_images))
|
||||
.route("/images/config", get(images_config))
|
||||
.route("/images/{id}", delete(delete_image))
|
||||
.route("/images/{id}/file", get(image_file))
|
||||
.route("/ws", get(ws))
|
||||
.layer(middleware::from_fn_with_state(app.clone(), auth));
|
||||
.layer(middleware::from_fn_with_state(app.clone(), auth::require_user))
|
||||
// Prompts may carry images (base64); the default limit is 2 MB.
|
||||
.layer(axum::extract::DefaultBodyLimit::max(50 * 1024 * 1024));
|
||||
|
||||
let web_dir = std::env::var("CODECRAFT_WEB").map(PathBuf::from).unwrap_or_else(|_| {
|
||||
["../web/dist", "./web/dist", "/app/web"]
|
||||
@@ -73,64 +117,42 @@ pub fn router(app: AppState) -> Router {
|
||||
.unwrap_or_else(|| PathBuf::from("./web/dist"))
|
||||
});
|
||||
let index = web_dir.join("index.html");
|
||||
let static_files = tower_http::services::ServeDir::new(&web_dir)
|
||||
.fallback(tower_http::services::ServeFile::new(index));
|
||||
let static_files = tower_http::services::ServeDir::new(&web_dir).fallback(tower_http::services::ServeFile::new(index));
|
||||
|
||||
Router::new().nest("/api", api).fallback_service(static_files).with_state(app)
|
||||
}
|
||||
|
||||
/// Bearer token, `token` query parameter (WebSocket) or `codecraft_token` cookie.
|
||||
async fn auth(State(app): State<AppState>, req: Request, next: Next) -> Response {
|
||||
let expected = app.config().auth_token();
|
||||
if expected.is_empty() || req.uri().path().ends_with("/info") {
|
||||
return next.run(req).await;
|
||||
/// Load a session the user may at least read.
|
||||
fn readable(app: &App, user: &User, id: &str) -> Result<Session, ApiError> {
|
||||
let s = app.store.get_session(id)?.ok_or_else(|| not_found("session"))?;
|
||||
if app.access(user, id) == Access::None {
|
||||
return Err(forbidden());
|
||||
}
|
||||
let header_token = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.map(String::from);
|
||||
let query_token = req.uri().query().and_then(|q| {
|
||||
q.split('&').find_map(|kv| kv.strip_prefix("token=")).map(String::from)
|
||||
});
|
||||
let cookie_token = req
|
||||
.headers()
|
||||
.get(header::COOKIE)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|c| c.split(';').find_map(|kv| kv.trim().strip_prefix("codecraft_token=")).map(String::from));
|
||||
let ok = [header_token, query_token, cookie_token]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.any(|t| constant_time_eq(t.as_bytes(), expected.as_bytes()));
|
||||
if ok {
|
||||
next.run(req).await
|
||||
} else {
|
||||
ApiError(StatusCode::UNAUTHORIZED, "unauthorized".into()).into_response()
|
||||
Ok(s)
|
||||
}
|
||||
|
||||
/// Load a session the user may control (owner or admin).
|
||||
fn writable(app: &App, user: &User, id: &str) -> Result<Session, ApiError> {
|
||||
let s = readable(app, user, id)?;
|
||||
if app.access(user, id) != Access::Full {
|
||||
return Err(ApiError(StatusCode::FORBIDDEN, "this session is shared with you read-only".into()));
|
||||
}
|
||||
Ok(s)
|
||||
}
|
||||
|
||||
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
|
||||
a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
|
||||
}
|
||||
// --- config (admin) -----------------------------------------------------------
|
||||
|
||||
// --- info & config ----------------------------------------------------------
|
||||
|
||||
async fn info(State(app): State<AppState>) -> Json<Value> {
|
||||
let cfg = app.config();
|
||||
Json(json!({
|
||||
"name": "CodeCraft",
|
||||
"version": env!("CARGO_PKG_VERSION"),
|
||||
"auth_required": !cfg.auth_token().is_empty(),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn get_config(State(app): State<AppState>) -> ApiResult {
|
||||
async fn get_config(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
let cfg = app.config();
|
||||
let default_model = if cfg.server.default_model.is_empty() { app.registry().first_model() } else { Some(cfg.server.default_model.clone()) };
|
||||
if !user.is_admin() {
|
||||
// Users only need defaults, not the configuration.
|
||||
return Ok(Json(json!({"default_model": default_model, "currency": cfg.server.currency})));
|
||||
}
|
||||
Ok(Json(json!({
|
||||
"config": config::redacted_json(&cfg),
|
||||
"workspace": app.workspace(),
|
||||
"default_model": if cfg.server.default_model.is_empty() { app.registry().first_model() } else { Some(cfg.server.default_model.clone()) },
|
||||
"default_model": default_model,
|
||||
"currency": cfg.server.currency,
|
||||
})))
|
||||
}
|
||||
@@ -140,14 +162,16 @@ struct PatchBody {
|
||||
patches: Vec<ConfigPatch>,
|
||||
}
|
||||
|
||||
async fn patch_config(State(app): State<AppState>, Json(body): Json<PatchBody>) -> ApiResult {
|
||||
async fn patch_config(Extension(user): Extension<User>, State(app): State<AppState>, Json(body): Json<PatchBody>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
// Masked secrets coming back from the UI must not overwrite real values.
|
||||
let patches: Vec<ConfigPatch> = body.patches.into_iter().filter(|p| p.value.as_str() != Some("••••••••")).collect();
|
||||
app.patch_config(&patches).await?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
async fn get_config_raw(State(app): State<AppState>) -> ApiResult {
|
||||
async fn get_config_raw(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
Ok(Json(json!({"text": app.config_text()?})))
|
||||
}
|
||||
|
||||
@@ -156,12 +180,14 @@ struct RawBody {
|
||||
text: String,
|
||||
}
|
||||
|
||||
async fn put_config_raw(State(app): State<AppState>, Json(body): Json<RawBody>) -> ApiResult {
|
||||
async fn put_config_raw(Extension(user): Extension<User>, State(app): State<AppState>, Json(body): Json<RawBody>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
app.save_config(&body.text).await?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
async fn reload_providers(State(app): State<AppState>) -> ApiResult {
|
||||
async fn reload_providers(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
app.reload_providers().await;
|
||||
Ok(Json(app.registry().models_json()))
|
||||
}
|
||||
@@ -175,33 +201,92 @@ struct CwdQuery {
|
||||
cwd: Option<String>,
|
||||
}
|
||||
|
||||
fn cwd_or_workspace(app: &App, cwd: Option<String>) -> PathBuf {
|
||||
cwd.map(PathBuf::from).unwrap_or_else(|| app.workspace())
|
||||
/// Root and cwd for listing agents/commands; falls back to the user's first root.
|
||||
fn scope(app: &App, user: &User, cwd: Option<String>) -> (PathBuf, PathBuf) {
|
||||
let owner = Some(user.id.as_str());
|
||||
let roots = app.user_roots(owner);
|
||||
let cwd = cwd.map(PathBuf::from).filter(|c| roots.iter().any(|r| c.starts_with(r))).unwrap_or_else(|| roots[0].clone());
|
||||
(app.root_for(owner, &cwd), cwd)
|
||||
}
|
||||
|
||||
async fn agents(State(app): State<AppState>, Query(q): Query<CwdQuery>) -> Json<Value> {
|
||||
let cwd = cwd_or_workspace(&app, q.cwd);
|
||||
Json(json!(prompts::agents(&app, &cwd)))
|
||||
async fn agents(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<CwdQuery>) -> Json<Value> {
|
||||
let (root, cwd) = scope(&app, &user, q.cwd);
|
||||
Json(json!(prompts::agents(&app, &root, &cwd)))
|
||||
}
|
||||
|
||||
async fn commands(State(app): State<AppState>, Query(q): Query<CwdQuery>) -> Json<Value> {
|
||||
let cwd = cwd_or_workspace(&app, q.cwd);
|
||||
async fn commands(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<CwdQuery>) -> Json<Value> {
|
||||
let (root, cwd) = scope(&app, &user, q.cwd);
|
||||
Json(json!({
|
||||
"commands": prompts::commands(&app, &cwd),
|
||||
"skills": prompts::skills(&app, &cwd),
|
||||
"commands": prompts::commands(&app, &root, &cwd),
|
||||
"skills": prompts::skills(&app, &root, &cwd),
|
||||
}))
|
||||
}
|
||||
|
||||
// --- shared folders (admin) ---------------------------------------------------
|
||||
|
||||
async fn list_folders(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
Ok(Json(json!({"workspace": app.workspace(), "folders": app.store.folders()?})))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct FolderBody {
|
||||
/// Path relative to the workspace (or absolute inside it). Created if missing.
|
||||
path: String,
|
||||
#[serde(default)]
|
||||
users: Vec<String>,
|
||||
}
|
||||
|
||||
async fn create_folder(Extension(user): Extension<User>, State(app): State<AppState>, Json(body): Json<FolderBody>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
let ws = app.workspace();
|
||||
let rel = body.path.trim().trim_start_matches(ws.to_str().unwrap_or("")).trim_matches('/');
|
||||
if rel.is_empty() || rel.split('/').any(|c| c.is_empty() || c == "." || c == ".." || c.starts_with('.')) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "invalid folder path".into()));
|
||||
}
|
||||
let path = ws.join(rel);
|
||||
std::fs::create_dir_all(&path)?;
|
||||
let path = std::fs::canonicalize(&path)?;
|
||||
if !path.starts_with(&ws) || path == ws {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "folder must be inside the workspace".into()));
|
||||
}
|
||||
let id = app.store.add_folder(&path.display().to_string())?;
|
||||
app.store.set_folder_access(&id, &body.users)?;
|
||||
Ok(Json(json!({"id": id, "path": path})))
|
||||
}
|
||||
|
||||
async fn delete_folder(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
// Only the share is removed; the directory and its files stay.
|
||||
app.store.delete_folder(&id)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AccessBody {
|
||||
users: Vec<String>,
|
||||
}
|
||||
|
||||
async fn set_folder_access(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<AccessBody>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
app.store.set_folder_access(&id, &body.users)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
// --- sessions ---------------------------------------------------------------
|
||||
|
||||
async fn list_sessions(State(app): State<AppState>) -> ApiResult {
|
||||
let sessions = app.store.list_sessions()?;
|
||||
let list: Vec<Value> = sessions
|
||||
async fn list_sessions(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
let filter = if user.is_admin() { None } else { Some(user.id.as_str()) };
|
||||
let list: Vec<Value> = app
|
||||
.store
|
||||
.list_sessions(filter)?
|
||||
.into_iter()
|
||||
.map(|s| {
|
||||
let busy = app.is_busy(&s.id);
|
||||
let shared = s.owner_id.as_deref() != Some(user.id.as_str());
|
||||
let mut v = json!(s);
|
||||
v["busy"] = json!(busy);
|
||||
v["shared"] = json!(shared);
|
||||
v
|
||||
})
|
||||
.collect();
|
||||
@@ -216,29 +301,34 @@ struct CreateSession {
|
||||
permission_mode: Option<String>,
|
||||
}
|
||||
|
||||
fn validate_cwd(app: &App, cwd: Option<String>) -> Result<PathBuf, ApiError> {
|
||||
let ws = app.workspace();
|
||||
/// Resolve a directory and check it lies inside one of the user's roots.
|
||||
fn validate_cwd(app: &App, user: &User, cwd: Option<String>) -> Result<PathBuf, ApiError> {
|
||||
let roots = app.user_roots(Some(&user.id));
|
||||
let p = match cwd.filter(|c| !c.trim().is_empty()) {
|
||||
Some(c) if std::path::Path::new(&c).is_absolute() => PathBuf::from(c),
|
||||
Some(c) => ws.join(c),
|
||||
None => ws.clone(),
|
||||
Some(c) => roots[0].join(c),
|
||||
None => roots[0].clone(),
|
||||
};
|
||||
let p = std::fs::canonicalize(&p).map_err(|_| ApiError(StatusCode::BAD_REQUEST, format!("directory does not exist: {}", p.display())))?;
|
||||
if !p.starts_with(&ws) || !p.is_dir() {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "directory must be inside the workspace".into()));
|
||||
if !p.is_dir() || !roots.iter().any(|r| p.starts_with(r)) {
|
||||
return Err(ApiError(StatusCode::FORBIDDEN, "directory is outside your folders".into()));
|
||||
}
|
||||
Ok(p)
|
||||
}
|
||||
|
||||
async fn create_session(State(app): State<AppState>, Json(body): Json<CreateSession>) -> ApiResult {
|
||||
let cwd = validate_cwd(&app, body.cwd)?;
|
||||
async fn create_session(Extension(user): Extension<User>, State(app): State<AppState>, Json(body): Json<CreateSession>) -> ApiResult {
|
||||
let cwd = validate_cwd(&app, &user, body.cwd)?;
|
||||
let cfg = app.config();
|
||||
let model = body
|
||||
.model
|
||||
.filter(|m| !m.is_empty())
|
||||
.or_else(|| Some(cfg.server.default_model.clone()).filter(|m| !m.is_empty()))
|
||||
.or_else(|| app.registry().first_model())
|
||||
.ok_or_else(|| ApiError(StatusCode::BAD_REQUEST, "no model configured – add a provider in the settings".into()))?;
|
||||
.ok_or_else(|| ApiError(StatusCode::BAD_REQUEST, "no model configured – ask an admin to add a provider".into()))?;
|
||||
let mode = body.permission_mode.unwrap_or_else(|| "auto".into());
|
||||
if !crate::permission::MODES.contains(&mode.as_str()) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "invalid permission mode".into()));
|
||||
}
|
||||
let now = now_ms();
|
||||
let s = Session {
|
||||
id: new_id(),
|
||||
@@ -246,8 +336,9 @@ async fn create_session(State(app): State<AppState>, Json(body): Json<CreateSess
|
||||
cwd: cwd.display().to_string(),
|
||||
model,
|
||||
agent: body.agent.unwrap_or_else(|| "build".into()),
|
||||
permission_mode: body.permission_mode.unwrap_or_else(|| "default".into()),
|
||||
permission_mode: mode,
|
||||
parent_id: None,
|
||||
owner_id: Some(user.id.clone()),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
todos: Value::Null,
|
||||
@@ -257,11 +348,13 @@ async fn create_session(State(app): State<AppState>, Json(body): Json<CreateSess
|
||||
Ok(Json(json!(s)))
|
||||
}
|
||||
|
||||
async fn get_session(State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
let s = app.store.get_session(&id)?.ok_or_else(|| not_found("session"))?;
|
||||
async fn get_session(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
let s = readable(&app, &user, &id)?;
|
||||
let messages = app.store.messages(&id)?;
|
||||
let busy = app.is_busy(&id);
|
||||
Ok(Json(json!({"session": s, "messages": messages, "busy": busy})))
|
||||
let can_write = app.access(&user, &id) == Access::Full;
|
||||
let queue = app.store.queue_list(&id)?;
|
||||
Ok(Json(json!({"session": s, "messages": messages, "busy": busy, "can_write": can_write, "queue": queue})))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -272,8 +365,8 @@ struct UpdateSession {
|
||||
permission_mode: Option<String>,
|
||||
}
|
||||
|
||||
async fn update_session(State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<UpdateSession>) -> ApiResult {
|
||||
let mut s = app.store.get_session(&id)?.ok_or_else(|| not_found("session"))?;
|
||||
async fn update_session(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<UpdateSession>) -> ApiResult {
|
||||
let mut s = writable(&app, &user, &id)?;
|
||||
if let Some(t) = body.title {
|
||||
s.title = t;
|
||||
}
|
||||
@@ -285,7 +378,7 @@ async fn update_session(State(app): State<AppState>, Path(id): Path<String>, Jso
|
||||
s.agent = a;
|
||||
}
|
||||
if let Some(p) = body.permission_mode {
|
||||
if !["default", "accept_edits", "bypass"].contains(&p.as_str()) {
|
||||
if !crate::permission::MODES.contains(&p.as_str()) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "invalid permission mode".into()));
|
||||
}
|
||||
s.permission_mode = p;
|
||||
@@ -293,10 +386,12 @@ async fn update_session(State(app): State<AppState>, Path(id): Path<String>, Jso
|
||||
Ok(Json(json!(app.save_session(s)?)))
|
||||
}
|
||||
|
||||
async fn delete_session(State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
async fn delete_session(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
writable(&app, &user, &id)?;
|
||||
// Notify before deleting: visibility is checked against the stored session.
|
||||
app.emit(Event::SessionDeleted { session_id: id.clone() });
|
||||
runner::abort(&app, &id);
|
||||
app.store.delete_session(&id)?;
|
||||
app.emit(Event::SessionDeleted { session_id: id });
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
@@ -307,8 +402,8 @@ struct PromptBody {
|
||||
images: Vec<ImageData>,
|
||||
}
|
||||
|
||||
async fn prompt(State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<PromptBody>) -> ApiResult {
|
||||
let session = app.store.get_session(&id)?.ok_or_else(|| not_found("session"))?;
|
||||
async fn prompt(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<PromptBody>) -> ApiResult {
|
||||
let session = writable(&app, &user, &id)?;
|
||||
let text = body.text.trim().to_string();
|
||||
// Built-in commands.
|
||||
if text == "/compact" {
|
||||
@@ -319,9 +414,10 @@ async fn prompt(State(app): State<AppState>, Path(id): Path<String>, Json(body):
|
||||
return Ok(Json(json!({"ok": true, "command": "compact"})));
|
||||
}
|
||||
let cwd = PathBuf::from(&session.cwd);
|
||||
let root = app.root_for(session.owner_id.as_deref(), &cwd);
|
||||
let mut prompt_text = text.clone();
|
||||
if text.starts_with('/') {
|
||||
if let Some((expanded, agent)) = prompts::expand_command(&app, &cwd, &text) {
|
||||
if let Some((expanded, agent)) = prompts::expand_command(&app, &root, &cwd, &text) {
|
||||
prompt_text = expanded;
|
||||
if let Some(a) = agent {
|
||||
let mut s = session.clone();
|
||||
@@ -334,10 +430,32 @@ async fn prompt(State(app): State<AppState>, Path(id): Path<String>, Json(body):
|
||||
for image in body.images {
|
||||
parts.push(Part::Image { image });
|
||||
}
|
||||
runner::start_prompt(app.clone(), id, parts)?;
|
||||
match runner::submit(app.clone(), id, parts, "user")? {
|
||||
runner::Submitted::Started => Ok(Json(json!({"ok": true}))),
|
||||
runner::Submitted::Queued(item) => Ok(Json(json!({"ok": true, "queued": item}))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_queue(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
readable(&app, &user, &id)?;
|
||||
Ok(Json(json!(app.store.queue_list(&id)?)))
|
||||
}
|
||||
|
||||
async fn delete_queue_item(Extension(user): Extension<User>, State(app): State<AppState>, Path((id, item)): Path<(String, String)>) -> ApiResult {
|
||||
writable(&app, &user, &id)?;
|
||||
if !app.store.queue_remove(&id, &item)? {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "already delivered to the agent".into()));
|
||||
}
|
||||
let items = app.store.queue_list(&id)?;
|
||||
app.emit(Event::QueueUpdated { session_id: id, items: json!(items) });
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
async fn flush_queue(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
writable(&app, &user, &id)?;
|
||||
Ok(Json(json!({"started": runner::flush_queue(app.clone(), &id)?})))
|
||||
}
|
||||
|
||||
async fn do_compact(app: AppState, id: String) {
|
||||
let cancel = tokio_util::sync::CancellationToken::new();
|
||||
app.with_runtime(&id, |r| r.cancel = Some(cancel.clone()));
|
||||
@@ -349,12 +467,14 @@ async fn do_compact(app: AppState, id: String) {
|
||||
app.emit(Event::SessionStatus { session_id: id, busy: false });
|
||||
}
|
||||
|
||||
async fn abort(State(app): State<AppState>, Path(id): Path<String>) -> Json<Value> {
|
||||
async fn abort(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
writable(&app, &user, &id)?;
|
||||
runner::abort(&app, &id);
|
||||
Json(json!({"ok": true}))
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
async fn compact(State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
async fn compact(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
writable(&app, &user, &id)?;
|
||||
if app.is_busy(&id) {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "session is busy".into()));
|
||||
}
|
||||
@@ -368,33 +488,58 @@ struct RewindBody {
|
||||
}
|
||||
|
||||
/// Remove a message and everything after it (edit & resend). Files are not restored.
|
||||
async fn rewind(State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<RewindBody>) -> ApiResult {
|
||||
async fn rewind(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<RewindBody>) -> ApiResult {
|
||||
let mut s = writable(&app, &user, &id)?;
|
||||
if app.is_busy(&id) {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "session is busy".into()));
|
||||
}
|
||||
app.store.truncate_from(&id, &body.message_id)?;
|
||||
let s = app.store.get_session(&id)?.ok_or_else(|| not_found("session"))?;
|
||||
let mut s = s;
|
||||
s.last_usage = None;
|
||||
app.save_session(s)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
async fn session_stats(State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
async fn session_stats(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
readable(&app, &user, &id)?;
|
||||
let mut v = app.store.session_stats(&id)?;
|
||||
v["currency"] = json!(app.config().server.currency);
|
||||
Ok(Json(v))
|
||||
}
|
||||
|
||||
async fn get_shares(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
writable(&app, &user, &id)?;
|
||||
Ok(Json(json!(app.store.shares(&id)?)))
|
||||
}
|
||||
|
||||
async fn put_shares(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<AccessBody>) -> ApiResult {
|
||||
let s = writable(&app, &user, &id)?;
|
||||
if s.parent_id.is_some() {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "share the parent session instead".into()));
|
||||
}
|
||||
let users: Vec<String> = body.users.into_iter().filter(|u| Some(u) != s.owner_id.as_ref()).collect();
|
||||
app.set_shares(&id, &users)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
// --- interaction ------------------------------------------------------------
|
||||
|
||||
/// Only users who control the session may answer its permission requests and questions.
|
||||
fn check_pending(app: &App, user: &User, id: &str) -> Result<(), ApiError> {
|
||||
let sid = app.pending_session(id).ok_or_else(|| not_found("request"))?;
|
||||
if app.access(user, &sid) != Access::Full {
|
||||
return Err(forbidden());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct PermissionBody {
|
||||
reply: PermissionReply,
|
||||
message: Option<String>,
|
||||
}
|
||||
|
||||
async fn reply_permission(State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<PermissionBody>) -> ApiResult {
|
||||
async fn reply_permission(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<PermissionBody>) -> ApiResult {
|
||||
check_pending(&app, &user, &id)?;
|
||||
app.reply_permission(&id, body.reply, body.message)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
@@ -405,37 +550,55 @@ struct QuestionBody {
|
||||
answers: Option<Map<String, Value>>,
|
||||
}
|
||||
|
||||
async fn answer_question(State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<QuestionBody>) -> ApiResult {
|
||||
async fn answer_question(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>, Json(body): Json<QuestionBody>) -> ApiResult {
|
||||
check_pending(&app, &user, &id)?;
|
||||
app.answer_question(&id, body.answers)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
async fn pending(State(app): State<AppState>) -> Json<Value> {
|
||||
Json(json!(app.pending_events()))
|
||||
async fn pending(Extension(user): Extension<User>, State(app): State<AppState>) -> Json<Value> {
|
||||
let list: Vec<Event> = app.pending_events().into_iter().filter(|e| app.event_visible(&user, e)).collect();
|
||||
Json(json!(list))
|
||||
}
|
||||
|
||||
async fn mcp_status(State(app): State<AppState>) -> Json<Value> {
|
||||
Json(app.mcp.status_json())
|
||||
async fn mcp_status(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
Ok(Json(app.mcp.status_json()))
|
||||
}
|
||||
|
||||
async fn mcp_reload(State(app): State<AppState>) -> Json<Value> {
|
||||
async fn mcp_reload(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
let mcp = app.mcp.clone();
|
||||
let servers = app.config().mcp.clone();
|
||||
let weak = Arc::downgrade(&app);
|
||||
mcp.start_all(weak, servers).await;
|
||||
Json(app.mcp.status_json())
|
||||
Ok(Json(app.mcp.status_json()))
|
||||
}
|
||||
|
||||
// --- files ------------------------------------------------------------------
|
||||
|
||||
/// The directories the user may work in (own folder, shared folders).
|
||||
async fn list_roots(Extension(user): Extension<User>, State(app): State<AppState>) -> Json<Value> {
|
||||
let ws = app.workspace();
|
||||
let roots: Vec<Value> = app
|
||||
.user_roots(Some(&user.id))
|
||||
.into_iter()
|
||||
.map(|r| {
|
||||
let rel = r.strip_prefix(&ws).map(|p| p.display().to_string()).unwrap_or_default();
|
||||
json!({"path": r, "name": if rel.is_empty() { "/".to_string() } else { rel }})
|
||||
})
|
||||
.collect();
|
||||
Json(json!(roots))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct DirQuery {
|
||||
path: Option<String>,
|
||||
}
|
||||
|
||||
/// Subdirectories for the project picker (inside the workspace only).
|
||||
async fn list_dirs(State(app): State<AppState>, Query(q): Query<DirQuery>) -> ApiResult {
|
||||
let dir = validate_cwd(&app, q.path)?;
|
||||
/// Subdirectories for the project picker (inside the user's roots only).
|
||||
async fn list_dirs(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<DirQuery>) -> ApiResult {
|
||||
let dir = validate_cwd(&app, &user, q.path)?;
|
||||
let mut dirs = Vec::new();
|
||||
for e in std::fs::read_dir(&dir)?.flatten() {
|
||||
let name = e.file_name().to_string_lossy().to_string();
|
||||
@@ -445,15 +608,125 @@ async fn list_dirs(State(app): State<AppState>, Query(q): Query<DirQuery>) -> Ap
|
||||
dirs.push(json!({"name": name, "path": e.path(), "git": e.path().join(".git").exists()}));
|
||||
}
|
||||
dirs.sort_by(|a, b| a["name"].as_str().cmp(&b["name"].as_str()));
|
||||
let ws = app.workspace();
|
||||
let root = app.root_for(Some(&user.id), &dir);
|
||||
Ok(Json(json!({
|
||||
"path": dir,
|
||||
"parent": if dir == ws { Value::Null } else { json!(dir.parent()) },
|
||||
"workspace": ws,
|
||||
"parent": if dir == root { Value::Null } else { json!(dir.parent()) },
|
||||
"root": root,
|
||||
"workspace": app.workspace(),
|
||||
"dirs": dirs,
|
||||
})))
|
||||
}
|
||||
|
||||
// --- budget -------------------------------------------------------------------
|
||||
|
||||
/// Budgets concern the whole installation: admins only.
|
||||
async fn budget(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
Ok(Json(crate::budget::overview(&app).await?))
|
||||
}
|
||||
|
||||
async fn provider_prices(Extension(user): Extension<User>, State(app): State<AppState>, Path(name): Path<String>) -> ApiResult {
|
||||
require_admin(&user)?;
|
||||
Ok(Json(crate::budget::discover_prices(&app, &name).await?))
|
||||
}
|
||||
|
||||
// --- images -------------------------------------------------------------------
|
||||
|
||||
async fn images_config(State(app): State<AppState>) -> Json<Value> {
|
||||
let cfg = app.config();
|
||||
let ic = &cfg.images;
|
||||
Json(json!({
|
||||
"enabled": ic.enabled && !crate::config::expand_env(&ic.api_key).is_empty(),
|
||||
"default_model": ic.default_model,
|
||||
"output_dir": ic.output_dir,
|
||||
"sizes": crate::images::SIZES,
|
||||
"models": ic.models.iter().map(|m| json!({"id": m.id, "name": m.name.clone().unwrap_or_else(|| m.id.clone()), "price": m.price})).collect::<Vec<_>>(),
|
||||
"currency": cfg.server.currency,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn list_images(Extension(user): Extension<User>, State(app): State<AppState>) -> ApiResult {
|
||||
Ok(Json(json!(app.store.images(Some(&user.id), 200)?)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ImageBody {
|
||||
prompt: String,
|
||||
model: Option<String>,
|
||||
size: Option<String>,
|
||||
#[serde(default = "one")]
|
||||
count: u32,
|
||||
/// Target directory (absolute inside the user's folders); default: first folder/output_dir.
|
||||
directory: Option<String>,
|
||||
}
|
||||
|
||||
fn one() -> u32 {
|
||||
1
|
||||
}
|
||||
|
||||
async fn create_images(Extension(user): Extension<User>, State(app): State<AppState>, Json(body): Json<ImageBody>) -> ApiResult {
|
||||
if body.prompt.trim().is_empty() {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "prompt is empty".into()));
|
||||
}
|
||||
let roots = app.user_roots(Some(&user.id));
|
||||
let dir = match body.directory.filter(|d| !d.trim().is_empty()) {
|
||||
Some(d) => {
|
||||
let p = crate::tools::normalize(&PathBuf::from(d));
|
||||
if !roots.iter().any(|r| p.starts_with(r)) {
|
||||
return Err(ApiError(StatusCode::FORBIDDEN, "directory is outside your folders".into()));
|
||||
}
|
||||
p
|
||||
}
|
||||
None => roots[0].join(&app.config().images.output_dir),
|
||||
};
|
||||
let records = crate::images::generate(
|
||||
&app,
|
||||
crate::images::GenerateRequest {
|
||||
prompt: body.prompt.trim().to_string(),
|
||||
model: body.model,
|
||||
size: body.size,
|
||||
count: body.count,
|
||||
dir,
|
||||
user_id: Some(user.id.clone()),
|
||||
session_id: None,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(json!(records)))
|
||||
}
|
||||
|
||||
fn own_image(app: &App, user: &User, id: &str) -> Result<crate::store::ImageRecord, ApiError> {
|
||||
let img = app.store.image(id)?.ok_or_else(|| not_found("image"))?;
|
||||
if !user.is_admin() && img.user_id.as_deref() != Some(user.id.as_str()) {
|
||||
return Err(ApiError(StatusCode::FORBIDDEN, "not your image".into()));
|
||||
}
|
||||
Ok(img)
|
||||
}
|
||||
|
||||
async fn image_file(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> Result<Response, ApiError> {
|
||||
let img = own_image(&app, &user, &id)?;
|
||||
let bytes = tokio::fs::read(&img.path).await.map_err(|_| not_found("image file"))?;
|
||||
let ctype = match std::path::Path::new(&img.path).extension().and_then(|e| e.to_str()) {
|
||||
Some("jpg") | Some("jpeg") => "image/jpeg",
|
||||
Some("webp") => "image/webp",
|
||||
_ => "image/png",
|
||||
};
|
||||
Ok(([(axum::http::header::CONTENT_TYPE, ctype), (axum::http::header::CACHE_CONTROL, "private, max-age=86400")], bytes).into_response())
|
||||
}
|
||||
|
||||
async fn delete_image(Extension(user): Extension<User>, State(app): State<AppState>, Path(id): Path<String>) -> ApiResult {
|
||||
let img = own_image(&app, &user, &id)?;
|
||||
// Only delete files inside the owner's folders.
|
||||
let roots = app.user_roots(img.user_id.as_deref());
|
||||
let p = PathBuf::from(&img.path);
|
||||
if roots.iter().any(|r| p.starts_with(r)) {
|
||||
let _ = tokio::fs::remove_file(&p).await;
|
||||
}
|
||||
app.store.delete_image(&id)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
// --- statistics -------------------------------------------------------------
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -462,12 +735,15 @@ struct StatsQuery {
|
||||
to: Option<i64>,
|
||||
/// Minutes east of UTC (JS: -new Date().getTimezoneOffset()).
|
||||
tz: Option<i64>,
|
||||
/// Admins: user id to filter by (omitted = all users).
|
||||
user: Option<String>,
|
||||
}
|
||||
|
||||
async fn stats(State(app): State<AppState>, Query(q): Query<StatsQuery>) -> ApiResult {
|
||||
async fn stats(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<StatsQuery>) -> ApiResult {
|
||||
let to = q.to.unwrap_or_else(|| now_ms() + 1);
|
||||
let from = q.from.unwrap_or(to - 30 * 86_400_000);
|
||||
let mut v = app.store.global_stats(from, to, q.tz.unwrap_or(0))?;
|
||||
let filter = if user.is_admin() { q.user.filter(|u| !u.is_empty()) } else { Some(user.id.clone()) };
|
||||
let mut v = app.store.global_stats(from, to, q.tz.unwrap_or(0), filter.as_deref())?;
|
||||
v["currency"] = json!(app.config().server.currency);
|
||||
v["from"] = json!(from);
|
||||
v["to"] = json!(to);
|
||||
@@ -476,23 +752,25 @@ async fn stats(State(app): State<AppState>, Query(q): Query<StatsQuery>) -> ApiR
|
||||
|
||||
// --- websocket --------------------------------------------------------------
|
||||
|
||||
async fn ws(State(app): State<AppState>, upgrade: WebSocketUpgrade) -> Response {
|
||||
upgrade.on_upgrade(move |socket| ws_loop(app, socket))
|
||||
async fn ws(Extension(user): Extension<User>, State(app): State<AppState>, upgrade: WebSocketUpgrade) -> Response {
|
||||
upgrade.on_upgrade(move |socket| ws_loop(app, user, socket))
|
||||
}
|
||||
|
||||
async fn ws_loop(app: AppState, mut socket: WebSocket) {
|
||||
async fn ws_loop(app: AppState, user: User, mut socket: WebSocket) {
|
||||
let mut rx = app.subscribe();
|
||||
for ev in app.pending_events() {
|
||||
if send_event(&mut socket, &ev).await.is_err() {
|
||||
if app.event_visible(&user, &ev) && send_event(&mut socket, &ev).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
let mut ping = tokio::time::interval(std::time::Duration::from_secs(25));
|
||||
// Re-check the login regularly so disabled users lose their live connection.
|
||||
let mut recheck = tokio::time::interval(std::time::Duration::from_secs(60));
|
||||
loop {
|
||||
tokio::select! {
|
||||
ev = rx.recv() => match ev {
|
||||
Ok(ev) => {
|
||||
if send_event(&mut socket, &ev).await.is_err() {
|
||||
if app.event_visible(&user, &ev) && send_event(&mut socket, &ev).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -514,6 +792,11 @@ async fn ws_loop(app: AppState, mut socket: WebSocket) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ = recheck.tick() => {
|
||||
if !matches!(app.store.user(&user.id), Ok(Some(u)) if !u.disabled) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+155
-3
@@ -10,7 +10,7 @@ use serde_json::{Map, Value};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex, RwLock};
|
||||
use tokio::sync::{broadcast, oneshot};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, oneshot};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize)]
|
||||
@@ -39,8 +39,6 @@ pub struct Runtime {
|
||||
pub cancel: Option<CancellationToken>,
|
||||
pub read_files: HashSet<PathBuf>,
|
||||
pub approved: HashSet<String>,
|
||||
/// Prompts sent while the agent was busy; injected at the next step.
|
||||
pub queue: Vec<String>,
|
||||
}
|
||||
|
||||
pub struct App {
|
||||
@@ -53,6 +51,24 @@ pub struct App {
|
||||
runtime: Mutex<HashMap<String, Runtime>>,
|
||||
permissions: Mutex<HashMap<String, PendingPermission>>,
|
||||
questions: Mutex<HashMap<String, PendingQuestion>>,
|
||||
/// Subagent concurrency limits: key → (size, semaphore). Recreated when the size changes.
|
||||
subagent_limits: Mutex<HashMap<String, (usize, Arc<Semaphore>)>>,
|
||||
/// One-time code for creating the first admin (None once users exist).
|
||||
pub setup_code: Mutex<Option<String>>,
|
||||
pub oidc_pending: Mutex<HashMap<String, crate::auth::OidcPending>>,
|
||||
pub oidc_discovery: Mutex<Option<(String, Value)>>,
|
||||
/// session id → (root session id, owner) – both never change.
|
||||
owners: Mutex<HashMap<String, (String, Option<String>)>>,
|
||||
/// root session id → users it is shared with.
|
||||
shares: Mutex<HashMap<String, Vec<String>>>,
|
||||
}
|
||||
|
||||
/// What a user may do with a session.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Access {
|
||||
None,
|
||||
Read,
|
||||
Full,
|
||||
}
|
||||
|
||||
impl App {
|
||||
@@ -72,7 +88,14 @@ impl App {
|
||||
runtime: Mutex::default(),
|
||||
permissions: Mutex::default(),
|
||||
questions: Mutex::default(),
|
||||
subagent_limits: Mutex::default(),
|
||||
setup_code: Mutex::default(),
|
||||
oidc_pending: Mutex::default(),
|
||||
oidc_discovery: Mutex::default(),
|
||||
owners: Mutex::default(),
|
||||
shares: Mutex::default(),
|
||||
});
|
||||
crate::auth::bootstrap(&app)?;
|
||||
app.clone().start_mcp();
|
||||
Ok(app)
|
||||
}
|
||||
@@ -97,6 +120,81 @@ impl App {
|
||||
std::fs::canonicalize(&ws).unwrap_or(ws)
|
||||
}
|
||||
|
||||
/// Directories a session owner may work in: admins (and legacy sessions
|
||||
/// without owner) the whole workspace, users their own folder plus shared folders.
|
||||
pub fn user_roots(&self, owner_id: Option<&str>) -> Vec<PathBuf> {
|
||||
let ws = self.workspace();
|
||||
let Some(user) = owner_id.and_then(|id| self.store.user(id).ok().flatten()) else {
|
||||
return vec![ws];
|
||||
};
|
||||
if user.is_admin() {
|
||||
return vec![ws];
|
||||
}
|
||||
let own = ws.join(&user.username);
|
||||
let _ = std::fs::create_dir_all(&own);
|
||||
let mut roots = vec![own];
|
||||
for p in self.store.folders_for_user(&user.id).unwrap_or_default() {
|
||||
let p = PathBuf::from(p);
|
||||
if p.starts_with(&ws) && p.is_dir() {
|
||||
roots.push(p);
|
||||
}
|
||||
}
|
||||
roots
|
||||
}
|
||||
|
||||
/// The root that contains `cwd` (most specific first).
|
||||
pub fn root_for(&self, owner_id: Option<&str>, cwd: &Path) -> PathBuf {
|
||||
let mut roots = self.user_roots(owner_id);
|
||||
roots.sort_by_key(|r| std::cmp::Reverse(r.as_os_str().len()));
|
||||
roots.iter().find(|r| cwd.starts_with(r)).cloned().unwrap_or_else(|| roots.pop().unwrap_or_else(|| self.workspace()))
|
||||
}
|
||||
|
||||
fn owner_of(&self, session_id: &str) -> Option<(String, Option<String>)> {
|
||||
if let Some(v) = self.owners.lock().unwrap().get(session_id) {
|
||||
return Some(v.clone());
|
||||
}
|
||||
let root = self.root_session_id(session_id);
|
||||
let owner = self.store.get_session(&root).ok().flatten()?.owner_id;
|
||||
let v = (root, owner);
|
||||
self.owners.lock().unwrap().insert(session_id.to_string(), v.clone());
|
||||
Some(v)
|
||||
}
|
||||
|
||||
fn shared_with(&self, root: &str) -> Vec<String> {
|
||||
if let Some(v) = self.shares.lock().unwrap().get(root) {
|
||||
return v.clone();
|
||||
}
|
||||
let v = self.store.shares(root).unwrap_or_default();
|
||||
self.shares.lock().unwrap().insert(root.to_string(), v.clone());
|
||||
v
|
||||
}
|
||||
|
||||
pub fn set_shares(&self, root: &str, users: &[String]) -> Result<()> {
|
||||
self.store.set_shares(root, users)?;
|
||||
self.shares.lock().unwrap().insert(root.to_string(), users.to_vec());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn access(&self, user: &crate::store::User, session_id: &str) -> Access {
|
||||
let Some((root, owner)) = self.owner_of(session_id) else { return Access::None };
|
||||
if user.is_admin() || owner.as_deref() == Some(user.id.as_str()) {
|
||||
Access::Full
|
||||
} else if self.shared_with(&root).contains(&user.id) {
|
||||
Access::Read
|
||||
} else {
|
||||
Access::None
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether an event may be sent to this user's connections.
|
||||
pub fn event_visible(&self, user: &crate::store::User, ev: &Event) -> bool {
|
||||
match ev.session_id() {
|
||||
None => user.is_admin(),
|
||||
Some(_) if user.is_admin() => true,
|
||||
Some(id) => self.access(user, id) != Access::None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn data_dir(&self) -> PathBuf {
|
||||
self.config().data_dir()
|
||||
}
|
||||
@@ -197,6 +295,54 @@ impl App {
|
||||
id
|
||||
}
|
||||
|
||||
fn limit(&self, key: &str, size: usize) -> Arc<Semaphore> {
|
||||
let mut m = self.subagent_limits.lock().unwrap();
|
||||
match m.get(key) {
|
||||
Some((n, sem)) if *n == size => sem.clone(),
|
||||
_ => {
|
||||
let sem = Arc::new(Semaphore::new(size));
|
||||
m.insert(key.to_string(), (size, sem.clone()));
|
||||
sem
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wait until a subagent of this type may run (per-agent limit, then the global
|
||||
/// limit – always in this order). Returns None if cancelled while waiting, and
|
||||
/// whether the call had to wait.
|
||||
pub async fn subagent_slot(
|
||||
&self,
|
||||
agent: &str,
|
||||
max_parallel: Option<usize>,
|
||||
session_id: &str,
|
||||
cancel: &CancellationToken,
|
||||
) -> (Option<Vec<OwnedSemaphorePermit>>, bool) {
|
||||
let mut sems = Vec::new();
|
||||
if let Some(n) = max_parallel.filter(|n| *n > 0) {
|
||||
sems.push((format!("subagent:{agent}"), self.limit(&format!("agent:{agent}"), n)));
|
||||
}
|
||||
let global = self.config().server.max_parallel_subagents;
|
||||
if global > 0 {
|
||||
sems.push(("subagent:*".to_string(), self.limit("*", global)));
|
||||
}
|
||||
let mut waited = false;
|
||||
let mut permits = Vec::new();
|
||||
for (reason, sem) in sems {
|
||||
if sem.available_permits() == 0 {
|
||||
waited = true;
|
||||
self.emit(Event::Waiting { session_id: session_id.to_string(), reason: Some(reason) });
|
||||
}
|
||||
tokio::select! {
|
||||
p = sem.acquire_owned() => match p {
|
||||
Ok(p) => permits.push(p),
|
||||
Err(_) => return (None, waited),
|
||||
},
|
||||
_ = cancel.cancelled() => return (None, waited),
|
||||
}
|
||||
}
|
||||
(Some(permits), waited)
|
||||
}
|
||||
|
||||
// --- user interaction -------------------------------------------------
|
||||
|
||||
/// Forward questions to the UI and wait. None = dismissed.
|
||||
@@ -227,6 +373,12 @@ impl App {
|
||||
Ok(res)
|
||||
}
|
||||
|
||||
/// Session of a pending permission request or question.
|
||||
pub fn pending_session(&self, id: &str) -> Option<String> {
|
||||
let from_perm = self.permissions.lock().unwrap().get(id).and_then(|p| p.event.session_id().map(String::from));
|
||||
from_perm.or_else(|| self.questions.lock().unwrap().get(id).and_then(|q| q.event.session_id().map(String::from)))
|
||||
}
|
||||
|
||||
pub fn answer_question(&self, id: &str, answers: Option<Map<String, Value>>) -> Result<()> {
|
||||
let p = self.questions.lock().unwrap().remove(id).ok_or_else(|| anyhow!("question not pending"))?;
|
||||
let _ = p.tx.send(answers);
|
||||
|
||||
@@ -0,0 +1,612 @@
|
||||
//! Users, sign-in (local passwords + OpenID Connect) and access control.
|
||||
|
||||
use crate::api::{ApiError, ApiResult};
|
||||
use crate::app::App;
|
||||
use crate::config::expand_env;
|
||||
use crate::store::User;
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use argon2::{Argon2, PasswordHasher, PasswordVerifier};
|
||||
use axum::extract::{Path, Query, Request, State};
|
||||
use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Redirect, Response};
|
||||
use axum::{Extension, Json};
|
||||
use base64::Engine;
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::Arc;
|
||||
|
||||
pub const COOKIE: &str = "codecraft_session";
|
||||
const MIN_PASSWORD: usize = 8;
|
||||
|
||||
// --- primitives ---------------------------------------------------------------
|
||||
|
||||
pub fn random_token(bytes: usize) -> String {
|
||||
let mut buf = vec![0u8; bytes];
|
||||
rand::fill(&mut buf[..]);
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(buf)
|
||||
}
|
||||
|
||||
pub fn sha256_hex(s: &str) -> String {
|
||||
Sha256::digest(s.as_bytes()).iter().map(|b| format!("{b:02x}")).collect()
|
||||
}
|
||||
|
||||
pub fn hash_password(password: &str) -> Result<String> {
|
||||
let mut salt = [0u8; 16];
|
||||
rand::fill(&mut salt);
|
||||
let hash = Argon2::default()
|
||||
.hash_password_with_salt(password.as_bytes(), &salt)
|
||||
.map_err(|e| anyhow!("hashing failed: {e}"))?;
|
||||
Ok(hash.to_string())
|
||||
}
|
||||
|
||||
pub fn verify_password(password: &str, hash: &str) -> bool {
|
||||
Argon2::default().verify_password(password.as_bytes(), hash).is_ok()
|
||||
}
|
||||
|
||||
fn valid_username(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& name.len() <= 64
|
||||
&& name.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.'))
|
||||
&& !name.starts_with('.')
|
||||
}
|
||||
|
||||
/// Usernames become directory names; derive a safe one from an IdP claim.
|
||||
fn sanitize_username(raw: &str) -> String {
|
||||
let local = raw.split('@').next().unwrap_or(raw);
|
||||
let s: String = local
|
||||
.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.') { c.to_ascii_lowercase() } else { '_' })
|
||||
.collect();
|
||||
s.trim_start_matches('.').chars().take(64).collect()
|
||||
}
|
||||
|
||||
// --- cookies & middleware -------------------------------------------------------
|
||||
|
||||
fn cookie_secure(app: &App, headers: &HeaderMap) -> bool {
|
||||
match app.config().auth.cookie_secure.as_str() {
|
||||
"true" => true,
|
||||
"false" => false,
|
||||
_ => {
|
||||
app.config().auth.public_url.starts_with("https://")
|
||||
|| headers.get("x-forwarded-proto").and_then(|v| v.to_str().ok()) == Some("https")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn session_cookie(app: &App, headers: &HeaderMap, token: &str, max_age_secs: i64) -> HeaderValue {
|
||||
let secure = if cookie_secure(app, headers) { "; Secure" } else { "" };
|
||||
HeaderValue::from_str(&format!("{COOKIE}={token}; Path=/; HttpOnly; SameSite=Lax; Max-Age={max_age_secs}{secure}"))
|
||||
.expect("valid cookie")
|
||||
}
|
||||
|
||||
pub fn token_from_headers(headers: &HeaderMap) -> Option<String> {
|
||||
if let Some(t) = headers
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
{
|
||||
return Some(t.trim().to_string());
|
||||
}
|
||||
headers
|
||||
.get_all(header::COOKIE)
|
||||
.iter()
|
||||
.filter_map(|v| v.to_str().ok())
|
||||
.flat_map(|c| c.split(';'))
|
||||
.find_map(|kv| kv.trim().strip_prefix(&format!("{COOKIE}=")).map(String::from))
|
||||
}
|
||||
|
||||
pub fn current_user(app: &App, headers: &HeaderMap) -> Option<User> {
|
||||
let token = token_from_headers(headers)?;
|
||||
app.store.user_by_token(&sha256_hex(&token)).ok().flatten()
|
||||
}
|
||||
|
||||
/// Paths reachable without a login.
|
||||
fn is_public(path: &str) -> bool {
|
||||
["/info", "/auth/login", "/auth/setup", "/auth/logout", "/auth/oidc/start", "/auth/oidc/callback"]
|
||||
.iter()
|
||||
.any(|p| path.ends_with(p))
|
||||
}
|
||||
|
||||
pub async fn require_user(State(app): State<Arc<App>>, mut req: Request, next: Next) -> Response {
|
||||
if is_public(req.uri().path()) {
|
||||
return next.run(req).await;
|
||||
}
|
||||
match current_user(&app, req.headers()) {
|
||||
Some(user) => {
|
||||
req.extensions_mut().insert(user);
|
||||
next.run(req).await
|
||||
}
|
||||
None => ApiError(StatusCode::UNAUTHORIZED, "unauthorized".into()).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn require_admin(user: &User) -> Result<(), ApiError> {
|
||||
if user.is_admin() { Ok(()) } else { Err(ApiError(StatusCode::FORBIDDEN, "admin only".into())) }
|
||||
}
|
||||
|
||||
fn issue_login(app: &App, headers: &HeaderMap, user: &User) -> Result<HeaderValue> {
|
||||
let token = random_token(32);
|
||||
let days = app.config().auth.session_days.max(1) as i64;
|
||||
let ua = headers.get(header::USER_AGENT).and_then(|v| v.to_str().ok()).unwrap_or("");
|
||||
app.store.add_token(&sha256_hex(&token), &user.id, crate::types::now_ms() + days * 86_400_000, &ua.chars().take(200).collect::<String>())?;
|
||||
Ok(session_cookie(app, headers, &token, days * 86_400))
|
||||
}
|
||||
|
||||
// --- bootstrap ------------------------------------------------------------------
|
||||
|
||||
/// Create the first admin from the environment, or generate a setup code.
|
||||
pub fn bootstrap(app: &App) -> Result<()> {
|
||||
if app.store.user_count()? > 0 {
|
||||
return Ok(());
|
||||
}
|
||||
if let (Ok(user), Ok(pass)) = (std::env::var("CODECRAFT_ADMIN_USER"), std::env::var("CODECRAFT_ADMIN_PASSWORD")) {
|
||||
if valid_username(&user) && pass.len() >= MIN_PASSWORD {
|
||||
let u = app.store.create_user(&user, &user, "admin", Some(&hash_password(&pass)?), None)?;
|
||||
let n = app.store.adopt_orphans(&u.id)?;
|
||||
tracing::info!("created admin '{user}' from the environment ({n} existing sessions assigned)");
|
||||
return Ok(());
|
||||
}
|
||||
tracing::warn!("CODECRAFT_ADMIN_USER/PASSWORD ignored: invalid username or password shorter than {MIN_PASSWORD}");
|
||||
}
|
||||
let code = random_token(9);
|
||||
tracing::warn!("no users yet – open the web UI and create the first admin with setup code: {code}");
|
||||
*app.setup_code.lock().unwrap() = Some(code);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// --- endpoints: session -------------------------------------------------------
|
||||
|
||||
pub async fn info(State(app): State<Arc<App>>, headers: HeaderMap) -> Json<Value> {
|
||||
let cfg = app.config();
|
||||
let user = current_user(&app, &headers);
|
||||
Json(json!({
|
||||
"name": "CodeCraft",
|
||||
"version": env!("CARGO_PKG_VERSION"),
|
||||
"setup_required": app.setup_code.lock().unwrap().is_some(),
|
||||
"oidc": if cfg.auth.oidc.enabled { json!({"label": cfg.auth.oidc.button_label}) } else { Value::Null },
|
||||
"user": user,
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct LoginBody {
|
||||
username: String,
|
||||
password: String,
|
||||
}
|
||||
|
||||
pub async fn login(State(app): State<Arc<App>>, headers: HeaderMap, Json(body): Json<LoginBody>) -> ApiResult<Response> {
|
||||
let found = app.store.user_by_name(body.username.trim())?;
|
||||
// Verify even for unknown users so timing does not reveal which names exist.
|
||||
let hash = found.as_ref().and_then(|(_, h)| h.clone());
|
||||
let ok = verify_password(&body.password, hash.as_deref().unwrap_or("$argon2id$v=19$m=19456,t=2,p=1$c29tZXNhbHRzb21lc2FsdA$Zm9vYmFyZm9vYmFyZm9vYmFyZm9vYmFyZm9vYmFy"));
|
||||
let Some((user, _)) = found.filter(|(u, _)| ok && hash.is_some() && !u.disabled) else {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(400)).await;
|
||||
return Err(ApiError(StatusCode::UNAUTHORIZED, "invalid username or password".into()));
|
||||
};
|
||||
let cookie = issue_login(&app, &headers, &user)?;
|
||||
Ok(([(header::SET_COOKIE, cookie)], Json(json!(user))).into_response())
|
||||
}
|
||||
|
||||
pub async fn logout(State(app): State<Arc<App>>, headers: HeaderMap) -> Response {
|
||||
if let Some(t) = token_from_headers(&headers) {
|
||||
let _ = app.store.delete_token(&sha256_hex(&t));
|
||||
}
|
||||
([(header::SET_COOKIE, session_cookie(&app, &headers, "", 0))], Json(json!({"ok": true}))).into_response()
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SetupBody {
|
||||
code: String,
|
||||
username: String,
|
||||
password: String,
|
||||
}
|
||||
|
||||
pub async fn setup(State(app): State<Arc<App>>, headers: HeaderMap, Json(body): Json<SetupBody>) -> ApiResult<Response> {
|
||||
let expected = app.setup_code.lock().unwrap().clone();
|
||||
let Some(expected) = expected else {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "setup already completed".into()));
|
||||
};
|
||||
if body.code.trim() != expected {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(400)).await;
|
||||
return Err(ApiError(StatusCode::UNAUTHORIZED, "invalid setup code".into()));
|
||||
}
|
||||
let name = body.username.trim();
|
||||
if !valid_username(name) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "invalid username".into()));
|
||||
}
|
||||
if body.password.len() < MIN_PASSWORD {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, format!("password must have at least {MIN_PASSWORD} characters")));
|
||||
}
|
||||
let user = app.store.create_user(name, name, "admin", Some(&hash_password(&body.password)?), None)?;
|
||||
app.store.adopt_orphans(&user.id)?;
|
||||
*app.setup_code.lock().unwrap() = None;
|
||||
let cookie = issue_login(&app, &headers, &user)?;
|
||||
Ok(([(header::SET_COOKIE, cookie)], Json(json!(user))).into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct PasswordBody {
|
||||
current: Option<String>,
|
||||
new: String,
|
||||
}
|
||||
|
||||
/// Change the own password (OIDC-only users may set one without `current`).
|
||||
pub async fn change_password(State(app): State<Arc<App>>, Extension(user): Extension<User>, headers: HeaderMap, Json(body): Json<PasswordBody>) -> ApiResult {
|
||||
if body.new.len() < MIN_PASSWORD {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, format!("password must have at least {MIN_PASSWORD} characters")));
|
||||
}
|
||||
if let Some((_, Some(hash))) = app.store.user_by_name(&user.username)? {
|
||||
if !verify_password(body.current.as_deref().unwrap_or(""), &hash) {
|
||||
return Err(ApiError(StatusCode::UNAUTHORIZED, "current password is wrong".into()));
|
||||
}
|
||||
}
|
||||
app.store.update_user(&user.id, None, None, None, Some(&hash_password(&body.new)?), None)?;
|
||||
// Sign out all other devices.
|
||||
app.store.delete_user_tokens(&user.id)?;
|
||||
let _ = issue_login(&app, &headers, &user)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
// --- endpoints: user management (admin) ------------------------------------------
|
||||
|
||||
pub async fn list_users(Extension(me): Extension<User>, State(app): State<Arc<App>>) -> ApiResult {
|
||||
require_admin(&me)?;
|
||||
Ok(Json(json!(app.store.users()?)))
|
||||
}
|
||||
|
||||
/// Minimal user list for sharing dialogs (all signed-in users).
|
||||
pub async fn directory(State(app): State<Arc<App>>) -> ApiResult {
|
||||
let list: Vec<Value> = app
|
||||
.store
|
||||
.users()?
|
||||
.into_iter()
|
||||
.filter(|u| !u.disabled)
|
||||
.map(|u| json!({"id": u.id, "username": u.username, "display_name": u.display_name}))
|
||||
.collect();
|
||||
Ok(Json(json!(list)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct CreateUserBody {
|
||||
username: String,
|
||||
#[serde(default)]
|
||||
display_name: String,
|
||||
password: Option<String>,
|
||||
#[serde(default = "default_role")]
|
||||
role: String,
|
||||
}
|
||||
|
||||
fn default_role() -> String {
|
||||
"user".into()
|
||||
}
|
||||
|
||||
fn check_role(role: &str) -> Result<(), ApiError> {
|
||||
if role == "admin" || role == "user" { Ok(()) } else { Err(ApiError(StatusCode::BAD_REQUEST, "role must be admin or user".into())) }
|
||||
}
|
||||
|
||||
pub async fn create_user(Extension(me): Extension<User>, State(app): State<Arc<App>>, Json(body): Json<CreateUserBody>) -> ApiResult {
|
||||
require_admin(&me)?;
|
||||
let name = body.username.trim();
|
||||
if !valid_username(name) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "invalid username (letters, digits, - _ .)".into()));
|
||||
}
|
||||
check_role(&body.role)?;
|
||||
let hash = match body.password.as_deref().filter(|p| !p.is_empty()) {
|
||||
Some(p) if p.len() < MIN_PASSWORD => {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, format!("password must have at least {MIN_PASSWORD} characters")));
|
||||
}
|
||||
Some(p) => Some(hash_password(p)?),
|
||||
None => None,
|
||||
};
|
||||
if app.store.user_by_name(name)?.is_some() {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "username already exists".into()));
|
||||
}
|
||||
let display = if body.display_name.trim().is_empty() { name } else { body.display_name.trim() };
|
||||
let user = app.store.create_user(name, display, &body.role, hash.as_deref(), None)?;
|
||||
Ok(Json(json!(user)))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct UpdateUserBody {
|
||||
display_name: Option<String>,
|
||||
role: Option<String>,
|
||||
disabled: Option<bool>,
|
||||
password: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn update_user(Extension(me): Extension<User>, State(app): State<Arc<App>>, Path(id): Path<String>, Json(body): Json<UpdateUserBody>) -> ApiResult {
|
||||
require_admin(&me)?;
|
||||
let target = app.store.user(&id)?.ok_or_else(|| ApiError(StatusCode::NOT_FOUND, "user not found".into()))?;
|
||||
if let Some(r) = &body.role {
|
||||
check_role(r)?;
|
||||
}
|
||||
let loses_admin = target.is_admin() && !target.disabled && (body.role.as_deref().is_some_and(|r| r != "admin") || body.disabled == Some(true));
|
||||
if loses_admin && app.store.admin_count()? <= 1 {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "the last admin cannot be demoted or disabled".into()));
|
||||
}
|
||||
let hash = match body.password.as_deref().filter(|p| !p.is_empty()) {
|
||||
Some(p) if p.len() < MIN_PASSWORD => {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, format!("password must have at least {MIN_PASSWORD} characters")));
|
||||
}
|
||||
Some(p) => Some(hash_password(p)?),
|
||||
None => None,
|
||||
};
|
||||
app.store.update_user(&id, body.display_name.as_deref(), body.role.as_deref(), body.disabled, hash.as_deref(), None)?;
|
||||
if hash.is_some() {
|
||||
app.store.delete_user_tokens(&id)?;
|
||||
}
|
||||
Ok(Json(json!(app.store.user(&id)?)))
|
||||
}
|
||||
|
||||
pub async fn delete_user(Extension(me): Extension<User>, State(app): State<Arc<App>>, Path(id): Path<String>) -> ApiResult {
|
||||
require_admin(&me)?;
|
||||
let target = app.store.user(&id)?.ok_or_else(|| ApiError(StatusCode::NOT_FOUND, "user not found".into()))?;
|
||||
if target.id == me.id {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "you cannot delete yourself".into()));
|
||||
}
|
||||
if target.is_admin() && app.store.admin_count()? <= 1 {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "the last admin cannot be deleted".into()));
|
||||
}
|
||||
// Sessions stay (statistics); they become visible to admins only.
|
||||
app.store.delete_user(&id)?;
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
// --- OpenID Connect -----------------------------------------------------------------
|
||||
|
||||
pub struct OidcPending {
|
||||
nonce: String,
|
||||
verifier: String,
|
||||
created: std::time::Instant,
|
||||
}
|
||||
|
||||
async fn discovery(app: &App) -> Result<Value> {
|
||||
let cfg = app.config();
|
||||
let issuer = expand_env(&cfg.auth.oidc.issuer);
|
||||
if let Some((iss, doc)) = app.oidc_discovery.lock().unwrap().clone() {
|
||||
if iss == issuer {
|
||||
return Ok(doc);
|
||||
}
|
||||
}
|
||||
let url = format!("{}/.well-known/openid-configuration", issuer.trim_end_matches('/'));
|
||||
let doc: Value = crate::provider::http_client().get(&url).send().await?.error_for_status()?.json().await?;
|
||||
*app.oidc_discovery.lock().unwrap() = Some((issuer, doc.clone()));
|
||||
Ok(doc)
|
||||
}
|
||||
|
||||
fn redirect_uri(app: &App) -> Result<String> {
|
||||
let base = app.config().auth.public_url.trim_end_matches('/').to_string();
|
||||
if base.is_empty() {
|
||||
bail!("auth.public_url must be set for single sign-on");
|
||||
}
|
||||
Ok(format!("{base}/api/auth/oidc/callback"))
|
||||
}
|
||||
|
||||
fn error_page(message: &str) -> Response {
|
||||
let msg = message.replace('&', "&").replace('<', "<").replace('>', ">");
|
||||
(
|
||||
StatusCode::FORBIDDEN,
|
||||
[(header::CONTENT_TYPE, "text/html; charset=utf-8")],
|
||||
format!("<!doctype html><meta charset=utf-8><title>CodeCraft</title><body style=\"font-family:sans-serif;padding:2em\"><h2>Sign-in failed</h2><p>{msg}</p><p><a href=\"/\">Back</a></p>"),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub async fn oidc_start(State(app): State<Arc<App>>) -> Response {
|
||||
match oidc_start_inner(&app).await {
|
||||
Ok(url) => Redirect::to(&url).into_response(),
|
||||
Err(e) => error_page(&format!("{e:#}")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn oidc_start_inner(app: &App) -> Result<String> {
|
||||
let cfg = app.config();
|
||||
if !cfg.auth.oidc.enabled {
|
||||
bail!("single sign-on is not enabled");
|
||||
}
|
||||
let doc = discovery(app).await?;
|
||||
let endpoint = doc["authorization_endpoint"].as_str().context("no authorization_endpoint")?;
|
||||
let state = random_token(24);
|
||||
let nonce = random_token(24);
|
||||
let verifier = random_token(48);
|
||||
let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
|
||||
{
|
||||
let mut pending = app.oidc_pending.lock().unwrap();
|
||||
pending.retain(|_, p| p.created.elapsed().as_secs() < 600);
|
||||
pending.insert(state.clone(), OidcPending { nonce: nonce.clone(), verifier, created: std::time::Instant::now() });
|
||||
}
|
||||
let mut url = reqwest::Url::parse(endpoint)?;
|
||||
url.query_pairs_mut()
|
||||
.append_pair("response_type", "code")
|
||||
.append_pair("client_id", &expand_env(&cfg.auth.oidc.client_id))
|
||||
.append_pair("redirect_uri", &redirect_uri(app)?)
|
||||
.append_pair("scope", &cfg.auth.oidc.scopes.join(" "))
|
||||
.append_pair("state", &state)
|
||||
.append_pair("nonce", &nonce)
|
||||
.append_pair("code_challenge", &challenge)
|
||||
.append_pair("code_challenge_method", "S256");
|
||||
Ok(url.to_string())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct CallbackQuery {
|
||||
code: Option<String>,
|
||||
state: Option<String>,
|
||||
error: Option<String>,
|
||||
error_description: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn oidc_callback(State(app): State<Arc<App>>, headers: HeaderMap, Query(q): Query<CallbackQuery>) -> Response {
|
||||
if let Some(e) = q.error {
|
||||
return error_page(&format!("{e}: {}", q.error_description.unwrap_or_default()));
|
||||
}
|
||||
match oidc_callback_inner(&app, q.code.unwrap_or_default(), q.state.unwrap_or_default()).await {
|
||||
Ok(user) => match issue_login(&app, &headers, &user) {
|
||||
Ok(cookie) => ([(header::SET_COOKIE, cookie)], Redirect::to("/")).into_response(),
|
||||
Err(e) => error_page(&format!("{e:#}")),
|
||||
},
|
||||
Err(e) => {
|
||||
tracing::warn!("oidc sign-in failed: {e:#}");
|
||||
error_page(&format!("{e:#}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_jwt_claims(jwt: &str) -> Result<Value> {
|
||||
let payload = jwt.split('.').nth(1).context("malformed id_token")?;
|
||||
let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(payload.trim_end_matches('='))
|
||||
.context("malformed id_token payload")?;
|
||||
Ok(serde_json::from_slice(&bytes)?)
|
||||
}
|
||||
|
||||
fn claim_list(v: &Value) -> Vec<String> {
|
||||
match v {
|
||||
Value::Array(a) => a.iter().filter_map(|x| x.as_str().map(String::from)).collect(),
|
||||
Value::String(s) => s.split([',', ' ']).filter(|x| !x.is_empty()).map(String::from).collect(),
|
||||
_ => vec![],
|
||||
}
|
||||
}
|
||||
|
||||
async fn oidc_callback_inner(app: &App, code: String, state: String) -> Result<User> {
|
||||
let cfg = app.config();
|
||||
let oc = &cfg.auth.oidc;
|
||||
let pending = app.oidc_pending.lock().unwrap().remove(&state).context("unknown or expired sign-in attempt – please try again")?;
|
||||
let doc = discovery(app).await?;
|
||||
let token_endpoint = doc["token_endpoint"].as_str().context("no token_endpoint")?;
|
||||
let client_id = expand_env(&oc.client_id);
|
||||
let resp = crate::provider::http_client()
|
||||
.post(token_endpoint)
|
||||
.form(&[
|
||||
("grant_type", "authorization_code"),
|
||||
("code", code.as_str()),
|
||||
("redirect_uri", redirect_uri(app)?.as_str()),
|
||||
("client_id", client_id.as_str()),
|
||||
("client_secret", expand_env(&oc.client_secret).as_str()),
|
||||
("code_verifier", pending.verifier.as_str()),
|
||||
])
|
||||
.send()
|
||||
.await?;
|
||||
if !resp.status().is_success() {
|
||||
bail!("token exchange failed: HTTP {} {}", resp.status(), resp.text().await.unwrap_or_default().chars().take(300).collect::<String>());
|
||||
}
|
||||
let tokens: Value = resp.json().await?;
|
||||
// The id_token comes straight from the token endpoint over TLS, which
|
||||
// authenticates the issuer (OIDC Core 3.1.3.7); we still check its claims.
|
||||
let mut claims = decode_jwt_claims(tokens["id_token"].as_str().context("no id_token in response")?)?;
|
||||
let issuer = expand_env(&oc.issuer);
|
||||
if claims["iss"].as_str().map(|s| s.trim_end_matches('/')) != Some(issuer.trim_end_matches('/')) {
|
||||
bail!("id_token issuer mismatch");
|
||||
}
|
||||
let aud_ok = match &claims["aud"] {
|
||||
Value::String(a) => *a == client_id,
|
||||
Value::Array(a) => a.iter().any(|x| x.as_str() == Some(&client_id)),
|
||||
_ => false,
|
||||
};
|
||||
if !aud_ok {
|
||||
bail!("id_token audience mismatch");
|
||||
}
|
||||
if claims["nonce"].as_str() != Some(pending.nonce.as_str()) {
|
||||
bail!("id_token nonce mismatch");
|
||||
}
|
||||
if claims["exp"].as_i64().unwrap_or(0) < chrono::Utc::now().timestamp() {
|
||||
bail!("id_token expired");
|
||||
}
|
||||
// Groups and names are often only in the userinfo response.
|
||||
if let (Some(ep), Some(at)) = (doc["userinfo_endpoint"].as_str(), tokens["access_token"].as_str()) {
|
||||
if let Ok(r) = crate::provider::http_client().get(ep).bearer_auth(at).send().await {
|
||||
if let Ok(info) = r.json::<Value>().await {
|
||||
if info["sub"] == claims["sub"] {
|
||||
for (k, v) in info.as_object().cloned().unwrap_or_default() {
|
||||
if claims.get(&k).is_none() {
|
||||
claims[&k] = v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let sub = claims["sub"].as_str().context("no sub claim")?.to_string();
|
||||
let groups = claim_list(&claims[&oc.groups_claim]);
|
||||
if !oc.allowed_group.is_empty() && !groups.contains(&oc.allowed_group) {
|
||||
bail!("your account is not allowed to use CodeCraft");
|
||||
}
|
||||
let is_admin = !oc.admin_group.is_empty() && groups.contains(&oc.admin_group);
|
||||
let display = claims["name"].as_str().unwrap_or("").to_string();
|
||||
|
||||
let user = if let Some(u) = app.store.user_by_oidc(&sub)? {
|
||||
u
|
||||
} else {
|
||||
let raw = claims[&oc.username_claim]
|
||||
.as_str()
|
||||
.or_else(|| claims["preferred_username"].as_str())
|
||||
.or_else(|| claims["email"].as_str())
|
||||
.unwrap_or(&sub);
|
||||
let name = sanitize_username(raw);
|
||||
if !valid_username(&name) {
|
||||
bail!("cannot derive a valid username from the '{}' claim", oc.username_claim);
|
||||
}
|
||||
match app.store.user_by_name(&name)? {
|
||||
Some((existing, _)) if oc.link_by_username && existing.oidc_sub_free() => {
|
||||
app.store.update_user(&existing.id, None, None, None, None, Some(&sub))?;
|
||||
existing
|
||||
}
|
||||
Some(_) => bail!("the username '{name}' is already taken by a local account"),
|
||||
None if oc.auto_create => {
|
||||
let first = app.store.user_count()? == 0;
|
||||
let role = if is_admin || first { "admin" } else { "user" };
|
||||
let u = app.store.create_user(&name, if display.is_empty() { &name } else { &display }, role, None, Some(&sub))?;
|
||||
if first {
|
||||
app.store.adopt_orphans(&u.id)?;
|
||||
*app.setup_code.lock().unwrap() = None;
|
||||
}
|
||||
u
|
||||
}
|
||||
None => bail!("no CodeCraft account exists for '{name}'"),
|
||||
}
|
||||
};
|
||||
if user.disabled {
|
||||
bail!("this account is disabled");
|
||||
}
|
||||
// Keep the role in sync with the IdP group (never removing the last admin).
|
||||
if !oc.admin_group.is_empty() && user.is_admin() != is_admin && (is_admin || app.store.admin_count()? > 1) {
|
||||
app.store.update_user(&user.id, None, Some(if is_admin { "admin" } else { "user" }), None, None, None)?;
|
||||
}
|
||||
Ok(app.store.user(&user.id)?.context("user vanished")?)
|
||||
}
|
||||
|
||||
impl User {
|
||||
fn oidc_sub_free(&self) -> bool {
|
||||
!self.oidc
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn password_roundtrip() {
|
||||
let h = hash_password("correct horse").unwrap();
|
||||
assert!(verify_password("correct horse", &h));
|
||||
assert!(!verify_password("wrong", &h));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usernames() {
|
||||
assert!(valid_username("zink"));
|
||||
assert!(!valid_username("../etc"));
|
||||
assert!(!valid_username(".hidden"));
|
||||
assert_eq!(sanitize_username("Richard.Zink@example.org"), "richard.zink");
|
||||
assert_eq!(sanitize_username("a/b c"), "a_b_c");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jwt_claims() {
|
||||
let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(br#"{"sub":"1","aud":"x"}"#);
|
||||
let v = decode_jwt_claims(&format!("h.{payload}.s")).unwrap();
|
||||
assert_eq!(v["sub"], "1");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
//! Budgets: monthly limits tracked from recorded costs, plus the real balance
|
||||
//! for providers that expose one (OpenRouter, fal.ai), and price discovery.
|
||||
|
||||
use crate::app::App;
|
||||
use crate::config::{ProviderConfig, expand_env};
|
||||
use anyhow::{Context, Result, bail};
|
||||
use chrono::{Datelike, TimeZone};
|
||||
use serde_json::{Value, json};
|
||||
use std::time::Duration;
|
||||
|
||||
/// Start of the current month in local time (ms).
|
||||
pub fn month_start() -> i64 {
|
||||
let now = chrono::Local::now();
|
||||
chrono::Local
|
||||
.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
|
||||
.single()
|
||||
.map(|d| d.timestamp_millis())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn is_openrouter(p: &ProviderConfig) -> bool {
|
||||
p.base_url.contains("openrouter.ai")
|
||||
}
|
||||
|
||||
/// Month-to-date cost from Anthropic's cost report (organisation-wide, USD).
|
||||
async fn anthropic_cost(p: &ProviderConfig) -> Result<Value> {
|
||||
let admin = expand_env(&p.admin_key);
|
||||
let start = chrono::DateTime::from_timestamp_millis(month_start()).unwrap_or_default();
|
||||
let end = chrono::Utc::now() + chrono::Duration::days(1);
|
||||
let base = if p.base_url.is_empty() { "https://api.anthropic.com/v1".to_string() } else { expand_env(&p.base_url).trim_end_matches('/').to_string() };
|
||||
let mut total_cents = 0.0;
|
||||
let mut page: Option<String> = None;
|
||||
for _ in 0..10 {
|
||||
let mut url = reqwest::Url::parse(&format!("{base}/organizations/cost_report"))?;
|
||||
url.query_pairs_mut()
|
||||
.append_pair("starting_at", &start.format("%Y-%m-%dT00:00:00Z").to_string())
|
||||
.append_pair("ending_at", &end.format("%Y-%m-%dT00:00:00Z").to_string())
|
||||
.append_pair("limit", "31");
|
||||
if let Some(pg) = &page {
|
||||
url.query_pairs_mut().append_pair("page", pg);
|
||||
}
|
||||
let resp = crate::provider::http_client()
|
||||
.get(url)
|
||||
.header("x-api-key", &admin)
|
||||
.header("anthropic-version", "2026-10-21")
|
||||
.timeout(Duration::from_secs(20))
|
||||
.send()
|
||||
.await?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(crate::provider::http_error(resp).await);
|
||||
}
|
||||
let v: Value = resp.json().await?;
|
||||
for bucket in v["data"].as_array().cloned().unwrap_or_default() {
|
||||
for r in bucket["results"].as_array().cloned().unwrap_or_default() {
|
||||
// Amounts are decimal strings in cents.
|
||||
total_cents += r["amount"].as_str().and_then(|a| a.parse::<f64>().ok()).or_else(|| r["amount"].as_f64()).unwrap_or(0.0);
|
||||
}
|
||||
}
|
||||
if v["has_more"].as_bool() != Some(true) {
|
||||
break;
|
||||
}
|
||||
page = v["next_page"].as_str().map(String::from);
|
||||
}
|
||||
Ok(json!({"currency": "USD", "spent_month": total_cents / 100.0, "scope": "organization"}))
|
||||
}
|
||||
|
||||
/// Balance or cost as reported by the provider, if it has an API for it.
|
||||
async fn remote_balance(p: &ProviderConfig) -> Option<Result<Value>> {
|
||||
if p.kind == crate::config::ProviderKind::Anthropic && !expand_env(&p.admin_key).is_empty() {
|
||||
return Some(anthropic_cost(p).await);
|
||||
}
|
||||
if !is_openrouter(p) {
|
||||
return None;
|
||||
}
|
||||
Some(
|
||||
async {
|
||||
let base = expand_env(&p.base_url).trim_end_matches('/').to_string();
|
||||
let resp = crate::provider::http_client()
|
||||
.get(format!("{base}/key"))
|
||||
.bearer_auth(expand_env(&p.api_key))
|
||||
.timeout(Duration::from_secs(15))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
let v: Value = resp.json().await?;
|
||||
let d = &v["data"];
|
||||
Ok(json!({
|
||||
"currency": "USD",
|
||||
"limit": d["limit"],
|
||||
"used": d["usage"],
|
||||
"remaining": d.get("limit_remaining").cloned().unwrap_or_else(|| match (d["limit"].as_f64(), d["usage"].as_f64()) {
|
||||
(Some(l), Some(u)) => json!(l - u),
|
||||
_ => Value::Null,
|
||||
}),
|
||||
}))
|
||||
}
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn fal_balance(admin_key: &str) -> Result<Value> {
|
||||
let resp = crate::provider::http_client()
|
||||
.get("https://api.fal.ai/v1/account/billing?expand=credits")
|
||||
.header("Authorization", format!("Key {admin_key}"))
|
||||
.timeout(Duration::from_secs(15))
|
||||
.send()
|
||||
.await?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(crate::provider::http_error(resp).await);
|
||||
}
|
||||
let v: Value = resp.json().await?;
|
||||
Ok(json!({
|
||||
"currency": v["credits"]["currency"],
|
||||
"remaining": v["credits"]["current_balance"],
|
||||
}))
|
||||
}
|
||||
|
||||
fn entry(name: &str, budget: Option<f64>, spent: f64, warn: f64) -> Value {
|
||||
let share = budget.filter(|b| *b > 0.0).map(|b| spent / b);
|
||||
json!({
|
||||
"provider": name,
|
||||
"budget": budget,
|
||||
"spent": spent,
|
||||
"remaining": budget.map(|b| b - spent),
|
||||
"share": share,
|
||||
"warning": share.is_some_and(|s| s >= warn),
|
||||
"exceeded": share.is_some_and(|s| s >= 1.0),
|
||||
})
|
||||
}
|
||||
|
||||
/// Budget overview for the current month.
|
||||
pub async fn overview(app: &App) -> Result<Value> {
|
||||
let cfg = app.config();
|
||||
let spent = app.store.cost_by_provider(month_start())?;
|
||||
let warn = cfg.server.budget_warning;
|
||||
let mut rows = Vec::new();
|
||||
let mut tasks = Vec::new();
|
||||
for (name, p) in &cfg.providers {
|
||||
if !p.enabled {
|
||||
continue;
|
||||
}
|
||||
rows.push(entry(name, p.monthly_budget, spent.get(name).copied().unwrap_or(0.0), warn));
|
||||
let p = p.clone();
|
||||
tasks.push(tokio::spawn(async move { remote_balance(&p).await }));
|
||||
}
|
||||
for (row, task) in rows.iter_mut().zip(tasks) {
|
||||
if let Ok(Some(res)) = task.await {
|
||||
row["remote"] = match res {
|
||||
Ok(v) => v,
|
||||
Err(e) => json!({"error": format!("{e:#}")}),
|
||||
};
|
||||
}
|
||||
}
|
||||
if cfg.images.enabled {
|
||||
let mut row = entry("fal", cfg.images.monthly_budget, spent.get("fal").copied().unwrap_or(0.0), warn);
|
||||
let admin = expand_env(&cfg.images.admin_key);
|
||||
if !admin.is_empty() {
|
||||
row["remote"] = match fal_balance(&admin).await {
|
||||
Ok(v) => v,
|
||||
Err(e) => json!({"error": format!("{e:#}")}),
|
||||
};
|
||||
}
|
||||
row["images"] = json!(true);
|
||||
rows.push(row);
|
||||
}
|
||||
let total: f64 = spent.values().sum();
|
||||
let global = (cfg.server.monthly_budget > 0.0).then_some(cfg.server.monthly_budget);
|
||||
Ok(json!({
|
||||
"month_start": month_start(),
|
||||
"currency": cfg.server.currency,
|
||||
"total": entry("total", global, total, warn),
|
||||
"providers": rows,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Per-token prices from the provider's model list, converted to per 1M tokens.
|
||||
pub async fn discover_prices(app: &App, provider: &str) -> Result<Value> {
|
||||
let cfg = app.config();
|
||||
let p = cfg.providers.get(provider).context("unknown provider")?;
|
||||
if !is_openrouter(p) {
|
||||
bail!("this provider does not publish prices via its API – enter them manually");
|
||||
}
|
||||
let base = expand_env(&p.base_url).trim_end_matches('/').to_string();
|
||||
let v: Value = crate::provider::http_client()
|
||||
.get(format!("{base}/models"))
|
||||
.timeout(Duration::from_secs(30))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?
|
||||
.json()
|
||||
.await?;
|
||||
let per_m = |x: &Value| x.as_str().and_then(|s| s.parse::<f64>().ok()).or_else(|| x.as_f64()).map(|f| (f * 1_000_000.0 * 10_000.0).round() / 10_000.0);
|
||||
let list: Vec<Value> = v["data"]
|
||||
.as_array()
|
||||
.cloned()
|
||||
.unwrap_or_default()
|
||||
.iter()
|
||||
.map(|m| {
|
||||
json!({
|
||||
"id": m["id"],
|
||||
"context": m["context_length"],
|
||||
"input_price": per_m(&m["pricing"]["prompt"]),
|
||||
"output_price": per_m(&m["pricing"]["completion"]),
|
||||
"cache_read_price": per_m(&m["pricing"]["input_cache_read"]),
|
||||
"cache_write_price": per_m(&m["pricing"]["input_cache_write"]),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
Ok(json!({"currency": "USD", "models": list}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[test]
|
||||
fn budget_entry_flags() {
|
||||
let e = super::entry("x", Some(10.0), 8.5, 0.8);
|
||||
assert_eq!(e["warning"], true);
|
||||
assert_eq!(e["exceeded"], false);
|
||||
assert_eq!(e["remaining"], 1.5);
|
||||
let none = super::entry("y", None, 3.0, 0.8);
|
||||
assert_eq!(none["warning"], false);
|
||||
}
|
||||
}
|
||||
+185
-12
@@ -19,6 +19,118 @@ pub struct Config {
|
||||
pub permissions: PermissionConfig,
|
||||
pub agents: BTreeMap<String, AgentConfig>,
|
||||
pub craft: CraftConfig,
|
||||
pub auth: AuthConfig,
|
||||
pub images: ImagesConfig,
|
||||
}
|
||||
|
||||
/// Image generation via fal.ai.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct ImagesConfig {
|
||||
pub enabled: bool,
|
||||
pub base_url: String,
|
||||
pub api_key: String,
|
||||
pub default_model: String,
|
||||
/// Directory (relative to the session/project directory) for generated images.
|
||||
pub output_dir: String,
|
||||
pub models: Vec<ImageModel>,
|
||||
/// Monthly budget for image generation.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub monthly_budget: Option<f64>,
|
||||
/// Optional fal.ai Admin key to query the account balance.
|
||||
#[serde(default, skip_serializing_if = "String::is_empty")]
|
||||
pub admin_key: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, Default)]
|
||||
pub struct ImageModel {
|
||||
pub id: String,
|
||||
#[serde(default)]
|
||||
pub name: Option<String>,
|
||||
/// Price per image (statistics).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub price: Option<f64>,
|
||||
/// "image_size" (FLUX style enums) or "aspect_ratio" ("16:9" style).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub size_param: Option<String>,
|
||||
/// Extra request fields sent as-is.
|
||||
#[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
|
||||
pub options: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl Default for ImagesConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
base_url: "https://fal.run".into(),
|
||||
api_key: "${FAL_KEY}".into(),
|
||||
default_model: "fal-ai/flux/schnell".into(),
|
||||
output_dir: "generated".into(),
|
||||
models: vec![ImageModel { id: "fal-ai/flux/schnell".into(), name: Some("FLUX.1 schnell".into()), price: Some(0.003), ..Default::default() }],
|
||||
monthly_budget: None,
|
||||
admin_key: String::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct AuthConfig {
|
||||
/// Lifetime of a login in days.
|
||||
pub session_days: u32,
|
||||
/// External URL, e.g. https://codecraft.example.org (needed for OIDC).
|
||||
pub public_url: String,
|
||||
/// "auto" (secure when served via https), "true" or "false".
|
||||
pub cookie_secure: String,
|
||||
pub oidc: OidcConfig,
|
||||
}
|
||||
|
||||
impl Default for AuthConfig {
|
||||
fn default() -> Self {
|
||||
Self { session_days: 30, public_url: String::new(), cookie_secure: "auto".into(), oidc: OidcConfig::default() }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
#[serde(default)]
|
||||
pub struct OidcConfig {
|
||||
pub enabled: bool,
|
||||
pub issuer: String,
|
||||
pub client_id: String,
|
||||
pub client_secret: String,
|
||||
pub scopes: Vec<String>,
|
||||
/// Claim used as username.
|
||||
pub username_claim: String,
|
||||
pub groups_claim: String,
|
||||
/// Members of this group become admins (empty = roles are managed in CodeCraft).
|
||||
pub admin_group: String,
|
||||
/// Only members of this group may sign in (empty = everybody the IdP lets through).
|
||||
pub allowed_group: String,
|
||||
/// Create users on first login.
|
||||
pub auto_create: bool,
|
||||
/// Link an IdP login to an existing local user with the same username.
|
||||
/// Only enable if you trust the IdP's usernames.
|
||||
pub link_by_username: bool,
|
||||
pub button_label: String,
|
||||
}
|
||||
|
||||
impl Default for OidcConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
issuer: String::new(),
|
||||
client_id: String::new(),
|
||||
client_secret: String::new(),
|
||||
scopes: vec!["openid".into(), "profile".into(), "email".into(), "groups".into()],
|
||||
username_claim: "preferred_username".into(),
|
||||
groups_claim: "groups".into(),
|
||||
admin_group: String::new(),
|
||||
allowed_group: String::new(),
|
||||
auto_create: true,
|
||||
link_by_username: false,
|
||||
button_label: "Single sign-on".into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
@@ -27,12 +139,19 @@ pub struct ServerConfig {
|
||||
pub bind: String,
|
||||
pub workspace: String,
|
||||
pub data_dir: String,
|
||||
pub auth_token: String,
|
||||
pub default_model: String,
|
||||
pub small_model: String,
|
||||
pub auto_compact_threshold: f64,
|
||||
/// Currency label for cost statistics (display only).
|
||||
pub currency: String,
|
||||
/// Max. subagents running at the same time (all sessions).
|
||||
pub max_parallel_subagents: usize,
|
||||
/// MCP servers with more tools than this are loaded on demand (load = "auto").
|
||||
pub mcp_direct_limit: usize,
|
||||
/// Monthly budget for all providers together (0 = none).
|
||||
pub monthly_budget: f64,
|
||||
/// Warn when this share of a budget is used (0.0–1.0).
|
||||
pub budget_warning: f64,
|
||||
}
|
||||
|
||||
impl Default for ServerConfig {
|
||||
@@ -41,11 +160,14 @@ impl Default for ServerConfig {
|
||||
bind: "127.0.0.1:8787".into(),
|
||||
workspace: ".".into(),
|
||||
data_dir: "./data".into(),
|
||||
auth_token: String::new(),
|
||||
default_model: String::new(),
|
||||
small_model: String::new(),
|
||||
auto_compact_threshold: 0.85,
|
||||
currency: "USD".into(),
|
||||
max_parallel_subagents: 4,
|
||||
mcp_direct_limit: 15,
|
||||
monthly_budget: 0.0,
|
||||
budget_warning: 0.8,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -74,6 +196,16 @@ pub struct ProviderConfig {
|
||||
pub models: Vec<ModelEntry>,
|
||||
#[serde(default = "yes")]
|
||||
pub enabled: bool,
|
||||
/// Max. simultaneous requests to this provider/host (None = unlimited).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_concurrency: Option<usize>,
|
||||
/// Monthly budget in the configured currency (tracked from recorded costs).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub monthly_budget: Option<f64>,
|
||||
/// Optional admin key to read the provider's own cost report
|
||||
/// (Anthropic: sk-ant-admin01-…, organisation-wide).
|
||||
#[serde(default, skip_serializing_if = "String::is_empty")]
|
||||
pub admin_key: String,
|
||||
}
|
||||
|
||||
fn yes() -> bool {
|
||||
@@ -112,6 +244,12 @@ pub struct ModelConfig {
|
||||
pub cache_read_price: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cache_write_price: Option<f64>,
|
||||
/// Max. simultaneous requests to this model (None = unlimited).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_concurrency: Option<usize>,
|
||||
/// Extra request body fields sent as-is (e.g. temperature, reasoning_budget_tokens).
|
||||
#[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
|
||||
pub options: serde_json::Map<String, serde_json::Value>,
|
||||
}
|
||||
|
||||
impl ModelConfig {
|
||||
@@ -151,6 +289,9 @@ pub struct McpConfig {
|
||||
pub headers: BTreeMap<String, String>,
|
||||
#[serde(default = "yes")]
|
||||
pub enabled: bool,
|
||||
/// "auto" (default), "always" or "on_demand": how the server's tools reach the model.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub load: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
|
||||
@@ -180,6 +321,9 @@ pub struct AgentConfig {
|
||||
pub tools: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub prompt: Option<String>,
|
||||
/// Max. instances of this subagent running at the same time.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_parallel: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
@@ -224,10 +368,6 @@ impl Config {
|
||||
pub fn data_dir(&self) -> PathBuf {
|
||||
PathBuf::from(expand_env(&self.server.data_dir))
|
||||
}
|
||||
|
||||
pub fn auth_token(&self) -> String {
|
||||
expand_env(&self.server.auth_token)
|
||||
}
|
||||
}
|
||||
|
||||
/// Locate the config file: $CODECRAFT_CONFIG, else ./data/config.toml (or /data in Docker).
|
||||
@@ -273,9 +413,12 @@ pub fn apply_patches(text: &str, patches: &[ConfigPatch]) -> Result<String> {
|
||||
let (last, parents) = patch.path.split_last().unwrap();
|
||||
let mut table: &mut toml_edit::Table = doc.as_table_mut();
|
||||
for key in parents {
|
||||
let entry = table
|
||||
.entry(key)
|
||||
.or_insert_with(|| toml_edit::Item::Table(toml_edit::Table::new()));
|
||||
let entry = table.entry(key).or_insert_with(|| {
|
||||
let mut t = toml_edit::Table::new();
|
||||
// No empty "[agents]" header when only "[agents.x]" is needed.
|
||||
t.set_implicit(true);
|
||||
toml_edit::Item::Table(t)
|
||||
});
|
||||
// Inline tables are converted so we can descend into them.
|
||||
if let Some(inline) = entry.as_inline_table().cloned() {
|
||||
*entry = toml_edit::Item::Table(inline.into_table());
|
||||
@@ -327,6 +470,11 @@ fn json_to_toml(v: &serde_json::Value) -> Option<toml_edit::Value> {
|
||||
n.as_f64()?.into()
|
||||
}
|
||||
}
|
||||
J::String(s) if s.contains('\n') && !s.contains("'''") && !s.contains('\r') => {
|
||||
// Multi-line text (prompts) as a readable literal block.
|
||||
let body = if s.ends_with('\n') { s.clone() } else { format!("{s}\n") };
|
||||
format!("'''\n{body}'''").parse::<toml_edit::Value>().unwrap_or_else(|_| s.as_str().into())
|
||||
}
|
||||
J::String(s) => s.as_str().into(),
|
||||
J::Array(a) => {
|
||||
let mut arr = toml_edit::Array::new();
|
||||
@@ -369,12 +517,20 @@ pub fn redacted_json(cfg: &Config) -> serde_json::Value {
|
||||
};
|
||||
if let Some(ps) = v.get_mut("providers").and_then(|p| p.as_object_mut()) {
|
||||
for p in ps.values_mut() {
|
||||
if let Some(k) = p.get_mut("api_key") {
|
||||
mask(k);
|
||||
for field in ["api_key", "admin_key"] {
|
||||
if let Some(k) = p.get_mut(field) {
|
||||
mask(k);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(t) = v.pointer_mut("/server/auth_token") {
|
||||
if let Some(t) = v.pointer_mut("/auth/oidc/client_secret") {
|
||||
mask(t);
|
||||
}
|
||||
if let Some(t) = v.pointer_mut("/images/api_key") {
|
||||
mask(t);
|
||||
}
|
||||
if let Some(t) = v.pointer_mut("/images/admin_key") {
|
||||
mask(t);
|
||||
}
|
||||
v
|
||||
@@ -408,6 +564,23 @@ mod tests {
|
||||
assert!(cfg.providers["gemini"].enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patch_multiline_prompt_roundtrips() {
|
||||
let prompt = "You write docs.\nUse \"quotes\" and \\backslashes.";
|
||||
let out = apply_patches(
|
||||
EXAMPLE_CONFIG,
|
||||
&[ConfigPatch {
|
||||
path: vec!["agents".into(), "writer".into(), "prompt".into()],
|
||||
value: serde_json::json!(prompt),
|
||||
}],
|
||||
)
|
||||
.unwrap();
|
||||
assert!(out.contains("prompt = '''\nYou write docs."));
|
||||
assert!(!out.contains("[agents]\n"));
|
||||
let cfg = Config::parse(&out).unwrap();
|
||||
assert_eq!(cfg.agents["writer"].prompt.as_deref().unwrap().trim_end(), prompt);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patch_adds_provider_section() {
|
||||
let out = apply_patches(
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
//! File manager API: browse, edit, upload, move and delete inside the user's folders.
|
||||
|
||||
use crate::api::{ApiError, ApiResult};
|
||||
use crate::app::App;
|
||||
use crate::store::User;
|
||||
use axum::body::Bytes;
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::{StatusCode, header};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::{Extension, Json};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
type AppState = Arc<App>;
|
||||
|
||||
const MAX_TEXT: u64 = 2 * 1024 * 1024;
|
||||
|
||||
fn bad(msg: impl Into<String>) -> ApiError {
|
||||
ApiError(StatusCode::BAD_REQUEST, msg.into())
|
||||
}
|
||||
|
||||
/// Resolve a path and make sure it (or, for new files, its parent) lies inside
|
||||
/// one of the user's folders. Symlinks are resolved before the check.
|
||||
pub fn checked_path(app: &App, user: &User, raw: &str, must_exist: bool) -> Result<PathBuf, ApiError> {
|
||||
let roots = app.user_roots(Some(&user.id));
|
||||
let p = crate::tools::normalize(Path::new(raw.trim()));
|
||||
if !p.is_absolute() {
|
||||
return Err(bad("path must be absolute"));
|
||||
}
|
||||
let real = if p.exists() {
|
||||
std::fs::canonicalize(&p).map_err(|e| bad(e.to_string()))?
|
||||
} else {
|
||||
if must_exist {
|
||||
return Err(ApiError(StatusCode::NOT_FOUND, format!("{} does not exist", p.display())));
|
||||
}
|
||||
let parent = p.parent().ok_or_else(|| bad("invalid path"))?;
|
||||
let parent = std::fs::canonicalize(parent).map_err(|_| bad("parent directory does not exist"))?;
|
||||
parent.join(p.file_name().ok_or_else(|| bad("invalid path"))?)
|
||||
};
|
||||
if !roots.iter().any(|r| real.starts_with(r)) {
|
||||
return Err(ApiError(StatusCode::FORBIDDEN, "path is outside your folders".into()));
|
||||
}
|
||||
Ok(real)
|
||||
}
|
||||
|
||||
fn is_root(app: &App, user: &User, p: &Path) -> bool {
|
||||
app.user_roots(Some(&user.id)).iter().any(|r| r == p)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct PathQuery {
|
||||
path: String,
|
||||
#[serde(default)]
|
||||
hidden: bool,
|
||||
}
|
||||
|
||||
pub async fn list(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<PathQuery>) -> ApiResult {
|
||||
let dir = checked_path(&app, &user, &q.path, true)?;
|
||||
if !dir.is_dir() {
|
||||
return Err(bad("not a directory"));
|
||||
}
|
||||
let mut entries = Vec::new();
|
||||
for e in std::fs::read_dir(&dir)?.flatten() {
|
||||
let name = e.file_name().to_string_lossy().to_string();
|
||||
if !q.hidden && name.starts_with('.') {
|
||||
continue;
|
||||
}
|
||||
let meta = e.metadata().ok();
|
||||
let is_dir = e.file_type().map(|t| t.is_dir()).unwrap_or(false);
|
||||
entries.push(json!({
|
||||
"name": name,
|
||||
"path": e.path(),
|
||||
"is_dir": is_dir,
|
||||
"size": meta.as_ref().map(|m| m.len()).unwrap_or(0),
|
||||
"mtime": meta.and_then(|m| m.modified().ok()).and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()).map(|d| d.as_millis() as u64),
|
||||
}));
|
||||
}
|
||||
entries.sort_by(|a, b| {
|
||||
b["is_dir"].as_bool().cmp(&a["is_dir"].as_bool()).then_with(|| a["name"].as_str().unwrap_or("").to_lowercase().cmp(&b["name"].as_str().unwrap_or("").to_lowercase()))
|
||||
});
|
||||
let parent = if is_root(&app, &user, &dir) { Value::Null } else { json!(dir.parent()) };
|
||||
let repo = crate::git::repo_root(&dir);
|
||||
Ok(Json(json!({"path": dir, "parent": parent, "entries": entries, "repo": repo})))
|
||||
}
|
||||
|
||||
pub async fn read(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<PathQuery>) -> ApiResult {
|
||||
let p = checked_path(&app, &user, &q.path, true)?;
|
||||
let meta = std::fs::metadata(&p)?;
|
||||
if meta.is_dir() {
|
||||
return Err(bad("is a directory"));
|
||||
}
|
||||
let mtime = meta.modified().ok().and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()).map(|d| d.as_millis() as u64);
|
||||
if meta.len() > MAX_TEXT {
|
||||
return Ok(Json(json!({"path": p, "binary": true, "too_large": true, "size": meta.len(), "mtime": mtime})));
|
||||
}
|
||||
let bytes = tokio::fs::read(&p).await?;
|
||||
let binary = bytes.iter().take(8000).any(|b| *b == 0) || std::str::from_utf8(&bytes).is_err();
|
||||
Ok(Json(json!({
|
||||
"path": p,
|
||||
"binary": binary,
|
||||
"size": meta.len(),
|
||||
"mtime": mtime,
|
||||
"content": if binary { Value::Null } else { json!(String::from_utf8_lossy(&bytes)) },
|
||||
})))
|
||||
}
|
||||
|
||||
fn mime_for(p: &Path) -> &'static str {
|
||||
match p.extension().and_then(|e| e.to_str()).map(|e| e.to_ascii_lowercase()).as_deref() {
|
||||
Some("png") => "image/png",
|
||||
Some("jpg") | Some("jpeg") => "image/jpeg",
|
||||
Some("gif") => "image/gif",
|
||||
Some("webp") => "image/webp",
|
||||
Some("svg") => "image/svg+xml",
|
||||
Some("pdf") => "application/pdf",
|
||||
Some("txt") | Some("md") | Some("log") => "text/plain; charset=utf-8",
|
||||
_ => "application/octet-stream",
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct RawQuery {
|
||||
path: String,
|
||||
#[serde(default)]
|
||||
download: bool,
|
||||
}
|
||||
|
||||
/// File bytes for previews and downloads. SVGs are served as attachments only,
|
||||
/// so scripts inside them never run in the app's origin.
|
||||
pub async fn raw(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<RawQuery>) -> Result<Response, ApiError> {
|
||||
let p = checked_path(&app, &user, &q.path, true)?;
|
||||
if p.is_dir() {
|
||||
return Err(bad("is a directory"));
|
||||
}
|
||||
let bytes = tokio::fs::read(&p).await?;
|
||||
let mime = mime_for(&p);
|
||||
let name = p.file_name().and_then(|n| n.to_str()).unwrap_or("file").replace('"', "");
|
||||
let disposition = if q.download || mime == "image/svg+xml" || mime == "application/octet-stream" {
|
||||
format!("attachment; filename=\"{name}\"")
|
||||
} else {
|
||||
format!("inline; filename=\"{name}\"")
|
||||
};
|
||||
Ok((
|
||||
[
|
||||
(header::CONTENT_TYPE, mime.to_string()),
|
||||
(header::CONTENT_DISPOSITION, disposition),
|
||||
(header::X_CONTENT_TYPE_OPTIONS, "nosniff".to_string()),
|
||||
(header::CONTENT_SECURITY_POLICY, "default-src 'none'; img-src 'self' data:; style-src 'unsafe-inline'".to_string()),
|
||||
],
|
||||
bytes,
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct WriteBody {
|
||||
path: String,
|
||||
content: String,
|
||||
/// mtime the client loaded; a newer file on disk is a conflict.
|
||||
expected_mtime: Option<u64>,
|
||||
#[serde(default)]
|
||||
force: bool,
|
||||
}
|
||||
|
||||
pub async fn write(Extension(user): Extension<User>, State(app): State<AppState>, Json(b): Json<WriteBody>) -> ApiResult {
|
||||
let p = checked_path(&app, &user, &b.path, false)?;
|
||||
if p.is_dir() {
|
||||
return Err(bad("is a directory"));
|
||||
}
|
||||
if let (Some(expected), false) = (b.expected_mtime, b.force) {
|
||||
let current = std::fs::metadata(&p).ok().and_then(|m| m.modified().ok()).and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()).map(|d| d.as_millis() as u64);
|
||||
if current.is_some_and(|c| c > expected) {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "the file was changed on disk since you opened it".into()));
|
||||
}
|
||||
}
|
||||
tokio::fs::write(&p, b.content.as_bytes()).await?;
|
||||
let mtime = std::fs::metadata(&p).ok().and_then(|m| m.modified().ok()).and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()).map(|d| d.as_millis() as u64);
|
||||
Ok(Json(json!({"ok": true, "mtime": mtime})))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct CreateBody {
|
||||
path: String,
|
||||
#[serde(default)]
|
||||
dir: bool,
|
||||
}
|
||||
|
||||
pub async fn create(Extension(user): Extension<User>, State(app): State<AppState>, Json(b): Json<CreateBody>) -> ApiResult {
|
||||
let p = checked_path(&app, &user, &b.path, false)?;
|
||||
if p.exists() {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "already exists".into()));
|
||||
}
|
||||
if b.dir {
|
||||
tokio::fs::create_dir(&p).await?;
|
||||
} else {
|
||||
tokio::fs::write(&p, b"").await?;
|
||||
}
|
||||
Ok(Json(json!({"ok": true, "path": p})))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct MoveBody {
|
||||
from: String,
|
||||
to: String,
|
||||
#[serde(default)]
|
||||
copy: bool,
|
||||
}
|
||||
|
||||
fn copy_recursive(from: &Path, to: &Path) -> std::io::Result<()> {
|
||||
if from.is_dir() {
|
||||
std::fs::create_dir_all(to)?;
|
||||
for e in std::fs::read_dir(from)? {
|
||||
let e = e?;
|
||||
copy_recursive(&e.path(), &to.join(e.file_name()))?;
|
||||
}
|
||||
Ok(())
|
||||
} else {
|
||||
std::fs::copy(from, to).map(|_| ())
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn move_or_copy(Extension(user): Extension<User>, State(app): State<AppState>, Json(b): Json<MoveBody>) -> ApiResult {
|
||||
let from = checked_path(&app, &user, &b.from, true)?;
|
||||
let to = checked_path(&app, &user, &b.to, false)?;
|
||||
if is_root(&app, &user, &from) {
|
||||
return Err(bad("a root folder cannot be moved"));
|
||||
}
|
||||
if to.exists() {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "target already exists".into()));
|
||||
}
|
||||
if to.starts_with(&from) {
|
||||
return Err(bad("cannot move a folder into itself"));
|
||||
}
|
||||
let (f, t) = (from.clone(), to.clone());
|
||||
tokio::task::spawn_blocking(move || if b.copy { copy_recursive(&f, &t) } else { std::fs::rename(&f, &t) }).await??;
|
||||
Ok(Json(json!({"ok": true, "path": to})))
|
||||
}
|
||||
|
||||
pub async fn remove(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<PathQuery>) -> ApiResult {
|
||||
let p = checked_path(&app, &user, &q.path, true)?;
|
||||
if is_root(&app, &user, &p) {
|
||||
return Err(bad("a root folder cannot be deleted"));
|
||||
}
|
||||
if p.is_dir() {
|
||||
tokio::fs::remove_dir_all(&p).await?;
|
||||
} else {
|
||||
tokio::fs::remove_file(&p).await?;
|
||||
}
|
||||
Ok(Json(json!({"ok": true})))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct UploadQuery {
|
||||
/// Target file path (absolute).
|
||||
path: String,
|
||||
}
|
||||
|
||||
/// Upload one file as the raw request body.
|
||||
pub async fn upload(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<UploadQuery>, body: Bytes) -> ApiResult {
|
||||
let p = checked_path(&app, &user, &q.path, false)?;
|
||||
if p.is_dir() {
|
||||
return Err(bad("is a directory"));
|
||||
}
|
||||
tokio::fs::write(&p, &body).await?;
|
||||
Ok(Json(json!({"ok": true, "path": p, "size": body.len()})))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SearchQuery {
|
||||
path: String,
|
||||
q: String,
|
||||
}
|
||||
|
||||
/// File name search below a folder (respects .gitignore).
|
||||
pub async fn search(Extension(user): Extension<User>, State(app): State<AppState>, Query(q): Query<SearchQuery>) -> ApiResult {
|
||||
let root = checked_path(&app, &user, &q.path, true)?;
|
||||
let needle = q.q.to_lowercase();
|
||||
if needle.trim().is_empty() {
|
||||
return Ok(Json(json!([])));
|
||||
}
|
||||
let hits = tokio::task::spawn_blocking(move || {
|
||||
let mut out = Vec::new();
|
||||
for e in ignore::WalkBuilder::new(&root).hidden(true).build().flatten() {
|
||||
if e.depth() == 0 {
|
||||
continue;
|
||||
}
|
||||
let name = e.file_name().to_string_lossy().to_lowercase();
|
||||
if name.contains(&needle) {
|
||||
out.push(json!({"path": e.path(), "is_dir": e.file_type().is_some_and(|t| t.is_dir())}));
|
||||
if out.len() >= 200 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
})
|
||||
.await?;
|
||||
Ok(Json(json!(hits)))
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
//! Git client for the file manager. Runs the git CLI without a shell; file
|
||||
//! names always follow `--`, refs and URLs are validated.
|
||||
|
||||
use crate::api::{ApiError, ApiResult};
|
||||
use crate::app::App;
|
||||
use crate::files::checked_path;
|
||||
use crate::store::User;
|
||||
use axum::extract::State;
|
||||
use axum::http::StatusCode;
|
||||
use axum::{Extension, Json};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Value, json};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::process::Command;
|
||||
|
||||
type AppState = Arc<App>;
|
||||
|
||||
/// Nearest directory containing `.git`.
|
||||
pub fn repo_root(dir: &Path) -> Option<PathBuf> {
|
||||
let mut d = dir.to_path_buf();
|
||||
loop {
|
||||
if d.join(".git").exists() {
|
||||
return Some(d);
|
||||
}
|
||||
if !d.pop() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Out {
|
||||
ok: bool,
|
||||
stdout: String,
|
||||
stderr: String,
|
||||
}
|
||||
|
||||
async fn git(dir: &Path, args: &[&str], timeout: u64) -> Result<Out, ApiError> {
|
||||
let fut = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(dir)
|
||||
.env("GIT_TERMINAL_PROMPT", "0")
|
||||
.env("GIT_PAGER", "cat")
|
||||
.env("LC_ALL", "C")
|
||||
.kill_on_drop(true)
|
||||
.output();
|
||||
let out = tokio::time::timeout(Duration::from_secs(timeout), fut)
|
||||
.await
|
||||
.map_err(|_| ApiError(StatusCode::GATEWAY_TIMEOUT, "git timed out".into()))?
|
||||
.map_err(|e| ApiError(StatusCode::INTERNAL_SERVER_ERROR, format!("cannot run git: {e}")))?;
|
||||
Ok(Out {
|
||||
ok: out.status.success(),
|
||||
stdout: String::from_utf8_lossy(&out.stdout).into_owned(),
|
||||
stderr: String::from_utf8_lossy(&out.stderr).into_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
fn fail(o: Out) -> ApiError {
|
||||
let msg = if o.stderr.trim().is_empty() { o.stdout } else { o.stderr };
|
||||
ApiError(StatusCode::BAD_REQUEST, msg.trim().to_string())
|
||||
}
|
||||
|
||||
async fn git_ok(dir: &Path, args: &[&str], timeout: u64) -> Result<String, ApiError> {
|
||||
let o = git(dir, args, timeout).await?;
|
||||
if o.ok { Ok(o.stdout) } else { Err(fail(o)) }
|
||||
}
|
||||
|
||||
fn valid_ref(name: &str) -> bool {
|
||||
!name.is_empty()
|
||||
&& !name.starts_with('-')
|
||||
&& !name.contains("..")
|
||||
&& name.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '/' | '-' | '_' | '.'))
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct GitBody {
|
||||
/// Repository directory, or any path inside it.
|
||||
path: String,
|
||||
action: String,
|
||||
#[serde(default)]
|
||||
files: Vec<String>,
|
||||
#[serde(default)]
|
||||
staged: bool,
|
||||
message: Option<String>,
|
||||
#[serde(default)]
|
||||
amend: bool,
|
||||
branch: Option<String>,
|
||||
#[serde(default)]
|
||||
create: bool,
|
||||
rev: Option<String>,
|
||||
url: Option<String>,
|
||||
limit: Option<u32>,
|
||||
}
|
||||
|
||||
/// Repo-relative paths for the given absolute file paths (all checked).
|
||||
fn rel_files(app: &App, user: &User, repo: &Path, files: &[String]) -> Result<Vec<String>, ApiError> {
|
||||
files
|
||||
.iter()
|
||||
.map(|f| {
|
||||
let p = if Path::new(f).is_absolute() { PathBuf::from(f) } else { repo.join(f) };
|
||||
let p = checked_path(app, user, &p.to_string_lossy(), false)?;
|
||||
p.strip_prefix(repo)
|
||||
.map(|r| r.to_string_lossy().into_owned())
|
||||
.map_err(|_| ApiError(StatusCode::BAD_REQUEST, "file is outside the repository".into()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn parse_status(raw: &str) -> Value {
|
||||
let mut branch = Value::Null;
|
||||
let mut upstream = Value::Null;
|
||||
let (mut ahead, mut behind) = (0, 0);
|
||||
let mut files = Vec::new();
|
||||
let mut entries = raw.split('\0').filter(|s| !s.is_empty());
|
||||
while let Some(e) = entries.next() {
|
||||
if let Some(h) = e.strip_prefix("## ") {
|
||||
// "main...origin/main [ahead 1, behind 2]" or "No commits yet on main"
|
||||
let (names, info) = h.split_once(" [").map(|(a, b)| (a, b.trim_end_matches(']'))).unwrap_or((h, ""));
|
||||
let names = names.trim_start_matches("No commits yet on ");
|
||||
let (b, u) = names.split_once("...").map(|(a, b)| (a, Some(b))).unwrap_or((names, None));
|
||||
branch = json!(b);
|
||||
upstream = json!(u);
|
||||
for part in info.split(", ") {
|
||||
if let Some(n) = part.strip_prefix("ahead ") {
|
||||
ahead = n.parse().unwrap_or(0);
|
||||
}
|
||||
if let Some(n) = part.strip_prefix("behind ") {
|
||||
behind = n.parse().unwrap_or(0);
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if e.len() < 4 {
|
||||
continue;
|
||||
}
|
||||
let (x, y) = (&e[0..1], &e[1..2]);
|
||||
let path = &e[3..];
|
||||
let mut item = json!({"path": path, "index": x, "worktree": y});
|
||||
// Renames/copies carry the original path as the next entry.
|
||||
if x == "R" || x == "C" {
|
||||
item["from"] = json!(entries.next());
|
||||
}
|
||||
files.push(item);
|
||||
}
|
||||
json!({"branch": branch, "upstream": upstream, "ahead": ahead, "behind": behind, "files": files})
|
||||
}
|
||||
|
||||
pub async fn handle(Extension(user): Extension<User>, State(app): State<AppState>, Json(b): Json<GitBody>) -> ApiResult {
|
||||
let target = checked_path(&app, &user, &b.path, b.action != "clone")?;
|
||||
|
||||
// Actions that create a repository.
|
||||
match b.action.as_str() {
|
||||
"init" => {
|
||||
if !target.is_dir() {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "not a directory".into()));
|
||||
}
|
||||
let out = git_ok(&target, &["init", "-q"], 30).await?;
|
||||
return Ok(Json(json!({"ok": true, "output": out})));
|
||||
}
|
||||
"clone" => {
|
||||
let url = b.url.as_deref().unwrap_or("").trim().to_string();
|
||||
let allowed = url.starts_with("https://") || url.starts_with("http://") || url.starts_with("ssh://") || url.starts_with("git@");
|
||||
if !allowed || url.contains(char::is_whitespace) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "unsupported repository URL".into()));
|
||||
}
|
||||
if target.exists() {
|
||||
return Err(ApiError(StatusCode::CONFLICT, "target folder already exists".into()));
|
||||
}
|
||||
let parent = target.parent().unwrap_or(Path::new("/")).to_path_buf();
|
||||
let name = target.file_name().and_then(|n| n.to_str()).unwrap_or("repo").to_string();
|
||||
let o = git(&parent, &["clone", "--", &url, &name], 600).await?;
|
||||
if !o.ok {
|
||||
return Err(fail(o));
|
||||
}
|
||||
return Ok(Json(json!({"ok": true, "path": target, "output": o.stderr})));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let start = if target.is_dir() { target.clone() } else { target.parent().unwrap_or(&target).to_path_buf() };
|
||||
let repo = repo_root(&start).ok_or_else(|| ApiError(StatusCode::BAD_REQUEST, "not a git repository".into()))?;
|
||||
// The repository itself must be inside the user's folders.
|
||||
checked_path(&app, &user, &repo.to_string_lossy(), true)?;
|
||||
|
||||
let result = match b.action.as_str() {
|
||||
"status" => {
|
||||
let raw = git_ok(&repo, &["status", "--porcelain=v1", "-b", "-z", "--untracked-files=all"], 60).await?;
|
||||
let mut v = parse_status(&raw);
|
||||
v["repo"] = json!(repo);
|
||||
v
|
||||
}
|
||||
"diff" => {
|
||||
let files = rel_files(&app, &user, &repo, &b.files)?;
|
||||
let mut args = vec!["diff", "--no-color", "--no-ext-diff"];
|
||||
if b.staged {
|
||||
args.push("--cached");
|
||||
}
|
||||
args.push("--");
|
||||
args.extend(files.iter().map(String::as_str));
|
||||
let mut out = git_ok(&repo, &args, 60).await?;
|
||||
// Untracked files have no diff: show them as new.
|
||||
if out.is_empty() && !b.staged && files.len() == 1 {
|
||||
let o = git(&repo, &["diff", "--no-color", "--no-index", "--", "/dev/null", &files[0]], 60).await?;
|
||||
out = o.stdout;
|
||||
}
|
||||
out.truncate(out.char_indices().nth(600_000).map(|(i, _)| i).unwrap_or(out.len()));
|
||||
json!({"diff": out})
|
||||
}
|
||||
"stage" | "unstage" | "discard" => {
|
||||
let files = rel_files(&app, &user, &repo, &b.files)?;
|
||||
if files.is_empty() {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "no files".into()));
|
||||
}
|
||||
let mut args: Vec<&str> = match b.action.as_str() {
|
||||
"stage" => vec!["add", "-A"],
|
||||
"unstage" => vec!["restore", "--staged"],
|
||||
_ => vec!["restore", "--worktree"],
|
||||
};
|
||||
args.push("--");
|
||||
args.extend(files.iter().map(String::as_str));
|
||||
git_ok(&repo, &args, 60).await?;
|
||||
json!({"ok": true})
|
||||
}
|
||||
"commit" => {
|
||||
let msg = b.message.as_deref().unwrap_or("").trim().to_string();
|
||||
if msg.is_empty() && !b.amend {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "commit message is empty".into()));
|
||||
}
|
||||
// Use the repository's identity if configured, otherwise the CodeCraft user.
|
||||
let has_identity = git(&repo, &["config", "user.email"], 10).await.map(|o| o.ok && !o.stdout.trim().is_empty()).unwrap_or(false);
|
||||
let name = format!("user.name={}", if user.display_name.is_empty() { &user.username } else { &user.display_name });
|
||||
let email = format!("user.email={}@users.codecraft.local", user.username);
|
||||
let mut args: Vec<&str> = Vec::new();
|
||||
if !has_identity {
|
||||
args.extend(["-c", &name, "-c", &email]);
|
||||
}
|
||||
args.push("commit");
|
||||
if b.amend {
|
||||
args.push("--amend");
|
||||
if msg.is_empty() {
|
||||
args.push("--no-edit");
|
||||
}
|
||||
}
|
||||
if !msg.is_empty() {
|
||||
args.extend(["-m", &msg]);
|
||||
}
|
||||
let out = git_ok(&repo, &args, 60).await?;
|
||||
json!({"ok": true, "output": out})
|
||||
}
|
||||
"log" => {
|
||||
let n = b.limit.unwrap_or(50).min(500).to_string();
|
||||
let raw = git(&repo, &["log", "-n", &n, "--pretty=format:%H%x1f%an%x1f%at%x1f%s%x1f%D%x1e"], 60).await?;
|
||||
let commits: Vec<Value> = raw
|
||||
.stdout
|
||||
.split('\u{1e}')
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.map(|c| {
|
||||
let f: Vec<&str> = c.trim_start_matches('\n').split('\u{1f}').collect();
|
||||
json!({
|
||||
"hash": f.first(), "author": f.get(1),
|
||||
"time": f.get(2).and_then(|t| t.parse::<i64>().ok()).map(|t| t * 1000),
|
||||
"subject": f.get(3), "refs": f.get(4),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
json!({"commits": commits})
|
||||
}
|
||||
"show" => {
|
||||
let rev = b.rev.as_deref().unwrap_or("");
|
||||
if rev.len() < 4 || !rev.chars().all(|c| c.is_ascii_hexdigit()) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "invalid commit".into()));
|
||||
}
|
||||
let mut out = git_ok(&repo, &["show", "--no-color", "--stat", "--patch", rev], 60).await?;
|
||||
out.truncate(out.char_indices().nth(600_000).map(|(i, _)| i).unwrap_or(out.len()));
|
||||
json!({"diff": out})
|
||||
}
|
||||
"branches" => {
|
||||
let raw = git_ok(&repo, &["branch", "-a", "--format=%(refname:short)%1f%(HEAD)%1f%(upstream:short)"], 30).await?;
|
||||
let list: Vec<Value> = raw
|
||||
.lines()
|
||||
.filter(|l| !l.trim().is_empty())
|
||||
.map(|l| {
|
||||
let f: Vec<&str> = l.split('\u{1f}').collect();
|
||||
json!({"name": f.first(), "current": f.get(1) == Some(&"*"), "upstream": f.get(2).filter(|u| !u.is_empty())})
|
||||
})
|
||||
.collect();
|
||||
json!({"branches": list})
|
||||
}
|
||||
"checkout" => {
|
||||
let branch = b.branch.as_deref().unwrap_or("");
|
||||
if !valid_ref(branch) {
|
||||
return Err(ApiError(StatusCode::BAD_REQUEST, "invalid branch name".into()));
|
||||
}
|
||||
let args: Vec<&str> = if b.create { vec!["switch", "-c", branch] } else { vec!["switch", branch] };
|
||||
let o = git(&repo, &args, 60).await?;
|
||||
if !o.ok {
|
||||
return Err(fail(o));
|
||||
}
|
||||
json!({"ok": true, "output": o.stderr})
|
||||
}
|
||||
"fetch" | "pull" | "push" => {
|
||||
let args: Vec<&str> = match b.action.as_str() {
|
||||
"fetch" => vec!["fetch", "--prune"],
|
||||
"pull" => vec!["pull", "--ff-only"],
|
||||
_ => {
|
||||
let has_upstream = git(&repo, &["rev-parse", "--abbrev-ref", "@{u}"], 10).await.map(|o| o.ok).unwrap_or(false);
|
||||
if has_upstream { vec!["push"] } else { vec!["push", "-u", "origin", "HEAD"] }
|
||||
}
|
||||
};
|
||||
let o = git(&repo, &args, 300).await?;
|
||||
if !o.ok {
|
||||
return Err(fail(o));
|
||||
}
|
||||
json!({"ok": true, "output": format!("{}{}", o.stdout, o.stderr)})
|
||||
}
|
||||
other => return Err(ApiError(StatusCode::BAD_REQUEST, format!("unknown action '{other}'"))),
|
||||
};
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parses_porcelain_status() {
|
||||
let raw = "## main...origin/main [ahead 2, behind 1]\0 M src/a.rs\0A new.txt\0R b.rs\0old.rs\0?? x y.txt\0";
|
||||
let v = parse_status(raw);
|
||||
assert_eq!(v["branch"], "main");
|
||||
assert_eq!(v["upstream"], "origin/main");
|
||||
assert_eq!(v["ahead"], 2);
|
||||
assert_eq!(v["behind"], 1);
|
||||
let files = v["files"].as_array().unwrap();
|
||||
assert_eq!(files.len(), 4);
|
||||
assert_eq!(files[2]["from"], "old.rs");
|
||||
assert_eq!(files[3]["path"], "x y.txt");
|
||||
assert_eq!(parse_status("## No commits yet on main\0")["branch"], "main");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refs_are_validated() {
|
||||
assert!(valid_ref("feature/login-2"));
|
||||
assert!(!valid_ref("--upload-pack=x"));
|
||||
assert!(!valid_ref("a..b"));
|
||||
assert!(!valid_ref("a b"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
//! Image generation via fal.ai – as an agent tool and for the manual image page.
|
||||
|
||||
use crate::app::App;
|
||||
use crate::config::{ImageModel, expand_env};
|
||||
use crate::store::{ImageRecord, UsageRecord};
|
||||
use crate::tools::{Tool, ToolCtx, ToolOutput, arg_str};
|
||||
use crate::types::{ImageData, new_id, now_ms};
|
||||
use anyhow::{Context, Result, anyhow, bail};
|
||||
use async_trait::async_trait;
|
||||
use base64::Engine;
|
||||
use serde_json::{Value, json};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Duration;
|
||||
|
||||
/// Size presets offered in the UI and to the model (FLUX-style names).
|
||||
pub const SIZES: &[&str] = &["square_hd", "square", "landscape_4_3", "landscape_16_9", "portrait_4_3", "portrait_16_9"];
|
||||
|
||||
fn aspect_ratio(size: &str) -> &'static str {
|
||||
match size {
|
||||
"landscape_4_3" => "4:3",
|
||||
"landscape_16_9" => "16:9",
|
||||
"portrait_4_3" => "3:4",
|
||||
"portrait_16_9" => "9:16",
|
||||
_ => "1:1",
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GenerateRequest {
|
||||
pub prompt: String,
|
||||
pub model: Option<String>,
|
||||
pub size: Option<String>,
|
||||
pub count: u32,
|
||||
/// Directory the files are written to.
|
||||
pub dir: PathBuf,
|
||||
pub user_id: Option<String>,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
fn slug(s: &str) -> String {
|
||||
let mut out: String = s
|
||||
.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' })
|
||||
.collect();
|
||||
while out.contains("--") {
|
||||
out = out.replace("--", "-");
|
||||
}
|
||||
out.trim_matches('-').chars().take(40).collect::<String>().trim_end_matches('-').to_string()
|
||||
}
|
||||
|
||||
fn model_config(app: &App, id: &str) -> ImageModel {
|
||||
app.config()
|
||||
.images
|
||||
.models
|
||||
.iter()
|
||||
.find(|m| m.id == id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| ImageModel { id: id.to_string(), ..Default::default() })
|
||||
}
|
||||
|
||||
async fn download(url: &str) -> Result<(Vec<u8>, String)> {
|
||||
if let Some(rest) = url.strip_prefix("data:") {
|
||||
let (meta, data) = rest.split_once(',').context("malformed data URI")?;
|
||||
let ctype = meta.split(';').next().unwrap_or("image/png").to_string();
|
||||
return Ok((base64::engine::general_purpose::STANDARD.decode(data)?, ctype));
|
||||
}
|
||||
let resp = crate::provider::http_client().get(url).timeout(Duration::from_secs(120)).send().await?.error_for_status()?;
|
||||
let ctype = resp.headers().get(reqwest::header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).unwrap_or("image/png").to_string();
|
||||
Ok((resp.bytes().await?.to_vec(), ctype))
|
||||
}
|
||||
|
||||
/// Generate images, store them in `req.dir` and record them (history + cost).
|
||||
pub async fn generate(app: &App, req: GenerateRequest) -> Result<Vec<ImageRecord>> {
|
||||
let cfg = app.config();
|
||||
let ic = &cfg.images;
|
||||
if !ic.enabled {
|
||||
bail!("image generation is disabled – enable [images] in the settings");
|
||||
}
|
||||
let key = expand_env(&ic.api_key);
|
||||
if key.is_empty() {
|
||||
bail!("no fal.ai API key configured (images.api_key, e.g. ${{FAL_KEY}})");
|
||||
}
|
||||
let model_id = req.model.clone().filter(|m| !m.is_empty()).unwrap_or_else(|| ic.default_model.clone());
|
||||
let model = model_config(app, &model_id);
|
||||
let size = req.size.clone().filter(|s| SIZES.contains(&s.as_str())).unwrap_or_else(|| "landscape_4_3".into());
|
||||
let count = req.count.clamp(1, 4);
|
||||
|
||||
let mut body = json!({"prompt": req.prompt, "num_images": count, "output_format": "png"});
|
||||
match model.size_param.as_deref() {
|
||||
Some("aspect_ratio") => body["aspect_ratio"] = json!(aspect_ratio(&size)),
|
||||
Some("none") => {}
|
||||
_ => body["image_size"] = json!(size),
|
||||
}
|
||||
for (k, v) in &model.options {
|
||||
body[k] = v.clone();
|
||||
}
|
||||
|
||||
let url = format!("{}/{}", expand_env(&ic.base_url).trim_end_matches('/'), model_id.trim_matches('/'));
|
||||
let resp = crate::provider::http_client()
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Key {key}"))
|
||||
.timeout(Duration::from_secs(300))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.context("fal.ai request failed")?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(crate::provider::http_error(resp).await);
|
||||
}
|
||||
let out: Value = resp.json().await?;
|
||||
let images = out["images"].as_array().cloned().or_else(|| out["image"].as_object().map(|o| vec![Value::Object(o.clone())])).unwrap_or_default();
|
||||
if images.is_empty() {
|
||||
bail!("fal.ai returned no images: {}", out.to_string().chars().take(300).collect::<String>());
|
||||
}
|
||||
|
||||
tokio::fs::create_dir_all(&req.dir).await?;
|
||||
let stamp = chrono::Local::now().format("%Y%m%d-%H%M%S");
|
||||
let base = slug(&req.prompt);
|
||||
let mut records = Vec::new();
|
||||
for (i, img) in images.iter().enumerate() {
|
||||
let src = img["url"].as_str().context("image without url")?;
|
||||
let (bytes, ctype) = download(src).await?;
|
||||
let ext = match ctype.as_str() {
|
||||
"image/jpeg" => "jpg",
|
||||
"image/webp" => "webp",
|
||||
_ => "png",
|
||||
};
|
||||
let name = if images.len() > 1 { format!("{stamp}-{base}-{}.{ext}", i + 1) } else { format!("{stamp}-{base}.{ext}") };
|
||||
let path = req.dir.join(name);
|
||||
tokio::fs::write(&path, &bytes).await?;
|
||||
let rec = ImageRecord {
|
||||
id: new_id(),
|
||||
user_id: req.user_id.clone(),
|
||||
session_id: req.session_id.clone(),
|
||||
prompt: req.prompt.clone(),
|
||||
model: model_id.clone(),
|
||||
path: path.display().to_string(),
|
||||
width: img["width"].as_i64().unwrap_or(0),
|
||||
height: img["height"].as_i64().unwrap_or(0),
|
||||
created_at: now_ms(),
|
||||
};
|
||||
app.store.add_image(&rec)?;
|
||||
records.push(rec);
|
||||
}
|
||||
|
||||
// Costs go into the normal statistics (provider "fal").
|
||||
let cost = model.price.unwrap_or(0.0) * records.len() as f64;
|
||||
let session = req.session_id.clone().unwrap_or_default();
|
||||
let _ = app.store.add_usage(&UsageRecord {
|
||||
ts: now_ms(),
|
||||
session_id: session.clone(),
|
||||
root_session_id: if session.is_empty() { String::new() } else { app.root_session_id(&session) },
|
||||
provider: "fal".into(),
|
||||
model: model_id,
|
||||
agent: "image".into(),
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_read_tokens: 0,
|
||||
cache_write_tokens: 0,
|
||||
ttft_ms: 0,
|
||||
duration_ms: 0,
|
||||
cost,
|
||||
error: false,
|
||||
});
|
||||
Ok(records)
|
||||
}
|
||||
|
||||
/// Downscaled copy for the model's context.
|
||||
fn preview_for_model(path: &Path) -> Option<ImageData> {
|
||||
let img = image::open(path).ok()?;
|
||||
let img = img.resize(1024, 1024, image::imageops::FilterType::Triangle);
|
||||
let mut buf = Vec::new();
|
||||
img.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png).ok()?;
|
||||
Some(ImageData { media_type: "image/png".into(), data: base64::engine::general_purpose::STANDARD.encode(buf) })
|
||||
}
|
||||
|
||||
pub struct GenerateImage;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GenerateImage {
|
||||
fn name(&self) -> &str {
|
||||
"generate_image"
|
||||
}
|
||||
fn description(&self) -> String {
|
||||
"Generate images from a text prompt (fal.ai) and save them in the project. Use for \
|
||||
illustrations, icons, backgrounds or placeholder photos the user needs. Write a detailed \
|
||||
visual prompt in English (subject, style, composition, colours, lighting). Each image \
|
||||
costs money – generate only what is needed. Returns the file paths and a preview."
|
||||
.into()
|
||||
}
|
||||
fn schema(&self) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"size": {"type": "string", "enum": SIZES},
|
||||
"count": {"type": "integer", "minimum": 1, "maximum": 4},
|
||||
"directory": {"type": "string", "description": "Target directory relative to the working directory (default: images.output_dir)"},
|
||||
"model": {"type": "string", "description": "fal.ai model id (optional, default from the settings)"}
|
||||
},
|
||||
"required": ["prompt"]
|
||||
})
|
||||
}
|
||||
fn title(&self, input: &Value) -> String {
|
||||
format!("Generate image: {}", input["prompt"].as_str().unwrap_or("").chars().take(80).collect::<String>())
|
||||
}
|
||||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||||
let prompt = arg_str(&input, "prompt")?.to_string();
|
||||
let dir_arg = input["directory"].as_str().map(String::from).unwrap_or_else(|| ctx.app.config().images.output_dir.clone());
|
||||
let dir = ctx.resolve(&dir_arg)?;
|
||||
let owner = ctx.app.session(&ctx.session_id).ok().and_then(|s| s.owner_id);
|
||||
let records = generate(
|
||||
&ctx.app,
|
||||
GenerateRequest {
|
||||
prompt,
|
||||
model: input["model"].as_str().map(String::from),
|
||||
size: input["size"].as_str().map(String::from),
|
||||
count: input["count"].as_u64().unwrap_or(1) as u32,
|
||||
dir,
|
||||
user_id: owner,
|
||||
session_id: Some(ctx.session_id.clone()),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow!("{e:#}"))?;
|
||||
let mut text = String::from("Generated:\n");
|
||||
let mut images = Vec::new();
|
||||
for r in &records {
|
||||
let rel = Path::new(&r.path).strip_prefix(&ctx.cwd).map(|p| p.display().to_string()).unwrap_or_else(|_| r.path.clone());
|
||||
text.push_str(&format!("- {rel} ({}×{})\n", r.width, r.height));
|
||||
if let Some(p) = preview_for_model(Path::new(&r.path)) {
|
||||
images.push(p);
|
||||
}
|
||||
}
|
||||
let ids: Vec<&str> = records.iter().map(|r| r.id.as_str()).collect();
|
||||
Ok(ToolOutput { content: text, images, meta: Some(json!({"image_ids": ids})), ..Default::default() })
|
||||
}
|
||||
}
|
||||
+5
-3
@@ -2,7 +2,12 @@
|
||||
|
||||
mod api;
|
||||
mod app;
|
||||
mod auth;
|
||||
mod budget;
|
||||
mod config;
|
||||
mod files;
|
||||
mod git;
|
||||
mod images;
|
||||
mod mcp;
|
||||
mod permission;
|
||||
mod prompts;
|
||||
@@ -30,9 +35,6 @@ async fn main() -> Result<()> {
|
||||
let (cfg, _) = config::load_or_init(&path)?;
|
||||
tracing::info!("config: {}", path.display());
|
||||
tracing::info!("workspace: {}", cfg.workspace().display());
|
||||
if cfg.auth_token().is_empty() {
|
||||
tracing::warn!("auth_token is empty – the API is unprotected. Only use this on localhost.");
|
||||
}
|
||||
let bind = cfg.server.bind.clone();
|
||||
let app = app::App::new(path, cfg).await?;
|
||||
let router = api::router(app);
|
||||
|
||||
+144
-10
@@ -30,6 +30,9 @@ pub struct McpToolInfo {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub input_schema: Value,
|
||||
/// The server marks the tool as read-only (annotations.readOnlyHint).
|
||||
#[serde(default)]
|
||||
pub read_only: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
@@ -183,6 +186,7 @@ impl McpServer {
|
||||
name: t["name"].as_str().unwrap_or_default().to_string(),
|
||||
description: t["description"].as_str().unwrap_or_default().to_string(),
|
||||
input_schema: t.get("inputSchema").cloned().unwrap_or(json!({"type": "object"})),
|
||||
read_only: t["annotations"]["readOnlyHint"].as_bool().unwrap_or(false),
|
||||
});
|
||||
}
|
||||
cursor = res["nextCursor"].as_str().map(String::from);
|
||||
@@ -498,20 +502,147 @@ impl McpManager {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn tools(&self) -> Vec<Arc<dyn Tool>> {
|
||||
/// All MCP tools, split into tools given to the model directly and tools
|
||||
/// that are only reachable through mcp_search / mcp_call.
|
||||
pub fn tools(&self, cfg: &crate::config::Config) -> (Vec<Arc<dyn Tool>>, Vec<Arc<McpTool>>) {
|
||||
let servers = self.servers.lock().unwrap();
|
||||
let mut out: Vec<Arc<dyn Tool>> = Vec::new();
|
||||
let mut direct: Vec<Arc<dyn Tool>> = Vec::new();
|
||||
let mut deferred = Vec::new();
|
||||
for (sname, server) in servers.iter() {
|
||||
for t in server.tools.lock().unwrap().iter() {
|
||||
out.push(Arc::new(McpTool {
|
||||
full_name: tool_name(sname, &t.name),
|
||||
server: server.clone(),
|
||||
info: t.clone(),
|
||||
}));
|
||||
let tools = server.tools.lock().unwrap();
|
||||
let load = cfg.mcp.get(sname).and_then(|m| m.load.clone()).unwrap_or_else(|| "auto".into());
|
||||
let on_demand = match load.as_str() {
|
||||
"always" => false,
|
||||
"on_demand" => true,
|
||||
_ => tools.len() > cfg.server.mcp_direct_limit,
|
||||
};
|
||||
for t in tools.iter() {
|
||||
let tool = Arc::new(McpTool { full_name: tool_name(sname, &t.name), server: server.clone(), info: t.clone() });
|
||||
if on_demand {
|
||||
deferred.push(tool);
|
||||
} else {
|
||||
direct.push(tool);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
(direct, deferred)
|
||||
}
|
||||
|
||||
/// Short overview of on-demand servers for the system prompt.
|
||||
pub fn deferred_summary(&self, cfg: &crate::config::Config) -> Vec<(String, usize, Vec<String>)> {
|
||||
let (_, deferred) = self.tools(cfg);
|
||||
let mut by_server: BTreeMap<String, Vec<String>> = BTreeMap::new();
|
||||
for t in deferred {
|
||||
by_server.entry(t.server.name.clone()).or_default().push(t.info.name.clone());
|
||||
}
|
||||
by_server
|
||||
.into_iter()
|
||||
.map(|(s, names)| {
|
||||
let n = names.len();
|
||||
(s, n, names.into_iter().take(12).collect())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Search the on-demand MCP tools by keywords.
|
||||
pub struct McpSearch;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for McpSearch {
|
||||
fn name(&self) -> &str {
|
||||
"mcp_search"
|
||||
}
|
||||
fn description(&self) -> String {
|
||||
"Find tools of the on-demand MCP servers listed in the system prompt. Returns matching tools \
|
||||
with their input schema; call them with mcp_call. Search with a few keywords (e.g. \
|
||||
\"container restart\", \"light living room\")."
|
||||
.into()
|
||||
}
|
||||
fn schema(&self) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Keywords"},
|
||||
"server": {"type": "string", "description": "Only search this server (optional)"}
|
||||
},
|
||||
"required": ["query"]
|
||||
})
|
||||
}
|
||||
fn read_only(&self) -> bool {
|
||||
true
|
||||
}
|
||||
fn title(&self, input: &Value) -> String {
|
||||
format!("Search MCP tools: {}", input["query"].as_str().unwrap_or(""))
|
||||
}
|
||||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||||
let query = input["query"].as_str().unwrap_or("").to_lowercase();
|
||||
let server = input["server"].as_str();
|
||||
let words: Vec<&str> = query.split_whitespace().collect();
|
||||
let (_, deferred) = ctx.app.mcp.tools(&ctx.app.config());
|
||||
let mut scored: Vec<(usize, &Arc<McpTool>)> = deferred
|
||||
.iter()
|
||||
.filter(|t| server.is_none_or(|s| t.server.name == s))
|
||||
.map(|t| {
|
||||
let name = t.info.name.to_lowercase();
|
||||
let desc = t.info.description.to_lowercase();
|
||||
let score: usize = words.iter().map(|w| if name.contains(w) { 3 } else if desc.contains(w) { 1 } else { 0 }).sum();
|
||||
(score, t)
|
||||
})
|
||||
.filter(|(s, _)| *s > 0 || words.is_empty())
|
||||
.collect();
|
||||
scored.sort_by(|a, b| b.0.cmp(&a.0));
|
||||
if scored.is_empty() {
|
||||
return Ok(ToolOutput::text("No matching tools. Try other keywords or another server."));
|
||||
}
|
||||
let mut out = String::new();
|
||||
for (_, t) in scored.iter().take(8) {
|
||||
let desc: String = t.info.description.chars().take(300).collect();
|
||||
out.push_str(&format!("## {}\n{}\nInput schema: {}\n\n", t.full_name, desc, t.info.input_schema));
|
||||
}
|
||||
if scored.len() > 8 {
|
||||
out.push_str(&format!("({} more matches – refine the query)\n", scored.len() - 8));
|
||||
}
|
||||
Ok(ToolOutput::text(out))
|
||||
}
|
||||
}
|
||||
|
||||
/// Call an on-demand MCP tool found with mcp_search. Permissions are checked
|
||||
/// against the target tool by the runner.
|
||||
pub struct McpCall;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for McpCall {
|
||||
fn name(&self) -> &str {
|
||||
"mcp_call"
|
||||
}
|
||||
fn description(&self) -> String {
|
||||
"Call an MCP tool found with mcp_search, by its full name (mcp__<server>__<tool>).".into()
|
||||
}
|
||||
fn schema(&self) -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tool": {"type": "string", "description": "Full tool name, e.g. mcp__dockhand__list_containers"},
|
||||
"arguments": {"type": "object", "description": "Arguments matching the tool's input schema"}
|
||||
},
|
||||
"required": ["tool"]
|
||||
})
|
||||
}
|
||||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||||
let target = resolve_call_target(ctx, &input)?;
|
||||
target.run(input.get("arguments").cloned().unwrap_or(json!({})), ctx).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the MCP tool an mcp_call refers to.
|
||||
pub fn resolve_call_target(ctx: &ToolCtx, input: &Value) -> Result<Arc<McpTool>> {
|
||||
let name = input["tool"].as_str().unwrap_or("");
|
||||
let (_, deferred) = ctx.app.mcp.tools(&ctx.app.config());
|
||||
deferred
|
||||
.into_iter()
|
||||
.find(|t| t.full_name == name)
|
||||
.ok_or_else(|| anyhow!("unknown MCP tool '{name}' – use mcp_search to find tool names"))
|
||||
}
|
||||
|
||||
/// Tool names must match ^[a-zA-Z0-9_-]{1,64}$ for most providers.
|
||||
@@ -523,7 +654,7 @@ pub fn tool_name(server: &str, tool: &str) -> String {
|
||||
}
|
||||
|
||||
pub struct McpTool {
|
||||
full_name: String,
|
||||
pub full_name: String,
|
||||
server: Arc<McpServer>,
|
||||
info: McpToolInfo,
|
||||
}
|
||||
@@ -546,6 +677,9 @@ impl Tool for McpTool {
|
||||
fn title(&self, _input: &Value) -> String {
|
||||
format!("{} → {}", self.server.name, self.info.name)
|
||||
}
|
||||
fn read_only(&self) -> bool {
|
||||
self.info.read_only
|
||||
}
|
||||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||||
self.server.call_tool(&ctx.session_id, &self.info.name, input, ctx.cancel.clone()).await
|
||||
}
|
||||
|
||||
+117
-2
@@ -81,6 +81,70 @@ pub fn bash_prefix(cmd: &str) -> String {
|
||||
cmd.split_whitespace().take(2).collect::<Vec<_>>().join(" ")
|
||||
}
|
||||
|
||||
/// Session permission modes.
|
||||
pub const MODES: &[&str] = &["default", "auto", "accept_edits", "bypass"];
|
||||
|
||||
/// Commands that only read. Checked per simple command (after splitting).
|
||||
const SAFE_COMMANDS: &[&str] = &[
|
||||
"ls", "ll", "cat", "head", "tail", "wc", "pwd", "echo", "printf", "which", "whereis", "type", "file", "stat",
|
||||
"du", "df", "tree", "grep", "egrep", "fgrep", "rg", "ag", "sort", "uniq", "cut", "tr", "diff", "cmp", "basename",
|
||||
"dirname", "realpath", "readlink", "date", "uname", "whoami", "id", "hostname", "nproc", "free", "uptime", "ps",
|
||||
"jq", "yq", "md5sum", "sha1sum", "sha256sum", "column", "nl", "fold", "true", "false", "test", "[", "less", "more",
|
||||
];
|
||||
|
||||
const SAFE_GIT: &[&str] = &[
|
||||
"status", "diff", "log", "show", "rev-parse", "ls-files", "blame", "describe", "grep", "shortlog", "reflog", "remote",
|
||||
];
|
||||
|
||||
/// Strip harmless redirections so they don't count as file writes.
|
||||
fn strip_redirections(cmd: &str) -> String {
|
||||
cmd.replace("2>&1", "").replace("2>/dev/null", "").replace(">/dev/null", "").replace("> /dev/null", "")
|
||||
}
|
||||
|
||||
/// A simple command that only reads.
|
||||
pub fn is_safe_command(part: &str) -> bool {
|
||||
let cleaned = strip_redirections(part);
|
||||
if cleaned.contains('>') {
|
||||
return false;
|
||||
}
|
||||
let words: Vec<String> = shell_words::split(&cleaned).unwrap_or_else(|_| cleaned.split_whitespace().map(String::from).collect());
|
||||
// Leading VAR=value assignments are not safe (could change behaviour); keep it strict.
|
||||
let Some(cmd) = words.first() else { return true };
|
||||
let args: Vec<&str> = words[1..].iter().map(String::as_str).collect();
|
||||
if args.len() == 1 && matches!(args[0], "--version" | "-V" | "--help" | "-h") {
|
||||
return true;
|
||||
}
|
||||
match cmd.as_str() {
|
||||
c if SAFE_COMMANDS.contains(&c) => true,
|
||||
"find" => !args.iter().any(|a| matches!(*a, "-exec" | "-execdir" | "-delete" | "-ok" | "-okdir" | "-fprint" | "-fprintf" | "-fls")),
|
||||
"sed" => !args.iter().any(|a| a.starts_with("-i") || a.starts_with("--in-place")),
|
||||
"git" => match args.first() {
|
||||
Some(sub) if SAFE_GIT.contains(sub) => !(*sub == "remote" && args.len() > 1 && args[1] != "-v"),
|
||||
Some(&"branch") => !args.iter().any(|a| matches!(*a, "-d" | "-D" | "-m" | "-M" | "--delete" | "--move")),
|
||||
Some(&"stash") => args.get(1) == Some(&"list"),
|
||||
Some(&"tag") => args.len() == 1 || args.get(1) == Some(&"-l") || args.get(1) == Some(&"--list"),
|
||||
Some(&"config") => args.iter().any(|a| *a == "--get" || *a == "--list" || *a == "-l"),
|
||||
_ => false,
|
||||
},
|
||||
"cargo" => matches!(args.first(), Some(&"check") | Some(&"metadata") | Some(&"tree") | Some(&"--list")),
|
||||
"npm" | "pnpm" | "yarn" => matches!(args.first(), Some(&"ls") | Some(&"list") | Some(&"outdated") | Some(&"view") | Some(&"why")),
|
||||
"pip" | "pip3" => matches!(args.first(), Some(&"list") | Some(&"show") | Some(&"freeze")),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Commands that always need confirmation unless everything is allowed.
|
||||
pub fn is_dangerous(cmd: &str) -> bool {
|
||||
let c = format!(" {} ", cmd.replace('\n', " "));
|
||||
const PATTERNS: &[&str] = &[
|
||||
" rm -r", " rm -f", " rm -R", " rm --recursive", " sudo ", " su ", " git push", " git reset --hard", " git clean",
|
||||
" git checkout -- ", " git restore ", " git rebase", " git commit --amend", " chmod -R", " chown -R", " dd ",
|
||||
" mkfs", " shutdown", " reboot", " poweroff", " kill ", " pkill ", " killall ", " docker ", " docker-compose ",
|
||||
" systemctl ", " npm publish", " cargo publish", " | sh", " | bash", "|sh ", "|bash ", " > /dev/sd", " crontab ",
|
||||
];
|
||||
PATTERNS.iter().any(|p| c.contains(p))
|
||||
}
|
||||
|
||||
pub struct Evaluator<'a> {
|
||||
pub cfg: &'a PermissionConfig,
|
||||
pub mode: &'a str,
|
||||
@@ -102,7 +166,17 @@ impl Evaluator<'_> {
|
||||
if self.mode == "bypass" {
|
||||
return Decision::Allow;
|
||||
}
|
||||
return bash.unwrap_or(base);
|
||||
if let Some(d) = bash {
|
||||
return d;
|
||||
}
|
||||
// Read-only commands never need confirmation.
|
||||
if split_command(cmd).is_some_and(|parts| !parts.is_empty() && parts.iter().all(|p| is_safe_command(p))) {
|
||||
return Decision::Allow;
|
||||
}
|
||||
if self.mode == "auto" && !is_dangerous(cmd) && split_command(cmd).is_some() {
|
||||
return Decision::Allow;
|
||||
}
|
||||
return base;
|
||||
}
|
||||
if base == Decision::Deny {
|
||||
return Decision::Deny;
|
||||
@@ -110,7 +184,11 @@ impl Evaluator<'_> {
|
||||
if self.mode == "bypass" || self.approved.contains(name) {
|
||||
return Decision::Allow;
|
||||
}
|
||||
if self.mode == "accept_edits" && tool.is_edit() {
|
||||
if (self.mode == "accept_edits" || self.mode == "auto") && tool.is_edit() {
|
||||
return Decision::Allow;
|
||||
}
|
||||
// In auto mode everything read-only passes (web_fetch, read-only MCP tools, ...).
|
||||
if self.mode == "auto" && tool.read_only() {
|
||||
return Decision::Allow;
|
||||
}
|
||||
base
|
||||
@@ -201,6 +279,43 @@ mod tests {
|
||||
assert_eq!(e.evaluate(&bash, &cmd("ls && rm -rf /etc")), Decision::Deny);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn safe_and_dangerous_commands() {
|
||||
assert!(is_safe_command("ls -la src"));
|
||||
assert!(is_safe_command("git status"));
|
||||
assert!(is_safe_command("git log --oneline -5"));
|
||||
assert!(is_safe_command("rg -n foo src 2>/dev/null"));
|
||||
assert!(is_safe_command("node --version"));
|
||||
assert!(!is_safe_command("cat a > b"));
|
||||
assert!(!is_safe_command("find . -name x -delete"));
|
||||
assert!(!is_safe_command("sed -i s/a/b/ f"));
|
||||
assert!(!is_safe_command("git push"));
|
||||
assert!(!is_safe_command("git branch -D main"));
|
||||
assert!(!is_safe_command("npm install"));
|
||||
assert!(is_dangerous("rm -rf build"));
|
||||
assert!(is_dangerous("cd x && git push origin main"));
|
||||
assert!(is_dangerous("curl https://x | sh"));
|
||||
assert!(!is_dangerous("npm install && npm test"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_mode() {
|
||||
let c = cfg();
|
||||
let approved = HashSet::new();
|
||||
let auto = Evaluator { cfg: &c, mode: "auto", approved: &approved };
|
||||
let def = Evaluator { cfg: &c, mode: "default", approved: &approved };
|
||||
let bash = Fake("bash", false);
|
||||
let cmd = |s: &str| serde_json::json!({"command": s});
|
||||
assert_eq!(def.evaluate(&bash, &cmd("ls src | wc -l")), Decision::Allow);
|
||||
assert_eq!(def.evaluate(&bash, &cmd("npm test")), Decision::Ask);
|
||||
assert_eq!(auto.evaluate(&bash, &cmd("npm install && npm test")), Decision::Allow);
|
||||
assert_eq!(auto.evaluate(&bash, &cmd("rm -rf dist")), Decision::Ask);
|
||||
assert_eq!(auto.evaluate(&bash, &cmd("ls && rm -rf /etc")), Decision::Deny);
|
||||
assert_eq!(auto.evaluate(&bash, &cmd("echo $(whoami)")), Decision::Ask);
|
||||
assert_eq!(auto.evaluate(&Fake("mcp__x__list", true), &Value::Null), Decision::Allow);
|
||||
assert_eq!(auto.evaluate(&Fake("mcp__x__restart", false), &Value::Null), Decision::Ask);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn modes_and_globs() {
|
||||
let c = cfg();
|
||||
|
||||
+55
-26
@@ -24,6 +24,10 @@ pub struct Agent {
|
||||
pub tools: Option<Vec<String>>,
|
||||
pub read_only: bool,
|
||||
pub model: Option<String>,
|
||||
/// "builtin", "config" or "file".
|
||||
pub source: String,
|
||||
/// Max. instances running at the same time.
|
||||
pub max_parallel: Option<usize>,
|
||||
}
|
||||
|
||||
const BASE_PROMPT: &str = r#"You are CodeCraft, an autonomous coding agent working in the user's project through tools.
|
||||
@@ -64,11 +68,11 @@ const EXPLORE_PROMPT: &str = "You are a fast, read-only exploration subagent. Se
|
||||
pub fn builtin_agents(clean_room: bool) -> Vec<Agent> {
|
||||
let craft = CRAFT_PROMPT.replace("{CLEAN_ROOM}", if clean_room { CLEAN_ROOM_ON } else { CLEAN_ROOM_OFF });
|
||||
vec![
|
||||
Agent { name: "build".into(), description: "Full access: implement, edit, run".into(), mode: "primary".into(), prompt: BUILD_PROMPT.into(), tools: None, read_only: false, model: None },
|
||||
Agent { name: "plan".into(), description: "Read-only: analyse and propose a plan".into(), mode: "primary".into(), prompt: PLAN_PROMPT.into(), tools: None, read_only: true, model: None },
|
||||
Agent { name: "craft".into(), description: "Study a target, build your own version, check and refine".into(), mode: "primary".into(), prompt: craft, tools: None, read_only: false, model: None },
|
||||
Agent { name: "general".into(), description: "General-purpose subagent for multi-step tasks".into(), mode: "subagent".into(), prompt: GENERAL_PROMPT.into(), tools: None, read_only: false, model: None },
|
||||
Agent { name: "explore".into(), description: "Read-only subagent for fast codebase searches".into(), mode: "subagent".into(), prompt: EXPLORE_PROMPT.into(), tools: None, read_only: true, model: None },
|
||||
Agent { name: "build".into(), description: "Full access: implement, edit, run".into(), mode: "primary".into(), prompt: BUILD_PROMPT.into(), tools: None, read_only: false, model: None, source: "builtin".into(), max_parallel: None },
|
||||
Agent { name: "plan".into(), description: "Read-only: analyse and propose a plan".into(), mode: "primary".into(), prompt: PLAN_PROMPT.into(), tools: None, read_only: true, model: None, source: "builtin".into(), max_parallel: None },
|
||||
Agent { name: "craft".into(), description: "Study a target, build your own version, check and refine".into(), mode: "primary".into(), prompt: craft, tools: None, read_only: false, model: None, source: "builtin".into(), max_parallel: None },
|
||||
Agent { name: "general".into(), description: "General-purpose subagent for multi-step tasks".into(), mode: "subagent".into(), prompt: GENERAL_PROMPT.into(), tools: None, read_only: false, model: None, source: "builtin".into(), max_parallel: None },
|
||||
Agent { name: "explore".into(), description: "Read-only subagent for fast codebase searches".into(), mode: "subagent".into(), prompt: EXPLORE_PROMPT.into(), tools: None, read_only: true, model: None, source: "builtin".into(), max_parallel: None },
|
||||
]
|
||||
}
|
||||
|
||||
@@ -91,9 +95,9 @@ pub fn frontmatter(text: &str) -> (BTreeMap<String, String>, String) {
|
||||
}
|
||||
|
||||
/// Resource directories, lowest priority first: data dir, workspace, project.
|
||||
fn resource_dirs(app: &App, cwd: &Path, kind: &str) -> Vec<PathBuf> {
|
||||
fn resource_dirs(app: &App, root: &Path, cwd: &Path, kind: &str) -> Vec<PathBuf> {
|
||||
let mut dirs = vec![app.data_dir().join(kind)];
|
||||
let ws = app.workspace();
|
||||
let ws = root.to_path_buf();
|
||||
let mut chain = vec![ws.clone()];
|
||||
if cwd != ws {
|
||||
chain.push(cwd.to_path_buf());
|
||||
@@ -105,12 +109,12 @@ fn resource_dirs(app: &App, cwd: &Path, kind: &str) -> Vec<PathBuf> {
|
||||
dirs
|
||||
}
|
||||
|
||||
pub fn agents(app: &App, cwd: &Path) -> Vec<Agent> {
|
||||
pub fn agents(app: &App, root: &Path, cwd: &Path) -> Vec<Agent> {
|
||||
let cfg = app.config();
|
||||
let mut map: BTreeMap<String, Agent> =
|
||||
builtin_agents(cfg.craft.clean_room).into_iter().map(|a| (a.name.clone(), a)).collect();
|
||||
// Markdown agents (Claude Code format).
|
||||
for dir in resource_dirs(app, cwd, "agents") {
|
||||
for dir in resource_dirs(app, root, cwd, "agents") {
|
||||
let Ok(entries) = std::fs::read_dir(&dir) else { continue };
|
||||
for e in entries.flatten() {
|
||||
let p = e.path();
|
||||
@@ -129,6 +133,8 @@ pub fn agents(app: &App, cwd: &Path) -> Vec<Agent> {
|
||||
tools,
|
||||
read_only: false,
|
||||
model: meta.get("model").filter(|m| m.contains('/')).cloned(),
|
||||
source: "file".into(),
|
||||
max_parallel: meta.get("max_parallel").and_then(|v| v.parse().ok()),
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -142,23 +148,25 @@ pub fn agents(app: &App, cwd: &Path) -> Vec<Agent> {
|
||||
prompt: a.prompt.clone().or(base.as_ref().map(|b| b.prompt.clone())).unwrap_or_default(),
|
||||
tools: a.tools.clone().or(base.as_ref().and_then(|b| b.tools.clone())),
|
||||
read_only: base.as_ref().is_some_and(|b| b.read_only),
|
||||
model: a.model.clone().or(base.and_then(|b| b.model)),
|
||||
model: a.model.clone().filter(|m| !m.is_empty()).or(base.as_ref().and_then(|b| b.model.clone())),
|
||||
source: "config".into(),
|
||||
max_parallel: a.max_parallel.or(base.as_ref().and_then(|b| b.max_parallel)),
|
||||
});
|
||||
}
|
||||
map.into_values().collect()
|
||||
}
|
||||
|
||||
pub fn find_agent(app: &App, cwd: &Path, name: &str) -> Agent {
|
||||
agents(app, cwd)
|
||||
pub fn find_agent(app: &App, root: &Path, cwd: &Path, name: &str) -> Agent {
|
||||
agents(app, root, cwd)
|
||||
.into_iter()
|
||||
.find(|a| a.name == name)
|
||||
.unwrap_or_else(|| builtin_agents(true).remove(0))
|
||||
}
|
||||
|
||||
/// AGENTS.md / CLAUDE.md files from the data dir, workspace root down to cwd.
|
||||
fn instruction_files(app: &App, cwd: &Path) -> Vec<(PathBuf, String)> {
|
||||
fn instruction_files(app: &App, root: &Path, cwd: &Path) -> Vec<(PathBuf, String)> {
|
||||
let mut dirs = vec![app.data_dir()];
|
||||
let ws = app.workspace();
|
||||
let ws = root.to_path_buf();
|
||||
let mut d = cwd.to_path_buf();
|
||||
let mut chain = Vec::new();
|
||||
loop {
|
||||
@@ -183,28 +191,49 @@ fn instruction_files(app: &App, cwd: &Path) -> Vec<(PathBuf, String)> {
|
||||
out
|
||||
}
|
||||
|
||||
pub fn system_prompt(app: &App, agent: &Agent, cwd: &Path, model: &str) -> String {
|
||||
pub fn system_prompt(app: &App, agent: &Agent, root: &Path, cwd: &Path, model: &str, permission_mode: &str) -> String {
|
||||
let mut s = String::from(BASE_PROMPT);
|
||||
s.push_str("\n\n");
|
||||
s.push_str(&agent.prompt);
|
||||
if permission_mode == "auto" && !agent.read_only {
|
||||
s.push_str("\n\n# Working autonomously\nThe user does not want to approve every step. Before you start a larger change (several files, installing dependencies, migrations, anything hard to undo), confirm the plan ONCE with ask_user in a short summary. After the user agrees, carry out the whole plan without asking again – reading, editing, building and testing are pre-approved. Only stop to ask when something unexpected changes the plan. Small, obvious fixes need no confirmation. Destructive or external actions (deleting, git push, deployments, changing other systems via MCP) still trigger an approval prompt automatically.");
|
||||
}
|
||||
let git = cwd.join(".git").exists();
|
||||
s.push_str(&format!(
|
||||
"\n\n# Environment\n- Working directory: {}\n- Workspace root: {}\n- Git repository: {}\n- Platform: {} ({})\n- Model: {model}\n- Date: {}\n",
|
||||
"\n\n# Environment\n- Working directory: {}\n- Project root: {}\n- Git repository: {}\n- Platform: {} ({})\n- Model: {model}\n- Date: {}\n",
|
||||
cwd.display(),
|
||||
app.workspace().display(),
|
||||
root.display(),
|
||||
if git { "yes" } else { "no" },
|
||||
std::env::consts::OS,
|
||||
std::env::consts::ARCH,
|
||||
chrono::Local::now().format("%Y-%m-%d"),
|
||||
));
|
||||
let skills = skills(app, cwd);
|
||||
if !agent.read_only || agent.name == "plan" {
|
||||
let subs: Vec<Agent> = agents(app, root, cwd).into_iter().filter(|a| a.mode != "primary").collect();
|
||||
if !subs.is_empty() {
|
||||
s.push_str("\n# Subagents\nAvailable for the task tool (subagent_type). Delegate work that fits a subagent's description – especially writing, summarizing and translating – and use background=true for tasks whose result you do not need right away.\n");
|
||||
for a in &subs {
|
||||
let model = a.model.as_deref().map(|m| format!(" [model: {m}]")).unwrap_or_default();
|
||||
let ro = if a.read_only { " (read-only)" } else { "" };
|
||||
s.push_str(&format!("- {}{ro}{model}: {}\n", a.name, a.description));
|
||||
}
|
||||
}
|
||||
}
|
||||
let deferred = app.mcp.deferred_summary(&app.config());
|
||||
if !deferred.is_empty() {
|
||||
s.push_str("\n# On-demand MCP servers\nThese servers have many tools that are not listed directly. Find tools with mcp_search, then call them with mcp_call:\n");
|
||||
for (server, n, names) in &deferred {
|
||||
s.push_str(&format!("- {server} ({n} tools, e.g. {})\n", names.join(", ")));
|
||||
}
|
||||
}
|
||||
let skills = skills(app, root, cwd);
|
||||
if !skills.is_empty() {
|
||||
s.push_str("\n# Skills\nLoad a skill with the skill tool when the task matches its description:\n");
|
||||
for sk in &skills {
|
||||
s.push_str(&format!("- {}: {}\n", sk.name, sk.description));
|
||||
}
|
||||
}
|
||||
let files = instruction_files(app, cwd);
|
||||
let files = instruction_files(app, root, cwd);
|
||||
if !files.is_empty() {
|
||||
s.push_str("\n# Project instructions\nThese instructions come from the user's instruction files and OVERRIDE default behaviour.\n");
|
||||
for (p, text) in files {
|
||||
@@ -226,9 +255,9 @@ pub struct Command {
|
||||
pub agent: Option<String>,
|
||||
}
|
||||
|
||||
pub fn commands(app: &App, cwd: &Path) -> Vec<Command> {
|
||||
pub fn commands(app: &App, root: &Path, cwd: &Path) -> Vec<Command> {
|
||||
let mut map = BTreeMap::new();
|
||||
for dir in resource_dirs(app, cwd, "commands") {
|
||||
for dir in resource_dirs(app, root, cwd, "commands") {
|
||||
let Ok(entries) = std::fs::read_dir(&dir) else { continue };
|
||||
for e in entries.flatten() {
|
||||
let p = e.path();
|
||||
@@ -249,10 +278,10 @@ pub fn commands(app: &App, cwd: &Path) -> Vec<Command> {
|
||||
}
|
||||
|
||||
/// Expand "/name args" into the command's prompt. Returns None if not a known command.
|
||||
pub fn expand_command(app: &App, cwd: &Path, text: &str) -> Option<(String, Option<String>)> {
|
||||
pub fn expand_command(app: &App, root: &Path, cwd: &Path, text: &str) -> Option<(String, Option<String>)> {
|
||||
let rest = text.strip_prefix('/')?;
|
||||
let (name, args) = rest.split_once(char::is_whitespace).unwrap_or((rest, ""));
|
||||
let cmd = commands(app, cwd).into_iter().find(|c| c.name == name)?;
|
||||
let cmd = commands(app, root, cwd).into_iter().find(|c| c.name == name)?;
|
||||
let args = args.trim();
|
||||
let mut body = cmd.body.replace("$ARGUMENTS", args);
|
||||
for (i, a) in args.split_whitespace().enumerate() {
|
||||
@@ -274,9 +303,9 @@ pub struct Skill {
|
||||
pub path: PathBuf,
|
||||
}
|
||||
|
||||
pub fn skills(app: &App, cwd: &Path) -> Vec<Skill> {
|
||||
pub fn skills(app: &App, root: &Path, cwd: &Path) -> Vec<Skill> {
|
||||
let mut map = BTreeMap::new();
|
||||
for dir in resource_dirs(app, cwd, "skills") {
|
||||
for dir in resource_dirs(app, root, cwd, "skills") {
|
||||
let Ok(entries) = std::fs::read_dir(&dir) else { continue };
|
||||
for e in entries.flatten() {
|
||||
let p = e.path().join("SKILL.md");
|
||||
@@ -307,7 +336,7 @@ impl Tool for SkillTool {
|
||||
}
|
||||
async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
|
||||
let name = arg_str(&input, "name")?;
|
||||
let Some(sk) = skills(&ctx.app, &ctx.cwd).into_iter().find(|s| s.name == name) else {
|
||||
let Some(sk) = skills(&ctx.app, &ctx.root(), &ctx.cwd).into_iter().find(|s| s.name == name) else {
|
||||
return Ok(ToolOutput::error(format!("unknown skill '{name}'")));
|
||||
};
|
||||
let text = std::fs::read_to_string(&sk.path)?;
|
||||
|
||||
@@ -165,6 +165,9 @@ impl Provider for Anthropic {
|
||||
body["output_config"] = json!({"effort": effort});
|
||||
}
|
||||
|
||||
for (k, v) in &req.model.options {
|
||||
body[k] = v.clone();
|
||||
}
|
||||
let resp = self.request(reqwest::Method::POST, "/messages").json(&body).send().await?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(http_error(resp).await);
|
||||
|
||||
@@ -105,6 +105,13 @@ pub struct ProviderEntry {
|
||||
#[derive(Default)]
|
||||
pub struct Registry {
|
||||
pub providers: BTreeMap<String, ProviderEntry>,
|
||||
/// Concurrency limits keyed by "provider" and "provider/model".
|
||||
limits: BTreeMap<String, Arc<tokio::sync::Semaphore>>,
|
||||
}
|
||||
|
||||
/// Held while a request runs; dropping it frees the slot(s).
|
||||
pub struct Slot {
|
||||
_permits: Vec<tokio::sync::OwnedSemaphorePermit>,
|
||||
}
|
||||
|
||||
pub fn http_client() -> reqwest::Client {
|
||||
@@ -121,6 +128,7 @@ impl Registry {
|
||||
pub async fn build(cfg: &Config) -> Self {
|
||||
let client = http_client();
|
||||
let mut providers = BTreeMap::new();
|
||||
let mut limits = BTreeMap::new();
|
||||
for (name, pc) in &cfg.providers {
|
||||
if !pc.enabled {
|
||||
continue;
|
||||
@@ -131,6 +139,14 @@ impl Registry {
|
||||
ProviderKind::Openai => Arc::new(openai::OpenAi::new(client.clone(), ep)),
|
||||
};
|
||||
let mut models: Vec<ModelConfig> = pc.models.iter().map(|m| m.resolved()).collect();
|
||||
if let Some(n) = pc.max_concurrency.filter(|n| *n > 0) {
|
||||
limits.insert(name.clone(), Arc::new(tokio::sync::Semaphore::new(n)));
|
||||
}
|
||||
for m in &models {
|
||||
if let Some(n) = m.max_concurrency.filter(|n| *n > 0) {
|
||||
limits.insert(format!("{name}/{}", m.id), Arc::new(tokio::sync::Semaphore::new(n)));
|
||||
}
|
||||
}
|
||||
if models.is_empty() {
|
||||
match tokio::time::timeout(std::time::Duration::from_secs(15), provider.list_models()).await {
|
||||
Ok(Ok(m)) => models = m,
|
||||
@@ -143,7 +159,28 @@ impl Registry {
|
||||
ProviderEntry { name: name.clone(), provider, models },
|
||||
);
|
||||
}
|
||||
Self { providers }
|
||||
Self { providers, limits }
|
||||
}
|
||||
|
||||
/// Wait for a free slot on the model and its provider. Always acquired in the
|
||||
/// same order (model, then provider), so concurrent requests cannot deadlock.
|
||||
pub async fn acquire(&self, spec: &str) -> Slot {
|
||||
let provider = spec.split_once('/').map(|(p, _)| p).unwrap_or(spec);
|
||||
let mut permits = Vec::new();
|
||||
for key in [spec, provider] {
|
||||
if let Some(sem) = self.limits.get(key) {
|
||||
if let Ok(p) = sem.clone().acquire_owned().await {
|
||||
permits.push(p);
|
||||
}
|
||||
}
|
||||
}
|
||||
Slot { _permits: permits }
|
||||
}
|
||||
|
||||
/// True if a request to this model would have to wait.
|
||||
pub fn is_saturated(&self, spec: &str) -> bool {
|
||||
let provider = spec.split_once('/').map(|(p, _)| p).unwrap_or(spec);
|
||||
[spec, provider].iter().any(|k| self.limits.get(*k).is_some_and(|s| s.available_permits() == 0))
|
||||
}
|
||||
|
||||
/// Resolve "provider/model". Unknown models of a known provider are allowed
|
||||
|
||||
@@ -107,6 +107,55 @@ pub fn to_wire_messages(system: &str, messages: &[Message]) -> Vec<Value> {
|
||||
out
|
||||
}
|
||||
|
||||
/// Splits streamed content into text and inline `<think>…</think>` reasoning.
|
||||
/// Tags may be split across chunks, so a possible partial tag is held back.
|
||||
#[derive(Default)]
|
||||
pub struct ThinkSplitter {
|
||||
in_think: bool,
|
||||
pending: String,
|
||||
}
|
||||
|
||||
const OPEN: &str = "<think>";
|
||||
const CLOSE: &str = "</think>";
|
||||
|
||||
impl ThinkSplitter {
|
||||
/// Returns (is_reasoning, text) pieces that are safe to emit.
|
||||
pub fn push(&mut self, chunk: &str) -> Vec<(bool, String)> {
|
||||
self.pending.push_str(chunk);
|
||||
let mut out = Vec::new();
|
||||
loop {
|
||||
let tag = if self.in_think { CLOSE } else { OPEN };
|
||||
if let Some(pos) = self.pending.find(tag) {
|
||||
let before: String = self.pending[..pos].to_string();
|
||||
if !before.is_empty() {
|
||||
out.push((self.in_think, before));
|
||||
}
|
||||
self.pending.drain(..pos + tag.len());
|
||||
self.in_think = !self.in_think;
|
||||
continue;
|
||||
}
|
||||
// Keep a trailing prefix of the tag (e.g. "<thi") until the next chunk.
|
||||
let keep = (1..tag.len()).rev().find(|n| self.pending.ends_with(&tag[..*n])).unwrap_or(0);
|
||||
let emit_len = self.pending.len() - keep;
|
||||
if emit_len > 0 {
|
||||
let text: String = self.pending.drain(..emit_len).collect();
|
||||
out.push((self.in_think, text));
|
||||
}
|
||||
break;
|
||||
}
|
||||
// Leading whitespace after </think> is noise.
|
||||
out.retain(|(_, t)| !t.is_empty());
|
||||
out
|
||||
}
|
||||
|
||||
pub fn finish(&mut self) -> Vec<(bool, String)> {
|
||||
if self.pending.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
vec![(self.in_think, std::mem::take(&mut self.pending))]
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct PendingCall {
|
||||
id: String,
|
||||
@@ -145,6 +194,9 @@ impl Provider for OpenAi {
|
||||
body["reasoning_effort"] = json!(effort);
|
||||
}
|
||||
|
||||
for (k, v) in &req.model.options {
|
||||
body[k] = v.clone();
|
||||
}
|
||||
let rb = self.client.post(format!("{}/chat/completions", self.ep.base_url)).json(&body);
|
||||
let resp = self.ep.apply(rb).send().await?;
|
||||
if !resp.status().is_success() {
|
||||
@@ -153,6 +205,7 @@ impl Provider for OpenAi {
|
||||
|
||||
let mut events = resp.bytes_stream().eventsource();
|
||||
let mut calls: BTreeMap<u64, PendingCall> = BTreeMap::new();
|
||||
let mut think = ThinkSplitter::default();
|
||||
let mut usage = Usage::default();
|
||||
let mut stop_reason = String::from("stop");
|
||||
|
||||
@@ -182,7 +235,11 @@ impl Provider for OpenAi {
|
||||
let Some(choice) = data["choices"].get(0) else { continue };
|
||||
let d = &choice["delta"];
|
||||
if let Some(t) = d["content"].as_str().filter(|t| !t.is_empty()) {
|
||||
tx.send(StreamEvent::TextDelta(t.into())).await.ok();
|
||||
// Models served by llama.cpp & co. may inline their reasoning as <think>…</think>.
|
||||
for (is_reasoning, chunk) in think.push(t) {
|
||||
let ev = if is_reasoning { StreamEvent::ReasoningDelta(chunk) } else { StreamEvent::TextDelta(chunk) };
|
||||
tx.send(ev).await.ok();
|
||||
}
|
||||
}
|
||||
let r = d["reasoning_content"].as_str().or_else(|| d["reasoning"].as_str());
|
||||
if let Some(r) = r.filter(|r| !r.is_empty()) {
|
||||
@@ -211,6 +268,10 @@ impl Provider for OpenAi {
|
||||
}
|
||||
}
|
||||
|
||||
for (is_reasoning, chunk) in think.finish() {
|
||||
let ev = if is_reasoning { StreamEvent::ReasoningDelta(chunk) } else { StreamEvent::TextDelta(chunk) };
|
||||
tx.send(ev).await.ok();
|
||||
}
|
||||
for (i, pc) in calls {
|
||||
let input = if pc.args.trim().is_empty() {
|
||||
json!({})
|
||||
@@ -249,3 +310,34 @@ impl Provider for OpenAi {
|
||||
Ok(models)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::ThinkSplitter;
|
||||
|
||||
fn run(chunks: &[&str]) -> (String, String) {
|
||||
let mut s = ThinkSplitter::default();
|
||||
let mut text = String::new();
|
||||
let mut reasoning = String::new();
|
||||
let mut parts: Vec<(bool, String)> = chunks.iter().flat_map(|c| s.push(c)).collect();
|
||||
parts.extend(s.finish());
|
||||
for (r, t) in parts {
|
||||
if r { reasoning.push_str(&t) } else { text.push_str(&t) }
|
||||
}
|
||||
(text, reasoning)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn splits_think_tags_across_chunks() {
|
||||
let (text, reasoning) = run(&["<thi", "nk>plan ", "the fix</th", "ink>Here is ", "the answer."]);
|
||||
assert_eq!(reasoning, "plan the fix");
|
||||
assert_eq!(text, "Here is the answer.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plain_text_and_lookalikes_pass_through() {
|
||||
let (text, reasoning) = run(&["a < b and <thin", "g> stays"]);
|
||||
assert_eq!(text, "a < b and <thing> stays");
|
||||
assert!(reasoning.is_empty());
|
||||
}
|
||||
}
|
||||
+299
-66
@@ -20,30 +20,55 @@ use tokio_util::sync::CancellationToken;
|
||||
const MAX_STEPS: usize = 250;
|
||||
const MAX_SUBAGENT_DEPTH: usize = 3;
|
||||
|
||||
/// Start processing a user prompt. If the session is busy the prompt is queued
|
||||
/// and injected at the next step.
|
||||
/// What happened to a prompt.
|
||||
pub enum Submitted {
|
||||
Started,
|
||||
Queued(crate::store::QueueItem),
|
||||
}
|
||||
|
||||
fn emit_queue(app: &App, session_id: &str) {
|
||||
let items = app.store.queue_list(session_id).unwrap_or_default();
|
||||
app.emit(Event::QueueUpdated { session_id: session_id.into(), items: serde_json::json!(items) });
|
||||
}
|
||||
|
||||
/// Take all queued messages (one user message's worth of parts).
|
||||
fn take_queue(app: &App, session_id: &str) -> Vec<Part> {
|
||||
let items = app.store.queue_take(session_id).unwrap_or_default();
|
||||
if !items.is_empty() {
|
||||
emit_queue(app, session_id);
|
||||
}
|
||||
items.into_iter().flat_map(|i| i.parts).collect()
|
||||
}
|
||||
|
||||
/// Start processing a user prompt. If the session is busy, the prompt (with
|
||||
/// attachments) is queued and delivered at the agent's next step.
|
||||
pub fn start_prompt(app: Arc<App>, session_id: String, parts: Vec<Part>) -> Result<()> {
|
||||
submit(app, session_id, parts, "user").map(|_| ())
|
||||
}
|
||||
|
||||
pub fn submit(app: Arc<App>, session_id: String, parts: Vec<Part>, source: &str) -> Result<Submitted> {
|
||||
// Queueing and the end-of-run check run under the same runtime lock, so a
|
||||
// message is never lost or delivered twice.
|
||||
let cancel = CancellationToken::new();
|
||||
let queued = app.with_runtime(&session_id, |r| -> Result<Option<crate::store::QueueItem>> {
|
||||
if r.cancel.is_some() {
|
||||
let item = crate::store::QueueItem { id: new_id(), parts: parts.clone(), source: source.into(), created_at: now_ms() };
|
||||
app.store.queue_add(&session_id, &item)?;
|
||||
Ok(Some(item))
|
||||
} else {
|
||||
r.cancel = Some(cancel.clone());
|
||||
Ok(None)
|
||||
}
|
||||
})?;
|
||||
if let Some(item) = queued {
|
||||
emit_queue(&app, &session_id);
|
||||
return Ok(Submitted::Queued(item));
|
||||
}
|
||||
let text_only: String = parts
|
||||
.iter()
|
||||
.filter_map(|p| if let Part::Text { text } = p { Some(text.as_str()) } else { None })
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
let queued = app.with_runtime(&session_id, |r| {
|
||||
if r.cancel.is_some() && parts.iter().all(|p| matches!(p, Part::Text { .. })) {
|
||||
r.queue.push(text_only.clone());
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
});
|
||||
if queued {
|
||||
return Ok(());
|
||||
}
|
||||
if app.is_busy(&session_id) {
|
||||
bail!("session is busy");
|
||||
}
|
||||
let cancel = CancellationToken::new();
|
||||
app.with_runtime(&session_id, |r| r.cancel = Some(cancel.clone()));
|
||||
app.emit(Event::SessionStatus { session_id: session_id.clone(), busy: true });
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -59,17 +84,31 @@ pub fn start_prompt(app: Arc<App>, session_id: String, parts: Vec<Part>) -> Resu
|
||||
tracing::warn!("session {session_id}: {e:#}");
|
||||
app.emit(Event::Error { session_id: session_id.clone(), message: format!("{e:#}") });
|
||||
}
|
||||
// Prompts queued after the last step start a new turn.
|
||||
// Messages queued after the last step start a new turn – unless the user
|
||||
// stopped the agent: then they stay queued until sent explicitly.
|
||||
let leftover = app.with_runtime(&session_id, |r| {
|
||||
r.cancel = None;
|
||||
std::mem::take(&mut r.queue)
|
||||
if cancel.is_cancelled() { Vec::new() } else { take_queue(&app, &session_id) }
|
||||
});
|
||||
app.emit(Event::SessionStatus { session_id: session_id.clone(), busy: false });
|
||||
if !leftover.is_empty() && !cancel.is_cancelled() {
|
||||
let _ = start_prompt(app, session_id, vec![Part::Text { text: leftover.join("\n\n") }]);
|
||||
if !leftover.is_empty() {
|
||||
let _ = start_prompt(app, session_id, leftover);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
Ok(Submitted::Started)
|
||||
}
|
||||
|
||||
/// Send queued messages now (after a stop or a server restart).
|
||||
pub fn flush_queue(app: Arc<App>, session_id: &str) -> Result<bool> {
|
||||
if app.is_busy(session_id) {
|
||||
return Ok(false);
|
||||
}
|
||||
let parts = take_queue(&app, session_id);
|
||||
if parts.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
start_prompt(app, session_id.to_string(), parts)?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub fn abort(app: &App, session_id: &str) {
|
||||
@@ -77,17 +116,72 @@ pub fn abort(app: &App, session_id: &str) {
|
||||
if let Some(c) = &r.cancel {
|
||||
c.cancel();
|
||||
}
|
||||
r.queue.clear();
|
||||
});
|
||||
}
|
||||
|
||||
/// Messages since the last compaction summary.
|
||||
/// Messages since the last compaction summary. Tool calls left without a
|
||||
/// result (run interrupted, e.g. by a server restart) get a synthetic result:
|
||||
/// providers reject histories with unanswered tool calls.
|
||||
fn active_history(all: Vec<Message>) -> Vec<Message> {
|
||||
let start = all
|
||||
.iter()
|
||||
.rposition(|m| m.parts.iter().any(|p| matches!(p, Part::Summary { .. })))
|
||||
.unwrap_or(0);
|
||||
all.into_iter().skip(start).collect()
|
||||
let msgs: Vec<Message> = all.into_iter().skip(start).collect();
|
||||
let mut out: Vec<Message> = Vec::with_capacity(msgs.len());
|
||||
for (i, m) in msgs.iter().enumerate() {
|
||||
out.push(m.clone());
|
||||
if m.role != Role::Assistant {
|
||||
continue;
|
||||
}
|
||||
let answered: std::collections::HashSet<&str> = msgs
|
||||
.get(i + 1)
|
||||
.filter(|n| n.role == Role::User)
|
||||
.map(|n| {
|
||||
n.parts
|
||||
.iter()
|
||||
.filter_map(|p| if let Part::ToolResult { call_id, .. } = p { Some(call_id.as_str()) } else { None })
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let missing: Vec<Part> = m
|
||||
.parts
|
||||
.iter()
|
||||
.filter_map(|p| match p {
|
||||
Part::ToolCall { id, name, .. } if !answered.contains(id.as_str()) => {
|
||||
Some(tool_result(id, name, "Not executed: the run was interrupted (e.g. server restart).".into(), true, None))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
if missing.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// Results must directly follow the call: merge into the next user message or insert one.
|
||||
match msgs.get(i + 1) {
|
||||
Some(n) if n.role == Role::User => {
|
||||
// The next message is pushed in the next iteration; prepend there.
|
||||
let mut next = n.clone();
|
||||
let mut parts = missing;
|
||||
parts.extend(next.parts);
|
||||
next.parts = parts;
|
||||
// Replace by pushing now and skipping the original.
|
||||
out.push(next);
|
||||
}
|
||||
_ => out.push(Message::new(Role::User, missing)),
|
||||
}
|
||||
}
|
||||
// Remove duplicates created by the merge above (same id pushed twice: keep the merged one).
|
||||
let mut dedup: Vec<Message> = Vec::with_capacity(out.len());
|
||||
for m in out {
|
||||
if let Some(last) = dedup.last() {
|
||||
if last.id == m.id {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
dedup.push(m);
|
||||
}
|
||||
dedup
|
||||
}
|
||||
|
||||
fn normalize_tool_name(s: &str) -> String {
|
||||
@@ -95,8 +189,17 @@ fn normalize_tool_name(s: &str) -> String {
|
||||
}
|
||||
|
||||
pub fn available_tools(app: &App, agent: &Agent) -> Vec<Arc<dyn Tool>> {
|
||||
let cfg = app.config();
|
||||
let mut all = tools::builtin_tools();
|
||||
all.extend(app.mcp.tools());
|
||||
if !cfg.images.enabled {
|
||||
all.retain(|t| t.name() != "generate_image");
|
||||
}
|
||||
let (direct, deferred) = app.mcp.tools(&cfg);
|
||||
all.extend(direct);
|
||||
if !deferred.is_empty() {
|
||||
all.push(Arc::new(crate::mcp::McpSearch));
|
||||
all.push(Arc::new(crate::mcp::McpCall));
|
||||
}
|
||||
all.into_iter()
|
||||
.filter(|t| !agent.read_only || t.read_only())
|
||||
// Craft tools only where they are useful.
|
||||
@@ -107,7 +210,7 @@ pub fn available_tools(app: &App, agent: &Agent) -> Vec<Arc<dyn Tool>> {
|
||||
let n = normalize_tool_name(t.name());
|
||||
list.iter().any(|a| {
|
||||
let a = normalize_tool_name(a);
|
||||
a == n || (t.name().starts_with("mcp__") && a == "mcp")
|
||||
a == n || (t.name().starts_with("mcp") && a == "mcp")
|
||||
})
|
||||
}
|
||||
})
|
||||
@@ -119,6 +222,17 @@ struct Step {
|
||||
stop_reason: String,
|
||||
}
|
||||
|
||||
/// Boxed with an explicit `Send` bound: subagents recurse into the agent loop,
|
||||
/// and the declared type breaks the cycle for the compiler.
|
||||
fn run_session_boxed(
|
||||
app: Arc<App>,
|
||||
session_id: String,
|
||||
cancel: CancellationToken,
|
||||
depth: usize,
|
||||
) -> futures::future::BoxFuture<'static, Result<String>> {
|
||||
Box::pin(async move { run_session(app, &session_id, cancel, depth).await })
|
||||
}
|
||||
|
||||
/// Run the loop for a session until the model stops. Returns the final text.
|
||||
pub async fn run_session(app: Arc<App>, session_id: &str, cancel: CancellationToken, depth: usize) -> Result<String> {
|
||||
let mut final_text = String::new();
|
||||
@@ -128,7 +242,9 @@ pub async fn run_session(app: Arc<App>, session_id: &str, cancel: CancellationTo
|
||||
}
|
||||
let session = app.session(session_id)?;
|
||||
let cwd = PathBuf::from(&session.cwd);
|
||||
let agent = prompts::find_agent(&app, &cwd, &session.agent);
|
||||
let roots = app.user_roots(session.owner_id.as_deref());
|
||||
let root = app.root_for(session.owner_id.as_deref(), &cwd);
|
||||
let agent = prompts::find_agent(&app, &root, &cwd, &session.agent);
|
||||
let model_spec = agent.model.clone().unwrap_or_else(|| session.model.clone());
|
||||
let (provider, model) = app.registry().resolve(&model_spec)?;
|
||||
|
||||
@@ -144,7 +260,7 @@ pub async fn run_session(app: Arc<App>, session_id: &str, cancel: CancellationTo
|
||||
let tools = available_tools(&app, &agent);
|
||||
let history = active_history(app.store.messages(session_id)?);
|
||||
let req = ChatRequest {
|
||||
system: prompts::system_prompt(&app, &agent, &cwd, &model_spec),
|
||||
system: prompts::system_prompt(&app, &agent, &root, &cwd, &model_spec, &session.permission_mode),
|
||||
messages: history,
|
||||
tools: tools.iter().map(|t| t.def()).collect(),
|
||||
max_tokens: model.max_output.unwrap_or(16_000),
|
||||
@@ -203,6 +319,7 @@ pub async fn run_session(app: Arc<App>, session_id: &str, cancel: CancellationTo
|
||||
cwd: cwd.clone(),
|
||||
cancel: cancel.clone(),
|
||||
agent: agent.name.clone(),
|
||||
roots: roots.clone(),
|
||||
};
|
||||
let by_name = |n: &str| tools.iter().find(|t| t.name() == n).cloned();
|
||||
let all_read_only = calls.iter().all(|(_, n, _)| by_name(n).is_some_and(|t| t.read_only()));
|
||||
@@ -217,11 +334,8 @@ pub async fn run_session(app: Arc<App>, session_id: &str, cancel: CancellationTo
|
||||
results.push(execute_call(&app, &session, &base, by_name(name), id, name, input.clone(), depth).await);
|
||||
}
|
||||
}
|
||||
// Inject prompts the user sent while we were working.
|
||||
let queued = app.with_runtime(session_id, |r| std::mem::take(&mut r.queue));
|
||||
for q in queued {
|
||||
results.push(Part::Text { text: q });
|
||||
}
|
||||
// Deliver messages the user queued while we were working (with attachments).
|
||||
results.extend(take_queue(&app, session_id));
|
||||
save_user_parts(&app, session_id, results)?;
|
||||
}
|
||||
Ok(final_text)
|
||||
@@ -265,12 +379,6 @@ async fn execute_call(
|
||||
return tool_result(call_id, name, "Maximum subagent depth reached.".into(), true, None);
|
||||
}
|
||||
|
||||
// Permission check. Approvals are tracked on the top-level session.
|
||||
let root = app.root_session_id(&session.id);
|
||||
let cfg = app.config();
|
||||
let decision = app.with_runtime(&root, |r| {
|
||||
Evaluator { cfg: &cfg.permissions, mode: &session.permission_mode, approved: &r.approved }.evaluate(tool.as_ref(), &input)
|
||||
});
|
||||
let ctx = ToolCtx {
|
||||
app: app.clone(),
|
||||
session_id: base.session_id.clone(),
|
||||
@@ -278,10 +386,33 @@ async fn execute_call(
|
||||
cwd: base.cwd.clone(),
|
||||
cancel: base.cancel.clone(),
|
||||
agent: base.agent.clone(),
|
||||
roots: base.roots.clone(),
|
||||
};
|
||||
// mcp_call is judged (and approved) as the MCP tool it calls.
|
||||
let call_name = name;
|
||||
let (tool, input, name): (Arc<dyn Tool>, Value, String) = if name == "mcp_call" {
|
||||
match crate::mcp::resolve_call_target(&ctx, &input) {
|
||||
Ok(target) => {
|
||||
let args = input.get("arguments").cloned().unwrap_or(serde_json::json!({}));
|
||||
let n = target.full_name.clone();
|
||||
(target, args, n)
|
||||
}
|
||||
Err(e) => return tool_result(call_id, call_name, format!("Error: {e:#}"), true, None),
|
||||
}
|
||||
} else {
|
||||
(tool, input, name.to_string())
|
||||
};
|
||||
let name = name.as_str();
|
||||
|
||||
// Permission check. Approvals are tracked on the top-level session.
|
||||
let root = app.root_session_id(&session.id);
|
||||
let cfg = app.config();
|
||||
let decision = app.with_runtime(&root, |r| {
|
||||
Evaluator { cfg: &cfg.permissions, mode: &session.permission_mode, approved: &r.approved }.evaluate(tool.as_ref(), &input)
|
||||
});
|
||||
match decision {
|
||||
Decision::Deny => {
|
||||
return tool_result(call_id, name, "Denied by the permission rules in the configuration.".into(), true, None);
|
||||
return tool_result(call_id, call_name, "Denied by the permission rules in the configuration.".into(), true, None);
|
||||
}
|
||||
Decision::Ask => {
|
||||
let diff = tool.preview(&input, &ctx).await;
|
||||
@@ -294,7 +425,7 @@ async fn execute_call(
|
||||
if let Some(m) = message.filter(|m| !m.trim().is_empty()) {
|
||||
text.push_str(&format!(" User feedback: {m}"));
|
||||
}
|
||||
return tool_result(call_id, name, text, true, None);
|
||||
return tool_result(call_id, call_name, text, true, None);
|
||||
}
|
||||
PermissionReply::Always => app.with_runtime(&root, |r| {
|
||||
if name == "bash" {
|
||||
@@ -321,40 +452,76 @@ async fn execute_call(
|
||||
match res {
|
||||
Ok(out) => Part::ToolResult {
|
||||
call_id: call_id.into(),
|
||||
name: name.into(),
|
||||
name: call_name.into(),
|
||||
content: out.content,
|
||||
images: out.images,
|
||||
is_error: out.is_error,
|
||||
meta: out.meta,
|
||||
},
|
||||
Err(e) => tool_result(call_id, name, format!("Error: {e:#}"), true, None),
|
||||
Err(e) => tool_result(call_id, call_name, format!("Error: {e:#}"), true, None),
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_task(app: &Arc<App>, ctx: &ToolCtx, input: Value, depth: usize) -> Result<tools::ToolOutput> {
|
||||
let prompt = tools::arg_str(&input, "prompt")?.to_string();
|
||||
let description = input["description"].as_str().unwrap_or("subagent").to_string();
|
||||
let agent = input["subagent_type"].as_str().unwrap_or("general").to_string();
|
||||
let (child, text) = spawn_subagent(app.clone(), ctx, &agent, &description, prompt, depth + 1).await?;
|
||||
Ok(tools::ToolOutput::text(text).with_meta(serde_json::json!({"session_id": child, "agent": agent})))
|
||||
let agent_name = input["subagent_type"].as_str().unwrap_or("general").to_string();
|
||||
let background = input["background"].as_bool().unwrap_or(false);
|
||||
let (child, agent) = create_child(app, ctx, &agent_name, &description, prompt)?;
|
||||
let meta = serde_json::json!({
|
||||
"session_id": child.id, "agent": agent.name, "model": child.model, "background": background
|
||||
});
|
||||
|
||||
if background {
|
||||
let app = app.clone();
|
||||
let parent_id = ctx.session_id.clone();
|
||||
let child_id = child.id.clone();
|
||||
let agent_name = agent.name.clone();
|
||||
let desc = description.clone();
|
||||
// Background tasks outlive the current turn; they are stopped via their own session.
|
||||
let cancel = CancellationToken::new();
|
||||
tokio::spawn(async move {
|
||||
let res = run_child(app.clone(), &child, &agent, cancel, depth + 1).await;
|
||||
let body = match res {
|
||||
Ok(text) => text,
|
||||
Err(e) => format!("The background task failed: {e:#}"),
|
||||
};
|
||||
let text = format!(
|
||||
"<background-task-result agent=\"{agent_name}\" session=\"{child_id}\" description=\"{desc}\">\n{body}\n</background-task-result>"
|
||||
);
|
||||
// Injected into the running turn, or starts a new turn when the parent is idle.
|
||||
if app.store.get_session(&parent_id).ok().flatten().is_some() {
|
||||
if let Err(e) = submit(app.clone(), parent_id.clone(), vec![Part::Text { text }], "background") {
|
||||
tracing::warn!("could not deliver background result: {e:#}");
|
||||
}
|
||||
}
|
||||
});
|
||||
return Ok(tools::ToolOutput::text(format!(
|
||||
"Started background task '{description}' with subagent '{}' (model {}). Its result will arrive \
|
||||
automatically as a message when it finishes – continue with other work and do not poll.",
|
||||
meta["agent"].as_str().unwrap_or(""),
|
||||
meta["model"].as_str().unwrap_or("")
|
||||
))
|
||||
.with_meta(meta));
|
||||
}
|
||||
|
||||
let text = run_child(app.clone(), &child, &agent, ctx.cancel.child_token(), depth + 1).await?;
|
||||
Ok(tools::ToolOutput::text(text).with_meta(meta))
|
||||
}
|
||||
|
||||
/// Entry point used by the task tool when called outside the runner.
|
||||
pub async fn run_subagent(app: Arc<App>, ctx: &ToolCtx, agent: &str, description: &str, prompt: String) -> Result<(String, String)> {
|
||||
spawn_subagent(app, ctx, agent, description, prompt, 1).await
|
||||
let (child, agent) = create_child(&app, ctx, agent, description, prompt)?;
|
||||
let text = run_child(app, &child, &agent, ctx.cancel.child_token(), 1).await?;
|
||||
Ok((child.id, text))
|
||||
}
|
||||
|
||||
async fn spawn_subagent(
|
||||
app: Arc<App>,
|
||||
ctx: &ToolCtx,
|
||||
agent_name: &str,
|
||||
description: &str,
|
||||
prompt: String,
|
||||
depth: usize,
|
||||
) -> Result<(String, String)> {
|
||||
/// Create the subagent session with its first message.
|
||||
fn create_child(app: &Arc<App>, ctx: &ToolCtx, agent_name: &str, description: &str, prompt: String) -> Result<(Session, Agent)> {
|
||||
let parent = app.session(&ctx.session_id)?;
|
||||
let parent_agent = prompts::find_agent(&app, &ctx.cwd, &parent.agent);
|
||||
let agents = prompts::agents(&app, &ctx.cwd);
|
||||
let root = ctx.root();
|
||||
let parent_agent = prompts::find_agent(app, &root, &ctx.cwd, &parent.agent);
|
||||
let agents = prompts::agents(app, &root, &ctx.cwd);
|
||||
let mut agent = agents
|
||||
.iter()
|
||||
.find(|a| a.name == agent_name && a.mode != "primary")
|
||||
@@ -367,6 +534,9 @@ async fn spawn_subagent(
|
||||
if parent_agent.read_only && !agent.read_only {
|
||||
agent = agents.into_iter().find(|a| a.name == "explore").unwrap_or(agent);
|
||||
}
|
||||
if let Some(m) = &agent.model {
|
||||
app.registry().resolve(m).map_err(|e| anyhow!("subagent '{}' has an unusable model: {e:#}", agent.name))?;
|
||||
}
|
||||
let now = now_ms();
|
||||
let child = Session {
|
||||
id: new_id(),
|
||||
@@ -376,6 +546,7 @@ async fn spawn_subagent(
|
||||
agent: agent.name.clone(),
|
||||
permission_mode: parent.permission_mode.clone(),
|
||||
parent_id: Some(parent.id.clone()),
|
||||
owner_id: parent.owner_id.clone(),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
todos: Value::Null,
|
||||
@@ -385,11 +556,34 @@ async fn spawn_subagent(
|
||||
let msg = Message::new(Role::User, vec![Part::Text { text: prompt }]);
|
||||
app.store.add_message(&child.id, &msg)?;
|
||||
app.emit(Event::MessageAdded { session_id: child.id.clone(), message: msg });
|
||||
Ok((child, agent))
|
||||
}
|
||||
|
||||
/// Run a subagent session to completion, honouring the subagent limits.
|
||||
async fn run_child(app: Arc<App>, child: &Session, agent: &Agent, cancel: CancellationToken, depth: usize) -> Result<String> {
|
||||
// Registered like a normal run so the user can stop it from its own session.
|
||||
app.with_runtime(&child.id, |r| r.cancel = Some(cancel.clone()));
|
||||
app.emit(Event::SessionStatus { session_id: child.id.clone(), busy: true });
|
||||
let res = Box::pin(run_session(app.clone(), &child.id, ctx.cancel.child_token(), depth)).await;
|
||||
// Only first-level subagents count against the limits: nested subagents would
|
||||
// otherwise deadlock waiting for a slot their own parent holds.
|
||||
let _slot = if depth <= 1 {
|
||||
let (slot, waited) = app.subagent_slot(&agent.name, agent.max_parallel, &child.id, &cancel).await;
|
||||
if waited {
|
||||
app.emit(Event::Waiting { session_id: child.id.clone(), reason: None });
|
||||
}
|
||||
slot
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let res = if cancel.is_cancelled() {
|
||||
Err(anyhow!("aborted"))
|
||||
} else {
|
||||
run_session_boxed(app.clone(), child.id.clone(), cancel, depth).await
|
||||
};
|
||||
app.with_runtime(&child.id, |r| r.cancel = None);
|
||||
app.emit(Event::SessionStatus { session_id: child.id.clone(), busy: false });
|
||||
let text = res?;
|
||||
Ok((child.id, if text.trim().is_empty() { "(subagent returned no text)".into() } else { text }))
|
||||
Ok(if text.trim().is_empty() { "(subagent returned no text)".into() } else { text })
|
||||
}
|
||||
|
||||
fn is_retryable(e: &anyhow::Error) -> bool {
|
||||
@@ -415,7 +609,20 @@ async fn stream_step(
|
||||
let p = provider.clone();
|
||||
let r = req.clone();
|
||||
let c = cancel.clone();
|
||||
let handle = tokio::spawn(async move { p.stream(r, tx, c).await });
|
||||
// Wait for a free slot on the model/host; not counted as time to first token.
|
||||
let registry = app.registry();
|
||||
if registry.is_saturated(model_spec) {
|
||||
app.emit(Event::Waiting { session_id: session.id.clone(), reason: Some(format!("model:{model_spec}")) });
|
||||
}
|
||||
let slot = tokio::select! {
|
||||
s = registry.acquire(model_spec) => s,
|
||||
_ = cancel.cancelled() => bail!("aborted"),
|
||||
};
|
||||
app.emit(Event::Waiting { session_id: session.id.clone(), reason: None });
|
||||
let handle = tokio::spawn(async move {
|
||||
let _slot = slot;
|
||||
p.stream(r, tx, c).await
|
||||
});
|
||||
|
||||
let started = Instant::now();
|
||||
let mut ttft: Option<u128> = None;
|
||||
@@ -516,9 +723,13 @@ async fn stream_step(
|
||||
}
|
||||
|
||||
/// Collect a complete, tool-free answer (titles, summaries).
|
||||
async fn complete_text(provider: Arc<dyn Provider>, req: ChatRequest, cancel: CancellationToken) -> Result<(String, Usage)> {
|
||||
async fn complete_text(app: &App, spec: &str, provider: Arc<dyn Provider>, req: ChatRequest, cancel: CancellationToken) -> Result<(String, Usage)> {
|
||||
let slot = app.registry().acquire(spec).await;
|
||||
let (tx, mut rx) = mpsc::channel(256);
|
||||
let handle = tokio::spawn(async move { provider.stream(req, tx, cancel).await });
|
||||
let handle = tokio::spawn(async move {
|
||||
let _slot = slot;
|
||||
provider.stream(req, tx, cancel).await
|
||||
});
|
||||
let mut text = String::new();
|
||||
let mut usage = Usage::default();
|
||||
while let Some(ev) = rx.recv().await {
|
||||
@@ -579,7 +790,7 @@ fn maybe_generate_title(app: Arc<App>, session_id: &str, text: &str) {
|
||||
reasoning: false,
|
||||
model: model.clone(),
|
||||
};
|
||||
match complete_text(provider, req, CancellationToken::new()).await {
|
||||
match complete_text(&app, &spec, provider, req, CancellationToken::new()).await {
|
||||
Ok((title, usage)) => {
|
||||
record_side_usage(&app, &session, &spec, &model, &usage, "title");
|
||||
let title = title.lines().find(|l| !l.trim().is_empty()).unwrap_or("").trim().trim_matches('"').to_string();
|
||||
@@ -625,7 +836,7 @@ pub async fn compact(app: Arc<App>, session_id: &str, cancel: CancellationToken)
|
||||
reasoning: false,
|
||||
model: model.clone(),
|
||||
};
|
||||
let (summary, usage) = complete_text(provider, req, cancel).await?;
|
||||
let (summary, usage) = complete_text(&app, &spec, provider, req, cancel).await?;
|
||||
record_side_usage(&app, &session, &spec, &model, &usage, "compaction");
|
||||
if summary.trim().is_empty() {
|
||||
bail!("compaction produced an empty summary");
|
||||
@@ -638,3 +849,25 @@ pub async fn compact(app: Arc<App>, session_id: &str, cancel: CancellationToken)
|
||||
app.save_session(s)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn call(id: &str) -> Part {
|
||||
Part::ToolCall { id: id.into(), name: "bash".into(), input: serde_json::json!({}), extra: None }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dangling_tool_calls_get_results() {
|
||||
let a = Message::new(Role::Assistant, vec![call("1"), call("2")]);
|
||||
let r = Message::new(Role::User, vec![tool_result("1", "bash", "ok".into(), false, None)]);
|
||||
let a2 = Message::new(Role::Assistant, vec![call("3")]);
|
||||
let h = active_history(vec![Message::new(Role::User, vec![Part::Text { text: "hi".into() }]), a, r, a2]);
|
||||
// user, assistant(1,2), user(result 2 + result 1), assistant(3), user(result 3)
|
||||
assert_eq!(h.len(), 5);
|
||||
let ids = |m: &Message| m.parts.iter().filter_map(|p| if let Part::ToolResult { call_id, .. } = p { Some(call_id.clone()) } else { None }).collect::<Vec<_>>();
|
||||
assert_eq!(ids(&h[2]), vec!["2", "1"]);
|
||||
assert_eq!(ids(&h[4]), vec!["3"]);
|
||||
}
|
||||
}
|
||||
+473
-18
@@ -50,7 +50,74 @@ impl Store {
|
||||
error INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS usage_ts ON usage(ts);
|
||||
CREATE INDEX IF NOT EXISTS usage_root ON usage(root_session_id);",
|
||||
CREATE INDEX IF NOT EXISTS usage_root ON usage(root_session_id);
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
id TEXT PRIMARY KEY,
|
||||
username TEXT NOT NULL UNIQUE COLLATE NOCASE,
|
||||
display_name TEXT NOT NULL DEFAULT '',
|
||||
password_hash TEXT,
|
||||
role TEXT NOT NULL DEFAULT 'user',
|
||||
oidc_sub TEXT UNIQUE,
|
||||
disabled INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS auth_tokens (
|
||||
token_hash TEXT PRIMARY KEY,
|
||||
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
created_at INTEGER NOT NULL,
|
||||
expires_at INTEGER NOT NULL,
|
||||
user_agent TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS folders (
|
||||
id TEXT PRIMARY KEY,
|
||||
path TEXT NOT NULL UNIQUE,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS folder_access (
|
||||
folder_id TEXT NOT NULL REFERENCES folders(id) ON DELETE CASCADE,
|
||||
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
PRIMARY KEY (folder_id, user_id)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS images (
|
||||
id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
session_id TEXT,
|
||||
prompt TEXT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
path TEXT NOT NULL,
|
||||
width INTEGER NOT NULL DEFAULT 0,
|
||||
height INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS images_user ON images(user_id, created_at);
|
||||
CREATE TABLE IF NOT EXISTS queue (
|
||||
id TEXT PRIMARY KEY,
|
||||
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
||||
data TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS queue_session ON queue(session_id, created_at);
|
||||
CREATE TABLE IF NOT EXISTS session_shares (
|
||||
session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
|
||||
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
PRIMARY KEY (session_id, user_id)
|
||||
);",
|
||||
)?;
|
||||
// Columns added after the first release.
|
||||
for (table, column, ddl) in [
|
||||
("sessions", "owner_id", "ALTER TABLE sessions ADD COLUMN owner_id TEXT"),
|
||||
("usage", "user_id", "ALTER TABLE usage ADD COLUMN user_id TEXT"),
|
||||
] {
|
||||
let exists: bool = conn
|
||||
.prepare(&format!("SELECT 1 FROM pragma_table_info('{table}') WHERE name = ?1"))?
|
||||
.exists([column])?;
|
||||
if !exists {
|
||||
conn.execute(ddl, [])?;
|
||||
}
|
||||
}
|
||||
conn.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS sessions_owner ON sessions(owner_id);
|
||||
CREATE INDEX IF NOT EXISTS usage_user ON usage(user_id);",
|
||||
)?;
|
||||
Ok(Self { conn: Mutex::new(conn) })
|
||||
}
|
||||
@@ -58,9 +125,10 @@ impl Store {
|
||||
pub fn save_session(&self, s: &Session) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO sessions (id, data, parent_id, updated_at) VALUES (?1, ?2, ?3, ?4)
|
||||
ON CONFLICT(id) DO UPDATE SET data = excluded.data, updated_at = excluded.updated_at",
|
||||
params![s.id, serde_json::to_string(s)?, s.parent_id, s.updated_at],
|
||||
"INSERT INTO sessions (id, data, parent_id, updated_at, owner_id) VALUES (?1, ?2, ?3, ?4, ?5)
|
||||
ON CONFLICT(id) DO UPDATE SET data = excluded.data, updated_at = excluded.updated_at,
|
||||
owner_id = excluded.owner_id",
|
||||
params![s.id, serde_json::to_string(s)?, s.parent_id, s.updated_at, s.owner_id],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -74,12 +142,15 @@ impl Store {
|
||||
}
|
||||
|
||||
/// Top-level sessions (subagent sessions are reachable through their parent).
|
||||
pub fn list_sessions(&self) -> Result<Vec<Session>> {
|
||||
/// `user = None` lists all sessions (admins); otherwise own and shared ones.
|
||||
pub fn list_sessions(&self, user: Option<&str>) -> Result<Vec<Session>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT data FROM sessions WHERE parent_id IS NULL ORDER BY updated_at DESC LIMIT 500",
|
||||
"SELECT data FROM sessions WHERE parent_id IS NULL AND (?1 IS NULL OR owner_id = ?1
|
||||
OR id IN (SELECT session_id FROM session_shares WHERE user_id = ?1))
|
||||
ORDER BY updated_at DESC LIMIT 500",
|
||||
)?;
|
||||
let rows = stmt.query_map([], |r| r.get::<_, String>(0))?;
|
||||
let rows = stmt.query_map([user], |r| r.get::<_, String>(0))?;
|
||||
let mut out = Vec::new();
|
||||
for row in rows {
|
||||
out.push(serde_json::from_str(&row?)?);
|
||||
@@ -197,8 +268,9 @@ impl Store {
|
||||
conn.execute(
|
||||
"INSERT INTO usage (ts, session_id, root_session_id, provider, model, agent,
|
||||
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens,
|
||||
ttft_ms, duration_ms, cost, error)
|
||||
VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14)",
|
||||
ttft_ms, duration_ms, cost, error, user_id)
|
||||
VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,
|
||||
(SELECT owner_id FROM sessions WHERE id = ?3))",
|
||||
params![
|
||||
u.ts, u.session_id, u.root_session_id, u.provider, u.model, u.agent,
|
||||
u.input_tokens as i64, u.output_tokens as i64, u.cache_read_tokens as i64,
|
||||
@@ -232,21 +304,22 @@ impl Store {
|
||||
|
||||
/// Global statistics in [from, to) (ms). `tz_offset_min` shifts day/hour buckets
|
||||
/// into the viewer's local time.
|
||||
pub fn global_stats(&self, from: i64, to: i64, tz_offset_min: i64) -> Result<Value> {
|
||||
pub fn global_stats(&self, from: i64, to: i64, tz_offset_min: i64, user: Option<&str>) -> Result<Value> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let shift = tz_offset_min * 60_000;
|
||||
let total = conn.query_row(
|
||||
&format!("SELECT {AGG} FROM usage WHERE ts >= ?1 AND ts < ?2"),
|
||||
params![from, to],
|
||||
&format!("SELECT {AGG} FROM usage WHERE ts >= ?1 AND ts < ?2 AND (?3 IS NULL OR user_id = ?3)"),
|
||||
params![from, to, user],
|
||||
|r| agg_row(r, 0),
|
||||
)?;
|
||||
|
||||
let grouped = |key_sql: &str, key_name: &str| -> Result<Vec<Value>> {
|
||||
let mut stmt = conn.prepare(&format!(
|
||||
"SELECT {key_sql}, {AGG} FROM usage WHERE ts >= ?1 AND ts < ?2 AND ?3 = ?3 GROUP BY 1 ORDER BY 1"
|
||||
"SELECT {key_sql}, {AGG} FROM usage WHERE ts >= ?1 AND ts < ?2 AND ?3 = ?3
|
||||
AND (?4 IS NULL OR user_id = ?4) GROUP BY 1 ORDER BY 1"
|
||||
))?;
|
||||
let rows = stmt
|
||||
.query_map(params![from, to, shift], |r| {
|
||||
.query_map(params![from, to, shift, user], |r| {
|
||||
let mut v = agg_row(r, 1)?;
|
||||
v[key_name] = match r.get_ref(0)? {
|
||||
rusqlite::types::ValueRef::Integer(i) => json!(i),
|
||||
@@ -264,6 +337,10 @@ impl Store {
|
||||
let by_model = grouped("provider || '/' || model", "model")?;
|
||||
let by_provider = grouped("provider", "provider")?;
|
||||
let by_agent = grouped("agent", "agent")?;
|
||||
let by_user = grouped(
|
||||
"COALESCE((SELECT username FROM users WHERE users.id = usage.user_id), '(none)')",
|
||||
"user",
|
||||
)?;
|
||||
// Heatmap: weekday (0 = Sunday) x hour of day.
|
||||
let heatmap = grouped(
|
||||
"CAST(strftime('%w', (ts + ?3) / 1000, 'unixepoch') AS INTEGER) * 24 + \
|
||||
@@ -274,9 +351,10 @@ impl Store {
|
||||
let mut stmt = conn.prepare(&format!(
|
||||
"SELECT {day_sql}, provider || '/' || model, SUM(cost),
|
||||
SUM(input_tokens + cache_read_tokens + cache_write_tokens), SUM(output_tokens)
|
||||
FROM usage WHERE ts >= ?1 AND ts < ?2 GROUP BY 1, 2 ORDER BY 1"
|
||||
FROM usage WHERE ts >= ?1 AND ts < ?2 AND (?4 IS NULL OR user_id = ?4)
|
||||
GROUP BY 1, 2 ORDER BY 1"
|
||||
))?;
|
||||
stmt.query_map(params![from, to, shift], |r| {
|
||||
stmt.query_map(params![from, to, shift, user], |r| {
|
||||
Ok(json!({
|
||||
"day": r.get::<_, String>(0)?,
|
||||
"model": r.get::<_, String>(1)?,
|
||||
@@ -292,10 +370,10 @@ impl Store {
|
||||
let mut stmt = conn.prepare(&format!(
|
||||
"SELECT u.root_session_id, COALESCE(json_extract(s.data, '$.title'), '(deleted)'), {AGG}
|
||||
FROM usage u LEFT JOIN sessions s ON s.id = u.root_session_id
|
||||
WHERE u.ts >= ?1 AND u.ts < ?2
|
||||
WHERE u.ts >= ?1 AND u.ts < ?2 AND (?3 IS NULL OR u.user_id = ?3)
|
||||
GROUP BY 1 ORDER BY SUM(u.cost) DESC, COUNT(*) DESC LIMIT 15"
|
||||
))?;
|
||||
stmt.query_map(params![from, to], |r| {
|
||||
stmt.query_map(params![from, to, user], |r| {
|
||||
let mut v = agg_row(r, 2)?;
|
||||
v["session_id"] = json!(r.get::<_, String>(0)?);
|
||||
v["title"] = json!(r.get::<_, String>(1)?);
|
||||
@@ -310,8 +388,385 @@ impl Store {
|
||||
"by_model": by_model,
|
||||
"by_provider": by_provider,
|
||||
"by_agent": by_agent,
|
||||
"by_user": by_user,
|
||||
"heatmap": heatmap,
|
||||
"top_sessions": top_sessions,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
// --- users & auth -------------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct User {
|
||||
pub id: String,
|
||||
pub username: String,
|
||||
pub display_name: String,
|
||||
pub role: String,
|
||||
pub has_password: bool,
|
||||
pub oidc: bool,
|
||||
pub disabled: bool,
|
||||
pub created_at: i64,
|
||||
}
|
||||
|
||||
impl User {
|
||||
pub fn is_admin(&self) -> bool {
|
||||
self.role == "admin"
|
||||
}
|
||||
}
|
||||
|
||||
const USER_COLS: &str =
|
||||
"id, username, display_name, role, password_hash IS NOT NULL, oidc_sub IS NOT NULL, disabled, created_at";
|
||||
|
||||
fn user_row(r: &rusqlite::Row) -> rusqlite::Result<User> {
|
||||
Ok(User {
|
||||
id: r.get(0)?,
|
||||
username: r.get(1)?,
|
||||
display_name: r.get(2)?,
|
||||
role: r.get(3)?,
|
||||
has_password: r.get(4)?,
|
||||
oidc: r.get(5)?,
|
||||
disabled: r.get::<_, i64>(6)? != 0,
|
||||
created_at: r.get(7)?,
|
||||
})
|
||||
}
|
||||
|
||||
impl Store {
|
||||
pub fn user_count(&self) -> Result<i64> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn.query_row("SELECT COUNT(*) FROM users", [], |r| r.get(0))?)
|
||||
}
|
||||
|
||||
pub fn create_user(&self, username: &str, display_name: &str, role: &str, password_hash: Option<&str>, oidc_sub: Option<&str>) -> Result<User> {
|
||||
let id = crate::types::new_id();
|
||||
{
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO users (id, username, display_name, password_hash, role, oidc_sub, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
|
||||
params![id, username, display_name, password_hash, role, oidc_sub, now_ms()],
|
||||
)?;
|
||||
}
|
||||
self.user(&id)?.ok_or_else(|| anyhow::anyhow!("user vanished"))
|
||||
}
|
||||
|
||||
pub fn user(&self, id: &str) -> Result<Option<User>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn.query_row(&format!("SELECT {USER_COLS} FROM users WHERE id = ?1"), [id], user_row).optional()?)
|
||||
}
|
||||
|
||||
pub fn user_by_name(&self, username: &str) -> Result<Option<(User, Option<String>)>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn
|
||||
.query_row(
|
||||
&format!("SELECT {USER_COLS}, password_hash FROM users WHERE username = ?1"),
|
||||
[username],
|
||||
|r| Ok((user_row(r)?, r.get::<_, Option<String>>(8)?)),
|
||||
)
|
||||
.optional()?)
|
||||
}
|
||||
|
||||
pub fn user_by_oidc(&self, sub: &str) -> Result<Option<User>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn.query_row(&format!("SELECT {USER_COLS} FROM users WHERE oidc_sub = ?1"), [sub], user_row).optional()?)
|
||||
}
|
||||
|
||||
pub fn users(&self) -> Result<Vec<User>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare(&format!("SELECT {USER_COLS} FROM users ORDER BY username"))?;
|
||||
let rows = stmt.query_map([], user_row)?.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
/// Update selected fields; None leaves a field unchanged.
|
||||
pub fn update_user(&self, id: &str, display_name: Option<&str>, role: Option<&str>, disabled: Option<bool>, password_hash: Option<&str>, oidc_sub: Option<&str>) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"UPDATE users SET display_name = COALESCE(?2, display_name), role = COALESCE(?3, role),
|
||||
disabled = COALESCE(?4, disabled), password_hash = COALESCE(?5, password_hash),
|
||||
oidc_sub = COALESCE(?6, oidc_sub) WHERE id = ?1",
|
||||
params![id, display_name, role, disabled.map(|d| d as i64), password_hash, oidc_sub],
|
||||
)?;
|
||||
if disabled == Some(true) {
|
||||
conn.execute("DELETE FROM auth_tokens WHERE user_id = ?1", [id])?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_user(&self, id: &str) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM users WHERE id = ?1", [id])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn admin_count(&self) -> Result<i64> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn.query_row("SELECT COUNT(*) FROM users WHERE role = 'admin' AND disabled = 0", [], |r| r.get(0))?)
|
||||
}
|
||||
|
||||
pub fn add_token(&self, token_hash: &str, user_id: &str, expires_at: i64, user_agent: &str) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM auth_tokens WHERE expires_at < ?1", [now_ms()])?;
|
||||
conn.execute(
|
||||
"INSERT INTO auth_tokens (token_hash, user_id, created_at, expires_at, user_agent) VALUES (?1, ?2, ?3, ?4, ?5)",
|
||||
params![token_hash, user_id, now_ms(), expires_at, user_agent],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The (enabled) user behind a valid login token.
|
||||
pub fn user_by_token(&self, token_hash: &str) -> Result<Option<User>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn
|
||||
.query_row(
|
||||
&format!(
|
||||
"SELECT {} FROM auth_tokens t JOIN users u ON u.id = t.user_id
|
||||
WHERE t.token_hash = ?1 AND t.expires_at > ?2 AND u.disabled = 0",
|
||||
USER_COLS.split(", ").map(|c| format!("u.{c}")).collect::<Vec<_>>().join(", ")
|
||||
),
|
||||
params![token_hash, now_ms()],
|
||||
user_row,
|
||||
)
|
||||
.optional()?)
|
||||
}
|
||||
|
||||
pub fn delete_token(&self, token_hash: &str) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM auth_tokens WHERE token_hash = ?1", [token_hash])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_user_tokens(&self, user_id: &str) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM auth_tokens WHERE user_id = ?1", [user_id])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn shares(&self, session_id: &str) -> Result<Vec<String>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare("SELECT user_id FROM session_shares WHERE session_id = ?1")?;
|
||||
let rows = stmt.query_map([session_id], |r| r.get(0))?.collect::<rusqlite::Result<Vec<String>>>()?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
pub fn set_shares(&self, session_id: &str, users: &[String]) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM session_shares WHERE session_id = ?1", [session_id])?;
|
||||
for u in users {
|
||||
conn.execute("INSERT OR IGNORE INTO session_shares (session_id, user_id) VALUES (?1, ?2)", params![session_id, u])?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Assign sessions without owner (created before user management) to a user.
|
||||
pub fn adopt_orphans(&self, user_id: &str) -> Result<usize> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let ids: Vec<(String, String)> = {
|
||||
let mut stmt = conn.prepare("SELECT id, data FROM sessions WHERE owner_id IS NULL")?;
|
||||
stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?.collect::<rusqlite::Result<Vec<_>>>()?
|
||||
};
|
||||
for (id, data) in &ids {
|
||||
let mut v: Value = serde_json::from_str(data)?;
|
||||
v["owner_id"] = json!(user_id);
|
||||
conn.execute("UPDATE sessions SET owner_id = ?2, data = ?3 WHERE id = ?1", params![id, user_id, v.to_string()])?;
|
||||
}
|
||||
conn.execute(
|
||||
"UPDATE usage SET user_id = ?1 WHERE user_id IS NULL",
|
||||
[user_id],
|
||||
)?;
|
||||
Ok(ids.len())
|
||||
}
|
||||
}
|
||||
|
||||
// --- shared folders -----------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct Folder {
|
||||
pub id: String,
|
||||
/// Absolute path inside the workspace.
|
||||
pub path: String,
|
||||
pub users: Vec<String>,
|
||||
pub created_at: i64,
|
||||
}
|
||||
|
||||
impl Store {
|
||||
pub fn folders(&self) -> Result<Vec<Folder>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT f.id, f.path, f.created_at, COALESCE(group_concat(a.user_id), '')
|
||||
FROM folders f LEFT JOIN folder_access a ON a.folder_id = f.id GROUP BY f.id ORDER BY f.path",
|
||||
)?;
|
||||
let rows = stmt
|
||||
.query_map([], |r| {
|
||||
let users: String = r.get(3)?;
|
||||
Ok(Folder {
|
||||
id: r.get(0)?,
|
||||
path: r.get(1)?,
|
||||
created_at: r.get(2)?,
|
||||
users: users.split(',').filter(|u| !u.is_empty()).map(String::from).collect(),
|
||||
})
|
||||
})?
|
||||
.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
pub fn add_folder(&self, path: &str) -> Result<String> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let id = crate::types::new_id();
|
||||
conn.execute("INSERT INTO folders (id, path, created_at) VALUES (?1, ?2, ?3)", params![id, path, now_ms()])?;
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
pub fn delete_folder(&self, id: &str) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM folders WHERE id = ?1", [id])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn set_folder_access(&self, id: &str, users: &[String]) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM folder_access WHERE folder_id = ?1", [id])?;
|
||||
for u in users {
|
||||
conn.execute("INSERT OR IGNORE INTO folder_access (folder_id, user_id) VALUES (?1, ?2)", params![id, u])?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Paths of folders shared with a user.
|
||||
pub fn folders_for_user(&self, user_id: &str) -> Result<Vec<String>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT f.path FROM folders f JOIN folder_access a ON a.folder_id = f.id WHERE a.user_id = ?1 ORDER BY f.path",
|
||||
)?;
|
||||
let rows = stmt.query_map([user_id], |r| r.get(0))?.collect::<rusqlite::Result<Vec<String>>>()?;
|
||||
Ok(rows)
|
||||
}
|
||||
}
|
||||
|
||||
// --- generated images ---------------------------------------------------------
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ImageRecord {
|
||||
pub id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub session_id: Option<String>,
|
||||
pub prompt: String,
|
||||
pub model: String,
|
||||
pub path: String,
|
||||
pub width: i64,
|
||||
pub height: i64,
|
||||
pub created_at: i64,
|
||||
}
|
||||
|
||||
fn image_row(r: &rusqlite::Row) -> rusqlite::Result<ImageRecord> {
|
||||
Ok(ImageRecord {
|
||||
id: r.get(0)?,
|
||||
user_id: r.get(1)?,
|
||||
session_id: r.get(2)?,
|
||||
prompt: r.get(3)?,
|
||||
model: r.get(4)?,
|
||||
path: r.get(5)?,
|
||||
width: r.get(6)?,
|
||||
height: r.get(7)?,
|
||||
created_at: r.get(8)?,
|
||||
})
|
||||
}
|
||||
|
||||
const IMAGE_COLS: &str = "id, user_id, session_id, prompt, model, path, width, height, created_at";
|
||||
|
||||
impl Store {
|
||||
pub fn add_image(&self, img: &ImageRecord) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
&format!("INSERT INTO images ({IMAGE_COLS}) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9)"),
|
||||
params![img.id, img.user_id, img.session_id, img.prompt, img.model, img.path, img.width, img.height, img.created_at],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Latest images; `user = None` lists everybody's (admins).
|
||||
pub fn images(&self, user: Option<&str>, limit: i64) -> Result<Vec<ImageRecord>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare(&format!(
|
||||
"SELECT {IMAGE_COLS} FROM images WHERE ?1 IS NULL OR user_id = ?1 ORDER BY created_at DESC LIMIT ?2"
|
||||
))?;
|
||||
let rows = stmt.query_map(params![user, limit], image_row)?.collect::<rusqlite::Result<Vec<_>>>()?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
pub fn image(&self, id: &str) -> Result<Option<ImageRecord>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn.query_row(&format!("SELECT {IMAGE_COLS} FROM images WHERE id = ?1"), [id], image_row).optional()?)
|
||||
}
|
||||
|
||||
pub fn delete_image(&self, id: &str) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute("DELETE FROM images WHERE id = ?1", [id])?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Store {
|
||||
/// Cost per provider since `from` (ms).
|
||||
pub fn cost_by_provider(&self, from: i64) -> Result<std::collections::BTreeMap<String, f64>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare("SELECT provider, COALESCE(SUM(cost), 0) FROM usage WHERE ts >= ?1 GROUP BY provider")?;
|
||||
let rows = stmt.query_map([from], |r| Ok((r.get::<_, String>(0)?, r.get::<_, f64>(1)?)))?;
|
||||
let mut out = std::collections::BTreeMap::new();
|
||||
for r in rows {
|
||||
let (p, c) = r?;
|
||||
out.insert(p, c);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
// --- message queue -------------------------------------------------------------
|
||||
|
||||
/// A message waiting to be delivered to a busy session.
|
||||
#[derive(Debug, Clone, Serialize, serde::Deserialize)]
|
||||
pub struct QueueItem {
|
||||
pub id: String,
|
||||
pub parts: Vec<crate::types::Part>,
|
||||
/// "user" or "background" (results of background subagents).
|
||||
pub source: String,
|
||||
pub created_at: i64,
|
||||
}
|
||||
|
||||
impl Store {
|
||||
pub fn queue_add(&self, session_id: &str, item: &QueueItem) -> Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO queue (id, session_id, data, created_at) VALUES (?1, ?2, ?3, ?4)",
|
||||
params![item.id, session_id, serde_json::to_string(item)?, item.created_at],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn queue_list(&self, session_id: &str) -> Result<Vec<QueueItem>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare("SELECT data FROM queue WHERE session_id = ?1 ORDER BY created_at, rowid")?;
|
||||
let rows = stmt.query_map([session_id], |r| r.get::<_, String>(0))?;
|
||||
let mut out = Vec::new();
|
||||
for r in rows {
|
||||
out.push(serde_json::from_str(&r?)?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Remove one item; false if it was already delivered.
|
||||
pub fn queue_remove(&self, session_id: &str, id: &str) -> Result<bool> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
Ok(conn.execute("DELETE FROM queue WHERE session_id = ?1 AND id = ?2", params![session_id, id])? > 0)
|
||||
}
|
||||
|
||||
/// Take all items for delivery.
|
||||
pub fn queue_take(&self, session_id: &str) -> Result<Vec<QueueItem>> {
|
||||
let items = self.queue_list(session_id)?;
|
||||
let conn = self.conn.lock().unwrap();
|
||||
for it in &items {
|
||||
conn.execute("DELETE FROM queue WHERE id = ?1", [&it.id])?;
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
}
|
||||
@@ -143,7 +143,8 @@ impl Tool for Task {
|
||||
"Launch a subagent to handle an independent, multi-step task in its own context (e.g. a \
|
||||
broad codebase search, a focused review). The subagent only knows what you put in the \
|
||||
prompt – write a complete, self-contained brief. Its final message is returned to you. \
|
||||
Several task calls in one response run in parallel."
|
||||
Several task calls in one response run in parallel (subject to the configured limits). \
|
||||
With background=true you continue working and receive the result later."
|
||||
.into()
|
||||
}
|
||||
fn schema(&self) -> Value {
|
||||
@@ -152,7 +153,8 @@ impl Tool for Task {
|
||||
"properties": {
|
||||
"description": {"type": "string", "description": "Short (3-5 words) task description"},
|
||||
"prompt": {"type": "string", "description": "The full brief for the subagent"},
|
||||
"subagent_type": {"type": "string", "description": "Subagent to use (default: general)"}
|
||||
"subagent_type": {"type": "string", "description": "Subagent to use (default: general)"},
|
||||
"background": {"type": "boolean", "description": "Run in the background: returns immediately, the result arrives later as a message. Use for work you don't need right away (e.g. writing docs, summaries, translations)."}
|
||||
},
|
||||
"required": ["description", "prompt"]
|
||||
})
|
||||
|
||||
+21
-4
@@ -38,6 +38,7 @@ impl ToolOutput {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ToolCtx {
|
||||
pub app: Arc<App>,
|
||||
pub session_id: String,
|
||||
@@ -46,6 +47,8 @@ pub struct ToolCtx {
|
||||
pub cancel: CancellationToken,
|
||||
/// Agent mode running this tool (subagents inherit restrictions).
|
||||
pub agent: String,
|
||||
/// Directories this session may access (the owner's roots).
|
||||
pub roots: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
impl ToolCtx {
|
||||
@@ -55,11 +58,24 @@ impl ToolCtx {
|
||||
let p = p.trim();
|
||||
let raw = if Path::new(p).is_absolute() { PathBuf::from(p) } else { self.cwd.join(p) };
|
||||
let norm = normalize(&raw);
|
||||
let ws = normalize(&self.app.workspace());
|
||||
if !norm.starts_with(&ws) && !norm.starts_with("/tmp") {
|
||||
bail!("path {} is outside the workspace {}", norm.display(), ws.display());
|
||||
if norm.starts_with("/tmp") || self.roots.iter().any(|r| norm.starts_with(normalize(r))) {
|
||||
// Symlinks must not lead out of the allowed directories either.
|
||||
if let Ok(real) = std::fs::canonicalize(&norm) {
|
||||
if !real.starts_with("/tmp") && !self.roots.iter().any(|r| real.starts_with(r)) {
|
||||
bail!("path {} resolves outside your directories", norm.display());
|
||||
}
|
||||
}
|
||||
return Ok(norm);
|
||||
}
|
||||
Ok(norm)
|
||||
let allowed: Vec<String> = self.roots.iter().map(|r| r.display().to_string()).collect();
|
||||
bail!("path {} is outside your directories ({})", norm.display(), allowed.join(", "))
|
||||
}
|
||||
|
||||
/// The root directory containing the working directory.
|
||||
pub fn root(&self) -> PathBuf {
|
||||
let mut roots = self.roots.clone();
|
||||
roots.sort_by_key(|r| std::cmp::Reverse(r.as_os_str().len()));
|
||||
roots.iter().find(|r| self.cwd.starts_with(r)).cloned().unwrap_or_else(|| self.cwd.clone())
|
||||
}
|
||||
|
||||
pub fn progress(&self, output: impl Into<String>) {
|
||||
@@ -131,6 +147,7 @@ pub fn builtin_tools() -> Vec<Arc<dyn Tool>> {
|
||||
Arc::new(craft::BrowserSnapshot),
|
||||
Arc::new(craft::ImageCompare),
|
||||
Arc::new(crate::prompts::SkillTool),
|
||||
Arc::new(crate::images::GenerateImage),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@@ -115,6 +115,9 @@ pub struct Session {
|
||||
pub permission_mode: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub parent_id: Option<String>,
|
||||
/// User who owns the session (subagent sessions inherit it).
|
||||
#[serde(default)]
|
||||
pub owner_id: Option<String>,
|
||||
pub created_at: i64,
|
||||
pub updated_at: i64,
|
||||
#[serde(default)]
|
||||
@@ -160,7 +163,36 @@ pub enum Event {
|
||||
QuestionRequest { id: String, session_id: String, root_session_id: String, source: String, questions: Vec<Question> },
|
||||
QuestionResolved { id: String, session_id: String },
|
||||
Error { session_id: String, message: String },
|
||||
/// The session waits for a free model or subagent slot (reason None = no longer waiting).
|
||||
/// reason: "model:<provider/model>" or "subagent:<name>".
|
||||
Waiting { session_id: String, reason: Option<String> },
|
||||
McpStatus { servers: Value },
|
||||
/// Messages waiting for a busy session (full list).
|
||||
QueueUpdated { session_id: String, items: Value },
|
||||
}
|
||||
|
||||
impl Event {
|
||||
pub fn session_id(&self) -> Option<&str> {
|
||||
use Event::*;
|
||||
match self {
|
||||
SessionUpdated { session } => Some(&session.id),
|
||||
SessionDeleted { session_id }
|
||||
| SessionStatus { session_id, .. }
|
||||
| MessageAdded { session_id, .. }
|
||||
| TextDelta { session_id, .. }
|
||||
| ReasoningDelta { session_id, .. }
|
||||
| ToolStarted { session_id, .. }
|
||||
| ToolProgress { session_id, .. }
|
||||
| PermissionRequest { session_id, .. }
|
||||
| PermissionResolved { session_id, .. }
|
||||
| QuestionRequest { session_id, .. }
|
||||
| QuestionResolved { session_id, .. }
|
||||
| Error { session_id, .. }
|
||||
| Waiting { session_id, .. }
|
||||
| QueueUpdated { session_id, .. } => Some(session_id),
|
||||
McpStatus { .. } => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_id() -> String {
|
||||
|
||||
Reference in new issue
Block a user