Compare commits
9 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| c3327b9ef6 | |||
| d751055452 | |||
| b0ce83a937 | |||
| 3bb1272c6d | |||
| 5501558354 | |||
| f0a76fd82c | |||
| 2506713fad | |||
| a060666f6f | |||
| 4192798ceb |
Generated
+36
@@ -390,6 +390,7 @@ dependencies = [
|
|||||||
"anyhow",
|
"anyhow",
|
||||||
"argon2",
|
"argon2",
|
||||||
"async-stream",
|
"async-stream",
|
||||||
|
"async-trait",
|
||||||
"axum 0.8.9",
|
"axum 0.8.9",
|
||||||
"axum-test",
|
"axum-test",
|
||||||
"base64",
|
"base64",
|
||||||
@@ -407,6 +408,7 @@ dependencies = [
|
|||||||
"mime",
|
"mime",
|
||||||
"mime_guess",
|
"mime_guess",
|
||||||
"once_cell",
|
"once_cell",
|
||||||
|
"open-agent-sdk",
|
||||||
"parking_lot",
|
"parking_lot",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest",
|
"reqwest",
|
||||||
@@ -1825,6 +1827,31 @@ version = "1.70.2"
|
|||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
|
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]]
|
[[package]]
|
||||||
name = "openssl"
|
name = "openssl"
|
||||||
version = "0.10.80"
|
version = "0.10.80"
|
||||||
@@ -3811,6 +3838,15 @@ dependencies = [
|
|||||||
"windows-targets 0.52.6",
|
"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]]
|
[[package]]
|
||||||
name = "windows-sys"
|
name = "windows-sys"
|
||||||
version = "0.61.2"
|
version = "0.61.2"
|
||||||
|
|||||||
@@ -78,6 +78,12 @@ hex = "0.4"
|
|||||||
include_dir = "0.7"
|
include_dir = "0.7"
|
||||||
mime_guess = "2"
|
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]
|
[dev-dependencies]
|
||||||
rstest = "0.24"
|
rstest = "0.24"
|
||||||
wiremock = "0.6"
|
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 admin;
|
||||||
|
mod agent;
|
||||||
mod auth;
|
mod auth;
|
||||||
mod checkin;
|
mod checkin;
|
||||||
mod comfy_cloud_adapter;
|
mod comfy_cloud_adapter;
|
||||||
@@ -57,6 +58,7 @@ pub fn build_router(state: AppState) -> Router {
|
|||||||
.merge(comfy_cloud_adapter::routes())
|
.merge(comfy_cloud_adapter::routes())
|
||||||
.merge(money_printer::routes())
|
.merge(money_printer::routes())
|
||||||
.merge(comfyui_agent::routes())
|
.merge(comfyui_agent::routes())
|
||||||
|
.merge(agent::routes())
|
||||||
.merge(digital_human::routes())
|
.merge(digital_human::routes())
|
||||||
.merge(comfyui::routes())
|
.merge(comfyui::routes())
|
||||||
.merge(crate::workflow::routes::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 announcement;
|
||||||
|
pub mod agent_workbench;
|
||||||
pub mod billing;
|
pub mod billing;
|
||||||
pub mod builtin_keys;
|
pub mod builtin_keys;
|
||||||
pub mod comfyui_process;
|
pub mod comfyui_process;
|
||||||
|
|||||||
@@ -42,6 +42,7 @@ const Models = lazy(() => import('./pages/Models.js'));
|
|||||||
const About = lazy(() => import('./pages/About.js'));
|
const About = lazy(() => import('./pages/About.js'));
|
||||||
const ModelTester = lazy(() => import('./pages/ModelTester.js'));
|
const ModelTester = lazy(() => import('./pages/ModelTester.js'));
|
||||||
const ComfyUIAgent = lazy(() => import('./pages/ComfyUIAgent.js'));
|
const ComfyUIAgent = lazy(() => import('./pages/ComfyUIAgent.js'));
|
||||||
|
const AgentWorkbench = lazy(() => import('./pages/AgentWorkbench.js'));
|
||||||
const PictureAgent = lazy(() => import('./pages/PictureAgent.js'));
|
const PictureAgent = lazy(() => import('./pages/PictureAgent.js'));
|
||||||
const DigitalHuman = lazy(() => import('./pages/DigitalHuman.js'));
|
const DigitalHuman = lazy(() => import('./pages/DigitalHuman.js'));
|
||||||
const MoneyPrinter = lazy(() => import('./pages/MoneyPrinter.js'));
|
const MoneyPrinter = lazy(() => import('./pages/MoneyPrinter.js'));
|
||||||
@@ -626,6 +627,7 @@ export const sidebarGroups: SidebarGroup[] = [
|
|||||||
{ to: '/digital-human', label: '短视频工厂/数字人', icon: navIcons.digitalHuman },
|
{ to: '/digital-human', label: '短视频工厂/数字人', icon: navIcons.digitalHuman },
|
||||||
{ to: '/workflows', label: '工作流', icon: navIcons.workflow },
|
{ to: '/workflows', label: '工作流', icon: navIcons.workflow },
|
||||||
{ to: '/comfyui-agent', label: '视频助手', icon: navIcons.assistant },
|
{ to: '/comfyui-agent', label: '视频助手', icon: navIcons.assistant },
|
||||||
|
{ to: '/agent-workbench', label: 'AI Agent', icon: navIcons.assistant },
|
||||||
{ to: '/workspace', label: '工作空间', icon: navIcons.workspace },
|
{ 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/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/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/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> },
|
{ 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="/workflows" element={<Workflows />} />
|
||||||
<Route path="/video-agent" element={<Navigate to="/workflows" replace />} />
|
<Route path="/video-agent" element={<Navigate to="/workflows" replace />} />
|
||||||
<Route path="/comfyui-agent" element={<ComfyUIAgent />} />
|
<Route path="/comfyui-agent" element={<ComfyUIAgent />} />
|
||||||
|
<Route path="/agent-workbench" element={<AgentWorkbench />} />
|
||||||
<Route path="/workspace" element={<Workspace />} />
|
<Route path="/workspace" element={<Workspace />} />
|
||||||
<Route path="/about" element={<About />} />
|
<Route path="/about" element={<About />} />
|
||||||
<Route path="/users" element={<UserManagement />} />
|
<Route path="/users" element={<UserManagement />} />
|
||||||
@@ -1690,6 +1694,7 @@ function AppShell() {
|
|||||||
<Route path="/user/workflows" element={<Workflows />} />
|
<Route path="/user/workflows" element={<Workflows />} />
|
||||||
<Route path="/user/video-agent" element={<Navigate to="/user/workflows" replace />} />
|
<Route path="/user/video-agent" element={<Navigate to="/user/workflows" replace />} />
|
||||||
<Route path="/user/comfyui-agent" element={<ComfyUIAgent />} />
|
<Route path="/user/comfyui-agent" element={<ComfyUIAgent />} />
|
||||||
|
<Route path="/user/agent-workbench" element={<AgentWorkbench />} />
|
||||||
<Route path="/user/workspace" element={<Workspace />} />
|
<Route path="/user/workspace" element={<Workspace />} />
|
||||||
<Route path="/user/about" element={<About />} />
|
<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