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

crates/goose-server/src/commands/agent.rs

@@ -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

@@ -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

@@ -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

@@ -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![];