From 7883225bd17cd432bc24ec8bde5a0a66c9bfea36 Mon Sep 17 00:00:00 2001
From: Dorien Koelemeijer <dkoelemeijer@squareup.com>
Date: Fri, 30 Jan 2026 12:02:27 +0700
Subject: [PATCH 1/5] Make sure command injection can be enabled without prompt
 injection using classifier models + small UI settings update

---
 crates/goose/src/security/mod.rs              |  10 +-
 .../settings/security/SecurityToggle.tsx      | 117 +++++++++---------
 2 files changed, 67 insertions(+), 60 deletions(-)

diff --git a/crates/goose/src/security/mod.rs b/crates/goose/src/security/mod.rs
index 780a0a45c45c..88b5d21487f1 100644
--- a/crates/goose/src/security/mod.rs
+++ b/crates/goose/src/security/mod.rs
@@ -43,9 +43,15 @@ impl SecurityManager {
     fn is_ml_scanning_enabled(&self) -> bool {
         let config = Config::global();
 
-        config
+        let prompt_enabled = config
             .get_param::<bool>("SECURITY_PROMPT_CLASSIFIER_ENABLED")
-            .unwrap_or(false)
+            .unwrap_or(false);
+
+        let command_enabled = config
+            .get_param::<bool>("SECURITY_COMMAND_CLASSIFIER_ENABLED")
+            .unwrap_or(false);
+
+        prompt_enabled || command_enabled
     }
 
     pub async fn analyze_tool_requests(
diff --git a/ui/desktop/src/components/settings/security/SecurityToggle.tsx b/ui/desktop/src/components/settings/security/SecurityToggle.tsx
index 42444ec31caf..775f4d7e54d7 100644
--- a/ui/desktop/src/components/settings/security/SecurityToggle.tsx
+++ b/ui/desktop/src/components/settings/security/SecurityToggle.tsx
@@ -271,17 +271,73 @@ export const SecurityToggle = () => {
             />
           </div>
 
-          {/* ML Detection Toggle */}
+          {/* Command Injection Detection Toggle */}
           <div className="border-t border-border-default pt-4">
             <div className="flex items-center justify-between py-2 hover:bg-background-muted rounded-lg transition-all">
               <div>
                 <h4
                   className={`text-sm font-medium ${enabled ? 'text-text-default' : 'text-text-muted'}`}
                 >
-                  Enable ML-Based Detection
+                  Enable Command Injection ML Detection
+                </h4>
+                <p className="text-xs text-text-muted max-w-md mt-[2px]">
+                  Use ML models to detect malicious shell commands (recommended)
+                </p>
+              </div>
+              <div className="flex items-center">
+                <Switch
+                  checked={effectiveCommandClassifierEnabled}
+                  onCheckedChange={handleCommandClassifierToggle}
+                  disabled={!enabled}
+                  variant="mono"
+                />
+              </div>
+            </div>
+
+            {hasCommandModel ? (
+              enabled &&
+              effectiveCommandClassifierEnabled && (
+                <div className="text-sm text-gray-700 dark:text-gray-300 mt-2">
+                  ✓ Command classifier active (auto-configured from environment)
+                </div>
+              )
+            ) : (
+              <div
+                className={`overflow-hidden transition-all duration-300 ease-in-out ${
+                  enabled && effectiveCommandClassifierEnabled
+                    ? 'max-h-[32rem] opacity-100 mt-3'
+                    : 'max-h-0 opacity-0'
+                }`}
+              >
+                <div className={enabled && effectiveCommandClassifierEnabled ? '' : 'opacity-50'}>
+                  <ClassifierEndpointInputs
+                    endpointValue={commandEndpointInput}
+                    tokenValue={commandTokenInput}
+                    onEndpointChange={setCommandEndpointInput}
+                    onTokenChange={setCommandTokenInput}
+                    onEndpointBlur={handleCommandEndpointChange}
+                    onTokenBlur={handleCommandTokenChange}
+                    disabled={!enabled || !effectiveCommandClassifierEnabled}
+                    endpointPlaceholder="https://example.com/classify"
+                    tokenPlaceholder="token..."
+                    endpointDescription="Enter the full URL for your command injection classification service"
+                  />
+                </div>
+              </div>
+            )}
+          </div>
+
+          {/* Prompt Injection Detection Toggle */}
+          <div className="border-t border-border-default pt-4">
+            <div className="flex items-center justify-between py-2 hover:bg-background-muted rounded-lg transition-all">
+              <div>
+                <h4
+                  className={`text-sm font-medium ${enabled ? 'text-text-default' : 'text-text-muted'}`}
+                >
+                  Enable Prompt Injection ML Detection
                 </h4>
                 <p className="text-xs text-text-muted max-w-md mt-[2px]">
-                  Use machine learning models for more accurate detection
+                  Use ML models to detect potential prompt injection in your chat
                 </p>
               </div>
               <div className="flex items-center">
@@ -348,61 +404,6 @@ export const SecurityToggle = () => {
               </div>
             </div>
           </div>
-
-          <div className="border-t border-border-default pt-4">
-            <div className="flex items-center justify-between py-2 hover:bg-background-muted rounded-lg transition-all">
-              <div>
-                <h4
-                  className={`text-sm font-medium ${enabled ? 'text-text-default' : 'text-text-muted'}`}
-                >
-                  Enable Command Injection ML Detection
-                </h4>
-                <p className="text-xs text-text-muted max-w-md mt-[2px]">
-                  Use ML models to detect malicious shell commands
-                </p>
-              </div>
-              <div className="flex items-center">
-                <Switch
-                  checked={effectiveCommandClassifierEnabled}
-                  onCheckedChange={handleCommandClassifierToggle}
-                  disabled={!enabled}
-                  variant="mono"
-                />
-              </div>
-            </div>
-
-            {hasCommandModel ? (
-              enabled &&
-              effectiveCommandClassifierEnabled && (
-                <div className="text-sm text-gray-700 dark:text-gray-300 mt-2">
-                  ✓ Command classifier active (auto-configured from environment)
-                </div>
-              )
-            ) : (
-              <div
-                className={`overflow-hidden transition-all duration-300 ease-in-out ${
-                  enabled && effectiveCommandClassifierEnabled
-                    ? 'max-h-[32rem] opacity-100 mt-3'
-                    : 'max-h-0 opacity-0'
-                }`}
-              >
-                <div className={enabled && effectiveCommandClassifierEnabled ? '' : 'opacity-50'}>
-                  <ClassifierEndpointInputs
-                    endpointValue={commandEndpointInput}
-                    tokenValue={commandTokenInput}
-                    onEndpointChange={setCommandEndpointInput}
-                    onTokenChange={setCommandTokenInput}
-                    onEndpointBlur={handleCommandEndpointChange}
-                    onTokenBlur={handleCommandTokenChange}
-                    disabled={!enabled || !effectiveCommandClassifierEnabled}
-                    endpointPlaceholder="https://example.com/classify"
-                    tokenPlaceholder="token..."
-                    endpointDescription="Enter the full URL for your command injection classification service"
-                  />
-                </div>
-              </div>
-            )}
-          </div>
         </div>
       </div>
     </div>

From a71bdcf7b56962e71cf0afceef7f68a01fecb9be Mon Sep 17 00:00:00 2001
From: Dorien Koelemeijer <dkoelemeijer@squareup.com>
Date: Fri, 30 Jan 2026 12:25:17 +0700
Subject: [PATCH 2/5] Small update combine confidence function

---
 crates/goose/src/security/scanner.rs | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/crates/goose/src/security/scanner.rs b/crates/goose/src/security/scanner.rs
index 2cb7612b23c5..6ce6f46ecf3c 100644
--- a/crates/goose/src/security/scanner.rs
+++ b/crates/goose/src/security/scanner.rs
@@ -149,18 +149,21 @@ impl PromptInjectionScanner {
         tracing::info!(
             "Classifier Results - Command: {:.3}, Prompt: {:.3}, Threshold: {:.3}",
             tool_result.confidence,
-            context_result.confidence,
+            context_result.ml_confidence.unwrap_or(0.0),
             threshold
         );
 
-        let final_confidence =
-            self.combine_confidences(tool_result.confidence, context_result.confidence);
+        let final_confidence = self.combine_confidences(
+            tool_result.confidence,
+            context_result.ml_confidence, // Use ml_confidence (Option<f32>) instead of confidence
+        );
 
         tracing::info!(
             tool_confidence = %tool_result.confidence,
-            context_confidence = %context_result.confidence,
+            context_confidence = ?context_result.ml_confidence,
             final_confidence = %final_confidence,
-            has_ml = tool_result.ml_confidence.is_some(),
+            has_command_ml = tool_result.ml_confidence.is_some(),
+            has_prompt_ml = context_result.ml_confidence.is_some(),
             has_patterns = !tool_result.pattern_matches.is_empty(),
             threshold = %threshold,
             malicious = final_confidence >= threshold,
@@ -239,7 +242,11 @@ impl PromptInjectionScanner {
         })
     }
 
-    fn combine_confidences(&self, tool_confidence: f32, context_confidence: f32) -> f32 {
+    fn combine_confidences(&self, tool_confidence: f32, context_confidence: Option<f32>) -> f32 {
+        let Some(context_confidence) = context_confidence else {
+            return tool_confidence;
+        };
+
         // If tool is safe, context is not taken into account
         if tool_confidence < 0.3 {
             return tool_confidence;

From 3996d2c6e26cc557471fab6b2e250f730f0e93eb Mon Sep 17 00:00:00 2001
From: Dorien Koelemeijer <dkoelemeijer@squareup.com>
Date: Fri, 30 Jan 2026 12:36:27 +0700
Subject: [PATCH 3/5] fix

---
 crates/goose/src/security/scanner.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/crates/goose/src/security/scanner.rs b/crates/goose/src/security/scanner.rs
index 6ce6f46ecf3c..cc9b26b7d5db 100644
--- a/crates/goose/src/security/scanner.rs
+++ b/crates/goose/src/security/scanner.rs
@@ -155,7 +155,7 @@ impl PromptInjectionScanner {
 
         let final_confidence = self.combine_confidences(
             tool_result.confidence,
-            context_result.ml_confidence, // Use ml_confidence (Option<f32>) instead of confidence
+            context_result.ml_confidence, 
         );
 
         tracing::info!(

From 9cd8d5bfe410631e338a18606b8bb78c5de03b7c Mon Sep 17 00:00:00 2001
From: Dorien Koelemeijer <dkoelemeijer@squareup.com>
Date: Fri, 30 Jan 2026 12:48:19 +0700
Subject: [PATCH 4/5] fix

---
 crates/goose/src/security/scanner.rs | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/crates/goose/src/security/scanner.rs b/crates/goose/src/security/scanner.rs
index cc9b26b7d5db..dca32e1ed323 100644
--- a/crates/goose/src/security/scanner.rs
+++ b/crates/goose/src/security/scanner.rs
@@ -153,10 +153,8 @@ impl PromptInjectionScanner {
             threshold
         );
 
-        let final_confidence = self.combine_confidences(
-            tool_result.confidence,
-            context_result.ml_confidence, 
-        );
+        let final_confidence =
+            self.combine_confidences(tool_result.confidence, context_result.ml_confidence);
 
         tracing::info!(
             tool_confidence = %tool_result.confidence,

From d7563e0368237b130dd7cdf3e98e68fa8b6c2f6f Mon Sep 17 00:00:00 2001
From: Dorien Koelemeijer <dkoelemeijer@squareup.com>
Date: Fri, 30 Jan 2026 12:49:57 +0700
Subject: [PATCH 5/5] fix

---
 ui/desktop/src/components/settings/security/SecurityToggle.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ui/desktop/src/components/settings/security/SecurityToggle.tsx b/ui/desktop/src/components/settings/security/SecurityToggle.tsx
index 775f4d7e54d7..58858f465f37 100644
--- a/ui/desktop/src/components/settings/security/SecurityToggle.tsx
+++ b/ui/desktop/src/components/settings/security/SecurityToggle.tsx
@@ -281,7 +281,7 @@ export const SecurityToggle = () => {
                   Enable Command Injection ML Detection
                 </h4>
                 <p className="text-xs text-text-muted max-w-md mt-[2px]">
-                  Use ML models to detect malicious shell commands (recommended)
+                  Use ML models to detect malicious shell commands
                 </p>
               </div>
               <div className="flex items-center">