diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 1100eec..abac8b6 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -507,6 +507,47 @@ struct ChatHistory { messages: Vec, } +// New branching structures +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Branch { + id: String, + name: String, + created_at: i64, // Unix timestamp in milliseconds + #[serde(default)] + parent_branch_id: Option, + #[serde(default)] + diverge_at_index: usize, // Message index in parent where this branch diverged +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct BranchedChatHistory { + #[serde(default = "default_branches")] + branches: Vec, + #[serde(default = "default_active_branch")] + active_branch_id: String, + #[serde(default)] + branch_messages: std::collections::HashMap>, +} + +fn default_branches() -> Vec { + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as i64; + + vec![Branch { + id: "main".to_string(), + name: "Main".to_string(), + created_at: timestamp, + parent_branch_id: None, + diverge_at_index: 0, + }] +} + +fn default_active_branch() -> String { + "main".to_string() +} + #[derive(Debug, Serialize, Deserialize)] struct ChatRequest { model: String, @@ -951,21 +992,61 @@ fn save_config(config: &ApiConfig) -> Result<(), String> { Ok(()) } -fn load_history(character_id: &str) -> ChatHistory { +// Load branched history (with backward compatibility) +fn load_branched_history(character_id: &str) -> BranchedChatHistory { let path = get_character_history_path(character_id); - if let Ok(contents) = fs::read_to_string(path) { - let mut history: ChatHistory = serde_json::from_str(&contents).unwrap_or(ChatHistory { messages: vec![] }); - // Migrate old messages to new format - for msg in &mut history.messages { - msg.migrate(); + if let Ok(contents) = fs::read_to_string(&path) { + // Try to load as new branched format first + if let Ok(mut branched) = serde_json::from_str::(&contents) { + // Migrate old messages to new format + for messages in branched.branch_messages.values_mut() { + for msg in messages { + msg.migrate(); + } + } + return branched; } - history - } else { - ChatHistory { messages: vec![] } + + // Fall back to old linear format and migrate + if let Ok(mut old_history) = serde_json::from_str::(&contents) { + for msg in &mut old_history.messages { + msg.migrate(); + } + + // Convert to branched format + let mut branch_messages = HashMap::new(); + branch_messages.insert("main".to_string(), old_history.messages); + + return BranchedChatHistory { + branches: default_branches(), + active_branch_id: "main".to_string(), + branch_messages, + }; + } + } + + // Return empty history with main branch + let mut branch_messages = HashMap::new(); + branch_messages.insert("main".to_string(), vec![]); + + BranchedChatHistory { + branches: default_branches(), + active_branch_id: "main".to_string(), + branch_messages, } } -fn save_history(character_id: &str, history: &ChatHistory) -> Result<(), String> { +// Legacy function - returns active branch messages +fn load_history(character_id: &str) -> ChatHistory { + let branched = load_branched_history(character_id); + let messages = branched.branch_messages + .get(&branched.active_branch_id) + .cloned() + .unwrap_or_default(); + ChatHistory { messages } +} + +fn save_branched_history(character_id: &str, history: &BranchedChatHistory) -> Result<(), String> { let path = get_character_history_path(character_id); if let Some(parent) = path.parent() { fs::create_dir_all(parent).map_err(|e| e.to_string())?; @@ -975,6 +1056,13 @@ fn save_history(character_id: &str, history: &ChatHistory) -> Result<(), String> Ok(()) } +// Legacy function - saves to active branch +fn save_history(character_id: &str, history: &ChatHistory) -> Result<(), String> { + let mut branched = load_branched_history(character_id); + branched.branch_messages.insert(branched.active_branch_id.clone(), history.messages.clone()); + save_branched_history(character_id, &branched) +} + fn load_character(character_id: &str) -> Option { let path = get_character_path(character_id); if let Ok(contents) = fs::read_to_string(path) { @@ -3299,6 +3387,153 @@ async fn export_character_card(app_handle: tauri::AppHandle, character_id: Strin Ok(output_path.to_string_lossy().to_string()) } +// ============================================================================ +// Branch Management Commands +// ============================================================================ + +#[tauri::command] +fn create_branch(message_index: usize, branch_name: String) -> Result { + let character = get_active_character(); + let mut branched = load_branched_history(&character.id); + + // Generate new branch ID + let branch_id = Uuid::new_v4().to_string(); + + // Get current branch messages + let current_messages = branched.branch_messages + .get(&branched.active_branch_id) + .ok_or_else(|| "Active branch not found".to_string())?; + + // Validate message index + if message_index > current_messages.len() { + return Err(format!("Invalid message index: {}", message_index)); + } + + // Create new branch + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as i64; + + let new_branch = Branch { + id: branch_id.clone(), + name: branch_name, + created_at: timestamp, + parent_branch_id: Some(branched.active_branch_id.clone()), + diverge_at_index: message_index, + }; + + // Copy messages up to divergence point + let branch_messages: Vec = current_messages[..message_index].to_vec(); + + // Add branch and its messages + branched.branches.push(new_branch.clone()); + branched.branch_messages.insert(branch_id.clone(), branch_messages); + + // Save and return + save_branched_history(&character.id, &branched)?; + Ok(new_branch) +} + +#[tauri::command] +fn switch_branch(branch_id: String) -> Result, String> { + let character = get_active_character(); + let mut branched = load_branched_history(&character.id); + + // Verify branch exists + if !branched.branch_messages.contains_key(&branch_id) { + return Err(format!("Branch '{}' not found", branch_id)); + } + + // Switch active branch + branched.active_branch_id = branch_id.clone(); + save_branched_history(&character.id, &branched)?; + + // Return messages for the new active branch + let messages = branched.branch_messages + .get(&branch_id) + .cloned() + .unwrap_or_default(); + + Ok(messages) +} + +#[tauri::command] +fn delete_branch(branch_id: String) -> Result<(), String> { + let character = get_active_character(); + let mut branched = load_branched_history(&character.id); + + // Cannot delete main branch + if branch_id == "main" { + return Err("Cannot delete main branch".to_string()); + } + + // Cannot delete active branch + if branch_id == branched.active_branch_id { + return Err("Cannot delete active branch. Switch to another branch first.".to_string()); + } + + // Remove branch + branched.branches.retain(|b| b.id != branch_id); + branched.branch_messages.remove(&branch_id); + + // Also remove any child branches that depended on this one + let mut branches_to_remove = vec![]; + for branch in &branched.branches { + if branch.parent_branch_id.as_ref() == Some(&branch_id) { + branches_to_remove.push(branch.id.clone()); + } + } + + for child_id in branches_to_remove { + branched.branches.retain(|b| b.id != child_id); + branched.branch_messages.remove(&child_id); + } + + save_branched_history(&character.id, &branched)?; + Ok(()) +} + +#[tauri::command] +fn list_branches() -> Result, String> { + let character = get_active_character(); + let branched = load_branched_history(&character.id); + Ok(branched.branches) +} + +#[tauri::command] +fn rename_branch(branch_id: String, new_name: String) -> Result<(), String> { + let character = get_active_character(); + let mut branched = load_branched_history(&character.id); + + // Find and rename the branch + let branch = branched.branches.iter_mut() + .find(|b| b.id == branch_id) + .ok_or_else(|| format!("Branch '{}' not found", branch_id))?; + + branch.name = new_name; + save_branched_history(&character.id, &branched)?; + Ok(()) +} + +#[tauri::command] +fn get_active_branch_id() -> Result { + let character = get_active_character(); + let branched = load_branched_history(&character.id); + Ok(branched.active_branch_id) +} + +#[tauri::command] +fn get_branch_info(branch_id: String) -> Result { + let character = get_active_character(); + let branched = load_branched_history(&character.id); + + branched.branches.iter() + .find(|b| b.id == branch_id) + .cloned() + .ok_or_else(|| format!("Branch '{}' not found", branch_id)) +} + #[cfg_attr(mobile, tauri::mobile_entry_point)] pub fn run() { tauri::Builder::default() @@ -3362,7 +3597,14 @@ pub fn run() { update_world_info_entry, delete_world_info_entry, export_world_info, - import_world_info + import_world_info, + create_branch, + switch_branch, + delete_branch, + list_branches, + rename_branch, + get_active_branch_id, + get_branch_info ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); diff --git a/src/main.js b/src/main.js index 7ebb177..e609e26 100644 --- a/src/main.js +++ b/src/main.js @@ -481,6 +481,16 @@ function addMessage(content, isUser = false, skipActions = false, timestamp = nu editBtn.addEventListener('click', () => handleEditMessage(messageDiv, content)); actionsDiv.appendChild(editBtn); + // Branch button + const branchBtn = document.createElement('button'); + branchBtn.className = 'message-action-btn'; + branchBtn.innerHTML = ` + + `; + branchBtn.title = 'Create branch from here'; + branchBtn.addEventListener('click', () => handleCreateBranch(messageDiv)); + actionsDiv.appendChild(branchBtn); + // Copy message button const copyBtn = document.createElement('button'); copyBtn.className = 'message-action-btn'; @@ -557,6 +567,16 @@ function addMessage(content, isUser = false, skipActions = false, timestamp = nu continueBtn.addEventListener('click', () => handleContinueMessage(messageDiv)); actionsDiv.appendChild(continueBtn); + // Branch button + const branchBtn = document.createElement('button'); + branchBtn.className = 'message-action-btn'; + branchBtn.innerHTML = ` + + `; + branchBtn.title = 'Create branch from here'; + branchBtn.addEventListener('click', () => handleCreateBranch(messageDiv)); + actionsDiv.appendChild(branchBtn); + // Pin button const pinBtn = document.createElement('button'); pinBtn.className = 'message-action-btn message-pin-btn'; @@ -1012,6 +1032,95 @@ async function handleDeleteMessage(messageDiv) { } } +// Handle creating a new branch from a message +async function handleCreateBranch(messageDiv) { + const allMessages = Array.from(messagesContainer.querySelectorAll('.message')); + const messageIndex = allMessages.indexOf(messageDiv); + + if (messageIndex === -1) { + console.error('Message not found in list'); + return; + } + + // Prompt for branch name + const branchName = prompt('Enter a name for the new branch:', `Branch ${Date.now()}`); + if (!branchName || branchName.trim() === '') { + return; + } + + try { + // Create branch from this message index + const branch = await invoke('create_branch', { + messageIndex: messageIndex + 1, // +1 because we want to branch AFTER this message + branchName: branchName.trim() + }); + + setStatus(`Created branch: ${branch.name}`, 'success'); + + // Switch to the new branch + await handleSwitchBranch(branch.id); + } catch (error) { + console.error('Failed to create branch:', error); + setStatus(`Branch creation failed: ${error}`, 'error'); + } +} + +// Handle switching to a different branch +async function handleSwitchBranch(branchId) { + try { + // Switch the active branch + const messages = await invoke('switch_branch', { branchId }); + + // Clear current messages + messagesContainer.innerHTML = ''; + + // Reload messages from the new branch (same pattern as loadChatHistory) + messages.forEach((msg) => { + const messageDiv = addMessage(msg.content, msg.role === 'user', false, msg.timestamp); + + // Apply pinned state + if (msg.pinned && messageDiv) { + messageDiv.classList.add('pinned'); + const pinBtn = messageDiv.querySelector('.message-pin-btn'); + if (pinBtn) { + pinBtn.classList.add('active'); + pinBtn.title = 'Unpin message'; + } + } + + // Apply hidden state + if (msg.hidden && messageDiv) { + messageDiv.classList.add('hidden-message'); + const hideBtn = messageDiv.querySelector('.message-hide-btn'); + if (hideBtn) { + hideBtn.classList.add('active'); + hideBtn.title = 'Unhide message'; + hideBtn.innerHTML = ` + + `; + } + } + + // Update swipe controls for assistant messages with swipe info + if (msg.role === 'assistant' && messageDiv && msg.swipes && msg.swipes.length > 0) { + updateSwipeControls(messageDiv, msg.current_swipe || 0, msg.swipes.length); + } + }); + + // Update token count + await updateTokenCount(); + + // Update branch indicator + await updateBranchIndicator(); + + setStatus('Switched branch', 'success'); + setTimeout(() => setStatus('Ready'), 2000); + } catch (error) { + console.error('Failed to switch branch:', error); + setStatus(`Branch switch failed: ${error}`, 'error'); + } +} + // Handle toggling message pin status async function handleTogglePin(messageDiv) { const allMessages = Array.from(messagesContainer.querySelectorAll('.message')); @@ -1714,6 +1823,134 @@ document.addEventListener('click', (e) => { } }); +// Update branch indicator in header +async function updateBranchIndicator() { + try { + const branchId = await invoke('get_active_branch_id'); + const branches = await invoke('list_branches'); + const activeBranch = branches.find(b => b.id === branchId); + + const featureBadges = document.getElementById('feature-badges'); + + // Remove existing branch badge + const existingBadge = featureBadges.querySelector('.branch-badge'); + if (existingBadge) { + existingBadge.remove(); + } + + // Add new branch badge if not on main + if (activeBranch && activeBranch.id !== 'main') { + const branchBadge = document.createElement('div'); + branchBadge.className = 'branch-badge'; + branchBadge.innerHTML = ` + + + + ${activeBranch.name} + `; + branchBadge.title = `Branch: ${activeBranch.name} (click to manage branches)`; + branchBadge.style.cursor = 'pointer'; + branchBadge.addEventListener('click', openBranchManager); + featureBadges.appendChild(branchBadge); + } + } catch (error) { + console.error('Failed to update branch indicator:', error); + } +} + +// Open branch manager modal +async function openBranchManager() { + try { + const branches = await invoke('list_branches'); + const activeBranchId = await invoke('get_active_branch_id'); + + // Create modal + const modal = document.createElement('div'); + modal.className = 'branch-manager-modal'; + modal.innerHTML = ` +
+
+
+

