Skip to content

Commit 741796e

Browse files
DevanshuNEUDevanshuNEU
authored
Merge pull request #160 from DevanshuNEU/fix/154-indexing-bugs
fix(indexing): WebSocket race condition, NaN stats, file streaming (#154)
2 parents 6bf8439 + 85f229f commit 741796e

8 files changed

Lines changed: 116 additions & 62 deletions

File tree

‎backend/routes/ws_playground.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,9 @@ async def websocket_playground_index(websocket: WebSocket, job_id: str):
152152
job_id=job_id[:12],
153153
event_type=event_type
154154
)
155+
# Small delay to ensure client processes message before close
156+
# This prevents race condition where onclose fires before onmessage
157+
await asyncio.sleep(0.2)
155158
break
156159

157160
except json.JSONDecodeError:

‎backend/services/anonymous_indexer.py‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,9 @@ def to_dict(self) -> dict:
4343
@dataclass
4444
class JobStats:
4545
"""Final stats for completed job."""
46-
files_indexed: int = 0
47-
functions_found: int = 0
48-
time_taken_seconds: float = 0
46+
files_processed: int = 0
47+
functions_indexed: int = 0
48+
indexing_time_seconds: float = 0
4949

5050
def to_dict(self) -> dict:
5151
return asdict(self)
@@ -345,12 +345,18 @@ async def run_indexing_job(
345345
job_manager.update_status(job_id, JobStatus.PROCESSING)
346346

347347
# Progress callback for real-time updates
348-
async def progress_callback(files_processed: int, functions_found: int, total: int):
348+
async def progress_callback(
349+
files_processed: int,
350+
functions_found: int,
351+
total: int,
352+
current_file: Optional[str] = None
353+
):
349354
job_manager.update_progress(
350355
job_id,
351356
files_processed=files_processed,
352357
functions_found=functions_found,
353-
files_total=total
358+
files_total=total,
359+
current_file=current_file
354360
)
355361

356362
# Run indexing with timeout
@@ -370,9 +376,9 @@ async def progress_callback(files_processed: int, functions_found: int, total: i
370376
# --- Step 3: Mark complete ---
371377
elapsed = time.time() - start_time
372378
stats = JobStats(
373-
files_indexed=file_count,
374-
functions_found=total_functions,
375-
time_taken_seconds=round(elapsed, 2)
379+
files_processed=file_count,
380+
functions_indexed=total_functions,
381+
indexing_time_seconds=round(elapsed, 2)
376382
)
377383

378384
job_manager.update_status(

‎backend/services/indexer_optimized.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -672,8 +672,11 @@ async def index_repository_with_progress(
672672

673673
files_processed = min(i + self.FILE_BATCH_SIZE, total_files)
674674

675-
# Send progress update
676-
await progress_callback(files_processed, len(all_functions_data), total_files)
675+
# Get the last file in this batch for display
676+
current_file = batch[-1].name if batch else None
677+
678+
# Send progress update with current file
679+
await progress_callback(files_processed, len(all_functions_data), total_files, current_file)
677680

678681
logger.debug("Processing files",
679682
processed=files_processed,

‎backend/tests/test_anonymous_indexing.py‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -199,9 +199,9 @@ def test_update_status_completed_with_stats(self, job_manager, mock_redis):
199199
})
200200

201201
stats = JobStats(
202-
files_indexed=100,
203-
functions_found=500,
204-
time_taken_seconds=45.5
202+
files_processed=100,
203+
functions_indexed=500,
204+
indexing_time_seconds=45.5
205205
)
206206

207207
result = job_manager.update_status(
@@ -259,13 +259,13 @@ def test_job_progress_none_excluded(self):
259259
def test_job_stats_to_dict(self):
260260
"""JobStats converts to dict correctly."""
261261
stats = JobStats(
262-
files_indexed=100,
263-
functions_found=500,
264-
time_taken_seconds=45.5
262+
files_processed=100,
263+
functions_indexed=500,
264+
indexing_time_seconds=45.5
265265
)
266266
d = stats.to_dict()
267-
assert d["files_indexed"] == 100
268-
assert d["time_taken_seconds"] == 45.5
267+
assert d["files_processed"] == 100
268+
assert d["indexing_time_seconds"] == 45.5
269269

270270

271271
# =============================================================================
@@ -608,9 +608,9 @@ def test_completed_job_returns_repo_id(self, mock_job_class, client):
608608
"created_at": "2024-01-01T00:00:00Z",
609609
"updated_at": "2024-01-01T00:01:00Z",
610610
"stats": {
611-
"files_indexed": 100,
612-
"functions_found": 500,
613-
"time_taken_seconds": 45.2
611+
"files_processed": 100,
612+
"functions_indexed": 500,
613+
"indexing_time_seconds": 45.2
614614
}
615615
}
616616
mock_job_class.return_value = mock_job_manager
@@ -621,7 +621,7 @@ def test_completed_job_returns_repo_id(self, mock_job_class, client):
621621
data = response.json()
622622
assert data["status"] == "completed"
623623
assert data["repo_id"] == "anon_idx_test123456"
624-
assert data["stats"]["files_indexed"] == 100
624+
assert data["stats"]["files_processed"] == 100
625625

626626
@patch('routes.playground.AnonymousIndexingJob')
627627
def test_failed_job_returns_error(self, mock_job_class, client):

‎backend/tests/test_ws_playground.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def test_websocket_handles_already_completed_job(self):
6666
"job_id": "idx_test123",
6767
"status": "completed",
6868
"repo_id": "anon_test123",
69-
"stats": {"files_indexed": 100, "functions_found": 500}
69+
"stats": {"files_processed": 100, "functions_indexed": 500}
7070
})
7171

7272
with patch('routes.ws_playground.redis_client', mock_redis):
@@ -172,9 +172,9 @@ def test_completed_event_includes_stats(self):
172172
job_manager = AnonymousIndexingJob(mock_redis)
173173

174174
stats = JobStats(
175-
files_indexed=100,
176-
functions_found=500,
177-
time_taken_seconds=45.2
175+
files_processed=100,
176+
functions_indexed=500,
177+
indexing_time_seconds=45.2
178178
)
179179

180180
job_manager.update_status(
@@ -189,7 +189,7 @@ def test_completed_event_includes_stats(self):
189189

190190
assert event_data["type"] == "completed"
191191
assert event_data["repo_id"] == "anon_test123"
192-
assert event_data["stats"]["functions_found"] == 500
192+
assert event_data["stats"]["functions_indexed"] == 500
193193

194194
def test_processing_status_skips_duplicate_publish(self):
195195
"""PROCESSING status should not publish (handled by update_progress)."""

‎frontend/src/components/playground/HeroPlayground.tsx‎

Lines changed: 33 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -169,11 +169,13 @@ export function HeroPlayground({
169169
}, [state]);
170170

171171
// Determine visibility states
172+
const hasError = state.status === 'error' || wsHasError;
172173
const showDemoSelector = mode === 'demo';
173-
const showUrlInput = mode === 'custom' && !['indexing', 'ready'].includes(state.status) && !showCelebration && !wsIsCompleted;
174-
const showValidation = mode === 'custom' && ['validating', 'valid', 'invalid'].includes(state.status) && !showCelebration;
175-
const showIndexing = mode === 'custom' && state.status === 'indexing' && !showCelebration && !wsIsCompleted;
174+
const showUrlInput = mode === 'custom' && !['indexing', 'ready', 'error'].includes(state.status) && !showCelebration && !wsIsCompleted && !wsHasError;
175+
const showValidation = mode === 'custom' && ['validating', 'valid', 'invalid'].includes(state.status) && !showCelebration && !hasError;
176+
const showIndexing = mode === 'custom' && state.status === 'indexing' && !showCelebration && !wsIsCompleted && !wsHasError;
176177
const showReady = (mode === 'custom' && state.status === 'ready') || (wsIsCompleted && !showCelebration);
178+
const showError = mode === 'custom' && hasError && !showCelebration;
177179
const isSearchDisabled = mode === 'custom' && state.status !== 'ready' && !wsIsCompleted;
178180

179181
// Can search?
@@ -350,6 +352,34 @@ export function HeroPlayground({
350352
</button>
351353
</motion.div>
352354
)}
355+
356+
{/* Error State - Inside AnimatePresence to prevent simultaneous renders */}
357+
{showError && (
358+
<motion.div
359+
key="error"
360+
initial={{ opacity: 0, y: 10 }}
361+
animate={{ opacity: 1, y: 0 }}
362+
exit={{ opacity: 0, y: -10 }}
363+
className="mb-4 px-4 py-3 rounded-xl bg-amber-500/10 border border-amber-500/20"
364+
>
365+
<div className="flex items-start gap-3">
366+
<span className="text-amber-400 text-lg">⚠️</span>
367+
<div className="flex-1">
368+
<p className="text-amber-200 text-sm font-medium mb-1">Something went wrong</p>
369+
<p className="text-zinc-400 text-sm">
370+
{state.status === 'error' ? state.message : wsError || 'An unexpected error occurred'}
371+
</p>
372+
<button
373+
type="button"
374+
onClick={handleIndexAnother}
375+
className="mt-3 text-sm text-amber-400 hover:text-amber-300 transition-colors font-medium"
376+
>
377+
← Try again with a different repository
378+
</button>
379+
</div>
380+
</div>
381+
</motion.div>
382+
)}
353383
</AnimatePresence>
354384

355385
{/* Search Box */}
@@ -407,26 +437,6 @@ export function HeroPlayground({
407437
</>
408438
)}
409439

410-
{/* Error State */}
411-
{(state.status === 'error' || wsHasError) && (
412-
<motion.div
413-
initial={{ opacity: 0, y: 10 }}
414-
animate={{ opacity: 1, y: 0 }}
415-
className="mt-4 px-4 py-3 rounded-xl bg-red-500/10 border border-red-500/20"
416-
>
417-
<p className="text-red-300 text-sm">
418-
{state.status === 'error' ? state.message : wsError}
419-
</p>
420-
<button
421-
type="button"
422-
onClick={handleIndexAnother}
423-
className="mt-2 text-xs text-red-400 hover:text-red-300 transition-colors"
424-
>
425-
Try again
426-
</button>
427-
</motion.div>
428-
)}
429-
430440
{/* Upgrade CTA (when limit reached) */}
431441
{remaining <= 0 && (
432442
<div className="mt-6 p-4 rounded-xl bg-gradient-to-r from-indigo-600/20 to-purple-600/20 border border-indigo-500/30 text-center">

‎frontend/src/components/playground/IndexingProgress.tsx‎

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -267,10 +267,19 @@ export function IndexingProgress({
267267
}: IndexingProgressProps) {
268268
const { percent, filesProcessed, filesTotal, currentFile, functionsFound } = progress;
269269

270-
// Estimate remaining time
270+
// Safe values to prevent NaN/division by zero
271+
const safePercent = Number.isFinite(percent) ? Math.max(0, Math.min(100, percent)) : 0;
272+
const safeFilesProcessed = Number.isFinite(filesProcessed) ? Math.max(0, filesProcessed) : 0;
273+
const safeFilesTotal = Number.isFinite(filesTotal) ? Math.max(0, filesTotal) : 0;
274+
const safeFunctionsFound = Number.isFinite(functionsFound) ? Math.max(0, functionsFound) : 0;
275+
276+
// Is this the initial "starting" state?
277+
const isStarting = safeFilesTotal === 0 || (safePercent === 0 && safeFilesProcessed === 0);
278+
279+
// Estimate remaining time (only when we have real data)
271280
const estimatedRemaining = (() => {
272-
if (percent <= 0 || filesProcessed <= 0) return null;
273-
const remainingFiles = Math.ceil((filesProcessed / percent) * (100 - percent));
281+
if (isStarting || safePercent <= 0 || safeFilesProcessed <= 0) return null;
282+
const remainingFiles = Math.ceil((safeFilesProcessed / safePercent) * (100 - safePercent));
274283
return Math.max(1, Math.ceil(remainingFiles * 0.15));
275284
})();
276285

@@ -287,7 +296,7 @@ export function IndexingProgress({
287296
initial={{ opacity: 0, y: 20 }}
288297
animate={{ opacity: 1, y: 0 }}
289298
role="status"
290-
aria-label={`Indexing ${repoName || 'repository'}: ${percent}% complete`}
299+
aria-label={`Indexing ${repoName || 'repository'}: ${safePercent}% complete`}
291300
>
292301
{/* Header */}
293302
<div className="px-5 py-4 border-b border-zinc-800/50 bg-gradient-to-r from-zinc-900 to-zinc-900/50">
@@ -308,12 +317,12 @@ export function IndexingProgress({
308317
</div>
309318
<motion.span
310319
className="text-3xl font-bold text-indigo-400"
311-
key={percent}
320+
key={safePercent}
312321
initial={{ scale: 1.3 }}
313322
animate={{ scale: 1 }}
314323
transition={{ type: 'spring', stiffness: 300 }}
315324
>
316-
{percent}%
325+
{safePercent}%
317326
</motion.span>
318327
</div>
319328

@@ -323,16 +332,29 @@ export function IndexingProgress({
323332

324333
{/* Progress bar */}
325334
<div className="px-5 py-4">
326-
<GlowingProgress value={percent} />
335+
<GlowingProgress value={safePercent} />
327336
</div>
328337

329338
{/* Stats grid */}
330339
<div className="px-5 py-4 bg-zinc-900/50 border-t border-b border-zinc-800/50">
331340
<div className="grid grid-cols-4 gap-4">
332-
<StatCard label="Files" value={`${filesProcessed}/${filesTotal}`} />
333-
<StatCard label="Functions" value={functionsFound} highlight />
334-
<StatCard label="Remaining" value={estimatedRemaining ? `~${estimatedRemaining}s` : '—'} />
335-
<StatCard label="Speed" value={filesProcessed > 0 ? `${Math.round(functionsFound / filesProcessed)}/file` : '—'} />
341+
<StatCard
342+
label="Files"
343+
value={isStarting ? 'Starting...' : `${safeFilesProcessed}/${safeFilesTotal}`}
344+
/>
345+
<StatCard
346+
label="Functions"
347+
value={isStarting ? '—' : safeFunctionsFound}
348+
highlight={!isStarting && safeFunctionsFound > 0}
349+
/>
350+
<StatCard
351+
label="Remaining"
352+
value={estimatedRemaining ? `~${estimatedRemaining}s` : '—'}
353+
/>
354+
<StatCard
355+
label="Speed"
356+
value={safeFilesProcessed > 0 ? `${Math.round(safeFunctionsFound / safeFilesProcessed)}/file` : '—'}
357+
/>
336358
</div>
337359
</div>
338360

‎frontend/src/hooks/useIndexingWebSocket.ts‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,8 +248,18 @@ export function useIndexingWebSocket(
248248
if (jobId) {
249249
connect(jobId);
250250
} else {
251+
// Only cleanup connection, DON'T reset state!
252+
// This preserves completedStats when jobId becomes null after completion
251253
cleanup();
252-
setState(INITIAL_STATE);
254+
// Only reset if we were never completed (e.g., user navigated away during indexing)
255+
setState(prev => {
256+
if (prev.phase === 'completed') {
257+
// Keep completed state - just disconnect
258+
return { ...prev, connectionState: 'disconnected' };
259+
}
260+
// Reset if we were mid-indexing (user cancelled, navigated away, etc.)
261+
return INITIAL_STATE;
262+
});
253263
}
254264
return cleanup;
255265
}, [jobId, connect, cleanup]);

0 commit comments

Comments
 (0)