Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
220 changes: 220 additions & 0 deletions src/agent/dispatcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,7 @@ impl Agent {
user_tz,
turn_usage: std::sync::Mutex::new(TurnUsageSummary::default()),
cached_tool_permissions: std::sync::Mutex::new(None),
cached_admin_tool_policy: tokio::sync::OnceCell::new(),
};

// If /skill-name mentions were expanded, rewrite the last user message
Expand Down Expand Up @@ -350,6 +351,7 @@ struct ChatDelegate<'a> {
turn_usage: std::sync::Mutex<TurnUsageSummary>,
cached_tool_permissions:
std::sync::Mutex<Option<std::collections::HashMap<String, PermissionState>>>,
cached_admin_tool_policy: crate::tools::permissions::AdminToolPolicyCache,
}

impl ChatDelegate<'_> {
Expand Down Expand Up @@ -422,6 +424,22 @@ impl<'a> LoopDelegate for ChatDelegate<'a> {
tool_defs
};

// Apply admin tool policy first so admin-disabled tools are removed
// before per-user permission filtering and session auto-approval.
let is_admin = self.tenant.identity().role.is_admin();
let admin_policy = crate::tools::permissions::load_cached_admin_tool_policy(
self.agent.store(),
&self.cached_admin_tool_policy,
)
.await;
let tool_defs = crate::tools::permissions::filter_admin_disabled_tools(
tool_defs,
self.agent.config.multi_tenant,
is_admin,
self.tenant.user_id(),
admin_policy,
);

// Apply per-user tool permission filtering.
//
// Load tool_permissions from the per-user DB settings store (same
Expand Down Expand Up @@ -2497,6 +2515,56 @@ mod tests {
}
}

#[derive(Default)]
struct RecordingToolsProvider {
seen_tools: std::sync::Mutex<Vec<Vec<String>>>,
}

#[async_trait]
impl LlmProvider for RecordingToolsProvider {
fn model_name(&self) -> &str {
"recording-tools"
}

fn cost_per_token(&self) -> (Decimal, Decimal) {
(Decimal::ZERO, Decimal::ZERO)
}

async fn complete(
&self,
_request: CompletionRequest,
) -> Result<CompletionResponse, crate::error::LlmError> {
Ok(CompletionResponse {
content: "ok".to_string(),
input_tokens: 0,
output_tokens: 1,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}

async fn complete_with_tools(
&self,
request: ToolCompletionRequest,
) -> Result<ToolCompletionResponse, crate::error::LlmError> {
let names: Vec<String> = request.tools.iter().map(|t| t.name.clone()).collect();
self.seen_tools
.lock()
.expect("recording tools mutex poisoned")
.push(names);
Ok(ToolCompletionResponse {
content: Some("ok".to_string()),
tool_calls: Vec::new(),
input_tokens: 0,
output_tokens: 1,
finish_reason: FinishReason::Stop,
cache_read_input_tokens: 0,
cache_creation_input_tokens: 0,
})
}
}

/// Helper to build a test Agent with a custom LLM provider and
/// `max_tool_iterations` override.
fn make_test_agent_with_llm(llm: Arc<dyn LlmProvider>, max_tool_iterations: usize) -> Agent {
Expand Down Expand Up @@ -2611,6 +2679,158 @@ mod tests {
);
}

#[cfg(feature = "libsql")]
#[tokio::test]
async fn test_admin_policy_filter_happens_before_auto_approval_and_llm_call() {
use crate::agent::session::Session;
use crate::channels::IncomingMessage;
use crate::llm::ChatMessage;
use crate::tools::builtin::{EchoTool, TimeTool};
use crate::tools::permissions::{ADMIN_SETTINGS_USER_ID, ADMIN_TOOL_POLICY_KEY};
use tokio::sync::Mutex;

let (db, _tmp_dir) = crate::testing::test_db().await;
db.set_setting(
"member-user",
"tool_permissions",
&serde_json::json!({
"echo": "always_allow",
"time": "always_allow"
}),
)
.await
.expect("failed to seed member tool permissions");
db.set_setting(
ADMIN_SETTINGS_USER_ID,
ADMIN_TOOL_POLICY_KEY,
&serde_json::json!({
"disabled_tools": ["echo"]
}),
)
.await
.expect("failed to seed admin tool policy");

let llm = Arc::new(RecordingToolsProvider::default());
let llm_for_assert = Arc::clone(&llm);
let tools = Arc::new(ToolRegistry::new());
tools.register_sync(Arc::new(EchoTool));
tools.register_sync(Arc::new(TimeTool));

let deps = AgentDeps {
owner_id: "default".to_string(),
store: Some(db),
llm: llm as Arc<dyn LlmProvider>,
cheap_llm: None,
safety: Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
})),
tools,
workspace: None,
extension_manager: None,
skill_registry: None,
skill_catalog: None,
skills_config: SkillsConfig::default(),
hooks: Arc::new(HookRegistry::new()),
cost_guard: Arc::new(CostGuard::new(CostGuardConfig::default())),
sse_tx: None,
http_interceptor: None,
transcription: None,
document_extraction: None,
auth_manager: None,
sandbox_readiness: crate::agent::routine_engine::SandboxReadiness::DisabledByConfig,
builder: None,
llm_backend: "nearai".to_string(),
tenant_rates: Arc::new(crate::tenant::TenantRateRegistry::new(4, 3)),
};

