More fixes

Author: cte

src/services/deep-research/DeepResearchService.ts

@@ -20,6 +20,7 @@ import {
 	ResearchQuery,
 } from "./types"
 import { truncatePrompt, trimPrompt } from "./utils/prompt"
+import { getTreeSize } from "./utils/progress"
 
 export class DeepResearchService {
 	public readonly providerId: string
@@ -183,63 +184,78 @@ export class DeepResearchService {
 		const results = await Promise.all(
 			queries.map(({ query, researchGoal }) =>
 				limit(async () => {
-					try {
-						if (this.isAborted()) {
-							return { learnings, visitedUrls }
-						}
+					if (this.isAborted()) {
+						return { learnings, visitedUrls }
+					}
+
+					let result: SearchResponse
 
-						const result = await this.firecrawl.search(query, {
+					try {
+						result = await this.firecrawl.search(query, {
 							timeout: 15000,
 							limit: 5,
 							scrapeOptions: { formats: ["markdown"] },
 						})
+					} catch (e) {
+						const text = e instanceof Error ? e.message : "Unknown error"
+						console.log(`[deepResearch] error = ${text}`)
 
-						const newUrls = result.data
-							.map(({ url }) => url)
-							.filter((url): url is string => url !== undefined)
-
-						const newBreadth = Math.ceil(breadth / 2)
-						const newDepth = depth - 1
-						const newLearnings = await this.extractLearnings({ query, result, breadth: newBreadth })
-						const allLearnings = [...learnings, ...newLearnings.learnings]
-						const allUrls = [...visitedUrls, ...newUrls]
-						onExtractedLearnings({ ...newLearnings, urls: newUrls })
-
-						this.progress.completedQueries = this.progress.completedQueries + 1
-						onProgressUpdated()
-
-						if (newDepth === 0) {
-							return { learnings: allLearnings, visitedUrls: allUrls }
-						}
-
-						console.log(`[deepResearch] researching deeper, breadth: ${newBreadth}, depth: ${newDepth}`)
-
-						const nextQuery = trimPrompt(`
-							Previous research goal: ${researchGoal}
-							Follow-up research directions: ${newLearnings.followUpQuestions.map((q) => `\n${q}`).join("")}
-						`)
-
-						return this.deepResearch({
-							query: nextQuery,
-							breadth: newBreadth,
-							depth: newDepth,
-							learnings: allLearnings,
-							visitedUrls: allUrls,
-							onProgressUpdated,
-							onGeneratedQueries,
-							onExtractedLearnings,
+						await this.postMessage({
+							type: "research.error",
+							text: `Encountered an error while crawling "${query}": ${text}`,
 						})
+
+						return { learnings, visitedUrls }
+					}
+
+					const newUrls = result.data.map(({ url }) => url).filter((url): url is string => url !== undefined)
+
+					const newBreadth = Math.ceil(breadth / 2)
+					const newDepth = depth - 1
+					let newLearnings: ResearchLearnings
+
+					try {
+						newLearnings = await this.extractLearnings({ query, result, breadth: newBreadth })
 					} catch (e) {
 						const text = e instanceof Error ? e.message : "Unknown error"
 						console.log(`[deepResearch] error = ${text}`)
 
 						await this.postMessage({
 							type: "research.error",
-							text: `Encountered an error while researching "${query}": ${text}`,
+							text: `Encountered an error while extracting learnings from "${query}": ${text}`,
 						})
 
-						return { learnings: [], visitedUrls: [] }
+						return { learnings, visitedUrls }
 					}
+
+					const allLearnings = [...learnings, ...newLearnings.learnings]
+					const allUrls = [...visitedUrls, ...newUrls]
+					onExtractedLearnings({ ...newLearnings, urls: newUrls })
+
+					this.progress.completedQueries = this.progress.completedQueries + 1
+					onProgressUpdated()
+
+					if (newDepth === 0) {
+						return { learnings: allLearnings, visitedUrls: allUrls }
+					}
+
+					console.log(`[deepResearch] researching deeper, breadth: ${newBreadth}, depth: ${newDepth}`)
+
+					const nextQuery = trimPrompt(`
+						Previous research goal: ${researchGoal}
+						Follow-up research directions: ${newLearnings.followUpQuestions.map((q) => `\n${q}`).join("")}
+					`)
+
+					return this.deepResearch({
+						query: nextQuery,
+						breadth: newBreadth,
+						depth: newDepth,
+						learnings: allLearnings,
+						visitedUrls: allUrls,
+						onProgressUpdated,
+						onGeneratedQueries,
+						onExtractedLearnings,
+					})
 				}),
 			),
 		)
@@ -326,7 +342,7 @@ export class DeepResearchService {
 				schema,
 			})
 
-			console.log(`[generateQueries] generated ${queries.length} queries`, queries)
+			console.log(`[generateQueries] generated ${queries.length} (out of ${breadth}) queries`, queries)
 
 			return queries.slice(0, breadth)
 		} catch (error) {
@@ -530,15 +546,14 @@ export class DeepResearchService {
 				}),
 			})
 
-		// Calculate total expected queries across all depth levels.
-		// At each level, the breadth is halved, so level 1 has full breadth,
-		// level 2 has breadth/2, level 3 has breadth/4, etc.
-		for (let i = this.depth; i > 0; i--) {
-			this.progress.expectedQueries += Math.ceil(this.breadth / Math.pow(2, this.depth - i))
-		}
-
+		this.progress.expectedQueries = getTreeSize({ breadth: this.breadth, depth: this.depth })
 		onProgressUpdated()
 
+		console.log(`[transitionToResearch] query = ${query}`)
+		console.log(`[transitionToResearch] breadth = ${this.breadth}`)
+		console.log(`[transitionToResearch] depth = ${this.depth}`)
+		console.log(`[transitionToResearch] expectedQueries = ${this.progress.expectedQueries}`)
+
 		const { learnings, visitedUrls } = await this.withLoading(
 			() =>
 				this.deepResearch({

src/services/deep-research/__tests__/utils/progress.test.ts

@@ -0,0 +1,19 @@
+// npx jest src/services/deep-research/__tests__/utils/progress.test.ts
+
+import { getTreeSize } from "../../utils/progress"
+
+describe("getTreeSize", () => {
+	it("should calculate the correct number of expected queries", () => {
+		// Typical cases.
+		expect(getTreeSize({ breadth: 2, depth: 0 })).toBe(2)
+		expect(getTreeSize({ breadth: 3, depth: 1 })).toBe(7)
+		expect(getTreeSize({ breadth: 4, depth: 2 })).toBe(12)
+		expect(getTreeSize({ breadth: 5, depth: 3 })).toBe(27)
+
+		// Try each minimum and maximum for breadth and depth.
+		expect(getTreeSize({ breadth: 1, depth: 0 })).toBe(1)
+		expect(getTreeSize({ breadth: 10, depth: 0 })).toBe(10)
+		expect(getTreeSize({ breadth: 1, depth: 9 })).toBe(1023)
+		expect(getTreeSize({ breadth: 10, depth: 9 })).toBe(1056)
+	})
+})

src/services/deep-research/utils/progress.ts

@@ -0,0 +1,17 @@
+// Calculate total expected queries across all depth levels.
+// At each level, the breadth is halved, so level 1 has full breadth,
+// level 2 has breadth/2, level 3 has breadth/4, etc.
+// For breadth = 4, depth = 2, the expected queries are:
+// D2: 2^2 * 1 = 4
+// D1: 2^1 * 2 = 4
+// D0: 2^0 * 4 = 4
+// Total: 12
+export const getTreeSize = ({ breadth, depth }: { breadth: number; depth: number }) => {
+	let value = 0
+
+	for (let i = depth; i >= 0; i--) {
+		value = value + Math.pow(2, i) * Math.ceil(breadth / Math.pow(2, i))
+	}
+
+	return value
+}

webview-ui/src/features/deep-research/DeepResearch.tsx

@@ -23,7 +23,7 @@ export const DeepResearch = ({ isHidden }: DeepResearchProps) => {
 
 	return (
 		<div
-			className={cn("lex flex-col items-center justify-center h-full gap-4 overflow-y-auto py-4", {
+			className={cn("flex flex-col items-center justify-center h-full min-w-64 gap-4 overflow-auto py-4", {
 				hidden: isHidden,
 			})}>
 			<GetStarted />