Skip to content

make agent manager singleton #4880

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yingjiehe-xyz merged 5 commits into from
Sep 29, 2025
Merged

make agent manager singleton #4880

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
bebce3d
make agent manager singleton
yingjiehe-xyz Sep 29, 2025
f429e8c
make agent manager singleton
yingjiehe-xyz Sep 29, 2025
4ef0142
fix singleton
yingjiehe-xyz Sep 29, 2025
e97d816
fix tests
yingjiehe-xyz Sep 29, 2025
586b2e2
resolve comments
yingjiehe-xyz Sep 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion crates/goose-server/src/state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())),
Expand Down
33 changes: 30 additions & 3 deletions crates/goose/src/execution/manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,25 +12,41 @@ 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>,
default_provider: Arc<RwLock<Option<Arc<dyn crate::providers::base::Provider>>>>,
}

impl AgentManager {
pub async fn new(max_sessions: Option<usize>) -> Result<Self> {
/// Reset the global singleton - ONLY for testing
pub fn reset_for_test() {
Copy link
Collaborator

@DOsinga DOsinga Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you do want to annotate this with test only though

Copy link
Contributor Author

@yingjiehe-xyz yingjiehe-xyz Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it work here, we can use #[cfg(test)] if the tests are agent class are in the same package, otherwise, tests cannot find the reset_for_test

see error:

no function or associated item named `reset_for_test` found for struct `AgentManager` in the current scope

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()
.join("schedule.json");

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 {
Expand All @@ -44,6 +60,17 @@ impl AgentManager {
Ok(manager)
}

/// Get or initialize the singleton instance
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))
}
Expand Down
166 changes: 75 additions & 91 deletions crates/goose/tests/execution_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Copy link
Collaborator

@DOsinga DOsinga Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we need to reset these? I'd imagine you only want to reset for testing for situations where you actually want to test that a new singleton does somethign different

Copy link
Contributor Author

@yingjiehe-xyz yingjiehe-xyz Sep 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it is better to reset every time to keep the tests isolated from each other

Copy link
Collaborator

@DOsinga DOsinga Sep 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah, no. you should isolate these tests from the rest of the system, but if you write tests for a singleton and then reset the singleton all the time, what are you testing?

Copy link
Contributor Author

@yingjiehe-xyz yingjiehe-xyz Sep 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you mean that we want to cover the singleton access crossing multiple test cases? Unit tests should only cover each simple functionalities of the methods, we have a test case to test the multiple access as the same time, and others for the main logic, to make all these tests isolated, we need such reset.

let manager = AgentManager::instance().await.unwrap();

let session1 = uuid::Uuid::new_v4().to_string();
let session2 = uuid::Uuid::new_v4().to_string();
Expand All @@ -51,35 +53,15 @@ mod execution_tests {
.unwrap();

assert!(Arc::ptr_eq(&agent1, &agent1_again));
}

#[tokio::test]
async fn test_session_limit() {
let manager = AgentManager::new(Some(3)).await.unwrap();

let sessions: Vec<_> = (0..3).map(|i| format!("session-{}", i)).collect();

for session in &sessions {
manager
.get_or_create_agent(session.clone(), SessionExecutionMode::chat())
.await
.unwrap();
}

// Create a new session after cleanup
let new_session = "new-session".to_string();
let _new_agent = manager
.get_or_create_agent(new_session, SessionExecutionMode::chat())
.await
.unwrap();

assert_eq!(manager.session_count().await, 3);
assert!(!manager.has_session(&sessions[0]).await);
AgentManager::reset_for_test();
}

#[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
Expand All @@ -92,11 +74,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![];
Expand All @@ -121,11 +107,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
Expand All @@ -142,13 +132,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
Expand Down Expand Up @@ -181,44 +175,24 @@ 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]
#[serial]
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());
Expand All @@ -234,14 +208,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!(
Expand All @@ -263,59 +241,65 @@ mod execution_tests {
.unwrap();

assert!(manager.has_session(&session).await);

AgentManager::reset_for_test();
}

#[tokio::test]
async fn test_eviction_updates_last_used() {
// Test that accessing a session updates its last_used timestamp
// and affects eviction order
let manager = AgentManager::new(Some(2)).await.unwrap();
#[serial]
async fn test_remove_nonexistent_session_error() {
// Test that removing a non-existent session returns an error
AgentManager::reset_for_test();
let manager = AgentManager::instance().await.unwrap();
let session = String::from("never-created");

let session1 = String::from("session-1");
let session2 = String::from("session-2");
let result = manager.remove_session(&session).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("not found"));

manager
.get_or_create_agent(session1.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
AgentManager::reset_for_test();
}

// Small delay to ensure different timestamps
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
#[tokio::test]
#[serial]
async fn test_singleton_instance() {
AgentManager::reset_for_test();

manager
.get_or_create_agent(session2.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
let instance1 = AgentManager::instance().await.unwrap();
let instance2 = AgentManager::instance().await.unwrap();

// Access session1 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)
.await
.unwrap();
assert!(Arc::ptr_eq(&instance1, &instance2));

// Now create a third session - should evict session2 (least recently used)
let session3 = String::from("session-3");
manager
.get_or_create_agent(session3.clone(), SessionExecutionMode::Interactive)
let session_id = String::from("singleton-test");
instance1
.get_or_create_agent(session_id.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!(instance2.has_session(&session_id).await);

AgentManager::reset_for_test();
}

#[tokio::test]
async fn test_remove_nonexistent_session_error() {
// Test that removing a non-existent session returns an error
let manager = AgentManager::new(None).await.unwrap();
let session = String::from("never-created");
#[serial]
async fn test_singleton_persistence_across_calls() {
AgentManager::reset_for_test();

let result = manager.remove_session(&session).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("not found"));
let session_id = String::from("persistence-test");
{
let manager = AgentManager::instance().await.unwrap();
manager
.get_or_create_agent(session_id.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
}

{
let manager = AgentManager::instance().await.unwrap();
assert!(manager.has_session(&session_id).await);
}

AgentManager::reset_for_test();
}
}
Loading