let agent = Agent::new(
AgentConfig {
name: "test-agent".to_string(),
max_parallel_jobs: 1,
job_timeout: Duration::from_secs(60),
stuck_threshold: Duration::from_secs(60),
repair_check_interval: Duration::from_secs(30),
max_repair_attempts: 1,
use_planning: false,
session_idle_timeout: Duration::from_secs(300),
allow_local_tools: false,
max_cost_per_day_cents: None,
max_actions_per_hour: None,
max_cost_per_user_per_day_cents: None,
max_tool_iterations: 5,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_jobs_per_user: None,
max_tokens_per_job: 0,
multi_tenant: true,
max_llm_concurrent_per_user: None,
max_jobs_concurrent_per_user: None,
engine_v2: false,
},
deps,
Arc::new(ChannelManager::new()),
None,
None,
None,
Some(Arc::new(ContextManager::new(1))),
None,
);

let session = Arc::new(Mutex::new(Session::new("member-user")));
let thread_id = {
let mut sess = session.lock().await;
sess.create_thread(Some("admin-policy")).id
};
let tenant = agent.tenant_ctx("member-user").await;
let message = IncomingMessage::new("test", "member-user", "hello");
let initial_messages = vec![ChatMessage::user("hello")];

let result = agent
.run_agentic_loop(
&message,
tenant,
Arc::clone(&session),
thread_id,
initial_messages,
)
.await;
assert!(result.is_ok(), "dispatcher run failed");

// admin-disabled tools must not remain auto-approved in session
let sess = session.lock().await;
assert!(
!sess.is_tool_auto_approved("echo"),
"echo is admin-disabled and must not be auto-approved"
);
assert!(
sess.is_tool_auto_approved("time"),
"time should remain auto-approved"
);
drop(sess);

// LLM should never see admin-disabled tools in available_tools.
let calls = llm_for_assert
.seen_tools
.lock()
.expect("recording tools mutex poisoned")
.clone();
assert!(
!calls.is_empty(),
"LLM should have been called at least once"
);
assert!(
!calls[0].iter().any(|name| name == "echo"),
"admin-disabled tool leaked into LLM tool list: {:?}",
calls[0]
);
assert!(
calls[0].iter().any(|name| name == "time"),
"expected non-disabled tool to remain available: {:?}",
calls[0]
);
}

/// Verify that the max_iterations guard terminates the loop even when the
/// LLM always returns tool calls and those calls succeed.
#[tokio::test]
Expand Down
1 change: 1 addition & 0 deletions src/agent/scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,7 @@ impl Scheduler {
sse_tx: self.sse_tx.clone(),
approval_context,
http_interceptor: self.http_interceptor.clone(),
multi_tenant: self.config.multi_tenant,
};
let worker = Worker::new(job_id, deps);

