tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Prompt injection detection (simplified - only pattern matching) (#4237)

Browse Source 
Signed-off-by: Dorien Koelemeijer <dkoelemeijer@squareup.com>

merging as looks good, lets keep an eye on it.
This commit is contained in:
dorien-koelemeijer
2025-09-10 10:41:30 +10:00
committed by GitHub
parent 9bb1bb530c
commit 916ba902dc
20 changed files with 2043 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+92 -48
View File
@@ -1,4 +1,4 @@
use std::collections::{HashMap, HashSet};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@@ -31,18 +31,21 @@ use crate::agents::tool_route_manager::ToolRouteManager;
use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
use crate::agents::types::SessionConfig;
use crate::agents::types::{FrontendTool, ToolResultReceiver};
use crate::config::{Config, ExtensionConfigManager, PermissionManager};
use crate::config::{Config, ExtensionConfigManager};
use crate::context_mgmt::auto_compact;
use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation};
use crate::permission::permission_judge::{check_tool_permissions, PermissionCheckResult};
use crate::permission::permission_inspector::PermissionInspector;
use crate::permission::permission_judge::PermissionCheckResult;
use crate::permission::PermissionConfirmation;
use crate::providers::base::Provider;
use crate::providers::errors::ProviderError;
use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
use crate::scheduler_trait::SchedulerTrait;
use crate::security::security_inspector::SecurityInspector;
use crate::session;
use crate::session::extension_data::ExtensionState;
use crate::tool_monitor::{ToolCall, ToolMonitor};
use crate::tool_inspection::ToolInspectionManager;
use crate::tool_monitor::RepetitionInspector;
use crate::utils::is_token_cancelled;
use mcp_core::ToolResult;
use regex::Regex;
@@ -81,8 +84,6 @@ pub struct ToolCategorizeResult {
pub frontend_requests: Vec<ToolRequest>,
pub remaining_requests: Vec<ToolRequest>,
pub filtered_response: Message,
pub readonly_tools: HashSet<String>,
pub regular_tools: HashSet<String>,
}
/// The main goose Agent
@@ -99,10 +100,11 @@ pub struct Agent {
pub(super) confirmation_rx: Mutex<mpsc::Receiver<(String, PermissionConfirmation)>>,
pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult<Vec<Content>>)>,
pub(super) tool_result_rx: ToolResultReceiver,
pub(super) tool_monitor: Arc<Mutex<Option<ToolMonitor>>>,
pub(super) tool_route_manager: ToolRouteManager,
pub(super) scheduler_service: Mutex<Option<Arc<dyn SchedulerTrait>>>,
pub(super) retry_manager: RetryManager,
pub(super) tool_inspection_manager: ToolInspectionManager,
pub(super) autopilot: Mutex<AutoPilot>,
}
@@ -160,9 +162,6 @@ impl Agent {
let (confirm_tx, confirm_rx) = mpsc::channel(32);
let (tool_tx, tool_rx) = mpsc::channel(32);
let tool_monitor = Arc::new(Mutex::new(None));
let retry_manager = RetryManager::with_tool_monitor(tool_monitor.clone());
Self {
provider: Mutex::new(None),
extension_manager: ExtensionManager::new(),
@@ -176,17 +175,33 @@ impl Agent {
confirmation_rx: Mutex::new(confirm_rx),
tool_result_tx: tool_tx,
tool_result_rx: Arc::new(Mutex::new(tool_rx)),
tool_monitor,
tool_route_manager: ToolRouteManager::new(),
scheduler_service: Mutex::new(None),
retry_manager,
retry_manager: RetryManager::new(),
tool_inspection_manager: Self::create_default_tool_inspection_manager(),
autopilot: Mutex::new(AutoPilot::new()),
}
}
pub async fn configure_tool_monitor(&self, max_repetitions: Option<u32>) {
let mut tool_monitor = self.tool_monitor.lock().await;
*tool_monitor = Some(ToolMonitor::new(max_repetitions));
/// Create a tool inspection manager with default inspectors
fn create_default_tool_inspection_manager() -> ToolInspectionManager {
let mut tool_inspection_manager = ToolInspectionManager::new();
// Add security inspector (highest priority - runs first)
tool_inspection_manager.add_inspector(Box::new(SecurityInspector::new()));
// Add permission inspector (medium-high priority)
// Note: mode will be updated dynamically based on session config
tool_inspection_manager.add_inspector(Box::new(PermissionInspector::new(
"smart_approve".to_string(),
std::collections::HashSet::new(), // readonly tools - will be populated from extension manager
std::collections::HashSet::new(), // regular tools - will be populated from extension manager
)));
// Add repetition inspector (lower priority - basic repetition checking)
tool_inspection_manager.add_inspector(Box::new(RepetitionInspector::new(None)));
tool_inspection_manager
}
/// Reset the retry attempts counter to 0
@@ -247,6 +262,11 @@ impl Agent {
let (tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt().await?;
let goose_mode = Self::determine_goose_mode(session.as_ref(), config);
// Update permission inspector mode to match the session mode
self.tool_inspection_manager
.update_permission_inspector_mode(goose_mode.clone())
.await;
Ok(ReplyContext {
messages: conversation,
tools,
@@ -261,10 +281,8 @@ impl Agent {
async fn categorize_tools(
&self,
response: &Message,
tools: &[rmcp::model::Tool],
_tools: &[rmcp::model::Tool],
) -> ToolCategorizeResult {
let (readonly_tools, regular_tools) = Self::categorize_tools_by_annotation(tools);
// Categorize tool requests
let (frontend_requests, remaining_requests, filtered_response) =
self.categorize_tool_requests(response).await;
@@ -273,8 +291,6 @@ impl Agent {
frontend_requests,
remaining_requests,
filtered_response,
readonly_tools,
regular_tools,
}
}
@@ -378,22 +394,6 @@ impl Agent {
cancellation_token: Option<CancellationToken>,
session: &Option<SessionConfig>,
) -> (String, Result<ToolCallResult, ErrorData>) {
// Check if this tool call should be allowed based on repetition monitoring
if let Some(monitor) = self.tool_monitor.lock().await.as_mut() {
let tool_call_info = ToolCall::new(tool_call.name.clone(), tool_call.arguments.clone());
if !monitor.check_tool_call(tool_call_info) {
return (
request_id,
Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
"Tool call rejected: exceeded maximum allowed repetitions".to_string(),
None,
)),
);
}
}
if tool_call.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME {
let result = self
.handle_schedule_management(tool_call.arguments, request_id.clone())
@@ -1119,8 +1119,6 @@ impl Agent {
frontend_requests,
remaining_requests,
filtered_response,
readonly_tools,
regular_tools,
} = self.categorize_tools(&response, &tools).await;
let requests_to_record: Vec<ToolRequest> = frontend_requests.iter().chain(remaining_requests.iter()).cloned().collect();
self.tool_route_manager
@@ -1159,16 +1157,40 @@ impl Agent {
);
}
} else {
let mut permission_manager = PermissionManager::default();
let (permission_check_result, enable_extension_request_ids) =
check_tool_permissions(
// Run all tool inspectors (security, repetition, permission, etc.)
let inspection_results = self.tool_inspection_manager
.inspect_tools(
&remaining_requests,
&mode,
readonly_tools.clone(),
regular_tools.clone(),
&mut permission_manager,
self.provider().await?,
).await;
messages.messages(),
)
.await?;
// Process inspection results into permission decisions using the permission inspector
let permission_check_result = self.tool_inspection_manager
.process_inspection_results_with_permission_inspector(
&remaining_requests,
&inspection_results,
)
.unwrap_or_else(|| {
// Fallback if permission inspector not found - default to needs approval
let mut result = PermissionCheckResult {
approved: vec![],
needs_approval: vec![],
denied: vec![],
};
result.needs_approval.extend(remaining_requests.iter().cloned());
result
});
// Track extension requests for special handling
let mut enable_extension_request_ids = vec![];
for request in &remaining_requests {
if let Ok(tool_call) = &request.tool_call {
if tool_call.name == PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME {
enable_extension_request_ids.push(request.id.clone());
}
}
}
let mut tool_futures = self.handle_approved_and_denied_tools(
&permission_check_result,
@@ -1183,9 +1205,9 @@ impl Agent {
let mut tool_approval_stream = self.handle_approval_tool_requests(
&permission_check_result.needs_approval,
tool_futures_arc.clone(),
&mut permission_manager,
message_tool_response.clone(),
cancel_token.clone(),
&inspection_results,
);
while let Some(msg) = tool_approval_stream.try_next().await? {
@@ -1675,6 +1697,28 @@ mod tests {
assert!(todo_read.is_some(), "TODO read tool should be present");
assert!(todo_write.is_some(), "TODO write tool should be present");
Ok(())
}
#[tokio::test]
async fn test_tool_inspection_manager_has_all_inspectors() -> Result<()> {
let agent = Agent::new();
// Verify that the tool inspection manager has all expected inspectors
let inspector_names = agent.tool_inspection_manager.inspector_names();
assert!(
inspector_names.contains(&"repetition"),
"Tool inspection manager should contain repetition inspector"
);
assert!(
inspector_names.contains(&"permission"),
"Tool inspection manager should contain permission inspector"
);
assert!(
inspector_names.contains(&"security"),
"Tool inspection manager should contain security inspector"
);
Ok(())
}
crates/goose/src/agents/reply_parts.rs
-23
View File
@@ -1,5 +1,4 @@
use anyhow::Result;
use std::collections::HashSet;
use std::sync::Arc;
use async_stream::try_stream;
@@ -85,28 +84,6 @@ impl Agent {
Ok((tools, toolshim_tools, system_prompt))
}
/// Categorize tools based on their annotations
/// Returns:
/// - read_only_tools: Tools with read-only annotations
/// - non_read_tools: Tools without read-only annotations
pub(crate) fn categorize_tools_by_annotation(
tools: &[Tool],
) -> (HashSet<String>, HashSet<String>) {
tools
.iter()
.fold((HashSet::new(), HashSet::new()), |mut acc, tool| {
match &tool.annotations {
Some(annotations) if annotations.read_only_hint.unwrap_or(false) => {
acc.0.insert(tool.name.to_string());
}
_ => {
acc.1.insert(tool.name.to_string());
}
}
acc
})
}
/// Generate a response from the LLM provider
/// Handles toolshim transformations if needed
pub(crate) async fn generate_response_from_provider(
crates/goose/src/agents/retry.rs
+13 -11
View File
@@ -13,7 +13,7 @@ use crate::agents::types::{
use crate::config::Config;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
use crate::tool_monitor::ToolMonitor;
use crate::tool_monitor::RepetitionInspector;
/// Result of a retry logic evaluation
#[derive(Debug, Clone, PartialEq)]
@@ -39,8 +39,8 @@ const GOOSE_RECIPE_ON_FAILURE_TIMEOUT_SECONDS: &str = "GOOSE_RECIPE_ON_FAILURE_T
pub struct RetryManager {
/// Current number of retry attempts
attempts: Arc<Mutex<u32>>,
/// Optional tool monitor for reset operations
tool_monitor: Option<Arc<Mutex<Option<ToolMonitor>>>>,
/// Optional repetition inspector for reset operations
repetition_inspector: Option<Arc<Mutex<Option<RepetitionInspector>>>>,
}
impl Default for RetryManager {
@@ -54,15 +54,17 @@ impl RetryManager {
pub fn new() -> Self {
Self {
attempts: Arc::new(Mutex::new(0)),
tool_monitor: None,
repetition_inspector: None,
}
}
/// Create a new retry manager with tool monitor
pub fn with_tool_monitor(tool_monitor: Arc<Mutex<Option<ToolMonitor>>>) -> Self {
/// Create a new retry manager with repetition inspector
pub fn with_repetition_inspector(
repetition_inspector: Arc<Mutex<Option<RepetitionInspector>>>,
) -> Self {
Self {
attempts: Arc::new(Mutex::new(0)),
tool_monitor: Some(tool_monitor),
repetition_inspector: Some(repetition_inspector),
}
}
@@ -71,10 +73,10 @@ impl RetryManager {
let mut attempts = self.attempts.lock().await;
*attempts = 0;
// Reset tool monitor if available
if let Some(monitor) = &self.tool_monitor {
if let Some(monitor) = monitor.lock().await.as_mut() {
monitor.reset();
// Reset repetition inspector if available
if let Some(inspector) = &self.repetition_inspector {
if let Some(inspector) = inspector.lock().await.as_mut() {
inspector.reset();
}
}
}
crates/goose/src/agents/tool_execution.rs
+18 -5
View File
@@ -8,7 +8,6 @@ use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use crate::config::permission::PermissionLevel;
use crate::config::PermissionManager;
use crate::permission::Permission;
use mcp_core::ToolResult;
use rmcp::model::{Content, ServerNotification};
@@ -51,18 +50,29 @@ impl Agent {
&'a self,
tool_requests: &'a [ToolRequest],
tool_futures: Arc<Mutex<Vec<(String, ToolStream)>>>,
permission_manager: &'a mut PermissionManager,
message_tool_response: Arc<Mutex<Message>>,
cancellation_token: Option<CancellationToken>,
inspection_results: &'a [crate::tool_inspection::InspectionResult],
) -> BoxStream<'a, anyhow::Result<Message>> {
try_stream! {
for request in tool_requests {
for request in tool_requests.iter() {
if let Ok(tool_call) = request.tool_call.clone() {
// Find the corresponding inspection result for this tool request
let security_message = inspection_results.iter()
.find(|result| result.tool_request_id == request.id)
.and_then(|result| {
if let crate::tool_inspection::InspectionAction::RequireApproval(Some(message)) = &result.action {
Some(message.clone())
} else {
None
}
});
let confirmation = Message::user().with_tool_confirmation_request(
request.id.clone(),
tool_call.name.clone(),
tool_call.arguments.clone(),
Some("Goose would like to call the above tool. Allow? (y/n):".to_string()),
security_message,
);
yield confirmation;
@@ -84,8 +94,11 @@ impl Agent {
),
}));
// Update the shared permission manager when user selects "Always Allow"
if confirmation.permission == Permission::AlwaysAllow {
permission_manager.update_user_permission(&tool_call.name, PermissionLevel::AlwaysAllow);
self.tool_inspection_manager
.update_permission_manager(&tool_call.name, PermissionLevel::AlwaysAllow)
.await;
}
} else {
// User declined - add declined response