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

make agent manager singleton (#4880)

Browse Source
This commit is contained in:
Yingjie He
2025-09-29 13:09:55 -07:00
committed by GitHub
parent 41d32c04b1
commit 1c424303c2
3 changed files with 95 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/state.rs
+1 -1
View File
@@ -18,7 +18,7 @@ pub struct AppState {
impl AppState {
pub async fn new() -> anyhow::Result<Arc<AppState>> {
let agent_manager = Arc::new(AgentManager::new(None).await?);
let agent_manager = AgentManager::instance().await?;
Ok(Arc::new(Self {
agent_manager,
recipe_file_hash_map: Arc::new(Mutex::new(HashMap::new())),
crates/goose/src/execution/manager.rs
+29 -3
View File
@@ -12,9 +12,13 @@ use etcetera::{choose_app_strategy, AppStrategy};
use lru::LruCache;
use std::num::NonZeroUsize;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::sync::{OnceCell, RwLock};
use tracing::{debug, info, warn};
const DEFAULT_MAX_SESSION: usize = 100;
static AGENT_MANAGER: OnceCell<Arc<AgentManager>> = OnceCell::const_new();
pub struct AgentManager {
sessions: Arc<RwLock<LruCache<String, Arc<Agent>>>>,
scheduler: Arc<dyn SchedulerTrait>,
@@ -22,7 +26,19 @@ pub struct AgentManager {
}
impl AgentManager {
pub async fn new(max_sessions: Option<usize>) -> Result<Self> {
/// Reset the global singleton - ONLY for testing
pub fn reset_for_test() {
unsafe {
// Cast away the const to get mutable access
// This is safe in test context where we control execution with #[serial]
let cell_ptr = &AGENT_MANAGER as *const OnceCell<Arc<AgentManager>>
as *mut OnceCell<Arc<AgentManager>>;
let _ = (*cell_ptr).take();
}
}
// Private constructor - prevents direct instantiation in production
async fn new(max_sessions: Option<usize>) -> Result<Self> {
// Construct scheduler with the standard goose-server path
let schedule_file_path = choose_app_strategy(APP_STRATEGY.clone())?
.data_dir()
@@ -30,7 +46,7 @@ impl AgentManager {
let scheduler = SchedulerFactory::create(schedule_file_path).await?;
let capacity = NonZeroUsize::new(max_sessions.unwrap_or(100))
let capacity = NonZeroUsize::new(max_sessions.unwrap_or(DEFAULT_MAX_SESSION))
.unwrap_or_else(|| NonZeroUsize::new(100).unwrap());
let manager = Self {
@@ -44,6 +60,16 @@ impl AgentManager {
Ok(manager)
}
pub async fn instance() -> Result<Arc<Self>> {
AGENT_MANAGER
.get_or_try_init(|| async {
let manager = Self::new(Some(DEFAULT_MAX_SESSION)).await?;
Ok(Arc::new(manager))
})
.await
.cloned()
}
pub async fn scheduler(&self) -> Result<Arc<dyn SchedulerTrait>> {
Ok(Arc::clone(&self.scheduler))
}
crates/goose/tests/execution_tests.rs
+65 -58
View File
@@ -25,8 +25,10 @@ mod execution_tests {
}
#[tokio::test]
#[serial]
async fn test_session_isolation() {
let manager = AgentManager::new(None).await.unwrap();
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let session1 = uuid::Uuid::new_v4().to_string();
let session2 = uuid::Uuid::new_v4().to_string();
@@ -51,13 +53,17 @@ mod execution_tests {
.unwrap();
assert!(Arc::ptr_eq(&agent1, &agent1_again));
AgentManager::reset_for_test();
}
#[tokio::test]
#[serial]
async fn test_session_limit() {
let manager = AgentManager::new(Some(3)).await.unwrap();
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let sessions: Vec<_> = (0..3).map(|i| format!("session-{}", i)).collect();
let sessions: Vec<_> = (0..100).map(|i| format!("session-{}", i)).collect();
for session in &sessions {
manager
@@ -73,13 +79,14 @@ mod execution_tests {
.await
.unwrap();
assert_eq!(manager.session_count().await, 3);
assert!(!manager.has_session(&sessions[0]).await);
assert_eq!(manager.session_count().await, 100);
}
#[tokio::test]
#[serial]
async fn test_remove_session() {
let manager = AgentManager::new(None).await.unwrap();
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let session = String::from("remove-test");
manager
@@ -92,11 +99,15 @@ mod execution_tests {
assert!(!manager.has_session(&session).await);
assert!(manager.remove_session(&session).await.is_err());
AgentManager::reset_for_test();
}
#[tokio::test]
#[serial]
async fn test_concurrent_access() {
let manager = Arc::new(AgentManager::new(None).await.unwrap());
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let session = String::from("concurrent-test");
let mut handles = vec![];
@@ -121,11 +132,15 @@ mod execution_tests {
}
assert_eq!(manager.session_count().await, 1);
AgentManager::reset_for_test();
}
#[tokio::test]
#[serial]
async fn test_different_modes_same_session() {
let manager = AgentManager::new(None).await.unwrap();
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let session_id = String::from("mode-test");
// Create initial agent
@@ -142,13 +157,17 @@ mod execution_tests {
.unwrap();
assert!(Arc::ptr_eq(&agent1, &agent2));
AgentManager::reset_for_test();
}
#[tokio::test]
#[serial]
async fn test_concurrent_session_creation_race_condition() {
// Test that concurrent attempts to create the same new session ID
// result in only one agent being created (tests double-check pattern)
let manager = Arc::new(AgentManager::new(None).await.unwrap());
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let session_id = String::from("race-condition-test");
// Spawn multiple tasks trying to create the same NEW session simultaneously
@@ -181,30 +200,8 @@ mod execution_tests {
// Only one session should exist
assert_eq!(manager.session_count().await, 1);
}
#[tokio::test]
async fn test_edge_case_max_sessions_one() {
let manager = AgentManager::new(Some(1)).await.unwrap();
let session1 = String::from("only-session");
manager
.get_or_create_agent(session1.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
assert_eq!(manager.session_count().await, 1);
// Creating second session should evict the first
let session2 = String::from("new-session");
manager
.get_or_create_agent(session2.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
assert!(!manager.has_session(&session1).await);
assert!(manager.has_session(&session2).await);
assert_eq!(manager.session_count().await, 1);
AgentManager::reset_for_test();
}
#[tokio::test]
@@ -212,13 +209,15 @@ mod execution_tests {
async fn test_configure_default_provider() {
use std::env;
AgentManager::reset_for_test();
let original_provider = env::var("GOOSE_DEFAULT_PROVIDER").ok();
let original_model = env::var("GOOSE_DEFAULT_MODEL").ok();
env::set_var("GOOSE_DEFAULT_PROVIDER", "openai");
env::set_var("GOOSE_DEFAULT_MODEL", "gpt-4o-mini");
let manager = AgentManager::new(None).await.unwrap();
let manager = AgentManager::instance().await.unwrap();
let result = manager.configure_default_provider().await;
assert!(result.is_ok());
@@ -234,14 +233,18 @@ mod execution_tests {
} else {
env::remove_var("GOOSE_DEFAULT_MODEL");
}
AgentManager::reset_for_test();
}
#[tokio::test]
#[serial]
async fn test_set_default_provider() {
use goose::providers::testprovider::TestProvider;
use std::sync::Arc;
let manager = AgentManager::new(None).await.unwrap();
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
// Create a test provider for replaying (doesn't need inner provider)
let temp_file = format!(
@@ -263,59 +266,63 @@ mod execution_tests {
.unwrap();
assert!(manager.has_session(&session).await);
AgentManager::reset_for_test();
}
#[tokio::test]
#[serial]
async fn test_eviction_updates_last_used() {
AgentManager::reset_for_test();
// Test that accessing a session updates its last_used timestamp
// and affects eviction order
let manager = AgentManager::new(Some(2)).await.unwrap();
let manager = AgentManager::instance().await.unwrap();
let session1 = String::from("session-1");
let session2 = String::from("session-2");
let sessions: Vec<_> = (0..100).map(|i| format!("session-{}", i)).collect();
manager
.get_or_create_agent(session1.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
for session in &sessions {
manager
.get_or_create_agent(session.clone(), SessionExecutionMode::chat())
.await
.unwrap();
// Small delay to ensure different timestamps
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
}
// Small delay to ensure different timestamps
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
manager
.get_or_create_agent(session2.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
// Access session1 again to update its last_used
// Access the first session again to update its last_used
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
manager
.get_or_create_agent(session1.clone(), SessionExecutionMode::Interactive)
.get_or_create_agent(sessions[0].clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
// Now create a third session - should evict session2 (least recently used)
let session3 = String::from("session-3");
// Now create a 101st session - should evict session2 (least recently used)
let session101 = String::from("session-101");
manager
.get_or_create_agent(session3.clone(), SessionExecutionMode::Interactive)
.get_or_create_agent(session101.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
// session1 should still exist (recently accessed)
// session2 should be evicted (least recently used)
assert!(manager.has_session(&session1).await);
assert!(!manager.has_session(&session2).await);
assert!(manager.has_session(&session3).await);
assert!(manager.has_session(&sessions[0]).await);
assert!(!manager.has_session(&sessions[1]).await);
assert!(manager.has_session(&session101).await);
AgentManager::reset_for_test();
}
#[tokio::test]
#[serial]
async fn test_remove_nonexistent_session_error() {
// Test that removing a non-existent session returns an error
let manager = AgentManager::new(None).await.unwrap();
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let session = String::from("never-created");
let result = manager.remove_session(&session).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("not found"));
AgentManager::reset_for_test();
}
}