Compare commits
9 Commits
4a12610662
...
c3327b9ef6
| Author | SHA1 | Date | |
|---|---|---|---|
| c3327b9ef6 | |||
| d751055452 | |||
| b0ce83a937 | |||
| 3bb1272c6d | |||
| 5501558354 | |||
| f0a76fd82c | |||
| 2506713fad | |||
| a060666f6f | |||
| 4192798ceb |
Generated
+36
@@ -390,6 +390,7 @@ dependencies = [
|
||||
"anyhow",
|
||||
"argon2",
|
||||
"async-stream",
|
||||
"async-trait",
|
||||
"axum 0.8.9",
|
||||
"axum-test",
|
||||
"base64",
|
||||
@@ -407,6 +408,7 @@ dependencies = [
|
||||
"mime",
|
||||
"mime_guess",
|
||||
"once_cell",
|
||||
"open-agent-sdk",
|
||||
"parking_lot",
|
||||
"regex",
|
||||
"reqwest",
|
||||
@@ -1825,6 +1827,31 @@ version = "1.70.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
|
||||
|
||||
[[package]]
|
||||
name = "open-agent-sdk"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"base64",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"futures",
|
||||
"glob",
|
||||
"pin-project-lite",
|
||||
"regex",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_yaml",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tokio-util",
|
||||
"tracing",
|
||||
"uuid",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "openssl"
|
||||
version = "0.10.80"
|
||||
@@ -3811,6 +3838,15 @@ dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.59.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
|
||||
dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.61.2"
|
||||
|
||||
@@ -78,6 +78,12 @@ hex = "0.4"
|
||||
include_dir = "0.7"
|
||||
mime_guess = "2"
|
||||
|
||||
# Embedded AI agent engine (MIT, vendored as a path dependency). Provides the
|
||||
# agentic loop, built-in tools, permission callback and cost tracking that the
|
||||
# Agent Workbench (src/routes/agent.rs) drives through BoosAPI's own provider.
|
||||
open-agent-sdk = { path = "../open-agent-sdk-rust" }
|
||||
async-trait = "0.1"
|
||||
|
||||
[dev-dependencies]
|
||||
rstest = "0.24"
|
||||
wiremock = "0.6"
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
-- Agent Workbench sessions: per-session token / cost accounting and metadata.
|
||||
--
|
||||
-- One row per agent conversation session. Token counts and cost accumulate
|
||||
-- across the turns of a session; `updated_at` advances on every `/api/agent/chat`
|
||||
-- round-trip. `user_id` is nullable so admin (global-token) sessions are allowed
|
||||
-- without a user row, but every session is still isolated by its own sandbox
|
||||
-- directory keyed on `id`.
|
||||
CREATE TABLE IF NOT EXISTS agent_sessions (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
user_id INTEGER,
|
||||
title TEXT,
|
||||
model TEXT NOT NULL DEFAULT '',
|
||||
permission_mode TEXT NOT NULL DEFAULT 'default',
|
||||
input_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
output_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
cost_usd REAL NOT NULL DEFAULT 0,
|
||||
num_turns INTEGER NOT NULL DEFAULT 0,
|
||||
message_count INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS agent_sessions_user_id_idx
|
||||
ON agent_sessions (user_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS agent_sessions_updated_at_idx
|
||||
ON agent_sessions (updated_at);
|
||||
@@ -0,0 +1,480 @@
|
||||
//! Agent Workbench HTTP surface.
|
||||
//!
|
||||
//! A general-purpose, server-side AI agent exposed to the SPA over SSE. The
|
||||
//! agentic loop comes from the `open-agent-sdk` crate; every LLM call is routed
|
||||
//! back through BoosAPI's own OpenAI-compatible `/v1` endpoint using the
|
||||
//! built-in downstream key, so agent usage flows through the same aggregation
|
||||
//! and accounting pipeline as any other consumer (and can use any aggregated
|
||||
//! model).
|
||||
//!
|
||||
//! Routes (all require a Bearer token — admin token, proxy token, or a user JWT):
|
||||
//! - `POST /api/agent/chat` Stream a turn as SSE (text / tool_use /
|
||||
//! tool_result / approval_request / result / error / done).
|
||||
//! - `POST /api/agent/approve` Resolve a pending tool approval.
|
||||
//! - `GET /api/agent/status` Engine status + a session's accounting.
|
||||
//!
|
||||
//! ## Safety
|
||||
//! Each session runs inside an isolated sandbox directory
|
||||
//! (`data/agent_sessions/{id}/`). The SDK permission callback ([`build_can_use_tool`])
|
||||
//! enforces two gates on every tool: (1) path arguments must stay within the
|
||||
//! sandbox, and (2) writes/executes require explicit user approval unless the
|
||||
//! session mode opts out. The default mode requires approval for anything that
|
||||
//! mutates state or runs a command.
|
||||
|
||||
use crate::auth::jwt::validate_jwt;
|
||||
use crate::services::agent_workbench::{
|
||||
approval::{self, ApprovalOutcome, SettledState},
|
||||
bash_sandbox::SandboxedBashTool,
|
||||
permission::{classify_tool_risk, decide_permission, PermissionAction, PermissionMode},
|
||||
sandbox,
|
||||
session_store::{self, UsageDelta},
|
||||
};
|
||||
use crate::AppState;
|
||||
use axum::{
|
||||
extract::{DefaultBodyLimit, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::sse::{Event, Sse},
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use futures::stream::Stream;
|
||||
use open_agent_sdk::api::openai::OpenAIProvider;
|
||||
use open_agent_sdk::types::{ContentBlock, PermissionDecision};
|
||||
use open_agent_sdk::{Agent, AgentOptions, LLMProvider, SDKMessage};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use std::convert::Infallible;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
|
||||
use tokio_stream::wrappers::UnboundedReceiverStream;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Max agentic turns per chat request (bounds runaway loops / cost).
|
||||
const MAX_TURNS: u32 = 20;
|
||||
/// How long a tool approval waits for the user before auto-denying.
|
||||
const APPROVAL_TIMEOUT: Duration = Duration::from_secs(300);
|
||||
/// Body limit for a chat request (images/attachments can be large).
|
||||
const AGENT_CHAT_BODY_LIMIT_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
const SYSTEM_PROMPT: &str = "You are BoosAPI's built-in engineering agent. You can read and write files \
|
||||
and run shell commands, but only inside the current working directory, which is an \
|
||||
isolated sandbox for this session. Be concise. Explain what you are about to do before \
|
||||
using a tool that writes files or runs commands. Never attempt to access paths outside \
|
||||
the working directory.";
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/agent/chat", post(chat_handler))
|
||||
.route("/api/agent/approve", post(approve_handler))
|
||||
.route("/api/agent/status", get(status_handler))
|
||||
.layer(DefaultBodyLimit::max(AGENT_CHAT_BODY_LIMIT_BYTES))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatRequest {
|
||||
/// The user's prompt for this turn.
|
||||
message: String,
|
||||
/// Existing session id to continue; a new one is minted when absent.
|
||||
#[serde(default)]
|
||||
session_id: Option<String>,
|
||||
/// Aggregated model id; falls back to the configured agent model.
|
||||
#[serde(default)]
|
||||
model: Option<String>,
|
||||
/// Permission mode: "default" | "accept_edits" | "full_access".
|
||||
#[serde(default)]
|
||||
permission_mode: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ApproveRequest {
|
||||
request_id: String,
|
||||
approved: bool,
|
||||
}
|
||||
|
||||
/// Resolved caller identity for an agent request.
|
||||
#[derive(Debug, Clone)]
|
||||
struct AgentAuth {
|
||||
/// `None` for admin/global-token callers; `Some(id)` for user JWTs.
|
||||
user_id: Option<i64>,
|
||||
}
|
||||
|
||||
fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
|
||||
headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.map(str::trim)
|
||||
.filter(|t| !t.is_empty())
|
||||
}
|
||||
|
||||
/// Authenticate the caller. Admin token / proxy token map to an admin
|
||||
/// (user-less) session; a valid user/admin JWT maps to that user's id.
|
||||
async fn resolve_auth(
|
||||
headers: &HeaderMap,
|
||||
state: &AppState,
|
||||
) -> Result<AgentAuth, (StatusCode, Json<Value>)> {
|
||||
let unauthorized = || {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
Json(json!({ "error": "Unauthorized" })),
|
||||
)
|
||||
};
|
||||
|
||||
let token = extract_bearer_token(headers).ok_or_else(unauthorized)?;
|
||||
|
||||
if token == state.config.auth.auth_token || token == state.runtime_proxy_token() {
|
||||
return Ok(AgentAuth { user_id: None });
|
||||
}
|
||||
|
||||
let claims = validate_jwt(&state.config.auth.jwt_secret, token).map_err(|_| unauthorized())?;
|
||||
if claims.role != "admin" && claims.role != "user" {
|
||||
return Err((StatusCode::FORBIDDEN, Json(json!({ "error": "Forbidden" }))));
|
||||
}
|
||||
let user_id = claims
|
||||
.sub
|
||||
.parse::<i64>()
|
||||
.map_err(|_| (StatusCode::UNAUTHORIZED, Json(json!({ "error": "Invalid user id" }))))?;
|
||||
Ok(AgentAuth { user_id: Some(user_id) })
|
||||
}
|
||||
|
||||
/// Send one SSE event; returns false if the receiver (client) has gone away.
|
||||
fn send_sse(tx: &UnboundedSender<Result<Event, Infallible>>, payload: Value) -> bool {
|
||||
tx.send(Ok(Event::default().data(payload.to_string()))).is_ok()
|
||||
}
|
||||
|
||||
/// Resolve BoosAPI's own base URL (no `/v1` suffix — the OpenAI provider appends
|
||||
/// `/v1/chat/completions`).
|
||||
fn boosapi_base(state: &AppState) -> String {
|
||||
let host = match state.config.server.host.as_str() {
|
||||
"0.0.0.0" | "::" => "127.0.0.1",
|
||||
other => other,
|
||||
};
|
||||
format!("http://{}:{}", host, state.config.server.port)
|
||||
}
|
||||
|
||||
/// Build the SDK permission callback that enforces sandbox path safety and the
|
||||
/// approval flow. The callback runs inside the agent loop for every tool call.
|
||||
///
|
||||
/// Order of checks:
|
||||
/// 1. Sandbox: any path argument that escapes `cwd` is hard-denied immediately.
|
||||
/// 2. Risk + mode: read-only (and, per mode, writes) auto-allow; otherwise emit
|
||||
/// an `approval_request` SSE event and block until the user responds (or the
|
||||
/// approval times out, which denies).
|
||||
fn build_can_use_tool(
|
||||
sse_tx: UnboundedSender<Result<Event, Infallible>>,
|
||||
cwd: std::path::PathBuf,
|
||||
mode: PermissionMode,
|
||||
) -> open_agent_sdk::types::CanUseToolFn {
|
||||
Arc::new(move |tool_name: &str, input: &Value| {
|
||||
let sse_tx = sse_tx.clone();
|
||||
let cwd = cwd.clone();
|
||||
let tool_name = tool_name.to_string();
|
||||
let input = input.clone();
|
||||
Box::pin(async move {
|
||||
// Gate 1: sandbox path safety (always enforced, even in full access).
|
||||
if let Err(reason) = sandbox::check_tool_paths(&tool_name, &input, &cwd) {
|
||||
send_sse(
|
||||
&sse_tx,
|
||||
json!({
|
||||
"type": "tool_denied",
|
||||
"data": { "name": tool_name, "reason": reason }
|
||||
}),
|
||||
);
|
||||
return PermissionDecision::Deny(reason);
|
||||
}
|
||||
|
||||
// Gate 2: risk + permission mode.
|
||||
let risk = classify_tool_risk(&tool_name);
|
||||
match decide_permission(mode, risk) {
|
||||
PermissionAction::AutoAllow => PermissionDecision::Allow,
|
||||
PermissionAction::RequireApproval => {
|
||||
let request_id = Uuid::new_v4().to_string();
|
||||
// `_guard` removes the registry entry when this future ends
|
||||
// (normal completion or drop on session abort) — no leak.
|
||||
let (entry, _guard) = approval::register(&request_id);
|
||||
send_sse(
|
||||
&sse_tx,
|
||||
json!({
|
||||
"type": "approval_request",
|
||||
"data": {
|
||||
"request_id": request_id,
|
||||
"name": tool_name,
|
||||
"input": input,
|
||||
"risk": format!("{risk:?}"),
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
// Single source of truth: the settled atomic state decides
|
||||
// allow vs deny. `resolve` reports back to the frontend
|
||||
// whether its CAS won, so the two can never disagree.
|
||||
let settled = approval::wait_for_decision(&entry, APPROVAL_TIMEOUT).await;
|
||||
let (decision, wire) = match settled {
|
||||
SettledState::Approved => (PermissionDecision::Allow, "approved"),
|
||||
SettledState::Denied => {
|
||||
(PermissionDecision::Deny("Denied by user".to_string()), "denied")
|
||||
}
|
||||
SettledState::TimedOut => (
|
||||
PermissionDecision::Deny("Approval timed out".to_string()),
|
||||
"timed_out",
|
||||
),
|
||||
};
|
||||
// Tell the client how this prompt actually resolved, so a
|
||||
// late click (which the /approve call already reported as
|
||||
// ok:false) can dismiss the modal with the real outcome.
|
||||
send_sse(
|
||||
&sse_tx,
|
||||
json!({
|
||||
"type": "approval_resolved",
|
||||
"data": { "request_id": request_id, "decision": wire }
|
||||
}),
|
||||
);
|
||||
decision
|
||||
}
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
async fn chat_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<ChatRequest>,
|
||||
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (StatusCode, Json<Value>)> {
|
||||
let auth = resolve_auth(&headers, &state).await?;
|
||||
|
||||
if body.message.trim().is_empty() {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "error": "message is required" })),
|
||||
));
|
||||
}
|
||||
|
||||
let session_id = body
|
||||
.session_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| Uuid::new_v4().to_string());
|
||||
let mode = PermissionMode::from_wire(body.permission_mode.as_deref().unwrap_or("default"));
|
||||
let model = body
|
||||
.model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
|
||||
|
||||
// Per-session sandbox working directory.
|
||||
let cwd = match sandbox::ensure_session_dir(&session_id) {
|
||||
Ok(dir) => dir,
|
||||
Err(e) => {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({ "error": format!("failed to create session sandbox: {e}") })),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
// Persist (idempotently) the session row before streaming.
|
||||
if let Err(e) = session_store::ensure_session(
|
||||
&state.db,
|
||||
&session_id,
|
||||
auth.user_id,
|
||||
&model,
|
||||
&format!("{mode:?}").to_lowercase(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(%e, "failed to ensure agent session row");
|
||||
}
|
||||
|
||||
// Route the agent's LLM calls back through BoosAPI's own aggregation using
|
||||
// the built-in downstream key.
|
||||
let base_url = boosapi_base(&state);
|
||||
let internal_key = state.builtin_keys.comfyui_agent.1.clone();
|
||||
|
||||
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
|
||||
|
||||
let db = state.db.clone();
|
||||
let prompt = body.message;
|
||||
let session_for_event = session_id.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
// First event: hand the (possibly newly minted) session id to the client.
|
||||
send_sse(
|
||||
&tx,
|
||||
json!({ "type": "session", "data": { "session_id": session_for_event, "model": model } }),
|
||||
);
|
||||
|
||||
let provider: Arc<dyn LLMProvider> = Arc::new(OpenAIProvider::new(
|
||||
reqwest::Client::new(),
|
||||
internal_key,
|
||||
base_url,
|
||||
Default::default(),
|
||||
));
|
||||
|
||||
let can_use_tool = build_can_use_tool(tx.clone(), cwd.clone(), mode);
|
||||
|
||||
// Replace the SDK's built-in Bash (a bare subprocess) with an
|
||||
// OS-sandboxed one confined to this session's directory. Registered as
|
||||
// a custom tool, it overrides the built-in by name; review modes run it
|
||||
// under bwrap, FullAccess runs it unconfined with audit logging.
|
||||
let sandboxed_bash: Arc<dyn open_agent_sdk::types::Tool> =
|
||||
Arc::new(SandboxedBashTool::new(cwd.clone(), mode));
|
||||
|
||||
let agent = Agent::new(AgentOptions {
|
||||
provider: Some(provider),
|
||||
model: Some(model.clone()),
|
||||
cwd: Some(cwd.to_string_lossy().to_string()),
|
||||
system_prompt: Some(SYSTEM_PROMPT.to_string()),
|
||||
max_turns: Some(MAX_TURNS),
|
||||
can_use_tool: Some(can_use_tool),
|
||||
custom_tools: vec![sandboxed_bash],
|
||||
session_id: Some(session_id.clone()),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
|
||||
let mut agent = match agent {
|
||||
Ok(a) => a,
|
||||
Err(e) => {
|
||||
send_sse(&tx, json!({ "type": "error", "data": { "message": e } }));
|
||||
send_sse(&tx, json!({ "type": "done" }));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let (mut rx_msg, handle) = agent.query(&prompt).await;
|
||||
|
||||
while let Some(event) = rx_msg.recv().await {
|
||||
// Stop early if the client disconnected.
|
||||
if tx.is_closed() {
|
||||
break;
|
||||
}
|
||||
forward_sdk_event(&tx, &db, &session_id, event).await;
|
||||
}
|
||||
let _ = handle.await;
|
||||
agent.close().await;
|
||||
|
||||
send_sse(&tx, json!({ "type": "done" }));
|
||||
});
|
||||
|
||||
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
|
||||
}
|
||||
|
||||
/// Translate one SDK message into SSE event(s), and persist usage on Result.
|
||||
async fn forward_sdk_event(
|
||||
tx: &UnboundedSender<Result<Event, Infallible>>,
|
||||
db: &sqlx::SqlitePool,
|
||||
session_id: &str,
|
||||
event: SDKMessage,
|
||||
) {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
for block in &message.content {
|
||||
match block {
|
||||
ContentBlock::Text { text } if !text.is_empty() => {
|
||||
send_sse(tx, json!({ "type": "text", "data": { "text": text } }));
|
||||
}
|
||||
ContentBlock::ToolUse { id, name, input } => {
|
||||
send_sse(
|
||||
tx,
|
||||
json!({ "type": "tool_use", "data": { "id": id, "name": name, "input": input } }),
|
||||
);
|
||||
}
|
||||
ContentBlock::Thinking { thinking, .. } if !thinking.is_empty() => {
|
||||
send_sse(tx, json!({ "type": "thinking", "data": { "thinking": thinking } }));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
SDKMessage::ToolResult { tool_use_id, tool_name, content, is_error } => {
|
||||
send_sse(
|
||||
tx,
|
||||
json!({
|
||||
"type": "tool_result",
|
||||
"data": { "tool_use_id": tool_use_id, "name": tool_name, "content": content, "is_error": is_error }
|
||||
}),
|
||||
);
|
||||
}
|
||||
SDKMessage::Result { text, usage, num_turns, cost_usd, .. } => {
|
||||
// Persist this turn's accounting (user prompt + assistant reply ≈ 2 msgs).
|
||||
let delta = UsageDelta {
|
||||
input_tokens: usage.input_tokens as i64,
|
||||
output_tokens: usage.output_tokens as i64,
|
||||
cost_usd,
|
||||
num_turns: num_turns as i64,
|
||||
};
|
||||
if let Err(e) = session_store::record_turn(db, session_id, delta, 2).await {
|
||||
tracing::warn!(%e, "failed to record agent turn usage");
|
||||
}
|
||||
send_sse(
|
||||
tx,
|
||||
json!({
|
||||
"type": "result",
|
||||
"data": {
|
||||
"text": text,
|
||||
"num_turns": num_turns,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
"total_tokens": usage.input_tokens + usage.output_tokens,
|
||||
"cost_usd": cost_usd,
|
||||
}
|
||||
}),
|
||||
);
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
send_sse(tx, json!({ "type": "error", "data": { "message": message } }));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
async fn approve_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<ApproveRequest>,
|
||||
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
|
||||
// Auth required, but any authenticated caller may answer the prompt for
|
||||
// their own running session (approvals are keyed by an unguessable uuid).
|
||||
resolve_auth(&headers, &state).await?;
|
||||
let outcome = if body.approved { ApprovalOutcome::Approve } else { ApprovalOutcome::Deny };
|
||||
let delivered = approval::resolve(&body.request_id, outcome);
|
||||
Ok(Json(json!({ "ok": delivered })))
|
||||
}
|
||||
|
||||
async fn status_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
axum::extract::Query(q): axum::extract::Query<StatusQuery>,
|
||||
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
|
||||
let auth = resolve_auth(&headers, &state).await?;
|
||||
let session = match q.session_id.as_deref().filter(|s| !s.is_empty()) {
|
||||
Some(sid) => session_store::get_session(&state.db, sid, auth.user_id)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({ "error": e.to_string() })),
|
||||
)
|
||||
})?,
|
||||
None => None,
|
||||
};
|
||||
Ok(Json(json!({
|
||||
"ok": true,
|
||||
"engine": "open-agent-sdk",
|
||||
"default_model": state.config.comfyui.agent_model,
|
||||
"session": session,
|
||||
})))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct StatusQuery {
|
||||
#[serde(default)]
|
||||
session_id: Option<String>,
|
||||
}
|
||||
@@ -22,6 +22,7 @@ const _: &[u8] = include_bytes!(concat!(
|
||||
));
|
||||
|
||||
mod admin;
|
||||
mod agent;
|
||||
mod auth;
|
||||
mod checkin;
|
||||
mod comfy_cloud_adapter;
|
||||
@@ -57,6 +58,7 @@ pub fn build_router(state: AppState) -> Router {
|
||||
.merge(comfy_cloud_adapter::routes())
|
||||
.merge(money_printer::routes())
|
||||
.merge(comfyui_agent::routes())
|
||||
.merge(agent::routes())
|
||||
.merge(digital_human::routes())
|
||||
.merge(comfyui::routes())
|
||||
.merge(crate::workflow::routes::routes())
|
||||
|
||||
@@ -0,0 +1,241 @@
|
||||
//! Pending-approval registry for the Agent Workbench.
|
||||
//!
|
||||
//! When a session runs in a mode that requires approval for a risky tool, the
|
||||
//! permission callback registers a pending approval, emits an `approval_request`
|
||||
//! SSE event, and waits for either the user's decision or a timeout. The
|
||||
//! `POST /api/agent/approve` endpoint resolves the pending entry.
|
||||
//!
|
||||
//! ## Single source of truth
|
||||
//!
|
||||
//! Each entry carries an [`AtomicU8`] state machine (`Pending → Approved /
|
||||
//! Denied / TimedOut`) plus a [`Notify`]. The decision is made by exactly one
|
||||
//! `compare_exchange` from `Pending`: whichever side wins (the approver or the
|
||||
//! timeout) settles the state; everyone else observes a losing CAS and no-ops.
|
||||
//!
|
||||
//! This closes the race that a bare `oneshot` channel had, where the HTTP
|
||||
//! response (derived from `Sender::send`'s bool) and the loop's allow/deny
|
||||
//! decision (derived from which `select`/`timeout` arm fired) were two
|
||||
//! independent signals that could disagree. Now both read the same atomic:
|
||||
//! `resolve` reports back whether *its* CAS won, and the callback decides purely
|
||||
//! from the settled state — so the frontend's `ok` and the executed decision can
|
||||
//! never contradict each other.
|
||||
|
||||
use once_cell::sync::Lazy;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU8, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::Notify;
|
||||
|
||||
/// Settled decision states. `Pending` is the only non-terminal value.
|
||||
const PENDING: u8 = 0;
|
||||
const APPROVED: u8 = 1;
|
||||
const DENIED: u8 = 2;
|
||||
const TIMED_OUT: u8 = 3;
|
||||
|
||||
/// User's decision on a pending tool approval.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ApprovalOutcome {
|
||||
Approve,
|
||||
Deny,
|
||||
}
|
||||
|
||||
/// Final, settled state of an approval — the single source of truth read by the
|
||||
/// permission callback to decide allow vs deny.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum SettledState {
|
||||
Approved,
|
||||
Denied,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
/// Shared approval entry: an atomic state machine plus a wakeup primitive.
|
||||
#[derive(Debug)]
|
||||
pub struct ApprovalEntry {
|
||||
state: AtomicU8,
|
||||
notify: Notify,
|
||||
}
|
||||
|
||||
impl ApprovalEntry {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
state: AtomicU8::new(PENDING),
|
||||
notify: Notify::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to settle this entry to a terminal state from `Pending`.
|
||||
/// Returns `true` iff this caller's CAS won the race.
|
||||
fn try_settle(&self, target: u8) -> bool {
|
||||
self.state
|
||||
.compare_exchange(PENDING, target, Ordering::SeqCst, Ordering::SeqCst)
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
fn settled(&self) -> Option<SettledState> {
|
||||
match self.state.load(Ordering::SeqCst) {
|
||||
APPROVED => Some(SettledState::Approved),
|
||||
DENIED => Some(SettledState::Denied),
|
||||
TIMED_OUT => Some(SettledState::TimedOut),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static REGISTRY: Lazy<Mutex<HashMap<String, Arc<ApprovalEntry>>>> =
|
||||
Lazy::new(|| Mutex::new(HashMap::new()));
|
||||
|
||||
/// RAII guard that removes its entry from the registry when dropped, so a
|
||||
/// pending approval is cleaned up whether the callback finishes normally or its
|
||||
/// future is dropped (session aborted/disconnected). No leak on any path.
|
||||
pub struct PendingGuard {
|
||||
request_id: String,
|
||||
}
|
||||
|
||||
impl Drop for PendingGuard {
|
||||
fn drop(&mut self) {
|
||||
REGISTRY.lock().unwrap().remove(&self.request_id);
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a new pending approval. Returns the shared entry (await its
|
||||
/// [`wait_settled`](ApprovalEntry)/notify and read its settled state) plus a
|
||||
/// [`PendingGuard`] that the caller must hold for the lifetime of the wait; when
|
||||
/// the guard drops, the registry entry is removed.
|
||||
pub fn register(request_id: &str) -> (Arc<ApprovalEntry>, PendingGuard) {
|
||||
let entry = Arc::new(ApprovalEntry::new());
|
||||
REGISTRY
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(request_id.to_string(), entry.clone());
|
||||
(
|
||||
entry,
|
||||
PendingGuard {
|
||||
request_id: request_id.to_string(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Wait until the entry is settled or `timeout` elapses, then return the settled
|
||||
/// state. On timeout this claims `TimedOut` (if still pending) so the result is
|
||||
/// always a terminal state read from the single atomic.
|
||||
pub async fn wait_for_decision(
|
||||
entry: &ApprovalEntry,
|
||||
timeout: std::time::Duration,
|
||||
) -> SettledState {
|
||||
// Create the notified future *before* checking state to avoid a lost wakeup.
|
||||
let notified = entry.notify.notified();
|
||||
|
||||
// Fast path: already settled by a resolve that raced ahead of us.
|
||||
if let Some(s) = entry.settled() {
|
||||
return s;
|
||||
}
|
||||
|
||||
if tokio::time::timeout(timeout, notified).await.is_err() {
|
||||
// Timer fired first — claim TimedOut if nobody settled in the meantime.
|
||||
entry.try_settle(TIMED_OUT);
|
||||
}
|
||||
|
||||
// Single source of truth: whatever the atomic settled to. If a wakeup was
|
||||
// lost but the state is set, we still read the correct decision; if somehow
|
||||
// still pending (shouldn't happen — we just CAS'd), treat as timed out.
|
||||
entry.settled().unwrap_or(SettledState::TimedOut)
|
||||
}
|
||||
|
||||
/// Resolve a pending approval with the user's decision.
|
||||
///
|
||||
/// Returns `true` iff this call's `compare_exchange` won — i.e. it is the call
|
||||
/// that actually settled the decision. Returns `false` truthfully when the id is
|
||||
/// unknown (already cleaned up) or when the timeout (or another resolve) already
|
||||
/// claimed the entry, so the frontend is never told its vote took effect when it
|
||||
/// did not.
|
||||
pub fn resolve(request_id: &str, outcome: ApprovalOutcome) -> bool {
|
||||
let entry = REGISTRY.lock().unwrap().get(request_id).cloned();
|
||||
let Some(entry) = entry else {
|
||||
return false;
|
||||
};
|
||||
let target = match outcome {
|
||||
ApprovalOutcome::Approve => APPROVED,
|
||||
ApprovalOutcome::Deny => DENIED,
|
||||
};
|
||||
if entry.try_settle(target) {
|
||||
entry.notify.notify_one();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
|
||||
#[tokio::test]
|
||||
async fn approve_before_wait_is_observed() {
|
||||
let (entry, _guard) = register("r1");
|
||||
assert!(resolve("r1", ApprovalOutcome::Approve));
|
||||
let s = wait_for_decision(&entry, Duration::from_secs(5)).await;
|
||||
assert_eq!(s, SettledState::Approved);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deny_is_observed() {
|
||||
let (entry, _guard) = register("r2");
|
||||
assert!(resolve("r2", ApprovalOutcome::Deny));
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_secs(5)).await,
|
||||
SettledState::Denied
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn timeout_settles_timed_out_and_later_resolve_loses() {
|
||||
let (entry, _guard) = register("r3");
|
||||
// Very short timeout → times out with no resolve.
|
||||
let s = wait_for_decision(&entry, Duration::from_millis(20)).await;
|
||||
assert_eq!(s, SettledState::TimedOut);
|
||||
// A late approve must report false: the CAS already went to TimedOut.
|
||||
assert!(!resolve("r3", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn first_cas_wins_second_resolve_loses() {
|
||||
let (_entry, _guard) = register("r4");
|
||||
assert!(resolve("r4", ApprovalOutcome::Approve));
|
||||
assert!(!resolve("r4", ApprovalOutcome::Deny));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_unknown_id_is_false() {
|
||||
assert!(!resolve("nope", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn guard_drop_removes_entry() {
|
||||
{
|
||||
let (_entry, _guard) = register("r5");
|
||||
} // guard dropped here
|
||||
// Entry gone → resolve can't find it.
|
||||
assert!(!resolve("r5", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_resolve_and_timeout_never_both_win() {
|
||||
// Stress the boundary: race a resolve against a near-zero timeout many
|
||||
// times and assert the frontend-ok and the settled decision always agree.
|
||||
for i in 0..2000 {
|
||||
let id = format!("race-{i}");
|
||||
let (entry, _guard) = register(&id);
|
||||
let id2 = id.clone();
|
||||
let resolver = tokio::spawn(async move { resolve(&id2, ApprovalOutcome::Approve) });
|
||||
let decided = wait_for_decision(&entry, Duration::from_nanos(1)).await;
|
||||
let won = resolver.await.unwrap();
|
||||
// The single-atomic invariant: ok==true <=> decision==Approved.
|
||||
match (won, decided) {
|
||||
(true, SettledState::Approved) => {}
|
||||
(false, SettledState::TimedOut) => {}
|
||||
other => panic!("inconsistent outcome: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,307 @@
|
||||
//! OS-level sandboxed Bash tool for the Agent Workbench.
|
||||
//!
|
||||
//! The open-agent-sdk's built-in `Bash` tool runs commands as a plain
|
||||
//! subprocess with only `current_dir` set — no OS boundary. On a server that
|
||||
//! means an agent's shell can `cat /etc/shadow`, `cd /`, write anywhere the
|
||||
//! process uid can, and reach the network. This tool replaces it (same name +
|
||||
//! schema, registered via `custom_tools` so it overrides the built-in) and
|
||||
//! confines the command with [bubblewrap](https://github.com/containers/bubblewrap):
|
||||
//!
|
||||
//! - **Default / AcceptEdits** (review modes): the command runs inside a
|
||||
//! bwrap namespace with the system mounted read-only, the session sandbox
|
||||
//! directory bind-mounted read-write at its real path, `/etc`, `/root`,
|
||||
//! `/home` and every other host path *absent*, and the network unshared.
|
||||
//! Filesystem reads/writes outside the session dir and all network egress are
|
||||
//! therefore impossible at the kernel level, not merely discouraged.
|
||||
//! - **FullAccess** (explicitly dangerous): the command runs unconfined, the
|
||||
//! same behaviour the bare SDK had, but every command is audit-logged with a
|
||||
//! loud warning.
|
||||
//!
|
||||
//! **Fail-closed:** in a review mode, if `bwrap` is unavailable the command is
|
||||
//! refused rather than silently run unconfined — we never fabricate a boundary.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use tokio::process::Command;
|
||||
|
||||
use crate::services::agent_workbench::permission::PermissionMode;
|
||||
use open_agent_sdk::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
const DEFAULT_TIMEOUT_MS: u64 = 120_000;
|
||||
const MAX_TIMEOUT_MS: u64 = 600_000;
|
||||
const MAX_OUTPUT_SIZE: usize = 100_000;
|
||||
|
||||
/// Read-only host paths exposed inside the sandbox so common tooling works.
|
||||
/// Deliberately excludes `/etc`, `/root`, `/home`, `/var`, `/proc`-host, etc.
|
||||
const RO_SYSTEM_PATHS: &[&str] = &["/usr", "/bin", "/sbin", "/lib", "/lib64", "/opt"];
|
||||
|
||||
/// A Bash tool that confines commands to the session sandbox via bubblewrap.
|
||||
pub struct SandboxedBashTool {
|
||||
/// Canonical session sandbox directory (bind-mounted rw inside bwrap).
|
||||
session_dir: PathBuf,
|
||||
/// Session permission mode — decides confined vs unconfined.
|
||||
mode: PermissionMode,
|
||||
/// Resolved `bwrap` binary path, if present on this host.
|
||||
bwrap: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl SandboxedBashTool {
|
||||
pub fn new(session_dir: PathBuf, mode: PermissionMode) -> Self {
|
||||
Self {
|
||||
session_dir,
|
||||
mode,
|
||||
bwrap: find_bwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether this mode confines bash (everything except FullAccess).
|
||||
fn is_confined(&self) -> bool {
|
||||
!matches!(self.mode, PermissionMode::FullAccess)
|
||||
}
|
||||
}
|
||||
|
||||
/// Locate the `bwrap` binary (PATH + well-known locations).
|
||||
fn find_bwrap() -> Option<PathBuf> {
|
||||
for cand in ["/usr/bin/bwrap", "/bin/bwrap", "/usr/local/bin/bwrap"] {
|
||||
let p = Path::new(cand);
|
||||
if p.is_file() {
|
||||
return Some(p.to_path_buf());
|
||||
}
|
||||
}
|
||||
std::env::var_os("PATH").and_then(|paths| {
|
||||
std::env::split_paths(&paths)
|
||||
.map(|d| d.join("bwrap"))
|
||||
.find(|p| p.is_file())
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the bwrap argument vector that confines `command` to `session_dir`.
|
||||
///
|
||||
/// Pure and unit-testable. The session directory is bound read-write at its own
|
||||
/// real absolute path and becomes the working directory; the rest of the
|
||||
/// filesystem is either read-only (system paths) or absent. `--unshare-all`
|
||||
/// (which includes the network namespace) plus `--die-with-parent` and
|
||||
/// `--new-session` complete the isolation.
|
||||
pub fn build_bwrap_args(session_dir: &Path, command: &str) -> Vec<String> {
|
||||
let sd = session_dir.to_string_lossy().to_string();
|
||||
let mut args: Vec<String> = Vec::new();
|
||||
|
||||
for p in RO_SYSTEM_PATHS {
|
||||
// `--ro-bind-try` tolerates paths that don't exist on this host.
|
||||
args.push("--ro-bind-try".into());
|
||||
args.push((*p).into());
|
||||
args.push((*p).into());
|
||||
}
|
||||
// Minimal virtual filesystems.
|
||||
args.push("--dev".into());
|
||||
args.push("/dev".into());
|
||||
args.push("--proc".into());
|
||||
args.push("/proc".into());
|
||||
args.push("--tmpfs".into());
|
||||
args.push("/tmp".into());
|
||||
// The session sandbox, bound rw at its real path, then made the cwd.
|
||||
args.push("--bind".into());
|
||||
args.push(sd.clone());
|
||||
args.push(sd.clone());
|
||||
args.push("--chdir".into());
|
||||
args.push(sd);
|
||||
// Full namespace isolation incl. network; tear down with the parent.
|
||||
args.push("--unshare-all".into());
|
||||
args.push("--die-with-parent".into());
|
||||
args.push("--new-session".into());
|
||||
// Run the command under a shell inside the sandbox.
|
||||
args.push("/bin/bash".into());
|
||||
args.push("-c".into());
|
||||
args.push(command.into());
|
||||
|
||||
args
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SandboxedBashTool {
|
||||
fn name(&self) -> &str {
|
||||
"Bash"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Executes a bash command inside an OS-level sandbox confined to the \
|
||||
session working directory (read-only system, no network) unless the \
|
||||
session is in full-access mode. Use for shell operations."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"command".to_string(),
|
||||
json!({ "type": "string", "description": "The command to execute" }),
|
||||
),
|
||||
(
|
||||
"timeout".to_string(),
|
||||
json!({ "type": "number", "description": "Optional timeout in milliseconds (max 600000)" }),
|
||||
),
|
||||
(
|
||||
"description".to_string(),
|
||||
json!({ "type": "string", "description": "Clear description of what this command does" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["command".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let command = input
|
||||
.get("command")
|
||||
.and_then(|c| c.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'command' field".to_string()))?;
|
||||
|
||||
let timeout_ms = input
|
||||
.get("timeout")
|
||||
.and_then(|t| t.as_u64())
|
||||
.unwrap_or(DEFAULT_TIMEOUT_MS)
|
||||
.min(MAX_TIMEOUT_MS);
|
||||
|
||||
// Build the runner per mode.
|
||||
let (program, args): (String, Vec<String>) = if self.is_confined() {
|
||||
let Some(bwrap) = &self.bwrap else {
|
||||
// Fail closed: never downgrade a review-mode session to an
|
||||
// unconfined shell just because the sandbox tool is missing.
|
||||
return Ok(ToolResult::error(
|
||||
"Bash is disabled: the OS sandbox (bwrap) is not available on \
|
||||
this host, and this session runs in a review mode that requires \
|
||||
confinement. Re-run in full-access mode (explicitly unsafe) or \
|
||||
install bubblewrap."
|
||||
.to_string(),
|
||||
));
|
||||
};
|
||||
(
|
||||
bwrap.to_string_lossy().to_string(),
|
||||
build_bwrap_args(&self.session_dir, command),
|
||||
)
|
||||
} else {
|
||||
// FullAccess: unconfined, but audit every command loudly.
|
||||
tracing::warn!(
|
||||
target: "boosapi::agent::bash",
|
||||
mode = "full_access",
|
||||
command = command,
|
||||
"UNCONFINED agent bash command (full-access mode)"
|
||||
);
|
||||
(
|
||||
"/bin/bash".to_string(),
|
||||
vec!["-c".to_string(), command.to_string()],
|
||||
)
|
||||
};
|
||||
|
||||
// Unconfined mode honours the SDK's working dir; confined mode's cwd is
|
||||
// set inside bwrap via --chdir, so the outer cwd is irrelevant.
|
||||
let outer_cwd = if self.is_confined() {
|
||||
self.session_dir.clone()
|
||||
} else {
|
||||
PathBuf::from(&context.working_dir)
|
||||
};
|
||||
|
||||
let fut = run_capture(&program, &args, &outer_cwd);
|
||||
match tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), fut).await {
|
||||
Ok(Ok((stdout, stderr, code))) => Ok(format_result(stdout, stderr, code)),
|
||||
Ok(Err(e)) => Ok(ToolResult::error(format!("Command failed: {e}"))),
|
||||
Err(_) => Ok(ToolResult::error(format!(
|
||||
"Command timed out after {timeout_ms}ms"
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_capture(
|
||||
program: &str,
|
||||
args: &[String],
|
||||
cwd: &Path,
|
||||
) -> Result<(String, String, i32), std::io::Error> {
|
||||
let output = Command::new(program)
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await?;
|
||||
Ok((
|
||||
String::from_utf8_lossy(&output.stdout).to_string(),
|
||||
String::from_utf8_lossy(&output.stderr).to_string(),
|
||||
output.status.code().unwrap_or(-1),
|
||||
))
|
||||
}
|
||||
|
||||
fn format_result(stdout: String, stderr: String, exit_code: i32) -> ToolResult {
|
||||
let mut result = String::new();
|
||||
if !stdout.is_empty() {
|
||||
result.push_str(&stdout);
|
||||
}
|
||||
if !stderr.is_empty() {
|
||||
if !result.is_empty() {
|
||||
result.push('\n');
|
||||
}
|
||||
result.push_str("STDERR:\n");
|
||||
result.push_str(&stderr);
|
||||
}
|
||||
if result.len() > MAX_OUTPUT_SIZE {
|
||||
result.truncate(MAX_OUTPUT_SIZE);
|
||||
result.push_str("\n... (output truncated)");
|
||||
}
|
||||
if exit_code != 0 {
|
||||
result.push_str(&format!("\n\nExit code: {exit_code}"));
|
||||
}
|
||||
if result.is_empty() {
|
||||
result = "(no output)".to_string();
|
||||
}
|
||||
if exit_code != 0 {
|
||||
ToolResult::error(result)
|
||||
} else {
|
||||
ToolResult::text(result)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn args_bind_session_dir_and_isolate() {
|
||||
let sd = Path::new("/srv/data/agent_sessions/abc");
|
||||
let args = build_bwrap_args(sd, "echo hi");
|
||||
let joined = args.join(" ");
|
||||
// Session dir is bound rw and made the cwd.
|
||||
assert!(joined.contains("--bind /srv/data/agent_sessions/abc /srv/data/agent_sessions/abc"));
|
||||
assert!(joined.contains("--chdir /srv/data/agent_sessions/abc"));
|
||||
// Network + everything else unshared, parent-bound.
|
||||
assert!(args.iter().any(|a| a == "--unshare-all"));
|
||||
assert!(args.iter().any(|a| a == "--die-with-parent"));
|
||||
// /etc is never bound.
|
||||
assert!(!joined.contains("/etc"));
|
||||
// System paths are read-only.
|
||||
assert!(joined.contains("--ro-bind-try /usr /usr"));
|
||||
// The command is the last argument, run under bash -c.
|
||||
assert_eq!(args.last().unwrap(), "echo hi");
|
||||
let n = args.len();
|
||||
assert_eq!(args[n - 3], "/bin/bash");
|
||||
assert_eq!(args[n - 2], "-c");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn confinement_follows_mode() {
|
||||
let sd = PathBuf::from("/tmp/x");
|
||||
assert!(SandboxedBashTool::new(sd.clone(), PermissionMode::Default).is_confined());
|
||||
assert!(SandboxedBashTool::new(sd.clone(), PermissionMode::AcceptEdits).is_confined());
|
||||
assert!(!SandboxedBashTool::new(sd, PermissionMode::FullAccess).is_confined());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn name_overrides_builtin_bash() {
|
||||
let t = SandboxedBashTool::new(PathBuf::from("/tmp/x"), PermissionMode::Default);
|
||||
assert_eq!(t.name(), "Bash");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
//! Agent Workbench: a general-purpose, server-side AI agent exposed over SSE.
|
||||
//!
|
||||
//! The agent loop is provided by the `open-agent-sdk` crate. This module owns
|
||||
//! the BoosAPI-specific concerns layered on top of it:
|
||||
//!
|
||||
//! - [`permission`]: tool risk classification + the per-mode decision matrix.
|
||||
//! - [`sandbox`]: per-session isolated working directory + path validation.
|
||||
//! - [`bash_sandbox`]: OS-level (bubblewrap) confinement for the Bash tool.
|
||||
//! - [`approval`]: process-global registry bridging the SSE approval round-trip.
|
||||
//! - [`session_store`]: persistence of per-session token/cost accounting.
|
||||
//!
|
||||
//! The HTTP surface (SSE chat, approval callback, status) lives in
|
||||
//! `crate::routes::agent`.
|
||||
|
||||
pub mod approval;
|
||||
pub mod bash_sandbox;
|
||||
pub mod permission;
|
||||
pub mod sandbox;
|
||||
pub mod session_store;
|
||||
@@ -0,0 +1,176 @@
|
||||
//! Tool risk classification and permission decision matrix for the Agent Workbench.
|
||||
//!
|
||||
//! Ported and adapted from the AQBot agent crate. Deliberately SDK-free so it
|
||||
//! can be unit-tested in isolation; the route layer maps [`PermissionAction`]
|
||||
//! onto the SDK's `PermissionDecision`.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Risk level of a tool, inferred from its name.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RiskLevel {
|
||||
/// Reads only; never mutates the filesystem or runs commands.
|
||||
ReadOnly,
|
||||
/// Mutates files (write/edit/create/delete/...).
|
||||
Write,
|
||||
/// Executes arbitrary commands (bash/shell/...).
|
||||
Execute,
|
||||
}
|
||||
|
||||
/// Permission mode for an agent session.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum PermissionMode {
|
||||
/// Standard review: read auto-allowed, write/execute require approval.
|
||||
Default,
|
||||
/// Auto-accept edits: read + write auto-allowed, execute requires approval.
|
||||
AcceptEdits,
|
||||
/// Everything auto-allowed (no approvals).
|
||||
FullAccess,
|
||||
}
|
||||
|
||||
impl PermissionMode {
|
||||
/// Parse a mode from a wire string. Unknown values fall back to the safest
|
||||
/// mode ([`PermissionMode::Default`]).
|
||||
pub fn from_wire(s: &str) -> Self {
|
||||
match s {
|
||||
"accept_edits" | "acceptEdits" => Self::AcceptEdits,
|
||||
"full_access" | "fullAccess" => Self::FullAccess,
|
||||
_ => Self::Default,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the permission system decides for a given tool invocation.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PermissionAction {
|
||||
/// Allow without prompting the user.
|
||||
AutoAllow,
|
||||
/// Pause and require explicit user approval before running.
|
||||
RequireApproval,
|
||||
}
|
||||
|
||||
/// Classify a tool's risk level from its name (case-insensitive).
|
||||
///
|
||||
/// Conservative by construction: anything that looks like it executes a command
|
||||
/// is [`RiskLevel::Execute`], anything that looks like it mutates state is
|
||||
/// [`RiskLevel::Write`], and everything else is [`RiskLevel::ReadOnly`].
|
||||
pub fn classify_tool_risk(tool_name: &str) -> RiskLevel {
|
||||
let n = tool_name.to_lowercase();
|
||||
|
||||
if matches!(n.as_str(), "bash" | "shell" | "run_command" | "execute")
|
||||
|| n.contains("exec")
|
||||
|| n.contains("bash")
|
||||
|| n.contains("shell")
|
||||
|| n.contains("command")
|
||||
|| n.contains("run")
|
||||
|| n.contains("terminal")
|
||||
|| n.contains("spawn")
|
||||
|| n.contains("process")
|
||||
{
|
||||
return RiskLevel::Execute;
|
||||
}
|
||||
|
||||
if matches!(
|
||||
n.as_str(),
|
||||
"write" | "edit" | "create" | "delete" | "rename" | "patch" | "move" | "mkdir" | "remove"
|
||||
) || n.contains("write")
|
||||
|| n.contains("edit")
|
||||
|| n.contains("create")
|
||||
|| n.contains("delete")
|
||||
|| n.contains("patch")
|
||||
|| n.contains("apply")
|
||||
{
|
||||
return RiskLevel::Write;
|
||||
}
|
||||
|
||||
RiskLevel::ReadOnly
|
||||
}
|
||||
|
||||
/// Decide what to do for a tool invocation given the session mode and the tool's
|
||||
/// risk level.
|
||||
pub fn decide_permission(mode: PermissionMode, risk: RiskLevel) -> PermissionAction {
|
||||
match (mode, risk) {
|
||||
(PermissionMode::Default, RiskLevel::ReadOnly) => PermissionAction::AutoAllow,
|
||||
(PermissionMode::Default, _) => PermissionAction::RequireApproval,
|
||||
|
||||
(PermissionMode::AcceptEdits, RiskLevel::ReadOnly) => PermissionAction::AutoAllow,
|
||||
(PermissionMode::AcceptEdits, RiskLevel::Write) => PermissionAction::AutoAllow,
|
||||
(PermissionMode::AcceptEdits, RiskLevel::Execute) => PermissionAction::RequireApproval,
|
||||
|
||||
(PermissionMode::FullAccess, _) => PermissionAction::AutoAllow,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn classifies_execute_tools() {
|
||||
for t in ["Bash", "bash", "shell", "run_command", "execute", "RunInTerminal"] {
|
||||
assert_eq!(classify_tool_risk(t), RiskLevel::Execute, "tool={t}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_write_tools() {
|
||||
for t in ["Write", "Edit", "create_file", "delete_file", "patch", "ApplyPatch"] {
|
||||
assert_eq!(classify_tool_risk(t), RiskLevel::Write, "tool={t}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_read_tools() {
|
||||
for t in ["Read", "Glob", "Grep", "WebFetch", "list_dir"] {
|
||||
assert_eq!(classify_tool_risk(t), RiskLevel::ReadOnly, "tool={t}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_mode_requires_approval_for_write_and_execute() {
|
||||
assert_eq!(
|
||||
decide_permission(PermissionMode::Default, RiskLevel::ReadOnly),
|
||||
PermissionAction::AutoAllow
|
||||
);
|
||||
assert_eq!(
|
||||
decide_permission(PermissionMode::Default, RiskLevel::Write),
|
||||
PermissionAction::RequireApproval
|
||||
);
|
||||
assert_eq!(
|
||||
decide_permission(PermissionMode::Default, RiskLevel::Execute),
|
||||
PermissionAction::RequireApproval
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accept_edits_allows_write_but_not_execute() {
|
||||
assert_eq!(
|
||||
decide_permission(PermissionMode::AcceptEdits, RiskLevel::Write),
|
||||
PermissionAction::AutoAllow
|
||||
);
|
||||
assert_eq!(
|
||||
decide_permission(PermissionMode::AcceptEdits, RiskLevel::Execute),
|
||||
PermissionAction::RequireApproval
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn full_access_allows_everything() {
|
||||
for r in [RiskLevel::ReadOnly, RiskLevel::Write, RiskLevel::Execute] {
|
||||
assert_eq!(
|
||||
decide_permission(PermissionMode::FullAccess, r),
|
||||
PermissionAction::AutoAllow
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mode_parsing_falls_back_to_default() {
|
||||
assert_eq!(PermissionMode::from_wire("accept_edits"), PermissionMode::AcceptEdits);
|
||||
assert_eq!(PermissionMode::from_wire("full_access"), PermissionMode::FullAccess);
|
||||
assert_eq!(PermissionMode::from_wire("garbage"), PermissionMode::Default);
|
||||
assert_eq!(PermissionMode::from_wire(""), PermissionMode::Default);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,326 @@
|
||||
//! Per-session filesystem sandbox for the Agent Workbench.
|
||||
//!
|
||||
//! Every agent session is confined to an isolated working directory
|
||||
//! (`data/agent_sessions/{session_id}/`). Tool path arguments are validated to
|
||||
//! resolve *within* that directory before the tool is allowed to run, so one
|
||||
//! tenant's agent can never read or write another tenant's files (or escape to
|
||||
//! the host filesystem via `..`/symlinks/absolute paths).
|
||||
//!
|
||||
//! Adapted from the AQBot agent crate's `security.rs`. SDK-free: the route layer
|
||||
//! turns an [`Err`] into a permission denial.
|
||||
|
||||
use serde_json::Value;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Root directory that holds all per-session sandboxes, relative to the server's
|
||||
/// working directory.
|
||||
pub const SESSIONS_ROOT: &str = "data/agent_sessions";
|
||||
|
||||
/// Compute the sandbox directory for a session id (not yet created on disk).
|
||||
pub fn session_dir(session_id: &str) -> PathBuf {
|
||||
Path::new(SESSIONS_ROOT).join(sanitize_session_id(session_id))
|
||||
}
|
||||
|
||||
/// Create the sandbox directory for a session if missing, returning its
|
||||
/// canonical absolute path (suitable for use as the agent cwd).
|
||||
pub fn ensure_session_dir(session_id: &str) -> std::io::Result<PathBuf> {
|
||||
let dir = session_dir(session_id);
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
dir.canonicalize()
|
||||
}
|
||||
|
||||
/// Strip anything that could let a session id traverse outside SESSIONS_ROOT.
|
||||
/// Keeps only `[A-Za-z0-9_-]`; everything else becomes `_`.
|
||||
fn sanitize_session_id(session_id: &str) -> String {
|
||||
let cleaned: String = session_id
|
||||
.chars()
|
||||
.map(|c| if c.is_ascii_alphanumeric() || c == '-' || c == '_' { c } else { '_' })
|
||||
.collect();
|
||||
if cleaned.is_empty() {
|
||||
"default".to_string()
|
||||
} else {
|
||||
cleaned
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate that `path` (absolute or relative to `cwd`) resolves to a location
|
||||
/// within `cwd`. Returns the resolved path on success, or an error message
|
||||
/// describing the escape attempt.
|
||||
///
|
||||
/// For existing paths the target is canonicalized (resolving symlinks). For
|
||||
/// not-yet-existing paths the nearest existing ancestor is canonicalized and
|
||||
/// checked, so writes to new files inside the sandbox are permitted while writes
|
||||
/// that would escape it are rejected.
|
||||
///
|
||||
/// Symlinks are handled explicitly with `lstat` (not `exists()`, which follows
|
||||
/// links): any component that is a symlink — including a *dangling* one whose
|
||||
/// target does not exist yet — has its target resolved and is rejected if it
|
||||
/// points outside the sandbox. This stops an approved `ln -s /outside x` from
|
||||
/// being written through to escape (a write to `x` would otherwise follow the
|
||||
/// link). Note this closes the escape, not the check-vs-open TOCTOU race; the
|
||||
/// actual open happens inside the tool and would need `O_NOFOLLOW`/`openat2`
|
||||
/// there to be fully atomic.
|
||||
pub fn validate_path_within_cwd(path: &str, cwd: &Path) -> Result<PathBuf, String> {
|
||||
let cwd_canonical = cwd
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("cannot resolve sandbox directory '{}': {e}", cwd.display()))?;
|
||||
|
||||
let target = Path::new(path);
|
||||
let resolved = if target.is_absolute() {
|
||||
target.to_path_buf()
|
||||
} else {
|
||||
cwd_canonical.join(target)
|
||||
};
|
||||
|
||||
// Final component: if it is a symlink (dangling or not), trust only its
|
||||
// target. lstat does not follow the link, so a dangling symlink is detected
|
||||
// here instead of slipping into the not-yet-existing branch below.
|
||||
if let Ok(meta) = std::fs::symlink_metadata(&resolved) {
|
||||
if meta.file_type().is_symlink() {
|
||||
let tgt = symlink_real_target(&resolved)?;
|
||||
return if tgt.starts_with(&cwd_canonical) {
|
||||
Ok(resolved)
|
||||
} else {
|
||||
Err(format!(
|
||||
"path '{path}' is a symlink that escapes the session sandbox"
|
||||
))
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if resolved.exists() {
|
||||
let canonical = resolved
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("cannot resolve path '{path}': {e}"))?;
|
||||
return if canonical.starts_with(&cwd_canonical) {
|
||||
Ok(canonical)
|
||||
} else {
|
||||
Err(format!(
|
||||
"path '{path}' resolves outside the session sandbox"
|
||||
))
|
||||
};
|
||||
}
|
||||
|
||||
// Not-yet-existing path: walk up to the nearest existing ancestor. Use
|
||||
// lstat so an intermediate *dangling* symlink directory is caught (an
|
||||
// existing symlink ancestor with a live target is already caught by
|
||||
// canonicalize, but a dangling one would otherwise be skipped over).
|
||||
let mut ancestor = resolved.parent();
|
||||
while let Some(dir) = ancestor {
|
||||
match std::fs::symlink_metadata(dir) {
|
||||
Ok(meta) if meta.file_type().is_symlink() => {
|
||||
let tgt = symlink_real_target(dir)?;
|
||||
return if tgt.starts_with(&cwd_canonical) {
|
||||
Ok(resolved)
|
||||
} else {
|
||||
Err(format!(
|
||||
"path '{path}' traverses a symlink that escapes the session sandbox"
|
||||
))
|
||||
};
|
||||
}
|
||||
Ok(_) => {
|
||||
// Existing real directory/file: canonicalize and bound-check.
|
||||
let anc_canonical = dir
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("cannot resolve ancestor of '{path}': {e}"))?;
|
||||
return if anc_canonical.starts_with(&cwd_canonical) {
|
||||
Ok(resolved)
|
||||
} else {
|
||||
Err(format!("path '{path}' is outside the session sandbox"))
|
||||
};
|
||||
}
|
||||
// Component does not exist in any form yet: keep walking up.
|
||||
Err(_) => {}
|
||||
}
|
||||
ancestor = dir.parent();
|
||||
}
|
||||
Err(format!("cannot verify path '{path}' against the session sandbox"))
|
||||
}
|
||||
|
||||
/// Resolve where a (possibly dangling) symlink points, as a real absolute path,
|
||||
/// without requiring the target to exist: read the link, make it absolute
|
||||
/// relative to the link's parent, then canonicalize its nearest existing
|
||||
/// ancestor and re-append the missing remainder.
|
||||
fn symlink_real_target(link: &Path) -> Result<PathBuf, String> {
|
||||
let raw = std::fs::read_link(link)
|
||||
.map_err(|e| format!("cannot read symlink '{}': {e}", link.display()))?;
|
||||
let abs = if raw.is_absolute() {
|
||||
raw
|
||||
} else {
|
||||
link.parent().unwrap_or_else(|| Path::new("/")).join(raw)
|
||||
};
|
||||
real_location(&abs)
|
||||
.ok_or_else(|| format!("cannot resolve symlink target '{}'", abs.display()))
|
||||
}
|
||||
|
||||
/// Real absolute location of a path that may not fully exist: canonicalize the
|
||||
/// nearest existing ancestor (resolving any symlinks/`..` in the existing part)
|
||||
/// and re-append the non-existing remainder lexically.
|
||||
fn real_location(p: &Path) -> Option<PathBuf> {
|
||||
if let Ok(c) = p.canonicalize() {
|
||||
return Some(c);
|
||||
}
|
||||
let mut remainder: Vec<std::ffi::OsString> = Vec::new();
|
||||
let mut cur = p.to_path_buf();
|
||||
loop {
|
||||
let parent = cur.parent()?;
|
||||
if let Some(name) = cur.file_name() {
|
||||
remainder.push(name.to_os_string());
|
||||
}
|
||||
if let Ok(mut base) = parent.canonicalize() {
|
||||
while let Some(name) = remainder.pop() {
|
||||
base.push(name);
|
||||
}
|
||||
return Some(base);
|
||||
}
|
||||
cur = parent.to_path_buf();
|
||||
}
|
||||
}
|
||||
|
||||
/// Check whether a tool's path-bearing arguments stay within the sandbox.
|
||||
///
|
||||
/// Returns `Ok(())` when the tool either takes no path arguments or all of its
|
||||
/// path arguments resolve within `cwd`; returns `Err(reason)` on the first
|
||||
/// argument that escapes. Bash/shell are intentionally not path-checked here —
|
||||
/// they are gated by the approval flow and run with the sandbox dir as cwd.
|
||||
pub fn check_tool_paths(tool_name: &str, input: &Value, cwd: &Path) -> Result<(), String> {
|
||||
let n = tool_name.to_lowercase();
|
||||
match n.as_str() {
|
||||
"read" | "read_file" | "write" | "write_file" | "edit" | "edit_file" | "create"
|
||||
| "create_file" | "delete" | "delete_file" | "rename" | "list_dir" | "listdir"
|
||||
| "multiedit" => {
|
||||
check_field(input, &["path", "file_path", "file"], cwd)?;
|
||||
check_field(input, &["new_path", "destination", "dest"], cwd)?;
|
||||
Ok(())
|
||||
}
|
||||
"glob" | "glob_search" | "grep" | "search" | "ripgrep" => {
|
||||
// Optional explicit base path; the pattern itself is relative to cwd.
|
||||
check_field(input, &["path", "dir", "directory"], cwd)?;
|
||||
Ok(())
|
||||
}
|
||||
// Command execution and unknown tools: no path arg to validate here.
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
fn check_field(input: &Value, keys: &[&str], cwd: &Path) -> Result<(), String> {
|
||||
for key in keys {
|
||||
if let Some(p) = input.get(*key).and_then(|v| v.as_str()) {
|
||||
validate_path_within_cwd(p, cwd)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn tmp_sandbox() -> PathBuf {
|
||||
let base = std::env::temp_dir().join(format!("agent_sb_{}", uuid::Uuid::new_v4().simple()));
|
||||
std::fs::create_dir_all(&base).unwrap();
|
||||
base.canonicalize().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_session_id() {
|
||||
assert_eq!(sanitize_session_id("../../etc"), "______etc");
|
||||
assert_eq!(sanitize_session_id("abc-123_DEF"), "abc-123_DEF");
|
||||
assert_eq!(sanitize_session_id(""), "default");
|
||||
assert_eq!(sanitize_session_id("a/b"), "a_b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_path_inside_sandbox() {
|
||||
let sb = tmp_sandbox();
|
||||
assert!(validate_path_within_cwd("notes.txt", &sb).is_ok());
|
||||
assert!(validate_path_within_cwd("sub/dir/new.txt", &sb).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_parent_traversal() {
|
||||
let sb = tmp_sandbox();
|
||||
assert!(validate_path_within_cwd("../escape.txt", &sb).is_err());
|
||||
assert!(validate_path_within_cwd("../../etc/passwd", &sb).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_absolute_outside() {
|
||||
let sb = tmp_sandbox();
|
||||
assert!(validate_path_within_cwd("/etc/passwd", &sb).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_tool_paths_blocks_escaping_write() {
|
||||
let sb = tmp_sandbox();
|
||||
let bad = json!({"path": "../../secret.txt", "content": "x"});
|
||||
assert!(check_tool_paths("write", &bad, &sb).is_err());
|
||||
|
||||
let ok = json!({"path": "out.txt", "content": "x"});
|
||||
assert!(check_tool_paths("write", &ok, &sb).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_tool_paths_ignores_bash() {
|
||||
let sb = tmp_sandbox();
|
||||
let input = json!({"command": "ls -la"});
|
||||
assert!(check_tool_paths("bash", &input, &sb).is_ok());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn rejects_dangling_symlink_escape() {
|
||||
// Models an approved `ln -s /etc/cron.d/pwn evil` inside the sandbox,
|
||||
// whose target does not exist yet. A later Write to `evil` must NOT be
|
||||
// allowed to follow the link out of the sandbox.
|
||||
let sb = tmp_sandbox();
|
||||
std::os::unix::fs::symlink("/etc/cron.d/pwn_dangling", sb.join("evil")).unwrap();
|
||||
let res = validate_path_within_cwd("evil", &sb);
|
||||
assert!(res.is_err(), "dangling symlink escape must be rejected: {res:?}");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn rejects_dangling_symlink_to_absolute_outside() {
|
||||
let sb = tmp_sandbox();
|
||||
std::os::unix::fs::symlink("/tmp/definitely_outside_xyz", sb.join("link")).unwrap();
|
||||
assert!(validate_path_within_cwd("link", &sb).is_err());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn rejects_relative_dangling_symlink_escape() {
|
||||
// `evil -> ../../escape` (target missing) must be rejected too.
|
||||
let sb = tmp_sandbox();
|
||||
std::os::unix::fs::symlink("../../escape_target", sb.join("evil")).unwrap();
|
||||
assert!(validate_path_within_cwd("evil", &sb).is_err());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn allows_dangling_symlink_pointing_inside() {
|
||||
// A symlink whose (missing) target is still inside the sandbox is fine.
|
||||
let sb = tmp_sandbox();
|
||||
std::os::unix::fs::symlink(sb.join("inside_target"), sb.join("ok_link")).unwrap();
|
||||
assert!(validate_path_within_cwd("ok_link", &sb).is_ok());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn rejects_write_through_intermediate_dangling_symlink_dir() {
|
||||
// `evildir -> /etc` (dangling), then write to `evildir/pwn`.
|
||||
let sb = tmp_sandbox();
|
||||
std::os::unix::fs::symlink("/etc/nonexistent_dir_xyz", sb.join("evildir")).unwrap();
|
||||
assert!(validate_path_within_cwd("evildir/pwn", &sb).is_err());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn check_tool_paths_blocks_dangling_symlink_write() {
|
||||
let sb = tmp_sandbox();
|
||||
std::os::unix::fs::symlink("/etc/cron.d/pwn2", sb.join("evil")).unwrap();
|
||||
let bad = json!({"path": "evil", "content": "* * * * * root sh"});
|
||||
assert!(check_tool_paths("write", &bad, &sb).is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
//! Persistence for Agent Workbench session accounting.
|
||||
//!
|
||||
//! Each agent session has one row in `agent_sessions` tracking cumulative token
|
||||
//! usage and cost across its turns. A session is created (idempotently) on its
|
||||
//! first chat request and its counters are advanced after every completed turn,
|
||||
//! using the usage/cost reported by the SDK `Result` event.
|
||||
|
||||
use crate::error::{AppError, AppResult};
|
||||
use serde::Serialize;
|
||||
use sqlx::SqlitePool;
|
||||
|
||||
/// One agent session's persisted accounting state.
|
||||
#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
|
||||
pub struct AgentSession {
|
||||
pub id: String,
|
||||
pub user_id: Option<i64>,
|
||||
pub title: Option<String>,
|
||||
pub model: String,
|
||||
pub permission_mode: String,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub total_tokens: i64,
|
||||
pub cost_usd: f64,
|
||||
pub num_turns: i64,
|
||||
pub message_count: i64,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
/// Token/cost delta produced by one completed agent turn.
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct UsageDelta {
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cost_usd: f64,
|
||||
pub num_turns: i64,
|
||||
}
|
||||
|
||||
/// Create the session row if it does not exist yet. Idempotent: an existing
|
||||
/// session keeps its counters and metadata untouched.
|
||||
pub async fn ensure_session(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
user_id: Option<i64>,
|
||||
model: &str,
|
||||
permission_mode: &str,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"INSERT INTO agent_sessions (id, user_id, model, permission_mode) \
|
||||
VALUES (?, ?, ?, ?) \
|
||||
ON CONFLICT(id) DO NOTHING",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(user_id)
|
||||
.bind(model)
|
||||
.bind(permission_mode)
|
||||
.execute(db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Advance a session's cumulative counters by one turn's usage delta and bump
|
||||
/// `updated_at`. The user-visible message count is incremented by `messages`
|
||||
/// (typically 2: one user prompt + one assistant reply).
|
||||
pub async fn record_turn(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
delta: UsageDelta,
|
||||
messages: i64,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE agent_sessions SET \
|
||||
input_tokens = input_tokens + ?, \
|
||||
output_tokens = output_tokens + ?, \
|
||||
total_tokens = total_tokens + ?, \
|
||||
cost_usd = cost_usd + ?, \
|
||||
num_turns = num_turns + ?, \
|
||||
message_count = message_count + ?, \
|
||||
updated_at = datetime('now') \
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(delta.input_tokens)
|
||||
.bind(delta.output_tokens)
|
||||
.bind(delta.input_tokens + delta.output_tokens)
|
||||
.bind(delta.cost_usd)
|
||||
.bind(delta.num_turns)
|
||||
.bind(messages)
|
||||
.bind(session_id)
|
||||
.execute(db)
|
||||
.await
|
||||
.map_err(AppError::Database)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fetch a single session by id (scoped to `user_id` when provided, so a user
|
||||
/// cannot read another tenant's session). Admin callers pass `None`.
|
||||
pub async fn get_session(
|
||||
db: &SqlitePool,
|
||||
session_id: &str,
|
||||
user_id: Option<i64>,
|
||||
) -> AppResult<Option<AgentSession>> {
|
||||
let row = match user_id {
|
||||
Some(uid) => {
|
||||
sqlx::query_as::<_, AgentSession>(
|
||||
"SELECT * FROM agent_sessions WHERE id = ? AND user_id = ?",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(uid)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
}
|
||||
None => {
|
||||
sqlx::query_as::<_, AgentSession>("SELECT * FROM agent_sessions WHERE id = ?")
|
||||
.bind(session_id)
|
||||
.fetch_optional(db)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(AppError::Database)?;
|
||||
Ok(row)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(include_str!(
|
||||
"../../../migrations/20260622000001_agent_sessions.sql"
|
||||
))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ensure_is_idempotent_and_accumulates() {
|
||||
let db = pool().await;
|
||||
ensure_session(&db, "s1", Some(7), "gpt-4o", "default")
|
||||
.await
|
||||
.unwrap();
|
||||
// Second ensure must not reset counters.
|
||||
record_turn(
|
||||
&db,
|
||||
"s1",
|
||||
UsageDelta { input_tokens: 100, output_tokens: 40, cost_usd: 0.5, num_turns: 3 },
|
||||
2,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
ensure_session(&db, "s1", Some(7), "gpt-4o", "default")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let s = get_session(&db, "s1", Some(7)).await.unwrap().unwrap();
|
||||
assert_eq!(s.input_tokens, 100);
|
||||
assert_eq!(s.output_tokens, 40);
|
||||
assert_eq!(s.total_tokens, 140);
|
||||
assert!((s.cost_usd - 0.5).abs() < 1e-9);
|
||||
assert_eq!(s.num_turns, 3);
|
||||
assert_eq!(s.message_count, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn record_turn_adds_across_turns() {
|
||||
let db = pool().await;
|
||||
ensure_session(&db, "s2", None, "sonnet", "default").await.unwrap();
|
||||
for _ in 0..3 {
|
||||
record_turn(
|
||||
&db,
|
||||
"s2",
|
||||
UsageDelta { input_tokens: 10, output_tokens: 5, cost_usd: 0.01, num_turns: 1 },
|
||||
2,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let s = get_session(&db, "s2", None).await.unwrap().unwrap();
|
||||
assert_eq!(s.total_tokens, 45);
|
||||
assert_eq!(s.num_turns, 3);
|
||||
assert_eq!(s.message_count, 6);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_session_is_user_scoped() {
|
||||
let db = pool().await;
|
||||
ensure_session(&db, "s3", Some(1), "m", "default").await.unwrap();
|
||||
// Wrong user cannot see it.
|
||||
assert!(get_session(&db, "s3", Some(2)).await.unwrap().is_none());
|
||||
// Owner can.
|
||||
assert!(get_session(&db, "s3", Some(1)).await.unwrap().is_some());
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod announcement;
|
||||
pub mod agent_workbench;
|
||||
pub mod billing;
|
||||
pub mod builtin_keys;
|
||||
pub mod comfyui_process;
|
||||
|
||||
@@ -42,6 +42,7 @@ const Models = lazy(() => import('./pages/Models.js'));
|
||||
const About = lazy(() => import('./pages/About.js'));
|
||||
const ModelTester = lazy(() => import('./pages/ModelTester.js'));
|
||||
const ComfyUIAgent = lazy(() => import('./pages/ComfyUIAgent.js'));
|
||||
const AgentWorkbench = lazy(() => import('./pages/AgentWorkbench.js'));
|
||||
const PictureAgent = lazy(() => import('./pages/PictureAgent.js'));
|
||||
const DigitalHuman = lazy(() => import('./pages/DigitalHuman.js'));
|
||||
const MoneyPrinter = lazy(() => import('./pages/MoneyPrinter.js'));
|
||||
@@ -626,6 +627,7 @@ export const sidebarGroups: SidebarGroup[] = [
|
||||
{ to: '/digital-human', label: '短视频工厂/数字人', icon: navIcons.digitalHuman },
|
||||
{ to: '/workflows', label: '工作流', icon: navIcons.workflow },
|
||||
{ to: '/comfyui-agent', label: '视频助手', icon: navIcons.assistant },
|
||||
{ to: '/agent-workbench', label: 'AI Agent', icon: navIcons.assistant },
|
||||
{ to: '/workspace', label: '工作空间', icon: navIcons.workspace },
|
||||
],
|
||||
},
|
||||
@@ -658,6 +660,7 @@ export const userSidebarGroups = [
|
||||
{ to: '/user/digital-human', label: '短视频工厂/数字人', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M12 12a4 4 0 100-8 4 4 0 000 8z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 21a8 8 0 0116 0M17 9l3 2v4l-3 2M7 9l-3 2v4l3 2" /></svg> },
|
||||
{ to: '/user/workflows', label: '工作流', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 6h5v5H4zM15 4h5v5h-5zM15 15h5v5h-5zM9 8h6M17.5 9v6M9 11l6 6" /></svg> },
|
||||
{ to: '/user/comfyui-agent', label: '视频助手', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M14.752 11.168l-3.197-2.132A1 1 0 0010 9.87v4.263a1 1 0 001.555.832l3.197-2.132a1 1 0 000-1.664z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M21 12a9 9 0 11-18 0 9 9 0 0118 0z" /></svg> },
|
||||
{ to: '/user/agent-workbench', label: 'AI Agent', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M9.75 17L9 20l-1 1h8l-1-1-.75-3M3 13h18M5 17h14a2 2 0 002-2V5a2 2 0 00-2-2H5a2 2 0 00-2 2v10a2 2 0 002 2z" /></svg> },
|
||||
{ to: '/user/workspace', label: '工作空间', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 7a2 2 0 012-2h4l2 2h6a2 2 0 012 2v8a2 2 0 01-2 2H6a2 2 0 01-2-2V7z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M8 13h8M8 16h5" /></svg> },
|
||||
],
|
||||
},
|
||||
@@ -1667,6 +1670,7 @@ function AppShell() {
|
||||
<Route path="/workflows" element={<Workflows />} />
|
||||
<Route path="/video-agent" element={<Navigate to="/workflows" replace />} />
|
||||
<Route path="/comfyui-agent" element={<ComfyUIAgent />} />
|
||||
<Route path="/agent-workbench" element={<AgentWorkbench />} />
|
||||
<Route path="/workspace" element={<Workspace />} />
|
||||
<Route path="/about" element={<About />} />
|
||||
<Route path="/users" element={<UserManagement />} />
|
||||
@@ -1690,6 +1694,7 @@ function AppShell() {
|
||||
<Route path="/user/workflows" element={<Workflows />} />
|
||||
<Route path="/user/video-agent" element={<Navigate to="/user/workflows" replace />} />
|
||||
<Route path="/user/comfyui-agent" element={<ComfyUIAgent />} />
|
||||
<Route path="/user/agent-workbench" element={<AgentWorkbench />} />
|
||||
<Route path="/user/workspace" element={<Workspace />} />
|
||||
<Route path="/user/about" element={<About />} />
|
||||
|
||||
|
||||
@@ -0,0 +1,482 @@
|
||||
import { useCallback, useEffect, useRef, useState } from 'react';
|
||||
import { useLocation } from 'react-router-dom';
|
||||
import { getSessionToken } from '../authSession.js';
|
||||
import { useToast } from '../components/Toast.js';
|
||||
|
||||
// Agent Workbench — a minimal chat UI over POST /api/agent/chat (SSE).
|
||||
// Streams assistant text, renders tool-call cards, surfaces interactive
|
||||
// approval prompts for risky tools, and shows running session token/cost.
|
||||
|
||||
type PermissionMode = 'default' | 'accept_edits' | 'full_access';
|
||||
|
||||
type ToolCall = {
|
||||
id: string;
|
||||
name: string;
|
||||
input: unknown;
|
||||
result?: string;
|
||||
isError?: boolean;
|
||||
denied?: string;
|
||||
};
|
||||
|
||||
type ChatMessage = {
|
||||
role: 'user' | 'assistant';
|
||||
text: string;
|
||||
tools: ToolCall[];
|
||||
};
|
||||
|
||||
type ApprovalPrompt = {
|
||||
requestId: string;
|
||||
name: string;
|
||||
input: unknown;
|
||||
risk: string;
|
||||
};
|
||||
|
||||
type SessionUsage = {
|
||||
inputTokens: number;
|
||||
outputTokens: number;
|
||||
totalTokens: number;
|
||||
costUsd: number;
|
||||
numTurns: number;
|
||||
};
|
||||
|
||||
type SseEvent = { type: string; data?: Record<string, unknown> };
|
||||
|
||||
const MODE_LABELS: Record<PermissionMode, string> = {
|
||||
default: '标准审核',
|
||||
accept_edits: '自动接受编辑',
|
||||
full_access: '完全放行',
|
||||
};
|
||||
|
||||
/** Read a `data:`-framed SSE stream, invoking `onEvent` per JSON line. */
|
||||
async function readSse(response: Response, onEvent: (e: SseEvent) => void): Promise<void> {
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) throw new Error('浏览器无法读取流式响应');
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
for (;;) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop() || '';
|
||||
for (const raw of lines) {
|
||||
const line = raw.trim();
|
||||
if (!line.startsWith('data:')) continue;
|
||||
const data = line.slice(5).trim();
|
||||
if (!data) continue;
|
||||
try {
|
||||
onEvent(JSON.parse(data) as SseEvent);
|
||||
} catch {
|
||||
/* ignore malformed frame */
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function formatCost(usd: number): string {
|
||||
return `$${usd.toFixed(4)}`;
|
||||
}
|
||||
export default function AgentWorkbench() {
|
||||
const toast = useToast();
|
||||
const location = useLocation();
|
||||
const audience = location.pathname.startsWith('/user/') ? 'user' : 'admin';
|
||||
|
||||
const [messages, setMessages] = useState<ChatMessage[]>([]);
|
||||
const [input, setInput] = useState('');
|
||||
const [running, setRunning] = useState(false);
|
||||
const [mode, setMode] = useState<PermissionMode>('default');
|
||||
const [sessionId, setSessionId] = useState<string | null>(null);
|
||||
const [usage, setUsage] = useState<SessionUsage | null>(null);
|
||||
const [approval, setApproval] = useState<ApprovalPrompt | null>(null);
|
||||
|
||||
const scrollRef = useRef<HTMLDivElement | null>(null);
|
||||
useEffect(() => {
|
||||
scrollRef.current?.scrollTo({ top: scrollRef.current.scrollHeight, behavior: 'smooth' });
|
||||
}, [messages, approval]);
|
||||
|
||||
const authToken = useCallback(
|
||||
() => getSessionToken(audience, localStorage) ?? '',
|
||||
[audience],
|
||||
);
|
||||
|
||||
// Mutate the last assistant message in place as stream events arrive.
|
||||
const updateAssistant = useCallback((mut: (m: ChatMessage) => void) => {
|
||||
setMessages((prev) => {
|
||||
if (prev.length === 0) return prev;
|
||||
const next = prev.slice();
|
||||
const last = { ...next[next.length - 1] };
|
||||
if (last.role !== 'assistant') return prev;
|
||||
last.tools = last.tools.slice();
|
||||
mut(last);
|
||||
next[next.length - 1] = last;
|
||||
return next;
|
||||
});
|
||||
}, []);
|
||||
|
||||
const resolveApproval = useCallback(
|
||||
async (requestId: string, approved: boolean) => {
|
||||
setApproval(null);
|
||||
try {
|
||||
await fetch('/api/agent/approve', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Authorization: `Bearer ${authToken()}`,
|
||||
},
|
||||
body: JSON.stringify({ request_id: requestId, approved }),
|
||||
});
|
||||
} catch {
|
||||
toast.error('提交审批失败');
|
||||
}
|
||||
},
|
||||
[authToken, toast],
|
||||
);
|
||||
|
||||
const handleEvent = useCallback(
|
||||
(e: SseEvent) => {
|
||||
const d = e.data ?? {};
|
||||
switch (e.type) {
|
||||
case 'session':
|
||||
if (typeof d.session_id === 'string') setSessionId(d.session_id);
|
||||
break;
|
||||
case 'text':
|
||||
if (typeof d.text === 'string') {
|
||||
const chunk = d.text;
|
||||
updateAssistant((m) => {
|
||||
m.text += chunk;
|
||||
});
|
||||
}
|
||||
break;
|
||||
case 'tool_use':
|
||||
updateAssistant((m) => {
|
||||
m.tools.push({
|
||||
id: String(d.id ?? ''),
|
||||
name: String(d.name ?? ''),
|
||||
input: d.input,
|
||||
});
|
||||
});
|
||||
break;
|
||||
case 'tool_result':
|
||||
updateAssistant((m) => {
|
||||
const tc = m.tools.find((t) => t.id === d.tool_use_id);
|
||||
if (tc) {
|
||||
tc.result = String(d.content ?? '');
|
||||
tc.isError = Boolean(d.is_error);
|
||||
}
|
||||
});
|
||||
break;
|
||||
case 'tool_denied':
|
||||
updateAssistant((m) => {
|
||||
const tc = m.tools.find((t) => t.name === d.name && !t.result && !t.denied);
|
||||
if (tc) tc.denied = String(d.reason ?? '已拒绝');
|
||||
});
|
||||
break;
|
||||
case 'approval_request':
|
||||
setApproval({
|
||||
requestId: String(d.request_id ?? ''),
|
||||
name: String(d.name ?? ''),
|
||||
input: d.input,
|
||||
risk: String(d.risk ?? ''),
|
||||
});
|
||||
break;
|
||||
case 'result':
|
||||
setUsage({
|
||||
inputTokens: Number(d.input_tokens ?? 0),
|
||||
outputTokens: Number(d.output_tokens ?? 0),
|
||||
totalTokens: Number(d.total_tokens ?? 0),
|
||||
costUsd: Number(d.cost_usd ?? 0),
|
||||
numTurns: Number(d.num_turns ?? 0),
|
||||
});
|
||||
break;
|
||||
case 'error':
|
||||
toast.error(String(d.message ?? '执行出错'));
|
||||
updateAssistant((m) => {
|
||||
m.text += `\n\n[错误: ${String(d.message ?? '')}]`;
|
||||
});
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
},
|
||||
[toast, updateAssistant],
|
||||
);
|
||||
|
||||
const send = useCallback(async () => {
|
||||
const prompt = input.trim();
|
||||
if (!prompt || running) return;
|
||||
const token = authToken();
|
||||
if (!token) {
|
||||
toast.error('未登录或会话已过期');
|
||||
return;
|
||||
}
|
||||
|
||||
setRunning(true);
|
||||
setApproval(null);
|
||||
setMessages((prev) => [
|
||||
...prev,
|
||||
{ role: 'user', text: prompt, tools: [] },
|
||||
{ role: 'assistant', text: '', tools: [] },
|
||||
]);
|
||||
setInput('');
|
||||
|
||||
try {
|
||||
const response = await fetch('/api/agent/chat', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Authorization: `Bearer ${token}`,
|
||||
},
|
||||
body: JSON.stringify({
|
||||
message: prompt,
|
||||
session_id: sessionId,
|
||||
permission_mode: mode,
|
||||
}),
|
||||
});
|
||||
if (!response.ok) {
|
||||
const detail = await response.text().catch(() => '');
|
||||
throw new Error(detail || `HTTP ${response.status}`);
|
||||
}
|
||||
await readSse(response, handleEvent);
|
||||
} catch (err) {
|
||||
toast.error(`请求失败: ${err instanceof Error ? err.message : String(err)}`);
|
||||
updateAssistant((m) => {
|
||||
m.text += '\n\n[连接中断]';
|
||||
});
|
||||
} finally {
|
||||
setRunning(false);
|
||||
setApproval(null);
|
||||
}
|
||||
}, [authToken, handleEvent, input, mode, running, sessionId, toast, updateAssistant]);
|
||||
|
||||
const newSession = useCallback(() => {
|
||||
if (running) return;
|
||||
setMessages([]);
|
||||
setSessionId(null);
|
||||
setUsage(null);
|
||||
setApproval(null);
|
||||
}, [running]);
|
||||
return (
|
||||
<div style={{ display: 'flex', flexDirection: 'column', height: '100%', gap: 12, padding: 16 }}>
|
||||
{/* Header */}
|
||||
<div style={{ display: 'flex', alignItems: 'center', justifyContent: 'space-between', gap: 12, flexWrap: 'wrap' }}>
|
||||
<div>
|
||||
<h2 style={{ margin: 0, fontSize: 18, fontWeight: 600 }}>AI Agent 工作台</h2>
|
||||
<p style={{ margin: '4px 0 0', fontSize: 12, color: 'var(--text-secondary, #888)' }}>
|
||||
服务端内置 agent,可读写文件 / 执行命令(沙箱隔离 + 权限审批)
|
||||
</p>
|
||||
</div>
|
||||
<div style={{ display: 'flex', alignItems: 'center', gap: 8, flexWrap: 'wrap' }}>
|
||||
<label style={{ fontSize: 12, color: 'var(--text-secondary, #888)' }}>权限模式</label>
|
||||
<select
|
||||
value={mode}
|
||||
disabled={running}
|
||||
onChange={(e) => setMode(e.target.value as PermissionMode)}
|
||||
style={{ padding: '4px 8px', borderRadius: 6 }}
|
||||
>
|
||||
{(Object.keys(MODE_LABELS) as PermissionMode[]).map((m) => (
|
||||
<option key={m} value={m}>{MODE_LABELS[m]}</option>
|
||||
))}
|
||||
</select>
|
||||
<button onClick={newSession} disabled={running} style={{ padding: '4px 12px', borderRadius: 6 }}>
|
||||
新会话
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Usage bar */}
|
||||
{usage && (
|
||||
<div style={{ display: 'flex', gap: 16, fontSize: 12, color: 'var(--text-secondary, #888)', flexWrap: 'wrap' }}>
|
||||
<span>会话: <code>{sessionId?.slice(0, 8) ?? '—'}</code></span>
|
||||
<span>轮次: {usage.numTurns}</span>
|
||||
<span>tokens: {usage.totalTokens} (in {usage.inputTokens} / out {usage.outputTokens})</span>
|
||||
<span>成本: {formatCost(usage.costUsd)}</span>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Messages */}
|
||||
<div
|
||||
ref={scrollRef}
|
||||
style={{
|
||||
flex: 1,
|
||||
overflowY: 'auto',
|
||||
border: '1px solid var(--border-color, #2a2a2a)',
|
||||
borderRadius: 8,
|
||||
padding: 12,
|
||||
display: 'flex',
|
||||
flexDirection: 'column',
|
||||
gap: 12,
|
||||
minHeight: 240,
|
||||
}}
|
||||
>
|
||||
{messages.length === 0 && (
|
||||
<div style={{ color: 'var(--text-secondary, #888)', fontSize: 13, textAlign: 'center', marginTop: 40 }}>
|
||||
输入任务,让 agent 在隔离沙箱中读写文件、执行命令。危险操作会请求你的审批。
|
||||
</div>
|
||||
)}
|
||||
{messages.map((m, i) => (
|
||||
<div key={i} style={{ display: 'flex', flexDirection: 'column', gap: 6 }}>
|
||||
<div style={{ fontSize: 12, fontWeight: 600, color: m.role === 'user' ? 'var(--primary, #4a9eff)' : '#22c55e' }}>
|
||||
{m.role === 'user' ? '你' : 'Agent'}
|
||||
</div>
|
||||
{m.text && (
|
||||
<div style={{ whiteSpace: 'pre-wrap', fontSize: 14, lineHeight: 1.6 }}>{m.text}</div>
|
||||
)}
|
||||
{m.tools.map((t) => (
|
||||
<ToolCard key={t.id} tool={t} />
|
||||
))}
|
||||
{m.role === 'assistant' && !m.text && m.tools.length === 0 && running && i === messages.length - 1 && (
|
||||
<div style={{ fontSize: 13, color: 'var(--text-secondary, #888)' }}>思考中…</div>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
|
||||
{/* Approval modal */}
|
||||
{approval && (
|
||||
<ApprovalModal
|
||||
prompt={approval}
|
||||
onApprove={() => resolveApproval(approval.requestId, true)}
|
||||
onDeny={() => resolveApproval(approval.requestId, false)}
|
||||
/>
|
||||
)}
|
||||
|
||||
{/* Composer */}
|
||||
<div style={{ display: 'flex', gap: 8 }}>
|
||||
<textarea
|
||||
value={input}
|
||||
disabled={running}
|
||||
onChange={(e) => setInput(e.target.value)}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key === 'Enter' && (e.metaKey || e.ctrlKey)) {
|
||||
e.preventDefault();
|
||||
void send();
|
||||
}
|
||||
}}
|
||||
placeholder="描述任务… (Ctrl/Cmd + Enter 发送)"
|
||||
rows={3}
|
||||
style={{ flex: 1, padding: 10, borderRadius: 8, resize: 'vertical', fontFamily: 'inherit', fontSize: 14 }}
|
||||
/>
|
||||
<button
|
||||
onClick={() => void send()}
|
||||
disabled={running || !input.trim()}
|
||||
style={{
|
||||
padding: '0 20px',
|
||||
borderRadius: 8,
|
||||
background: running ? 'var(--border-color, #333)' : 'var(--primary, #4a9eff)',
|
||||
color: '#fff',
|
||||
border: 'none',
|
||||
cursor: running ? 'default' : 'pointer',
|
||||
}}
|
||||
>
|
||||
{running ? '运行中…' : '发送'}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function ToolCard({ tool }: { tool: ToolCall }) {
|
||||
const border = tool.denied
|
||||
? '#f59e0b'
|
||||
: tool.isError
|
||||
? '#ef4444'
|
||||
: tool.result
|
||||
? '#22c55e'
|
||||
: 'var(--border-color, #2a2a2a)';
|
||||
return (
|
||||
<div style={{ border: `1px solid ${border}`, borderRadius: 8, padding: 10, fontSize: 12, background: 'var(--surface, rgba(255,255,255,0.02))' }}>
|
||||
<div style={{ fontWeight: 600, marginBottom: 4, display: 'flex', alignItems: 'center', gap: 6 }}>
|
||||
<span>🔧 {tool.name}</span>
|
||||
{!tool.result && !tool.denied && <span style={{ color: 'var(--text-secondary, #888)' }}>· 运行中</span>}
|
||||
{tool.denied && <span style={{ color: '#f59e0b' }}>· 已拒绝</span>}
|
||||
</div>
|
||||
<pre style={{ margin: '4px 0', whiteSpace: 'pre-wrap', wordBreak: 'break-all', color: 'var(--text-secondary, #aaa)' }}>
|
||||
{typeof tool.input === 'string' ? tool.input : JSON.stringify(tool.input, null, 2)}
|
||||
</pre>
|
||||
{tool.denied && <div style={{ color: '#f59e0b' }}>{tool.denied}</div>}
|
||||
{tool.result && (
|
||||
<pre
|
||||
style={{
|
||||
margin: '6px 0 0',
|
||||
padding: 8,
|
||||
borderRadius: 6,
|
||||
background: 'rgba(0,0,0,0.25)',
|
||||
whiteSpace: 'pre-wrap',
|
||||
wordBreak: 'break-all',
|
||||
maxHeight: 200,
|
||||
overflowY: 'auto',
|
||||
color: tool.isError ? '#ef4444' : 'inherit',
|
||||
}}
|
||||
>
|
||||
{tool.result}
|
||||
</pre>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function ApprovalModal({
|
||||
prompt,
|
||||
onApprove,
|
||||
onDeny,
|
||||
}: {
|
||||
prompt: ApprovalPrompt;
|
||||
onApprove: () => void;
|
||||
onDeny: () => void;
|
||||
}) {
|
||||
return (
|
||||
<div
|
||||
style={{
|
||||
position: 'fixed',
|
||||
inset: 0,
|
||||
background: 'rgba(0,0,0,0.5)',
|
||||
display: 'flex',
|
||||
alignItems: 'center',
|
||||
justifyContent: 'center',
|
||||
zIndex: 1000,
|
||||
}}
|
||||
>
|
||||
<div
|
||||
style={{
|
||||
background: 'var(--bg, #1a1a1a)',
|
||||
border: '1px solid var(--border-color, #333)',
|
||||
borderRadius: 12,
|
||||
padding: 20,
|
||||
width: 'min(520px, 92vw)',
|
||||
maxHeight: '80vh',
|
||||
overflowY: 'auto',
|
||||
}}
|
||||
>
|
||||
<h3 style={{ margin: '0 0 8px', fontSize: 16 }}>需要审批:{prompt.name}</h3>
|
||||
<p style={{ margin: '0 0 12px', fontSize: 13, color: 'var(--text-secondary, #888)' }}>
|
||||
Agent 请求执行一个 <strong>{prompt.risk}</strong> 级别的工具。请确认参数后决定是否允许。
|
||||
</p>
|
||||
<pre
|
||||
style={{
|
||||
background: 'rgba(0,0,0,0.3)',
|
||||
borderRadius: 8,
|
||||
padding: 12,
|
||||
fontSize: 12,
|
||||
whiteSpace: 'pre-wrap',
|
||||
wordBreak: 'break-all',
|
||||
maxHeight: 280,
|
||||
overflowY: 'auto',
|
||||
}}
|
||||
>
|
||||
{typeof prompt.input === 'string' ? prompt.input : JSON.stringify(prompt.input, null, 2)}
|
||||
</pre>
|
||||
<div style={{ display: 'flex', justifyContent: 'flex-end', gap: 8, marginTop: 16 }}>
|
||||
<button onClick={onDeny} style={{ padding: '8px 16px', borderRadius: 8 }}>
|
||||
拒绝
|
||||
</button>
|
||||
<button
|
||||
onClick={onApprove}
|
||||
style={{ padding: '8px 16px', borderRadius: 8, background: '#22c55e', color: '#fff', border: 'none' }}
|
||||
>
|
||||
允许执行
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
//! Integration tests for the Agent Workbench approval registry.
|
||||
//!
|
||||
//! These run against the library's public API (not the in-crate `#[cfg(test)]`
|
||||
//! modules), so they compile and execute independently of unrelated test
|
||||
//! fixtures elsewhere in the crate. They pin the core safety invariant of the
|
||||
//! atomic approval state machine:
|
||||
//!
|
||||
//! the value `/api/agent/approve` reports back (`resolve` -> bool) and the
|
||||
//! decision the agent loop executes (`wait_for_decision` -> SettledState)
|
||||
//! are derived from the *same* atomic, so they can never disagree.
|
||||
|
||||
use boosapi::services::agent_workbench::approval::{
|
||||
register, resolve, wait_for_decision, ApprovalOutcome, SettledState,
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
#[tokio::test]
|
||||
async fn approve_then_wait_observes_approved() {
|
||||
let (entry, _guard) = register("it-approve");
|
||||
assert!(resolve("it-approve", ApprovalOutcome::Approve));
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_secs(5)).await,
|
||||
SettledState::Approved
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deny_then_wait_observes_denied() {
|
||||
let (entry, _guard) = register("it-deny");
|
||||
assert!(resolve("it-deny", ApprovalOutcome::Deny));
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_secs(5)).await,
|
||||
SettledState::Denied
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn timeout_denies_and_late_click_is_truthfully_rejected() {
|
||||
let (entry, _guard) = register("it-timeout");
|
||||
// No resolve arrives before the (tiny) timeout.
|
||||
assert_eq!(
|
||||
wait_for_decision(&entry, Duration::from_millis(20)).await,
|
||||
SettledState::TimedOut
|
||||
);
|
||||
// A late click must be told ok:false — the timeout already claimed it.
|
||||
assert!(!resolve("it-timeout", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_id_resolves_false() {
|
||||
assert!(!resolve("it-unknown", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn guard_drop_cleans_up_registry() {
|
||||
{
|
||||
let (_entry, _guard) = register("it-guard");
|
||||
}
|
||||
assert!(!resolve("it-guard", ApprovalOutcome::Approve));
|
||||
}
|
||||
|
||||
/// The invariant under contention: race a `resolve(Approve)` against a
|
||||
/// near-zero timeout thousands of times. The frontend-reported success
|
||||
/// (`resolve` -> true) must hold if and only if the executed decision is
|
||||
/// `Approved`; otherwise the decision must be `TimedOut`. The mixed state
|
||||
/// (frontend told "approved" while the loop denied, or vice versa) must never
|
||||
/// occur — that was the bug the atomic state machine fixes.
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn resolve_and_timeout_never_disagree() {
|
||||
for i in 0..5000 {
|
||||
let id = format!("it-race-{i}");
|
||||
let (entry, _guard) = register(&id);
|
||||
let id2 = id.clone();
|
||||
let resolver = tokio::spawn(async move { resolve(&id2, ApprovalOutcome::Approve) });
|
||||
let decided = wait_for_decision(&entry, Duration::from_nanos(1)).await;
|
||||
let frontend_ok = resolver.await.unwrap();
|
||||
match (frontend_ok, decided) {
|
||||
(true, SettledState::Approved) => {}
|
||||
(false, SettledState::TimedOut) => {}
|
||||
other => panic!("inconsistent outcome at i={i}: {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
//! Integration tests for OS-level (bubblewrap) confinement of the agent Bash
|
||||
//! tool. These actually invoke `SandboxedBashTool::call` and assert that, in a
|
||||
//! review mode, the kernel boundary blocks filesystem and network escapes while
|
||||
//! legitimate in-sandbox work succeeds.
|
||||
//!
|
||||
//! The escape tests require a working `bwrap`; when it is absent the confined
|
||||
//! path fails closed (asserted separately) and the live-escape tests skip with
|
||||
//! a printed notice rather than producing a false pass.
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
use boosapi::services::agent_workbench::bash_sandbox::{build_bwrap_args, SandboxedBashTool};
|
||||
use boosapi::services::agent_workbench::permission::PermissionMode;
|
||||
use open_agent_sdk::types::{Tool, ToolUseContext};
|
||||
use serde_json::json;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
fn bwrap_available() -> bool {
|
||||
["/usr/bin/bwrap", "/bin/bwrap", "/usr/local/bin/bwrap"]
|
||||
.iter()
|
||||
.any(|p| Path::new(p).is_file())
|
||||
}
|
||||
|
||||
fn tmp_session() -> PathBuf {
|
||||
let base =
|
||||
std::env::temp_dir().join(format!("agent_bash_it_{}", uuid::Uuid::new_v4().simple()));
|
||||
std::fs::create_dir_all(&base).unwrap();
|
||||
base.canonicalize().unwrap()
|
||||
}
|
||||
|
||||
async fn run(tool: &SandboxedBashTool, command: &str) -> (String, bool) {
|
||||
let ctx = ToolUseContext::new(std::env::temp_dir().to_string_lossy().to_string());
|
||||
let res = tool
|
||||
.call(json!({ "command": command }), &ctx)
|
||||
.await
|
||||
.expect("tool call ok");
|
||||
(res.get_text(), res.is_error)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_args_pin_session_dir_and_isolation() {
|
||||
let args = build_bwrap_args(Path::new("/srv/sess/abc"), "id");
|
||||
let j = args.join(" ");
|
||||
assert!(j.contains("--bind /srv/sess/abc /srv/sess/abc"));
|
||||
assert!(j.contains("--chdir /srv/sess/abc"));
|
||||
assert!(args.iter().any(|a| a == "--unshare-all"));
|
||||
assert!(!j.contains("/etc"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_mode_blocks_reading_etc_passwd() {
|
||||
if !bwrap_available() {
|
||||
eprintln!("SKIP: bwrap not available");
|
||||
return;
|
||||
}
|
||||
let tool = SandboxedBashTool::new(tmp_session(), PermissionMode::Default);
|
||||
let (out, _is_err) = run(&tool, "cat /etc/passwd").await;
|
||||
// The host's /etc is never mounted, so the read cannot succeed.
|
||||
assert!(
|
||||
!out.contains("root:") && !out.contains(":/bin/"),
|
||||
"should not see host passwd contents, got: {out}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_mode_blocks_listing_root() {
|
||||
if !bwrap_available() {
|
||||
eprintln!("SKIP: bwrap not available");
|
||||
return;
|
||||
}
|
||||
let tool = SandboxedBashTool::new(tmp_session(), PermissionMode::Default);
|
||||
let (out, _) = run(&tool, "ls /").await;
|
||||
// /etc, /root, /home must be absent from the sandbox root.
|
||||
assert!(!out.contains("etc"), "sandbox root must not expose /etc: {out}");
|
||||
assert!(!out.contains("root"), "sandbox root must not expose /root: {out}");
|
||||
assert!(!out.contains("home"), "sandbox root must not expose /home: {out}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_mode_blocks_writing_outside_sandbox() {
|
||||
if !bwrap_available() {
|
||||
eprintln!("SKIP: bwrap not available");
|
||||
return;
|
||||
}
|
||||
let tool = SandboxedBashTool::new(tmp_session(), PermissionMode::Default);
|
||||
let (_out, is_err) = run(&tool, "echo pwned > /etc/cron.d/x").await;
|
||||
assert!(is_err, "writing outside the sandbox must fail");
|
||||
// And the host file must not exist.
|
||||
assert!(!Path::new("/etc/cron.d/x").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_mode_blocks_network() {
|
||||
if !bwrap_available() {
|
||||
eprintln!("SKIP: bwrap not available");
|
||||
return;
|
||||
}
|
||||
let tool = SandboxedBashTool::new(tmp_session(), PermissionMode::Default);
|
||||
let (out, _) = run(&tool, "getent hosts example.com || echo NET-BLOCKED").await;
|
||||
assert!(out.contains("NET-BLOCKED"), "network should be unshared: {out}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_mode_allows_work_inside_sandbox() {
|
||||
if !bwrap_available() {
|
||||
eprintln!("SKIP: bwrap not available");
|
||||
return;
|
||||
}
|
||||
let sd = tmp_session();
|
||||
let tool = SandboxedBashTool::new(sd.clone(), PermissionMode::Default);
|
||||
let (out, is_err) = run(&tool, "echo hello > note.txt && cat note.txt").await;
|
||||
assert!(!is_err, "in-sandbox work should succeed: {out}");
|
||||
assert!(out.contains("hello"));
|
||||
// The file lands in the real session dir (bound rw).
|
||||
assert!(sd.join("note.txt").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_access_mode_is_unconfined() {
|
||||
// FullAccess is the explicitly-dangerous mode: it should NOT be sandboxed,
|
||||
// so it can read the host (this documents the intended escape hatch).
|
||||
let tool = SandboxedBashTool::new(tmp_session(), PermissionMode::FullAccess);
|
||||
let (out, _) = run(&tool, "test -e /etc && echo HOST-VISIBLE || echo no").await;
|
||||
assert!(out.contains("HOST-VISIBLE"), "full access should see host: {out}");
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
//! Integration test for the dangling-symlink write-escape fix.
|
||||
//!
|
||||
//! Runs against the library's public API so it compiles/executes independently
|
||||
//! of unrelated in-crate test fixtures. Reproduces the concrete escape an
|
||||
//! approved `ln -s` could otherwise open: a symlink whose target does not exist
|
||||
//! yet (so `exists()` is false and `canonicalize()` fails) must still be
|
||||
//! rejected when it points outside the per-session sandbox, instead of being
|
||||
//! treated as a brand-new in-sandbox path that a Write would follow out.
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
use boosapi::services::agent_workbench::sandbox::{check_tool_paths, validate_path_within_cwd};
|
||||
use serde_json::json;
|
||||
use std::os::unix::fs::symlink;
|
||||
use std::path::PathBuf;
|
||||
|
||||
fn tmp_sandbox() -> PathBuf {
|
||||
let base = std::env::temp_dir().join(format!("agent_it_sb_{}", uuid::Uuid::new_v4().simple()));
|
||||
std::fs::create_dir_all(&base).unwrap();
|
||||
base.canonicalize().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dangling_symlink_to_absolute_target_is_rejected() {
|
||||
let sb = tmp_sandbox();
|
||||
symlink("/etc/cron.d/pwn_it", sb.join("evil")).unwrap();
|
||||
let res = validate_path_within_cwd("evil", &sb);
|
||||
assert!(res.is_err(), "expected rejection, got {res:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dangling_relative_symlink_escape_is_rejected() {
|
||||
let sb = tmp_sandbox();
|
||||
symlink("../../escape_it", sb.join("evil")).unwrap();
|
||||
assert!(validate_path_within_cwd("evil", &sb).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_tool_through_dangling_symlink_is_blocked() {
|
||||
let sb = tmp_sandbox();
|
||||
symlink("/etc/cron.d/pwn_it2", sb.join("evil")).unwrap();
|
||||
let input = json!({ "path": "evil", "content": "* * * * * root sh -c id" });
|
||||
assert!(check_tool_paths("write", &input, &sb).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_through_intermediate_dangling_symlink_dir_is_blocked() {
|
||||
let sb = tmp_sandbox();
|
||||
symlink("/etc/nonexistent_it_dir", sb.join("evildir")).unwrap();
|
||||
let input = json!({ "path": "evildir/pwn", "content": "x" });
|
||||
assert!(check_tool_paths("write", &input, &sb).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dangling_symlink_pointing_inside_sandbox_is_allowed() {
|
||||
let sb = tmp_sandbox();
|
||||
symlink(sb.join("future_file"), sb.join("ok_link")).unwrap();
|
||||
assert!(validate_path_within_cwd("ok_link", &sb).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ordinary_new_file_still_allowed() {
|
||||
let sb = tmp_sandbox();
|
||||
assert!(validate_path_within_cwd("notes/today.md", &sb).is_ok());
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
//! TOCTOU analysis evidence for the agent sandbox.
|
||||
//!
|
||||
//! The file tools (Read/Write/Edit) open paths *inside the SDK process*, after
|
||||
//! BoosAPI's `can_use_tool` callback has validated the path. That leaves a
|
||||
//! check-vs-open window which, in the abstract, a concurrent writer could use to
|
||||
//! swap a validated name for a symlink that escapes the sandbox.
|
||||
//!
|
||||
//! Closing that window atomically would require replacing the SDK's in-process
|
||||
//! file tools with `openat2(RESOLVE_BENEATH)` implementations (tracked as
|
||||
//! follow-up). What this test pins down is whether the window is *reachable* at
|
||||
//! all under the current execution model: an exploit needs a second, concurrent
|
||||
//! mutator, and the only agent-controlled way to get concurrency is a
|
||||
//! backgrounded process. Bash now runs under bwrap with `--unshare-all`, which
|
||||
//! includes a PID namespace; the command is PID 1 of that namespace, so when it
|
||||
//! exits the kernel reaps every other process in the namespace — a backgrounded
|
||||
//! racer cannot outlive the single bash call that spawned it. Combined with
|
||||
//! sequential tool execution and per-session bind isolation, there is no
|
||||
//! surviving concurrent mutator to win the race.
|
||||
//!
|
||||
//! This test proves the reaping property empirically (it is the load-bearing
|
||||
//! fact); it does not claim the in-process open is atomic.
|
||||
|
||||
#![cfg(unix)]
|
||||
|
||||
use boosapi::services::agent_workbench::bash_sandbox::SandboxedBashTool;
|
||||
use boosapi::services::agent_workbench::permission::PermissionMode;
|
||||
use open_agent_sdk::types::{Tool, ToolUseContext};
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
|
||||
fn bwrap_available() -> bool {
|
||||
["/usr/bin/bwrap", "/bin/bwrap", "/usr/local/bin/bwrap"]
|
||||
.iter()
|
||||
.any(|p| Path::new(p).is_file())
|
||||
}
|
||||
|
||||
fn tmp_session() -> PathBuf {
|
||||
let base =
|
||||
std::env::temp_dir().join(format!("agent_toctou_{}", uuid::Uuid::new_v4().simple()));
|
||||
std::fs::create_dir_all(&base).unwrap();
|
||||
base.canonicalize().unwrap()
|
||||
}
|
||||
|
||||
/// A backgrounded process spawned by a sandboxed bash call must not survive the
|
||||
/// call: bwrap's PID namespace reaps it when the shell (PID 1) exits. This is
|
||||
/// what removes the concurrent mutator that a file-tool TOCTOU would need.
|
||||
#[tokio::test]
|
||||
async fn backgrounded_process_does_not_outlive_bash_call() {
|
||||
if !bwrap_available() {
|
||||
eprintln!("SKIP: bwrap not available");
|
||||
return;
|
||||
}
|
||||
let sd = tmp_session();
|
||||
let tool = SandboxedBashTool::new(sd.clone(), PermissionMode::Default);
|
||||
let ctx = ToolUseContext::new(std::env::temp_dir().to_string_lossy().to_string());
|
||||
|
||||
// Background a writer that would keep appending for ~2s if it survived,
|
||||
// then let the foreground shell exit immediately.
|
||||
let cmd = "( for i in $(seq 1 200); do echo x >> racer_marker; sleep 0.01; done ) & \
|
||||
echo backgrounded";
|
||||
let res = tool.call(json!({ "command": cmd }), &ctx).await.unwrap();
|
||||
assert!(!res.is_error, "spawn call should succeed: {}", res.get_text());
|
||||
|
||||
let marker = sd.join("racer_marker");
|
||||
let count_now = std::fs::read_to_string(&marker)
|
||||
.map(|s| s.lines().count())
|
||||
.unwrap_or(0);
|
||||
|
||||
// Wait well past the racer's would-be lifetime.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1200)).await;
|
||||
|
||||
let count_later = std::fs::read_to_string(&marker)
|
||||
.map(|s| s.lines().count())
|
||||
.unwrap_or(0);
|
||||
|
||||
assert_eq!(
|
||||
count_now, count_later,
|
||||
"backgrounded racer kept writing after the bash call returned \
|
||||
(count {count_now} -> {count_later}); the PID namespace did not reap it, \
|
||||
so a concurrent mutator could exist"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user