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

feat(goose-acp): enable parallel sessions with isolated agent state (#6392)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-01-15 06:19:17 +08:00
committed by GitHub
parent fb0eca2c36
commit 7d4a6bd8ff
86 changed files with 2594 additions and 1938 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/config/base.rs
-1
View File
@@ -1100,7 +1100,6 @@ mod tests {
}
#[test]
#[serial]
fn test_multiple_secrets() -> Result<(), ConfigError> {
let config = new_test_config();
crates/goose/src/config/permission.rs
+57 -66
View File
@@ -3,8 +3,14 @@ use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, LazyLock, RwLock};
use utoipa::ToSchema;
const PERMISSION_FILE: &str = "permission.yaml";
static PERMISSION_MANAGER: LazyLock<Arc<PermissionManager>> =
LazyLock::new(|| Arc::new(PermissionManager::new(Paths::config_dir())));
/// Enum representing the possible permission levels for a tool.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq, ToSchema)]
#[serde(rename_all = "snake_case")]
@@ -23,62 +29,46 @@ pub struct PermissionConfig {
}
/// PermissionManager manages permission configurations for various tools.
#[derive(Debug, Clone)]
#[derive(Debug)]
pub struct PermissionManager {
config_path: PathBuf, // Path to the permission configuration file
permission_map: HashMap<String, PermissionConfig>, // Mapping of permission names to configurations
config_path: PathBuf,
permission_map: RwLock<HashMap<String, PermissionConfig>>,
}
// Constants representing specific permission categories
const USER_PERMISSION: &str = "user";
const SMART_APPROVE_PERMISSION: &str = "smart_approve";
/// Implements the default constructor for `PermissionManager`.
impl Default for PermissionManager {
fn default() -> Self {
let config_path = Paths::config_dir().join("permission.yaml");
// Load the existing configuration file or create an empty map if the file doesn't exist
let permission_map = if config_path.exists() {
// Load the configuration file
impl PermissionManager {
pub fn new(config_dir: PathBuf) -> Self {
let permission_path = config_dir.join(PERMISSION_FILE);
let permission_map = if permission_path.exists() {
let file_contents =
fs::read_to_string(&config_path).expect("Failed to read permission.yaml");
fs::read_to_string(&permission_path).expect("Failed to read permission.yaml");
serde_yaml::from_str(&file_contents).unwrap_or_else(|_| HashMap::new())
} else {
HashMap::new() // No config file, create an empty map
// Consolidate directory creation for re-use in global singleton or ACP.
fs::create_dir_all(&config_dir).expect("Failed to create config directory");
HashMap::new()
};
PermissionManager {
config_path,
permission_map,
config_path: permission_path,
permission_map: RwLock::new(permission_map),
}
}
}
impl PermissionManager {
/// Creates a new `PermissionManager` with a specified config path.
pub fn new<P: AsRef<Path>>(config_path: P) -> Self {
let config_path = config_path.as_ref().to_path_buf();
// Load the existing configuration file or create an empty map if the file doesn't exist
let permission_map = if config_path.exists() {
// Load the configuration file
let file_contents =
fs::read_to_string(&config_path).expect("Failed to read permission.yaml");
serde_yaml::from_str(&file_contents).unwrap_or_else(|_| HashMap::new())
} else {
HashMap::new() // No config file, create an empty map
};
PermissionManager {
config_path,
permission_map,
}
pub fn instance() -> Arc<PermissionManager> {
Arc::clone(&PERMISSION_MANAGER)
}
/// Returns a list of all the names (keys) in the permission map.
pub fn get_permission_names(&self) -> Vec<String> {
self.permission_map.keys().cloned().collect()
self.permission_map
.read()
.unwrap()
.keys()
.cloned()
.collect()
}
/// Retrieves the user permission level for a specific tool.
@@ -98,8 +88,9 @@ impl PermissionManager {
/// Helper function to retrieve the permission level for a specific permission category and tool.
fn get_permission(&self, name: &str, principal_name: &str) -> Option<PermissionLevel> {
let map = self.permission_map.read().unwrap();
// Check if the permission category exists in the map
if let Some(permission_config) = self.permission_map.get(name) {
if let Some(permission_config) = map.get(name) {
// Check the permission levels for the given tool
if permission_config
.always_allow
@@ -122,23 +113,20 @@ impl PermissionManager {
}
/// Updates the user permission level for a specific tool.
pub fn update_user_permission(&mut self, principal_name: &str, level: PermissionLevel) {
pub fn update_user_permission(&self, principal_name: &str, level: PermissionLevel) {
self.update_permission(USER_PERMISSION, principal_name, level)
}
/// Updates the smart approve permission level for a specific tool.
pub fn update_smart_approve_permission(
&mut self,
principal_name: &str,
level: PermissionLevel,
) {
pub fn update_smart_approve_permission(&self, principal_name: &str, level: PermissionLevel) {
self.update_permission(SMART_APPROVE_PERMISSION, principal_name, level)
}
/// Helper function to update a permission level for a specific tool in a given permission category.
fn update_permission(&mut self, name: &str, principal_name: &str, level: PermissionLevel) {
fn update_permission(&self, name: &str, principal_name: &str, level: PermissionLevel) {
let mut map = self.permission_map.write().unwrap();
// Get or create a new PermissionConfig for the specified category
let permission_config = self.permission_map.entry(name.to_string()).or_default();
let permission_config = map.entry(name.to_string()).or_default();
// Remove the principal from all existing lists to avoid duplicates
permission_config
@@ -163,14 +151,15 @@ impl PermissionManager {
}
// Serialize the updated permission map and write it back to the config file
let yaml_content = serde_yaml::to_string(&self.permission_map)
.expect("Failed to serialize permission config");
let yaml_content =
serde_yaml::to_string(&*map).expect("Failed to serialize permission config");
fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml");
}
/// Removes all entries where the principal name starts with the given extension name.
pub fn remove_extension(&mut self, extension_name: &str) {
for permission_config in self.permission_map.values_mut() {
pub fn remove_extension(&self, extension_name: &str) {
let mut map = self.permission_map.write().unwrap();
for permission_config in map.values_mut() {
permission_config
.always_allow
.retain(|p| !p.starts_with(extension_name));
@@ -182,8 +171,8 @@ impl PermissionManager {
.retain(|p| !p.starts_with(extension_name));
}
let yaml_content = serde_yaml::to_string(&self.permission_map)
.expect("Failed to serialize permission config");
let yaml_content =
serde_yaml::to_string(&*map).expect("Failed to serialize permission config");
fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml");
}
}
@@ -191,25 +180,25 @@ impl PermissionManager {
#[cfg(test)]
mod tests {
use super::*;
use tempfile::NamedTempFile;
use tempfile::TempDir;
// Helper function to create a test instance of PermissionManager with a temp dir
fn create_test_permission_manager() -> PermissionManager {
let temp_file = NamedTempFile::new().unwrap();
let temp_path = temp_file.path();
PermissionManager::new(temp_path)
fn create_test_permission_manager() -> (PermissionManager, TempDir) {
let temp_dir = TempDir::new().unwrap();
let manager = PermissionManager::new(temp_dir.path().to_path_buf());
(manager, temp_dir)
}
#[test]
fn test_get_permission_names_empty() {
let manager = create_test_permission_manager();
let (manager, _temp_dir) = create_test_permission_manager();
assert!(manager.get_permission_names().is_empty());
}
#[test]
fn test_update_user_permission() {
let mut manager = create_test_permission_manager();
let (manager, _temp_dir) = create_test_permission_manager();
manager.update_user_permission("tool1", PermissionLevel::AlwaysAllow);
let permission = manager.get_user_permission("tool1");
@@ -218,7 +207,7 @@ mod tests {
#[test]
fn test_update_smart_approve_permission() {
let mut manager = create_test_permission_manager();
let (manager, _temp_dir) = create_test_permission_manager();
manager.update_smart_approve_permission("tool2", PermissionLevel::AskBefore);
let permission = manager.get_smart_approve_permission("tool2");
@@ -227,7 +216,7 @@ mod tests {
#[test]
fn test_get_permission_not_found() {
let manager = create_test_permission_manager();
let (manager, _temp_dir) = create_test_permission_manager();
let permission = manager.get_user_permission("non_existent_tool");
assert_eq!(permission, None);
@@ -235,7 +224,7 @@ mod tests {
#[test]
fn test_permission_levels() {
let mut manager = create_test_permission_manager();
let (manager, _temp_dir) = create_test_permission_manager();
manager.update_user_permission("tool4", PermissionLevel::AlwaysAllow);
manager.update_user_permission("tool5", PermissionLevel::AskBefore);
@@ -258,7 +247,7 @@ mod tests {
#[test]
fn test_permission_update_replaces_existing_level() {
let mut manager = create_test_permission_manager();
let (manager, _temp_dir) = create_test_permission_manager();
// Initially AlwaysAllow
manager.update_user_permission("tool7", PermissionLevel::AlwaysAllow);
@@ -275,7 +264,8 @@ mod tests {
);
// Ensure it's removed from other levels
let config = manager.permission_map.get(USER_PERMISSION).unwrap();
let map = manager.permission_map.read().unwrap();
let config = map.get(USER_PERMISSION).unwrap();
assert!(!config.always_allow.contains(&"tool7".to_string()));
assert!(!config.ask_before.contains(&"tool7".to_string()));
assert!(config.never_allow.contains(&"tool7".to_string()));
@@ -283,7 +273,7 @@ mod tests {
#[test]
fn test_remove_extension() {
let mut manager = create_test_permission_manager();
let (manager, _temp_dir) = create_test_permission_manager();
manager.update_user_permission("prefix__tool1", PermissionLevel::AlwaysAllow);
manager.update_user_permission("nonprefix__tool2", PermissionLevel::AlwaysAllow);
manager.update_user_permission("prefix__tool3", PermissionLevel::AskBefore);
@@ -291,7 +281,8 @@ mod tests {
// Remove entries starting with "prefix"
manager.remove_extension("prefix");
let config = manager.permission_map.get(USER_PERMISSION).unwrap();
let map = manager.permission_map.read().unwrap();
let config = map.get(USER_PERMISSION).unwrap();
// Verify entries with "prefix" are removed
assert!(!config.always_allow.contains(&"prefix__tool1".to_string()));