@@ -291,7 +291,7 @@ class ConversationStore:
291291
292292 def __init__ (self , db_path : Path ):
293293 self ._db_path = db_path
294- self ._lock = threading .Lock ()
294+ self ._lock = threading .RLock () # Use RLock to allow reentrant locking
295295 self ._init_db ()
296296
297297 # ------------------------------------------------------------------
@@ -524,6 +524,109 @@ def clear_session(self, session_id: str) -> None:
524524 finally :
525525 conn .close ()
526526
527+ def delete_message_pair (self , session_id : str , user_seq : int , delete_user : bool = True , cascade : bool = False ) -> int :
528+ """Delete a user message and/or its corresponding assistant reply.
529+
530+ The assistant reply is identified as all messages between user_seq
531+ and the next visible user message (or end of session).
532+
533+ Args:
534+ session_id: Session identifier.
535+ user_seq: The seq number of the user message.
536+ delete_user: If True (default), delete the user message too.
537+ If False, only delete assistant reply (for regenerate scenarios).
538+ cascade: If True, also delete all subsequent turns after this one.
539+ Used by edit-message which removes this turn and everything after.
540+
541+ Returns:
542+ Number of message rows deleted.
543+ """
544+ with self ._lock :
545+ conn = self ._connect ()
546+ try :
547+ with conn :
548+ # Verify this is a user message
549+ row = conn .execute (
550+ "SELECT role FROM messages WHERE session_id = ? AND seq = ?" ,
551+ (session_id , user_seq ),
552+ ).fetchone ()
553+ if not row or row [0 ] != "user" :
554+ return 0
555+
556+ if cascade :
557+ # Delete from this message to end of session
558+ start_seq = user_seq if delete_user else user_seq + 1
559+ end_seq_row = conn .execute (
560+ "SELECT MAX(seq) FROM messages WHERE session_id = ?" ,
561+ (session_id ,),
562+ ).fetchone ()
563+ end_seq = (end_seq_row [0 ] or user_seq ) + 1
564+ else :
565+ # Find the next visible user message seq (exclude tool_result)
566+ # Use batched query to avoid loading too many rows at once
567+ next_user_seq = None
568+ batch_size = 100
569+ offset = 0
570+ while True :
571+ batch = conn .execute (
572+ """
573+ SELECT seq, content FROM messages
574+ WHERE session_id = ? AND seq > ? AND role = 'user'
575+ ORDER BY seq ASC
576+ LIMIT ? OFFSET ?
577+ """ ,
578+ (session_id , user_seq , batch_size , offset ),
579+ ).fetchall ()
580+ if not batch :
581+ break
582+ for seq , content in batch :
583+ try :
584+ content_obj = json .loads (content )
585+ except Exception :
586+ content_obj = content
587+ if _is_visible_user_message (content_obj ):
588+ next_user_seq = seq
589+ break
590+ if next_user_seq is not None :
591+ break
592+ offset += batch_size
593+
594+ # Determine the end boundary for deletion
595+ if next_user_seq is not None :
596+ end_seq = next_user_seq
597+ else :
598+ end_seq_row = conn .execute (
599+ "SELECT MAX(seq) FROM messages WHERE session_id = ?" ,
600+ (session_id ,),
601+ ).fetchone ()
602+ end_seq = (end_seq_row [0 ] or user_seq ) + 1
603+
604+ # Determine the start boundary for deletion
605+ start_seq = user_seq if delete_user else user_seq + 1
606+
607+ # Delete messages from start_seq to end_seq (exclusive)
608+ cur = conn .execute (
609+ "DELETE FROM messages WHERE session_id = ? AND seq >= ? AND seq < ?" ,
610+ (session_id , start_seq , end_seq ),
611+ )
612+ deleted = cur .rowcount
613+
614+ # Update session msg_count
615+ conn .execute (
616+ """
617+ UPDATE sessions
618+ SET msg_count = (
619+ SELECT COUNT(*) FROM messages WHERE session_id = ?
620+ )
621+ WHERE session_id = ?
622+ """ ,
623+ (session_id , session_id ),
624+ )
625+
626+ return deleted
627+ finally :
628+ conn .close ()
629+
527630 def prune_scheduled_messages (
528631 self ,
529632 session_id : str ,
@@ -1053,3 +1156,4 @@ def get_conversation_store() -> ConversationStore:
10531156 _store_instance = ConversationStore (db_path )
10541157 logger .debug (f"[ConversationStore] Using shared DB at: { db_path } )
10551158 return _store_instance
1159+
0 commit comments