Add session to agents (#4216)

Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Jack Amadeo <jackamadeo@squareup.com>
Co-authored-by: Jack Amadeo <jackamadeo@block.xyz>

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

@@ -33,7 +33,7 @@ pub async fn run() -> Result<()> {
     let new_agent = Agent::new();
     let agent_ref = Arc::new(new_agent);
 
-    let app_state = state::AppState::new(agent_ref.clone(), secret_key.clone()).await;
+    let app_state = state::AppState::new(agent_ref.clone(), secret_key.clone());
 
     let schedule_file_path = choose_app_strategy(APP_STRATEGY.clone())?
         .data_dir()

crates/goose-server/src/openapi.rs

@@ -370,6 +370,8 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
         super::routes::config_management::upsert_permissions,
         super::routes::config_management::create_custom_provider,
         super::routes::config_management::remove_custom_provider,
+        super::routes::agent::start_agent,
+        super::routes::agent::resume_agent,
         super::routes::agent::get_tools,
         super::routes::agent::add_sub_recipes,
         super::routes::agent::extend_prompt,
@@ -486,6 +488,10 @@ impl<'__s> ToSchema<'__s> for AnnotatedSchema {
         super::routes::agent::UpdateProviderRequest,
         super::routes::agent::SessionConfigRequest,
         super::routes::agent::GetToolsQuery,
+        super::routes::agent::UpdateRouterToolSelectorRequest,
+        super::routes::agent::StartAgentRequest,
+        super::routes::agent::ResumeAgentRequest,
+        super::routes::agent::StartAgentResponse,
         super::routes::agent::ErrorResponse,
     ))
 )]

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

@@ -1,5 +1,6 @@
 use super::utils::verify_secret_key;
 use crate::state::AppState;
