gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-09-02 17:49:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Addressing vareity of security issues

Browse Source
This commit is contained in:
Illia Polosukhin
2026-02-05 10:40:24 -08:00
parent 0ab9643843
commit 3e6dfb8409
11 changed files with 602 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/agent_loop.rs
+85 -13
View File
@@ -424,6 +424,29 @@ impl Agent {
}
}
// Safety validation for user input
let validation = self.safety().validate_input(content);
if !validation.is_valid {
let details = validation
.errors
.iter()
.map(|e| format!("{}: {}", e.field, e.message))
.collect::<Vec<_>>()
.join("; ");
return Ok(SubmissionResult::error(format!(
"Input rejected by safety validation: {}",
details
)));
}
let violations = self.safety().check_policy(content);
if violations
.iter()
.any(|rule| rule.action == crate::safety::PolicyAction::Block)
{
return Ok(SubmissionResult::error("Input rejected by safety policy."));
}
// Handle explicit commands (starting with /) directly
// Everything else goes through the normal agentic loop with tools
let temp_message = IncomingMessage {
@@ -767,6 +790,22 @@ impl Agent {
name: tool_name.to_string(),
})?;
// Validate tool parameters
let validation = self.safety().validator().validate_tool_params(params);
if !validation.is_valid {
let details = validation
.errors
.iter()
.map(|e| format!("{}: {}", e.field, e.message))
.collect::<Vec<_>>()
.join("; ");
return Err(crate::error::ToolError::InvalidParameters {
name: tool_name.to_string(),
reason: format!("Invalid tool parameters: {}", details),
}
.into());
}
// Execute with timeout
let result = tokio::time::timeout(std::time::Duration::from_secs(60), async {
tool.execute(params.clone(), job_ctx).await
@@ -813,11 +852,22 @@ impl Agent {
title,
description,
category,
} => self.handle_create_job(title, description, category).await?,
MessageIntent::CheckJobStatus { job_id } => self.handle_check_status(job_id).await?,
MessageIntent::CancelJob { job_id } => self.handle_cancel_job(&job_id).await?,
MessageIntent::ListJobs { filter } => self.handle_list_jobs(filter).await?,
MessageIntent::HelpJob { job_id } => self.handle_help_job(&job_id).await?,
} => {
self.handle_create_job(&message.user_id, title, description, category)
.await?
}
MessageIntent::CheckJobStatus { job_id } => {
self.handle_check_status(&message.user_id, job_id).await?
}
MessageIntent::CancelJob { job_id } => {
self.handle_cancel_job(&message.user_id, &job_id).await?
}
MessageIntent::ListJobs { filter } => {
self.handle_list_jobs(&message.user_id, filter).await?
}
MessageIntent::HelpJob { job_id } => {
self.handle_help_job(&message.user_id, &job_id).await?
}
MessageIntent::Command { command, args } => {
match self.handle_command(&command, &args).await? {
Some(s) => s,
@@ -1223,6 +1273,7 @@ impl Agent {
async fn handle_create_job(
&self,
user_id: &str,
title: String,
description: String,
category: Option<String>,
@@ -1230,7 +1281,7 @@ impl Agent {
// Create job context
let job_id = self
.context_manager
.create_job(&title, &description)
.create_job_for_user(user_id, &title, &description)
.await?;
// Update category if provided
@@ -1263,13 +1314,20 @@ impl Agent {
))
}
async fn handle_check_status(&self, job_id: Option<String>) -> Result<String, Error> {
async fn handle_check_status(
&self,
user_id: &str,
job_id: Option<String>,
) -> Result<String, Error> {
match job_id {
Some(id) => {
let uuid = Uuid::parse_str(&id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
Ok(format!(
"Job: {}\nStatus: {:?}\nCreated: {}\nStarted: {}\nActual cost: {}",
@@ -1284,7 +1342,7 @@ impl Agent {
}
None => {
// Show summary of all jobs
let summary = self.context_manager.summary().await;
let summary = self.context_manager.summary_for(user_id).await;
Ok(format!(
"Jobs summary:\n Total: {}\n In Progress: {}\n Completed: {}\n Failed: {}\n Stuck: {}",
summary.total,
@@ -1297,17 +1355,26 @@ impl Agent {
}
}
async fn handle_cancel_job(&self, job_id: &str) -> Result<String, Error> {
async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
let uuid = Uuid::parse_str(job_id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
self.scheduler.stop(uuid).await?;
Ok(format!("Job {} has been cancelled.", job_id))
}
async fn handle_list_jobs(&self, _filter: Option<String>) -> Result<String, Error> {
let jobs = self.context_manager.all_jobs().await;
async fn handle_list_jobs(
&self,
user_id: &str,
_filter: Option<String>,
) -> Result<String, Error> {
let jobs = self.context_manager.all_jobs_for(user_id).await;
if jobs.is_empty() {
return Ok("No jobs found.".to_string());
@@ -1316,18 +1383,23 @@ impl Agent {
let mut output = String::from("Jobs:\n");
for job_id in jobs {
if let Ok(ctx) = self.context_manager.get_context(job_id).await {
output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state));
if ctx.user_id == user_id {
output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state));
}
}
}
Ok(output)
}
async fn handle_help_job(&self, job_id: &str) -> Result<String, Error> {
async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
let uuid = Uuid::parse_str(job_id)
.map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?;
let ctx = self.context_manager.get_context(uuid).await?;
if ctx.user_id != user_id {
return Err(crate::error::JobError::NotFound { id: uuid }.into());
}
if ctx.state == crate::context::JobState::Stuck {
// Attempt recovery
src/agent/scheduler.rs
+79 -4
View File
@@ -50,9 +50,9 @@ pub struct Scheduler {
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
/// Running jobs (main LLM-driven jobs).
jobs: RwLock<HashMap<Uuid, ScheduledJob>>,
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
/// Running sub-tasks (tool executions, background tasks).
subtasks: RwLock<HashMap<Uuid, ScheduledSubtask>>,
subtasks: Arc<RwLock<HashMap<Uuid, ScheduledSubtask>>>,
}
impl Scheduler {
@@ -72,8 +72,8 @@ impl Scheduler {
safety,
tools,
store,
jobs: RwLock::new(HashMap::new()),
subtasks: RwLock::new(HashMap::new()),
jobs: Arc::new(RwLock::new(HashMap::new())),
subtasks: Arc::new(RwLock::new(HashMap::new())),
}
}
@@ -137,6 +137,27 @@ impl Scheduler {
.await
.insert(job_id, ScheduledJob { handle, tx });
// Cleanup task for this job to avoid capacity leaks
let jobs = Arc::clone(&self.jobs);
tokio::spawn(async move {
loop {
let finished = {
let jobs_read = jobs.read().await;
match jobs_read.get(&job_id) {
Some(scheduled) => scheduled.handle.is_finished(),
None => true,
}
};
if finished {
jobs.write().await.remove(&job_id);
break;
}
tokio::time::sleep(Duration::from_secs(1)).await;
}
});
tracing::info!("Scheduled job {} for execution", job_id);
Ok(())
}
@@ -171,11 +192,13 @@ impl Scheduler {
} => {
let tools = self.tools.clone();
let context_manager = self.context_manager.clone();
let safety = self.safety.clone();
tokio::spawn(async move {
let result = Self::execute_tool_task(
tools,
context_manager,
safety,
tool_parent_id,
&tool_name,
params,
@@ -217,6 +240,27 @@ impl Scheduler {
},
);
// Cleanup task for subtask tracking
let subtasks = Arc::clone(&self.subtasks);
tokio::spawn(async move {
loop {
let finished = {
let subtasks_read = subtasks.read().await;
match subtasks_read.get(&task_id) {
Some(scheduled) => scheduled.handle.is_finished(),
None => true,
}
};
if finished {
subtasks.write().await.remove(&task_id);
break;
}
tokio::time::sleep(Duration::from_secs(1)).await;
}
});
tracing::debug!(
parent_id = %parent_id,
task_id = %task_id,
@@ -282,6 +326,7 @@ impl Scheduler {
async fn execute_tool_task(
tools: Arc<ToolRegistry>,
context_manager: Arc<ContextManager>,
safety: Arc<SafetyLayer>,
job_id: Uuid,
tool_name: &str,
params: serde_json::Value,
@@ -297,6 +342,36 @@ impl Scheduler {
// Get job context
let job_ctx: JobContext = context_manager.get_context(job_id).await?;
if job_ctx.state == JobState::Cancelled {
return Err(crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: "Job is cancelled".to_string(),
}
.into());
}
if tool.requires_approval() {
return Err(crate::error::ToolError::AuthRequired {
name: tool_name.to_string(),
}
.into());
}
// Validate tool parameters
let validation = safety.validator().validate_tool_params(&params);
if !validation.is_valid {
let details = validation
.errors
.iter()
.map(|e| format!("{}: {}", e.field, e.message))
.collect::<Vec<_>>()
.join("; ");
return Err(crate::error::ToolError::InvalidParameters {
name: tool_name.to_string(),
reason: format!("Invalid tool parameters: {}", details),
}
.into());
}
// Execute with timeout
let result = tokio::time::timeout(Duration::from_secs(60), async {
src/agent/worker.rs
+43 -7
View File
@@ -226,6 +226,14 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
}
}
// Check for cancellation
if let Ok(ctx) = self.context_manager().get_context(self.job_id).await {
if ctx.state == JobState::Cancelled {
tracing::info!("Worker for job {} detected cancellation", self.job_id);
return Ok(());
}
}
iteration += 1;
if iteration > max_iterations {
self.mark_stuck("Maximum iterations exceeded").await?;
@@ -335,6 +343,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
let params = selection.parameters.clone();
let tools = self.tools().clone();
let context_manager = self.context_manager().clone();
let safety = self.safety().clone();
let job_id = self.job_id;
let store = self.deps.store.clone();
@@ -342,6 +351,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
let result = Self::execute_tool_inner(
tools,
context_manager,
safety,
store,
job_id,
&tool_name,
@@ -360,6 +370,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
async fn execute_tool_inner(
tools: Arc<ToolRegistry>,
context_manager: Arc<ContextManager>,
safety: Arc<SafetyLayer>,
store: Option<Arc<Store>>,
job_id: Uuid,
tool_name: &str,
@@ -372,17 +383,39 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
name: tool_name.to_string(),
})?;
// Log warning if tool requires approval (autonomous jobs auto-approve for now)
// Tools requiring approval are blocked in autonomous jobs
if tool.requires_approval() {
tracing::warn!(
job_id = %job_id,
tool = %tool_name,
"Executing sensitive tool in autonomous job (auto-approved)"
);
return Err(crate::error::ToolError::AuthRequired {
name: tool_name.to_string(),
}
.into());
}
// Get job context for the tool
let job_ctx = context_manager.get_context(job_id).await?;
if job_ctx.state == JobState::Cancelled {
return Err(crate::error::ToolError::ExecutionFailed {
name: tool_name.to_string(),
reason: "Job is cancelled".to_string(),
}
.into());
}
// Validate tool parameters
let validation = safety.validator().validate_tool_params(params);
if !validation.is_valid {
let details = validation
.errors
.iter()
.map(|e| format!("{}: {}", e.field, e.message))
.collect::<Vec<_>>()
.join("; ");
return Err(crate::error::ToolError::InvalidParameters {
name: tool_name.to_string(),
reason: format!("Invalid tool parameters: {}", details),
}
.into());
}
// Execute with timeout and timing
let start = std::time::Instant::now();
@@ -395,7 +428,9 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
// Record action in memory and get the ActionRecord for persistence
let action = match &result {
Ok(Ok(output)) => {
let output_str = serde_json::to_string_pretty(&output.result).ok();
let output_str = serde_json::to_string_pretty(&output.result)
.ok()
.map(|s| safety.sanitize_tool_output(tool_name, &s).content);
context_manager
.update_memory(job_id, |mem| {
let rec = mem.create_action(tool_name, params.clone()).succeed(
@@ -625,6 +660,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
Self::execute_tool_inner(
self.tools().clone(),
self.context_manager().clone(),
self.safety().clone(),
self.deps.store.clone(),
self.job_id,
tool_name,