feat: optimistic updates for jj new with pre-allocated change IDs

+17 -2

apps/desktop/src-tauri/src/lib.rs

··· 274 274 watcher_manager.unwatch(&PathBuf::from(repo_path)) 275 275 } 276 276 277 + /// Generate change IDs for optimistic UI updates 277 278 #[tauri::command] 278 - async fn jj_new(repo_path: String, parent_change_ids: Vec<String>) -> Result<(), String> { 279 + async fn generate_change_ids(repo_path: String, count: usize) -> Result<Vec<String>, String> { 280 + let path = Path::new(&repo_path); 281 + let jj_repo = JjRepo::open(path).map_err(|e| format!("Failed to open repo: {}", e))?; 282 + jj_repo 283 + .generate_change_ids(count) 284 + .map_err(|e| format!("Failed to generate change IDs: {}", e)) 285 + } 286 + 287 + #[tauri::command] 288 + async fn jj_new( 289 + repo_path: String, 290 + parent_change_ids: Vec<String>, 291 + change_id: Option<String>, 292 + ) -> Result<String, String> { 279 293 let path = Path::new(&repo_path); 280 294 let mut jj_repo = JjRepo::open(path).map_err(|e| format!("Failed to open repo: {}", e))?; 281 295 jj_repo 282 - .new_revision(parent_change_ids) 296 + .new_revision(parent_change_ids, change_id) 283 297 .map_err(|e| format!("Failed to create new revision: {}", e)) 284 298 } 285 299 ··· 391 405 update_layout, 392 406 watch_repository, 393 407 unwatch_repository, 408 + generate_change_ids, 394 409 jj_new, 395 410 jj_edit, 396 411 jj_abandon,

+33 -6

apps/desktop/src-tauri/src/repo/jj.rs

