Limit to 3 questions, revise progress UI

examples/ag-ui-research/backend/agent.py

@@ -76,7 +76,7 @@ def create_agent(
 You work step-by-step with the user to conduct deep research on complex questions.
 
 Your workflow:
-1. When user asks a question, propose a research plan (3-5 sub-questions)
+1. When user asks a question, propose a research plan (exactly 3 sub-questions)
 2. Wait for user approval before proceeding to the search phase
 3. Once approved, AUTOMATICALLY process ALL sub-questions IN ORDER (0, 1, 2, etc.) WITHOUT pausing:
    - For each sub-question:
@@ -123,11 +123,11 @@ Show search scores, explain your reasoning, cite your sources, and involve the u
         )
 
         # Use LLM to decompose the question
-        decompose_prompt = f"""Break down this research question into 3-5 specific sub-questions that would help answer it comprehensively.
+        decompose_prompt = f"""Break down this research question into exactly 3 specific sub-questions that would help answer it comprehensively.
 
 Research Question: {question}
 
-Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?", ...]"""
+Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?", "Question 3?"]"""
 
         response = await ctx.deps.client.ask(decompose_prompt)
 
@@ -142,7 +142,7 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
                 q.strip().lstrip("0123456789.-) ")
                 for q in response.split("\n")
                 if q.strip()
-            ][:5]
+            ][:3]
 
         # Create plan
         plan = [
@@ -256,7 +256,7 @@ Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?",
             "type": search_type,
             "results": results,
         }
-        plan[question_id]["status"] = "done"
+        plan[question_id]["status"] = "searched"
         ctx.deps.state.status = f"Found {len(results)} results"
         print("[AGENT] Search complete, sending state snapshot")
 
@@ -360,6 +360,8 @@ Return a JSON array of insights with format:
         # Add to accumulated insights
         ctx.deps.state.insights.extend(new_insights)
 
+        # Mark question as fully done (searched + analyzed)
+        plan[question_id]["status"] = "done"
         ctx.deps.state.status = f"Extracted {len(new_insights)} insights"
         print("[AGENT] Insights extracted, sending state snapshot")
 

examples/ag-ui-research/frontend/components/StateDisplay.tsx

@@ -98,16 +98,11 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 		}));
 	};
 
-	// Phase indicator
+	// Calculate research progress
-	const phases = [
+	const completedQuestions = state.plan.filter((q) => q.status === "done").length;
-		"idle",
+	const totalQuestions = state.plan.length;
-		"planning",
+	const researchProgress =
-		"searching",
+		totalQuestions > 0 ? (completedQuestions / totalQuestions) * 100 : 0;
-		"analyzing",
-		"evaluating",
-		"done",
-	];
-	const currentPhaseIndex = phases.indexOf(state.phase);
 
 	return (
 		<div
@@ -117,7 +112,7 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 				gap: "1rem",
 			}}
 		>
-			{/* Phase Progress */}
+			{/* Current Phase & Status */}
 			<div
 				style={{
 					background: "white",
@@ -133,45 +128,103 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 						marginBottom: "0.5rem",
 					}}
 				>
-					Progress
+					Current Phase
 				</div>
-				<div style={{ display: "flex", gap: "0.5rem", alignItems: "center" }}>
+				<div style={{ display: "flex", gap: "0.75rem", alignItems: "center" }}>
-					{phases.slice(1).map((phase, idx) => (
+					<div
-						<div key={phase} style={{ display: "flex", alignItems: "center" }}>
+						style={{
-							<div
+							padding: "0.5rem 1rem",
+							background:
+								state.phase === "idle"
+									? "#e2e8f0"
+									: state.phase === "planning"
+										? "#fef3c7"
+										: state.phase === "searching"
+											? "#dbeafe"
+											: state.phase === "analyzing"
+												? "#e0e7ff"
+												: state.phase === "evaluating"
+													? "#fce7f3"
+													: "#d1fae5",
+							color:
+								state.phase === "idle"
+									? "#718096"
+									: state.phase === "planning"
+										? "#92400e"
+										: state.phase === "searching"
+											? "#1e40af"
+											: state.phase === "analyzing"
+												? "#3730a3"
+												: state.phase === "evaluating"
+													? "#9f1239"
+													: "#065f46",
+							borderRadius: "6px",
+							fontSize: "1rem",
+							fontWeight: "700",
+							textTransform: "capitalize",
+						}}
+					>
+						{state.phase}
+					</div>
+					{state.status && (
+						<div
+							style={{
+								fontSize: "0.875rem",
+								color: "#4a5568",
+							}}
+						>
+							{state.status}
+						</div>
+					)}
+				</div>
+				{/* Research Progress Bar */}
+				{totalQuestions > 0 && state.phase !== "idle" && (
+					<div style={{ marginTop: "1rem" }}>
+						<div
+							style={{
+								display: "flex",
+								justifyContent: "space-between",
+								alignItems: "center",
+								marginBottom: "0.5rem",
+							}}
+						>
+							<span
 								style={{
-									padding: "0.25rem 0.75rem",
-									background:
-										idx < currentPhaseIndex
-											? "#48bb78"
-											: idx === currentPhaseIndex
-												? "#4299e1"
-												: "#e2e8f0",
-									color:
-										idx < currentPhaseIndex || idx === currentPhaseIndex
-											? "white"
-											: "#718096",
-									borderRadius: "4px",
 									fontSize: "0.75rem",
-									fontWeight: "600",
+									color: "#718096",
-									textTransform: "capitalize",
 								}}
 							>
-								{phase}
+								Research Progress
-							</div>
+							</span>
-							{idx < phases.length - 2 && (
+							<span
-								<div
+								style={{
-									style={{
+									fontSize: "0.75rem",
-										width: "1rem",
+									fontWeight: "600",
-										height: "2px",
+									color: "#2d3748",
-										background: idx < currentPhaseIndex ? "#48bb78" : "#e2e8f0",
+								}}
-										margin: "0 0.25rem",
+							>
-									}}
+								{completedQuestions}/{totalQuestions} questions
-								/>
+							</span>
-							)}
 						</div>
-					))}
+						<div
-				</div>
+							style={{
+								height: "0.5rem",
+								background: "#e2e8f0",
+								borderRadius: "4px",
+								overflow: "hidden",
+							}}
+						>
+							<div
+								style={{
+									width: `${researchProgress}%`,
+									height: "100%",
+									background: "#48bb78",
+									transition: "width 0.3s ease",
+								}}
+							/>
+						</div>
+					</div>
+				)}
 			</div>
 
 			{/* Question */}
@@ -339,7 +392,7 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 												color:
 													item.status === "done"
 														? "#48bb78"
-														: item.status === "searching"
+														: item.status === "searching" || item.status === "searched"
 															? "#4299e1"
 															: "#a0aec0",
 												flexShrink: 0,
@@ -349,7 +402,9 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 												? "✓"
 												: item.status === "searching"
 													? "🔍"
-													: "⏳"}
+													: item.status === "searched"
+														? "📊"
+														: "⏳"}
 										</div>
 										<div style={{ flex: 1 }}>
 											<div