feat: implement memory storage for follow-up questions (#6553)

packages/types/src/codebase-index.ts

@@ -34,6 +34,8 @@ export const codebaseIndexConfigSchema = z.object({
 	// OpenAI Compatible specific fields
 	codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(),
 	codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(),
+	// Memory storage settings
+	memoryStorageEnabled: z.boolean().optional(),
 })
 
 export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>

packages/types/src/global-settings.ts

@@ -62,6 +62,10 @@ export const globalSettingsSchema = z.object({
 	alwaysAllowFollowupQuestions: z.boolean().optional(),
 	followupAutoApproveTimeoutMs: z.number().optional(),
 	alwaysAllowUpdateTodoList: z.boolean().optional(),
+
+	// Memory storage settings
+	memoryStorageEnabled: z.boolean().optional(),
+	memoryStorageAutoApprove: z.boolean().optional(),
 	allowedCommands: z.array(z.string()).optional(),
 	deniedCommands: z.array(z.string()).optional(),
 	commandExecutionTimeout: z.number().optional(),
@@ -240,6 +244,8 @@ export const EVALS_SETTINGS: RooCodeSettings = {
 	alwaysAllowFollowupQuestions: true,
 	alwaysAllowUpdateTodoList: true,
 	followupAutoApproveTimeoutMs: 0,
+	memoryStorageEnabled: false,
+	memoryStorageAutoApprove: false,
 	allowedCommands: ["*"],
 	commandExecutionTimeout: 20,
 	commandTimeoutAllowlist: [],

packages/types/src/tool.ts

@@ -28,6 +28,8 @@ export const toolNames = [
 	"use_mcp_tool",
 	"access_mcp_resource",
 	"ask_followup_question",
+	"ask_memory_aware_followup_question",
+	"search_memories",
 	"attempt_completion",
 	"switch_mode",
 	"new_task",

src/core/assistant-message/presentAssistantMessage.ts

@@ -21,6 +21,8 @@ import { executeCommandTool } from "../tools/executeCommandTool"
 import { useMcpToolTool } from "../tools/useMcpToolTool"
 import { accessMcpResourceTool } from "../tools/accessMcpResourceTool"
 import { askFollowupQuestionTool } from "../tools/askFollowupQuestionTool"
+import { askMemoryAwareFollowupQuestionTool } from "../tools/askMemoryAwareFollowupQuestionTool"
+import { searchMemoriesTool } from "../tools/searchMemoriesTool"
 import { switchModeTool } from "../tools/switchModeTool"
 import { attemptCompletionTool } from "../tools/attemptCompletionTool"
 import { newTaskTool } from "../tools/newTaskTool"
@@ -151,6 +153,10 @@ export async function presentAssistantMessage(cline: Task) {
 			break
 		}
 		case "tool_use":
+			// Get customModes early for use in toolDescription
+			const stateForDescription = await cline.providerRef.deref()?.getState()
+			const customModesForDescription = stateForDescription?.customModes ?? []
+
 			const toolDescription = (): string => {
 				switch (block.name) {
 					case "execute_command":
@@ -200,6 +206,10 @@ export async function presentAssistantMessage(cline: Task) {
 						return `[${block.name} for '${block.params.server_name}']`
 					case "ask_followup_question":
 						return `[${block.name} for '${block.params.question}']`
+					case "ask_memory_aware_followup_question":
+						return `[${block.name} for '${block.params.question}']`
+					case "search_memories":
+						return `[${block.name} for '${block.params.query}']`
 					case "attempt_completion":
 						return `[${block.name}]`
 					case "switch_mode":
@@ -211,7 +221,7 @@ export async function presentAssistantMessage(cline: Task) {
 					case "new_task": {
 						const mode = block.params.mode ?? defaultModeSlug
 						const message = block.params.message ?? "(no message)"
-						const modeName = getModeBySlug(mode, customModes)?.name ?? mode
+						const modeName = getModeBySlug(mode, customModesForDescription)?.name ?? mode
 						return `[${block.name} in ${modeName} mode: '${message}']`
 					}
 				}
@@ -504,6 +514,19 @@ export async function presentAssistantMessage(cline: Task) {
 						removeClosingTag,
 					)
 					break
+				case "ask_memory_aware_followup_question":
+					await askMemoryAwareFollowupQuestionTool(
+						cline,
+						block,
+						askApproval,
+						handleError,
+						pushToolResult,
+						removeClosingTag,
+					)
+					break
+				case "search_memories":
+					await searchMemoriesTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
+					break
 				case "switch_mode":
 					await switchModeTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
 					break

src/core/tools/askFollowupQuestionTool.ts

@@ -2,6 +2,10 @@ import { Task } from "../task/Task"
 import { ToolUse, AskApproval, HandleError, PushToolResult, RemoveClosingTag } from "../../shared/tools"
 import { formatResponse } from "../prompts/responses"
 import { parseXml } from "../../utils/xml"
+import { CodeIndexConfigManager } from "../../services/code-index/config-manager"
+import { CodeIndexServiceFactory } from "../../services/code-index/service-factory"
+import { MemoryStorageManager } from "../../services/memory-storage/MemoryStorageManager"
+import { CacheManager } from "../../services/code-index/cache-manager"
 
 export async function askFollowupQuestionTool(
 	cline: Task,
@@ -80,6 +84,49 @@ export async function askFollowupQuestionTool(
 			await cline.say("user_feedback", text ?? "", images)
 			pushToolResult(formatResponse.toolResult(`<answer>\n${text}\n</answer>`, images))
 
+			// Store memory if enabled
+			try {
+				const provider = cline.providerRef.deref()
+				if (provider && text) {
+					// Get the code index manager from the provider
+					const codeIndexManager = provider.codeIndexManager
+					if (codeIndexManager) {
+						// Create config manager and service factory
+						const configManager = new CodeIndexConfigManager(provider.contextProxy)
+						const cacheManager = new CacheManager(provider.context, cline.workspacePath)
+						const serviceFactory = new CodeIndexServiceFactory(
+							configManager,
+							cline.workspacePath,
+							cacheManager,
+						)
+
+						// Get or create the memory storage manager
+						const memoryManager = MemoryStorageManager.getInstance(
+							configManager,
+							serviceFactory,
+							cline.workspacePath,
+						)
+
+						// Store the memory if enabled
+						if (memoryManager.isEnabled()) {
+							const memoryService = await memoryManager.getMemoryStorageService()
+							if (memoryService) {
+								await memoryService.storeMemory(
+									question,
+									text,
+									follow_up_json.suggest,
+									cline.taskId,
+									await cline.getTaskMode(),
+								)
+							}
+						}
+					}
+				}
+			} catch (error) {
+				// Log error but don't fail the tool
+				console.error("Failed to store memory:", error)
+			}
+
 			return
 		}
 	} catch (error) {

src/core/tools/askMemoryAwareFollowupQuestionTool.ts

@@ -0,0 +1,183 @@
+import { Task } from "../task/Task"
+import { ToolUse, AskApproval, HandleError, PushToolResult, RemoveClosingTag } from "../../shared/tools"
+import { formatResponse } from "../prompts/responses"
+import { parseXml } from "../../utils/xml"
+import { CodeIndexConfigManager } from "../../services/code-index/config-manager"
+import { CodeIndexServiceFactory } from "../../services/code-index/service-factory"
+import { MemoryStorageManager } from "../../services/memory-storage/MemoryStorageManager"
+import { CacheManager } from "../../services/code-index/cache-manager"
+
+export async function askMemoryAwareFollowupQuestionTool(
+	cline: Task,
+	block: ToolUse,
+	askApproval: AskApproval,
+	handleError: HandleError,
+	pushToolResult: PushToolResult,
+	removeClosingTag: RemoveClosingTag,
+) {
+	const question: string | undefined = block.params.question
+	const follow_up: string | undefined = block.params.follow_up
+
+	try {
+		if (block.partial) {
+			await cline.ask("followup", removeClosingTag("question", question), block.partial).catch(() => {})
+			return
+		} else {
+			if (!question) {
+				cline.consecutiveMistakeCount++
+				cline.recordToolError("ask_memory_aware_followup_question")
+				pushToolResult(
+					await cline.sayAndCreateMissingParamError("ask_memory_aware_followup_question", "question"),
+				)
+				return
+			}
+
+			type Suggest = { answer: string; mode?: string }
+
+			let follow_up_json = {
+				question,
+				suggest: [] as Suggest[],
+			}
+
+			if (follow_up) {
+				// Define the actual structure returned by the XML parser
+				type ParsedSuggestion = string | { "#text": string; "@_mode"?: string }
+
+				let parsedSuggest: {
+					suggest: ParsedSuggestion[] | ParsedSuggestion
+				}
+
+				try {
+					parsedSuggest = parseXml(follow_up, ["suggest"]) as {
+						suggest: ParsedSuggestion[] | ParsedSuggestion
+					}
+				} catch (error) {
+					cline.consecutiveMistakeCount++
+					cline.recordToolError("ask_memory_aware_followup_question")
+					await cline.say("error", `Failed to parse operations: ${error.message}`)
+					pushToolResult(formatResponse.toolError("Invalid operations xml format"))
+					return
+				}
+
+				const rawSuggestions = Array.isArray(parsedSuggest?.suggest)
+					? parsedSuggest.suggest
+					: [parsedSuggest?.suggest].filter((sug): sug is ParsedSuggestion => sug !== undefined)
+
+				// Transform parsed XML to our Suggest format
+				const normalizedSuggest: Suggest[] = rawSuggestions.map((sug) => {
+					if (typeof sug === "string") {
+						// Simple string suggestion (no mode attribute)
+						return { answer: sug }
+					} else {
+						// XML object with text content and optional mode attribute
+						const result: Suggest = { answer: sug["#text"] }
+						if (sug["@_mode"]) {
+							result.mode = sug["@_mode"]
+						}
+						return result
+					}
+				})
+
+				follow_up_json.suggest = normalizedSuggest
+			}
+
+			// Get relevant memories before asking the question
+			let memoryContext = ""
+			try {
+				const provider = cline.providerRef.deref()
+				if (provider) {
+					const codeIndexManager = provider.codeIndexManager
+					if (codeIndexManager) {
+						const configManager = new CodeIndexConfigManager(provider.contextProxy)
+						const cacheManager = new CacheManager(provider.context, cline.workspacePath)
+						const serviceFactory = new CodeIndexServiceFactory(
+							configManager,
+							cline.workspacePath,
+							cacheManager,
+						)
+
+						const memoryManager = MemoryStorageManager.getInstance(
+							configManager,
+							serviceFactory,
+							cline.workspacePath,
+						)
+
+						if (memoryManager.isEnabled()) {
+							const memoryService = await memoryManager.getMemoryStorageService()
+							if (memoryService) {
+								// Search for relevant memories
+								const relevantMemories = await memoryService.searchMemories(question, 5)
+
+								// Format memories for context
+								if (relevantMemories.length > 0) {
+									memoryContext = "\n\nBased on previous interactions:\n"
+									relevantMemories.forEach((memory, index) => {
+										memoryContext += `${index + 1}. Q: ${memory.question}\n   A: ${memory.answer}\n`
+									})
+								}
+							}
+						}
+					}
+				}
+			} catch (error) {
+				console.error("Failed to retrieve memories:", error)
+				// Continue without memory context
+			}
+
+			// Add memory context to the question
+			const questionWithContext = question + memoryContext
+
+			cline.consecutiveMistakeCount = 0
+			const { text, images } = await cline.ask(
+				"followup",
+				JSON.stringify({ ...follow_up_json, question: questionWithContext }),
+				false,
+			)
+			await cline.say("user_feedback", text ?? "", images)
+			pushToolResult(formatResponse.toolResult(`<answer>\n${text}\n</answer>`, images))
+
+			// Store memory if enabled (same as askFollowupQuestionTool)
+			try {
+				const provider = cline.providerRef.deref()
+				if (provider && text) {
+					const codeIndexManager = provider.codeIndexManager
+					if (codeIndexManager) {
+						const configManager = new CodeIndexConfigManager(provider.contextProxy)
+						const cacheManager = new CacheManager(provider.context, cline.workspacePath)
+						const serviceFactory = new CodeIndexServiceFactory(
+							configManager,
+							cline.workspacePath,
+							cacheManager,
+						)
+
+						const memoryManager = MemoryStorageManager.getInstance(
+							configManager,
+							serviceFactory,
+							cline.workspacePath,
+						)
+
+						if (memoryManager.isEnabled()) {
+							const memoryService = await memoryManager.getMemoryStorageService()
+							if (memoryService) {
+								await memoryService.storeMemory(
+									question,
+									text,
+									follow_up_json.suggest,
+									cline.taskId,
+									await cline.getTaskMode(),
+								)
+							}
+						}
+					}
+				}
+			} catch (error) {
+				console.error("Failed to store memory:", error)
+			}
+
+			return
+		}
+	} catch (error) {
+		await handleError("asking memory-aware question", error)
+		return
+	}
+}