+use axum::response::IntoResponse;
 use axum::{
     extract::{Query, State},
     http::{HeaderMap, StatusCode},
@@ -7,20 +8,29 @@ use axum::{
     Json, Router,
 };
 use goose::config::PermissionManager;
+use goose::conversation::message::Message;
+use goose::conversation::Conversation;
 use goose::model::ModelConfig;
 use goose::providers::create;
-use goose::recipe::Response;
+use goose::recipe::{Recipe, Response};
+use goose::session;
+use goose::session::SessionMetadata;
 use goose::{
     agents::{extension::ToolInfo, extension_manager::get_parameter_names},
     config::permission::PermissionLevel,
 };
 use goose::{config::Config, recipe::SubRecipe};
 use serde::{Deserialize, Serialize};
+use std::path::PathBuf;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
+use tracing::error;
 
 #[derive(Deserialize, utoipa::ToSchema)]
 pub struct ExtendPromptRequest {
     extension: String,
+    #[allow(dead_code)]
+    session_id: String,
 }
 
 #[derive(Serialize, utoipa::ToSchema)]
@@ -31,6 +41,8 @@ pub struct ExtendPromptResponse {
 #[derive(Deserialize, utoipa::ToSchema)]
 pub struct AddSubRecipesRequest {
     sub_recipes: Vec<SubRecipe>,
+    #[allow(dead_code)]
+    session_id: String,
 }
 
 #[derive(Serialize, utoipa::ToSchema)]
@@ -42,16 +54,47 @@ 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,
+}
+
+#[derive(Deserialize, utoipa::ToSchema)]
+pub struct StartAgentRequest {
+    working_dir: String,
+    recipe: Option<Recipe>,
+}
+
+#[derive(Deserialize, utoipa::ToSchema)]
+pub struct ResumeAgentRequest {
+    session_id: String,
+}
+
+// This is the same as SessionHistoryResponse
+#[derive(Serialize, utoipa::ToSchema)]
+pub struct StartAgentResponse {
+    session_id: String,
+    metadata: SessionMetadata,
+    messages: Vec<Message>,
 }
 
 #[derive(Serialize, utoipa::ToSchema)]
@@ -59,6 +102,101 @@ pub struct ErrorResponse {
     error: String,
 }
 
+#[utoipa::path(
+    post,
+    path = "/agent/start",
+    request_body = StartAgentRequest,
+    responses(
+        (status = 200, description = "Agent started successfully", body = StartAgentResponse),
+        (status = 400, description = "Bad request - invalid working directory"),
+        (status = 401, description = "Unauthorized - invalid secret key"),
+        (status = 500, description = "Internal server error")
+    )
+)]
+async fn start_agent(
+    State(state): State<Arc<AppState>>,
+    headers: HeaderMap,
+    Json(payload): Json<StartAgentRequest>,
+) -> Result<Json<StartAgentResponse>, StatusCode> {
+    verify_secret_key(&headers, &state)?;
+
+    state.reset().await;
+
+    let session_id = session::generate_session_id();
+    let counter = state.session_counter.fetch_add(1, Ordering::SeqCst) + 1;
+
+    let metadata = SessionMetadata {
+        working_dir: PathBuf::from(&payload.working_dir),
+        description: format!("New session {}", counter),
+        schedule_id: None,
+        message_count: 0,
+        total_tokens: Some(0),
+        input_tokens: Some(0),
+        output_tokens: Some(0),
+        accumulated_total_tokens: Some(0),
+        accumulated_input_tokens: Some(0),
+        accumulated_output_tokens: Some(0),
+        extension_data: Default::default(),
+        recipe: payload.recipe,
+    };
+
+    let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
+        Ok(path) => path,
+        Err(_) => return Err(StatusCode::BAD_REQUEST),
+    };
+
+    let conversation = Conversation::empty();
+    session::storage::save_messages_with_metadata(&session_path, &metadata, &conversation)
+        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+
+    Ok(Json(StartAgentResponse {
+        session_id,
+        metadata,
+        messages: conversation.messages().clone(),
+    }))
+}
+
+#[utoipa::path(
+    post,
+    path = "/agent/resume",
+    request_body = ResumeAgentRequest,
+    responses(
+        (status = 200, description = "Agent started successfully", body = StartAgentResponse),
+        (status = 400, description = "Bad request - invalid working directory"),
+        (status = 401, description = "Unauthorized - invalid secret key"),
+        (status = 500, description = "Internal server error")
+    )
+)]
+async fn resume_agent(
+    State(state): State<Arc<AppState>>,
+    headers: HeaderMap,
+    Json(payload): Json<ResumeAgentRequest>,
+) -> Result<Json<StartAgentResponse>, StatusCode> {
+    verify_secret_key(&headers, &state)?;
+
+    let session_path =
+        match session::get_path(session::Identifier::Name(payload.session_id.clone())) {
+            Ok(path) => path,
+            Err(_) => return Err(StatusCode::BAD_REQUEST),
+        };
+
+    let metadata = session::read_metadata(&session_path).map_err(|_| StatusCode::NOT_FOUND)?;
+
+    let conversation = match session::read_messages(&session_path) {
+        Ok(messages) => messages,
+        Err(e) => {
+            error!("Failed to read session messages: {:?}", e);
+            return Err(StatusCode::NOT_FOUND);
+        }
+    };
+
+    Ok(Json(StartAgentResponse {
+        session_id: payload.session_id.clone(),
+        metadata,
+        messages: conversation.messages().clone(),
+    }))
+}
+
 #[utoipa::path(
     post,
     path = "/agent/add_sub_recipes",
@@ -76,10 +214,7 @@ async fn add_sub_recipes(
 ) -> Result<Json<AddSubRecipesResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    let agent = state.get_agent().await;
     agent.add_sub_recipes(payload.sub_recipes.clone()).await;
     Ok(Json(AddSubRecipesResponse { success: true }))
 }
@@ -101,10 +236,7 @@ async fn extend_prompt(
 ) -> Result<Json<ExtendPromptResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    let agent = state.get_agent().await;
     agent.extend_system_prompt(payload.extension.clone()).await;
     Ok(Json(ExtendPromptResponse { success: true }))
 }
