diff --git a/.env.example b/.env.example index 1200400d..5c21e995 100644 --- a/.env.example +++ b/.env.example @@ -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 diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index f5306644..99feed9d 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -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()), @@ -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 @@ -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()), @@ -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()), diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 99386d7f..85f3f6eb 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -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 @@ -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?; diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 3604cea9..19bfc8e5 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -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(()); } @@ -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( @@ -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( @@ -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 @@ -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" + ); + } } diff --git a/src/channels/web/handlers/jobs.rs b/src/channels/web/handlers/jobs.rs index 8a127243..5a94e055 100644 --- a/src/channels/web/handlers/jobs.rs +++ b/src/channels/web/handlers/jobs.rs @@ -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, diff --git a/src/config/agent.rs b/src/config/agent.rs index 096c141f..cb09707d 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -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 { @@ -50,6 +52,7 @@ impl AgentConfig { max_tool_iterations: 10, auto_approve_tools: true, default_timezone: "UTC".to_string(), + max_tokens_per_job: 0, } } @@ -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, + )?, }) } } diff --git a/src/settings.rs b/src/settings.rs index 836d1d2c..63535aef 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -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 { @@ -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, } } } diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 86142bff..33e1b262 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1573,46 +1573,18 @@ impl SetupWizard { } /// Fetch available models from the NEAR AI API. + /// + /// Uses [`build_nearai_model_fetch_config`] to construct the provider config, + /// which reads `NEARAI_API_KEY` from the environment when present. async fn fetch_nearai_models(&self) -> Vec { let session = match self.session_manager { Some(ref s) => Arc::clone(s), None => return vec![], }; - use crate::config::LlmConfig; use crate::llm::create_llm_provider; - let base_url = std::env::var("NEARAI_BASE_URL") - .unwrap_or_else(|_| "https://private.near.ai".to_string()); - let auth_base_url = std::env::var("NEARAI_AUTH_URL") - .unwrap_or_else(|_| "https://private.near.ai".to_string()); - - let config = LlmConfig { - backend: "nearai".to_string(), - session: crate::llm::session::SessionConfig { - auth_base_url, - session_path: crate::config::llm::default_session_path(), - }, - nearai: crate::config::NearAiConfig { - model: "dummy".to_string(), - cheap_model: None, - base_url, - api_key: None, - fallback_model: None, - max_retries: 3, - circuit_breaker_threshold: None, - circuit_breaker_recovery_secs: 30, - response_cache_enabled: false, - response_cache_ttl_secs: 3600, - response_cache_max_entries: 1000, - failover_cooldown_secs: 300, - failover_cooldown_threshold: 3, - smart_routing_cascade: true, - }, - provider: None, - bedrock: None, - request_timeout_secs: 120, - }; + let config = build_nearai_model_fetch_config(); match create_llm_provider(&config, session).await { Ok(provider) => match provider.list_models().await { @@ -3240,6 +3212,52 @@ async fn discover_wasm_channels(dir: &std::path::Path) -> Vec<(String, ChannelCa /// Mask an API key for display: show first 6 + last 4 chars. /// /// Uses char-based indexing to avoid panicking on multi-byte UTF-8. +/// Build the `LlmConfig` used by `fetch_nearai_models` to list available models. +/// +/// Reads `NEARAI_API_KEY` from the environment so that users who authenticated +/// via Cloud API key (option 4) don't get re-prompted during model selection. +fn build_nearai_model_fetch_config() -> crate::config::LlmConfig { + let base_url = + std::env::var("NEARAI_BASE_URL").unwrap_or_else(|_| "https://private.near.ai".to_string()); + let auth_base_url = + std::env::var("NEARAI_AUTH_URL").unwrap_or_else(|_| "https://private.near.ai".to_string()); + + // If the user authenticated via API key (option 4), the key is stored + // as an env var. Pass it through so `resolve_bearer_token()` doesn't + // re-trigger the interactive auth prompt. + let api_key = std::env::var("NEARAI_API_KEY") + .ok() + .filter(|k| !k.is_empty()) + .map(secrecy::SecretString::from); + + crate::config::LlmConfig { + backend: "nearai".to_string(), + session: crate::llm::session::SessionConfig { + auth_base_url, + session_path: crate::config::llm::default_session_path(), + }, + nearai: crate::config::NearAiConfig { + model: "dummy".to_string(), + cheap_model: None, + base_url, + api_key, + fallback_model: None, + max_retries: 3, + circuit_breaker_threshold: None, + circuit_breaker_recovery_secs: 30, + response_cache_enabled: false, + response_cache_ttl_secs: 3600, + response_cache_max_entries: 1000, + failover_cooldown_secs: 300, + failover_cooldown_threshold: 3, + smart_routing_cascade: true, + }, + provider: None, + bedrock: None, + request_timeout_secs: 120, + } +} + fn mask_api_key(key: &str) -> String { let chars: Vec = key.chars().collect(); if chars.len() < 12 { @@ -3640,6 +3658,14 @@ mod tests { } impl EnvGuard { + fn set(key: &'static str, value: &str) -> Self { + let original = std::env::var(key).ok(); + unsafe { + std::env::set_var(key, value); + } + Self { key, original } + } + fn clear(key: &'static str) -> Self { let original = std::env::var(key).ok(); unsafe { @@ -3827,4 +3853,49 @@ mod tests { }; assert!(settings.secrets_master_key_hex.is_some()); } + + /// Regression test for #799: `fetch_nearai_models` hardcoded `api_key: None`, + /// causing the auth prompt to re-appear during model selection when the user + /// had authenticated via NEAR AI Cloud API key (option 4). + #[test] + fn test_build_nearai_model_fetch_config_picks_up_api_key_env() { + use secrecy::ExposeSecret; + + let _guard = EnvGuard::set("NEARAI_API_KEY", "test-cloud-api-key-12345"); + + let config = build_nearai_model_fetch_config(); + assert!( + config.nearai.api_key.is_some(), + "config should include NEARAI_API_KEY from env" + ); + assert_eq!( + config.nearai.api_key.as_ref().unwrap().expose_secret(), + "test-cloud-api-key-12345" + ); + } + + /// Regression test for #799: when NEARAI_API_KEY is absent or empty, + /// the config should have `api_key: None` (session token path). + #[test] + fn test_build_nearai_model_fetch_config_none_when_no_api_key() { + let _guard = EnvGuard::clear("NEARAI_API_KEY"); + + let config = build_nearai_model_fetch_config(); + assert!( + config.nearai.api_key.is_none(), + "config should have no api_key when env var is absent" + ); + } + + /// Regression test for #799: empty NEARAI_API_KEY should be treated as absent. + #[test] + fn test_build_nearai_model_fetch_config_none_when_empty_api_key() { + let _guard = EnvGuard::set("NEARAI_API_KEY", ""); + + let config = build_nearai_model_fetch_config(); + assert!( + config.nearai.api_key.is_none(), + "config should have no api_key when env var is empty" + ); + } } diff --git a/src/tools/builtin/extension_tools.rs b/src/tools/builtin/extension_tools.rs index 7ba4ef0c..00c79548 100644 --- a/src/tools/builtin/extension_tools.rs +++ b/src/tools/builtin/extension_tools.rs @@ -451,8 +451,8 @@ impl Tool for ToolRemoveTool { } fn description(&self) -> &str { - "Remove an installed extension (channel, tool, or MCP server). \ - Unregisters tools and deletes configuration." + "Permanently remove an installed extension (channel, tool, or MCP server) from disk. \ + This action cannot be undone — the WASM binary and configuration files will be deleted." } fn parameters_schema(&self) -> serde_json::Value { @@ -492,7 +492,7 @@ impl Tool for ToolRemoveTool { } fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - ApprovalRequirement::UnlessAutoApproved + ApprovalRequirement::Always } } @@ -701,10 +701,38 @@ mod tests { assert_eq!(tool.name(), "tool_remove"); assert_eq!( tool.requires_approval(&serde_json::json!({})), - ApprovalRequirement::UnlessAutoApproved + ApprovalRequirement::Always ); } + #[test] + fn tool_remove_always_requires_approval_regardless_of_params() { + use crate::tools::tool::ApprovalRequirement; + let tool = ToolRemoveTool { + manager: test_manager_stub(), + }; + + let test_cases = vec![ + ("no params", serde_json::json!({})), + ("empty name", serde_json::json!({"name": ""})), + ("slack", serde_json::json!({"name": "slack"})), + ("github-cli", serde_json::json!({"name": "github-cli"})), + ( + "with extra fields", + serde_json::json!({"name": "tool", "extra": "field"}), + ), + ]; + + for (case_name, params) in test_cases { + assert_eq!( + tool.requires_approval(¶ms), + ApprovalRequirement::Always, + "tool_remove must always require approval for case: {}", + case_name + ); + } + } + #[test] fn test_tool_upgrade_schema() { use crate::tools::tool::ApprovalRequirement; diff --git a/src/tools/builtin/skill_tools.rs b/src/tools/builtin/skill_tools.rs index 84c889ae..a7581ac4 100644 --- a/src/tools/builtin/skill_tools.rs +++ b/src/tools/builtin/skill_tools.rs @@ -709,7 +709,8 @@ impl Tool for SkillRemoveTool { } fn description(&self) -> &str { - "Remove an installed skill by name. Only user-installed skills can be removed." + "Permanently remove an installed skill from disk. This action cannot be undone — \ + the skill files will be deleted." } fn parameters_schema(&self) -> serde_json::Value { @@ -770,7 +771,7 @@ impl Tool for SkillRemoveTool { } fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - ApprovalRequirement::UnlessAutoApproved + ApprovalRequirement::Always } } @@ -837,12 +838,41 @@ mod tests { assert_eq!(tool.name(), "skill_remove"); assert_eq!( tool.requires_approval(&serde_json::json!({})), - ApprovalRequirement::UnlessAutoApproved + ApprovalRequirement::Always ); let schema = tool.parameters_schema(); assert!(schema["properties"].get("name").is_some()); } + #[test] + fn skill_remove_always_requires_approval_regardless_of_params() { + use crate::tools::tool::ApprovalRequirement; + let tool = SkillRemoveTool::new(test_registry()); + + let test_cases = vec![ + ("no params", serde_json::json!({})), + ("empty name", serde_json::json!({"name": ""})), + ( + "deployment skill", + serde_json::json!({"name": "deployment"}), + ), + ("custom skill", serde_json::json!({"name": "custom-skill"})), + ( + "with extra fields", + serde_json::json!({"name": "skill", "extra": "field"}), + ), + ]; + + for (case_name, params) in test_cases { + assert_eq!( + tool.requires_approval(¶ms), + ApprovalRequirement::Always, + "skill_remove must always require approval for case: {}", + case_name + ); + } + } + #[test] fn test_validate_fetch_url_allows_https() { assert!(super::validate_fetch_url("https://clawhub.ai/api/v1/download?slug=foo").is_ok());