Next camp refactor live (#5706)

crates/goose-server/src/openapi.rs

@@ -368,6 +368,7 @@ derive_utoipa!(Icon as IconSchema);
         super::routes::session::export_session,
         super::routes::session::import_session,
         super::routes::session::update_session_user_recipe_values,
+        super::routes::session::edit_message,
         super::routes::schedule::create_schedule,
         super::routes::schedule::list_schedules,
         super::routes::schedule::delete_schedule,
@@ -414,6 +415,9 @@ derive_utoipa!(Icon as IconSchema);
         super::routes::session::UpdateSessionNameRequest,
         super::routes::session::UpdateSessionUserRecipeValuesRequest,
         super::routes::session::UpdateSessionUserRecipeValuesResponse,
+        super::routes::session::EditType,
+        super::routes::session::EditMessageRequest,
+        super::routes::session::EditMessageResponse,
         Message,
         MessageContent,
         MessageMetadata,

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

@@ -49,6 +49,31 @@ pub struct ImportSessionRequest {
     json: String,
 }
 
+#[derive(Debug, Deserialize, ToSchema)]
+#[serde(rename_all = "lowercase")]
+pub enum EditType {
+    Fork,
+    Edit,
+}
+
+#[derive(Deserialize, ToSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct EditMessageRequest {
+    timestamp: i64,
+    #[serde(default = "default_edit_type")]
+    edit_type: EditType,
+}
+
+fn default_edit_type() -> EditType {
+    EditType::Fork
+}
+
+#[derive(Serialize, ToSchema)]
+#[serde(rename_all = "camelCase")]
+pub struct EditMessageResponse {
+    session_id: String,
+}
+
 const MAX_NAME_LENGTH: usize = 200;
 
 #[utoipa::path(
@@ -307,6 +332,64 @@ async fn import_session(
     Ok(Json(session))
 }
 
+#[utoipa::path(
+    post,
+    path = "/sessions/{session_id}/edit_message",
+    request_body = EditMessageRequest,
+    params(
+        ("session_id" = String, Path, description = "Unique identifier for the session")
+    ),
+    responses(
+        (status = 200, description = "Session prepared for editing - frontend should submit the edited message", body = EditMessageResponse),
+        (status = 400, description = "Bad request - Invalid message timestamp"),
+        (status = 401, description = "Unauthorized - Invalid or missing API key"),
+        (status = 404, description = "Session or message not found"),
+        (status = 500, description = "Internal server error")
+    ),
+    security(
+        ("api_key" = [])
+    ),
+    tag = "Session Management"
+)]
+async fn edit_message(
+    Path(session_id): Path<String>,
+    Json(request): Json<EditMessageRequest>,
+) -> Result<Json<EditMessageResponse>, StatusCode> {
+    match request.edit_type {
+        EditType::Fork => {
+            let new_session = SessionManager::copy_session(&session_id, "(edited)".to_string())
+                .await
+                .map_err(|e| {
+                    tracing::error!("Failed to copy session: {}", e);
+                    StatusCode::INTERNAL_SERVER_ERROR
+                })?;
+
+            SessionManager::truncate_conversation(&new_session.id, request.timestamp)
+                .await
+                .map_err(|e| {
+                    tracing::error!("Failed to truncate conversation: {}", e);
+                    StatusCode::INTERNAL_SERVER_ERROR
+                })?;
+
+            Ok(Json(EditMessageResponse {
+                session_id: new_session.id,
+            }))
+        }
+        EditType::Edit => {
+            SessionManager::truncate_conversation(&session_id, request.timestamp)
+                .await
+                .map_err(|e| {
+                    tracing::error!("Failed to truncate conversation: {}", e);
+                    StatusCode::INTERNAL_SERVER_ERROR
+                })?;
+
+            Ok(Json(EditMessageResponse {
+                session_id: session_id.clone(),
+            }))
+        }
+    }
+}
+
 pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/sessions", get(list_sessions))
@@ -320,5 +403,6 @@ pub fn routes(state: Arc<AppState>) -> Router {
             "/sessions/{session_id}/user_recipe_values",
             put(update_session_user_recipe_values),
         )
+        .route("/sessions/{session_id}/edit_message", post(edit_message))
         .with_state(state)
 }

crates/goose/src/agents/agent.rs

@@ -805,9 +805,7 @@ impl Agent {
         } else {
             SessionManager::add_message(&session_config.id, &user_message).await?;
         }
-
         let session = SessionManager::get_session(&session_config.id, true).await?;
-
         let conversation = session
             .conversation
             .clone()

crates/goose/src/agents/tool_execution.rs

@@ -71,12 +71,14 @@ impl Agent {
                             }
                         });
 
-                    let confirmation = Message::user().with_tool_confirmation_request(
-                        request.id.clone(),
-                        tool_call.name.to_string().clone(),
-                        tool_call.arguments.clone().unwrap_or_default(),
-                        security_message,
-                    );
+                    let confirmation = Message::assistant()
+                        .with_tool_confirmation_request(
+                            request.id.clone(),
+                            tool_call.name.to_string().clone(),
+                            tool_call.arguments.clone().unwrap_or_default(),
+                            security_message,
+                        )
+                        .user_only();
                     yield confirmation;
 
                     let mut rx = self.confirmation_rx.lock().await;

crates/goose/src/session/session_manager.rs

@@ -307,6 +307,20 @@ impl SessionManager {
         Self::instance().await?.import_session(json).await
     }
 
+    pub async fn copy_session(session_id: &str, new_name: String) -> Result<Session> {
+        Self::instance()
+            .await?
+            .copy_session(session_id, new_name)
+            .await
+    }
+
+    pub async fn truncate_conversation(session_id: &str, timestamp: i64) -> Result<()> {
+        Self::instance()
+            .await?
+            .truncate_conversation(session_id, timestamp)
+            .await
+    }
+
     pub async fn maybe_update_name(id: &str, provider: Arc<dyn Provider>) -> Result<()> {
         let session = Self::get_session(id, true).await?;
 
@@ -1214,6 +1228,43 @@ impl SessionStorage {
         self.get_session(&session.id, true).await
     }
 
+    async fn copy_session(&self, session_id: &str, new_name: String) -> Result<Session> {
+        let original_session = self.get_session(session_id, true).await?;
+
+        let new_session = self
+            .create_session(
+                original_session.working_dir.clone(),
+                new_name,
+                original_session.session_type,
+            )
+            .await?;
+
+        let builder = SessionUpdateBuilder::new(new_session.id.clone())
+            .extension_data(original_session.extension_data)
+            .schedule_id(original_session.schedule_id)
+            .recipe(original_session.recipe)
+            .user_recipe_values(original_session.user_recipe_values);
+
+        self.apply_update(builder).await?;
+
+        if let Some(conversation) = original_session.conversation {
+            self.replace_conversation(&new_session.id, &conversation)
+                .await?;
+        }
+
+        self.get_session(&new_session.id, true).await
+    }
+
+    async fn truncate_conversation(&self, session_id: &str, timestamp: i64) -> Result<()> {
+        sqlx::query("DELETE FROM messages WHERE session_id = ? AND created_timestamp >= ?")
+            .bind(session_id)
+            .bind(timestamp)
+            .execute(&self.pool)
+            .await?;
+
+        Ok(())
+    }
+
     async fn search_chat_history(
         &self,
         query: &str,