Expand Down
4 changes: 2 additions & 2 deletions src/channels/web/handlers/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ pub async fn engine_mission_pause_handler(
AuthenticatedUser(user): AuthenticatedUser,
Path(id): Path<String>,
) -> Result<Json<EngineActionResponse>, (StatusCode, String)> {
let is_admin = user.role == "admin";
let is_admin = crate::ownership::UserRole::from_db_role(&user.role).is_admin();
crate::bridge::pause_engine_mission(&id, &user.user_id, is_admin)
.await
.map_err(|e| {
Expand All @@ -214,7 +214,7 @@ pub async fn engine_mission_resume_handler(
AuthenticatedUser(user): AuthenticatedUser,
Path(id): Path<String>,
) -> Result<Json<EngineActionResponse>, (StatusCode, String)> {
let is_admin = user.role == "admin";
let is_admin = crate::ownership::UserRole::from_db_role(&user.role).is_admin();
crate::bridge::resume_engine_mission(&id, &user.user_id, is_admin)
.await
.map_err(|e| {
Expand Down
1 change: 1 addition & 0 deletions src/channels/web/handlers/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ pub mod secrets;
pub mod skills;
pub mod system_prompt;
pub mod tokens;
pub mod tool_policy;
pub mod users;

// Modules not yet wired into server.rs router -- suppress dead_code until
Expand Down
106 changes: 106 additions & 0 deletions src/channels/web/handlers/tool_policy.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
//! Admin tool policy handlers.
//!
//! Allows an admin to define which tools are disabled for all non-admin users
//! or for specific users. The policy is stored in the settings table under the
//! well-known `__admin__` scope.
//!
//! dispatch-exempt: These endpoints access `state.store` directly (not through
//! the agentic tool pipeline) because they are admin-only infrastructure
//! operations gated behind `AdminUser` auth, consistent with the other admin
//! handlers in this module (users, secrets, tokens).

use std::sync::Arc;

use axum::{Json, extract::State, http::StatusCode};

use crate::channels::web::auth::AdminUser;
use crate::channels::web::server::GatewayState;
use crate::tools::permissions::{
ADMIN_SETTINGS_USER_ID, ADMIN_TOOL_POLICY_KEY, AdminToolPolicy, parse_admin_tool_policy,
validate_admin_tool_policy,
};

/// GET /api/admin/tool-policy — retrieve the current admin tool policy.
///
/// Only available in multi-tenant mode (returns 404 in single-user deployments).
pub async fn tool_policy_get_handler(
State(state): State<Arc<GatewayState>>,
AdminUser(_admin): AdminUser,
) -> Result<Json<AdminToolPolicy>, (StatusCode, String)> {
let pool = state.workspace_pool.as_ref(); // dispatch-exempt: gateway-mode probe, not a state mutation
if pool.is_none() {
return Err((
StatusCode::NOT_FOUND,
"Admin tool policy is only available in multi-tenant mode".to_string(),
));
}

let store = state.store.as_ref(); // dispatch-exempt: admin-only read of cross-tenant policy scope
let store = store.ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"Database not available".to_string(),
))?;

let policy = match store
.get_setting(ADMIN_SETTINGS_USER_ID, ADMIN_TOOL_POLICY_KEY)
.await
{
Ok(Some(value)) => parse_admin_tool_policy(value, "http_get").map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Stored admin tool policy is corrupt: {e}"),
)
})?,
Ok(None) => AdminToolPolicy::default(),
Err(e) => {
return Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string()));
}
};

Ok(Json(policy))
}

/// PUT /api/admin/tool-policy — replace the admin tool policy.
///
/// Body must be a JSON `AdminToolPolicy`. Tool names and user IDs are
/// validated for basic sanity (non-empty, reasonable length).
///
/// This endpoint is a full replacement with last-write-wins semantics.
/// Each PUT overwrites the previously stored policy; there is no merge/patch.
///
/// Only available in multi-tenant mode (returns 404 in single-user deployments).
pub async fn tool_policy_put_handler(
State(state): State<Arc<GatewayState>>,
AdminUser(_admin): AdminUser,
Json(policy): Json<AdminToolPolicy>,
) -> Result<Json<AdminToolPolicy>, (StatusCode, String)> {
let pool = state.workspace_pool.as_ref(); // dispatch-exempt: gateway-mode probe, not a state mutation
if pool.is_none() {
return Err((
StatusCode::NOT_FOUND,
"Admin tool policy is only available in multi-tenant mode".to_string(),
));
}

validate_admin_tool_policy(&policy).map_err(|error| (StatusCode::BAD_REQUEST, error))?;

let store = state.store.as_ref(); // dispatch-exempt: admin-only write to cross-tenant policy scope
let store = store.ok_or((
StatusCode::SERVICE_UNAVAILABLE,
"Database not available".to_string(),
))?;

let value = serde_json::to_value(&policy).map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to serialize policy: {e}"),
)
})?;

store
.set_setting(ADMIN_SETTINGS_USER_ID, ADMIN_TOOL_POLICY_KEY, &value)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

Ok(Json(policy))
}
Loading
Loading