910 lines
31 KiB
Rust
910 lines
31 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}/switch",
|
|
post(crate::package_transitions::switch),
|
|
)
|
|
.route(
|
|
"/api/bots/{id}/tools/{binding_id}/call",
|
|
post(crate::installed_mcp::invoke),
|
|
)
|
|
.route(
|
|
"/api/bots/{id}/tools/{binding_id}/discover",
|
|
post(crate::installed_mcp::discover),
|
|
)
|
|
.route(
|
|
"/api/bots/{id}/tools/install",
|
|
post(crate::installed_mcp::install),
|
|
)
|
|
.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> {
|
|
binding_list(state.pool(), &actor, &bot_id).await.map(Json)
|
|
}
|
|
|
|
#[derive(sqlx::FromRow)]
|
|
struct ListedBinding {
|
|
id: String,
|
|
package_id: String,
|
|
version: String,
|
|
sha256: String,
|
|
install_status: String,
|
|
binding_status: String,
|
|
archive: bool,
|
|
transition: Option<Value>,
|
|
}
|
|
|
|
async fn binding_list(pool: &sqlx::PgPool, actor: &Actor, bot_id: &str) -> Result<Value, ApiError> {
|
|
let rows: Vec<ListedBinding> = sqlx::query_as(
|
|
"SELECT b.id, p.package_id, p.version, p.sha256, p.status AS install_status, b.status AS binding_status, p.manifest IS NOT NULL AS archive, (SELECT jsonb_build_object('operationId',t.id,'version',target.version,'status',t.status) FROM package_transitions t JOIN tool_packages target ON target.id=t.target_id WHERE t.binding_id=b.id AND t.status IN ('pending','cancelling')) AS transition
|
|
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.computer_id=bot.computer_id
|
|
ORDER BY b.created_at",
|
|
)
|
|
.bind(bot_id)
|
|
.bind(&actor.space_id)
|
|
.bind(&actor.user_id)
|
|
.fetch_all(pool)
|
|
.await
|
|
.map_err(internal)?;
|
|
Ok(json!({
|
|
"tools": rows.into_iter().map(|row| json!({
|
|
"bindingId": row.id,
|
|
"archivePackage": row.archive,
|
|
"transition": row.transition,
|
|
"packageId": row.package_id,
|
|
"version": row.version,
|
|
"sha256": row.sha256,
|
|
"installStatus": row.install_status,
|
|
"bindingStatus": row.binding_status,
|
|
"ready": row.install_status == "installed" && row.binding_status == "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}))))?;
|
|
// Provisioning may have assigned a Computer; do not use the pre-boot row.
|
|
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 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 IN ('running','accepted','paused','cancelling','unknown')",
|
|
)
|
|
.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 IN ('running','accepted','paused','cancelling','unknown') 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, ..
|
|
} => {
|
|
// Retired identifiers stay retired; a new sample authorization
|
|
// gets a fresh binding instead of reviving an old identifier.
|
|
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) WHERE status IN ('installing','ready','switching','retiring') 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 p.computer_id=bot.computer_id 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. Execute the exact bytes that were hashed, not a second open of
|
|
/// the pathname; a single shell command alone does not prevent a file swap.
|
|
pub fn verified_run_argv(absolute_path: &str, sha256: &str) -> Vec<String> {
|
|
vec![
|
|
"python3".into(),
|
|
"-I".into(),
|
|
"-c".into(),
|
|
r#"import hashlib, os, stat, sys
|
|
digest, path = sys.argv[1:3]
|
|
fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
|
|
with os.fdopen(fd, 'rb') as source:
|
|
if not stat.S_ISREG(os.fstat(source.fileno()).st_mode):
|
|
sys.exit(97)
|
|
code = source.read(1048577)
|
|
if len(code) > 1048576 or hashlib.sha256(code).hexdigest() != digest:
|
|
print('package digest mismatch', file=sys.stderr)
|
|
sys.exit(97)
|
|
sys.argv = [path]
|
|
exec(compile(code, path, 'exec'), {'__name__': '__main__', '__file__': path})"#
|
|
.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 p.computer_id=bot.computer_id 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"})),
|
|
));
|
|
}
|
|
crate::package_transitions::cancel_for_binding(&state, &actor, &bot_id, &binding_id)
|
|
.await
|
|
.map_err(|error| (StatusCode::CONFLICT, Json(json!({"message":error}))))?;
|
|
crate::installed_mcp::revoke_runtime(&state, &actor, &bot_id, &binding_id)
|
|
.await
|
|
.map_err(|error| (StatusCode::CONFLICT, Json(json!({"message":error}))))?;
|
|
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, bool)> = sqlx::query_as(
|
|
"SELECT p.computer_id, p.package_id, p.version, b.status, p.manifest IS NOT NULL
|
|
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, archive)) = row else {
|
|
return Err((
|
|
StatusCode::NOT_FOUND,
|
|
Json(json!({"message":"binding not found"})),
|
|
));
|
|
};
|
|
if archive {
|
|
return Err((
|
|
StatusCode::CONFLICT,
|
|
Json(
|
|
json!({"message":"archive rollback requires a matching runtime binding; this endpoint supports sample packages only"}),
|
|
),
|
|
));
|
|
}
|
|
// 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)?;
|
|
crate::package_transitions::cancel_for_binding(&state, &actor, &bot_id, &binding_id)
|
|
.await
|
|
.map_err(|error| (StatusCode::CONFLICT, Json(json!({"message":error}))))?;
|
|
crate::installed_mcp::revoke_runtime(&state, &actor, &bot_id, &binding_id)
|
|
.await
|
|
.map_err(|error| (StatusCode::CONFLICT, Json(json!({"message":error}))))?;
|
|
let ready: i64 = sqlx::query_scalar(
|
|
"SELECT COUNT(*) FROM tool_bindings WHERE package_row_id=$1 AND status IN ('ready','installing','switching','retiring')",
|
|
)
|
|
.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 IN ('running','accepted','paused','cancelling','unknown')",
|
|
)
|
|
.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::*;
|
|
|
|
#[sqlx::test(migrations = "../../migrations")]
|
|
async fn package_list_reports_runtime_kind_and_excludes_old_computers(pool: sqlx::PgPool) {
|
|
for sql in [
|
|
"INSERT INTO users(id,name) VALUES ('u','owner'),('v','other')",
|
|
"INSERT INTO spaces(id,user_id,name) VALUES ('s','u','test')",
|
|
"INSERT INTO computers(id,space_id,user_id,scope,scope_key,home_key) VALUES ('c','s','u','team','team:s','home'),('d','s','u','dedicated','bot:d','private')",
|
|
"INSERT INTO bots(id,space_id,user_id,name,computer_id) VALUES ('a','s','u','A','c'),('b','s','u','B','c')",
|
|
"INSERT INTO tool_packages(id,computer_id,package_id,version,sha256,status,manifest) VALUES ('p','c','fixture','1.0.0','digest','installed','{}'),('sample','c','echo','0.0.1','digest','installed',NULL)",
|
|
"INSERT INTO tool_bindings(id,package_row_id,bot_id,status) VALUES ('binding','p','a','ready'),('echo','sample','a','revoked')",
|
|
] {
|
|
sqlx::query(sql).execute(&pool).await.unwrap();
|
|
}
|
|
let owner = Actor {
|
|
user_id: "u".into(),
|
|
space_id: "s".into(),
|
|
};
|
|
let other = Actor {
|
|
user_id: "v".into(),
|
|
space_id: "s".into(),
|
|
};
|
|
let result = binding_list(&pool, &owner, "a").await.unwrap();
|
|
let tools = result["tools"].as_array().unwrap();
|
|
assert_eq!(tools.len(), 2);
|
|
let archive = tools.iter().find(|t| t["bindingId"] == "binding").unwrap();
|
|
assert_eq!(archive["archivePackage"], true);
|
|
assert_eq!(archive["ready"], true);
|
|
let sample = tools.iter().find(|t| t["bindingId"] == "echo").unwrap();
|
|
assert_eq!(sample["archivePackage"], false);
|
|
assert_eq!(sample["ready"], false);
|
|
assert_eq!(
|
|
binding_list(&pool, &other, "a").await.unwrap()["tools"],
|
|
json!([])
|
|
);
|
|
assert_eq!(
|
|
binding_list(&pool, &owner, "b").await.unwrap()["tools"],
|
|
json!([])
|
|
);
|
|
sqlx::query("UPDATE bots SET computer_id='d' WHERE id='a'")
|
|
.execute(&pool)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(
|
|
binding_list(&pool, &owner, "a").await.unwrap()["tools"],
|
|
json!([])
|
|
);
|
|
}
|
|
|
|
#[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)"));
|
|
}
|
|
}
|