feat: implement AI Deep Research with real-time SSE support

- Add aiDeepResearchTool.ts for handling AI deep research requests
- Create AIDeepResearchService.ts for SSE communication with server
- Add AIDeepResearchBlock.tsx UI component for displaying research progress
- Update tool types and registration in shared/tools.ts
- Add ai_deep_research to ClineSayTool interface
- Add ai_deep_research_result to ClineSay types
- Update presentAssistantMessage.ts to handle the new tool
- Add UI integration in ChatRow.tsx
- Add translation keys for AI Deep Research
- Add comprehensive tests for aiDeepResearchTool

This implementation provides real-time streaming of AI research progress
including thinking, searching, reading, and analyzing states.
This commit is contained in:
Roo Code 2025-07-20 01:27:32 +00:00
parent 1b12108172
commit 08c6146b3b
11 changed files with 711 additions and 1 deletions

View file

@ -106,6 +106,7 @@ export const clineSays = [
"condense_context",
"condense_context_error",
"codebase_search_result",
"ai_deep_research_result",
"user_edit_todos",
] as const

View file

@ -34,6 +34,7 @@ export const toolNames = [
"fetch_instructions",
"codebase_search",
"update_todo_list",
"ai_deep_research",
] as const
export const toolNamesSchema = z.enum(toolNames)

View file

