chore: refactor read-write lock on agent (#2225)

Co-authored-by: Alice Hau <ahau@squareup.com>

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

@@ -1,6 +1,9 @@
+use std::sync::Arc;
+
 use crate::configuration;
 use crate::state;
 use anyhow::Result;
+use goose::agents::Agent;
 use tower_http::cors::{Any, CorsLayer};
 use tracing::info;
 
@@ -15,8 +18,10 @@ pub async fn run() -> Result<()> {
     let secret_key =
         std::env::var("GOOSE_SERVER__SECRET_KEY").unwrap_or_else(|_| "test".to_string());
 
-    // Create app state - agent will start as None
-    let state = state::AppState::new(secret_key.clone()).await?;
+    let new_agent = Agent::new();
+
+    // Create app state with agent
+    let state = state::AppState::new(Arc::new(new_agent), secret_key.clone()).await;
 
     // Create router with CORS support
     let cors = CorsLayer::new()

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

@@ -8,14 +8,15 @@ use axum::{
 };
 use goose::config::Config;
 use goose::config::PermissionManager;
-use goose::{agents::Agent, model::ModelConfig, providers};
+use goose::model::ModelConfig;
+use goose::providers::create;
 use goose::{
     agents::{extension::ToolInfo, extension_manager::get_parameter_names},
     config::permission::PermissionLevel,
 };
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
-use std::env;
+use std::sync::Arc;
 
 #[derive(Serialize)]
 struct VersionsResponse {
@@ -33,17 +34,6 @@ struct ExtendPromptResponse {
     success: bool,
 }
 
-#[derive(Deserialize)]
-struct CreateAgentRequest {
-    provider: String,
-    model: Option<String>,
-}
-
-#[derive(Serialize)]
-struct CreateAgentResponse {
-    version: String,
-}
-
 #[derive(Deserialize)]
 struct ProviderFile {
     name: String,
@@ -66,6 +56,12 @@ struct ProviderList {
     details: ProviderDetails,
 }
 
+#[derive(Deserialize)]
+struct UpdateProviderRequest {
+    provider: String,
+    model: Option<String>,
+}
+
 #[derive(Deserialize)]
 pub struct GetToolsQuery {
     extension_name: Option<String>,
@@ -82,53 +78,18 @@ async fn get_versions() -> Json<VersionsResponse> {
 }
 
 async fn extend_prompt(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(payload): Json<ExtendPromptRequest>,
 ) -> Result<Json<ExtendPromptResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    let mut agent = state.agent.write().await;
-    if let Some(ref mut agent) = *agent {
-        agent.extend_system_prompt(payload.extension).await;
-        Ok(Json(ExtendPromptResponse { success: true }))
-    } else {
-        Err(StatusCode::NOT_FOUND)
-    }
-}
-
-#[axum::debug_handler]
-async fn create_agent(
-    State(state): State<AppState>,
-    headers: HeaderMap,
-    Json(payload): Json<CreateAgentRequest>,
-) -> Result<Json<CreateAgentResponse>, StatusCode> {
-    verify_secret_key(&headers, &state)?;
-
-    // Set the environment variable for the model if provided
-    if let Some(model) = &payload.model {
-        let env_var_key = format!("{}_MODEL", payload.provider.to_uppercase());
-        env::set_var(env_var_key.clone(), model);
-        println!("Set environment variable: {}={}", env_var_key, model);
-    }
-
-    let config = Config::global();
-    let model = payload.model.unwrap_or_else(|| {
-        config
-            .get_param("GOOSE_MODEL")
-            .expect("Did not find a model on payload or in env")
-    });
-    let model_config = ModelConfig::new(model);
-    let provider =
-        providers::create(&payload.provider, model_config).expect("Failed to create provider");
-
-    let version = String::from("goose");
-    let new_agent = Agent::new(provider);
-
-    let mut agent = state.agent.write().await;
-    *agent = Some(new_agent);
-
-    Ok(Json(CreateAgentResponse { version }))
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    agent.extend_system_prompt(payload.extension.clone()).await;
+    Ok(Json(ExtendPromptResponse { success: true }))
 }
 
 async fn list_providers() -> Json<Vec<ProviderList>> {
@@ -168,7 +129,7 @@ async fn list_providers() -> Json<Vec<ProviderList>> {
     )
 )]
 async fn get_tools(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Query(query): Query<GetToolsQuery>,
 ) -> Result<Json<Vec<ToolInfo>>, StatusCode> {
@@ -176,8 +137,10 @@ async fn get_tools(
 
     let config = Config::global();
     let goose_mode = config.get_param("GOOSE_MODE").unwrap_or("auto".to_string());
-    let agent = state.agent.read().await;
-    let agent = agent.as_ref().ok_or(StatusCode::PRECONDITION_REQUIRED)?;
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
     let permission_manager = PermissionManager::default();
 
     let mut tools: Vec<ToolInfo> = agent
@@ -210,12 +173,56 @@ async fn get_tools(
     Ok(Json(tools))
 }
 
-pub fn routes(state: AppState) -> Router {
+#[utoipa::path(
+    post,
+    path = "/agent/update_provider",
+    responses(
+        (status = 200, description = "Update provider completed", body = String),
+        (status = 500, description = "Internal server error")
+    )
+)]
+async fn update_agent_provider(
+    State(state): State<Arc<AppState>>,
+    headers: HeaderMap,
+    Json(payload): Json<UpdateProviderRequest>,
+) -> Result<StatusCode, StatusCode> {
+    // Verify secret key
+    let secret_key = headers
+        .get("X-Secret-Key")
+        .and_then(|value| value.to_str().ok())
+        .ok_or(StatusCode::UNAUTHORIZED)?;
+
+    if secret_key != state.secret_key {
+        return Err(StatusCode::UNAUTHORIZED);
+    }
+
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+
+    let config = Config::global();
+    let model = payload.model.unwrap_or_else(|| {
+        config
+            .get_param("GOOSE_MODEL")
+            .expect("Did not find a model on payload or in env to update provider with")
+    });
+    let model_config = ModelConfig::new(model);
+    let new_provider = create(&payload.provider, model_config).unwrap();
+    agent
+        .update_provider(new_provider)
+        .await
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    Ok(StatusCode::OK)
+}
+
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/agent/versions", get(get_versions))
         .route("/agent/providers", get(list_providers))
         .route("/agent/prompt", post(extend_prompt))
         .route("/agent/tools", get(get_tools))
-        .route("/agent", post(create_agent))
+        .route("/agent/update_provider", post(update_agent_provider))
         .with_state(state)
 }

crates/goose-server/src/routes/config_management.rs

@@ -18,7 +18,7 @@ use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use serde_yaml;
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 use utoipa::ToSchema;
 
 #[derive(Serialize, ToSchema)]
@@ -89,7 +89,7 @@ pub struct UpsertPermissionsQuery {
     )
 )]
 pub async fn upsert_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(query): Json<UpsertConfigQuery>,
 ) -> Result<Json<Value>, StatusCode> {
@@ -116,7 +116,7 @@ pub async fn upsert_config(
     )
 )]
 pub async fn remove_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(query): Json<ConfigKeyQuery>,
 ) -> Result<Json<String>, StatusCode> {
@@ -148,7 +148,7 @@ pub async fn remove_config(
     )
 )]
 pub async fn read_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(query): Json<ConfigKeyQuery>,
 ) -> Result<Json<Value>, StatusCode> {
@@ -180,7 +180,7 @@ pub async fn read_config(
     )
 )]
 pub async fn get_extensions(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
 ) -> Result<Json<ExtensionResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
