Merge pull request #807 from nearai/staging-promote/83950d11-22884429853

chore: promote staging to main (2026-03-10 02:35 UTC)

.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
 

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()),

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?;

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"
+        );
+    }
 }

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,

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,
+            )?,
         })
     }
 }

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,
         }
     }
 }

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<String> {
         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<char> = 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"
+        );
+    }
 }

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(&params),
+                ApprovalRequirement::Always,
+                "tool_remove must always require approval for case: {}",
+                case_name
+            );
+        }
+    }
+
     #[test]
     fn test_tool_upgrade_schema() {
         use crate::tools::tool::ApprovalRequirement;

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(&params),
+                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());