··· 1 1 use anyhow::{Context, Result}; 2 - use jj_lib::backend::CommitId; 2 + use jj_lib::backend::{ChangeId, CommitId}; 3 3 use jj_lib::commit::Commit; 4 4 use jj_lib::config::ConfigSource; 5 5 use jj_lib::merged_tree::MergedTree; ··· 182 182 self.workspace.workspace_root() 183 183 } 184 184 185 - pub fn new_revision(&mut self, parent_change_ids: Vec<String>) -> Result<()> { 185 + /// Generate change IDs using jj-lib's RNG. Returns reverse-hex encoded IDs. 186 + pub fn generate_change_ids(&self, count: usize) -> Result<Vec<String>> { 187 + let repo = self.workspace.repo_loader().load_at_head()?; 188 + let rng = self.user_settings.get_rng(); 189 + let length = repo.store().change_id_length(); 190 + 191 + let ids: Vec<String> = (0..count) 192 + .map(|_| rng.new_change_id(length).reverse_hex()) 193 + .collect(); 194 + 195 + Ok(ids) 196 + } 197 + 198 + pub fn new_revision( 199 + &mut self, 200 + parent_change_ids: Vec<String>, 201 + change_id: Option<String>, 202 + ) -> Result<String> { 186 203 let repo = self.workspace.repo_loader().load_at_head()?; 187 204 let mut tx = repo.start_transaction(); 188 205 ··· 204 221 205 222 // Create new commit with parent commits and their tree (no changes) 206 223 let parent_commit_ids: Vec<_> = parent_commits.iter().map(|c| c.id().clone()).collect(); 207 - let new_commit = tx 208 - .repo_mut() 209 - .new_commit(parent_commit_ids, tree_id) 224 + let mut commit_builder = tx.repo_mut().new_commit(parent_commit_ids, tree_id); 225 + 226 + // Set pre-generated change ID if provided 227 + if let Some(ref cid) = change_id { 228 + let parsed = ChangeId::try_from_reverse_hex(cid) 229 + .context("Invalid change ID format")?; 230 + commit_builder = commit_builder.set_change_id(parsed); 231 + } 232 + 233 + let new_commit = commit_builder 210 234 .write() 211 235 .map_err(|e| anyhow::anyhow!("Failed to write commit: {}", e))?; 236 + 237 + // Get the actual change ID (either provided or generated) 238 + let actual_change_id = new_commit.change_id().reverse_hex(); 212 239 213 240 // Set as working copy 214 241 let workspace_name = self.workspace.workspace_name().to_owned(); ··· 233 260 .check_out(operation_id, Some(&old_tree_id), &new_commit) 234 261 .context("Failed to check out new commit")?; 235 262 236 - Ok(()) 263 + Ok(actual_change_id) 237 264 } 238 265 239 266 pub fn abandon_revision(&mut self, change_id: &str) -> Result<()> {

+134 -25

apps/desktop/src/db.ts

··· 1 - import { createCollection, createTransaction } from "@tanstack/db"; 1 + import { createCollection } from "@tanstack/db"; 2 2 import { QueryClient } from "@tanstack/query-core"; 3 3 import { queryCollectionOptions } from "@tanstack/query-db-collection"; 4 4 import { listen } from "@tauri-apps/api/event"; 5 + import { Effect } from "effect"; 5 6 import type { ChangedFile, Repository, Revision } from "@/tauri-commands"; 6 7 import { 8 + generateChangeIds, 7 9 getCommitRecency, 8 10 getRepositories, 9 11 getRevisionChanges, ··· 46 48 } 47 49 48 50 // ============================================================================ 51 + // Change ID Pool Collection (pre-allocated IDs for optimistic updates) 52 + // ============================================================================ 53 + 54 + const POOL_SIZE = 10; 55 + const POOL_REFILL_THRESHOLD = 3; 56 + 57 + interface ChangeIdPool { 58 + repoPath: string; 59 + ids: string[]; 60 + } 61 + 62 + function changeIdPoolQueryKey(repoPath: string) { 63 + return ["change-id-pool", repoPath] as const; 64 + } 65 + 66 + async function fetchChangeIdPool(repoPath: string): Promise<ChangeIdPool> { 67 + const ids = await generateChangeIds(repoPath, POOL_SIZE); 68 + return { repoPath, ids }; 69 + } 70 + 71 + /** Ensure the change ID pool is loaded. Call from router beforeLoad. */ 72 + export async function ensureChangeIdPool(repoPath: string): Promise<void> { 73 + await queryClient.ensureQueryData({ 74 + queryKey: changeIdPoolQueryKey(repoPath), 75 + queryFn: () => fetchChangeIdPool(repoPath), 76 + }); 77 + } 78 + 79 + /** Consume a change ID from the pool, triggering refill if needed */ 80 + function consumeChangeId(repoPath: string): string | null { 81 + const poolEntry = queryClient.getQueryData<ChangeIdPool>(changeIdPoolQueryKey(repoPath)); 82 + if (!poolEntry || poolEntry.ids.length === 0) return null; 83 + 84 + const [id, ...remaining] = poolEntry.ids; 85 + 86 + // Update the cache directly 87 + queryClient.setQueryData<ChangeIdPool>(changeIdPoolQueryKey(repoPath), { 88 + repoPath, 89 + ids: remaining, 90 + }); 91 + 92 + // Trigger refill if running low 93 + if (remaining.length < POOL_REFILL_THRESHOLD) { 94 + generateChangeIds(repoPath, POOL_SIZE).then((newIds) => { 95 + const current = queryClient.getQueryData<ChangeIdPool>(changeIdPoolQueryKey(repoPath)); 96 + const currentIds = current?.ids ?? []; 97 + queryClient.setQueryData<ChangeIdPool>(changeIdPoolQueryKey(repoPath), { 98 + repoPath, 99 + ids: [...currentIds, ...newIds], 100 + }); 101 + }); 102 + } 103 + 104 + return id; 105 + } 106 + 107 + // ============================================================================ 49 108 // Shared Repository Watcher (one per repo, invalidates all queries) 50 109 // ============================================================================ 51 110 ··· 83 142 // Repositories Collection 84 143 // ============================================================================ 85 144 145 + const repositoriesQueryKey = ["repositories"] as const; 146 + 86 147 export const repositoriesCollection = createCollection({ 87 148 ...queryCollectionOptions({ 88 149 queryClient, 89 - queryKey: ["repositories"], 150 + queryKey: repositoriesQueryKey, 90 151 queryFn: getRepositories, 91 152 getKey: (repository: Repository) => repository.id, 92 153 }), 93 154 }); 94 155 95 156 export type RepositoriesCollection = typeof repositoriesCollection; 157 + 158 + /** Ensure repositories are loaded. Returns the list. */ 159 + export async function ensureRepositories(): Promise<Repository[]> { 160 + return queryClient.ensureQueryData({ 161 + queryKey: repositoriesQueryKey, 162 + queryFn: getRepositories, 163 + }); 164 + } 96 165 97 166 export async function addRepository(collection: RepositoriesCollection, repository: Repository) { 98 167 // Optimistic update first ··· 174 243 175 244 const revisionCollections = new Map<string, ReturnType<typeof createRevisionsCollection>>(); 176 245 177 - function createRevisionsCollection(repoPath: string, preset?: string, customRevset?: string) { 246 + function createRevisionsCollection(repoPath: string, preset?: string) { 178 247 const limit = preset === "full_history" ? 10000 : 100; 179 248 180 249 // Set up the shared watcher (idempotent - increments refCount if already exists) ··· 183 252 return createCollection({ 184 253 ...queryCollectionOptions({ 185 254 queryClient, 186 - queryKey: ["revisions", repoPath, preset, customRevset], 187 - queryFn: () => getRevisions(repoPath, limit, customRevset, customRevset ? undefined : preset), 255 + queryKey: ["revisions", repoPath, preset], 256 + queryFn: () => getRevisions(repoPath, limit, undefined, preset), 188 257 getKey: getRevisionKey, 189 258 }), 190 259 }); ··· 192 261 193 262 export type RevisionsCollection = ReturnType<typeof createRevisionsCollection>; 194 263 195 - export function getRevisionsCollection(repoPath: string, preset?: string, customRevset?: string) { 196 - const cacheKey = `${repoPath}:${preset ?? "full_history"}:${customRevset ?? ""}`; 264 + export function getRevisionsCollection(repoPath: string, preset?: string) { 265 + const cacheKey = `${repoPath}:${preset ?? "full_history"}`; 197 266 let collection = revisionCollections.get(cacheKey); 198 267 if (!collection) { 199 - collection = createRevisionsCollection(repoPath, preset, customRevset); 268 + collection = createRevisionsCollection(repoPath, preset); 200 269 revisionCollections.set(cacheKey, collection); 201 270 } 202 271 return collection; ··· 241 310 }); 242 311 } 243 312 244 - export function newRevision(repoPath: string, parentChangeIds: string[]) { 313 + export function newRevision( 314 + collection: RevisionsCollection, 315 + repoPath: string, 316 + parentChangeIds: string[], 317 + parentRevision: Revision, 318 + currentWcRevision: Revision | null, 319 + ) { 245 320 const mutationId = `new-${Date.now()}-${Math.random()}`; 246 - console.log("[newRevision] start, mutationId:", mutationId); 321 + const preAllocatedChangeId = consumeChangeId(repoPath); 322 + 323 + // Create optimistic revision if we have a pre-allocated change ID 324 + let optimisticRevision: Revision | null = null; 325 + if (preAllocatedChangeId) { 326 + optimisticRevision = { 327 + commit_id: `pending-${preAllocatedChangeId}`, // Temporary, will be replaced 328 + change_id: preAllocatedChangeId, 329 + change_id_short: preAllocatedChangeId.slice(0, 8), // Approximate short ID 330 + parent_ids: [parentRevision.commit_id], 331 + parent_edges: [{ parent_id: parentRevision.commit_id, edge_type: "direct" as const }], 332 + description: "", 333 + author: parentRevision.author, // Inherit from parent 334 + timestamp: new Date().toISOString(), 335 + is_working_copy: true, 336 + is_immutable: false, 337 + is_mine: true, 338 + is_trunk: false, 339 + is_divergent: false, 340 + divergent_index: null, 341 + bookmarks: [], 342 + }; 247 343 248 - const tx = createTransaction({ 249 - mutationFn: async () => { 250 - await trackMutation(mutationId, jjNew(repoPath, parentChangeIds)); 251 - // Invalidate to get fresh data including the new revision 252 - await queryClient.invalidateQueries({ queryKey: ["revisions", repoPath] }); 253 - }, 254 - }); 344 + // Optimistic update: clear WC from current, insert new revision 345 + const updates: Revision[] = []; 346 + if (currentWcRevision) { 347 + updates.push({ ...currentWcRevision, is_working_copy: false }); 348 + } 349 + updates.push(optimisticRevision); 350 + collection.utils.writeUpsert(updates); 351 + } 255 352 256 - tx.mutate(() => { 257 - // No optimistic update - we don't know the new revision's ID 258 - // TanStack Query invalidation will add it to the collection 259 - }); 353 + // Fire backend call 354 + const program = Effect.tryPromise({ 355 + try: () => jjNew(repoPath, parentChangeIds, preAllocatedChangeId ?? undefined), 356 + catch: (error) => new Error(`Failed to create new revision: ${error}`), 357 + }).pipe(Effect.tapError((error) => Effect.logError("jjNew failed", error))); 260 358 261 - return tx; 359 + trackMutation(mutationId, Effect.runPromise(program)) 360 + .then(() => { 361 + // Invalidate to get authoritative data (correct commit_id, short_id, etc.) 362 + queryClient.invalidateQueries({ queryKey: ["revisions", repoPath] }); 363 + }) 364 + .catch((err) => { 365 + console.error("[newRevision] failed:", err); 366 + // Revert optimistic update 367 + if (optimisticRevision) { 368 + collection.utils.writeDelete(getRevisionKey(optimisticRevision)); 369 + if (currentWcRevision) { 370 + collection.utils.writeUpsert([{ ...currentWcRevision, is_working_copy: true }]); 371 + } 372 + } 373 + }); 262 374 } 263 375 264 376 export function abandonRevision( 265 377 collection: RevisionsCollection, 266 378 repoPath: string, 267 379 revision: Revision, 268 - _limit: number, 269 - _customRevset?: string, 270 - _preset?: string, 271 380 ) { 272 381 const mutationId = `abandon-${Date.now()}-${Math.random()}`; 273 382 console.log("[abandonRevision] abandoning:", revision.change_id_short, "mutationId:", mutationId);

+10 -3

apps/desktop/src/mocks/setup.ts

··· 918 918 get_revision_changes: (): ChangedFile[] => mockChangedFiles, 919 919 watch_repository: () => undefined, 920 920 unwatch_repository: () => undefined, 921 + generate_change_ids: (args) => { 922 + const count = (args.count as number) ?? 10; 923 + return Array.from({ length: count }, () => generateChangeId()); 924 + }, 921 925 jj_new: (args) => { 922 926 const parentChangeIds = args.parentChangeIds as string[]; 927 + const providedChangeId = args.changeId as string | null; 928 + 923 929 // Find parent revisions by change_id (handling short IDs) 924 930 const parentRevisions = parentChangeIds 925 931 .map((id) => ··· 936 942 mockRevisions[currentWcIndex] = { ...mockRevisions[currentWcIndex], is_working_copy: false }; 937 943 } 938 944 939 - // Create new revision 940 - const newChangeId = generateChangeId(); 945 + // Use provided change ID or generate new one 946 + const newChangeId = providedChangeId ?? generateChangeId(); 941 947 const newCommitId = `new${Date.now().toString(16).slice(-10)}`; 942 948 const newRevision: Omit<Revision, "change_id_short"> = { 943 949 commit_id: newCommitId, ··· 963 969 ]; 964 970 mockRevisions = calculateShortIds(allRevisionsRaw); 965 971 966 - return undefined; 972 + // Return the change ID (matching real backend behavior) 973 + return newChangeId; 967 974 }, 968 975 jj_edit: (args) => { 969 976 const changeId = args.changeId as string;

+10 -1

apps/desktop/src/routes/project.$projectId.tsx

··· 1 1 import { useLiveQuery } from "@tanstack/react-db"; 2 2 import { createRoute, Navigate, useParams } from "@tanstack/react-router"; 3 3 import { AppShell } from "@/components/AppShell"; 4 - import { repositoriesCollection } from "@/db"; 4 + import { ensureChangeIdPool, ensureRepositories, repositoriesCollection } from "@/db"; 5 5 import type { Repository } from "@/tauri-commands"; 6 6 import { Route as rootRoute } from "./__root"; 7 7 ··· 27 27 selectionAnchor: 28 28 typeof search.selectionAnchor === "string" ? search.selectionAnchor : undefined, 29 29 }; 30 + }, 31 + beforeLoad: async ({ params }) => { 32 + // Find repository path from projectId to pre-warm change ID pool 33 + const repositories = await ensureRepositories(); 34 + const repo = repositories.find((r) => r.id === params.projectId); 35 + if (repo) { 36 + // Ensure change ID pool is loaded before rendering 37 + await ensureChangeIdPool(repo.path); 38 + } 30 39 }, 31 40 component: ProjectComponent, 32 41 });

+10 -2

apps/desktop/src/tauri-commands.ts

··· 69 69 return invoke("unwatch_repository", { repoPath }); 70 70 } 71 71 72 - export async function jjNew(repoPath: string, parentChangeIds: string[]): Promise<void> { 73 - return invoke("jj_new", { repoPath, parentChangeIds }); 72 + export async function generateChangeIds(repoPath: string, count: number): Promise<string[]> { 73 + return invoke<string[]>("generate_change_ids", { repoPath, count }); 74 + } 75 + 76 + export async function jjNew( 77 + repoPath: string, 78 + parentChangeIds: string[], 79 + changeId?: string, 80 + ): Promise<string> { 81 + return invoke<string>("jj_new", { repoPath, parentChangeIds, changeId: changeId ?? null }); 74 82 } 75 83 76 84 export async function jjEdit(repoPath: string, changeId: string): Promise<void> {

Configure Feed

Configure Feed