787 lines
25 KiB
Rust
787 lines
25 KiB
Rust
|
|
//! 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<Value>);
|
||
|
|
|
||
|
|
pub fn router() -> Router<AppState> {
|
||
|
|
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<String>,
|
||
|
|
#[serde(default)]
|
||
|
|
pub version: Option<String>,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// 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<AppState>,
|
||
|
|
actor: Actor,
|
||
|
|
Path(bot_id): Path<String>,
|
||
|
|
) -> Result<Json<Value>, 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::<Vec<_>>()
|
||
|
|
})))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn install_sample(
|
||
|
|
State(state): State<AppState>,
|
||
|
|
actor: Actor,
|
||
|
|
Path(bot_id): Path<String>,
|
||
|
|
Json(body): Json<InstallBody>,
|
||
|
|
) -> Result<Json<Value>, 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<String>,
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn call_installed(
|
||
|
|
State(state): State<AppState>,
|
||
|
|
actor: Actor,
|
||
|
|
Path(bot_id): Path<String>,
|
||
|
|
Json(body): Json<CallBody>,
|
||
|
|
) -> Result<Json<Value>, 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<Option<BoundPackage>, 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<String> {
|
||
|
|
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<bool, sqlx::Error> {
|
||
|
|
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<lazyboy_control::CommandResult, String> {
|
||
|
|
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<AppState>,
|
||
|
|
actor: Actor,
|
||
|
|
Path((bot_id, binding_id)): Path<(String, String)>,
|
||
|
|
) -> Result<Json<Value>, 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<AppState>,
|
||
|
|
actor: Actor,
|
||
|
|
Path((bot_id, binding_id)): Path<(String, String)>,
|
||
|
|
) -> Result<Json<Value>, 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<String> = 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<AppState>,
|
||
|
|
actor: Actor,
|
||
|
|
Path((bot_id, binding_id)): Path<(String, String)>,
|
||
|
|
) -> Result<Json<Value>, 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<E: std::fmt::Display>(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::<Vec<_>>()
|
||
|
|
.join("\n")
|
||
|
|
.to_ascii_lowercase();
|
||
|
|
assert!(!statements.contains("unique (bots.computer_id)"));
|
||
|
|
assert!(!statements.contains("unique(bots.computer_id)"));
|
||
|
|
}
|
||
|
|
}
|