Add message edit and regenerate from any point in conversation

Backend: add truncate_to_user_message() on ContextManager, POST
/api/truncate/{session_id} endpoint, and session_manager.truncate()
helper to slice conversation history to before a given user message.

Frontend: add editAndRegenerate callback in useAgentChat that
truncates backend + frontend messages then resubmits via existing
chat flow. UserMessage gains an inline edit mode with pencil icon
on hover, TextField for editing, Enter/Escape to confirm/cancel.

agent/context_manager/manager.py

@@ -243,6 +243,25 @@ class ContextManager:
 
         return False
 
+    def truncate_to_user_message(self, user_message_index: int) -> bool:
+        """Truncate history to just before the Nth user message (0-indexed).
+
+        Removes that user message and everything after it.
+        System message (index 0) is never removed.
+
+        Returns True if the target user message was found and removed.
+        """
+        count = 0
+        for i, msg in enumerate(self.items):
+            if i == 0:
+                continue  # skip system message
+            if getattr(msg, "role", None) == "user":
+                if count == user_message_index:
+                    self.items = self.items[:i]
+                    return True
+                count += 1
+        return False
+
     async def compact(
         self, model_name: str, tool_specs: list[dict] | None = None
     ) -> None:

backend/models.py

@@ -54,6 +54,12 @@ class SubmitRequest(BaseModel):
     text: str
 
 
+class TruncateRequest(BaseModel):
+    """Request to truncate conversation history to before a specific user message."""
+
+    user_message_index: int
+
+
 class SessionResponse(BaseModel):
     """Response when creating a new session."""
 

backend/routes/agent.py

@@ -26,6 +26,7 @@ from models import (
     SessionInfo,
     SessionResponse,
     SubmitRequest,
+    TruncateRequest,
 )
 from session_manager import MAX_SESSIONS, SessionCapacityError, session_manager
 
@@ -438,6 +439,18 @@ async def undo_session(session_id: str, user: dict = Depends(get_current_user))
     return {"status": "undo_requested", "session_id": session_id}
 
 