@@ -213,7 +213,7 @@ pub async fn get_extensions(
     )
 )]
 pub async fn add_extension(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(extension_query): Json<ExtensionQuery>,
 ) -> Result<Json<String>, StatusCode> {
@@ -251,7 +251,7 @@ pub async fn add_extension(
     )
 )]
 pub async fn remove_extension(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     axum::extract::Path(name): axum::extract::Path<String>,
 ) -> Result<Json<String>, StatusCode> {
@@ -272,7 +272,7 @@ pub async fn remove_extension(
     )
 )]
 pub async fn read_all_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
 ) -> Result<Json<ConfigResponse>, StatusCode> {
     // Use the helper function to verify the secret key
@@ -297,7 +297,7 @@ pub async fn read_all_config(
     )
 )]
 pub async fn providers(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
 ) -> Result<Json<Vec<ProviderDetails>>, StatusCode> {
     verify_secret_key(&headers, &state)?;
@@ -332,7 +332,7 @@ pub async fn providers(
     )
 )]
 pub async fn init_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
 ) -> Result<Json<String>, StatusCode> {
     verify_secret_key(&headers, &state)?;
@@ -402,7 +402,7 @@ pub async fn init_config(
     )
 )]
 pub async fn upsert_permissions(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(query): Json<UpsertPermissionsQuery>,
 ) -> Result<Json<String>, StatusCode> {
@@ -435,7 +435,7 @@ pub static APP_STRATEGY: Lazy<AppStrategyArgs> = Lazy::new(|| AppStrategyArgs {
     )
 )]
 pub async fn backup_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
 ) -> Result<Json<String>, StatusCode> {
     verify_secret_key(&headers, &state)?;
@@ -466,7 +466,7 @@ pub async fn backup_config(
     }
 }
 
