From 1ebf901e75e2e15f9272aff0da5a2ec857e71627 Mon Sep 17 00:00:00 2001 From: Zane Staggs Date: Wed, 19 Nov 2025 10:16:22 -0700 Subject: [PATCH 1/5] Added editing and forking messages --- crates/goose-server/src/openapi.rs | 4 + crates/goose-server/src/routes/reply.rs | 6 +- crates/goose-server/src/routes/session.rs | 108 +++++++++++++++ crates/goose/src/agents/agent.rs | 18 ++- crates/goose/src/conversation/message.rs | 5 + crates/goose/src/session/session_manager.rs | 145 +++++++++++++++++++- ui/desktop/openapi.json | 112 +++++++++++++++ ui/desktop/src/App.tsx | 1 + ui/desktop/src/api/sdk.gen.ts | 13 +- ui/desktop/src/api/types.gen.ts | 61 ++++++++ ui/desktop/src/components/BaseChat.tsx | 59 +++++++- ui/desktop/src/components/UserMessage.tsx | 58 +++++--- ui/desktop/src/hooks/useChatStream.ts | 121 ++++++++++++++-- 13 files changed, 675 insertions(+), 36 deletions(-) diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index b3f9aedb14f6..19606111e11d 100644 --- a/crates/goose-server/src/openapi.rs +++ b/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, diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 5a340fa3d947..c763f0f6b0bc 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -85,6 +85,8 @@ pub struct ChatRequest { session_id: String, recipe_name: Option, recipe_version: Option, + #[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(), )) diff --git a/crates/goose-server/src/routes/session.rs b/crates/goose-server/src/routes/session.rs index 67a5bf388d59..66b8f88827ec 100644 --- a/crates/goose-server/src/routes/session.rs +++ b/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, + /// Conversation (either in new session for fork, or updated current session for edit) + conversation: Vec, +} + 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, + Json(request): Json, +) -> Result, 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) -> Router { Router::new() .route("/sessions", get(list_sessions)) @@ -320,5 +427,6 @@ pub fn routes(state: Arc) -> 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) } diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index c9f84e9571ba..2bbff1c432aa 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -776,6 +776,18 @@ impl Agent { user_message: Message, session_config: SessionConfig, cancel_token: Option, + ) -> Result>> { + 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, + skip_add_message: bool, ) -> Result>> { 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() diff --git a/crates/goose/src/conversation/message.rs b/crates/goose/src/conversation/message.rs index 2f18d038836f..2b034b0123e0 100644 --- a/crates/goose/src/conversation/message.rs +++ b/crates/goose/src/conversation/message.rs @@ -468,6 +468,8 @@ impl MessageMetadata { #[serde(rename_all = "camelCase")] pub struct Message { pub id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub row_id: Option, 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) -> 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(), diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 564911acbc1e..16604d87604c 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -290,6 +290,28 @@ impl SessionManager { Self::instance().await?.import_session(json).await } + pub async fn fork_session_at_message( + session_id: &str, + message_row_id: i64, + new_message_content: String, + ) -> Result { + Self::instance() + .await? + .fork_session_at_message(session_id, message_row_id, new_message_content) + .await + } + + pub async fn edit_message_in_place( + session_id: &str, + message_row_id: i64, + new_message_content: String, + ) -> Result { + Self::instance() + .await? + .edit_message_in_place(session_id, message_row_id, new_message_content) + .await + } + pub async fn maybe_update_name(id: &str, provider: Arc) -> Result<()> { let session = Self::get_session(id, true).await?; @@ -938,15 +960,15 @@ impl SessionStorage { } async fn get_conversation(&self, session_id: &str) -> Result { - let rows = sqlx::query_as::<_, (String, String, i64, Option)>( - "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp", + let rows = sqlx::query_as::<_, (i64, String, String, i64, Option)>( + "SELECT id, role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp", ) .bind(session_id) .fetch_all(&self.pool) .await?; let mut messages = Vec::new(); - for (idx, (role_str, content_json, created_timestamp, metadata_json)) in + for (idx, (row_id, role_str, content_json, created_timestamp, metadata_json)) in rows.into_iter().enumerate() { let role = match role_str.as_str() { @@ -962,6 +984,7 @@ impl SessionStorage { let mut message = Message::new(role, created_timestamp, content); message.metadata = metadata; + message.row_id = Some(row_id); message = message.with_id(format!("msg_{}_{}", session_id, idx)); messages.push(message); } @@ -1137,6 +1160,118 @@ impl SessionStorage { self.get_session(&session.id, true).await } + async fn fork_session_at_message( + &self, + session_id: &str, + message_row_id: i64, + new_message_content: String, + ) -> Result { + use crate::conversation::message::MessageContent; + + let original_session = self.get_session(session_id, false).await?; + let rows = sqlx::query_as::<_, (String, String, i64, Option)>( + "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? AND id < ? ORDER BY id", + ) + .bind(session_id) + .bind(message_row_id) + .fetch_all(&self.pool) + .await?; + + let mut messages = Vec::new(); + for (role_str, content_json, created_timestamp, metadata_json) in rows { + let role = match role_str.as_str() { + "user" => Role::User, + "assistant" => Role::Assistant, + _ => continue, + }; + + let content = serde_json::from_str(&content_json)?; + let metadata = metadata_json + .and_then(|json| serde_json::from_str(&json).ok()) + .unwrap_or_default(); + + let mut message = Message::new(role, created_timestamp, content); + message.metadata = metadata; + messages.push(message); + } + + let new_user_message = Message::new( + Role::User, + chrono::Utc::now().timestamp_millis(), + vec![MessageContent::text(&new_message_content)], + ); + messages.push(new_user_message); + + let new_session = self + .create_session( + original_session.working_dir.clone(), + format!("{} (edited)", original_session.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?; + + let conversation = Conversation::new_unvalidated(messages); + self.replace_conversation(&new_session.id, &conversation) + .await?; + + self.get_session(&new_session.id, true).await + } + + async fn edit_message_in_place( + &self, + session_id: &str, + message_row_id: i64, + new_message_content: String, + ) -> Result { + use crate::conversation::message::MessageContent; + + let rows = sqlx::query_as::<_, (String, String, i64, Option)>( + "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? AND id < ? ORDER BY id", + ) + .bind(session_id) + .bind(message_row_id) + .fetch_all(&self.pool) + .await?; + + let mut messages = Vec::new(); + for (role_str, content_json, created_timestamp, metadata_json) in rows { + let role = match role_str.as_str() { + "user" => Role::User, + "assistant" => Role::Assistant, + _ => continue, + }; + + let content = serde_json::from_str(&content_json)?; + let metadata = metadata_json + .and_then(|json| serde_json::from_str(&json).ok()) + .unwrap_or_default(); + + let mut message = Message::new(role, created_timestamp, content); + message.metadata = metadata; + messages.push(message); + } + + let new_user_message = Message::new( + Role::User, + chrono::Utc::now().timestamp_millis(), + vec![MessageContent::text(&new_message_content)], + ); + messages.push(new_user_message); + + let conversation = Conversation::new_unvalidated(messages); + self.replace_conversation(session_id, &conversation).await?; + + self.get_session(session_id, true).await + } + async fn search_chat_history( &self, query: &str, @@ -1193,6 +1328,7 @@ mod tests { &session.id, &Message { id: None, + row_id: None, role: Role::User, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text("hello world")], @@ -1207,6 +1343,7 @@ mod tests { &session.id, &Message { id: None, + row_id: None, role: Role::Assistant, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text("sup world?")], @@ -1300,6 +1437,7 @@ mod tests { &original.id, &Message { id: None, + row_id: None, role: Role::User, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text(USER_MESSAGE)], @@ -1314,6 +1452,7 @@ mod tests { &original.id, &Message { id: None, + row_id: None, role: Role::Assistant, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text(ASSISTANT_MESSAGE)], diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index e7e4f8a4f6da..d1267b6a6747 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -1916,6 +1916,64 @@ ] } }, + "/sessions/{session_id}/edit_message": { + "post": { + "tags": [ + "Session Management" + ], + "operationId": "edit_message", + "parameters": [ + { + "name": "session_id", + "in": "path", + "description": "Unique identifier for the session", + "required": true, + "schema": { + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/EditMessageRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "description": "Message edited successfully", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/EditMessageResponse" + } + } + } + }, + "400": { + "description": "Bad request - Invalid message ID or empty content" + }, + "401": { + "description": "Unauthorized - Invalid or missing API key" + }, + "404": { + "description": "Session or message not found" + }, + "500": { + "description": "Internal server error" + } + }, + "security": [ + { + "api_key": [] + } + ] + } + }, "/sessions/{session_id}/export": { "get": { "tags": [ @@ -2186,6 +2244,9 @@ }, "session_id": { "type": "string" + }, + "skip_add_user_message": { + "type": "boolean" } } }, @@ -2448,6 +2509,52 @@ } } }, + "EditMessageRequest": { + "type": "object", + "required": [ + "messageRowId", + "newContent" + ], + "properties": { + "editType": { + "$ref": "#/components/schemas/EditType" + }, + "messageRowId": { + "type": "integer", + "format": "int64" + }, + "newContent": { + "type": "string" + } + } + }, + "EditMessageResponse": { + "type": "object", + "required": [ + "conversation" + ], + "properties": { + "conversation": { + "type": "array", + "items": { + "$ref": "#/components/schemas/Message" + }, + "description": "Conversation (either in new session for fork, or updated current session for edit)" + }, + "newSessionId": { + "type": "string", + "description": "New session ID created from the fork (only present for fork edit_type)", + "nullable": true + } + } + }, + "EditType": { + "type": "string", + "enum": [ + "fork", + "edit" + ] + }, "EmbeddedResource": { "type": "object", "required": [ @@ -3119,6 +3226,11 @@ }, "role": { "$ref": "#/components/schemas/Role" + }, + "rowId": { + "type": "integer", + "format": "int64", + "nullable": true } } }, diff --git a/ui/desktop/src/App.tsx b/ui/desktop/src/App.tsx index 212f2b03af58..adeca871e596 100644 --- a/ui/desktop/src/App.tsx +++ b/ui/desktop/src/App.tsx @@ -178,6 +178,7 @@ const PairRouteWrapper = ({ return ( = Options2 & { /** @@ -522,6 +522,17 @@ export const getSession = (options: Option }); }; +export const editMessage = (options: Options) => { + return (options.client ?? client).post({ + url: '/sessions/{session_id}/edit_message', + ...options, + headers: { + 'Content-Type': 'application/json', + ...options.headers + } + }); +}; + export const exportSession = (options: Options) => { return (options.client ?? client).get({ url: '/sessions/{session_id}/export', diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 75de116e5690..a46cece49b9b 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -30,6 +30,7 @@ export type ChatRequest = { recipe_name?: string | null; recipe_version?: string | null; session_id: string; + skip_add_user_message?: boolean; }; export type CheckProviderRequest = { @@ -121,6 +122,25 @@ export type DeleteRecipeRequest = { id: string; }; +export type EditMessageRequest = { + editType?: EditType; + messageRowId: number; + newContent: string; +}; + +export type EditMessageResponse = { + /** + * Conversation (either in new session for fork, or updated current session for edit) + */ + conversation: Array; + /** + * New session ID created from the fork (only present for fork edit_type) + */ + newSessionId?: string | null; +}; + +export type EditType = 'fork' | 'edit'; + export type EmbeddedResource = { _meta?: { [key: string]: unknown; @@ -343,6 +363,7 @@ export type Message = { id?: string | null; metadata: MessageMetadata; role: Role; + rowId?: number | null; }; /** @@ -2417,6 +2438,46 @@ export type GetSessionResponses = { export type GetSessionResponse = GetSessionResponses[keyof GetSessionResponses]; +export type EditMessageData = { + body: EditMessageRequest; + path: { + /** + * Unique identifier for the session + */ + session_id: string; + }; + query?: never; + url: '/sessions/{session_id}/edit_message'; +}; + +export type EditMessageErrors = { + /** + * Bad request - Invalid message ID or empty content + */ + 400: unknown; + /** + * Unauthorized - Invalid or missing API key + */ + 401: unknown; + /** + * Session or message not found + */ + 404: unknown; + /** + * Internal server error + */ + 500: unknown; +}; + +export type EditMessageResponses = { + /** + * Message edited successfully + */ + 200: EditMessageResponse; +}; + +export type EditMessageResponse2 = EditMessageResponses[keyof EditMessageResponses]; + export type ExportSessionData = { body?: never; path: { diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index b004a1bcc5e3..8b42946be7e5 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -7,7 +7,7 @@ import React, { useRef, useState, } from 'react'; -import { useLocation } from 'react-router-dom'; +import { useLocation, useNavigate, useSearchParams } from 'react-router-dom'; import { SearchView } from './conversation/SearchView'; import LoadingGoose from './LoadingGoose'; import PopularChatTopics from './PopularChatTopics'; @@ -65,6 +65,8 @@ function BaseChatContent({ initialMessage, }: BaseChatProps) { const location = useLocation(); + const navigate = useNavigate(); + const [searchParams] = useSearchParams(); const scrollRef = useRef(null); const disableAnimation = location.state?.disableAnimation || false; @@ -85,11 +87,13 @@ function BaseChatContent({ const [sessionLoaded, setSessionLoaded] = useState(false); const [hasSubmittedInitialMessage, setHasSubmittedInitialMessage] = useState(false); + const [hasTriggeredAgentForFork, setHasTriggeredAgentForFork] = useState(false); const [isCreateRecipeModalOpen, setIsCreateRecipeModalOpen] = useState(false); useEffect(() => { setSessionLoaded(false); setHasSubmittedInitialMessage(false); + setHasTriggeredAgentForFork(false); }, [sessionId]); const handleSessionLoaded = useCallback(() => { @@ -106,6 +110,7 @@ function BaseChatContent({ setRecipeUserParams, tokenState, notifications, + onMessageUpdate, } = useChatStream({ sessionId, onStreamFinish, @@ -120,6 +125,29 @@ function BaseChatContent({ } }, [sessionLoaded, initialMessage, hasSubmittedInitialMessage, handleSubmit]); + useEffect(() => { + const shouldStartAgent = searchParams.get('shouldStartAgent') === 'true'; + + if ( + sessionLoaded && + shouldStartAgent && + messages.length > 0 && + !hasTriggeredAgentForFork && + chatState === ChatState.Idle + ) { + setHasTriggeredAgentForFork(true); + handleSubmit('', { isContinuation: true }); + } + }, [ + sessionLoaded, + searchParams, + messages.length, + hasTriggeredAgentForFork, + chatState, + handleSubmit, + messages, + ]); + const handleFormSubmit = (e: React.FormEvent) => { const customEvent = e as unknown as CustomEvent; const textValue = customEvent.detail?.value || ''; @@ -207,6 +235,34 @@ function BaseChatContent({ return () => window.removeEventListener('make-agent-from-chat', handleMakeAgent); }, []); + useEffect(() => { + const handleSessionForked = (event: Event) => { + const customEvent = event as CustomEvent<{ + newSessionId: string; + shouldStartAgent?: boolean; + }>; + const { newSessionId, shouldStartAgent } = customEvent.detail; + + const params = new URLSearchParams(); + params.set('resumeSessionId', newSessionId); + if (shouldStartAgent) { + params.set('shouldStartAgent', 'true'); + } + + navigate(`/pair?${params.toString()}`, { + state: { + disableAnimation: true, + }, + }); + }; + + window.addEventListener('session-forked', handleSessionForked); + + return () => { + window.removeEventListener('session-forked', handleSessionForked); + }; + }, [location.pathname, navigate]); + const handleRecipeCreated = (recipe: Recipe) => { toastSuccess({ title: 'Recipe created successfully!', @@ -234,6 +290,7 @@ function BaseChatContent({ isUserMessage={(m: Message) => m.role === 'user'} isStreamingMessage={chatState !== ChatState.Idle} onRenderingComplete={handleRenderingComplete} + onMessageUpdate={onMessageUpdate} /> ); diff --git a/ui/desktop/src/components/UserMessage.tsx b/ui/desktop/src/components/UserMessage.tsx index 1292d4fbeba6..87613e6f4cf1 100644 --- a/ui/desktop/src/components/UserMessage.tsx +++ b/ui/desktop/src/components/UserMessage.tsx @@ -11,7 +11,7 @@ import { Button } from './ui/button'; interface UserMessageProps { message: Message; - onMessageUpdate?: (messageId: string, newContent: string) => void; + onMessageUpdate?: (messageId: string, newContent: string, editType?: 'fork' | 'edit') => void; } export default function UserMessage({ message, onMessageUpdate }: UserMessageProps) { @@ -85,26 +85,26 @@ export default function UserMessage({ message, onMessageUpdate }: UserMessagePro window.electron.logInfo(`Content changed: ${newContent}`); }, []); - // Handle save action - const handleSave = useCallback(() => { - // Exit edit mode immediately - setIsEditing(false); - - // Check if content has actually changed - if (editContent !== displayText) { - // Validate content + const handleSave = useCallback( + (editType: 'fork' | 'edit' = 'fork') => { if (editContent.trim().length === 0) { setError('Message cannot be empty'); return; } - // Update the message content through the callback + setIsEditing(false); + + if (editContent.trim() === displayText.trim()) { + return; + } + if (onMessageUpdate && message.id) { - onMessageUpdate(message.id, editContent); + onMessageUpdate(message.id, editContent, editType); setHasBeenEdited(true); } - } - }, [editContent, displayText, onMessageUpdate, message.id]); + }, + [editContent, displayText, onMessageUpdate, message.id] + ); // Handle cancel action const handleCancel = useCallback(() => { @@ -177,13 +177,31 @@ export default function UserMessage({ message, onMessageUpdate }: UserMessagePro {error} )} -
- - +
+
+ Edit in Place updates this session •{' '} + Fork Session creates a new session +
+
+ + + +
) : ( diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index 0ace33eac7e5..152ec641beb0 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -29,12 +29,17 @@ interface UseChatStreamReturn { session?: Session; messages: Message[]; chatState: ChatState; - handleSubmit: (userMessage: string) => Promise; + handleSubmit: (userMessage: string, options?: { isContinuation?: boolean }) => Promise; setRecipeUserParams: (values: Record) => Promise; stopStreaming: () => void; sessionLoadError?: string; tokenState: TokenState; notifications: NotificationEvent[]; + onMessageUpdate: ( + messageId: string, + newContent: string, + editType?: 'fork' | 'edit' + ) => Promise; } function pushMessage(currentMessages: Message[], incomingMsg: Message): Message[] { @@ -167,7 +172,7 @@ export function useChatStream({ }, []); const onFinish = useCallback( - (error?: string): void => { + async (error?: string): Promise => { if (error) { setSessionLoadError(error); } @@ -180,10 +185,28 @@ export function useChatStream({ window.dispatchEvent(new CustomEvent('message-stream-finished')); } + // Reload the session to get auto generated/incremented rowIds for all messages + try { + const { getSession } = await import('../api'); + const response = await getSession({ + path: { session_id: sessionId }, + throwOnError: true, + }); + + if (response.data?.conversation) { + window.electron.logInfo( + `Reloaded session with ${response.data.conversation.length} messages with rowIds` + ); + updateMessages(response.data.conversation); + } + } catch (reloadError) { + window.electron.logInfo(`Failed to reload session: ${errorMessage(reloadError)}`); + } + setChatState(ChatState.Idle); onStreamFinish(); }, - [onStreamFinish, sessionId] + [onStreamFinish, sessionId, updateMessages] ); // Load session on mount or sessionId change @@ -239,22 +262,35 @@ export function useChatStream({ }, [sessionId, updateMessages, onSessionLoaded]); const handleSubmit = useCallback( - async (userMessage: string) => { + async (userMessage: string, options?: { isContinuation?: boolean }) => { // Guard: Don't submit if session hasn't been loaded yet if (!session || chatState === ChatState.LoadingConversation) { return; } - if (!userMessage.trim()) { + // Check if this is a continuation request (agent should respond to existing messages) + const isContinuationRequest = options?.isContinuation || false; + if (!isContinuationRequest && !userMessage.trim()) { return; } - if (messagesRef.current.length === 0) { + if (messagesRef.current.length === 0 && !isContinuationRequest) { window.dispatchEvent(new CustomEvent('session-created')); } - const currentMessages = [...messagesRef.current, createUserMessage(userMessage)]; - updateMessages(currentMessages); + // For continuation requests, don't add a new user message + const currentMessages = isContinuationRequest + ? [...messagesRef.current] + : [...messagesRef.current, createUserMessage(userMessage)]; + + window.electron.logInfo( + `handleSubmit: isContinuationRequest=${isContinuationRequest}, messagesRef.current.length=${messagesRef.current.length}, currentMessages.length=${currentMessages.length}` + ); + + // Don't update messages if it's a continuation request (messages are already correct) + if (!isContinuationRequest) { + updateMessages(currentMessages); + } setChatState(ChatState.Streaming); setNotifications([]); @@ -265,6 +301,7 @@ export function useChatStream({ body: { session_id: sessionId, messages: currentMessages, + skip_add_user_message: isContinuationRequest, // Tell backend not to add the message again }, throwOnError: true, signal: abortControllerRef.current.signal, @@ -335,6 +372,73 @@ export function useChatStream({ setChatState(ChatState.Idle); }, []); + const onMessageUpdate = useCallback( + async (messageId: string, newContent: string, editType: 'fork' | 'edit' = 'fork') => { + try { + const { editMessage } = await import('../api'); + const message = messagesRef.current.find((m) => m.id === messageId); + + if (!message) { + throw new Error(`Message with id ${messageId} not found in current messages`); + } + + if (!message.rowId) { + throw new Error( + `Message ${messageId} is missing rowId. This might be an old session - please reload the page and try again.` + ); + } + + const response = await editMessage({ + path: { + session_id: sessionId, + }, + body: { + messageRowId: message.rowId, + newContent, + editType, + }, + throwOnError: true, + }); + + if (editType === 'fork') { + if (response.data?.newSessionId && response.data?.conversation) { + const newSessionId = response.data.newSessionId; + const event = new CustomEvent('session-forked', { + detail: { + newSessionId, + shouldStartAgent: true, + }, + }); + window.dispatchEvent(event); + window.electron.logInfo(`Dispatched session-forked event`); + } else { + window.electron.logInfo( + `Unexpected fork response format: ${JSON.stringify(response.data)}` + ); + } + } else { + if (response.data?.conversation) { + updateMessages(response.data.conversation); + await handleSubmit('', { isContinuation: true }); + } else { + window.electron.logInfo( + `Unexpected edit response format: ${JSON.stringify(response.data)}` + ); + } + } + } catch (error) { + const errorMsg = errorMessage(error); + console.error('Failed to edit message:', error); + const { toastError } = await import('../toasts'); + toastError({ + title: 'Failed to edit message', + msg: errorMsg, + }); + } + }, + [sessionId, updateMessages, handleSubmit] + ); + const cached = resultsCache.get(sessionId); const maybe_cached_messages = session ? messages : cached?.messages || []; const maybe_cached_session = session ?? cached?.session; @@ -349,5 +453,6 @@ export function useChatStream({ setRecipeUserParams, tokenState, notifications, + onMessageUpdate, }; } From d2f575f26876b84d884f3dac8d63b9bf90b788ea Mon Sep 17 00:00:00 2001 From: Zane Staggs Date: Wed, 19 Nov 2025 19:24:03 -0700 Subject: [PATCH 2/5] change to use timestamp rather than row id --- crates/goose-server/src/routes/session.rs | 6 +- crates/goose/src/conversation/message.rs | 5 - crates/goose/src/session/session_manager.rs | 149 +++++++++----------- ui/desktop/openapi.json | 15 +- ui/desktop/src/api/types.gen.ts | 3 +- ui/desktop/src/hooks/useChatStream.ts | 28 +--- 6 files changed, 76 insertions(+), 130 deletions(-) diff --git a/crates/goose-server/src/routes/session.rs b/crates/goose-server/src/routes/session.rs index 66b8f88827ec..259900b0be37 100644 --- a/crates/goose-server/src/routes/session.rs +++ b/crates/goose-server/src/routes/session.rs @@ -60,7 +60,7 @@ pub enum EditType { #[derive(Deserialize, ToSchema)] #[serde(rename_all = "camelCase")] pub struct EditMessageRequest { - message_row_id: i64, + timestamp: i64, new_content: String, #[serde(default = "default_edit_type")] edit_type: EditType, @@ -370,7 +370,7 @@ async fn edit_message( EditType::Fork => { let new_session = SessionManager::fork_session_at_message( &session_id, - request.message_row_id, + request.timestamp, request.new_content, ) .await @@ -392,7 +392,7 @@ async fn edit_message( EditType::Edit => { let updated_session = SessionManager::edit_message_in_place( &session_id, - request.message_row_id, + request.timestamp, request.new_content, ) .await diff --git a/crates/goose/src/conversation/message.rs b/crates/goose/src/conversation/message.rs index 2b034b0123e0..2f18d038836f 100644 --- a/crates/goose/src/conversation/message.rs +++ b/crates/goose/src/conversation/message.rs @@ -468,8 +468,6 @@ impl MessageMetadata { #[serde(rename_all = "camelCase")] pub struct Message { pub id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub row_id: Option, pub role: Role, pub created: i64, #[serde(deserialize_with = "deserialize_sanitized_content")] @@ -481,7 +479,6 @@ impl Message { pub fn new(role: Role, created: i64, content: Vec) -> Self { Message { id: None, - row_id: None, role, created, content, @@ -496,7 +493,6 @@ impl Message { pub fn user() -> Self { Message { id: None, - row_id: None, role: Role::User, created: Utc::now().timestamp(), content: Vec::new(), @@ -508,7 +504,6 @@ impl Message { pub fn assistant() -> Self { Message { id: None, - row_id: None, role: Role::Assistant, created: Utc::now().timestamp(), content: Vec::new(), diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 16604d87604c..885ad97d19c0 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -290,25 +290,39 @@ impl SessionManager { Self::instance().await?.import_session(json).await } + pub async fn copy_session(session_id: &str, new_name: String) -> Result { + 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 fork_session_at_message( session_id: &str, - message_row_id: i64, + timestamp: i64, new_message_content: String, ) -> Result { Self::instance() .await? - .fork_session_at_message(session_id, message_row_id, new_message_content) + .fork_session_at_message(session_id, timestamp, new_message_content) .await } pub async fn edit_message_in_place( session_id: &str, - message_row_id: i64, + timestamp: i64, new_message_content: String, ) -> Result { Self::instance() .await? - .edit_message_in_place(session_id, message_row_id, new_message_content) + .edit_message_in_place(session_id, timestamp, new_message_content) .await } @@ -960,15 +974,15 @@ impl SessionStorage { } async fn get_conversation(&self, session_id: &str) -> Result { - let rows = sqlx::query_as::<_, (i64, String, String, i64, Option)>( - "SELECT id, role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp", + let rows = sqlx::query_as::<_, (String, String, i64, Option)>( + "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? ORDER BY timestamp", ) .bind(session_id) .fetch_all(&self.pool) .await?; let mut messages = Vec::new(); - for (idx, (row_id, role_str, content_json, created_timestamp, metadata_json)) in + for (idx, (role_str, content_json, created_timestamp, metadata_json)) in rows.into_iter().enumerate() { let role = match role_str.as_str() { @@ -984,7 +998,6 @@ impl SessionStorage { let mut message = Message::new(role, created_timestamp, content); message.metadata = metadata; - message.row_id = Some(row_id); message = message.with_id(format!("msg_{}_{}", session_id, idx)); messages.push(message); } @@ -1160,52 +1173,13 @@ impl SessionStorage { self.get_session(&session.id, true).await } - async fn fork_session_at_message( - &self, - session_id: &str, - message_row_id: i64, - new_message_content: String, - ) -> Result { - use crate::conversation::message::MessageContent; - - let original_session = self.get_session(session_id, false).await?; - let rows = sqlx::query_as::<_, (String, String, i64, Option)>( - "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? AND id < ? ORDER BY id", - ) - .bind(session_id) - .bind(message_row_id) - .fetch_all(&self.pool) - .await?; - - let mut messages = Vec::new(); - for (role_str, content_json, created_timestamp, metadata_json) in rows { - let role = match role_str.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - _ => continue, - }; - - let content = serde_json::from_str(&content_json)?; - let metadata = metadata_json - .and_then(|json| serde_json::from_str(&json).ok()) - .unwrap_or_default(); - - let mut message = Message::new(role, created_timestamp, content); - message.metadata = metadata; - messages.push(message); - } - - let new_user_message = Message::new( - Role::User, - chrono::Utc::now().timestamp_millis(), - vec![MessageContent::text(&new_message_content)], - ); - messages.push(new_user_message); + async fn copy_session(&self, session_id: &str, new_name: String) -> Result { + let original_session = self.get_session(session_id, true).await?; let new_session = self .create_session( original_session.working_dir.clone(), - format!("{} (edited)", original_session.name), + new_name, original_session.session_type, ) .await?; @@ -1218,57 +1192,68 @@ impl SessionStorage { self.apply_update(builder).await?; - let conversation = Conversation::new_unvalidated(messages); - self.replace_conversation(&new_session.id, &conversation) - .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 edit_message_in_place( + 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 fork_session_at_message( &self, session_id: &str, - message_row_id: i64, + timestamp: i64, new_message_content: String, ) -> Result { use crate::conversation::message::MessageContent; - let rows = sqlx::query_as::<_, (String, String, i64, Option)>( - "SELECT role, content_json, created_timestamp, metadata_json FROM messages WHERE session_id = ? AND id < ? ORDER BY id", - ) - .bind(session_id) - .bind(message_row_id) - .fetch_all(&self.pool) + let original_session = self.get_session(session_id, true).await?; + + let new_session = self + .copy_session(session_id, format!("{} (edited)", original_session.name)) .await?; - let mut messages = Vec::new(); - for (role_str, content_json, created_timestamp, metadata_json) in rows { - let role = match role_str.as_str() { - "user" => Role::User, - "assistant" => Role::Assistant, - _ => continue, - }; + self.truncate_conversation(&new_session.id, timestamp) + .await?; - let content = serde_json::from_str(&content_json)?; - let metadata = metadata_json - .and_then(|json| serde_json::from_str(&json).ok()) - .unwrap_or_default(); + let new_user_message = Message::new( + Role::User, + chrono::Utc::now().timestamp_millis(), + vec![MessageContent::text(&new_message_content)], + ); - let mut message = Message::new(role, created_timestamp, content); - message.metadata = metadata; - messages.push(message); - } + self.add_message(&new_session.id, &new_user_message).await?; + self.get_session(&new_session.id, true).await + } + + async fn edit_message_in_place( + &self, + session_id: &str, + timestamp: i64, + new_message_content: String, + ) -> Result { + use crate::conversation::message::MessageContent; + + self.truncate_conversation(session_id, timestamp).await?; let new_user_message = Message::new( Role::User, chrono::Utc::now().timestamp_millis(), vec![MessageContent::text(&new_message_content)], ); - messages.push(new_user_message); - - let conversation = Conversation::new_unvalidated(messages); - self.replace_conversation(session_id, &conversation).await?; + self.add_message(session_id, &new_user_message).await?; self.get_session(session_id, true).await } @@ -1328,7 +1313,6 @@ mod tests { &session.id, &Message { id: None, - row_id: None, role: Role::User, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text("hello world")], @@ -1343,7 +1327,6 @@ mod tests { &session.id, &Message { id: None, - row_id: None, role: Role::Assistant, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text("sup world?")], @@ -1437,7 +1420,6 @@ mod tests { &original.id, &Message { id: None, - row_id: None, role: Role::User, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text(USER_MESSAGE)], @@ -1452,7 +1434,6 @@ mod tests { &original.id, &Message { id: None, - row_id: None, role: Role::Assistant, created: chrono::Utc::now().timestamp_millis(), content: vec![MessageContent::text(ASSISTANT_MESSAGE)], diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index d1267b6a6747..b95f2571c656 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -2512,19 +2512,19 @@ "EditMessageRequest": { "type": "object", "required": [ - "messageRowId", + "timestamp", "newContent" ], "properties": { "editType": { "$ref": "#/components/schemas/EditType" }, - "messageRowId": { - "type": "integer", - "format": "int64" - }, "newContent": { "type": "string" + }, + "timestamp": { + "type": "integer", + "format": "int64" } } }, @@ -3226,11 +3226,6 @@ }, "role": { "$ref": "#/components/schemas/Role" - }, - "rowId": { - "type": "integer", - "format": "int64", - "nullable": true } } }, diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index a46cece49b9b..16b41834c4b7 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -124,8 +124,8 @@ export type DeleteRecipeRequest = { export type EditMessageRequest = { editType?: EditType; - messageRowId: number; newContent: string; + timestamp: number; }; export type EditMessageResponse = { @@ -363,7 +363,6 @@ export type Message = { id?: string | null; metadata: MessageMetadata; role: Role; - rowId?: number | null; }; /** diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index 152ec641beb0..cea1730dee9d 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -185,28 +185,10 @@ export function useChatStream({ window.dispatchEvent(new CustomEvent('message-stream-finished')); } - // Reload the session to get auto generated/incremented rowIds for all messages - try { - const { getSession } = await import('../api'); - const response = await getSession({ - path: { session_id: sessionId }, - throwOnError: true, - }); - - if (response.data?.conversation) { - window.electron.logInfo( - `Reloaded session with ${response.data.conversation.length} messages with rowIds` - ); - updateMessages(response.data.conversation); - } - } catch (reloadError) { - window.electron.logInfo(`Failed to reload session: ${errorMessage(reloadError)}`); - } - setChatState(ChatState.Idle); onStreamFinish(); }, - [onStreamFinish, sessionId, updateMessages] + [onStreamFinish, sessionId] ); // Load session on mount or sessionId change @@ -382,18 +364,12 @@ export function useChatStream({ throw new Error(`Message with id ${messageId} not found in current messages`); } - if (!message.rowId) { - throw new Error( - `Message ${messageId} is missing rowId. This might be an old session - please reload the page and try again.` - ); - } - const response = await editMessage({ path: { session_id: sessionId, }, body: { - messageRowId: message.rowId, + timestamp: message.created, newContent, editType, }, From c3f0ae5a96bfc6781ca9685d62f6ab51dbdce4c1 Mon Sep 17 00:00:00 2001 From: Zane Staggs Date: Wed, 19 Nov 2025 20:08:59 -0700 Subject: [PATCH 3/5] change to not use skip_add_user_message --- crates/goose-server/src/routes/session.rs | 70 +++++++-------------- crates/goose/src/session/session_manager.rs | 69 -------------------- ui/desktop/openapi.json | 25 ++------ ui/desktop/src/api/types.gen.ts | 14 +---- ui/desktop/src/components/BaseChat.tsx | 14 +++-- ui/desktop/src/hooks/useChatStream.ts | 48 +++++++------- 6 files changed, 64 insertions(+), 176 deletions(-) diff --git a/crates/goose-server/src/routes/session.rs b/crates/goose-server/src/routes/session.rs index 259900b0be37..39cc203970e3 100644 --- a/crates/goose-server/src/routes/session.rs +++ b/crates/goose-server/src/routes/session.rs @@ -9,7 +9,6 @@ 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}; @@ -61,7 +60,6 @@ pub enum EditType { #[serde(rename_all = "camelCase")] pub struct EditMessageRequest { timestamp: i64, - new_content: String, #[serde(default = "default_edit_type")] edit_type: EditType, } @@ -73,11 +71,7 @@ fn default_edit_type() -> EditType { #[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, - /// Conversation (either in new session for fork, or updated current session for edit) - conversation: Vec, + session_id: String, } const MAX_NAME_LENGTH: usize = 200; @@ -346,8 +340,8 @@ async fn import_session( ("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 = 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") @@ -361,54 +355,36 @@ async fn edit_message( Path(session_id): Path, Json(request): Json, ) -> Result, 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.timestamp, - request.new_content, - ) - .await - .map_err(|e| { - tracing::error!("Failed to fork session: {}", e); - StatusCode::INTERNAL_SERVER_ERROR - })?; + 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 + })?; - let conversation = new_session.conversation.ok_or_else(|| { - tracing::error!("Forked session has no conversation"); - 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 { - new_session_id: Some(new_session.id), - conversation: conversation.messages().to_vec(), + session_id: new_session.id, })) } EditType::Edit => { - let updated_session = SessionManager::edit_message_in_place( - &session_id, - request.timestamp, - 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 - })?; + 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 { - new_session_id: None, - conversation: conversation.messages().to_vec(), + session_id: session_id.clone(), })) } } diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 885ad97d19c0..b1ea049f3472 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -304,28 +304,6 @@ impl SessionManager { .await } - pub async fn fork_session_at_message( - session_id: &str, - timestamp: i64, - new_message_content: String, - ) -> Result { - Self::instance() - .await? - .fork_session_at_message(session_id, timestamp, new_message_content) - .await - } - - pub async fn edit_message_in_place( - session_id: &str, - timestamp: i64, - new_message_content: String, - ) -> Result { - Self::instance() - .await? - .edit_message_in_place(session_id, timestamp, new_message_content) - .await - } - pub async fn maybe_update_name(id: &str, provider: Arc) -> Result<()> { let session = Self::get_session(id, true).await?; @@ -1210,53 +1188,6 @@ impl SessionStorage { Ok(()) } - async fn fork_session_at_message( - &self, - session_id: &str, - timestamp: i64, - new_message_content: String, - ) -> Result { - use crate::conversation::message::MessageContent; - - let original_session = self.get_session(session_id, true).await?; - - let new_session = self - .copy_session(session_id, format!("{} (edited)", original_session.name)) - .await?; - - self.truncate_conversation(&new_session.id, timestamp) - .await?; - - let new_user_message = Message::new( - Role::User, - chrono::Utc::now().timestamp_millis(), - vec![MessageContent::text(&new_message_content)], - ); - - self.add_message(&new_session.id, &new_user_message).await?; - self.get_session(&new_session.id, true).await - } - - async fn edit_message_in_place( - &self, - session_id: &str, - timestamp: i64, - new_message_content: String, - ) -> Result { - use crate::conversation::message::MessageContent; - - self.truncate_conversation(session_id, timestamp).await?; - - let new_user_message = Message::new( - Role::User, - chrono::Utc::now().timestamp_millis(), - vec![MessageContent::text(&new_message_content)], - ); - - self.add_message(session_id, &new_user_message).await?; - self.get_session(session_id, true).await - } - async fn search_chat_history( &self, query: &str, diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index b95f2571c656..9700218b9cd9 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -1945,7 +1945,7 @@ }, "responses": { "200": { - "description": "Message edited successfully", + "description": "Session prepared for editing - frontend should submit the edited message", "content": { "application/json": { "schema": { @@ -1955,7 +1955,7 @@ } }, "400": { - "description": "Bad request - Invalid message ID or empty content" + "description": "Bad request - Invalid message timestamp" }, "401": { "description": "Unauthorized - Invalid or missing API key" @@ -2512,16 +2512,12 @@ "EditMessageRequest": { "type": "object", "required": [ - "timestamp", - "newContent" + "timestamp" ], "properties": { "editType": { "$ref": "#/components/schemas/EditType" }, - "newContent": { - "type": "string" - }, "timestamp": { "type": "integer", "format": "int64" @@ -2531,20 +2527,11 @@ "EditMessageResponse": { "type": "object", "required": [ - "conversation" + "sessionId" ], "properties": { - "conversation": { - "type": "array", - "items": { - "$ref": "#/components/schemas/Message" - }, - "description": "Conversation (either in new session for fork, or updated current session for edit)" - }, - "newSessionId": { - "type": "string", - "description": "New session ID created from the fork (only present for fork edit_type)", - "nullable": true + "sessionId": { + "type": "string" } } }, diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 16b41834c4b7..564f7b6fb07a 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -124,19 +124,11 @@ export type DeleteRecipeRequest = { export type EditMessageRequest = { editType?: EditType; - newContent: string; timestamp: number; }; export type EditMessageResponse = { - /** - * Conversation (either in new session for fork, or updated current session for edit) - */ - conversation: Array; - /** - * New session ID created from the fork (only present for fork edit_type) - */ - newSessionId?: string | null; + sessionId: string; }; export type EditType = 'fork' | 'edit'; @@ -2451,7 +2443,7 @@ export type EditMessageData = { export type EditMessageErrors = { /** - * Bad request - Invalid message ID or empty content + * Bad request - Invalid message timestamp */ 400: unknown; /** @@ -2470,7 +2462,7 @@ export type EditMessageErrors = { export type EditMessageResponses = { /** - * Message edited successfully + * Session prepared for editing - frontend should submit the edited message */ 200: EditMessageResponse; }; diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index 8b42946be7e5..69d3817e2b72 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -117,22 +117,20 @@ function BaseChatContent({ onSessionLoaded: handleSessionLoaded, }); - // Handle auto-submission when session is loaded and we have an initial message useEffect(() => { + const shouldStartAgent = searchParams.get('shouldStartAgent') === 'true'; if (sessionLoaded && initialMessage && !hasSubmittedInitialMessage) { setHasSubmittedInitialMessage(true); handleSubmit(initialMessage); + return; } - }, [sessionLoaded, initialMessage, hasSubmittedInitialMessage, handleSubmit]); - - useEffect(() => { - const shouldStartAgent = searchParams.get('shouldStartAgent') === 'true'; if ( sessionLoaded && shouldStartAgent && messages.length > 0 && !hasTriggeredAgentForFork && + !initialMessage && chatState === ChatState.Idle ) { setHasTriggeredAgentForFork(true); @@ -140,6 +138,8 @@ function BaseChatContent({ } }, [ sessionLoaded, + initialMessage, + hasSubmittedInitialMessage, searchParams, messages.length, hasTriggeredAgentForFork, @@ -240,8 +240,9 @@ function BaseChatContent({ const customEvent = event as CustomEvent<{ newSessionId: string; shouldStartAgent?: boolean; + editedMessage?: string; }>; - const { newSessionId, shouldStartAgent } = customEvent.detail; + const { newSessionId, shouldStartAgent, editedMessage } = customEvent.detail; const params = new URLSearchParams(); params.set('resumeSessionId', newSessionId); @@ -252,6 +253,7 @@ function BaseChatContent({ navigate(`/pair?${params.toString()}`, { state: { disableAnimation: true, + initialMessage: editedMessage, }, }); }; diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index cea1730dee9d..bf3b5d1b9dab 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -370,37 +370,37 @@ export function useChatStream({ }, body: { timestamp: message.created, - newContent, editType, }, throwOnError: true, }); + const targetSessionId = response.data?.sessionId; + if (!targetSessionId) { + throw new Error('No session ID returned from edit_message'); + } + if (editType === 'fork') { - if (response.data?.newSessionId && response.data?.conversation) { - const newSessionId = response.data.newSessionId; - const event = new CustomEvent('session-forked', { - detail: { - newSessionId, - shouldStartAgent: true, - }, - }); - window.dispatchEvent(event); - window.electron.logInfo(`Dispatched session-forked event`); - } else { - window.electron.logInfo( - `Unexpected fork response format: ${JSON.stringify(response.data)}` - ); - } + const event = new CustomEvent('session-forked', { + detail: { + newSessionId: targetSessionId, + shouldStartAgent: true, + editedMessage: newContent, + }, + }); + window.dispatchEvent(event); + window.electron.logInfo(`Dispatched session-forked event for session ${targetSessionId}`); } else { - if (response.data?.conversation) { - updateMessages(response.data.conversation); - await handleSubmit('', { isContinuation: true }); - } else { - window.electron.logInfo( - `Unexpected edit response format: ${JSON.stringify(response.data)}` - ); + const { getSession } = await import('../api'); + const sessionResponse = await getSession({ + path: { session_id: targetSessionId }, + throwOnError: true, + }); + + if (sessionResponse.data?.conversation) { + updateMessages(sessionResponse.data.conversation); } + await handleSubmit(newContent); } } catch (error) { const errorMsg = errorMessage(error); @@ -412,7 +412,7 @@ export function useChatStream({ }); } }, - [sessionId, updateMessages, handleSubmit] + [sessionId, handleSubmit, updateMessages] ); const cached = resultsCache.get(sessionId); From e5aa47e7ea179969031d5cc2540d18983264f204 Mon Sep 17 00:00:00 2001 From: Zane Staggs Date: Wed, 19 Nov 2025 20:39:35 -0700 Subject: [PATCH 4/5] change to not use skip_add_user_message --- crates/goose-server/src/routes/reply.rs | 6 +---- crates/goose/src/agents/agent.rs | 16 +----------- ui/desktop/openapi.json | 3 --- ui/desktop/src/api/types.gen.ts | 1 - ui/desktop/src/components/BaseChat.tsx | 2 +- ui/desktop/src/hooks/useChatStream.ts | 34 ++++++++++++------------- 6 files changed, 19 insertions(+), 43 deletions(-) diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index c763f0f6b0bc..5a340fa3d947 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -85,8 +85,6 @@ pub struct ChatRequest { session_id: String, recipe_name: Option, recipe_version: Option, - #[serde(default)] - skip_add_user_message: bool, } pub struct SseResponse { @@ -302,11 +300,10 @@ pub async fn reply( }; let mut stream = match agent - .reply_with_options( + .reply( user_message.clone(), session_config, Some(task_cancel.clone()), - request.skip_add_user_message, ) .await { @@ -539,7 +536,6 @@ mod tests { session_id: "test-session".to_string(), recipe_name: None, recipe_version: None, - skip_add_user_message: false, }) .unwrap(), )) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 2bbff1c432aa..004ad745580f 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -776,18 +776,6 @@ impl Agent { user_message: Message, session_config: SessionConfig, cancel_token: Option, - ) -> Result>> { - 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, - skip_add_message: bool, ) -> Result>> { let is_manual_compact = user_message.content.iter().any(|c| { if let MessageContent::Text(text) = c { @@ -797,9 +785,7 @@ impl Agent { } }); - if !skip_add_message { - SessionManager::add_message(&session_config.id, &user_message).await?; - } + SessionManager::add_message(&session_config.id, &user_message).await?; let session = SessionManager::get_session(&session_config.id, true).await?; let conversation = session diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index 9700218b9cd9..46570b5e6bfa 100644 --- a/ui/desktop/openapi.json +++ b/ui/desktop/openapi.json @@ -2244,9 +2244,6 @@ }, "session_id": { "type": "string" - }, - "skip_add_user_message": { - "type": "boolean" } } }, diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 564f7b6fb07a..f5d0ee9c9c02 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -30,7 +30,6 @@ export type ChatRequest = { recipe_name?: string | null; recipe_version?: string | null; session_id: string; - skip_add_user_message?: boolean; }; export type CheckProviderRequest = { diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index 69d3817e2b72..c4f887758421 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -134,7 +134,7 @@ function BaseChatContent({ chatState === ChatState.Idle ) { setHasTriggeredAgentForFork(true); - handleSubmit('', { isContinuation: true }); + handleSubmit(''); } }, [ sessionLoaded, diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index bf3b5d1b9dab..6273d672ff9c 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -29,7 +29,7 @@ interface UseChatStreamReturn { session?: Session; messages: Message[]; chatState: ChatState; - handleSubmit: (userMessage: string, options?: { isContinuation?: boolean }) => Promise; + handleSubmit: (userMessage: string) => Promise; setRecipeUserParams: (values: Record) => Promise; stopStreaming: () => void; sessionLoadError?: string; @@ -244,38 +244,37 @@ export function useChatStream({ }, [sessionId, updateMessages, onSessionLoaded]); const handleSubmit = useCallback( - async (userMessage: string, options?: { isContinuation?: boolean }) => { + async (userMessage: string) => { // Guard: Don't submit if session hasn't been loaded yet if (!session || chatState === ChatState.LoadingConversation) { return; } - // Check if this is a continuation request (agent should respond to existing messages) - const isContinuationRequest = options?.isContinuation || false; - if (!isContinuationRequest && !userMessage.trim()) { + const hasExistingMessages = messagesRef.current.length > 0; + const hasNewMessage = userMessage.trim().length > 0; + + // Don't submit if there's no message and no conversation to continue + if (!hasNewMessage && !hasExistingMessages) { return; } - if (messagesRef.current.length === 0 && !isContinuationRequest) { + // Emit session-created event for first message in a new session + if (!hasExistingMessages && hasNewMessage) { window.dispatchEvent(new CustomEvent('session-created')); } - // For continuation requests, don't add a new user message - const currentMessages = isContinuationRequest - ? [...messagesRef.current] - : [...messagesRef.current, createUserMessage(userMessage)]; - - window.electron.logInfo( - `handleSubmit: isContinuationRequest=${isContinuationRequest}, messagesRef.current.length=${messagesRef.current.length}, currentMessages.length=${currentMessages.length}` - ); + // Build message list: add new message if provided, otherwise continue with existing + const currentMessages = hasNewMessage + ? [...messagesRef.current, createUserMessage(userMessage)] + : [...messagesRef.current]; - // Don't update messages if it's a continuation request (messages are already correct) - if (!isContinuationRequest) { + // Update UI with new message before streaming + if (hasNewMessage) { updateMessages(currentMessages); } + setChatState(ChatState.Streaming); setNotifications([]); - abortControllerRef.current = new AbortController(); try { @@ -283,7 +282,6 @@ export function useChatStream({ body: { session_id: sessionId, messages: currentMessages, - skip_add_user_message: isContinuationRequest, // Tell backend not to add the message again }, throwOnError: true, signal: abortControllerRef.current.signal, From 5ad9b0df97416d6c6a1e4e7b79f4b4d39ecc6af0 Mon Sep 17 00:00:00 2001 From: Zane Staggs Date: Fri, 21 Nov 2025 09:47:33 -0700 Subject: [PATCH 5/5] simplify handlesubmit logic --- ui/desktop/src/components/BaseChat.tsx | 51 ++++++++------------------ 1 file changed, 15 insertions(+), 36 deletions(-) diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index c4f887758421..3eddac30a838 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -85,21 +85,14 @@ function BaseChatContent({ const onStreamFinish = useCallback(() => {}, []); - const [sessionLoaded, setSessionLoaded] = useState(false); - const [hasSubmittedInitialMessage, setHasSubmittedInitialMessage] = useState(false); - const [hasTriggeredAgentForFork, setHasTriggeredAgentForFork] = useState(false); const [isCreateRecipeModalOpen, setIsCreateRecipeModalOpen] = useState(false); + const hasAutoSubmittedRef = useRef(false); + // Reset auto-submit flag when session changes useEffect(() => { - setSessionLoaded(false); - setHasSubmittedInitialMessage(false); - setHasTriggeredAgentForFork(false); + hasAutoSubmittedRef.current = false; }, [sessionId]); - const handleSessionLoaded = useCallback(() => { - setSessionLoaded(true); - }, []); - const { session, messages, @@ -114,39 +107,25 @@ function BaseChatContent({ } = useChatStream({ sessionId, onStreamFinish, - onSessionLoaded: handleSessionLoaded, }); useEffect(() => { - const shouldStartAgent = searchParams.get('shouldStartAgent') === 'true'; - if (sessionLoaded && initialMessage && !hasSubmittedInitialMessage) { - setHasSubmittedInitialMessage(true); - handleSubmit(initialMessage); + if (!session || hasAutoSubmittedRef.current) { return; } - if ( - sessionLoaded && - shouldStartAgent && - messages.length > 0 && - !hasTriggeredAgentForFork && - !initialMessage && - chatState === ChatState.Idle - ) { - setHasTriggeredAgentForFork(true); + const shouldStartAgent = searchParams.get('shouldStartAgent') === 'true'; + + if (initialMessage) { + // Submit the initial message (e.g., from fork) + hasAutoSubmittedRef.current = true; + handleSubmit(initialMessage); + } else if (shouldStartAgent) { + // Trigger agent to continue with existing conversation + hasAutoSubmittedRef.current = true; handleSubmit(''); } - }, [ - sessionLoaded, - initialMessage, - hasSubmittedInitialMessage, - searchParams, - messages.length, - hasTriggeredAgentForFork, - chatState, - handleSubmit, - messages, - ]); + }, [session, initialMessage, searchParams, handleSubmit]); const handleFormSubmit = (e: React.FormEvent) => { const customEvent = e as unknown as CustomEvent; @@ -318,7 +297,7 @@ function BaseChatContent({ } const initialPrompt = - (initialMessage && !hasSubmittedInitialMessage ? initialMessage : '') || recipePrompt; + (initialMessage && !hasAutoSubmittedRef.current ? initialMessage : '') || recipePrompt; if (sessionLoadError) { return (