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
2 changes: 2 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,8 @@ AGENT_NAME=ironclaw
AGENT_MAX_PARALLEL_JOBS=5
AGENT_JOB_TIMEOUT_SECS=3600
AGENT_STUCK_THRESHOLD_SECS=300
# Maximum tokens per job (0 = unlimited, also settable via settings.json agent.max_tokens_per_job)
# AGENT_MAX_TOKENS_PER_JOB=0
# Enable planning phase before tool execution (default: true)
AGENT_USE_PLANNING=true

Expand Down
93 changes: 93 additions & 0 deletions src/agent/dispatcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1205,6 +1205,7 @@ mod tests {
max_tool_iterations: 50,
auto_approve_tools: false,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
},
deps,
Arc::new(ChannelManager::new()),
Expand Down Expand Up @@ -1263,6 +1264,96 @@ mod tests {
}
}

#[test]
fn test_always_approval_requirement_bypasses_session_auto_approve() {
// Regression test: even if tool is auto-approved in session,
// ApprovalRequirement::Always must still trigger approval.
use crate::tools::ApprovalRequirement;

let mut session = Session::new("user-1");
let tool_name = "tool_remove";

// Manually auto-approve tool_remove in this session
session.auto_approve_tool(tool_name);
assert!(
session.is_tool_auto_approved(tool_name),
"tool should be auto-approved"
);

// However, ApprovalRequirement::Always should always require approval
// This is verified by the dispatcher logic: Always => true (ignores session state)
let always_req = ApprovalRequirement::Always;
let requires_approval = match always_req {
ApprovalRequirement::Never => false,
ApprovalRequirement::UnlessAutoApproved => !session.is_tool_auto_approved(tool_name),
ApprovalRequirement::Always => true,
};

assert!(
requires_approval,
"ApprovalRequirement::Always must require approval even when tool is auto-approved"
);
}

#[test]
fn test_always_approval_requirement_vs_unless_auto_approved() {
// Verify the two requirements behave differently
use crate::tools::ApprovalRequirement;

let mut session = Session::new("user-2");
let tool_name = "http";

// Scenario 1: Tool is auto-approved
session.auto_approve_tool(tool_name);

// UnlessAutoApproved → doesn't require approval if auto-approved
let unless_req = ApprovalRequirement::UnlessAutoApproved;
let unless_needs = match unless_req {
ApprovalRequirement::Never => false,
ApprovalRequirement::UnlessAutoApproved => !session.is_tool_auto_approved(tool_name),
ApprovalRequirement::Always => true,
};
assert!(
!unless_needs,
"UnlessAutoApproved should not need approval when auto-approved"
);

// Always → always requires approval
let always_req = ApprovalRequirement::Always;
let always_needs = match always_req {
ApprovalRequirement::Never => false,
ApprovalRequirement::UnlessAutoApproved => !session.is_tool_auto_approved(tool_name),
ApprovalRequirement::Always => true,
};
assert!(
always_needs,
"Always must always require approval, even when auto-approved"
);

// Scenario 2: Tool is NOT auto-approved
let new_tool = "new_tool";
assert!(!session.is_tool_auto_approved(new_tool));

// UnlessAutoApproved → requires approval
let unless_needs = match unless_req {
ApprovalRequirement::Never => false,
ApprovalRequirement::UnlessAutoApproved => !session.is_tool_auto_approved(new_tool),
ApprovalRequirement::Always => true,
};
assert!(
unless_needs,
"UnlessAutoApproved should need approval when not auto-approved"
);

// Always → always requires approval
let always_needs = match always_req {
ApprovalRequirement::Never => false,
ApprovalRequirement::UnlessAutoApproved => !session.is_tool_auto_approved(new_tool),
ApprovalRequirement::Always => true,
};
assert!(always_needs, "Always must always require approval");
}

