//! Install a reviewed local package onto the bound Computer and bind it to a bot. use axum::extract::{Path, State}; use axum::http::StatusCode; use axum::routing::{get, post}; use axum::{Json, Router}; use lazyboy_contracts::ComputerMode; use lazyboy_control::{ AdapterContext, CommandRequest, ComputerRef, ReadyBinding, ReadyBindingChange, RollbackChange, SandboxProvider, ToolManifest, artifact_digest, package_gc_allowed, pin_running_jobs, plan_version_switch, previous_version, rollback_one_ready_binding, upsert_one_ready_binding, validate_manifest, }; use serde::Deserialize; use serde_json::{Value, json}; use sqlx::PgPool; use uuid::Uuid; use crate::computer::{self, adapter_context}; use crate::db::{Actor, parse_mode}; use crate::state::AppState; const ECHO_SOURCE: &str = include_str!("../../../scripts/sample-mcp/echo_server.py"); pub const ECHO_ID: &str = "lazyboy.example.echo"; pub const ECHO_VERSION: &str = "0.0.1"; pub const ECHO_VERSION_NEXT: &str = "0.0.2"; type ApiError = (StatusCode, Json); pub fn router() -> Router { Router::new() .route("/api/bots/{id}/tools", get(list_tools).post(install_sample)) .route("/api/bots/{id}/tools/call", post(call_installed)) .route( "/api/bots/{id}/tools/{binding_id}/revoke", post(revoke_binding), ) .route( "/api/bots/{id}/tools/{binding_id}/rollback", post(rollback_binding), ) .route( "/api/bots/{id}/tools/{binding_id}/remove", post(remove_binding), ) } #[derive(Deserialize)] struct InstallBody { #[serde(default, rename = "packageId", alias = "package_id")] pub package_id: Option, #[serde(default)] pub version: Option, } /// Immutable packages live on the Computer, not in the bot workspace folder. pub fn echo_relative_path(mode: ComputerMode, version: &str) -> String { match mode { ComputerMode::Team => format!("shared/tools/{ECHO_ID}/{version}/echo_server.py"), ComputerMode::Dedicated => format!("tools/{ECHO_ID}/{version}/echo_server.py"), } } fn echo_version(requested: Option<&str>) -> Result<&str, ApiError> { match requested.unwrap_or(ECHO_VERSION) { ECHO_VERSION => Ok(ECHO_VERSION), ECHO_VERSION_NEXT => Ok(ECHO_VERSION_NEXT), other => Err(( StatusCode::BAD_REQUEST, Json(json!({"message": format!("unsupported echo version {other}")})), )), } } pub fn tool_install_enabled() -> bool { std::env::var("LAZYBOY_TOOL_INSTALL_ENABLED") .map(|value| value != "false" && value != "0") .unwrap_or(true) } async fn list_tools( State(state): State, actor: Actor, Path(bot_id): Path, ) -> Result, ApiError> { let rows: Vec<(String, String, String, String, String, String)> = sqlx::query_as( "SELECT b.id, p.package_id, p.version, p.sha256, p.status, b.status FROM tool_bindings b JOIN tool_packages p ON p.id=b.package_row_id JOIN bots bot ON bot.id=b.bot_id WHERE b.bot_id=$1 AND bot.space_id=$2 AND bot.user_id=$3 ORDER BY b.created_at", ) .bind(&bot_id) .bind(&actor.space_id) .bind(&actor.user_id) .fetch_all(state.pool()) .await .map_err(internal)?; Ok(Json(json!({ "tools": rows.into_iter().map(|(id, pkg, ver, sha, install, bind)| json!({ "bindingId": id, "packageId": pkg, "version": ver, "sha256": sha, "installStatus": install, "bindingStatus": bind, "ready": install == "installed" && bind == "ready", "executionLocation": "assigned_computer", })).collect::>() }))) } async fn install_sample( State(state): State, actor: Actor, Path(bot_id): Path, Json(body): Json, ) -> Result, ApiError> { if !tool_install_enabled() { return Err(( StatusCode::FORBIDDEN, Json(json!({"message":"tool install is disabled"})), )); } let package_id = body.package_id.unwrap_or_else(|| ECHO_ID.into()); if package_id != ECHO_ID { return Err(( StatusCode::BAD_REQUEST, Json( json!({"message":"only lazyboy.example.echo is installable without a reviewed artifact"}), ), )); } let version = echo_version(body.version.as_deref())?.to_string(); let digest = artifact_digest(ECHO_SOURCE.as_bytes()); let manifest = ToolManifest { id: ECHO_ID.into(), version: version.clone(), sha256: digest.clone(), entrypoint: vec!["python3".into(), "./echo_server.py".into()], share_immutable_package: true, }; validate_manifest(&manifest).map_err(|error| { ( StatusCode::BAD_REQUEST, Json(json!({"message": format!("{error:?}")})), ) })?; let bot = state .db .get_bot(&actor, &bot_id) .await .map_err(internal)? .ok_or(( StatusCode::NOT_FOUND, Json(json!({"message":"bot not found"})), ))?; computer::boot_for(&state, &actor, &bot_id, false) .await .map_err(|error| (StatusCode::BAD_GATEWAY, Json(json!({"message": error}))))?; let computer = state .db .get_computer(bot.computer_id.as_deref().unwrap_or("")) .await .map_err(internal)? .ok_or(( StatusCode::NOT_FOUND, Json(json!({"message":"computer not found"})), ))?; let computer_ref = computer::computer_ref(&computer).ok_or(( StatusCode::CONFLICT, Json(json!({"message":"computer is not running"})), ))?; let mode = parse_mode(&computer.scope); let relative = echo_relative_path(mode, &version); state .sandbox .write_file( &computer_ref, &relative, ECHO_SOURCE.as_bytes(), &adapter_context(&actor, &bot_id, "tool-install"), ) .await .map_err(|error| { ( StatusCode::BAD_GATEWAY, Json(json!({"message": error.to_string()})), ) })?; let package_row = Uuid::new_v4().to_string(); sqlx::query( "INSERT INTO tool_packages (id, computer_id, package_id, version, sha256, status) VALUES ($1,$2,$3,$4,$5,'installed') ON CONFLICT (computer_id, package_id, version) DO UPDATE SET status='installed', sha256=EXCLUDED.sha256", ) .bind(&package_row) .bind(&computer.id) .bind(ECHO_ID) .bind(&version) .bind(&digest) .execute(state.pool()) .await .map_err(internal)?; let row_id: String = sqlx::query_scalar( "SELECT id FROM tool_packages WHERE computer_id=$1 AND package_id=$2 AND version=$3", ) .bind(&computer.id) .bind(ECHO_ID) .bind(&version) .fetch_one(state.pool()) .await .map_err(internal)?; let existing: Option<(String, String, String)> = sqlx::query_as( "SELECT b.id, b.package_row_id, p.version FROM tool_bindings b JOIN tool_packages p ON p.id=b.package_row_id WHERE b.bot_id=$1 AND p.package_id=$2 AND b.status='ready' ORDER BY b.created_at DESC LIMIT 1", ) .bind(&bot_id) .bind(ECHO_ID) .fetch_optional(state.pool()) .await .map_err(internal)?; let jobs: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM computer_jobs WHERE computer_id=$1 AND status='running'", ) .bind(&computer.id) .fetch_one(state.pool()) .await .unwrap_or(0); let current_version = existing .as_ref() .map(|row| row.2.as_str()) .unwrap_or(&version); let plan = plan_version_switch(current_version, &version, jobs.max(0) as usize); // Only this bot's jobs were started with the version being replaced; other // bots sharing the Computer keep their own bindings. if let Some(pin) = pin_running_jobs(current_version, &version, jobs.max(0) as usize) { let _ = sqlx::query( "UPDATE computer_jobs SET pin_version=$2 WHERE computer_id=$1 AND bot_id=$3 AND status='running' AND pin_version IS NULL", ) .bind(&computer.id) .bind(&pin) .bind(&bot_id) .execute(state.pool()) .await; } let ready = existing .as_ref() .map(|(id, package_row_id, ver)| ReadyBinding { id: id.clone(), package_row_id: package_row_id.clone(), version: ver.clone(), }); let change = upsert_one_ready_binding( ready.as_ref(), &row_id, &version, &Uuid::new_v4().to_string(), ); let binding_id = match &change { ReadyBindingChange::Insert { id, package_row_id, .. } => { // (package_row_id, bot_id) is unique: reinstalling after a revoke // revives the old row instead of failing on the constraint. let binding_id: String = sqlx::query_scalar( "INSERT INTO tool_bindings (id, package_row_id, bot_id, status) VALUES ($1,$2,$3,'ready') ON CONFLICT (package_row_id, bot_id) DO UPDATE SET status='ready' RETURNING id", ) .bind(id) .bind(package_row_id) .bind(&bot_id) .fetch_one(state.pool()) .await .map_err(internal)?; binding_id } ReadyBindingChange::Switch { id, package_row_id, .. } => { switch_binding(state.pool(), id, package_row_id, &bot_id) .await .map_err(internal)?; id.clone() } ReadyBindingChange::Keep { id } => id.clone(), }; Ok(Json(json!({ "packageId": ECHO_ID, "version": version, "bindingId": binding_id, "sha256": digest, "path": relative, "installStatus": "installed", "bindingStatus": "ready", "jobPin": plan.job_pin, "newRuns": plan.new_runs, "executionLocation": "assigned_computer", }))) } /// Point a ready binding at another package row. A revoked leftover for the /// same (package_row, bot) would trip the unique index, so it is dropped first. async fn switch_binding( pool: &PgPool, binding_id: &str, package_row_id: &str, bot_id: &str, ) -> Result<(), sqlx::Error> { let mut tx = pool.begin().await?; sqlx::query( "DELETE FROM tool_bindings WHERE package_row_id=$1 AND bot_id=$2 AND status<>'ready' AND id<>$3", ) .bind(package_row_id) .bind(bot_id) .bind(binding_id) .execute(&mut *tx) .await?; sqlx::query("UPDATE tool_bindings SET package_row_id=$2, status='ready' WHERE id=$1") .bind(binding_id) .bind(package_row_id) .execute(&mut *tx) .await?; tx.commit().await } #[derive(Deserialize)] struct CallBody { pub text: String, #[serde(default, rename = "packageId")] pub package_id: Option, } async fn call_installed( State(state): State, actor: Actor, Path(bot_id): Path, Json(body): Json, ) -> Result, ApiError> { let package_id = body.package_id.as_deref().unwrap_or(ECHO_ID); if package_id != ECHO_ID { return Err(( StatusCode::BAD_REQUEST, Json(json!({"message":"only lazyboy.example.echo can be called via this fixture"})), )); } if !package_is_bound(state.pool(), &actor, &bot_id, package_id) .await .map_err(internal)? { return Err(( StatusCode::FORBIDDEN, Json(json!({"message":"package not bound to this agent"})), )); } let bot = state .db .get_bot(&actor, &bot_id) .await .map_err(internal)? .ok_or(( StatusCode::NOT_FOUND, Json(json!({"message":"bot not found"})), ))?; let computer = state .db .get_computer(bot.computer_id.as_deref().unwrap_or("")) .await .map_err(internal)? .ok_or(( StatusCode::NOT_FOUND, Json(json!({"message":"computer not found"})), ))?; let computer_ref = computer::computer_ref(&computer).ok_or(( StatusCode::CONFLICT, Json(json!({"message":"computer is not running"})), ))?; let bound = bound_echo_relative(state.pool(), &actor, &bot_id, parse_mode(&computer.scope)) .await .map_err(internal)? .ok_or(( StatusCode::FORBIDDEN, Json(json!({"message":"package not bound to this agent"})), ))?; let result = exec_echo( state.sandbox.as_ref(), &computer_ref, &adapter_context(&actor, &bot_id, "computer-mcp"), &bound, &body.text, ) .await .map_err(|error| (StatusCode::BAD_GATEWAY, Json(json!({"message": error}))))?; Ok(Json(json!({ "exitCode": result.code, "stdout": result.stdout, "stderr": result.stderr, "executionLocation": "assigned_computer", }))) } /// The bound package as it must exist on the Computer: where it lives and the /// digest the installer recorded for it. #[derive(Debug, Clone, PartialEq, Eq)] pub struct BoundPackage { pub relative: String, pub sha256: String, } pub async fn bound_echo_relative( pool: &PgPool, actor: &Actor, bot_id: &str, mode: ComputerMode, ) -> Result, sqlx::Error> { let row: Option<(String, String)> = sqlx::query_as( "SELECT p.version, p.sha256 FROM tool_bindings b JOIN tool_packages p ON p.id=b.package_row_id JOIN bots bot ON bot.id=b.bot_id WHERE b.bot_id=$1 AND bot.space_id=$2 AND bot.user_id=$3 AND p.package_id=$4 AND b.status='ready' AND p.status='installed' ORDER BY b.created_at DESC LIMIT 1", ) .bind(bot_id) .bind(&actor.space_id) .bind(&actor.user_id) .bind(ECHO_ID) .fetch_optional(pool) .await?; Ok(row.map(|(version, sha256)| BoundPackage { relative: echo_relative_path(mode, &version), sha256, })) } /// Run the package only if the file on the Computer still hashes to what was /// installed. The check and the run are one command so nothing can swap the /// file in between; the digest and path travel as argv, never interpolated. pub fn verified_run_argv(absolute_path: &str, sha256: &str) -> Vec { vec![ "bash".into(), "-c".into(), r#"printf '%s %s\n' "$1" "$2" | sha256sum -c --status || { echo "package digest mismatch: $2" >&2; exit 97; }; exec python3 "$2""#.into(), "verify".into(), sha256.into(), absolute_path.into(), ] } pub async fn package_is_bound( pool: &PgPool, actor: &Actor, bot_id: &str, package_id: &str, ) -> Result { let bound: Option<(String,)> = sqlx::query_as( "SELECT p.id FROM tool_bindings b JOIN tool_packages p ON p.id=b.package_row_id JOIN bots bot ON bot.id=b.bot_id WHERE b.bot_id=$1 AND bot.space_id=$2 AND bot.user_id=$3 AND p.package_id=$4 AND b.status='ready' AND p.status='installed' LIMIT 1", ) .bind(bot_id) .bind(&actor.space_id) .bind(&actor.user_id) .bind(package_id) .fetch_optional(pool) .await?; Ok(bound.is_some()) } pub async fn exec_echo( sandbox: &dyn SandboxProvider, computer_ref: &ComputerRef, adapter: &AdapterContext, bound: &BoundPackage, text: &str, ) -> Result { let rpc = json!({ "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"arguments": {"text": text}}, }); sandbox .execute( computer_ref, CommandRequest { argv: verified_run_argv( &format!("/home/lazyboy/{}", bound.relative), &bound.sha256, ), cwd: None, timeout_ms: Some(10_000), stdin: Some(format!("{rpc}\n")), ..CommandRequest::default() }, adapter, ) .await .map_err(|error| error.to_string()) } async fn revoke_binding( State(state): State, actor: Actor, Path((bot_id, binding_id)): Path<(String, String)>, ) -> Result, ApiError> { let done = sqlx::query( "UPDATE tool_bindings b SET status='revoked' FROM bots bot WHERE b.id=$1 AND b.bot_id=$2 AND bot.id=b.bot_id AND bot.space_id=$3 AND bot.user_id=$4", ) .bind(&binding_id) .bind(&bot_id) .bind(&actor.space_id) .bind(&actor.user_id) .execute(state.pool()) .await .map_err(internal)?; if done.rows_affected() == 0 { return Err(( StatusCode::NOT_FOUND, Json(json!({"message":"binding not found"})), )); } Ok(Json(json!({"ok": true}))) } async fn rollback_binding( State(state): State, actor: Actor, Path((bot_id, binding_id)): Path<(String, String)>, ) -> Result, ApiError> { let row: Option<(String, String, String, String)> = sqlx::query_as( "SELECT p.computer_id, p.package_id, p.version, b.status FROM tool_bindings b JOIN tool_packages p ON p.id=b.package_row_id JOIN bots bot ON bot.id=b.bot_id WHERE b.id=$1 AND b.bot_id=$2 AND bot.space_id=$3 AND bot.user_id=$4", ) .bind(&binding_id) .bind(&bot_id) .bind(&actor.space_id) .bind(&actor.user_id) .fetch_optional(state.pool()) .await .map_err(internal)?; let Some((computer_id, package_id, current, binding_status)) = row else { return Err(( StatusCode::NOT_FOUND, Json(json!({"message":"binding not found"})), )); }; // Rolling back a revoked binding would silently re-enable the package. if binding_status != "ready" { return Err(( StatusCode::CONFLICT, Json(json!({"message": format!("binding is {binding_status}; install again instead")})), )); } let versions: Vec = sqlx::query_scalar( "SELECT version FROM tool_packages WHERE computer_id=$1 AND package_id=$2 AND status='installed' ORDER BY created_at", ) .bind(&computer_id) .bind(&package_id) .fetch_all(state.pool()) .await .map_err(internal)?; let refs: Vec<&str> = versions.iter().map(String::as_str).collect(); let Some(prev) = previous_version(&refs, ¤t) else { return Err(( StatusCode::CONFLICT, Json(json!({"message":"no previous version to rollback"})), )); }; let prev_row: String = sqlx::query_scalar( "SELECT id FROM tool_packages WHERE computer_id=$1 AND package_id=$2 AND version=$3", ) .bind(&computer_id) .bind(&package_id) .bind(prev) .fetch_one(state.pool()) .await .map_err(internal)?; let leftover: Option<(String,)> = sqlx::query_as( "SELECT id FROM tool_bindings WHERE package_row_id=$1 AND bot_id=$2 AND status='ready' AND id<>$3", ) .bind(&prev_row) .bind(&bot_id) .bind(&binding_id) .fetch_optional(state.pool()) .await .map_err(internal)?; let current_binding = ReadyBinding { id: binding_id.clone(), package_row_id: String::new(), version: current.clone(), }; match rollback_one_ready_binding( ¤t_binding, leftover.is_some(), Some(&prev_row), Some(prev), ) { RollbackChange::NoPrevious => { return Err(( StatusCode::CONFLICT, Json(json!({"message":"no previous version to rollback"})), )); } RollbackChange::Conflict => { sqlx::query( "UPDATE tool_bindings SET status='revoked' WHERE package_row_id=$1 AND bot_id=$2 AND status='ready' AND id<>$3", ) .bind(&prev_row) .bind(&bot_id) .bind(&binding_id) .execute(state.pool()) .await .map_err(internal)?; switch_binding(state.pool(), &binding_id, &prev_row, &bot_id) .await .map_err(internal)?; } RollbackChange::Switch { id, package_row_id, .. } => { switch_binding(state.pool(), &id, &package_row_id, &bot_id) .await .map_err(internal)?; } } Ok(Json(json!({ "ok": true, "version": prev, "previous": current, "executionLocation": "assigned_computer", }))) } async fn remove_binding( State(state): State, actor: Actor, Path((bot_id, binding_id)): Path<(String, String)>, ) -> Result, ApiError> { let row: Option<(String, String)> = sqlx::query_as( "SELECT p.id, p.computer_id FROM tool_bindings b JOIN tool_packages p ON p.id=b.package_row_id JOIN bots bot ON bot.id=b.bot_id WHERE b.id=$1 AND b.bot_id=$2 AND bot.space_id=$3 AND bot.user_id=$4", ) .bind(&binding_id) .bind(&bot_id) .bind(&actor.space_id) .bind(&actor.user_id) .fetch_optional(state.pool()) .await .map_err(internal)?; let Some((package_row, computer_id)) = row else { return Err(( StatusCode::NOT_FOUND, Json(json!({"message":"binding not found"})), )); }; sqlx::query( "UPDATE tool_bindings b SET status='revoked' FROM bots bot WHERE b.id=$1 AND b.bot_id=$2 AND bot.id=b.bot_id AND bot.space_id=$3 AND bot.user_id=$4", ) .bind(&binding_id) .bind(&bot_id) .bind(&actor.space_id) .bind(&actor.user_id) .execute(state.pool()) .await .map_err(internal)?; let ready: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM tool_bindings WHERE package_row_id=$1 AND status='ready'", ) .bind(&package_row) .fetch_one(state.pool()) .await .unwrap_or(0); let jobs: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM computer_jobs WHERE computer_id=$1 AND status='running'", ) .bind(&computer_id) .fetch_one(state.pool()) .await .unwrap_or(0); let removed = package_gc_allowed(ready.max(0) as usize, jobs.max(0) as usize); if removed { sqlx::query("UPDATE tool_packages SET status='removed' WHERE id=$1") .bind(&package_row) .execute(state.pool()) .await .map_err(internal)?; } Ok(Json(json!({ "ok": true, "packageRemoved": removed, "executionLocation": "assigned_computer", }))) } fn internal(error: E) -> ApiError { tracing::error!("tool install: {error}"); ( StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"message":"internal error"})), ) } #[cfg(test)] mod tests { use super::*; #[test] fn team_packages_live_on_the_shared_tree() { assert_eq!( echo_relative_path(ComputerMode::Team, ECHO_VERSION), "shared/tools/lazyboy.example.echo/0.0.1/echo_server.py" ); assert_eq!( echo_relative_path(ComputerMode::Dedicated, ECHO_VERSION_NEXT), "tools/lazyboy.example.echo/0.0.2/echo_server.py" ); } /// A package whose bytes no longer match the installed digest must not run. #[test] fn tampered_package_is_refused_before_it_runs() { let dir = std::env::temp_dir().join(format!("lazyboy-verify-{}", Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let path = dir.join("echo_server.py"); std::fs::write(&path, "print('ok')\n").unwrap(); let path = path.to_string_lossy().into_owned(); let good = artifact_digest(b"print('ok')\n"); let run = |digest: &str| { let argv = verified_run_argv(&path, digest); std::process::Command::new(&argv[0]) .args(&argv[1..]) .output() .unwrap() }; let ok = run(&good); assert_eq!(ok.status.code(), Some(0), "{ok:?}"); assert_eq!(String::from_utf8_lossy(&ok.stdout).trim(), "ok"); let bad = run(&artifact_digest(b"something else")); assert_eq!(bad.status.code(), Some(97), "{bad:?}"); assert!(bad.stdout.is_empty(), "tampered package must not execute"); assert!(String::from_utf8_lossy(&bad.stderr).contains("digest mismatch")); let _ = std::fs::remove_dir_all(&dir); } #[test] fn agent_computer_migration_does_not_unique_bot_computer_id() { let sql = include_str!("../../../migrations/022_agent_computer.sql"); assert!(sql.contains("do NOT add UNIQUE(bots.computer_id)")); let statements: String = sql .lines() .filter(|line| !line.trim_start().starts_with("--")) .collect::>() .join("\n") .to_ascii_lowercase(); assert!(!statements.contains("unique (bots.computer_id)")); assert!(!statements.contains("unique(bots.computer_id)")); } }