+@router.post("/truncate/{session_id}")
+async def truncate_session(
+    session_id: str, body: TruncateRequest, user: dict = Depends(get_current_user)
+) -> dict:
+    """Truncate conversation to before a specific user message."""
+    _check_session_access(session_id, user)
+    success = await session_manager.truncate(session_id, body.user_message_index)
+    if not success:
+        raise HTTPException(status_code=404, detail="Session not found, inactive, or message index out of range")
+    return {"status": "truncated", "session_id": session_id}
+
+
 @router.post("/compact/{session_id}")
 async def compact_session(
     session_id: str, user: dict = Depends(get_current_user)

backend/session_manager.py

@@ -319,6 +319,14 @@ class SessionManager:
         operation = Operation(op_type=OpType.UNDO)
         return await self.submit(session_id, operation)
 
+    async def truncate(self, session_id: str, user_message_index: int) -> bool:
+        """Truncate conversation to before a specific user message (direct, no queue)."""
+        async with self._lock:
+            agent_session = self.sessions.get(session_id)
+        if not agent_session or not agent_session.is_active:
+            return False
+        return agent_session.session.context_manager.truncate_to_user_message(user_message_index)
+
     async def compact(self, session_id: str) -> bool:
         """Compact context in a session."""
         operation = Operation(op_type=OpType.COMPACT)

frontend/src/components/Chat/MessageBubble.tsx

@@ -6,6 +6,7 @@ interface MessageBubbleProps {
   message: UIMessage;
   isLastTurn?: boolean;
   onUndoTurn?: () => void;
+  onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise<void>;
   isProcessing?: boolean;
   isStreaming?: boolean;
   approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
@@ -15,6 +16,7 @@ export default function MessageBubble({
   message,
   isLastTurn = false,
   onUndoTurn,
+  onEditAndRegenerate,
   isProcessing = false,
   isStreaming = false,
   approveTools,
@@ -25,6 +27,7 @@ export default function MessageBubble({
         message={message}
         isLastTurn={isLastTurn}
         onUndoTurn={onUndoTurn}
+        onEditAndRegenerate={onEditAndRegenerate}
         isProcessing={isProcessing}
       />
     );

frontend/src/components/Chat/MessageList.tsx

@@ -10,6 +10,7 @@ interface MessageListProps {
   isProcessing: boolean;
   approveTools: (approvals: Array<{ tool_call_id: string; approved: boolean; feedback?: string | null }>) => Promise<boolean>;
   onUndoLastTurn: () => void | Promise<void>;
+  onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise<void>;
 }
 
 function getGreeting(): string {
@@ -56,7 +57,7 @@ function WelcomeGreeting() {
   );
 }
 
-export default function MessageList({ messages, isProcessing, approveTools, onUndoLastTurn }: MessageListProps) {
+export default function MessageList({ messages, isProcessing, approveTools, onUndoLastTurn, onEditAndRegenerate }: MessageListProps) {
   const scrollContainerRef = useRef<HTMLDivElement>(null);
   const stickToBottom = useRef(true);
 
@@ -135,6 +136,7 @@ export default function MessageList({ messages, isProcessing, approveTools, onUn
               message={msg}
               isLastTurn={msg.id === lastUserMsgId}
               onUndoTurn={onUndoLastTurn}
+              onEditAndRegenerate={onEditAndRegenerate}
               isProcessing={isProcessing}
               isStreaming={isProcessing && msg.id === lastAssistantId}
               approveTools={approveTools}

frontend/src/components/Chat/UserMessage.tsx

@@ -1,5 +1,8 @@
-import { Box, Stack, Typography, IconButton, Tooltip } from '@mui/material';
+import { useState, useRef, useEffect } from 'react';
+import { Box, Stack, Typography, IconButton, Tooltip, TextField } from '@mui/material';
 import CloseIcon from '@mui/icons-material/Close';
+import EditIcon from '@mui/icons-material/Edit';
+import CheckIcon from '@mui/icons-material/Check';
 import type { UIMessage } from 'ai';
 import type { MessageMeta } from '@/types/agent';
 
@@ -7,6 +10,7 @@ interface UserMessageProps {
   message: UIMessage;
   isLastTurn?: boolean;
   onUndoTurn?: () => void;
+  onEditAndRegenerate?: (messageId: string, newText: string) => void | Promise<void>;
   isProcessing?: boolean;
 }
 
@@ -21,14 +25,57 @@ export default function UserMessage({
   message,
   isLastTurn = false,
   onUndoTurn,
+  onEditAndRegenerate,
   isProcessing = false,
 }: UserMessageProps) {
   const showUndo = isLastTurn && !isProcessing && !!onUndoTurn;
+  const showEdit = !isProcessing && !!onEditAndRegenerate;
   const text = extractText(message);
   const meta = message.metadata as MessageMeta | undefined;
   const timeStr = meta?.createdAt
     ? new Date(meta.createdAt).toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })
     : null;
+
+  const [isEditing, setIsEditing] = useState(false);
+  const [editText, setEditText] = useState(text);
+  const inputRef = useRef<HTMLTextAreaElement>(null);
+
+  useEffect(() => {
+    if (isEditing && inputRef.current) {
+      inputRef.current.focus();
+      inputRef.current.selectionStart = inputRef.current.value.length;
+    }
+  }, [isEditing]);
+
+  const handleStartEdit = () => {
+    setEditText(text);
+    setIsEditing(true);
+  };
+
+  const handleConfirmEdit = () => {
+    const trimmed = editText.trim();
+    if (!trimmed || trimmed === text) {
+      setIsEditing(false);
+      return;
+    }
+    setIsEditing(false);
+    onEditAndRegenerate?.(message.id, trimmed);
+  };
+
+  const handleCancelEdit = () => {
+    setIsEditing(false);
+    setEditText(text);
+  };
+
+  const handleKeyDown = (e: React.KeyboardEvent) => {
+    if (e.key === 'Enter' && !e.shiftKey) {
+      e.preventDefault();
+      handleConfirmEdit();
+    } else if (e.key === 'Escape') {
+      handleCancelEdit();
+    }
+  };
+
   return (
     <Stack
       direction="row"
@@ -36,35 +83,56 @@ export default function UserMessage({
       justifyContent="flex-end"
       alignItems="flex-start"
       sx={{
-        '& .undo-btn': {
+        '& .action-btn': {
           opacity: 0,
           transition: 'opacity 0.15s ease',
         },
-        '&:hover .undo-btn': {
+        '&:hover .action-btn': {
           opacity: 1,
         },
       }}
     >
-      {showUndo && (
-        <Box className="undo-btn" sx={{ display: 'flex', alignItems: 'center', mt: 0.75 }}>
-          <Tooltip title="Remove this turn" placement="left">
-            <IconButton
-              onClick={onUndoTurn}
-              size="small"
-              sx={{
-                width: 24,
-                height: 24,
-                color: 'var(--muted-text)',
-                '&:hover': {
-                  color: 'var(--accent-red)',
-                  bgcolor: 'rgba(244,67,54,0.08)',
-                },
-              }}
-            >
-              <CloseIcon sx={{ fontSize: 14 }} />
-            </IconButton>
-          </Tooltip>
-        </Box>
+      {!isEditing && (showUndo || showEdit) && (
+        <Stack className="action-btn" direction="row" spacing={0.25} sx={{ mt: 0.75 }}>
+          {showEdit && (
+            <Tooltip title="Edit & regenerate" placement="left">
+              <IconButton
+                onClick={handleStartEdit}
+                size="small"
+                sx={{
+                  width: 24,
+                  height: 24,
+                  color: 'var(--muted-text)',
+                  '&:hover': {
+                    color: 'var(--accent-yellow)',
+                    bgcolor: 'rgba(255,157,0,0.08)',
+                  },
+                }}
+              >
+                <EditIcon sx={{ fontSize: 14 }} />
+              </IconButton>
+            </Tooltip>
+          )}
+          {showUndo && (
+            <Tooltip title="Remove this turn" placement="left">
+              <IconButton
+                onClick={onUndoTurn}
+                size="small"
+                sx={{
+                  width: 24,
+                  height: 24,
+                  color: 'var(--muted-text)',
+                  '&:hover': {
+                    color: 'var(--accent-red)',
+                    bgcolor: 'rgba(244,67,54,0.08)',
+                  },
+                }}
+              >
+                <CloseIcon sx={{ fontSize: 14 }} />
+              </IconButton>
+            </Tooltip>
+          )}
+        </Stack>
       )}
 
       <Box
@@ -78,20 +146,66 @@ export default function UserMessage({
           border: '1px solid var(--border)',
         }}
       >
-        <Typography
-          variant="body1"
-          sx={{
-            fontSize: '0.925rem',
-            lineHeight: 1.65,
-            color: 'var(--text)',
-            whiteSpace: 'pre-wrap',
-            wordBreak: 'break-word',
-          }}
-        >
-          {text}
-        </Typography>
+        {isEditing ? (
+          <Stack spacing={1}>
+            <TextField
+              inputRef={inputRef}
+              multiline
+              fullWidth
+              value={editText}
+              onChange={(e) => setEditText(e.target.value)}
+              onKeyDown={handleKeyDown}
+              variant="outlined"
+              size="small"
+              sx={{
+                '& .MuiOutlinedInput-root': {
+                  fontFamily: 'inherit',
+                  fontSize: '0.925rem',
+                  lineHeight: 1.65,
+                  color: 'var(--text)',
+                  '& fieldset': { borderColor: 'var(--accent-yellow)', borderWidth: 1.5 },
+                  '&:hover fieldset': { borderColor: 'var(--accent-yellow)' },
+                  '&.Mui-focused fieldset': { borderColor: 'var(--accent-yellow)' },
+                },
+              }}
+            />
+            <Stack direction="row" spacing={0.5} justifyContent="flex-end">
+              <Tooltip title="Cancel (Esc)">
+                <IconButton
+                  onClick={handleCancelEdit}
+                  size="small"
+                  sx={{ color: 'var(--muted-text)', '&:hover': { color: 'var(--accent-red)' } }}
+                >
+                  <CloseIcon sx={{ fontSize: 16 }} />
+                </IconButton>
+              </Tooltip>
+              <Tooltip title="Confirm (Enter)">
+                <IconButton
+                  onClick={handleConfirmEdit}
+                  size="small"
+                  sx={{ color: 'var(--accent-green)', '&:hover': { bgcolor: 'rgba(47,204,113,0.1)' } }}
+                >
+                  <CheckIcon sx={{ fontSize: 16 }} />
+                </IconButton>
+              </Tooltip>
+            </Stack>
+          </Stack>
+        ) : (
+          <Typography
+            variant="body1"
+            sx={{
+              fontSize: '0.925rem',
+              lineHeight: 1.65,
+              color: 'var(--text)',
+              whiteSpace: 'pre-wrap',
+              wordBreak: 'break-word',
+            }}
+          >
+            {text}
+          </Typography>
+        )}
 
-        {timeStr && (
+        {timeStr && !isEditing && (
           <Typography
             variant="caption"
             sx={{ color: 'var(--muted-text)', mt: 0.5, display: 'block', textAlign: 'right', fontSize: '0.7rem' }}

frontend/src/components/SessionChat.tsx

@@ -24,7 +24,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
   const { isConnected, isProcessing, activityStatus, updateSession } = useAgentStore();
   const { updateSessionTitle } = useSessionStore();
 
-  const { messages, sendMessage, stop, status, undoLastTurn, approveTools } = useAgentChat({
+  const { messages, sendMessage, stop, status, undoLastTurn, editAndRegenerate, approveTools } = useAgentChat({
     sessionId,
     isActive,
     onReady: () => logger.log(`Session ${sessionId} ready`),
@@ -102,6 +102,7 @@ export default function SessionChat({ sessionId, isActive, onSessionDead }: Sess
         isProcessing={busy}
         approveTools={approveTools}
         onUndoLastTurn={undoLastTurn}
+        onEditAndRegenerate={editAndRegenerate}
       />
       <ChatInput
         onSend={handleSendMessage}

frontend/src/hooks/useAgentChat.ts

@@ -640,12 +640,52 @@ export function useAgentChat({ sessionId, isActive, onReady, onError, onSessionD
     apiFetch(`/api/interrupt/${sessionId}`, { method: 'POST' }).catch(() => {});
   }, [sessionId, updateSession]);
 
+  // -- Edit message + regenerate from that point ----------------------------
+  const editAndRegenerate = useCallback(async (messageId: string, newText: string) => {
+    try {
+      const msgs = chatActionsRef.current.messages;
+      const setMsgs = chatActionsRef.current.setMessages;
+      if (!setMsgs) return;
+
+      // Find the target message and compute user message index (0-indexed, skipping system)
+      const msgIndex = msgs.findIndex(m => m.id === messageId);
+      if (msgIndex < 0) return;
+
+      let userMsgIndex = 0;
+      for (let i = 0; i < msgIndex; i++) {
+        if (msgs[i].role === 'user') userMsgIndex++;
+      }
+
+      // 1. Truncate backend history
+      const res = await apiFetch(`/api/truncate/${sessionId}`, {
+        method: 'POST',
+        body: JSON.stringify({ user_message_index: userMsgIndex }),
+        headers: { 'Content-Type': 'application/json' },
+      });
+      if (!res.ok) {
+        logger.error('Truncate API returned', res.status);
+        return;
+      }
+
+      // 2. Truncate frontend messages
+      const truncated = msgs.slice(0, msgIndex);
+      setMsgs(truncated);
+      saveMessages(sessionId, truncated);
+
+      // 3. Send the edited message (reuses existing transport + /api/chat)
+      chat.sendMessage({ text: newText, metadata: { createdAt: new Date().toISOString() } });
+    } catch (e) {
+      logger.error('Edit and regenerate failed:', e);
+    }
+  }, [sessionId, chat]);
+
   return {
     messages: chat.messages,
     sendMessage: chat.sendMessage,
     stop,
     status: chat.status,
     undoLastTurn,
+    editAndRegenerate,
     approveTools,
   };
 }