#[test]
fn test_pending_approval_serialization_backcompat_without_deferred_calls() {
// PendingApproval from before the deferred_tool_calls field was added
Expand Down Expand Up @@ -1953,6 +2044,7 @@ mod tests {
max_tool_iterations,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
},
deps,
Arc::new(ChannelManager::new()),
Expand Down Expand Up @@ -2069,6 +2161,7 @@ mod tests {
max_tool_iterations: max_iter,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
},
deps,
Arc::new(ChannelManager::new()),
Expand Down
16 changes: 16 additions & 0 deletions src/agent/scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,13 @@ impl Scheduler {
.create_job_for_user(user_id, title, description)
.await?;

// Apply token budget from config, allowing per-job metadata override.
let max_tokens = metadata
.as_ref()
.and_then(|m| m.get("max_tokens"))
.and_then(|v| v.as_u64())
.unwrap_or(self.config.max_tokens_per_job);

// Apply metadata if provided
if let Some(meta) = metadata {
self.context_manager
Expand All @@ -169,6 +176,15 @@ impl Scheduler {
.await?;
}

// Set token budget (separate update to avoid overwriting metadata)
if max_tokens > 0 {
self.context_manager
.update_context(job_id, |ctx| {
ctx.max_tokens = max_tokens;
})
.await?;
}

// Persist to DB before scheduling so the worker's FK references are valid
if let Some(ref store) = self.store {
let ctx = self.context_manager.get_context(job_id).await?;
Expand Down
103 changes: 100 additions & 3 deletions src/agent/worker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -417,7 +417,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#

iteration += 1;
if iteration > max_iterations {
self.mark_stuck("Maximum iterations exceeded").await?;
self.mark_failed("Maximum iterations exceeded: job hit the iteration cap")
.await?;
return Ok(());
}

Expand All @@ -437,7 +438,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
"LLM rate limited during tool selection, backing off"
);
if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS {
self.mark_stuck("Persistent rate limiting").await?;
self.mark_failed("Persistent rate limiting: exceeded retry limit")
.await?;
return Ok(());
}
self.log_event(
Expand Down Expand Up @@ -467,7 +469,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
"LLM rate limited during respond_with_tools, backing off"
);
if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS {
self.mark_stuck("Persistent rate limiting").await?;
self.mark_failed("Persistent rate limiting: exceeded retry limit")
.await?;
return Ok(());
}
self.log_event(
Expand All @@ -483,6 +486,20 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
Err(e) => return Err(e.into()),
};

// Track token usage from LLM call against the job budget.
// NOTE: select_tools() also makes LLM calls but doesn't expose
// TokenUsage; only respond_with_tools() usage is tracked here.
let total_tokens = respond_output.usage.total() as u64;
if total_tokens > 0
&& let Err(msg) = self
.context_manager()
.update_context(self.job_id, |ctx| ctx.add_tokens(total_tokens))
.await?
{
self.mark_failed(&msg).await?;
return Ok(());
}

match respond_output.result {
RespondResult::Text(response) => {
// Check for explicit completion phrases. Use word-boundary
Expand Down Expand Up @@ -1762,4 +1779,84 @@ mod tests {
"Always tool should be allowed with permission"
);
}

#[tokio::test]
async fn test_token_budget_exceeded_fails_job() {
let worker = make_worker(vec![]).await;

// Transition to InProgress (required for mark_failed)
worker
.context_manager()
.update_context(worker.job_id, |ctx| {
ctx.transition_to(JobState::InProgress, None)
})
.await
.unwrap()
.unwrap();

// Set a token budget
worker
.context_manager()
.update_context(worker.job_id, |ctx| {
ctx.max_tokens = 100;
})
.await
.unwrap();

// Simulate adding tokens that exceed the budget
let budget_result = worker
.context_manager()
.update_context(worker.job_id, |ctx| ctx.add_tokens(200))
.await
.unwrap();

assert!(
budget_result.is_err(),
"Should return error when token budget exceeded"
);

// Verify that mark_failed transitions job to Failed
worker
.mark_failed(&budget_result.unwrap_err())
.await
.unwrap();
let ctx = worker
.context_manager()
.get_context(worker.job_id)
.await
.unwrap();
assert_eq!(ctx.state, JobState::Failed);
}

#[tokio::test]
async fn test_iteration_cap_marks_failed_not_stuck() {
let worker = make_worker(vec![]).await;

// Transition to InProgress (required for mark_failed)
worker
.context_manager()
.update_context(worker.job_id, |ctx| {
ctx.transition_to(JobState::InProgress, None)
})
.await
.unwrap()
.unwrap();

// Simulate what the execution loop does when max_iterations is exceeded
worker
.mark_failed("Maximum iterations exceeded: job hit the iteration cap")
.await
.unwrap();

let ctx = worker
.context_manager()
.get_context(worker.job_id)
.await
.unwrap();
assert_eq!(
ctx.state,
JobState::Failed,
"Iteration cap should transition to Failed, not Stuck"
);
}
}
16 changes: 15 additions & 1 deletion src/channels/web/handlers/jobs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -276,11 +276,25 @@ pub async fn jobs_cancel_handler(
})));
}

// Fall back to agent job cancellation via DB status update.
// Fall back to agent job cancellation: stop the worker via the scheduler
// (which updates the in-memory ContextManager AND aborts the task handle),
// then persist the status to the DB as a fallback.
if let Some(ref store) = state.store
&& let Ok(Some(job)) = store.get_job(job_id).await
{
if job.state.is_active() {
// Try to stop via scheduler (aborts the worker task + updates
// in-memory ContextManager). This is best-effort — the job may
// not be in the scheduler map if it already finished.
if let Some(ref slot) = state.scheduler
&& let Some(ref scheduler) = *slot.read().await
{
let _ = scheduler.stop(job_id).await;
}

// Always persist cancellation to the DB so the state is
// consistent even if the scheduler wasn't available or the
// job wasn't in its in-memory map.
store
.update_job_status(
job_id,
Expand Down
7 changes: 7 additions & 0 deletions src/config/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ pub struct AgentConfig {
pub auto_approve_tools: bool,
/// Default timezone for new sessions (IANA name, e.g. "America/New_York").
pub default_timezone: String,
/// Maximum tokens per job (0 = unlimited).
pub max_tokens_per_job: u64,
}

impl AgentConfig {
Expand All @@ -50,6 +52,7 @@ impl AgentConfig {
max_tool_iterations: 10,
auto_approve_tools: true,
default_timezone: "UTC".to_string(),
max_tokens_per_job: 0,
}
}

Expand Down Expand Up @@ -105,6 +108,10 @@ impl AgentConfig {
}
tz
},
max_tokens_per_job: parse_optional_env(
"AGENT_MAX_TOKENS_PER_JOB",
settings.agent.max_tokens_per_job,
)?,
})
}
}
Expand Down
5 changes: 5 additions & 0 deletions src/settings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -386,6 +386,10 @@ pub struct AgentSettings {
/// Default timezone for new sessions (IANA name, e.g. "America/New_York").
#[serde(default = "default_timezone")]
pub default_timezone: String,

/// Maximum tokens per job (0 = unlimited).
#[serde(default)]
pub max_tokens_per_job: u64,
}

fn default_agent_name() -> String {
Expand Down Expand Up @@ -442,6 +446,7 @@ impl Default for AgentSettings {
max_tool_iterations: default_max_tool_iterations(),
auto_approve_tools: false,
default_timezone: default_timezone(),
max_tokens_per_job: 0,
}
}
}
Expand Down
Loading
Loading