@ -32,6 +32,7 @@ import { formatResponse } from "../prompts/responses"
import { validateToolUse } from "../tools/validateToolUse"
import { Task } from "../task/Task"
import { codebaseSearchTool } from "../tools/codebaseSearchTool"
import { aiDeepResearchTool } from "../tools/aiDeepResearchTool"
import { experiments, EXPERIMENT_IDS } from "../../shared/experiments"
import { applyDiffToolLegacy } from "../tools/applyDiffTool"
@ -204,7 +205,9 @@ export async function presentAssistantMessage(cline: Task) {
return `[${block.name}]`
case "switch_mode":
return `[${block.name} to '${block.params.mode_slug}'${block.params.reason ? ` because: ${block.params.reason}` : ""}]`
case "codebase_search": // Add case for the new tool
case "codebase_search":
return `[${block.name} for '${block.params.query}']`
case "ai_deep_research":
return `[${block.name} for '${block.params.query}']`
case "update_todo_list":
return `[${block.name}]`
@ -462,6 +465,9 @@ export async function presentAssistantMessage(cline: Task) {
case "codebase_search":
await codebaseSearchTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
break
case "ai_deep_research":
await aiDeepResearchTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
break
case "list_code_definition_names":
await listCodeDefinitionNamesTool(
cline,

View file

@ -0,0 +1,231 @@
import { describe, it, expect, vi, beforeEach } from "vitest"
import { aiDeepResearchTool } from "../aiDeepResearchTool"
import { Task } from "../../task/Task"
import { AIDeepResearchService } from "../../../services/ai-deep-research/AIDeepResearchService"
// Mock the AIDeepResearchService
vi.mock("../../../services/ai-deep-research/AIDeepResearchService")
describe("aiDeepResearchTool", () => {
let mockCline: any
let mockAskApproval: any
let mockHandleError: any
let mockPushToolResult: any
let mockRemoveClosingTag: any
let mockPerformResearch: any
beforeEach(() => {
vi.clearAllMocks()
// Mock the Task instance
mockCline = {
say: vi.fn(),
ask: vi.fn().mockResolvedValue(undefined),
sayAndCreateMissingParamError: vi.fn().mockResolvedValue("Missing parameter error"),
consecutiveMistakeCount: 0,
providerRef: {
deref: vi.fn().mockReturnValue({
context: {},
}),
},
}
// Mock the callback functions
mockAskApproval = vi.fn().mockResolvedValue(true)
mockHandleError = vi.fn()
mockPushToolResult = vi.fn()
mockRemoveClosingTag = vi.fn((tag, content) => content || "")
// Mock AIDeepResearchService
mockPerformResearch = vi.fn().mockResolvedValue("Research completed successfully")
AIDeepResearchService.prototype.performResearch = mockPerformResearch
})
it("should handle missing query parameter", async () => {
const block = {
type: "tool_use" as const,
name: "ai_deep_research" as const,
params: {},
partial: false,
}
await aiDeepResearchTool(
mockCline,
block,
mockAskApproval,
mockHandleError,
mockPushToolResult,
mockRemoveClosingTag,
)
expect(mockCline.consecutiveMistakeCount).toBe(1)
expect(mockPushToolResult).toHaveBeenCalledWith("Missing parameter error")
expect(mockCline.sayAndCreateMissingParamError).toHaveBeenCalledWith("ai_deep_research", "query")
})
it("should handle partial block", async () => {
const block = {
type: "tool_use" as const,
name: "ai_deep_research" as const,
params: { query: "test query" },
partial: true,
}
await aiDeepResearchTool(
mockCline,
block,
mockAskApproval,
mockHandleError,
mockPushToolResult,
mockRemoveClosingTag,
)
expect(mockCline.ask).toHaveBeenCalledWith(
"tool",
JSON.stringify({
tool: "aiDeepResearch",
query: "test query",
}),
true,
)
expect(mockAskApproval).not.toHaveBeenCalled()
})
it("should handle user rejection", async () => {
mockAskApproval.mockResolvedValue(false)
const block = {
type: "tool_use" as const,
name: "ai_deep_research" as const,
params: { query: "test query" },
partial: false,
}
await aiDeepResearchTool(
mockCline,
block,
mockAskApproval,
mockHandleError,
mockPushToolResult,
mockRemoveClosingTag,
)
expect(mockAskApproval).toHaveBeenCalled()
expect(mockPushToolResult).toHaveBeenCalledWith("The user denied this operation.")
expect(mockPerformResearch).not.toHaveBeenCalled()
})
it("should perform research successfully", async () => {
const block = {
type: "tool_use" as const,
name: "ai_deep_research" as const,
params: { query: "test query" },
partial: false,
}
await aiDeepResearchTool(
mockCline,
block,
mockAskApproval,
mockHandleError,
mockPushToolResult,
mockRemoveClosingTag,
)
expect(mockAskApproval).toHaveBeenCalled()
expect(mockCline.say).toHaveBeenCalledWith(
"ai_deep_research_result",
expect.stringContaining('"status":"thinking"'),
)
expect(mockPerformResearch).toHaveBeenCalledWith("test query", expect.any(Object))
expect(mockPushToolResult).toHaveBeenCalledWith("Research completed successfully")
})
it("should handle errors during research", async () => {
const error = new Error("Research failed")
mockPerformResearch.mockRejectedValue(error)
const block = {
type: "tool_use" as const,
name: "ai_deep_research" as const,
params: { query: "test query" },
partial: false,
}
await aiDeepResearchTool(
mockCline,
block,
mockAskApproval,
mockHandleError,
mockPushToolResult,
mockRemoveClosingTag,
)
expect(mockHandleError).toHaveBeenCalledWith("ai_deep_research", error)
})
it("should handle missing context", async () => {
mockCline.providerRef.deref.mockReturnValue(null)
const block = {
type: "tool_use" as const,
name: "ai_deep_research" as const,
params: { query: "test query" },
partial: false,
}
await aiDeepResearchTool(
mockCline,
block,
mockAskApproval,
mockHandleError,
mockPushToolResult,
mockRemoveClosingTag,
)
expect(mockHandleError).toHaveBeenCalledWith(
"ai_deep_research",
expect.objectContaining({
message: "Extension context is not available.",
}),
)
})
it("should call all callbacks during research", async () => {
let capturedCallbacks: any = {}
mockPerformResearch.mockImplementation(async (query: string, callbacks: any) => {
capturedCallbacks = callbacks
// Simulate calling each callback
await callbacks.onThinking("Thinking about the query...")
await callbacks.onSearching("machine learning")
await callbacks.onReading("https://example.com/article")
await callbacks.onAnalyzing("Analyzing the content...")
await callbacks.onResult("Final research result")
return "Research completed successfully"
})
const block = {
type: "tool_use" as const,
name: "ai_deep_research" as const,
params: { query: "test query" },
partial: false,
}
await aiDeepResearchTool(
mockCline,
block,
mockAskApproval,
mockHandleError,
mockPushToolResult,
mockRemoveClosingTag,
)
// Verify all status updates were sent
const sayCalls = mockCline.say.mock.calls
expect(sayCalls.some((call: any[]) => call[1].includes('"status":"thinking"'))).toBe(true)
expect(sayCalls.some((call: any[]) => call[1].includes('"status":"searching"'))).toBe(true)
expect(sayCalls.some((call: any[]) => call[1].includes('"status":"reading"'))).toBe(true)
expect(sayCalls.some((call: any[]) => call[1].includes('"status":"analyzing"'))).toBe(true)
expect(sayCalls.some((call: any[]) => call[1].includes('"status":"completed"'))).toBe(true)
})
})

View file

@ -0,0 +1,123 @@
import { Task } from "../task/Task"
import { AIDeepResearchService } from "../../services/ai-deep-research/AIDeepResearchService"
import { formatResponse } from "../prompts/responses"
import { AskApproval, HandleError, PushToolResult, RemoveClosingTag, ToolUse } from "../../shared/tools"
import { ClineSayTool } from "../../shared/ExtensionMessage"
export async function aiDeepResearchTool(
cline: Task,
block: ToolUse,
askApproval: AskApproval,
handleError: HandleError,
pushToolResult: PushToolResult,
removeClosingTag: RemoveClosingTag,
) {
const toolName = "ai_deep_research"
// --- Parameter Extraction and Validation ---
let query: string | undefined = block.params.query
query = removeClosingTag("query", query)
const sharedMessageProps: ClineSayTool = {
tool: "aiDeepResearch",
query: query,
}
if (block.partial) {
await cline.ask("tool", JSON.stringify(sharedMessageProps), block.partial).catch(() => {})
return
}
if (!query) {
cline.consecutiveMistakeCount++
pushToolResult(await cline.sayAndCreateMissingParamError(toolName, "query"))
return
}
const didApprove = await askApproval("tool", JSON.stringify(sharedMessageProps))
if (!didApprove) {
pushToolResult(formatResponse.toolDenied())
return
}
cline.consecutiveMistakeCount = 0
// --- Core Logic ---
try {
const context = cline.providerRef.deref()?.context
if (!context) {
throw new Error("Extension context is not available.")
}
// Initialize the AI Deep Research Service
const service = new AIDeepResearchService(context)
// Send initial status to UI
const initialStatus = {
tool: "aiDeepResearch",
query: query,
status: "thinking",
content: "",
}
await cline.say("ai_deep_research_result", JSON.stringify(initialStatus))
// Start the research with SSE streaming
const result = await service.performResearch(query, {
onThinking: async (thought: string) => {
// Send thinking updates to UI
const thinkingStatus = {
tool: "aiDeepResearch",
query: query,
status: "thinking",
content: thought,
}
await cline.say("ai_deep_research_result", JSON.stringify(thinkingStatus))
},
onSearching: async (searchQuery: string) => {
// Send search status to UI
const searchStatus = {
tool: "aiDeepResearch",
query: query,
status: "searching",
content: searchQuery,
}
await cline.say("ai_deep_research_result", JSON.stringify(searchStatus))
},
onReading: async (url: string) => {
// Send reading status to UI
const readingStatus = {
tool: "aiDeepResearch",
query: query,
status: "reading",
content: url,
}
await cline.say("ai_deep_research_result", JSON.stringify(readingStatus))
},
onAnalyzing: async (content: string) => {
// Send analyzing status to UI
const analyzingStatus = {
tool: "aiDeepResearch",
query: query,
status: "analyzing",
content: content,
}
await cline.say("ai_deep_research_result", JSON.stringify(analyzingStatus))
},
onResult: async (finalResult: string) => {
// Send final result to UI
const resultStatus = {
tool: "aiDeepResearch",
query: query,
status: "completed",
content: finalResult,
}
await cline.say("ai_deep_research_result", JSON.stringify(resultStatus))
},
})
// Push the final result to the AI
pushToolResult(result)
} catch (error: any) {
await handleError(toolName, error)
}
}

View file

@ -0,0 +1,151 @@
import * as vscode from "vscode"
export interface AIDeepResearchCallbacks {
onThinking?: (thought: string) => Promise<void>
onSearching?: (query: string) => Promise<void>
onReading?: (url: string) => Promise<void>
onAnalyzing?: (content: string) => Promise<void>
onResult?: (result: string) => Promise<void>
}
export interface SSEEvent {
type: "thinking" | "searching" | "reading" | "analyzing" | "result" | "error"
content: string
}
export class AIDeepResearchService {
private context: vscode.ExtensionContext
private serverUrl: string
constructor(context: vscode.ExtensionContext) {
this.context = context
// Get server URL from configuration or use default
const config = vscode.workspace.getConfiguration("roo-code")
this.serverUrl = config.get<string>("aiDeepResearchServerUrl") || "https://node-deepresearch-ai.onrender.com"
}
async performResearch(query: string, callbacks: AIDeepResearchCallbacks): Promise<string> {
const endpoint = `${this.serverUrl}/v1/chat/completions`
const requestBody = {
model: "jina-deepsearch-v2",
messages: [
{
role: "user",
content: query,
},
],
stream: true,
}
try {
const response = await fetch(endpoint, {
method: "POST",
headers: {
"Content-Type": "application/json",
Accept: "text/event-stream",
},
body: JSON.stringify(requestBody),
})
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`)
}
if (!response.body) {
throw new Error("Response body is null")
}
// Process SSE stream manually
const reader = response.body.pipeThrough(new TextDecoderStream()).getReader()
let buffer = ""
let fullResult = ""
let currentThought = ""
while (true) {
const { done, value } = await reader.read()
if (done) break
buffer += value
const lines = buffer.split("\n")
buffer = lines.pop() || ""
for (const line of lines) {
if (line.trim() === "") continue
if (line.startsWith("data: ")) {
const data = line.slice(6)
if (data === "[DONE]") {
break
}
try {
const parsed = JSON.parse(data)
const content = parsed.choices?.[0]?.delta?.content
if (content) {
// Parse the content to determine the event type
const event = this.parseEventFromContent(content)
switch (event.type) {
case "thinking":
currentThought += event.content
if (callbacks.onThinking) {
await callbacks.onThinking(currentThought)
}
break
case "searching":
if (callbacks.onSearching) {
await callbacks.onSearching(event.content)
}
break
case "reading":
if (callbacks.onReading) {
await callbacks.onReading(event.content)
}
break
case "analyzing":
if (callbacks.onAnalyzing) {
await callbacks.onAnalyzing(event.content)
}
break
case "result":
fullResult += event.content
if (callbacks.onResult) {
await callbacks.onResult(fullResult)
}
break
}
}
} catch (error) {
console.error("Error parsing SSE data:", error)
}
}
}
}
return fullResult || "Research completed but no results were returned."
} catch (error) {
console.error("AI Deep Research error:", error)
throw error
}
}
private parseEventFromContent(content: string): SSEEvent {
// Simple parsing logic - in a real implementation, the server would send structured events
// For now, we'll use heuristics to determine the event type
if (content.includes("thinking") || content.includes("analyzing")) {
return { type: "thinking", content }
} else if (content.includes("searching") || content.includes("query")) {
return { type: "searching", content }
} else if (content.includes("reading") || content.includes("URL") || content.includes("http")) {
return { type: "reading", content }
} else if (content.includes("found") || content.includes("result")) {
return { type: "result", content }
} else {
// Default to thinking for general content
return { type: "thinking", content }
}
}
}

View file

@ -302,6 +302,7 @@ export interface ClineSayTool {
| "finishTask"
| "searchAndReplace"
| "insertContent"
| "aiDeepResearch"
path?: string
diff?: string
content?: string
@ -338,6 +339,7 @@ export interface ClineSayTool {
}>
}>
question?: string
status?: "thinking" | "searching" | "reading" | "analyzing" | "completed"
}
// Must keep in sync with system prompt.

View file

@ -164,6 +164,11 @@ export interface SearchAndReplaceToolUse extends ToolUse {
Partial<Pick<Record<ToolParamName, string>, "use_regex" | "ignore_case" | "start_line" | "end_line">>
}
export interface AiDeepResearchToolUse extends ToolUse {
name: "ai_deep_research"
params: Partial<Pick<Record<ToolParamName, string>, "query">>
}
// Define tool group configuration
export type ToolGroupConfig = {
tools: readonly string[]
@ -190,6 +195,7 @@ export const TOOL_DISPLAY_NAMES: Record<ToolName, string> = {
search_and_replace: "search and replace",
codebase_search: "codebase search",
update_todo_list: "update todo list",
ai_deep_research: "ai deep research",
} as const
// Define available tool groups.
@ -202,6 +208,7 @@ export const TOOL_GROUPS: Record<ToolGroup, ToolGroupConfig> = {
"list_files",
"list_code_definition_names",
"codebase_search",
"ai_deep_research",
],
},
edit: {

View file

@ -0,0 +1,134 @@
import React, { useState, useEffect } from "react"
import { useTranslation } from "react-i18next"
import { MagnifyingGlassIcon, ReaderIcon, LightningBoltIcon, CheckCircledIcon } from "@radix-ui/react-icons"
import MarkdownBlock from "../common/MarkdownBlock"
interface AIDeepResearchBlockProps {
query: string
status?: "thinking" | "searching" | "reading" | "analyzing" | "completed"
content?: string
result?: string
}
const AIDeepResearchBlock: React.FC<AIDeepResearchBlockProps> = ({ query, status, content, result }) => {
const { t } = useTranslation("chat")
const [isExpanded, setIsExpanded] = useState(true)
const [displayContent, setDisplayContent] = useState("")
useEffect(() => {
if (content) {
setDisplayContent(content)
}
}, [content])
const getStatusIcon = () => {
switch (status) {
case "thinking":
return <LightningBoltIcon className="w-4 h-4 animate-pulse" />
case "searching":
return <MagnifyingGlassIcon className="w-4 h-4 animate-spin" />
case "reading":
return <ReaderIcon className="w-4 h-4 animate-pulse" />
case "analyzing":
return <LightningBoltIcon className="w-4 h-4 animate-pulse" />
case "completed":
return <CheckCircledIcon className="w-4 h-4 text-green-500" />
default:
return null
}
}
const getStatusText = () => {
switch (status) {
case "thinking":
return t("aiDeepResearch.thinking", "Thinking...")
case "searching":
return t("aiDeepResearch.searching", "Searching the web...")
case "reading":
return t("aiDeepResearch.reading", "Reading sources...")
case "analyzing":
return t("aiDeepResearch.analyzing", "Analyzing information...")
case "completed":
return t("aiDeepResearch.completed", "Research completed")
default:
return t("aiDeepResearch.initializing", "Initializing research...")
}
}
return (
<div className="flex flex-col gap-2 my-2">
<div className="bg-vscode-editor-background border border-vscode-border rounded-xs overflow-hidden">
<div
className="flex items-center justify-between gap-2 px-3 py-2 cursor-pointer hover:bg-vscode-list-hoverBackground"
onClick={() => setIsExpanded(!isExpanded)}>
<div className="flex items-center gap-2 flex-1">
{getStatusIcon()}
<span className="font-medium text-vscode-foreground">
{t("aiDeepResearch.title", "AI Deep Research")}
</span>
<span className="text-vscode-descriptionForeground text-sm">{getStatusText()}</span>
</div>
<span className={`codicon codicon-chevron-${isExpanded ? "up" : "down"}`}></span>
</div>
{isExpanded && (
<div className="border-t border-vscode-border">
<div className="px-3 py-2">
<div className="text-sm text-vscode-descriptionForeground mb-2">
<strong>{t("aiDeepResearch.query", "Query")}:</strong> {query}
</div>
{status === "thinking" && displayContent && (
<div className="bg-vscode-editor-inactiveSelectionBackground rounded p-2 mb-2">
<div className="text-xs text-vscode-descriptionForeground mb-1">
{t("aiDeepResearch.thoughtProcess", "Thought Process")}
</div>
<div className="text-sm text-vscode-foreground whitespace-pre-wrap">
{displayContent}
</div>
</div>
)}
{status === "searching" && displayContent && (
<div className="text-sm text-vscode-foreground">
<span className="text-vscode-textLink-foreground">
{t("aiDeepResearch.searchingFor", "Searching for")}: {displayContent}
</span>
</div>
)}
{status === "reading" && displayContent && (
<div className="text-sm text-vscode-foreground">
<span className="text-vscode-textLink-foreground">
{t("aiDeepResearch.readingUrl", "Reading")}: {displayContent}
</span>
</div>
)}
{status === "analyzing" && displayContent && (
<div className="text-sm text-vscode-foreground">
<span className="text-vscode-descriptionForeground">
{t("aiDeepResearch.analyzingContent", "Analyzing content...")}
</span>
</div>
)}
{status === "completed" && result && (
<div className="mt-3">
<div className="text-xs text-vscode-descriptionForeground mb-2">
{t("aiDeepResearch.results", "Research Results")}
</div>
<div className="bg-vscode-editor-inactiveSelectionBackground rounded p-3">
<MarkdownBlock markdown={result} />
</div>
</div>
)}
</div>
</div>
)}
</div>
</div>
)
}
export default AIDeepResearchBlock

View file

@ -46,6 +46,7 @@ import { CommandExecutionError } from "./CommandExecutionError"
import { AutoApprovedRequestLimitWarning } from "./AutoApprovedRequestLimitWarning"
import { CondenseContextErrorRow, CondensingContextRow, ContextCondenseRow } from "./ContextCondenseRow"
import CodebaseSearchResultsDisplay from "./CodebaseSearchResultsDisplay"
import AIDeepResearchBlock from "./AIDeepResearchBlock"
interface ChatRowProps {
message: ClineMessage
@ -493,6 +494,15 @@ export const ChatRowContent = ({
</div>
)
}
case "aiDeepResearch":
return (
<AIDeepResearchBlock
query={tool.query || ""}
status={tool.status}
content={tool.content}
result={tool.content}
/>
)
case "updateTodoList" as any: {
const todos = (tool as any).todos || []
return (
@ -1201,6 +1211,35 @@ export const ChatRowContent = ({
const { results = [] } = parsed?.content || {}
return <CodebaseSearchResultsDisplay results={results} />
case "ai_deep_research_result":
let aiParsed: {
tool: string
query: string
status?: "thinking" | "searching" | "reading" | "analyzing" | "completed"
content?: string
} | null = null
try {
if (message.text) {
aiParsed = JSON.parse(message.text)
}
} catch (error) {
console.error("Failed to parse ai_deep_research_result content:", error)
}
if (!aiParsed || aiParsed.tool !== "aiDeepResearch") {
console.error("Invalid ai_deep_research_result content structure:", aiParsed)
return <div>Error displaying AI Deep Research results.</div>
}
return (
<AIDeepResearchBlock
query={aiParsed.query}
status={aiParsed.status}
content={aiParsed.content}
result={aiParsed.status === "completed" ? aiParsed.content : undefined}
/>
)
case "user_edit_todos":
return <UpdateTodoListToolBlock userEdited onChange={() => {}} />
default:

View file

@ -210,6 +210,21 @@
"didSearch_other": "Found {{count}} results",
"resultTooltip": "Similarity score: {{score}} (click to open file)"
},
"aiDeepResearch": {
"title": "AI Deep Research",
"thinking": "Thinking...",
"searching": "Searching the web...",
"reading": "Reading sources...",
"analyzing": "Analyzing information...",
"completed": "Research completed",
"initializing": "Initializing research...",
"query": "Query",
"thoughtProcess": "Thought Process",
"searchingFor": "Searching for",
"readingUrl": "Reading",
"analyzingContent": "Analyzing content...",
"results": "Research Results"
},
"commandOutput": "Command Output",
"response": "Response",
"arguments": "Arguments",