-pub fn routes(state: AppState) -> Router {
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/config", get(read_all_config))
         .route("/config/upsert", post(upsert_config))

crates/goose-server/src/routes/configs.rs

@@ -10,7 +10,7 @@ use http::{HeaderMap, StatusCode};
 use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 #[derive(Serialize)]
 struct ConfigResponse {
@@ -26,7 +26,7 @@ struct ConfigRequest {
 }
 
 async fn store_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(request): Json<ConfigRequest>,
 ) -> Result<Json<ConfigResponse>, StatusCode> {
@@ -148,7 +148,7 @@ pub struct GetConfigResponse {
 }
 
 pub async fn get_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Query(query): Query<GetConfigQuery>,
 ) -> Result<Json<GetConfigResponse>, StatusCode> {
@@ -174,7 +174,7 @@ struct DeleteConfigRequest {
 }
 
 async fn delete_config(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(request): Json<DeleteConfigRequest>,
 ) -> Result<StatusCode, StatusCode> {
@@ -193,7 +193,7 @@ async fn delete_config(
     }
 }
 
-pub fn routes(state: AppState) -> Router {
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/configs/providers", post(check_provider_configs))
         .route("/configs/get", get(get_config))

crates/goose-server/src/routes/context.rs

@@ -8,6 +8,7 @@ use axum::{
 };
 use goose::message::Message;
 use serde::{Deserialize, Serialize};
+use std::sync::Arc;
 
 // Direct message serialization for context mgmt request
 #[derive(Debug, Deserialize)]
@@ -26,15 +27,16 @@ pub struct ContextManageResponse {
 }
 
 async fn manage_context(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(request): Json<ContextManageRequest>,
 ) -> Result<Json<ContextManageResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    // Get a lock on the shared agent
-    let agent = state.agent.read().await;
-    let agent = agent.as_ref().ok_or(StatusCode::PRECONDITION_REQUIRED)?;
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
 
     let mut processed_messages: Vec<Message> = vec![];
     let mut token_counts: Vec<usize> = vec![];
@@ -57,7 +59,7 @@ async fn manage_context(
 }
 
 // Configure routes for this module
-pub fn routes(state: AppState) -> Router {
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/context/manage", post(manage_context))
         .with_state(state)

crates/goose-server/src/routes/extension.rs

@@ -1,5 +1,6 @@
 use std::env;
 use std::path::Path;
+use std::sync::Arc;
 use std::sync::OnceLock;
 
 use super::utils::verify_secret_key;
@@ -79,7 +80,7 @@ struct ExtensionResponse {
 
 /// Handler for adding a new extension configuration.
 async fn add_extension(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     raw: axum::extract::Json<serde_json::Value>,
 ) -> Result<Json<ExtensionResponse>, StatusCode> {
@@ -228,9 +229,11 @@ async fn add_extension(
         },
     };
 
-    // Acquire a lock on the agent and attempt to add the extension.
-    let mut agent = state.agent.write().await;
-    let agent = agent.as_mut().ok_or(StatusCode::PRECONDITION_REQUIRED)?;
+    // Get a reference to the agent
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
     let response = agent.add_extension(extension_config).await;
 
     // Respond with the result.
@@ -254,15 +257,17 @@ async fn add_extension(
 
 /// Handler for removing an extension by name
 async fn remove_extension(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(name): Json<String>,
 ) -> Result<Json<ExtensionResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    // Acquire a lock on the agent and attempt to remove the extension
-    let mut agent = state.agent.write().await;
-    let agent = agent.as_mut().ok_or(StatusCode::PRECONDITION_REQUIRED)?;
+    // Get a reference to the agent
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
     agent.remove_extension(&name).await;
 
     Ok(Json(ExtensionResponse {
@@ -272,7 +277,7 @@ async fn remove_extension(
 }
 
 /// Registers the extension management routes with the Axum router.
-pub fn routes(state: AppState) -> Router {
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/extensions/add", post(add_extension))
         .route("/extensions/remove", post(remove_extension))

crates/goose-server/src/routes/mod.rs

@@ -9,10 +9,12 @@ pub mod recipe;
 pub mod reply;
 pub mod session;
 pub mod utils;
+use std::sync::Arc;
+
 use axum::Router;
 
 // Function to configure all routes
-pub fn configure(state: crate::state::AppState) -> Router {
+pub fn configure(state: Arc<crate::state::AppState>) -> Router {
     Router::new()
         .merge(health::routes())
         .merge(reply::routes(state.clone()))
@@ -22,5 +24,5 @@ pub fn configure(state: crate::state::AppState) -> Router {
         .merge(configs::routes(state.clone()))
         .merge(config_management::routes(state.clone()))
         .merge(recipe::routes(state.clone()))
-        .merge(session::routes(state))
+        .merge(session::routes(state.clone()))
 }

crates/goose-server/src/routes/recipe.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use axum::{extract::State, http::StatusCode, routing::post, Json, Router};
 use goose::message::Message;
 use goose::recipe::Recipe;
@@ -34,17 +36,17 @@ pub struct CreateRecipeResponse {
 
 /// Create a Recipe configuration from the current state of an agent
 async fn create_recipe(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     Json(request): Json<CreateRecipeRequest>,
 ) -> Result<Json<CreateRecipeResponse>, (StatusCode, Json<CreateRecipeResponse>)> {
-    let agent = state.agent.read().await;
-    let agent = agent.as_ref().ok_or_else(|| {
-        let error_response = CreateRecipeResponse {
-            recipe: None,
-            error: Some("Agent not initialized".to_string()),
-        };
-        (StatusCode::PRECONDITION_REQUIRED, Json(error_response))
-    })?;
+    let error_response = CreateRecipeResponse {
+        recipe: None,
+        error: Some("Missing agent".to_string()),
+    };
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| (StatusCode::PRECONDITION_FAILED, Json(error_response)))?;
 
     // Create base recipe from agent state and messages
     let recipe_result = agent.create_recipe(request.messages).await;
@@ -82,7 +84,7 @@ async fn create_recipe(
     }
 }
 
-pub fn routes(state: AppState) -> Router {
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/recipe/create", post(create_recipe))
         .with_state(state)

crates/goose-server/src/routes/reply.rs

@@ -26,6 +26,7 @@ use std::{
     convert::Infallible,
     path::PathBuf,
     pin::Pin,
+    sync::Arc,
     task::{Context, Poll},
     time::Duration,
 };
@@ -101,7 +102,7 @@ async fn stream_event(
 }
 
 async fn handler(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(request): Json<ChatRequest>,
 ) -> Result<SseResponse, StatusCode> {
@@ -119,15 +120,34 @@ async fn handler(
         .session_id
         .unwrap_or_else(session::generate_session_id);
 
-    // Get a lock on the shared agent
-    let agent = state.agent.clone();
-
     // Spawn task to handle streaming
     tokio::spawn(async move {
-        let agent = agent.read().await;
-        let agent = match agent.as_ref() {
-            Some(agent) => agent,
-            None => {
+        let agent = state.get_agent().await;
+        let agent = match agent {
+            Ok(agent) => {
+                let provider = agent.provider().await;
+                match provider {
+                    Ok(_) => agent,
+                    Err(_) => {
+                        let _ = stream_event(
+                            MessageEvent::Error {
+                                error: "No provider configured".to_string(),
+                            },
+                            &tx,
+                        )
+                        .await;
+                        let _ = stream_event(
+                            MessageEvent::Finish {
+                                reason: "error".to_string(),
+                            },
+                            &tx,
+                        )
+                        .await;
+                        return;
+                    }
+                }
+            }
+            Err(_) => {
                 let _ = stream_event(
                     MessageEvent::Error {
                         error: "No agent configured".to_string(),
@@ -147,7 +167,7 @@ async fn handler(
         };
 
         // Get the provider first, before starting the reply stream
-        let provider = agent.provider();
+        let provider = agent.provider().await;
 
         let mut stream = match agent
             .reply(
@@ -204,7 +224,7 @@ async fn handler(
                             // Store messages and generate description in background
                             let session_path = session_path.clone();
                             let messages = all_messages.clone();
-                            let provider = provider.clone();
+                            let provider = Arc::clone(provider.as_ref().unwrap());
                             tokio::spawn(async move {
                                 if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
                                     tracing::error!("Failed to store session history: {:?}", e);
@@ -262,7 +282,7 @@ struct AskResponse {
 
 // Simple ask an AI for a response, non streaming
 async fn ask_handler(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(request): Json<AskRequest>,
 ) -> Result<Json<AskResponse>, StatusCode> {
@@ -275,12 +295,13 @@ async fn ask_handler(
         .session_id
         .unwrap_or_else(session::generate_session_id);
 
-    let agent = state.agent.clone();
-    let agent = agent.write().await;
-    let agent = agent.as_ref().ok_or(StatusCode::NOT_FOUND)?;
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
 
     // Get the provider first, before starting the reply stream
-    let provider = agent.provider();
+    let provider = agent.provider().await;
 
     // Create a single message for the prompt
     let messages = vec![Message::user().with_text(request.prompt)];
@@ -339,7 +360,7 @@ async fn ask_handler(
     // Store messages and generate description in background
     let session_path = session_path.clone();
     let messages = all_messages.clone();
-    let provider = provider.clone();
+    let provider = Arc::clone(provider.as_ref().unwrap());
     tokio::spawn(async move {
         if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
             tracing::error!("Failed to store session history: {:?}", e);
@@ -374,15 +395,16 @@ fn default_principal_type() -> PrincipalType {
     )
 )]
 pub async fn confirm_permission(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(request): Json<PermissionConfirmationRequest>,
 ) -> Result<Json<Value>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    let agent = state.agent.clone();
-    let agent = agent.read().await;
-    let agent = agent.as_ref().ok_or(StatusCode::NOT_FOUND)?;
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
 
     let permission = match request.action.as_str() {
         "always_allow" => Permission::AlwaysAllow,
@@ -410,7 +432,7 @@ struct ToolResultRequest {
 }
 
 async fn submit_tool_result(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     raw: axum::extract::Json<serde_json::Value>,
 ) -> Result<Json<Value>, StatusCode> {
@@ -435,14 +457,16 @@ async fn submit_tool_result(
         }
     };
 
-    let agent = state.agent.read().await;
-    let agent = agent.as_ref().ok_or(StatusCode::NOT_FOUND)?;
+    let agent = state
+        .get_agent()
+        .await
+        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
     agent.handle_tool_result(payload.id, payload.result).await;
     Ok(Json(json!({"status": "ok"})))
 }
 
 // Configure routes for this module
-pub fn routes(state: AppState) -> Router {
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/reply", post(handler))
         .route("/ask", post(ask_handler))
@@ -496,9 +520,7 @@ mod tests {
     mod integration_tests {
         use super::*;
         use axum::{body::Body, http::Request};
-        use std::collections::HashMap;
         use std::sync::Arc;
-        use tokio::sync::{Mutex, RwLock};
         use tower::ServiceExt;
 
         // This test requires tokio runtime
@@ -509,12 +531,9 @@ mod tests {
             let mock_provider = Arc::new(MockProvider {
                 model_config: mock_model_config,
             });
-            let agent = Agent::new(mock_provider);
-            let state = AppState {
-                config: Arc::new(Mutex::new(HashMap::new())),
-                agent: Arc::new(RwLock::new(Some(agent))),
-                secret_key: "test-secret".to_string(),
-            };
+            let agent = Agent::new();
+            let _ = agent.update_provider(mock_provider).await;
+            let state = AppState::new(Arc::new(agent), "test-secret".to_string()).await;
 
             // Build router
             let app = routes(state);

crates/goose-server/src/routes/session.rs

@@ -1,4 +1,6 @@
 use super::utils::verify_secret_key;
+use std::sync::Arc;
+
 use crate::state::AppState;
 use axum::{
     extract::{Path, State},
@@ -25,7 +27,7 @@ struct SessionHistoryResponse {
 
 // List all available sessions
 async fn list_sessions(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
 ) -> Result<Json<SessionListResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
@@ -38,7 +40,7 @@ async fn list_sessions(
 
 // Get a specific session's history
 async fn get_session_history(
-    State(state): State<AppState>,
+    State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Path(session_id): Path<String>,
 ) -> Result<Json<SessionHistoryResponse>, StatusCode> {
@@ -65,7 +67,7 @@ async fn get_session_history(
 }
 
 // Configure routes for this module
-pub fn routes(state: AppState) -> Router {
+pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/sessions", get(list_sessions))
         .route("/sessions/:session_id", get(get_session_history))

crates/goose-server/src/state.rs

@@ -1,25 +1,34 @@
-use anyhow::Result;
 use goose::agents::Agent;
-use serde_json::Value;
-use std::collections::HashMap;
 use std::sync::Arc;
-use tokio::sync::{Mutex, RwLock};
 
+/// Shared reference to an Agent that can be cloned cheaply
+/// without cloning the underlying Agent object
+pub type AgentRef = Arc<Agent>;
+
+/// Thread-safe container for an optional Agent reference
+/// Outer Arc: Allows multiple route handlers to access the same Mutex
+/// - Mutex provides exclusive access for updates
+/// - Option allows for the case where no agent exists yet
+///
 /// Shared application state
-#[allow(dead_code)]
 #[derive(Clone)]
 pub struct AppState {
-    pub agent: Arc<RwLock<Option<Agent>>>,
+    // agent: SharedAgentStore,
+    agent: Option<AgentRef>,
     pub secret_key: String,
-    pub config: Arc<Mutex<HashMap<String, Value>>>,
 }
 
 impl AppState {
-    pub async fn new(secret_key: String) -> Result<Self> {
-        Ok(Self {
-            agent: Arc::new(RwLock::new(None)),
+    pub async fn new(agent: AgentRef, secret_key: String) -> Arc<AppState> {
+        Arc::new(Self {
+            agent: Some(agent.clone()),
             secret_key,
-            config: Arc::new(Mutex::new(HashMap::new())),
         })
     }
+
+    pub async fn get_agent(&self) -> Result<Arc<Agent>, anyhow::Error> {
+        self.agent
+            .clone()
+            .ok_or_else(|| anyhow::anyhow!("Agent needs to be created first."))
+    }
 }