@@ -113,7 +245,8 @@ async fn extend_prompt(
     get,
     path = "/agent/tools",
     params(
-        ("extension_name" = Option<String>, Query, description = "Optional extension name to filter tools")
+        ("extension_name" = Option<String>, Query, description = "Optional extension name to filter tools"),
+        ("session_id" = String, Query, description = "Required session ID to scope tools to a specific session")
     ),
     responses(
         (status = 200, description = "Tools retrieved successfully", body = Vec<ToolInfo>),
@@ -131,10 +264,7 @@ async fn get_tools(
 
     let config = Config::global();
     let goose_mode = config.get_param("GOOSE_MODE").unwrap_or("auto".to_string());
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    let agent = state.get_agent().await;
     let permission_manager = PermissionManager::default();
 
     let mut tools: Vec<ToolInfo> = agent
@@ -186,31 +316,37 @@ async fn update_agent_provider(
     State(state): State<Arc<AppState>>,
     headers: HeaderMap,
     Json(payload): Json<UpdateProviderRequest>,
-) -> Result<StatusCode, StatusCode> {
-    verify_secret_key(&headers, &state)?;
-
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_e| StatusCode::PRECONDITION_FAILED)?;
+) -> Result<StatusCode, impl IntoResponse> {
+    verify_secret_key(&headers, &state).map_err(|e| (e, String::new()))?;
 
+    let agent = state.get_agent().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),
+        None => return Err((StatusCode::BAD_REQUEST, "No model specified".to_string())),
     };
 
-    let model_config = ModelConfig::new(&model).map_err(|_| StatusCode::BAD_REQUEST)?;
+    let model_config = ModelConfig::new(&model).map_err(|e| {
+        (
+            StatusCode::BAD_REQUEST,
+            format!("Invalid model config: {}", e),
+        )
+    })?;
+
+    let new_provider = create(&payload.provider, model_config).map_err(|e| {
+        (
+            StatusCode::BAD_REQUEST,
+            format!("Failed to create provider: {}", e),
+        )
+    })?;
 
-    let new_provider =
-        create(&payload.provider, model_config).map_err(|_| StatusCode::BAD_REQUEST)?;
     agent
         .update_provider(new_provider)
         .await
-        .map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
+        .map_err(|_e| (StatusCode::INTERNAL_SERVER_ERROR, String::new()))?;
 
     Ok(StatusCode::OK)
 }
