Skip to content

Commit bcf8989

Browse files
cursoragentvenkateshsakamuri-lab
cursoragent
and
venkateshsakamuri-lab
committed
fix(brain): treat NEEDS_ATTENTION as terminal across init UI
Coverage gate can stop the pipeline at NEEDS_ATTENTION; without this, onboarding and init polls stay in "running" forever. Shared initStage helpers keep COMPLETED/FAILED/NEEDS_ATTENTION/ERROR consistent. Co-authored-by: Venkat SF <venkatesh.sakamuri@stayflexi.com>
1 parent c1b7504 commit bcf8989

8 files changed

Lines changed: 81 additions & 21 deletions

File tree

‎src/components/BrainInitModal.jsx‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { useState, useEffect, useRef, useCallback } from 'react'
22
import { X, CheckCircle, Loader2, XCircle, RotateCcw, ArrowRight } from 'lucide-react'
33
import { connectionAPI } from '@/lib/api/client'
4+
import { isInitRunning, isInitTerminal } from '@/lib/initStage'
45

56
const STAGES = [
67
{ key: 'SCHEMA_SCAN', label: 'Scanning schema', desc: 'Reading tables, columns, and relationships' },
@@ -70,7 +71,7 @@ export default function BrainInitModal({
7071

7172
setHasExistingRun(hasRun)
7273

73-
if (currentStage && !['NONE', 'COMPLETED', 'FAILED', 'ERROR'].includes(currentStage)) {
74+
if (currentStage && isInitRunning(currentStage)) {
7475
setStatus('running')
7576
setActiveStage(currentStage)
7677
setDoneStages(completedStages.filter((stage) => stage !== currentStage))
@@ -123,7 +124,7 @@ export default function BrainInitModal({
123124
try {
124125
const nextStatus = await loadExistingStatus()
125126
const stage = nextStatus?.currentStage || nextStatus?.stage || ''
126-
if (!stage || ['COMPLETED', 'FAILED', 'ERROR', 'NONE'].includes(stage)) {
127+
if (!stage || (stage === 'NONE' || isInitTerminal(stage))) {
127128
stopPolling()
128129
}
129130
} catch {
@@ -163,7 +164,7 @@ export default function BrainInitModal({
163164
} else {
164165
void loadExistingStatus().then((nextStatus) => {
165166
const stage = nextStatus?.currentStage || nextStatus?.stage || ''
166-
if (stage && !['NONE', 'COMPLETED', 'FAILED', 'ERROR'].includes(stage)) {
167+
if (stage && isInitRunning(stage)) {
167168
startPolling()
168169
}
169170
})

‎src/components/InitProgressIndicator.js‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import useInitProgressStore, {
1616
} from "../lib/stores/useInitProgressStore";
1717
import { connectionAPI } from "../lib/api/client";
1818
import InitHistoryModal from "./InitHistoryModal";
19+
import { isInitTerminal } from "../lib/initStage";
1920

2021
const STAGE_LABELS = {
2122
SCHEMA_SCAN: "Scanning schema",
@@ -30,6 +31,7 @@ const STAGE_LABELS = {
3031
SEMANTIC_MODELING: "Modeling semantics",
3132
COMPLETED: "All set!",
3233
FAILED: "Initialization failed",
34+
NEEDS_ATTENTION: "Needs attention",
3335
};
3436

3537
const STAGE_DESCRIPTIONS = {
@@ -246,7 +248,7 @@ export function InitProgressIndicator({ connectionId }) {
246248
errorCountRef.current = 0;
247249
if (data) {
248250
setInitProgress(data);
249-
if (["COMPLETED", "FAILED"].includes(data.currentStage)) {
251+
if (isInitTerminal(data.currentStage)) {
250252
clearInterval(pollingRef.current);
251253
pollingRef.current = null;
252254
}
@@ -350,7 +352,7 @@ export function InitProgressIndicator({ connectionId }) {
350352
errorCountRef.current = 0;
351353
if (data) {
352354
setInitProgress(data);
353-
if (["COMPLETED", "FAILED"].includes(data.currentStage)) {
355+
if (isInitTerminal(data.currentStage)) {
354356
clearInterval(pollingRef.current);
355357
pollingRef.current = null;
356358
}
@@ -403,8 +405,8 @@ export function InitProgressIndicator({ connectionId }) {
403405
useEffect(() => {
404406
if (
405407
prevStageRef.current &&
406-
!["COMPLETED", "FAILED"].includes(prevStageRef.current) &&
407-
["COMPLETED", "FAILED"].includes(stage)
408+
!isInitTerminal(prevStageRef.current) &&
409+
isInitTerminal(stage)
408410
) {
409411
fetchHistory();
410412
}
@@ -433,6 +435,11 @@ export function InitProgressIndicator({ connectionId }) {
433435
chipLabel = "Brain ready";
434436
chipBg = "#f0fdf4";
435437
chipColor = "#15803d";
438+
} else if (stage === "NEEDS_ATTENTION") {
439+
chipIcon = <AlertCircle size={14} style={{ color: "#d97706" }} />;
440+
chipLabel = "Needs attention";
441+
chipBg = "#fffbeb";
442+
chipColor = "#b45309";
436443
} else if (stage === "FAILED") {
437444
chipIcon = <AlertCircle size={14} style={{ color: "#dc2626" }} />;
438445
chipLabel = "Init failed";
@@ -723,7 +730,7 @@ export function InitProgressIndicator({ connectionId }) {
723730
)}
724731

725732
{/* Action buttons — re-init + view history */}
726-
{(stage === "COMPLETED" || stage === "FAILED" || !stage) && (
733+
{(stage === "COMPLETED" || stage === "FAILED" || stage === "NEEDS_ATTENTION" || !stage) && (
727734
<div
728735
style={{
729736
display: "flex",

‎src/components/company-knowledge/BackgroundJobsTab.jsx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,8 @@ export default function BackgroundJobsTab({ connectionId }) {
146146
const initActive = initStatus
147147
&& initStatus.currentStage !== 'COMPLETED'
148148
&& initStatus.currentStage !== 'FAILED'
149+
&& initStatus.currentStage !== 'NEEDS_ATTENTION'
150+
&& initStatus.currentStage !== 'ERROR'
149151
&& !initStatus.completedAt
150152
const intervalMs = (anyRunning || initActive) ? 4000 : 15000
151153
const intervalId = window.setInterval(() => {

‎src/components/company-knowledge/CompanyKnowledgePanel.jsx‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,14 +244,22 @@ export default function CompanyKnowledgePanel({ connectionId }) {
244244
retry: false,
245245
refetchInterval: (query) => {
246246
const stage = query.state.data?.currentStage
247-
return stage && stage !== 'COMPLETED' && stage !== 'FAILED' ? 4000 : false
247+
// Keep polling only while a non-terminal stage is running
248+
// (COMPLETED / FAILED / NEEDS_ATTENTION / NONE stop).
249+
if (!stage || stage === 'NONE' || stage === 'COMPLETED'
250+
|| stage === 'FAILED' || stage === 'NEEDS_ATTENTION' || stage === 'ERROR') {
251+
return false
252+
}
253+
return 4000
248254
},
249255
})
250256

251257
useEffect(() => {
252258
if (userChoseTab || activeTab !== 'background-jobs') {
253259
return
254260
}
261+
// Only auto-advance on full COMPLETED — NEEDS_ATTENTION stays on Initialize
262+
// so the user sees coverage messaging and can re-init.
255263
if (initStatusQuery.data?.currentStage === 'COMPLETED') {
256264
setDefaultTab('schema-context')
257265
}

‎src/components/onboarding/StepBrainInit.jsx‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { useEffect, useState, useRef, useCallback } from 'react'
22
import { CheckCircle, Loader2, XCircle, AlertCircle, RotateCcw, ArrowRight } from 'lucide-react'
33
import { connectionAPI } from '@/lib/api/client'
4+
import { isInitComplete, isInitFailed, isInitNeedsAttention } from '@/lib/initStage'
45

56
const STAGES = [
67
{ key: 'SCHEMA_SCAN', label: 'Scanning schema', desc: 'Reading tables, columns, and relationships' },
@@ -31,11 +32,11 @@ export default function StepBrainInit({ connectionId, onComplete }) {
3132

3233
const stage = s.currentStage || s.stage || ''
3334
const prog = s.progressPercent ?? s.progress ?? 0
34-
const isCompleted = stage === 'COMPLETED' || prog >= 100
35-
const isFailed = stage === 'FAILED' || stage === 'ERROR'
3635
const timings = s.stageTimings || {}
3736

38-
if (isCompleted) {
37+
// Do NOT treat progress>=100 alone as complete — coverage can stop at
38+
// NEEDS_ATTENTION with a high percent while still incomplete.
39+
if (isInitComplete(stage)) {
3940
clearInterval(pollRef.current)
4041
setDoneStages(STAGES.map(st => st.key))
4142
setActiveStage(null)
@@ -45,7 +46,7 @@ export default function StepBrainInit({ connectionId, onComplete }) {
4546
return
4647
}
4748

48-
if (isFailed) {
49+
if (isInitNeedsAttention(stage) || isInitFailed(stage)) {
4950
clearInterval(pollRef.current)
5051
const lastAttempted = [...STAGES].reverse().find(st => timings[st.key]?.startedAt)
5152
const failedKey = lastAttempted?.key
@@ -55,7 +56,12 @@ export default function StepBrainInit({ connectionId, onComplete }) {
5556
setActiveStage(null)
5657
setProgress(prog)
5758
setStatus('error')
58-
setErrorMsg(s.errorMessage || 'Initialization failed')
59+
setErrorMsg(
60+
s.errorMessage
61+
|| (isInitNeedsAttention(stage)
62+
? 'Brain indexed only part of the schema. Fix grants or schema access, then retry.'
63+
: 'Initialization failed'),
64+
)
5965
return
6066
}
6167

‎src/components/tabs/Brain/BrainWorkspace.jsx‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
22
import { Clock3, Database, Loader2, Network, Play, RefreshCw, Sparkles } from 'lucide-react'
33
import BrainInitModal from '@/components/BrainInitModal'
44
import { connectionAPI } from '@/lib/api/client'
5+
import { isInitRunning, isInitFailed, isInitNeedsAttention } from '@/lib/initStage'
56
import { DEFAULT_ROLE_FILTERS, ERD3DErrorBoundary, SchemaDiagramFilter, SchemaERD3D } from './SchemaERD3D'
67
import styles from './BrainWorkspace.module.css'
78

@@ -23,6 +24,7 @@ const STAGE_LABELS = {
2324
...Object.fromEntries(STAGES.map((stage) => [stage.key, stage.label])),
2425
COMPLETED: 'Brain ready',
2526
FAILED: 'Initialization failed',
27+
NEEDS_ATTENTION: 'Needs attention',
2628
}
2729

2830
function stageLabel(stage) {
@@ -97,7 +99,7 @@ export default function BrainWorkspace({ connectionId }) {
9799
return undefined
98100
}
99101
const currentStage = status?.currentStage || status?.stage || 'NONE'
100-
const isRunning = currentStage && !['NONE', 'COMPLETED', 'FAILED'].includes(currentStage)
102+
const isRunning = isInitRunning(currentStage)
101103
const intervalMs = isRunning ? 4000 : 15000
102104
const intervalId = window.setInterval(() => {
103105
void loadWorkspace({ silent: true })
@@ -111,11 +113,11 @@ export default function BrainWorkspace({ connectionId }) {
111113
)
112114

113115
const currentStage = status?.currentStage || status?.stage || 'NONE'
114-
const isRunning = currentStage && !['NONE', 'COMPLETED', 'FAILED'].includes(currentStage)
116+
const isRunning = isInitRunning(currentStage)
115117
const isReady = currentStage === 'COMPLETED' || hasCompletedInit
116118
const progress = Math.max(0, Math.min(100, Number(status?.progressPercent ?? status?.progress ?? (isReady ? 100 : 0))))
117-
const ctaLabel = isRunning ? 'View current status' : isReady ? 'View completed stages' : currentStage === 'FAILED' ? 'View failure details' : 'Initialize Brain'
118-
const primaryAutoStart = !isRunning && !isReady && currentStage !== 'FAILED'
119+
const ctaLabel = isRunning ? 'View current status' : isReady ? 'View completed stages' : (isInitFailed(currentStage) || isInitNeedsAttention(currentStage)) ? 'View details' : 'Initialize Brain'
120+
const primaryAutoStart = !isRunning && !isReady && !isInitFailed(currentStage) && !isInitNeedsAttention(currentStage)
119121
const statusText = isRunning
120122
? `${stageLabel(currentStage)} · ${progress}%`
121123
: currentStage === 'FAILED'
@@ -308,14 +310,14 @@ export default function BrainWorkspace({ connectionId }) {
308310
{isRunning ? <Loader2 size={15} className={styles.spinningIcon} /> : <Sparkles size={15} />}
309311
{ctaLabel}
310312
</button>
311-
{(!isRunning && (isReady || currentStage === 'FAILED' || currentStage === 'NONE')) && (
313+
{(!isRunning && (isReady || isInitFailed(currentStage) || isInitNeedsAttention(currentStage) || currentStage === 'NONE')) && (
312314
<button
313315
className={styles.secondaryButton}
314316
onClick={() => openStatusModal(true)}
315317
disabled={loading}
316318
>
317319
<RefreshCw size={15} />
318-
{isReady ? 'Refresh Brain' : currentStage === 'FAILED' ? 'Retry Brain init' : 'Start Brain init'}
320+
{isReady ? 'Refresh Brain' : (isInitFailed(currentStage) || isInitNeedsAttention(currentStage)) ? 'Retry Brain init' : 'Start Brain init'}
319321
</button>
320322
)}
321323
{(!isRunning && !forceRebuilding) && (

‎src/lib/initStage.js‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
/**
2+
* Brain init stage helpers — keep in sync with
3+
* com.dbaagent.model.InitStage#isTerminal().
4+
*/
5+
6+
const TERMINAL_STAGES = new Set([
7+
'COMPLETED',
8+
'FAILED',
9+
'ERROR',
10+
'NEEDS_ATTENTION',
11+
])
12+
13+
/** Pipeline finished (success, failure, or incomplete coverage). Stop polling. */
14+
export function isInitTerminal(stage) {
15+
return !stage || TERMINAL_STAGES.has(stage)
16+
}
17+
18+
/** Still running a non-terminal stage (includes NONE? no — NONE is idle). */
19+
export function isInitRunning(stage) {
20+
return Boolean(stage) && stage !== 'NONE' && !TERMINAL_STAGES.has(stage)
21+
}
22+
23+
export function isInitComplete(stage) {
24+
return stage === 'COMPLETED'
25+
}
26+
27+
export function isInitNeedsAttention(stage) {
28+
return stage === 'NEEDS_ATTENTION'
29+
}
30+
31+
export function isInitFailed(stage) {
32+
return stage === 'FAILED' || stage === 'ERROR'
33+
}

‎src/lib/stores/useInitProgressStore.js‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { create } from "zustand";
22
import { shallow } from "zustand/shallow";
3+
import { isInitRunning } from "@/lib/initStage";
34

45
const useInitProgressStore = create((set) => ({
56
connectionId: null,
@@ -66,6 +67,6 @@ export const useIsInitActive = (connectionId) =>
6667
(s) =>
6768
!!s.stage &&
6869
s.connectionId === connectionId &&
69-
!["COMPLETED", "FAILED"].includes(s.stage),
70+
isInitRunning(s.stage),
7071
);
7172
export default useInitProgressStore;

0 commit comments

Comments
 (0)