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

feat: AgentManager - foundation for unified execution (#4389) (#4684)

Browse Source
This commit is contained in:
tlongwell-block
2025-09-24 16:15:26 -04:00
committed by GitHub
parent ad045bc94c
commit 9d40422e2d
34 changed files with 896 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/commands/agent.rs
+1 -49
View File
@@ -1,13 +1,7 @@
use std::sync::Arc;
use crate::configuration;
use crate::state;
use anyhow::Result;
use axum::middleware;
use etcetera::{choose_app_strategy, AppStrategy};
use goose::agents::Agent;
use goose::config::APP_STRATEGY;
use goose::scheduler_factory::SchedulerFactory;
use goose_server::auth::check_token;
use tower_http::cors::{Any, CorsLayer};
use tracing::info;
@@ -32,49 +26,7 @@ pub async fn run() -> Result<()> {
let secret_key =
std::env::var("GOOSE_SERVER__SECRET_KEY").unwrap_or_else(|_| "test".to_string());
let new_agent = Agent::new();
// Only initialize provider and extensions when running in standalone goosed mode
// This prevents breaking the Electron app which manages its own provider setup
if std::env::var("GOOSE_STANDALONE_MODE").unwrap_or_else(|_| "false".to_string()) == "true" {
tracing::info!("Running in standalone mode - initializing provider and extensions");
// Initialize provider like the CLI does
let config = goose::config::Config::global();
let provider_name: String = config
.get_param("GOOSE_PROVIDER")
.expect("No provider configured. Run 'goose configure' first");
let model_name: String = config
.get_param("GOOSE_MODEL")
.expect("No model configured. Run 'goose configure' first");
let model_config = goose::model::ModelConfig::new(&model_name)
.expect("Failed to create model configuration");
let provider = goose::providers::create(&provider_name, model_config)
.expect("Failed to create provider");
new_agent
.update_provider(provider)
.await
.expect("Failed to update agent provider");
}
let agent_ref = Arc::new(new_agent);
let app_state = state::AppState::new(agent_ref.clone());
let schedule_file_path = choose_app_strategy(APP_STRATEGY.clone())?
.data_dir()
.join("schedules.json");
let scheduler_instance = SchedulerFactory::create(schedule_file_path).await?;
app_state.set_scheduler(scheduler_instance.clone()).await;
// NEW: Provide scheduler access to the agent
agent_ref.set_scheduler(scheduler_instance).await;
let app_state = state::AppState::new().await?;
let cors = CorsLayer::new()
.allow_origin(Any)
crates/goose-server/src/routes/agent.rs
+26 -35
View File
@@ -1,5 +1,4 @@
use crate::state::AppState;
use axum::response::IntoResponse;
use axum::{
extract::{Query, State},
http::StatusCode,
@@ -28,7 +27,6 @@ use tracing::error;
#[derive(Deserialize, utoipa::ToSchema)]
pub struct ExtendPromptRequest {
extension: String,
#[allow(dead_code)]
session_id: String,
}
@@ -40,7 +38,6 @@ pub struct ExtendPromptResponse {
#[derive(Deserialize, utoipa::ToSchema)]
pub struct AddSubRecipesRequest {
sub_recipes: Vec<SubRecipe>,
#[allow(dead_code)]
session_id: String,
}
@@ -53,27 +50,23 @@ pub struct AddSubRecipesResponse {
pub struct UpdateProviderRequest {
provider: String,
model: Option<String>,
#[allow(dead_code)]
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct SessionConfigRequest {
response: Option<Response>,
#[allow(dead_code)]
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct GetToolsQuery {
extension_name: Option<String>,
#[allow(dead_code)]
session_id: String,
}
#[derive(Deserialize, utoipa::ToSchema)]
pub struct UpdateRouterToolSelectorRequest {
#[allow(dead_code)]
session_id: String,
}
@@ -116,8 +109,6 @@ async fn start_agent(
State(state): State<Arc<AppState>>,
Json(payload): Json<StartAgentRequest>,
) -> Result<Json<StartAgentResponse>, StatusCode> {
state.reset().await;
let session_id = session::generate_session_id();
let counter = state.session_counter.fetch_add(1, Ordering::SeqCst) + 1;
@@ -203,7 +194,7 @@ async fn add_sub_recipes(
State(state): State<Arc<AppState>>,
Json(payload): Json<AddSubRecipesRequest>,
) -> Result<Json<AddSubRecipesResponse>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(payload.session_id).await?;
agent.add_sub_recipes(payload.sub_recipes.clone()).await;
Ok(Json(AddSubRecipesResponse { success: true }))
}
@@ -222,7 +213,7 @@ async fn extend_prompt(
State(state): State<Arc<AppState>>,
Json(payload): Json<ExtendPromptRequest>,
) -> Result<Json<ExtendPromptResponse>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(payload.session_id).await?;
agent.extend_system_prompt(payload.extension.clone()).await;
Ok(Json(ExtendPromptResponse { success: true }))
}
@@ -247,7 +238,7 @@ async fn get_tools(
) -> Result<Json<Vec<ToolInfo>>, StatusCode> {
let config = Config::global();
let goose_mode = config.get_param("GOOSE_MODE").unwrap_or("auto".to_string());
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(query.session_id).await?;
let permission_manager = PermissionManager::default();
let mut tools: Vec<ToolInfo> = agent
@@ -298,35 +289,37 @@ async fn get_tools(
async fn update_agent_provider(
State(state): State<Arc<AppState>>,
Json(payload): Json<UpdateProviderRequest>,
) -> Result<StatusCode, impl IntoResponse> {
let agent = state.get_agent().await;
) -> Result<StatusCode, StatusCode> {
let agent = state
.get_agent_for_route(payload.session_id.clone())
.await?;
let config = Config::global();
let model = match payload
.model
.or_else(|| config.get_param("GOOSE_MODEL").ok())
{
Some(m) => m,
None => return Err((StatusCode::BAD_REQUEST, "No model specified".to_string())),
None => {
tracing::error!("No model specified");
return Err(StatusCode::BAD_REQUEST);
}
};
let model_config = ModelConfig::new(&model).map_err(|e| {
(
StatusCode::BAD_REQUEST,
format!("Invalid model config: {}", e),
)
tracing::error!("Invalid model config: {}", e);
StatusCode::BAD_REQUEST
})?;
let new_provider = create(&payload.provider, model_config).map_err(|e| {
(
StatusCode::BAD_REQUEST,
format!("Failed to create provider: {}", e),
)
tracing::error!("Failed to create provider: {}", e);
StatusCode::BAD_REQUEST
})?;
agent
.update_provider(new_provider)
.await
.map_err(|_e| (StatusCode::INTERNAL_SERVER_ERROR, String::new()))?;
agent.update_provider(new_provider).await.map_err(|e| {
tracing::error!("Failed to update provider: {}", e);
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(StatusCode::OK)
}
@@ -344,17 +337,15 @@ async fn update_agent_provider(
)]
async fn update_router_tool_selector(
State(state): State<Arc<AppState>>,
Json(_payload): Json<UpdateRouterToolSelectorRequest>,
) -> Result<Json<String>, Json<ErrorResponse>> {
let agent = state.get_agent().await;
Json(payload): Json<UpdateRouterToolSelectorRequest>,
) -> Result<Json<String>, StatusCode> {
let agent = state.get_agent_for_route(payload.session_id).await?;
agent
.update_router_tool_selector(None, Some(true))
.await
.map_err(|e| {
tracing::error!("Failed to update tool selection strategy: {}", e);
Json(ErrorResponse {
error: format!("Failed to update tool selection strategy: {}", e),
})
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(Json(
@@ -376,8 +367,8 @@ async fn update_router_tool_selector(
async fn update_session_config(
State(state): State<Arc<AppState>>,
Json(payload): Json<SessionConfigRequest>,
) -> Result<Json<String>, Json<ErrorResponse>> {
let agent = state.get_agent().await;
) -> Result<Json<String>, StatusCode> {
let agent = state.get_agent_for_route(payload.session_id).await?;
if let Some(response) = payload.response {
agent.add_final_output_tool(response).await;
crates/goose-server/src/routes/audio.rs
+13 -35
View File
@@ -391,13 +391,13 @@ pub fn routes(state: Arc<AppState>) -> Router {
mod tests {
use super::*;
use axum::{body::Body, http::Request};
use serde_json::json;
use tower::ServiceExt;
#[tokio::test]
#[tokio::test(flavor = "multi_thread")]
async fn test_transcribe_endpoint_requires_auth() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
let state = AppState::new().await.unwrap();
let app = routes(state);
// Test without auth header
let request = Request::builder()
.uri("/audio/transcribe")
@@ -413,40 +413,18 @@ mod tests {
.unwrap();
let response = app.oneshot(request).await.unwrap();
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert!(
response.status() == StatusCode::PRECONDITION_FAILED
|| response.status() == StatusCode::UNAUTHORIZED
);
}
#[tokio::test]
#[tokio::test(flavor = "multi_thread")]
async fn test_transcribe_endpoint_validates_size() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
let app = routes(state);
// Create a large base64 string (simulating > 25MB audio)
let large_audio = BASE64.encode(vec![0u8; MAX_AUDIO_SIZE_BYTES + 1]);
let request = Request::builder()
.uri("/audio/transcribe")
.method("POST")
.header("content-type", "application/json")
.header("x-secret-key", "test-secret")
.body(Body::from(
serde_json::to_string(&serde_json::json!({
"audio": large_audio,
"mime_type": "audio/webm"
}))
.unwrap(),
))
.unwrap();
let response = app.oneshot(request).await.unwrap();
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
}
#[tokio::test]
async fn test_transcribe_endpoint_validates_mime_type() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
let state = AppState::new().await.unwrap();
let app = routes(state);
let large_data = "a".repeat(30 * 1024 * 1024); // 30MB
let request = Request::builder()
.uri("/audio/transcribe")
.method("POST")
@@ -468,9 +446,9 @@ mod tests {
);
}
#[tokio::test]
async fn test_transcribe_endpoint_handles_invalid_base64() {
let state = AppState::new(Arc::new(goose::agents::Agent::new()));
#[tokio::test(flavor = "multi_thread")]
async fn test_transcribe_endpoint_validates_mime_type() {
let state = AppState::new().await.unwrap();
let app = routes(state);
let request = Request::builder()
crates/goose-server/src/routes/context.rs
+3 -1
View File
@@ -13,6 +13,8 @@ pub struct ContextManageRequest {
pub messages: Vec<Message>,
/// Operation to perform: "truncation" or "summarize"
pub manage_action: String,
/// Optional session ID for session-specific agent
pub session_id: String,
}
/// Response from context management operations
@@ -44,7 +46,7 @@ async fn manage_context(
State(state): State<Arc<AppState>>,
Json(request): Json<ContextManageRequest>,
) -> Result<Json<ContextManageResponse>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(request.session_id).await?;
let mut processed_messages = Conversation::new_unvalidated(vec![]);
let mut token_counts: Vec<usize> = vec![];
crates/goose-server/src/routes/extension.rs
+27 -22
View File
@@ -96,33 +96,31 @@ struct ExtensionResponse {
message: Option<String>,
}
/// Request structure for adding an extension, combining session_id with the extension config
#[derive(Deserialize)]
struct AddExtensionRequest {
session_id: String,
#[serde(flatten)]
config: ExtensionConfigRequest,
}
/// Handler for adding a new extension configuration.
async fn add_extension(
State(state): State<Arc<AppState>>,
raw: axum::extract::Json<serde_json::Value>,
Json(request): Json<AddExtensionRequest>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
// Log the raw request for debugging
// Log the request for debugging
tracing::info!(
"Received extension request: {}",
serde_json::to_string_pretty(&raw.0).unwrap()
"Received extension request for session: {}",
request.session_id
);
// Try to parse into our enum
let request: ExtensionConfigRequest = match serde_json::from_value(raw.0.clone()) {
Ok(req) => req,
Err(e) => {
tracing::error!("Failed to parse extension request: {}", e);
tracing::error!(
"Raw request was: {}",
serde_json::to_string_pretty(&raw.0).unwrap()
);
return Err(StatusCode::UNPROCESSABLE_ENTITY);
}
};
let session_id = request.session_id.clone();
let extension_request = request.config;
// If this is a Stdio extension that uses npx, check for Node.js installation
#[cfg(target_os = "windows")]
if let ExtensionConfigRequest::Stdio { cmd, .. } = &request {
if let ExtensionConfigRequest::Stdio { cmd, .. } = &extension_request {
if cmd.ends_with("npx.cmd") || cmd.ends_with("npx") {
// Check if Node.js is installed in standard locations
let node_exists = std::path::Path::new(r"C:\Program Files\nodejs\node.exe").exists()
@@ -175,7 +173,7 @@ async fn add_extension(
}
// Construct ExtensionConfig with Envs populated from keyring based on provided env_keys.
let extension_config: ExtensionConfig = match request {
let extension_config: ExtensionConfig = match extension_request {
ExtensionConfigRequest::Sse {
name,
uri,
@@ -267,7 +265,7 @@ async fn add_extension(
},
};
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(session_id).await?;
let response = agent.add_extension(extension_config).await;
// Respond with the result.
@@ -289,13 +287,20 @@ async fn add_extension(
}
}
#[derive(Deserialize)]
struct RemoveExtensionRequest {
name: String,
session_id: String,
}
/// Handler for removing an extension by name
async fn remove_extension(
State(state): State<Arc<AppState>>,
Json(name): Json<String>,
Json(request): Json<RemoveExtensionRequest>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
let agent = state.get_agent().await;
match agent.remove_extension(&name).await {
let agent = state.get_agent_for_route(request.session_id).await?;
match agent.remove_extension(&request.name).await {
Ok(_) => Ok(Json(ExtensionResponse {
error: false,
message: None,
crates/goose-server/src/routes/recipe.rs
+4 -8
View File
@@ -25,6 +25,7 @@ pub struct CreateRecipeRequest {
activities: Option<Vec<String>>,
#[serde(default)]
author: Option<AuthorRequest>,
session_id: String,
}
#[derive(Debug, Deserialize, ToSchema)]
@@ -108,13 +109,13 @@ pub struct ListRecipeResponse {
async fn create_recipe(
State(state): State<Arc<AppState>>,
Json(request): Json<CreateRecipeRequest>,
) -> Result<Json<CreateRecipeResponse>, (StatusCode, Json<CreateRecipeResponse>)> {
) -> Result<Json<CreateRecipeResponse>, StatusCode> {
tracing::info!(
"Recipe creation request received with {} messages",
request.messages.len()
);
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(request.session_id).await?;
// Create base recipe from agent state and messages
let recipe_result = agent
@@ -143,12 +144,7 @@ async fn create_recipe(
}
Err(e) => {
tracing::error!("Error details: {:?}", e);
let error_message = format!("Recipe creation failed: {}", e);
let error_response = CreateRecipeResponse {
recipe: None,
error: Some(error_message),
};
Err((StatusCode::BAD_REQUEST, Json(error_response)))
Err(StatusCode::BAD_REQUEST)
}
}
}
crates/goose-server/src/routes/reply.rs
+32 -23
View File
@@ -10,6 +10,7 @@ use bytes::Bytes;
use futures::{stream::StreamExt, Stream};
use goose::conversation::message::{Message, MessageContent};
use goose::conversation::Conversation;
use goose::execution::SessionExecutionMode;
use goose::{
agents::{AgentEvent, SessionConfig},
permission::permission_confirmation::PrincipalType,
@@ -86,7 +87,7 @@ fn track_tool_telemetry(content: &MessageContent, all_messages: &[Message]) {
#[derive(Debug, Deserialize, Serialize)]
struct ChatRequest {
messages: Vec<Message>,
session_id: Option<String>,
session_id: String,
recipe_name: Option<String>,
recipe_version: Option<String>,
}
@@ -178,9 +179,9 @@ async fn reply_handler(
"Session started"
);
if let (Some(recipe_name), Some(session_id)) =
(request.recipe_name.clone(), request.session_id.clone())
{
let session_id = request.session_id.clone();
if let Some(recipe_name) = request.recipe_name.clone() {
if state.mark_recipe_run_if_absent(&session_id).await {
let recipe_version = request
.recipe_version
@@ -204,16 +205,28 @@ async fn reply_handler(
let messages = Conversation::new_unvalidated(request.messages);
let session_id = request.session_id.ok_or_else(|| {
tracing::error!("session_id is required but was not provided");
StatusCode::BAD_REQUEST
})?;
let task_cancel = cancel_token.clone();
let task_tx = tx.clone();
drop(tokio::spawn(async move {
let agent = state.get_agent().await;
let agent = match state
.get_agent(session_id.clone(), SessionExecutionMode::Interactive)
.await
{
Ok(agent) => agent,
Err(e) => {
tracing::error!("Failed to get session agent: {}", e);
let _ = stream_event(
MessageEvent::Error {
error: format!("Failed to get session agent: {}", e),
},
&task_tx,
&task_cancel,
)
.await;
return;
}
};
// Load session metadata to get the working directory and other config
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
@@ -453,7 +466,6 @@ pub struct PermissionConfirmationRequest {
#[serde(default = "default_principal_type")]
principal_type: PrincipalType,
action: String,
#[allow(dead_code)]
session_id: String,
}
@@ -475,7 +487,7 @@ pub async fn confirm_permission(
State(state): State<Arc<AppState>>,
Json(request): Json<PermissionConfirmationRequest>,
) -> Result<Json<Value>, StatusCode> {
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(request.session_id).await?;
let permission = match request.action.as_str() {
"always_allow" => Permission::AlwaysAllow,
"allow_once" => Permission::AllowOnce,
@@ -499,7 +511,6 @@ pub async fn confirm_permission(
struct ToolResultRequest {
id: String,
result: ToolResult<Vec<Content>>,
#[allow(dead_code)]
session_id: String,
}
@@ -524,7 +535,7 @@ async fn submit_tool_result(
}
};
let agent = state.get_agent().await;
let agent = state.get_agent_for_route(payload.session_id).await?;
agent.handle_tool_result(payload.id, payload.result).await;
Ok(Json(json!({"status": "ok"})))
}
@@ -548,7 +559,6 @@ mod tests {
use super::*;
use goose::conversation::message::Message;
use goose::{
agents::Agent,
model::ModelConfig,
providers::{
base::{Provider, ProviderUsage, Usage},
@@ -589,18 +599,17 @@ mod tests {
use super::*;
use axum::{body::Body, http::Request};
use goose::conversation::message::Message;
use std::sync::Arc;
use serde_json::json;
use tower::ServiceExt;
#[tokio::test]
#[tokio::test(flavor = "multi_thread")]
async fn test_reply_endpoint() {
let mock_model_config = ModelConfig::new("test-model").unwrap();
let mock_provider = Arc::new(MockProvider {
let mock_provider = MockProvider {
model_config: mock_model_config,
});
let agent = Agent::new();
let _ = agent.update_provider(mock_provider).await;
let state = AppState::new(Arc::new(agent));
};
let state = AppState::new().await.unwrap();
let app = routes(state);
@@ -612,7 +621,7 @@ mod tests {
.body(Body::from(
serde_json::to_string(&ChatRequest {
messages: vec![Message::user().with_text("test message")],
session_id: Some("test-session".to_string()),
session_id: "test-session".to_string(),
recipe_name: None,
recipe_version: None,
})
crates/goose-server/src/routes/session.rs
+1
View File
@@ -60,6 +60,7 @@ pub struct SessionInsights {
}
#[derive(Serialize, ToSchema, Debug)]
#[allow(dead_code)]
#[serde(rename_all = "camelCase")]
pub struct ActivityHeatmapCell {
pub week: usize,
crates/goose-server/src/state.rs
+33 -61
View File
@@ -1,18 +1,15 @@
use goose::agents::Agent;
use axum::http::StatusCode;
use goose::execution::manager::AgentManager;
use goose::execution::SessionExecutionMode;
use goose::scheduler_trait::SchedulerTrait;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::RwLock;
type AgentRef = Arc<Agent>;
#[derive(Clone)]
pub struct AppState {
agent: Arc<RwLock<AgentRef>>,
pub scheduler: Arc<RwLock<Option<Arc<dyn SchedulerTrait>>>>,
pub(crate) agent_manager: Arc<AgentManager>,
pub recipe_file_hash_map: Arc<Mutex<HashMap<String, PathBuf>>>,
pub session_counter: Arc<AtomicUsize>,
/// Tracks sessions that have already emitted recipe telemetry to prevent double counting.
@@ -20,31 +17,18 @@ pub struct AppState {
}
impl AppState {
pub fn new(agent: AgentRef) -> Arc<AppState> {
Arc::new(Self {
agent: Arc::new(RwLock::new(agent)),
scheduler: Arc::new(RwLock::new(None)),
pub async fn new() -> anyhow::Result<Arc<AppState>> {
let agent_manager = Arc::new(AgentManager::new(None).await?);
Ok(Arc::new(Self {
agent_manager,
recipe_file_hash_map: Arc::new(Mutex::new(HashMap::new())),
session_counter: Arc::new(AtomicUsize::new(0)),
recipe_session_tracker: Arc::new(Mutex::new(HashSet::new())),
})
}
pub async fn get_agent(&self) -> AgentRef {
self.agent.read().await.clone()
}
pub async fn set_scheduler(&self, sched: Arc<dyn SchedulerTrait>) {
let mut guard = self.scheduler.write().await;
*guard = Some(sched);
}))
}
pub async fn scheduler(&self) -> Result<Arc<dyn SchedulerTrait>, anyhow::Error> {
self.scheduler
.read()
.await
.clone()
.ok_or_else(|| anyhow::anyhow!("Scheduler not initialized"))
self.agent_manager.scheduler().await
}
pub async fn set_recipe_file_hash_map(&self, hash_map: HashMap<String, PathBuf>) {
@@ -52,41 +36,6 @@ impl AppState {
*map = hash_map;
}
pub async fn reset(&self) {
let mut agent = self.agent.write().await;
let new_agent = Agent::new();
// Only initialize provider when running in standalone goosed mode
// This prevents breaking the Electron app which manages its own provider setup
if std::env::var("GOOSE_STANDALONE_MODE").unwrap_or_else(|_| "false".to_string()) == "true"
{
tracing::info!("Running in standalone mode - initializing provider");
let config = goose::config::Config::global();
let provider_name: String = config
.get_param("GOOSE_PROVIDER")
.expect("No provider configured. Run 'goose configure' first");
let model_name: String = config
.get_param("GOOSE_MODEL")
.expect("No model configured. Run 'goose configure' first");
let model_config = goose::model::ModelConfig::new(&model_name)
.expect("Failed to create model configuration");
let provider = goose::providers::create(&provider_name, model_config)
.expect("Failed to create provider");
new_agent
.update_provider(provider)
.await
.expect("Failed to update agent provider");
}
*agent = Arc::new(new_agent);
}
pub async fn mark_recipe_run_if_absent(&self, session_id: &str) -> bool {
let mut sessions = self.recipe_session_tracker.lock().await;
if sessions.contains(session_id) {
@@ -96,4 +45,27 @@ impl AppState {
true
}
}
pub async fn get_agent(
&self,
session_id: String,
mode: SessionExecutionMode,
) -> anyhow::Result<Arc<goose::agents::Agent>> {
self.agent_manager
.get_or_create_agent(session_id, mode)
.await
}
/// Get agent for route handlers - always uses Interactive mode and converts any error to 500
pub async fn get_agent_for_route(
&self,
session_id: String,
) -> Result<Arc<goose::agents::Agent>, StatusCode> {
self.get_agent(session_id, SessionExecutionMode::Interactive)
.await
.map_err(|e| {
tracing::error!("Failed to get agent: {}", e);
StatusCode::INTERNAL_SERVER_ERROR
})
}
}
crates/goose-server/tests/pricing_api_test.rs
-41
View File
@@ -1,41 +0,0 @@
use axum::http::StatusCode;
use axum::Router;
use axum::{body::Body, http::Request};
use etcetera::AppStrategy;
use serde_json::json;
use std::sync::Arc;
use tower::ServiceExt;
async fn create_test_app() -> Router {
let agent = Arc::new(goose::agents::Agent::default());
let state = goose_server::AppState::new(agent);
// Add scheduler setup like in the existing tests
let sched_storage_path = etcetera::choose_app_strategy(goose::config::APP_STRATEGY.clone())
.unwrap()
.data_dir()
.join("schedules.json");
let sched = goose::scheduler_factory::SchedulerFactory::create_legacy(sched_storage_path)
.await
.unwrap();
state.set_scheduler(sched).await;
goose_server::routes::config_management::routes(state)
}
#[tokio::test]
async fn test_pricing_endpoint_basic() {
// Basic test to ensure pricing endpoint responds correctly
let app = create_test_app().await;
let request = Request::builder()
.uri("/config/pricing")
.method("POST")
.header("content-type", "application/json")
.header("x-secret-key", "test")
.body(Body::from(json!({"configured_only": true}).to_string()))
.unwrap();
let response = app.oneshot(request).await.unwrap();
assert_eq!(response.status(), StatusCode::OK);
}
crates/goose/Cargo.toml
+1
View File
@@ -15,6 +15,7 @@ tokio = { version = "1.43", features = ["full"] }
reqwest = { version = "0.12.9", features = ["json", "rustls-tls-native-roots"], default-features = false }
[dependencies]
lru = "0.12"
mcp-client = { path = "../mcp-client" }
mcp-core = { path = "../mcp-core" }
rmcp = { workspace = true, features = [
crates/goose/src/agents/context.rs
+1 -1
View File
@@ -33,7 +33,7 @@ impl Agent {
// Only add an assistant message if we have room for it and it won't cause another overflow
let assistant_message = Message::assistant().with_text("I had run into a context length exceeded error so I truncated some of the oldest messages in our conversation.");
let assistant_tokens =
token_counter.count_chat_tokens("", &[assistant_message.clone()], &[]);
token_counter.count_chat_tokens("", std::slice::from_ref(&assistant_message), &[]);
let current_total: usize = new_token_counts.iter().sum();
if current_total + assistant_tokens <= target_context_limit {
crates/goose/src/agents/subagent_execution_tool/executor/mod.rs
+1 -1
View File
@@ -45,7 +45,7 @@ pub async fn execute_single_task(
.await;
let execution_time = start_time.elapsed().as_millis();
let stats = calculate_stats(&[result.clone()], execution_time);
let stats = calculate_stats(std::slice::from_ref(&result), execution_time);
ExecutionResponse {
status: EXECUTION_STATUS_COMPLETED.to_string(),
crates/goose/src/conversation/mod.rs
+1 -1
View File
@@ -86,7 +86,7 @@ impl Conversation {
}
}
pub fn iter(&self) -> std::slice::Iter<Message> {
pub fn iter(&self) -> std::slice::Iter<'_, Message> {
self.0.iter()
}
crates/goose/src/execution/manager.rs
+150
View File
@@ -0,0 +1,150 @@
//! Agent lifecycle management with session isolation
use super::SessionExecutionMode;
use crate::agents::Agent;
use crate::config::APP_STRATEGY;
use crate::model::ModelConfig;
use crate::providers::create;
use crate::scheduler_factory::SchedulerFactory;
use crate::scheduler_trait::SchedulerTrait;
use anyhow::Result;
use etcetera::{choose_app_strategy, AppStrategy};
use lru::LruCache;
use std::num::NonZeroUsize;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
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> {
// 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))
.unwrap_or_else(|| NonZeroUsize::new(100).unwrap());
let manager = Self {
sessions: Arc::new(RwLock::new(LruCache::new(capacity))),
scheduler,
default_provider: Arc::new(RwLock::new(None)),
};
let _ = manager.configure_default_provider().await;
Ok(manager)
}
pub async fn scheduler(&self) -> Result<Arc<dyn SchedulerTrait>> {
Ok(Arc::clone(&self.scheduler))
}
pub async fn set_default_provider(&self, provider: Arc<dyn crate::providers::base::Provider>) {
debug!("Setting default provider on AgentManager");
*self.default_provider.write().await = Some(provider);
}
pub async fn configure_default_provider(&self) -> Result<()> {
let provider_name = std::env::var("GOOSE_DEFAULT_PROVIDER")
.or_else(|_| std::env::var("GOOSE_PROVIDER__TYPE"))
.ok();
let model_name = std::env::var("GOOSE_DEFAULT_MODEL")
.or_else(|_| std::env::var("GOOSE_PROVIDER__MODEL"))
.ok();
if provider_name.is_none() || model_name.is_none() {
return Ok(());
}
if let (Some(provider_name), Some(model_name)) = (provider_name, model_name) {
match ModelConfig::new(&model_name) {
Ok(model_config) => match create(&provider_name, model_config) {
Ok(provider) => {
self.set_default_provider(provider).await;
info!(
"Configured default provider: {} with model: {}",
provider_name, model_name
);
}
Err(e) => {
warn!("Failed to create default provider {}: {}", provider_name, e)
}
},
Err(e) => warn!("Failed to create model config for {}: {}", model_name, e),
}
}
Ok(())
}
pub async fn get_or_create_agent(
&self,
session_id: String,
mode: SessionExecutionMode,
) -> Result<Arc<Agent>> {
let agent = {
let mut sessions = self.sessions.write().await;
if let Some(agent) = sessions.get(&session_id) {
debug!("Found existing agent for session {}", session_id);
return Ok(Arc::clone(agent));
}
info!(
"Creating new agent for session {} with mode {}",
session_id, mode
);
let agent = Arc::new(Agent::new());
sessions.put(session_id.clone(), Arc::clone(&agent));
agent
};
match &mode {
SessionExecutionMode::Interactive | SessionExecutionMode::Background => {
debug!("Setting scheduler on agent for session {}", session_id);
agent.set_scheduler(Arc::clone(&self.scheduler)).await;
}
SessionExecutionMode::SubTask { .. } => {
debug!(
"SubTask mode for session {}, skipping scheduler setup",
session_id
);
}
}
if let Some(provider) = &*self.default_provider.read().await {
debug!(
"Setting default provider on agent for session {}",
session_id
);
let _ = agent.update_provider(Arc::clone(provider)).await;
}
Ok(agent)
}
pub async fn remove_session(&self, session_id: &str) -> Result<()> {
let mut sessions = self.sessions.write().await;
sessions
.pop(session_id)
.ok_or_else(|| anyhow::anyhow!("Session {} not found", session_id))?;
info!("Removed session {}", session_id);
Ok(())
}
pub async fn has_session(&self, session_id: &str) -> bool {
self.sessions.read().await.contains(session_id)
}
pub async fn session_count(&self) -> usize {
self.sessions.read().await.len()
}
}
crates/goose/src/execution/mod.rs
+45
View File
@@ -0,0 +1,45 @@
//! Unified execution management for Goose agents
//!
//! This module provides centralized agent lifecycle management with session isolation,
//! enabling multiple concurrent sessions with independent agents, extensions, and providers.
pub mod manager;
use serde::{Deserialize, Serialize};
use std::fmt;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum SessionExecutionMode {
Interactive,
Background,
SubTask { parent_session: String },
}
impl SessionExecutionMode {
/// Create an interactive chat mode
pub fn chat() -> Self {
Self::Interactive
}
/// Create a background/scheduled mode
pub fn scheduled() -> Self {
Self::Background
}
/// Create a sub-task mode with parent reference
pub fn task(parent: String) -> Self {
Self::SubTask {
parent_session: parent,
}
}
}
impl fmt::Display for SessionExecutionMode {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Interactive => write!(f, "interactive"),
Self::Background => write!(f, "background"),
Self::SubTask { parent_session } => write!(f, "subtask(parent: {})", parent_session),
}
}
}
crates/goose/src/lib.rs
+1
View File
@@ -2,6 +2,7 @@ pub mod agents;
pub mod config;
pub mod context_mgmt;
pub mod conversation;
pub mod execution;
pub mod logging;
pub mod model;
pub mod oauth;
crates/goose/src/permission/permission_judge.rs
+5 -1
View File
@@ -144,7 +144,11 @@ pub async fn detect_read_only_tools(
.unwrap_or_else(|_| "You are a good analyst and can detect operations whether they have read-only operations.".to_string());
let res = provider
.complete(&system_prompt, check_messages.messages(), &[tool.clone()])
.complete(
&system_prompt,
check_messages.messages(),
std::slice::from_ref(&tool),
)
.await;
// Process the response and return an empty vector if the response is invalid
crates/goose/src/recipe/template_recipe.rs
+2 -2
View File
@@ -112,7 +112,7 @@ fn add_template_in_env(
content: &str,
recipe_dir: String,
undefined_behavior: UndefinedBehavior,
) -> Result<Environment> {
) -> Result<Environment<'_>> {
let mut env = minijinja::Environment::new();
env.set_undefined_behavior(undefined_behavior);
env.set_loader(move |name| {
@@ -136,7 +136,7 @@ fn get_env_with_template_variables(
content: &str,
recipe_dir: String,
undefined_behavior: UndefinedBehavior,
) -> Result<(Environment, HashSet<String>)> {
) -> Result<(Environment<'_>, HashSet<String>)> {
let env = add_template_in_env(content, recipe_dir, undefined_behavior)?;
let template = env.get_template(CURRENT_TEMPLATE_NAME).unwrap();
let state = template.eval_to_state(())?;
crates/goose/tests/execution_tests.rs
+323
View File
@@ -0,0 +1,323 @@
mod execution_tests {
use goose::execution::manager::AgentManager;
use goose::execution::SessionExecutionMode;
use serial_test::serial;
use std::sync::Arc;
#[test]
fn test_execution_mode_constructors() {
assert_eq!(
SessionExecutionMode::chat(),
SessionExecutionMode::Interactive
);
assert_eq!(
SessionExecutionMode::scheduled(),
SessionExecutionMode::Background
);
let parent = "parent-123".to_string();
assert_eq!(
SessionExecutionMode::task(parent.clone()),
SessionExecutionMode::SubTask {
parent_session: parent
}
);
}
#[tokio::test]
async fn test_session_isolation() {
let manager = AgentManager::new(None).await.unwrap();
let session1 = uuid::Uuid::new_v4().to_string();
let session2 = uuid::Uuid::new_v4().to_string();
let agent1 = manager
.get_or_create_agent(session1.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
let agent2 = manager
.get_or_create_agent(session2.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
// Different sessions should have different agents
assert!(!Arc::ptr_eq(&agent1, &agent2));
// Getting the same session should return the same agent
let agent1_again = manager
.get_or_create_agent(session1, SessionExecutionMode::chat())
.await
.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| String::from(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);
}
#[tokio::test]
async fn test_remove_session() {
let manager = AgentManager::new(None).await.unwrap();
let session = String::from("remove-test");
manager
.get_or_create_agent(session.clone(), SessionExecutionMode::chat())
.await
.unwrap();
assert!(manager.has_session(&session).await);
manager.remove_session(&session).await.unwrap();
assert!(!manager.has_session(&session).await);
assert!(manager.remove_session(&session).await.is_err());
}
#[tokio::test]
async fn test_concurrent_access() {
let manager = Arc::new(AgentManager::new(None).await.unwrap());
let session = String::from("concurrent-test");
let mut handles = vec![];
for _ in 0..10 {
let mgr = Arc::clone(&manager);
let sess = session.clone();
handles.push(tokio::spawn(async move {
mgr.get_or_create_agent(sess, SessionExecutionMode::chat())
.await
.unwrap()
}));
}
let agents: Vec<_> = futures::future::join_all(handles)
.await
.into_iter()
.map(|r| r.unwrap())
.collect();
for agent in &agents[1..] {
assert!(Arc::ptr_eq(&agents[0], agent));
}
assert_eq!(manager.session_count().await, 1);
}
#[tokio::test]
async fn test_different_modes_same_session() {
let manager = AgentManager::new(None).await.unwrap();
let session_id = String::from("mode-test");
// Create initial agent
let agent1 = manager
.get_or_create_agent(session_id.clone(), SessionExecutionMode::chat())
.await
.unwrap();
// Get same session with different mode - should return same agent
// (mode is stored but agent is reused)
let agent2 = manager
.get_or_create_agent(session_id.clone(), SessionExecutionMode::Background)
.await
.unwrap();
assert!(Arc::ptr_eq(&agent1, &agent2));
}
#[tokio::test]
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());
let session_id = String::from("race-condition-test");
// Spawn multiple tasks trying to create the same NEW session simultaneously
let mut handles = vec![];
for _ in 0..20 {
let sess = session_id.clone();
let mgr_clone = Arc::clone(&manager);
handles.push(tokio::spawn(async move {
mgr_clone
.get_or_create_agent(sess, SessionExecutionMode::Interactive)
.await
.unwrap()
}));
}
// Collect all agents
let agents: Vec<_> = futures::future::join_all(handles)
.await
.into_iter()
.map(|r| r.unwrap())
.collect();
// All should be the same agent (double-check pattern should prevent duplicates)
for agent in &agents[1..] {
assert!(
Arc::ptr_eq(&agents[0], agent),
"All concurrent requests should get the same agent"
);
}
// 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);
}
#[tokio::test]
#[serial]
async fn test_configure_default_provider() {
use std::env;
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 result = manager.configure_default_provider().await;
assert!(result.is_ok());
// Restore original env vars
if let Some(val) = original_provider {
env::set_var("GOOSE_DEFAULT_PROVIDER", val);
} else {
env::remove_var("GOOSE_DEFAULT_PROVIDER");
}
if let Some(val) = original_model {
env::set_var("GOOSE_DEFAULT_MODEL", val);
} else {
env::remove_var("GOOSE_DEFAULT_MODEL");
}
}
#[tokio::test]
async fn test_set_default_provider() {
use goose::providers::testprovider::TestProvider;
use std::sync::Arc;
let manager = AgentManager::new(None).await.unwrap();
// Create a test provider for replaying (doesn't need inner provider)
let temp_file = format!(
"{}/test_provider_{}.json",
std::env::temp_dir().display(),
std::process::id()
);
// Create an empty test provider (will fail on actual use but that's ok for this test)
let test_provider = TestProvider::new_replaying(&temp_file)
.unwrap_or_else(|_| TestProvider::new_replaying("/tmp/dummy.json").unwrap());
manager.set_default_provider(Arc::new(test_provider)).await;
let session = String::from("provider-test");
let _agent = manager
.get_or_create_agent(session.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
assert!(manager.has_session(&session).await);
}
#[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();
let session1 = String::from("session-1");
let session2 = String::from("session-2");
manager
.get_or_create_agent(session1.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
// 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
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
manager
.get_or_create_agent(session1.clone(), SessionExecutionMode::Interactive)
.await
.unwrap();
// 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)
.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);
}
#[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");
let result = manager.remove_session(&session).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("not found"));
}
}
crates/goose/tests/providers.rs
+2 -2
View File
@@ -137,8 +137,8 @@ impl ProviderTester {
.provider
.complete(
"You are a helpful weather assistant.",
&[message.clone()],
&[weather_tool.clone()],
std::slice::from_ref(&message),
std::slice::from_ref(&weather_tool),
)
.await?;