@@ -218,6 +354,7 @@ async fn update_agent_provider(
 #[utoipa::path(
     post,
     path = "/agent/update_router_tool_selector",
+    request_body = UpdateRouterToolSelectorRequest,
     responses(
         (status = 200, description = "Tool selection strategy updated successfully", body = String),
         (status = 401, description = "Unauthorized - invalid secret key"),
@@ -228,6 +365,7 @@ async fn update_agent_provider(
 async fn update_router_tool_selector(
     State(state): State<Arc<AppState>>,
     headers: HeaderMap,
+    Json(_payload): Json<UpdateRouterToolSelectorRequest>,
 ) -> Result<Json<String>, Json<ErrorResponse>> {
     verify_secret_key(&headers, &state).map_err(|_| {
         Json(ErrorResponse {
@@ -235,13 +373,7 @@ async fn update_router_tool_selector(
         })
     })?;
 
-    let agent = state.get_agent().await.map_err(|e| {
-        tracing::error!("Failed to get agent: {}", e);
-        Json(ErrorResponse {
-            error: format!("Failed to get agent: {}", e),
-        })
-    })?;
-
+    let agent = state.get_agent().await;
     agent
         .update_router_tool_selector(None, Some(true))
         .await
@@ -279,13 +411,7 @@ async fn update_session_config(
         })
     })?;
 
-    let agent = state.get_agent().await.map_err(|e| {
-        tracing::error!("Failed to get agent: {}", e);
-        Json(ErrorResponse {
-            error: format!("Failed to get agent: {}", e),
-        })
-    })?;
-
+    let agent = state.get_agent().await;
     if let Some(response) = payload.response {
         agent.add_final_output_tool(response).await;
 
@@ -300,6 +426,8 @@ async fn update_session_config(
 
 pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
+        .route("/agent/start", post(start_agent))
+        .route("/agent/resume", post(resume_agent))
         .route("/agent/prompt", post(extend_prompt))
         .route("/agent/tools", get(get_tools))
         .route("/agent/update_provider", post(update_agent_provider))

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

@@ -413,8 +413,7 @@ mod tests {
         let state = AppState::new(
             Arc::new(goose::agents::Agent::new()),
             "test-secret".to_string(),
-        )
-        .await;
+        );
         let app = routes(state);
 
         // Test without auth header
@@ -440,8 +439,7 @@ mod tests {
         let state = AppState::new(
             Arc::new(goose::agents::Agent::new()),
             "test-secret".to_string(),
-        )
-        .await;
+        );
         let app = routes(state);
 
         // Create a large base64 string (simulating > 25MB audio)
@@ -470,8 +468,7 @@ mod tests {
         let state = AppState::new(
             Arc::new(goose::agents::Agent::new()),
             "test-secret".to_string(),
-        )
-        .await;
+        );
         let app = routes(state);
 
         let request = Request::builder()
@@ -500,8 +497,7 @@ mod tests {
         let state = AppState::new(
             Arc::new(goose::agents::Agent::new()),
             "test-secret".to_string(),
-        )
-        .await;
+        );
         let app = routes(state);
 
         let request = Request::builder()

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

@@ -858,8 +858,7 @@ mod tests {
         let test_state = AppState::new(
             Arc::new(goose::agents::Agent::default()),
             "test".to_string(),
-        )
-        .await;
+        );
         let sched_storage_path = choose_app_strategy(APP_STRATEGY.clone())
             .unwrap()
             .data_dir()

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

@@ -53,10 +53,7 @@ async fn manage_context(
 ) -> Result<Json<ContextManageResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    let agent = state.get_agent().await;
 
     let mut processed_messages = Conversation::new_unvalidated(vec![]);
     let mut token_counts: Vec<usize> = vec![];

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

@@ -271,11 +271,7 @@ async fn add_extension(
         },
     };
 
-    // Get a reference to the agent
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    let agent = state.get_agent().await;
     let response = agent.add_extension(extension_config).await;
 
     // Respond with the result.
@@ -305,11 +301,7 @@ async fn remove_extension(
 ) -> Result<Json<ExtensionResponse>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    // Get a reference to the agent
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    let agent = state.get_agent().await;
     match agent.remove_extension(&name).await {
         Ok(_) => Ok(Json(ExtensionResponse {
             error: false,

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

@@ -116,16 +116,7 @@ async fn create_recipe(
         request.messages.len()
     );
 
-    let error_response = CreateRecipeResponse {
-        recipe: None,
-        error: Some("Missing agent".to_string()),
-    };
-    let agent = state.get_agent().await.map_err(|e| {
-        tracing::error!("Failed to get agent for recipe creation: {}", e);
-        (StatusCode::PRECONDITION_FAILED, Json(error_response))
-    })?;
-
-    tracing::debug!("Agent retrieved successfully, creating recipe from conversation");
+    let agent = state.get_agent().await;
 
     // Create base recipe from agent state and messages
     let recipe_result = agent
@@ -134,16 +125,12 @@ async fn create_recipe(
 
     match recipe_result {
         Ok(mut recipe) => {
-            tracing::info!("Recipe created successfully with title: '{}'", recipe.title);
-
-            // Update with user-provided metadata
             recipe.title = request.title;
             recipe.description = request.description;
             if request.activities.is_some() {
                 recipe.activities = request.activities
             };
 
-            // Add author if provided
             if let Some(author_req) = request.author {
                 recipe.author = Some(goose::recipe::Author {
                     contact: author_req.contact,
@@ -151,19 +138,13 @@ async fn create_recipe(
                 });
             }
 
-            tracing::debug!("Recipe metadata updated, returning success response");
-
             Ok(Json(CreateRecipeResponse {
                 recipe: Some(recipe),
                 error: None,
             }))
         }
         Err(e) => {
-            // Log the detailed error for debugging
-            tracing::error!("Recipe creation failed: {}", e);
             tracing::error!("Error details: {:?}", e);
-
-            // Return 400 Bad Request with error message
             let error_message = format!("Recipe creation failed: {}", e);
             let error_response = CreateRecipeResponse {
                 recipe: None,

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

@@ -26,7 +26,6 @@ use serde_json::json;
 use serde_json::Value;
 use std::{
     convert::Infallible,
-    path::PathBuf,
     pin::Pin,
     sync::Arc,
     task::{Context, Poll},
@@ -89,8 +88,6 @@ fn track_tool_telemetry(content: &MessageContent, all_messages: &[Message]) {
 struct ChatRequest {
     messages: Vec<Message>,
     session_id: Option<String>,
-    session_working_dir: String,
-    scheduled_job_id: Option<String>,
     recipe_name: Option<String>,
     recipe_version: Option<String>,
 }
@@ -203,22 +200,42 @@ async fn reply_handler(
     let cancel_token = CancellationToken::new();
 
     let messages = Conversation::new_unvalidated(request.messages);
-    let session_working_dir = request.session_working_dir.clone();
 
-    let session_id = request
-        .session_id
-        .unwrap_or_else(session::generate_session_id);
+    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();
 
-    std::mem::drop(tokio::spawn(async move {
-        let agent = match state.get_agent().await {
-            Ok(agent) => agent,
-            Err(_) => {
+    drop(tokio::spawn(async move {
+        let agent = state.get_agent().await;
+
+        // Load session metadata to get the working directory and other config
+        let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
+            Ok(path) => path,
+            Err(e) => {
+                tracing::error!("Failed to get session path for {}: {}", session_id, e);
                 let _ = stream_event(
                     MessageEvent::Error {
-                        error: "No agent configured".to_string(),
+                        error: format!("Failed to get session path: {}", e),
+                    },
+                    &task_tx,
+                    &cancel_token,
+                )
+                .await;
+                return;
+            }
+        };
+
+        let session_metadata = match session::read_metadata(&session_path) {
+            Ok(metadata) => metadata,
+            Err(e) => {
+                tracing::error!("Failed to read session metadata for {}: {}", session_id, e);
+                let _ = stream_event(
+                    MessageEvent::Error {
+                        error: format!("Failed to read session metadata: {}", e),
                     },
                     &task_tx,
                     &cancel_token,
@@ -230,8 +247,8 @@ async fn reply_handler(
 
         let session_config = SessionConfig {
             id: session::Identifier::Name(session_id.clone()),
-            working_dir: PathBuf::from(&session_working_dir),
-            schedule_id: request.scheduled_job_id.clone(),
+            working_dir: session_metadata.working_dir.clone(),
+            schedule_id: session_metadata.schedule_id.clone(),
             execution_mode: None,
             max_turns: None,
             retry_config: None,
@@ -240,7 +257,7 @@ async fn reply_handler(
         let mut stream = match agent
             .reply(
                 messages.clone(),
-                Some(session_config),
+                Some(session_config.clone()),
                 Some(task_cancel.clone()),
             )
             .await
@@ -344,12 +361,13 @@ async fn reply_handler(
                 let provider = Arc::clone(&provider);
                 let session_path_clone = session_path.to_path_buf();
                 let all_messages_clone = all_messages.clone();
+                let working_dir = session_config.working_dir.clone();
                 tokio::spawn(async move {
                     if let Err(e) = session::persist_messages(
                         &session_path_clone,
                         &all_messages_clone,
                         Some(provider),
-                        Some(PathBuf::from(&session_working_dir)),
+                        Some(working_dir),
                     )
                     .await
                     {
@@ -358,7 +376,6 @@ async fn reply_handler(
                 });
             }
         }
-
         let session_duration = session_start.elapsed();
 
         if let Ok(metadata) = session::read_metadata(&session_path) {
@@ -429,6 +446,8 @@ pub struct PermissionConfirmationRequest {
     #[serde(default = "default_principal_type")]
     principal_type: PrincipalType,
     action: String,
+    #[allow(dead_code)]
+    session_id: String,
 }
 
 fn default_principal_type() -> PrincipalType {
@@ -452,11 +471,7 @@ pub async fn confirm_permission(
 ) -> Result<Json<Value>, StatusCode> {
     verify_secret_key(&headers, &state)?;
 
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
-
+    let agent = state.get_agent().await;
     let permission = match request.action.as_str() {
         "always_allow" => Permission::AlwaysAllow,
         "allow_once" => Permission::AllowOnce,
@@ -480,6 +495,8 @@ pub async fn confirm_permission(
 struct ToolResultRequest {
     id: String,
     result: ToolResult<Vec<Content>>,
+    #[allow(dead_code)]
+    session_id: String,
 }
 
 async fn submit_tool_result(
@@ -506,10 +523,7 @@ async fn submit_tool_result(
         }
     };
 
-    let agent = state
-        .get_agent()
-        .await
-        .map_err(|_| StatusCode::PRECONDITION_FAILED)?;
+    let agent = state.get_agent().await;
     agent.handle_tool_result(payload.id, payload.result).await;
     Ok(Json(json!({"status": "ok"})))
 }
@@ -585,7 +599,7 @@ mod tests {
             });
             let agent = Agent::new();
             let _ = agent.update_provider(mock_provider).await;
-            let state = AppState::new(Arc::new(agent), "test-secret".to_string()).await;
+            let state = AppState::new(Arc::new(agent), "test-secret".to_string());
 
             let app = routes(state);
 
@@ -598,8 +612,6 @@ mod tests {
                     serde_json::to_string(&ChatRequest {
                         messages: vec![Message::user().with_text("test message")],
                         session_id: Some("test-session".to_string()),
-                        session_working_dir: "test-working-dir".to_string(),
-                        scheduled_job_id: None,
                         recipe_name: None,
                         recipe_version: None,
                     })

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

@@ -129,7 +129,7 @@ async fn get_session_history(
     let messages = match session::read_messages(&session_path) {
         Ok(messages) => messages,
         Err(e) => {
-            tracing::error!("Failed to read session messages: {:?}", e);
+            error!("Failed to read session messages: {:?}", e);
             return Err(StatusCode::NOT_FOUND);
         }
     };

crates/goose-server/src/state.rs

@@ -2,43 +2,45 @@ use goose::agents::Agent;
 use goose::scheduler_trait::SchedulerTrait;
 use std::collections::HashMap;
 use std::path::PathBuf;
+use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 use tokio::sync::Mutex;
+use tokio::sync::RwLock;
 
-pub type AgentRef = Arc<Agent>;
+type AgentRef = Arc<Agent>;
 
 #[derive(Clone)]
 pub struct AppState {
-    agent: Option<AgentRef>,
+    agent: Arc<RwLock<AgentRef>>,
     pub secret_key: String,
-    pub scheduler: Arc<Mutex<Option<Arc<dyn SchedulerTrait>>>>,
+    pub scheduler: Arc<RwLock<Option<Arc<dyn SchedulerTrait>>>>,
     pub recipe_file_hash_map: Arc<Mutex<HashMap<String, PathBuf>>>,
+    pub session_counter: Arc<AtomicUsize>,
 }
 
 impl AppState {
-    pub async fn new(agent: AgentRef, secret_key: String) -> Arc<AppState> {
+    pub fn new(agent: AgentRef, secret_key: String) -> Arc<AppState> {
         Arc::new(Self {
-            agent: Some(agent.clone()),
+            agent: Arc::new(RwLock::new(agent)),
             secret_key,
-            scheduler: Arc::new(Mutex::new(None)),
+            scheduler: Arc::new(RwLock::new(None)),
             recipe_file_hash_map: Arc::new(Mutex::new(HashMap::new())),
+            session_counter: Arc::new(AtomicUsize::new(0)),
         })
     }
 
-    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."))
+    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.lock().await;
+        let mut guard = self.scheduler.write().await;
         *guard = Some(sched);
     }
 
     pub async fn scheduler(&self) -> Result<Arc<dyn SchedulerTrait>, anyhow::Error> {
         self.scheduler
-            .lock()
+            .read()
             .await
             .clone()
             .ok_or_else(|| anyhow::anyhow!("Scheduler not initialized"))
@@ -48,4 +50,9 @@ impl AppState {
         let mut map = self.recipe_file_hash_map.lock().await;
         *map = hash_map;
     }
+
+    pub async fn reset(&self) {
+        let mut agent = self.agent.write().await;
+        *agent = Arc::new(Agent::new());
+    }
 }