Next Camp Live - Added editing messages functionality

crates/goose-server/src/openapi.rs

@@ -364,6 +364,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,
@@ -405,6 +406,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/reply.rs

@@ -85,6 +85,8 @@ pub struct ChatRequest {
     session_id: String,
     recipe_name: Option<String>,
     recipe_version: Option<String>,
+    #[serde(default)]
+    skip_add_user_message: bool,
 }
 
 pub struct SseResponse {
@@ -300,10 +302,11 @@ pub async fn reply(
         };
 
         let mut stream = match agent
-            .reply(
+            .reply_with_options(
                 user_message.clone(),
                 session_config,
                 Some(task_cancel.clone()),
+                request.skip_add_user_message,
             )
             .await
         {
@@ -536,6 +539,7 @@ mod tests {
                         session_id: "test-session".to_string(),
                         recipe_name: None,
                         recipe_version: None,
+                        skip_add_user_message: false,
                     })
                     .unwrap(),
                 ))

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

@@ -9,6 +9,7 @@ use axum::{
     routing::{delete, get, put},
     Json, Router,
 };
+use goose::conversation::message::Message;
 use goose::recipe::Recipe;
 use goose::session::session_manager::SessionInsights;
 use goose::session::{Session, SessionManager};
@@ -49,6 +50,36 @@ 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 {
+    message_row_id: i64,
+    new_content: String,
+    #[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 {
+    /// New session ID created from the fork (only present for fork edit_type)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    new_session_id: Option<String>,
+    /// Conversation (either in new session for fork, or updated current session for edit)
+    conversation: Vec<Message>,
+}
+
 const MAX_NAME_LENGTH: usize = 200;
 
 #[utoipa::path(
@@ -307,6 +338,82 @@ 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 = "Message edited successfully", body = EditMessageResponse),
+        (status = 400, description = "Bad request - Invalid message ID or empty content"),
+        (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> {
+    if request.new_content.trim().is_empty() {
+        tracing::warn!("edit_message: empty content provided");
+        return Err(StatusCode::BAD_REQUEST);
+    }
+
+    match request.edit_type {
+        EditType::Fork => {
+            let new_session = SessionManager::fork_session_at_message(
+                &session_id,
+                request.message_row_id,
+                request.new_content,
+            )
+            .await
+            .map_err(|e| {
+                tracing::error!("Failed to fork session: {}", e);
+                StatusCode::INTERNAL_SERVER_ERROR
+            })?;
+
+            let conversation = new_session.conversation.ok_or_else(|| {
+                tracing::error!("Forked session has no conversation");
+                StatusCode::INTERNAL_SERVER_ERROR
+            })?;
+
+            Ok(Json(EditMessageResponse {
+                new_session_id: Some(new_session.id),
+                conversation: conversation.messages().to_vec(),
+            }))
+        }
+        EditType::Edit => {
+            let updated_session = SessionManager::edit_message_in_place(
+                &session_id,
+                request.message_row_id,
+                request.new_content,
+            )
+            .await
+            .map_err(|e| {
+                tracing::error!("Failed to edit message in place: {}", e);
+                StatusCode::INTERNAL_SERVER_ERROR
+            })?;
+
+            let conversation = updated_session.conversation.ok_or_else(|| {
+                tracing::error!("Updated session has no conversation");
+                StatusCode::INTERNAL_SERVER_ERROR
+            })?;
+
+            Ok(Json(EditMessageResponse {
+                new_session_id: None,
+                conversation: conversation.messages().to_vec(),
+            }))
+        }
+    }
+}
+
 pub fn routes(state: Arc<AppState>) -> Router {
     Router::new()
         .route("/sessions", get(list_sessions))
@@ -320,5 +427,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

@@ -776,6 +776,18 @@ impl Agent {
         user_message: Message,
         session_config: SessionConfig,
         cancel_token: Option<CancellationToken>,
+    ) -> Result<BoxStream<'_, Result<AgentEvent>>> {
+        self.reply_with_options(user_message, session_config, cancel_token, false)
+            .await
+    }
+
+    #[instrument(skip(self, user_message, session_config), fields(user_message))]
+    pub async fn reply_with_options(
+        &self,
+        user_message: Message,
+        session_config: SessionConfig,
+        cancel_token: Option<CancellationToken>,
+        skip_add_message: bool,
     ) -> Result<BoxStream<'_, Result<AgentEvent>>> {
         let is_manual_compact = user_message.content.iter().any(|c| {
             if let MessageContent::Text(text) = c {
@@ -785,9 +797,11 @@ impl Agent {
             }
         });
 
-        SessionManager::add_message(&session_config.id, &user_message).await?;
-        let session = SessionManager::get_session(&session_config.id, true).await?;
+        if !skip_add_message {
+            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/conversation/message.rs

@@ -468,6 +468,8 @@ impl MessageMetadata {
 #[serde(rename_all = "camelCase")]
 pub struct Message {
     pub id: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub row_id: Option<i64>,
     pub role: Role,
     pub created: i64,
     #[serde(deserialize_with = "deserialize_sanitized_content")]
@@ -479,6 +481,7 @@ impl Message {
     pub fn new(role: Role, created: i64, content: Vec<MessageContent>) -> Self {
         Message {
             id: None,
+            row_id: None,
             role,
             created,
             content,
@@ -493,6 +496,7 @@ impl Message {
     pub fn user() -> Self {
         Message {
             id: None,
+            row_id: None,
             role: Role::User,
             created: Utc::now().timestamp(),
             content: Vec::new(),
@@ -504,6 +508,7 @@ impl Message {
     pub fn assistant() -> Self {
         Message {
             id: None,
+            row_id: None,
             role: Role::Assistant,
             created: Utc::now().timestamp(),
             content: Vec::new(),