Branch Manager

+ +
+
+ ${branches.map(branch => ` +
+
+
${branch.name}
+
${new Date(branch.created_at).toLocaleString()}
+
+
+ ${branch.id === activeBranchId ? 'Active' : ``} + ${branch.id !== 'main' ? `` : ''} + ${branch.id !== 'main' && branch.id !== activeBranchId ? `` : ''} +
+
+ `).join('')} +
+
+ `; + + document.body.appendChild(modal); + + // Add event listeners + modal.querySelector('#close-branch-manager').addEventListener('click', () => modal.remove()); + modal.querySelector('.branch-manager-overlay').addEventListener('click', () => modal.remove()); + + modal.querySelectorAll('.branch-switch-btn').forEach(btn => { + btn.addEventListener('click', async (e) => { + const branchId = e.target.dataset.branchId; + modal.remove(); + await handleSwitchBranch(branchId); + }); + }); + + modal.querySelectorAll('.branch-rename-btn').forEach(btn => { + btn.addEventListener('click', async (e) => { + const branchId = e.target.dataset.branchId; + const branch = branches.find(b => b.id === branchId); + const newName = prompt('Enter new branch name:', branch.name); + if (newName && newName.trim() !== '') { + try { + await invoke('rename_branch', { branchId, newName: newName.trim() }); + modal.remove(); + await updateBranchIndicator(); + setStatus('Branch renamed', 'success'); + } catch (error) { + setStatus(`Rename failed: ${error}`, 'error'); + } + } + }); + }); + + modal.querySelectorAll('.branch-delete-btn').forEach(btn => { + btn.addEventListener('click', async (e) => { + const branchId = e.target.dataset.branchId; + const branch = branches.find(b => b.id === branchId); + if (confirm(`Delete branch "${branch.name}"? This cannot be undone.`)) { + try { + await invoke('delete_branch', { branchId }); + modal.remove(); + await updateBranchIndicator(); + setStatus('Branch deleted', 'success'); + } catch (error) { + setStatus(`Delete failed: ${error}`, 'error'); + } + } + }); + }); + } catch (error) { + console.error('Failed to open branch manager:', error); + setStatus(`Failed to open branch manager: ${error}`, 'error'); + } +} + // Load characters and populate dropdown async function loadCharacters() { console.log('Loading characters...'); @@ -1748,6 +1985,7 @@ async function loadCharacters() { } await loadChatHistory(); + await updateBranchIndicator(); } catch (error) { console.error('Failed to load characters:', error); addMessage(`Failed to load characters: ${error}`, false); diff --git a/src/styles.css b/src/styles.css index ba49cf0..b5dc17e 100644 --- a/src/styles.css +++ b/src/styles.css @@ -211,6 +211,151 @@ body { color: var(--accent); } +/* Branch Badge */ +.branch-badge { + display: inline-flex; + align-items: center; + gap: 4px; + padding: 3px 8px; + background: rgba(34, 197, 94, 0.1); + border: 1px solid rgba(34, 197, 94, 0.3); + border-radius: 12px; + font-size: 11px; + font-weight: 500; + color: #22c55e; + white-space: nowrap; + transition: all 0.2s ease; +} + +.branch-badge:hover { + background: rgba(34, 197, 94, 0.15); + border-color: rgba(34, 197, 94, 0.5); +} + +.branch-badge svg { + width: 12px; + height: 12px; + opacity: 0.8; +} + +/* Branch Manager Modal */ +.branch-manager-modal { + position: fixed; + top: 0; + left: 0; + right: 0; + bottom: 0; + z-index: 10002; + display: flex; + align-items: center; + justify-content: center; +} + +.branch-manager-overlay { + position: absolute; + top: 0; + left: 0; + right: 0; + bottom: 0; + background: rgba(0, 0, 0, 0.7); + backdrop-filter: blur(4px); + -webkit-backdrop-filter: blur(4px); + cursor: pointer; +} + +.branch-manager-content { + position: relative; + z-index: 1; + background: var(--bg-primary); + border: 1px solid var(--border); + border-radius: 12px; + padding: 24px; + width: 90%; + max-width: 600px; + max-height: 80vh; + box-shadow: 0 20px 60px rgba(0, 0, 0, 0.5); + display: flex; + flex-direction: column; +} + +.branch-manager-header { + display: flex; + justify-content: space-between; + align-items: center; + margin-bottom: 20px; + padding-bottom: 16px; + border-bottom: 1px solid var(--border); +} + +.branch-manager-header h3 { + font-size: 20px; + font-weight: 600; + color: var(--text-primary); + margin: 0; +} + +.branch-list { + display: flex; + flex-direction: column; + gap: 12px; + overflow-y: auto; + max-height: 60vh; +} + +.branch-item { + background: var(--bg-secondary); + border: 1px solid var(--border); + border-radius: 8px; + padding: 12px; + display: flex; + justify-content: space-between; + align-items: center; + gap: 12px; + transition: all 0.2s ease; +} + +.branch-item:hover { + border-color: var(--accent); +} + +.branch-item.active { + background: rgba(99, 102, 241, 0.1); + border-color: var(--accent); +} + +.branch-info { + flex: 1; + min-width: 0; +} + +.branch-name { + font-size: 14px; + font-weight: 500; + color: var(--text-primary); + margin-bottom: 4px; +} + +.branch-meta { + font-size: 11px; + color: var(--text-secondary); +} + +.branch-actions { + display: flex; + align-items: center; + gap: 8px; + flex-shrink: 0; +} + +.branch-active-label { + font-size: 12px; + padding: 4px 12px; + background: var(--accent); + color: white; + border-radius: 6px; + font-weight: 500; +} + .icon-btn { width: 28px; height: 28px; @@ -371,19 +516,24 @@ body { /* Message action buttons */ .message-actions { position: absolute; - top: 8px; + top: -4px; display: flex; gap: 4px; opacity: 0; transition: opacity 0.2s ease; + background: rgba(0, 0, 0, 0.6); + padding: 4px; + border-radius: 8px; + backdrop-filter: blur(12px); + -webkit-backdrop-filter: blur(12px); } .message.user .message-actions { - right: 8px; + right: 4px; } .message.assistant .message-actions { - right: 8px; + right: 4px; } .message:hover .message-actions { @@ -391,26 +541,28 @@ body { } .message-action-btn { - width: 24px; - height: 24px; - border: none; - background: rgba(0, 0, 0, 0.5); - backdrop-filter: blur(8px); - -webkit-backdrop-filter: blur(8px); + width: 28px; + height: 28px; + border: 1px solid rgba(255, 255, 255, 0.15); + background: rgba(0, 0, 0, 0.85); + backdrop-filter: blur(12px); + -webkit-backdrop-filter: blur(12px); border-radius: 6px; - color: var(--text-secondary); + color: rgba(255, 255, 255, 0.8); cursor: pointer; display: flex; align-items: center; justify-content: center; transition: all 0.2s ease; padding: 0; + box-shadow: 0 2px 8px rgba(0, 0, 0, 0.4); } .message-action-btn:hover { - background: rgba(0, 0, 0, 0.7); - color: var(--text-primary); - transform: scale(1.1); + background: rgba(0, 0, 0, 0.95); + border-color: rgba(255, 255, 255, 0.3); + color: white; + transform: scale(1.05); } .message-action-btn:active { @@ -418,8 +570,8 @@ body { } .message-action-btn svg { - width: 14px; - height: 14px; + width: 16px; + height: 16px; } /* Enhanced message control buttons */