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/session.rs b/crates/goose-server/src/routes/session.rs index 67a5bf388d59..39cc203970e3 100644 --- a/crates/goose-server/src/routes/session.rs +++ b/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, + Json(request): Json, +) -> Result, 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) -> Router { Router::new() .route("/sessions", get(list_sessions)) @@ -320,5 +403,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..004ad745580f 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -786,8 +786,8 @@ impl Agent { }); SessionManager::add_message(&session_config.id, &user_message).await?; - let session = SessionManager::get_session(&session_config.id, true).await?; + let session = SessionManager::get_session(&session_config.id, true).await?; let conversation = session .conversation .clone() diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 564911acbc1e..b1ea049f3472 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -290,6 +290,20 @@ 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 maybe_update_name(id: &str, provider: Arc) -> Result<()> { let session = Self::get_session(id, true).await?; @@ -1137,6 +1151,43 @@ impl SessionStorage { self.get_session(&session.id, true).await } + 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(), + 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, diff --git a/ui/desktop/openapi.json b/ui/desktop/openapi.json index e7e4f8a4f6da..46570b5e6bfa 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": "Session prepared for editing - frontend should submit the edited message", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/EditMessageResponse" + } + } + } + }, + "400": { + "description": "Bad request - Invalid message timestamp" + }, + "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": [ @@ -2448,6 +2506,39 @@ } } }, + "EditMessageRequest": { + "type": "object", + "required": [ + "timestamp" + ], + "properties": { + "editType": { + "$ref": "#/components/schemas/EditType" + }, + "timestamp": { + "type": "integer", + "format": "int64" + } + } + }, + "EditMessageResponse": { + "type": "object", + "required": [ + "sessionId" + ], + "properties": { + "sessionId": { + "type": "string" + } + } + }, + "EditType": { + "type": "string", + "enum": [ + "fork", + "edit" + ] + }, "EmbeddedResource": { "type": "object", "required": [ 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..f5d0ee9c9c02 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -121,6 +121,17 @@ export type DeleteRecipeRequest = { id: string; }; +export type EditMessageRequest = { + editType?: EditType; + timestamp: number; +}; + +export type EditMessageResponse = { + sessionId: string; +}; + +export type EditType = 'fork' | 'edit'; + export type EmbeddedResource = { _meta?: { [key: string]: unknown; @@ -2417,6 +2428,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 timestamp + */ + 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 = { + /** + * Session prepared for editing - frontend should submit the edited message + */ + 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..3eddac30a838 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; @@ -83,19 +85,14 @@ function BaseChatContent({ const onStreamFinish = useCallback(() => {}, []); - const [sessionLoaded, setSessionLoaded] = useState(false); - const [hasSubmittedInitialMessage, setHasSubmittedInitialMessage] = useState(false); const [isCreateRecipeModalOpen, setIsCreateRecipeModalOpen] = useState(false); + const hasAutoSubmittedRef = useRef(false); + // Reset auto-submit flag when session changes useEffect(() => { - setSessionLoaded(false); - setHasSubmittedInitialMessage(false); + hasAutoSubmittedRef.current = false; }, [sessionId]); - const handleSessionLoaded = useCallback(() => { - setSessionLoaded(true); - }, []); - const { session, messages, @@ -106,19 +103,29 @@ function BaseChatContent({ setRecipeUserParams, tokenState, notifications, + onMessageUpdate, } = useChatStream({ sessionId, onStreamFinish, - onSessionLoaded: handleSessionLoaded, }); - // Handle auto-submission when session is loaded and we have an initial message useEffect(() => { - if (sessionLoaded && initialMessage && !hasSubmittedInitialMessage) { - setHasSubmittedInitialMessage(true); + if (!session || hasAutoSubmittedRef.current) { + return; + } + + 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, handleSubmit]); + }, [session, initialMessage, searchParams, handleSubmit]); const handleFormSubmit = (e: React.FormEvent) => { const customEvent = e as unknown as CustomEvent; @@ -207,6 +214,36 @@ function BaseChatContent({ return () => window.removeEventListener('make-agent-from-chat', handleMakeAgent); }, []); + useEffect(() => { + const handleSessionForked = (event: Event) => { + const customEvent = event as CustomEvent<{ + newSessionId: string; + shouldStartAgent?: boolean; + editedMessage?: string; + }>; + const { newSessionId, shouldStartAgent, editedMessage } = customEvent.detail; + + const params = new URLSearchParams(); + params.set('resumeSessionId', newSessionId); + if (shouldStartAgent) { + params.set('shouldStartAgent', 'true'); + } + + navigate(`/pair?${params.toString()}`, { + state: { + disableAnimation: true, + initialMessage: editedMessage, + }, + }); + }; + + 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 +271,7 @@ function BaseChatContent({ isUserMessage={(m: Message) => m.role === 'user'} isStreamingMessage={chatState !== ChatState.Idle} onRenderingComplete={handleRenderingComplete} + onMessageUpdate={onMessageUpdate} /> ); @@ -259,7 +297,7 @@ function BaseChatContent({ } const initialPrompt = - (initialMessage && !hasSubmittedInitialMessage ? initialMessage : '') || recipePrompt; + (initialMessage && !hasAutoSubmittedRef.current ? initialMessage : '') || recipePrompt; if (sessionLoadError) { return ( 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..6273d672ff9c 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -35,6 +35,11 @@ interface UseChatStreamReturn { 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); } @@ -245,19 +250,31 @@ export function useChatStream({ return; } - if (!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) { + // Emit session-created event for first message in a new session + if (!hasExistingMessages && hasNewMessage) { window.dispatchEvent(new CustomEvent('session-created')); } - const currentMessages = [...messagesRef.current, createUserMessage(userMessage)]; - updateMessages(currentMessages); + // Build message list: add new message if provided, otherwise continue with existing + const currentMessages = hasNewMessage + ? [...messagesRef.current, createUserMessage(userMessage)] + : [...messagesRef.current]; + + // Update UI with new message before streaming + if (hasNewMessage) { + updateMessages(currentMessages); + } + setChatState(ChatState.Streaming); setNotifications([]); - abortControllerRef.current = new AbortController(); try { @@ -335,6 +352,67 @@ 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`); + } + + const response = await editMessage({ + path: { + session_id: sessionId, + }, + body: { + timestamp: message.created, + editType, + }, + throwOnError: true, + }); + + const targetSessionId = response.data?.sessionId; + if (!targetSessionId) { + throw new Error('No session ID returned from edit_message'); + } + + if (editType === 'fork') { + 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 { + 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); + console.error('Failed to edit message:', error); + const { toastError } = await import('../toasts'); + toastError({ + title: 'Failed to edit message', + msg: errorMsg, + }); + } + }, + [sessionId, handleSubmit, updateMessages] + ); + const cached = resultsCache.get(sessionId); const maybe_cached_messages = session ? messages : cached?.messages || []; const maybe_cached_session = session ?? cached?.session; @@ -349,5 +427,6 @@ export function useChatStream({ setRecipeUserParams, tokenState, notifications, + onMessageUpdate, }; }