diff --git a/packages/types/src/message.ts b/packages/types/src/message.ts index 6f5162544d..f50142602d 100644 --- a/packages/types/src/message.ts +++ b/packages/types/src/message.ts @@ -367,6 +367,18 @@ export const tokenUsageSchema = z.object({ export type TokenUsage = z.infer +/** + * MentionSegment + */ + +/** A part of a message with a known origin. Only `user` parts may have mentions expanded. */ +export const mentionSegmentSchema = z.object({ + text: z.string(), + source: z.enum(["user", "generated"]), +}) + +export type MentionSegment = z.infer + /** * QueuedMessage */ @@ -376,6 +388,8 @@ export const queuedMessageSchema = z.object({ id: z.string(), text: z.string(), images: z.array(z.string()).optional(), + /** Origin of each part of `text`. Absent means all generated. */ + segments: z.array(mentionSegmentSchema).optional(), }) export type QueuedMessage = z.infer diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts index 8a006c0639..516cfe2479 100644 --- a/packages/types/src/vscode-extension-host.ts +++ b/packages/types/src/vscode-extension-host.ts @@ -5,7 +5,7 @@ import type { ProviderSettings, ProviderSettingsEntry } from "./provider-setting import type { HistoryItem } from "./history.js" import type { ModeConfig, PromptComponent } from "./mode.js" import type { Experiments } from "./experiment.js" -import type { ClineMessage, QueuedMessage } from "./message.js" +import type { ClineMessage, MentionSegment, QueuedMessage } from "./message.js" import type { MarketplaceItem, MarketplaceInstalledMetadata, InstallMarketplaceItemOptions } from "./marketplace.js" import type { TodoItem } from "./todo.js" import type { CloudUserInfo, CloudOrganizationMembership, OrganizationAllowList, ShareVisibility } from "./cloud.js" @@ -133,6 +133,8 @@ export interface ExtensionMessage { | "switchTab" | "toggleAutoApprove" invoke?: "newChat" | "sendMessage" | "primaryButtonClick" | "secondaryButtonClick" | "setChatBoxMessage" + /** With `invoke: "setChatBoxMessage"`: origin of each part of `text`. Absent means all generated. */ + segments?: MentionSegment[] /** * Partial state updates are allowed to reduce message size (e.g. omit large fields like taskHistory). * The webview is responsible for merging. @@ -817,6 +819,8 @@ export interface WebviewMessage { updatedSettings?: RooCodeSettings /** Task configuration applied via `createTask()`. */ taskConfiguration?: RooCodeSettings + /** Origin of each part of `text`. Absent or not matching `text` means all generated. */ + segments?: MentionSegment[] // Worktree properties worktreePath?: string worktreeBranch?: string diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts index 74949a8742..c2c733d178 100644 --- a/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-images.spec.ts @@ -381,13 +381,15 @@ describe("presentAssistantMessage - Image Handling in Native Tool Calling", () = await presentAssistantMessage(mockTask, baseSnapshot) - // Anthropic rejects user messages where image blocks sit between - // tool_result blocks, so images from both tools must be grouped - // after the last tool_result instead of interleaved. + // Images must come after all tool_result blocks (#1307). Feedback + // text blocks join the tail in push order. No image may appear + // before the second tool_result. expect(mockTask.userMessageContent.map((block: UserContentBlock) => block.type)).toEqual([ "tool_result", "tool_result", + "text", "image", + "text", "image", ]) expect( @@ -442,10 +444,14 @@ describe("presentAssistantMessage - Image Handling in Native Tool Calling", () = await presentAssistantMessage(mockTask, baseSnapshot) + // The image must come after the last tool_result (#1307), never + // between tool_result blocks. Feedback text joins the tail. expect(mockTask.userMessageContent.map((block: UserContentBlock) => block.type)).toEqual([ "tool_result", "tool_result", + "text", "image", + "text", ]) expect( mockTask.userMessageContent diff --git a/src/core/mentions/__tests__/processUserContentMentions.spec.ts b/src/core/mentions/__tests__/processUserContentMentions.spec.ts index 1995c36bc0..b8a6ed85c1 100644 --- a/src/core/mentions/__tests__/processUserContentMentions.spec.ts +++ b/src/core/mentions/__tests__/processUserContentMentions.spec.ts @@ -1,345 +1,177 @@ -// npx vitest core/mentions/__tests__/processUserContentMentions.spec.ts - -import { processUserContentMentions } from "../processUserContentMentions" +import type Anthropic from "@anthropic-ai/sdk" +import type { MentionSegment } from "@roo-code/types" +import { processUserContentMentions, userMentionContent, wrapUserMessage } from "../processUserContentMentions" import { parseMentions } from "../index" import { FileContextTracker } from "../../context-tracking/FileContextTracker" -// Mock the parseMentions function -vi.mock("../index", () => ({ - parseMentions: vi.fn(), -})) +vi.mock("../index", () => ({ parseMentions: vi.fn() })) -describe("processUserContentMentions", () => { - let mockFileContextTracker: FileContextTracker +const user = (text: string): MentionSegment => ({ text, source: "user" }) +const gen = (text: string): MentionSegment => ({ text, source: "generated" }) +describe("processUserContentMentions authorization", () => { beforeEach(() => { vi.clearAllMocks() - - mockFileContextTracker = {} as FileContextTracker - - // Default mock implementation - returns ParseMentionsResult object vi.mocked(parseMentions).mockImplementation(async (text) => ({ text: `parsed: ${text}`, - mode: undefined, - contentBlocks: [], + contentBlocks: [{ type: "file", path: "allowed", content: "file contents" }], + mode: "code", + slashCommandHelp: "command help", })) }) - describe("content processing", () => { - it("should process text blocks with tags", async () => { - const userContent = [ - { - type: "text" as const, - text: "Do something", - }, - ] - - const result = await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) - - expect(parseMentions).toHaveBeenCalled() - expect(result.content[0]).toEqual({ - type: "text", - text: "parsed: Do something", - }) - expect(result.mode).toBeUndefined() - }) - - it("should not process text blocks without user_message tags", async () => { - const userContent = [ - { - type: "text" as const, - text: "Regular text without special tags", - }, - ] - - const result = await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) - - expect(parseMentions).not.toHaveBeenCalled() - expect(result.content[0]).toEqual(userContent[0]) - expect(result.mode).toBeUndefined() - }) - - it("should process tool_result blocks with string content", async () => { - const userContent = [ - { - type: "tool_result" as const, - tool_use_id: "123", - content: "Tool feedback", - }, - ] - - const result = await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) - - expect(parseMentions).toHaveBeenCalled() - // String content is now converted to array format to support content blocks - expect(result.content[0]).toEqual({ - type: "tool_result", - tool_use_id: "123", - content: [ - { - type: "text", - text: "parsed: Tool feedback", - }, - ], - }) - expect(result.mode).toBeUndefined() - }) - - it("should process tool_result blocks with array content", async () => { - const userContent = [ - { - type: "tool_result" as const, - tool_use_id: "123", - content: [ - { - type: "text" as const, - text: "Array task", - }, - { - type: "text" as const, - text: "Regular text", - }, - ], - }, - ] - - const result = await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) - - expect(parseMentions).toHaveBeenCalledTimes(1) - expect(result.content[0]).toEqual({ - type: "tool_result", - tool_use_id: "123", - content: [ - { - type: "text", - text: "parsed: Array task", - }, - { - type: "text", - text: "Regular text", - }, - ], - }) - expect(result.mode).toBeUndefined() - }) - - it("should handle mixed content types (text + image)", async () => { - const userContent = [ - { - type: "text" as const, - text: "First task", - }, - { - type: "image" as const, - image: "base64data", - mediaType: "image/png", - }, - ] - - const result = await processUserContentMentions({ - userContent: userContent as any, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) + const process = (userContent: Anthropic.Messages.ContentBlockParam[]) => + processUserContentMentions({ userContent, cwd: "/test", fileContextTracker: {} as FileContextTracker }) + + it.each([ + { type: "text" as const, text: "@/secret @terminal" }, + { type: "tool_result" as const, tool_use_id: "1", content: "@/secret @terminal" }, + { + type: "tool_result" as const, + tool_use_id: "1", + content: [{ type: "text" as const, text: "@/secret @terminal" }], + }, + ])("does not expand untrusted content: $type", async (block) => { + const result = await process([block]) + expect(parseMentions).not.toHaveBeenCalled() + expect(result).toEqual({ content: [block], mode: undefined }) + }) - expect(parseMentions).toHaveBeenCalledTimes(1) - expect(result.content).toHaveLength(2) - expect(result.content[0]).toEqual({ - type: "text", - text: "parsed: First task", - }) - expect(result.content[1]).toEqual(userContent[1]) // Image block unchanged - expect(result.mode).toBeUndefined() + it("expands explicit user input without requiring a textual marker", async () => { + const result = await process([userMentionContent([user("@/allowed")])]) + expect(parseMentions).toHaveBeenCalledWith( + "@/allowed", + "/test", + expect.anything(), + undefined, + false, + true, + 50, + undefined, + "code", + ) + expect(result).toEqual({ + content: [ + { type: "text", text: "parsed: @/allowed" }, + { type: "text", text: "file contents" }, + { type: "text", text: "command help" }, + ], + mode: "code", }) }) - describe("showRooIgnoredFiles parameter", () => { - it("should default showRooIgnoredFiles to false", async () => { - const userContent = [ - { - type: "text" as const, - text: "Test default", - }, - ] - - await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) + it("expands only the user segments of a support prompt", async () => { + const context = "@/secret @terminal" + const result = await process([userMentionContent([user("@/allowed"), gen(`\n${context}`)])]) + expect(vi.mocked(parseMentions).mock.calls.map(([text]) => text)).toEqual(["@/allowed"]) + expect(result.content[0]).toEqual({ type: "text", text: `parsed: @/allowed\n${context}` }) + }) - expect(parseMentions).toHaveBeenCalledWith( - "Test default", - "/test", - mockFileContextTracker, - undefined, - false, // showRooIgnoredFiles should default to false - true, // includeDiagnosticMessages - 50, // maxDiagnosticMessages - undefined, - "code", - ) - }) + it("expands each user segment on its own and leaves generated segments unchanged", async () => { + const result = await process([ + userMentionContent([gen("a @/x "), user("@/one"), gen(" b @/y "), user("@/two")]), + ]) + expect(vi.mocked(parseMentions).mock.calls.map(([text]) => text)).toEqual(["@/one", "@/two"]) + expect(result.content[0]).toEqual({ type: "text", text: "a @/x parsed: @/one b @/y parsed: @/two" }) + }) - it("should respect showRooIgnoredFiles when explicitly set to false", async () => { - const userContent = [ - { - type: "text" as const, - text: "Test explicit false", - }, - ] + it("expands the same text only where it is a user segment", async () => { + const result = await process([userMentionContent([gen("@/same "), user("@/same")])]) + expect(vi.mocked(parseMentions).mock.calls.map(([text]) => text)).toEqual(["@/same"]) + expect(result.content[0]).toEqual({ type: "text", text: "@/same parsed: @/same" }) + }) - await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - showRooIgnoredFiles: false, - }) + it("does not expand all generated segments", async () => { + await process([userMentionContent([gen("@/secret")])]) + expect(parseMentions).not.toHaveBeenCalled() + }) - expect(parseMentions).toHaveBeenCalledWith( - "Test explicit false", - "/test", - mockFileContextTracker, - undefined, - false, - true, // includeDiagnosticMessages - 50, // maxDiagnosticMessages - undefined, - "code", - ) + it("wraps segments in tags that stay generated", async () => { + const result = await process([userMentionContent(wrapUserMessage([user("@/allowed")]))]) + expect(vi.mocked(parseMentions).mock.calls.map(([text]) => text)).toEqual(["@/allowed"]) + expect(result.content[0]).toEqual({ + type: "text", + text: "\nparsed: @/allowed\n", }) }) - describe("slash command content processing", () => { - it("should separate slash command content into a new block", async () => { - vi.mocked(parseMentions).mockResolvedValueOnce({ - text: "parsed text", - slashCommandHelp: "command help", - mode: undefined, - contentBlocks: [], - }) - - const userContent = [ - { - type: "text" as const, - text: "Run command", - }, - ] + it.each([ + { label: "spread copy", change: (b: Anthropic.TextBlockParam) => ({ ...b }) }, + { label: "structured clone", change: (b: Anthropic.TextBlockParam) => structuredClone(b) }, + { label: "JSON round trip", change: (b: Anthropic.TextBlockParam) => JSON.parse(JSON.stringify(b)) }, + { label: "edited text", change: (b: Anthropic.TextBlockParam) => Object.assign(b, { text: b.text + " @/x" }) }, + ])("does not trust a block after $label", async ({ change }) => { + await process([change(userMentionContent([user("@/allowed")]))]) + expect(parseMentions).not.toHaveBeenCalled() + }) - const result = await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) + it("does not trust all generated segments", async () => { + await process([userMentionContent([gen("context @/secret")])]) + expect(parseMentions).not.toHaveBeenCalled() + }) - expect(result.content).toHaveLength(2) - expect(result.content[0]).toEqual({ - type: "text", - text: "parsed text", - }) - expect(result.content[1]).toEqual({ - type: "text", - text: "command help", - }) + it("keeps labels out of the serialized block", () => { + expect(JSON.parse(JSON.stringify(userMentionContent([user("@/a"), gen("b")])))).toEqual({ + type: "text", + text: "@/ab", }) + }) - it("should include slash command content in tool_result string content", async () => { - vi.mocked(parseMentions).mockResolvedValueOnce({ - text: "parsed tool output", - slashCommandHelp: "command help", - mode: undefined, - contentBlocks: [], - }) - - const userContent = [ - { - type: "tool_result" as const, - tool_use_id: "123", - content: "Tool output", - }, - ] + it("trusts the same block after the array is copied", async () => { + const block = userMentionContent([user("@/allowed")]) + await process([...[block]].map((b) => b)) + expect(parseMentions).toHaveBeenCalledTimes(1) + }) - const result = await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) + it("passes image blocks through unchanged next to a user block", async () => { + const image: Anthropic.Messages.ImageBlockParam = { + type: "image", + source: { type: "base64", media_type: "image/png", data: "abc" }, + } + const result = await process([userMentionContent([user("@/allowed")]), image]) + expect(result.content[result.content.length - 1]).toBe(image) + }) - expect(result.content).toHaveLength(1) - expect(result.content[0]).toEqual({ - type: "tool_result", - tool_use_id: "123", - content: [ - { - type: "text", - text: "parsed tool output", - }, - { - type: "text", - text: "command help", - }, - ], - }) + it.each([ + { showRooIgnoredFiles: true, includeDiagnosticMessages: false, maxDiagnosticMessages: 3, currentMode: "ask" }, + { showRooIgnoredFiles: false, includeDiagnosticMessages: true, maxDiagnosticMessages: 50, currentMode: "code" }, + ])("passes the flags to parseMentions: %o", async (flags) => { + await processUserContentMentions({ + userContent: [userMentionContent([user("@/allowed")])], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + ...flags, }) + expect(parseMentions).toHaveBeenCalledWith( + "@/allowed", + "/test", + expect.anything(), + undefined, + flags.showRooIgnoredFiles, + flags.includeDiagnosticMessages, + flags.maxDiagnosticMessages, + undefined, + flags.currentMode, + ) + }) - it("should include slash command content in tool_result array content", async () => { - vi.mocked(parseMentions).mockResolvedValueOnce({ - text: "parsed array item", - slashCommandHelp: "command help", - mode: undefined, - contentBlocks: [], - }) - - const userContent = [ - { - type: "tool_result" as const, - tool_use_id: "123", - content: [ - { - type: "text" as const, - text: "Array item", - }, - ], - }, - ] - - const result = await processUserContentMentions({ - userContent, - cwd: "/test", - fileContextTracker: mockFileContextTracker, - }) + it("keeps the mode of the first slash command", async () => { + vi.mocked(parseMentions) + .mockResolvedValueOnce({ text: "one", contentBlocks: [], mode: "architect" }) + .mockResolvedValueOnce({ text: "two", contentBlocks: [], mode: "debug" }) + const result = await process([userMentionContent([user("/a"), gen(" "), user("/b")])]) + expect(result.mode).toBe("architect") + }) - expect(result.content).toHaveLength(1) - expect(result.content[0]).toEqual({ - type: "tool_result", - tool_use_id: "123", - content: [ - { - type: "text", - text: "parsed array item", - }, - { - type: "text", - text: "command help", - }, - ], - }) - }) + it("orders the extra blocks by segment after the text block", async () => { + vi.mocked(parseMentions) + .mockResolvedValueOnce({ text: "one", contentBlocks: [{ type: "file", path: "a", content: "A" }] }) + .mockResolvedValueOnce({ text: "two", contentBlocks: [{ type: "file", path: "b", content: "B" }] }) + const result = await process([userMentionContent([user("@/a"), gen(" "), user("@/b")])]) + expect(result.content).toEqual([ + { type: "text", text: "one two" }, + { type: "text", text: "A" }, + { type: "text", text: "B" }, + ]) }) }) diff --git a/src/core/mentions/processUserContentMentions.ts b/src/core/mentions/processUserContentMentions.ts index 0a8ad40dfa..90d9c255ed 100644 --- a/src/core/mentions/processUserContentMentions.ts +++ b/src/core/mentions/processUserContentMentions.ts @@ -1,4 +1,5 @@ import Anthropic from "@anthropic-ai/sdk" +import type { MentionSegment } from "@roo-code/types" import { parseMentions, MentionContentBlock } from "./index" import { FileContextTracker } from "../context-tracking/FileContextTracker" @@ -8,6 +9,26 @@ import type { SkillLookup } from "../../services/skills/skillInvocation" type TextPart = Anthropic.Messages.TextBlockParam type ImagePart = Anthropic.Messages.ImageBlockParam +// Origin labels stay in memory and cannot be forged by message text or restored history. +const segmentedBlocks = new WeakMap() + +/** Call only at a user-input boundary. Only `user` segments may have mentions expanded. */ +export function userMentionContent(segments: MentionSegment[]): TextPart { + const text = segments.map((s) => s.text).join("") + const block: TextPart = { type: "text", text } + segmentedBlocks.set(block, { text, segments }) + return block +} + +/** Wraps segments in the `` tags. The tags are generated text. */ +export function wrapUserMessage(segments: MentionSegment[]): MentionSegment[] { + return [ + { text: "\n", source: "generated" }, + ...segments, + { text: "\n", source: "generated" }, + ] +} + export interface ProcessUserContentMentionsResult { content: Anthropic.Messages.ContentBlockParam[] mode?: string // Mode from the first slash command that has one @@ -26,7 +47,7 @@ function contentBlocksToTextParts(contentBlocks: MentionContentBlock[]): TextPar } /** - * Process mentions in user content, specifically within task and feedback tags. + * Process mentions only in explicitly authorized user input. * * File/folder @ mentions are now returned as separate text blocks that * look like read_file tool results, making it clear to the model that @@ -57,61 +78,22 @@ export async function processUserContentMentions({ let commandMode: string | undefined // Process userContent array, which contains text and image parts. - // We need to apply parseMentions() to TextPart's text that contains "". + // Tool results and unmarked text never authorize resource access. const content = ( await Promise.all( userContent.map(async (block) => { - const shouldProcessMentions = (text: string) => text.includes("") - if (block.type === "text") { - if (shouldProcessMentions(block.text)) { - const result = await parseMentions( - block.text, - cwd, - fileContextTracker, - rooIgnoreController, - showRooIgnoredFiles, - includeDiagnosticMessages, - maxDiagnosticMessages, - skillsManager, - currentMode, - ) - // Capture the first mode found - if (!commandMode && result.mode) { - commandMode = result.mode - } - - // Build the blocks array: - // 1. User's text (with @ mentions replaced by clean paths) - // 2. File/folder content blocks (formatted like read_file results) - // 3. Slash command help (if any) - const blocks: Array = [ - { - ...block, - text: result.text, - }, - ] - - // Add file/folder content as separate blocks - if (result.contentBlocks.length > 0) { - blocks.push(...contentBlocksToTextParts(result.contentBlocks)) - } - - if (result.slashCommandHelp) { - blocks.push({ - type: "text" as const, - text: result.slashCommandHelp, - }) - } - return blocks - } - - return block - } else if (block.type === "tool_result") { - if (typeof block.content === "string") { - if (shouldProcessMentions(block.content)) { + const labeled = segmentedBlocks.get(block) + if (labeled && labeled.text === block.text) { + let text = "" + const extra: TextPart[] = [] + for (const segment of labeled.segments) { + if (segment.source !== "user") { + text += segment.text + continue + } const result = await parseMentions( - block.content, + segment.text, cwd, fileContextTracker, rooIgnoreController, @@ -125,89 +107,15 @@ export async function processUserContentMentions({ if (!commandMode && result.mode) { commandMode = result.mode } - - // Build content array with file blocks included - const contentParts: Array<{ type: "text"; text: string }> = [ - { - type: "text" as const, - text: result.text, - }, - ] - - // Add file/folder content blocks - for (const contentBlock of result.contentBlocks) { - contentParts.push({ - type: "text" as const, - text: contentBlock.content, - }) - } - + text += result.text + // File/folder content as separate blocks (formatted like read_file results) + extra.push(...contentBlocksToTextParts(result.contentBlocks)) if (result.slashCommandHelp) { - contentParts.push({ - type: "text" as const, - text: result.slashCommandHelp, - }) - } - - return { - ...block, - content: contentParts, + extra.push({ type: "text" as const, text: result.slashCommandHelp }) } } - - return block - } else if (Array.isArray(block.content)) { - const parsedContent = ( - await Promise.all( - block.content.map(async (contentBlock) => { - if (contentBlock.type === "text" && shouldProcessMentions(contentBlock.text)) { - const result = await parseMentions( - contentBlock.text, - cwd, - fileContextTracker, - rooIgnoreController, - showRooIgnoredFiles, - includeDiagnosticMessages, - maxDiagnosticMessages, - skillsManager, - currentMode, - ) - // Capture the first mode found - if (!commandMode && result.mode) { - commandMode = result.mode - } - - // Build blocks array with file content - const blocks: Array<{ type: "text"; text: string }> = [ - { - ...contentBlock, - text: result.text, - }, - ] - - // Add file/folder content blocks - for (const cb of result.contentBlocks) { - blocks.push({ - type: "text" as const, - text: cb.content, - }) - } - - if (result.slashCommandHelp) { - blocks.push({ - type: "text" as const, - text: result.slashCommandHelp, - }) - } - return blocks - } - - return contentBlock - }), - ) - ).flat() - - return { ...block, content: parsedContent } + const blocks: Array = [{ ...block, text }, ...extra] + return blocks } return block diff --git a/src/core/message-queue/MessageQueueService.ts b/src/core/message-queue/MessageQueueService.ts index 340fad6087..d87d149d68 100644 --- a/src/core/message-queue/MessageQueueService.ts +++ b/src/core/message-queue/MessageQueueService.ts @@ -2,7 +2,9 @@ import { EventEmitter } from "events" import { v4 as uuidv4 } from "uuid" -import { QueuedMessage } from "@roo-code/types" +import type { MentionSegment, QueuedMessage } from "@roo-code/types" + +import { applyEdit, generatedSegments } from "../../shared/mentionSegments" export interface QueueEvents { stateChanged: [messages: QueuedMessage[]] @@ -28,7 +30,7 @@ export class MessageQueueService extends EventEmitter { return { index, message: this._messages[index] } } - public addMessage(text: string, images?: string[]): QueuedMessage | undefined { + public addMessage(text: string, images?: string[], segments?: MentionSegment[]): QueuedMessage | undefined { if (!text && !images?.length) { return undefined } @@ -38,6 +40,7 @@ export class MessageQueueService extends EventEmitter { id: uuidv4(), text, images, + segments, } this._messages.push(message) @@ -66,6 +69,8 @@ export class MessageQueueService extends EventEmitter { return false } + // Only the changed text becomes user text. Unchanged generated text stays generated. + message.segments = applyEdit(message.segments ?? generatedSegments(message.text), message.text, text) message.timestamp = Date.now() message.text = text message.images = images diff --git a/src/core/message-queue/__tests__/MessageQueueService.spec.ts b/src/core/message-queue/__tests__/MessageQueueService.spec.ts index 567dca314b..2a2d7c7275 100644 --- a/src/core/message-queue/__tests__/MessageQueueService.spec.ts +++ b/src/core/message-queue/__tests__/MessageQueueService.spec.ts @@ -68,3 +68,44 @@ describe("MessageQueueService claims", () => { expect(queue.hasUnclaimed()).toBe(true) }) }) + +describe("MessageQueueService segments", () => { + it("stores the segments with the message", () => { + const queue = new MessageQueueService() + const segments = [ + { text: "ctx ", source: "generated" as const }, + { text: "mine", source: "user" as const }, + ] + expect(queue.addMessage("ctx mine", undefined, segments)?.segments).toEqual(segments) + }) + + it("keeps unchanged generated text generated when the user saves without edits", () => { + const queue = new MessageQueueService() + const message = queue.addMessage("ctx @/x", undefined, [{ text: "ctx @/x", source: "generated" }])! + + queue.updateMessage(message.id, "ctx @/x") + + expect(queue.messages[0]?.segments).toEqual([{ text: "ctx @/x", source: "generated" }]) + }) + + it("marks only the appended text as user text", () => { + const queue = new MessageQueueService() + const message = queue.addMessage("ctx @/x", undefined, [{ text: "ctx @/x", source: "generated" }])! + + queue.updateMessage(message.id, "ctx @/x mine") + + expect(queue.messages[0]?.segments).toEqual([ + { text: "ctx @/x", source: "generated" }, + { text: " mine", source: "user" }, + ]) + }) + + it("treats a message without segments as generated", () => { + const queue = new MessageQueueService() + const message = queue.addMessage("old @/x")! + + queue.updateMessage(message.id, "old @/x") + + expect(queue.messages[0]?.segments).toEqual([{ text: "old @/x", source: "generated" }]) + }) +}) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index aa6026e56b..9f38d6d24d 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -57,6 +57,7 @@ import { MAX_MCP_TOOLS_THRESHOLD, countEnabledMcpTools, providerIdentifiers, + type MentionSegment, } from "@roo-code/types" import { TelemetryService } from "@roo-code/telemetry" import { CloudService } from "@roo-code/cloud" @@ -134,7 +135,8 @@ import { checkpointRestore, checkpointDiff, } from "../checkpoints" -import { processUserContentMentions } from "../mentions/processUserContentMentions" +import { joinSegments, normalizeSegments, trimSegments } from "../../shared/mentionSegments" +import { processUserContentMentions, userMentionContent, wrapUserMessage } from "../mentions/processUserContentMentions" import { getMessagesSinceLastSummary, summarizeConversation, getEffectiveApiHistory } from "../condense" import { MessageQueueService } from "../message-queue/MessageQueueService" import { type AutoDenyDetail, AutoApprovalHandler, checkAutoApproval } from "../auto-approval" @@ -254,6 +256,8 @@ export interface TaskOptions extends CreateTaskOptions { checkpointTimeout?: number consecutiveMistakeLimit?: number task?: string + /** Origin of each part of `task`. Absent means all generated, so no mention expands. */ + segments?: MentionSegment[] images?: string[] historyItem?: HistoryItem experiments?: Record @@ -450,6 +454,7 @@ export class Task extends EventEmitter implements TaskLike { private askResponse?: ClineAskResponse private askResponseText?: string private askResponseImages?: string[] + private askResponseSegments?: MentionSegment[] /** * Structured detail of the automatic denial that resolved the current ask, * set when `checkAutoApproval` denies via policy (denylist, blanket @@ -693,7 +698,10 @@ export class Task extends EventEmitter implements TaskLike { // MessageManager for high-level message operations (lazy initialized) private _messageManager?: MessageManager + private readonly segments?: MentionSegment[] + constructor({ + segments, provider, apiConfiguration, enableCheckpoints = true, @@ -718,6 +726,8 @@ export class Task extends EventEmitter implements TaskLike { handoffExecutionContext, }: TaskOptions) { super() + this.segments = segments + this.resetAssistantMessagePersistence() if (startTask && !task && !images && !historyItem) { @@ -1251,7 +1261,7 @@ export class Task extends EventEmitter implements TaskLike { } private handleQueuedAskResponse(message: QueuedMessage, resolution: QueuedAskResolution): string | undefined { - this.handleWebviewAskResponse(resolution.response, message.text, message.images) + this.handleWebviewAskResponse(resolution.response, message.text, message.images, message.segments) if (resolution.requiresDurableAck) { return message.id } @@ -1921,6 +1931,8 @@ export class Task extends EventEmitter implements TaskLike { response: ClineAskResponse text?: string images?: string[] + /** Origin of each part of `text`. */ + segments?: MentionSegment[] queuedMessageId?: string /** * Present when the ask was resolved by an automatic (policy) denial @@ -2439,12 +2451,14 @@ export class Task extends EventEmitter implements TaskLike { response: this.askResponse!, text: this.askResponseText, images: this.askResponseImages, + segments: this.askResponseSegments, queuedMessageId, autoDenyDetail: this.pendingAutoDenyDetail, } this.askResponse = undefined this.askResponseText = undefined this.askResponseImages = undefined + this.askResponseSegments = undefined this.pendingAutoDenyDetail = undefined // Switch back to an active state. @@ -2467,13 +2481,19 @@ export class Task extends EventEmitter implements TaskLike { } } - handleWebviewAskResponse(askResponse: ClineAskResponse, text?: string, images?: string[]) { + handleWebviewAskResponse( + askResponse: ClineAskResponse, + text?: string, + images?: string[], + segments?: MentionSegment[], + ) { // Clear any pending auto-approval timeout when user responds this.cancelAutoApprovalTimeout() this.askResponse = askResponse this.askResponseText = text this.askResponseImages = images + this.askResponseSegments = segments // Create a checkpoint whenever the user sends a message. // Use allowEmpty=true to ensure a checkpoint is recorded even if there are no file changes. @@ -2614,9 +2634,12 @@ export class Task extends EventEmitter implements TaskLike { images?: string[], mode?: string, providerProfile?: string, + segments?: MentionSegment[], ): Promise { try { - text = (text ?? "").trim() + // Trim the text and its segments together so they stay equal. + const trimmed = trimSegments(normalizeSegments(text ?? "", segments)) + text = joinSegments(trimmed) images = images ?? [] if (text.length === 0 && images.length === 0) { @@ -2648,7 +2671,7 @@ export class Task extends EventEmitter implements TaskLike { // Handle the message directly instead of routing through the webview. // This avoids a race condition where the webview's message state hasn't // hydrated yet, causing it to interpret the message as a new task request. - this.handleWebviewAskResponse("messageResponse", text, images) + this.handleWebviewAskResponse("messageResponse", text, images, trimmed) } else { console.error("[Task#submitUserMessage] Provider reference lost") } @@ -3158,10 +3181,9 @@ export class Task extends EventEmitter implements TaskLike { // Task starting await this.initiateTaskLoop([ - { - type: "text", - text: `\n${task}\n`, - }, + // Task text can carry generated context. Without segments + // nothing expands. + userMentionContent(wrapUserMessage(normalizeSegments(task ?? "", this.segments))), ...imageBlocks, ]).catch((error) => { // Swallow loop rejection when the task was intentionally abandoned/aborted @@ -3301,7 +3323,7 @@ export class Task extends EventEmitter implements TaskLike { this.markInitialized() - const { response, text, images } = await this.ask(askType) // Calls `postStateToWebview`. + const { response, text, images, segments } = await this.ask(askType) // Calls `postStateToWebview`. let responseText: string | undefined let responseImages: string[] | undefined @@ -3423,10 +3445,8 @@ export class Task extends EventEmitter implements TaskLike { const newUserContent: Anthropic.Messages.ContentBlockParam[] = [...modifiedOldUserContent] if (responseText) { - newUserContent.push({ - type: "text", - text: `\n${responseText}\n`, - }) + // Only user segments expand. + newUserContent.push(userMentionContent(wrapUserMessage(normalizeSegments(responseText, segments)))) } if (responseImages && responseImages.length > 0) { @@ -6531,8 +6551,8 @@ export class Task extends EventEmitter implements TaskLike { const queued = this.messageQueueService.dequeueMessage() if (queued) { setTimeout(() => { - this.submitUserMessage(queued.text, queued.images).catch((err) => - console.error(`[Task] Failed to submit queued message:`, err), + this.submitUserMessage(queued.text, queued.images, undefined, undefined, queued.segments).catch( + (err) => console.error(`[Task] Failed to submit queued message:`, err), ) }, 0) } diff --git a/src/core/task/__tests__/Task.persistence.spec.ts b/src/core/task/__tests__/Task.persistence.spec.ts index b7c97a357e..e896d13a7c 100644 --- a/src/core/task/__tests__/Task.persistence.spec.ts +++ b/src/core/task/__tests__/Task.persistence.spec.ts @@ -16,6 +16,9 @@ import { TelemetryService } from "@roo-code/telemetry" import type { Anthropic } from "@anthropic-ai/sdk" import { PendingActionSettlementError, Task } from "../Task" +import { parseMentions } from "../../mentions" +import { processUserContentMentions } from "../../mentions/processUserContentMentions" +import type { FileContextTracker } from "../../context-tracking/FileContextTracker" import { ClineProvider } from "../../webview/ClineProvider" import { ContextProxy } from "../../config/ContextProxy" import { providerIdentifiers } from "@roo-code/types/provider-identifiers" @@ -2219,6 +2222,93 @@ describe("Task persistence", () => { }) }) + describe("resumeTaskFromHistory response mentions", () => { + const resumeWith = async (response: { text: string; segments?: unknown }) => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + historyItem: { + id: "resume-mentions", + number: 1, + ts: Date.now(), + task: "Resumed task", + tokensIn: 10, + tokensOut: 5, + totalCost: 0.001, + }, + startTask: false, + initialStatus: "interrupted", + }) + const initiate = vi.spyOn(getTaskPersistenceAccess(task), "initiateTaskLoop").mockResolvedValue(undefined) + vi.spyOn(task, "ask").mockResolvedValue({ response: "messageResponse", ...response } as never) + mockReadApiMessages.mockResolvedValue([{ role: "assistant", content: [{ type: "text", text: "Hi" }] }]) + mockProvider.taskHistoryStore.get = vi + .fn() + .mockReturnValue({ id: "resume-mentions", status: "interrupted" }) + + await getTaskPersistenceAccess(task).resumeTaskFromHistory() + + vi.mocked(parseMentions).mockClear() + const processed = await processUserContentMentions({ + userContent: initiate.mock.calls[0][0], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + return { processed, parsed: vi.mocked(parseMentions).mock.calls.map((call) => call[0]) } + } + + it("expands a user segment and keeps the wrapper", async () => { + const text = "read @/allowed please" + const { processed, parsed } = await resumeWith({ text, segments: [{ text, source: "user" }] }) + + expect(parsed).toEqual([text]) + expect(processed.content).toContainEqual({ + type: "text", + text: `\nprocessed: ${text}\n`, + }) + }) + + it("expands only the user segments of an interleaved answer", async () => { + const text = "@/a @/secret @/b" + const { processed, parsed } = await resumeWith({ + text, + segments: [ + { text: "@/a ", source: "user" }, + { text: "@/secret ", source: "generated" }, + { text: "@/b", source: "user" }, + ], + }) + + expect(parsed).toEqual(["@/a ", "@/b"]) + expect(processed.content).toContainEqual({ + type: "text", + text: "\nprocessed: @/a @/secret processed: @/b\n", + }) + }) + + it.each([ + { label: "omitted", segments: undefined }, + { label: "not matching the text", segments: [{ text: "other", source: "user" }] }, + { label: "matching but generated only", segments: [{ text: "read @/secret", source: "generated" }] }, + { label: "matching but of an invalid source", segments: [{ text: "read @/secret", source: "system" }] }, + ])("expands nothing when the segments are $label", async ({ segments }) => { + const { processed, parsed } = await resumeWith({ text: "read @/secret", segments }) + + expect(parsed).toEqual([]) + expect(processed.content).toContainEqual({ + type: "text", + text: "\nread @/secret\n", + }) + }) + + it("pushes no response block for an empty response", async () => { + const { processed, parsed } = await resumeWith({ text: "" }) + + expect(parsed).toEqual([]) + expect(JSON.stringify(processed.content)).not.toContain("") + }) + }) + describe("resumeTaskFromHistory", () => { it.each(["not_found", "invalid", "io_error"] as const)( "does not persist when hydration fails with %s", diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index c56d5291f6..eaf3065af0 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -32,7 +32,8 @@ import { ClineProvider } from "../../webview/ClineProvider" import { ApiStreamChunk } from "../../../api/transform/stream" import { OutputTokenLimitError } from "../../../api/providers/utils/output-token-limit-error" import { ContextProxy } from "../../config/ContextProxy" -import { processUserContentMentions } from "../../mentions/processUserContentMentions" +import { processUserContentMentions, userMentionContent } from "../../mentions/processUserContentMentions" +import { parseMentions } from "../../mentions" import { MultiSearchReplaceDiffStrategy } from "../../diff/strategies/multi-search-replace" import type { ApiMessage } from "../../task-persistence" import type { RequestPolicySnapshot } from "../../prompts/tools/effective-tool-policy" @@ -3511,7 +3512,11 @@ describe("Cline", () => { }) describe("processUserContentMentions", () => { - it("should process mentions in user_message tags", async () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it("expands only the authorized user input and leaves tool results unexpanded", async () => { const [cline, task] = Task.create({ provider: mockProvider, apiConfiguration: mockApiConfig, @@ -3523,10 +3528,14 @@ describe("Cline", () => { type: "text", text: "Regular text with 'some/path' (see below for file content)", } as const, - { - type: "text", - text: "Text with 'some/path' (see below for file content) in user_message tags", - } as const, + userMentionContent([ + { text: "", source: "generated" }, + { + text: "Text with 'some/path' (see below for file content) in user_message tags", + source: "user", + }, + { text: "", source: "generated" }, + ]), { type: "tool_result", tool_use_id: "test-id", @@ -3560,16 +3569,17 @@ describe("Cline", () => { "Regular text with 'some/path' (see below for file content)", ) - // Text within user_message tags should be processed + // Only the authorized input is replaced; the surrounding + // wrapper text stays intact. expect((processedContent[1] as Anthropic.TextBlockParam).text).toContain("processed:") expect((processedContent[1] as Anthropic.TextBlockParam).text).toContain( - "Text with 'some/path' (see below for file content) in user_message tags", + "processed: Text with 'some/path' (see below for file content) in user_message tags", ) - // user_message tag content should be processed + // Tool result markers cannot authorize mention expansion const toolResult1 = processedContent[2] as Anthropic.ToolResultBlockParam const content1 = Array.isArray(toolResult1.content) ? toolResult1.content[0] : toolResult1.content - expect((content1 as Anthropic.TextBlockParam).text).toContain("processed:") + expect((content1 as Anthropic.TextBlockParam).text).not.toContain("processed:") expect((content1 as Anthropic.TextBlockParam).text).toContain( "Check 'some/path' (see below for file content)", ) @@ -3584,6 +3594,64 @@ describe("Cline", () => { await cline.abortTask(true) await task.catch(() => {}) }) + + it.each([ + { + label: "all generated", + segments: [{ text: "read '@/secret' please", source: "generated" as const }], + }, + { label: "omitted", segments: undefined }, + { label: "not matching the text", segments: [{ text: "other", source: "user" as const }] }, + ])("does not expand task text when segments are $label", async ({ segments }) => { + const [cline, task] = Task.create({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "read '@/secret' please", + segments, + }) + + // Wait until the loop reached mention processing, then stop it. + await vi.waitFor(() => expect(processUserContentMentions).toHaveBeenCalled()) + await cline.abortTask(true) + await task.catch(() => {}) + + expect(parseMentions).not.toHaveBeenCalled() + }) + + it("expands the task text when the start authorizes it", async () => { + const taskText = "read '@/secret' please" + const [cline, task] = Task.create({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: taskText, + segments: [{ text: taskText, source: "user" }], + }) + + await vi.waitFor(() => expect(parseMentions).toHaveBeenCalled()) + await cline.abortTask(true) + await task.catch(() => {}) + + expect(vi.mocked(parseMentions).mock.calls[0]?.[0]).toBe(taskText) + }) + + it("expands only the authorized substring of generated task text", async () => { + const [cline, task] = Task.create({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "context @/context.txt\nchosen @/chosen.txt", + segments: [ + { text: "context @/context.txt\n", source: "generated" }, + { text: "chosen @/chosen.txt", source: "user" }, + ], + }) + + await vi.waitFor(() => expect(parseMentions).toHaveBeenCalled()) + await cline.abortTask(true) + await task.catch(() => {}) + + const texts = vi.mocked(parseMentions).mock.calls.map(([text]) => text) + expect(texts).toEqual(["chosen @/chosen.txt"]) + }) }) }) @@ -4204,11 +4272,55 @@ describe("Cline", () => { await task.submitUserMessage("test message", ["image1.png"]) // Verify handleWebviewAskResponse was called directly (not webview) - expect(handleResponseSpy).toHaveBeenCalledWith("messageResponse", "test message", ["image1.png"]) + expect(handleResponseSpy).toHaveBeenCalledWith( + "messageResponse", + "test message", + ["image1.png"], + [{ text: "test message", source: "generated" }], + ) // Should NOT route through webview anymore expect(mockProvider.postMessageToWebview).not.toHaveBeenCalled() }) + it("forwards the segments to the ask response", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "initial task", + startTask: false, + }) + const handleResponseSpy = vi.spyOn(task, "handleWebviewAskResponse") + + const segments = [ + { text: "a @/x ", source: "generated" as const }, + { text: "b", source: "user" as const }, + ] + await task.submitUserMessage("a @/x b", [], undefined, undefined, segments) + + expect(handleResponseSpy).toHaveBeenCalledWith("messageResponse", "a @/x b", [], segments) + }) + + it("trims the segments the same way as the text", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "initial task", + startTask: false, + }) + const handleResponseSpy = vi.spyOn(task, "handleWebviewAskResponse") + + await task.submitUserMessage(" @/chosen.txt ", [], undefined, undefined, [ + { text: " @/chosen.txt ", source: "user" }, + ]) + + expect(handleResponseSpy).toHaveBeenCalledWith( + "messageResponse", + "@/chosen.txt", + [], + [{ text: "@/chosen.txt", source: "user" }], + ) + }) + it("uses a mode selected through submitUserMessage in the next API request", async () => { vi.spyOn(mockProvider, "getState").mockResolvedValue(providerStateWith({ mode: "ask" })) vi.spyOn(mockProvider, "setMode").mockResolvedValue(undefined) @@ -4299,7 +4411,12 @@ describe("Cline", () => { task.clineMessages = [] await task.submitUserMessage("new task", ["image1.png"]) - expect(handleResponseSpy).toHaveBeenCalledWith("messageResponse", "new task", ["image1.png"]) + expect(handleResponseSpy).toHaveBeenCalledWith( + "messageResponse", + "new task", + ["image1.png"], + [{ text: "new task", source: "generated" }], + ) // Clear mock handleResponseSpy.mockClear() @@ -4315,7 +4432,12 @@ describe("Cline", () => { ] await task.submitUserMessage("follow-up message", ["image2.png"]) - expect(handleResponseSpy).toHaveBeenCalledWith("messageResponse", "follow-up message", ["image2.png"]) + expect(handleResponseSpy).toHaveBeenCalledWith( + "messageResponse", + "follow-up message", + ["image2.png"], + [{ text: "follow-up message", source: "generated" }], + ) }) it("should handle undefined provider gracefully", async () => { @@ -8076,7 +8198,7 @@ describe("Queued message processing after condense", () => { const submitSpy = vi.spyOn(task, "submitUserMessage").mockResolvedValue(undefined) // Queue a message during condensing - task.messageQueueService.addMessage("queued text", ["img1.png"]) + task.messageQueueService.addMessage("queued text", ["img1.png"], [{ text: "queued text", source: "user" }]) // Use fake timers to capture setTimeout(0) in processQueuedMessages vi.useFakeTimers() @@ -8086,7 +8208,9 @@ describe("Queued message processing after condense", () => { vi.runAllTimers() vi.useRealTimers() - expect(submitSpy).toHaveBeenCalledWith("queued text", ["img1.png"]) + expect(submitSpy).toHaveBeenCalledWith("queued text", ["img1.png"], undefined, undefined, [ + { text: "queued text", source: "user" }, + ]) expect(task.messageQueueService.isEmpty()).toBe(true) }) @@ -8122,7 +8246,7 @@ describe("Queued message processing after condense", () => { vi.runAllTimers() vi.useRealTimers() - expect(spyA).toHaveBeenCalledWith("A message", undefined) + expect(spyA).toHaveBeenCalledWith("A message", undefined, undefined, undefined, undefined) expect(spyB).not.toHaveBeenCalled() expect(taskB.messageQueueService.isEmpty()).toBe(false) @@ -8132,7 +8256,7 @@ describe("Queued message processing after condense", () => { vi.runAllTimers() vi.useRealTimers() - expect(spyB).toHaveBeenCalledWith("B message", undefined) + expect(spyB).toHaveBeenCalledWith("B message", undefined, undefined, undefined, undefined) expect(taskB.messageQueueService.isEmpty()).toBe(true) }) }) diff --git a/src/core/task/__tests__/ask-queued-message-drain.spec.ts b/src/core/task/__tests__/ask-queued-message-drain.spec.ts index 16fbb1c93e..a0b4ac91fa 100644 --- a/src/core/task/__tests__/ask-queued-message-drain.spec.ts +++ b/src/core/task/__tests__/ask-queued-message-drain.spec.ts @@ -53,6 +53,17 @@ describe("Task.ask queued message drain", () => { expect(result.text).toBe("picked answer") }) + it("forwards the queued message segments to the ask response", async () => { + const task = await createTask() + const segments = [{ text: "@/a", source: "user" as const }] + + const askPromise = task.ask("followup", "Q?", false) + task.messageQueueService.addMessage("@/a", undefined, segments) + + const result = await askPromise + expect(result.segments).toEqual(segments) + }) + it("does not consume queued messages for command_output asks", async () => { const task = await createTask() diff --git a/src/core/tools/AskFollowupQuestionTool.ts b/src/core/tools/AskFollowupQuestionTool.ts index 3baea7ed52..54c46a0dce 100644 --- a/src/core/tools/AskFollowupQuestionTool.ts +++ b/src/core/tools/AskFollowupQuestionTool.ts @@ -1,3 +1,5 @@ +import { userMentionContent, wrapUserMessage } from "../mentions/processUserContentMentions" +import { normalizeSegments } from "../../shared/mentionSegments" import { Task } from "../task/Task" import { formatResponse } from "../prompts/responses" import type { ToolUse } from "../../shared/tools" @@ -70,9 +72,11 @@ export class AskFollowupQuestionTool extends BaseTool<"ask_followup_question"> { } task.consecutiveMistakeCount = 0 - const { text, images } = await task.ask("followup", JSON.stringify(follow_up_json), false) + const { text, images, segments } = await task.ask("followup", JSON.stringify(follow_up_json), false) const safeText = text ?? "" await task.say("user_feedback", safeText, images) + // Only user segments expand. The tool result repeats the text by design. + task.userMessageContent.push(userMentionContent(wrapUserMessage(normalizeSegments(safeText, segments)))) pushToolResult(formatResponse.toolResult(`\n${safeText}\n`, images)) } catch (error) { await handleError("asking question", error as Error) diff --git a/src/core/tools/AttemptCompletionTool.ts b/src/core/tools/AttemptCompletionTool.ts index 4fe03ce94c..e8fe148d57 100644 --- a/src/core/tools/AttemptCompletionTool.ts +++ b/src/core/tools/AttemptCompletionTool.ts @@ -1,3 +1,5 @@ +import { userMentionContent, wrapUserMessage } from "../mentions/processUserContentMentions" +import { normalizeSegments } from "../../shared/mentionSegments" import * as vscode from "vscode" import { RooCodeEventName, type HistoryItem, type PendingTaskAction } from "@roo-code/types" @@ -218,7 +220,7 @@ export class AttemptCompletionTool extends BaseTool<"attempt_completion"> { task.flushTelemetryInstallment("attempt_completion") } - const { response, text, images, queuedMessageId } = await task.ask("completion_result", "", false) + const { response, text, images, queuedMessageId, segments } = await task.ask("completion_result", "", false) if (response === "yesButtonClicked") { // A stale history replay reruns this handler on a fresh Task instance for a @@ -245,6 +247,8 @@ export class AttemptCompletionTool extends BaseTool<"attempt_completion"> { } const feedbackText = `\n${text}\n` + // Only user segments expand. The tool result repeats the text by design. + task.userMessageContent.push(userMentionContent(wrapUserMessage(normalizeSegments(text ?? "", segments)))) pushToolResult(formatResponse.toolResult(feedbackText, images)) } catch (error) { await handleError("inspecting site", error as Error) diff --git a/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts b/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts index 965cb23e19..6af1b91ff5 100644 --- a/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts +++ b/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts @@ -6,23 +6,36 @@ import type { ToolCallbacks } from "../BaseTool" import { AskFollowupQuestionTool, askFollowupQuestionTool } from "../AskFollowupQuestionTool" import { formatResponse } from "../../prompts/responses" import { NativeToolCallParser } from "../../assistant-message/NativeToolCallParser" +import { processUserContentMentions } from "../../mentions/processUserContentMentions" +import { parseMentions } from "../../mentions/index" +import { FileContextTracker } from "../../context-tracking/FileContextTracker" + +vi.mock("../../mentions/index", () => ({ parseMentions: vi.fn() })) describe("AskFollowupQuestionTool", () => { let tool: AskFollowupQuestionTool let mockTask: Task let mockCallbacks: ToolCallbacks + // Kept outside the Task double so tests can re-stub answers without casts. + let askMock: ReturnType beforeEach(() => { vi.clearAllMocks() + vi.mocked(parseMentions).mockImplementation(async (text) => ({ + text: `parsed: ${text}`, + contentBlocks: [], + })) tool = new AskFollowupQuestionTool() + askMock = vi.fn().mockResolvedValue({ text: "User answer", images: [] }) mockTask = { + userMessageContent: [], consecutiveMistakeCount: 0, recordToolError: vi.fn(), didToolFailInCurrentTurn: false, sayAndCreateMissingParamError: vi.fn().mockResolvedValue("Missing parameter error"), - ask: vi.fn().mockResolvedValue({ text: "User answer", images: [] }), + ask: askMock, say: vi.fn().mockResolvedValue(undefined), } as unknown as Task @@ -197,6 +210,107 @@ describe("AskFollowupQuestionTool", () => { ) }) + it("pushes an authorized feedback block that mention processing expands", async () => { + const params = { question: "What?", follow_up: [{ text: "A" }] } + askMock.mockResolvedValue({ text: "@/allowed", images: [], segments: [{ text: "@/allowed", source: "user" }] }) + + await tool.execute(params, mockTask, mockCallbacks) + + // Deleting the push in AskFollowupQuestionTool fails this assertion. + expect(mockTask.userMessageContent).toHaveLength(1) + const block = mockTask.userMessageContent[0] + expect(block).toEqual({ type: "text", text: "\n@/allowed\n" }) + + // The pushed block authorizes expansion of its input. + const result = await processUserContentMentions({ + userContent: [block], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(vi.mocked(parseMentions).mock.calls[0]?.[0]).toBe("@/allowed") + expect(result.content).toEqual([{ type: "text", text: "\nparsed: @/allowed\n" }]) + }) + + it("expands only the user segment of a mixed answer", async () => { + const params = { question: "What?", follow_up: [{ text: "A" }] } + askMock.mockResolvedValue({ + text: "@/secret @/mine", + images: [], + segments: [ + { text: "@/secret ", source: "generated" }, + { text: "@/mine", source: "user" }, + ], + }) + + await tool.execute(params, mockTask, mockCallbacks) + + await processUserContentMentions({ + userContent: [mockTask.userMessageContent[0]], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(vi.mocked(parseMentions).mock.calls.map((call) => call[0])).toEqual(["@/mine"]) + }) + + it("keeps generated text between user segments and expands both user segments", async () => { + const params = { question: "What?", follow_up: [{ text: "A" }] } + askMock.mockResolvedValue({ + text: "@/a @/secret @/b", + images: [], + segments: [ + { text: "@/a ", source: "user" }, + { text: "@/secret ", source: "generated" }, + { text: "@/b", source: "user" }, + ], + }) + vi.mocked(parseMentions).mockImplementation(async (input) => ({ text: `parsed: ${input}`, contentBlocks: [] })) + + await tool.execute(params, mockTask, mockCallbacks) + + const result = await processUserContentMentions({ + userContent: [mockTask.userMessageContent[0]], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(result.content).toEqual([ + { type: "text", text: "\nparsed: @/a @/secret parsed: @/b\n" }, + ]) + }) + + it.each([ + { label: "missing", segments: undefined }, + { label: "not matching the text", segments: [{ text: "@/other", source: "user" }] }, + { label: "invalid", segments: [{ text: "@/secret", source: "system" }] }, + ])("expands nothing when the segments are $label", async ({ segments }) => { + const params = { question: "What?", follow_up: [{ text: "A" }] } + askMock.mockResolvedValue({ text: "@/secret", images: [], segments }) + + await tool.execute(params, mockTask, mockCallbacks) + + await processUserContentMentions({ + userContent: [mockTask.userMessageContent[0]], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(parseMentions).not.toHaveBeenCalled() + }) + + it("does not authorize the feedback block when the user answer is null", async () => { + const params = { question: "What?", follow_up: [{ text: "A" }] } + askMock.mockResolvedValue({ text: null, images: [] }) + + await tool.execute(params, mockTask, mockCallbacks) + + expect(mockTask.userMessageContent).toHaveLength(1) + const result = await processUserContentMentions({ + userContent: [mockTask.userMessageContent[0]], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(parseMentions).not.toHaveBeenCalled() + expect(result.content).toEqual([{ type: "text", text: "\n\n" }]) + }) + it("should pass images from user response to toolResult", async () => { const params = { question: "Look at this?", diff --git a/src/core/tools/__tests__/attemptCompletionTool.spec.ts b/src/core/tools/__tests__/attemptCompletionTool.spec.ts index 29580ef203..c1a5689034 100644 --- a/src/core/tools/__tests__/attemptCompletionTool.spec.ts +++ b/src/core/tools/__tests__/attemptCompletionTool.spec.ts @@ -31,6 +31,11 @@ import { attemptCompletionTool, AttemptCompletionCallbacks } from "../AttemptCom import { Task } from "../../task/Task" import { AskApproval, HandleError, PushToolResult } from "../../../shared/tools" import * as vscode from "vscode" +import { processUserContentMentions } from "../../mentions/processUserContentMentions" +import { parseMentions } from "../../mentions/index" +import { FileContextTracker } from "../../context-tracking/FileContextTracker" + +vi.mock("../../mentions/index", () => ({ parseMentions: vi.fn() })) describe("attemptCompletionTool", () => { let mockTask: Partial @@ -60,6 +65,7 @@ describe("attemptCompletionTool", () => { vi.mocked(vscode.workspace.getConfiguration).mockImplementation(mockGetConfiguration) mockTask = { + userMessageContent: [], consecutiveMistakeCount: 0, recordToolError: vi.fn(), todoList: undefined, @@ -1161,6 +1167,187 @@ describe("attemptCompletionTool", () => { expect(mockPushToolResult).toHaveBeenCalledTimes(1) }) + it("keeps generated text between user segments and expands both user segments", async () => { + const block: ToolUse<"attempt_completion"> = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Done" }, + nativeArgs: { result: "Done" }, + partial: false, + } + mockTask.ask = vi.fn().mockResolvedValue({ + response: "messageResponse", + text: "@/a @/secret @/b", + images: [], + segments: [ + { text: "@/a ", source: "user" }, + { text: "@/secret ", source: "generated" }, + { text: "@/b", source: "user" }, + ], + }) + vi.mocked(parseMentions).mockImplementation(async (input) => ({ + text: `parsed: ${input}`, + contentBlocks: [], + })) + + await attemptCompletionTool.handle(mockTask as Task, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + toolCallId: "", + askFinishSubTaskApproval: mockAskFinishSubTaskApproval, + toolDescription: mockToolDescription, + }) + + const result = await processUserContentMentions({ + userContent: [mockTask.userMessageContent![0]], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(result.content).toEqual([ + { type: "text", text: "\nparsed: @/a @/secret parsed: @/b\n" }, + ]) + }) + + it.each([ + { + label: "mixed origin", + segments: [ + { text: "@/secret ", source: "generated" }, + { text: "@/mine", source: "user" }, + ], + text: "@/secret @/mine", + expanded: ["@/mine"], + }, + { label: "missing", segments: undefined, text: "@/secret", expanded: [] }, + { + label: "not matching the text", + segments: [{ text: "@/other", source: "user" }], + text: "@/secret", + expanded: [], + }, + { + label: "invalid", + segments: [{ text: "@/secret", source: "system" }], + text: "@/secret", + expanded: [], + }, + ])("expands only authorized text when the segments are $label", async ({ segments, text, expanded }) => { + const block: ToolUse<"attempt_completion"> = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Done" }, + nativeArgs: { result: "Done" }, + partial: false, + } + mockTask.ask = vi.fn().mockResolvedValue({ response: "messageResponse", text, images: [], segments }) + vi.mocked(parseMentions).mockImplementation(async (input) => ({ + text: `parsed: ${input}`, + contentBlocks: [], + })) + vi.mocked(parseMentions).mockClear() + + await attemptCompletionTool.handle(mockTask as Task, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + toolCallId: "", + askFinishSubTaskApproval: mockAskFinishSubTaskApproval, + toolDescription: mockToolDescription, + }) + + await processUserContentMentions({ + userContent: [mockTask.userMessageContent![0]], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(vi.mocked(parseMentions).mock.calls.map((call) => call[0])).toEqual(expanded) + }) + + it("pushes an authorized feedback block that mention processing expands", async () => { + const block: ToolUse<"attempt_completion"> = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Done" }, + nativeArgs: { result: "Done" }, + partial: false, + } + mockTask.ask = vi.fn().mockResolvedValue({ + response: "messageResponse", + text: "@/allowed", + images: [], + segments: [{ text: "@/allowed", source: "user" }], + }) + vi.mocked(parseMentions).mockImplementation(async (text) => ({ + text: `parsed: ${text}`, + contentBlocks: [], + })) + vi.mocked(parseMentions).mockClear() + + await attemptCompletionTool.handle(mockTask as Task, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + toolCallId: "", + askFinishSubTaskApproval: mockAskFinishSubTaskApproval, + toolDescription: mockToolDescription, + }) + + // Deleting the push in AttemptCompletionTool fails this assertion. + expect(mockTask.userMessageContent).toHaveLength(1) + const block0 = mockTask.userMessageContent![0] + expect(block0).toEqual({ type: "text", text: "\n@/allowed\n" }) + + // The pushed block authorizes expansion of its input. + const result = await processUserContentMentions({ + userContent: [block0], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(vi.mocked(parseMentions).mock.calls[0]?.[0]).toBe("@/allowed") + expect(result.content).toEqual([ + { type: "text", text: "\nparsed: @/allowed\n" }, + ]) + }) + + it("does not authorize the feedback block when the feedback text is undefined", async () => { + const block: ToolUse<"attempt_completion"> = { + type: "tool_use", + name: "attempt_completion", + params: { result: "Done" }, + nativeArgs: { result: "Done" }, + partial: false, + } + mockTask.ask = vi.fn().mockResolvedValue({ + response: "messageResponse", + text: undefined, + images: [], + }) + vi.mocked(parseMentions).mockImplementation(async (text) => ({ + text: `parsed: ${text}`, + contentBlocks: [], + })) + vi.mocked(parseMentions).mockClear() + + await attemptCompletionTool.handle(mockTask as Task, block, { + askApproval: mockAskApproval, + handleError: mockHandleError, + pushToolResult: mockPushToolResult, + toolCallId: "", + askFinishSubTaskApproval: mockAskFinishSubTaskApproval, + toolDescription: mockToolDescription, + }) + + expect(mockTask.userMessageContent).toHaveLength(1) + const result = await processUserContentMentions({ + userContent: [mockTask.userMessageContent![0]], + cwd: "/test", + fileContextTracker: {} as FileContextTracker, + }) + expect(parseMentions).not.toHaveBeenCalled() + expect(result.content).toEqual([{ type: "text", text: "\n\n" }]) + }) + it("does not clear pending metadata when stale delegation continues without an action id", async () => { const block: ToolUse<"attempt_completion"> = { type: "tool_use", diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 6cec3bdbe7..63bb9eea06 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -90,6 +90,7 @@ import { CloudService, getRooCodeApiUrl } from "@roo-code/cloud" import { Package } from "../../shared/package" import { findLast } from "../../shared/array" import { supportPrompt } from "../../shared/support-prompt" +import { generatedSegments, joinSegments } from "../../shared/mentionSegments" import { GlobalFileNames } from "../../shared/globalFileNames" import { Mode, defaultModeSlug, getModeBySlug } from "../../shared/modes" import { experimentDefault } from "../../shared/experiments" @@ -138,7 +139,7 @@ import { PendingActionSettlementError, Task } from "../task/Task" import { webviewMessageHandler } from "./webviewMessageHandler" import type { WebviewFocusTracker } from "./WebviewFocusTracker" -import type { ClineMessage, TodoItem } from "@roo-code/types" +import type { ClineMessage, MentionSegment, TodoItem } from "@roo-code/types" import { type ApiMessage, readApiMessages, @@ -1068,19 +1069,22 @@ export class ClineProvider const { customSupportPrompts } = await this.getState() // TODO: Improve type safety for promptType. - const prompt = supportPrompt.create(promptType, params, customSupportPrompts) + // Only the typed input is a user segment. The template and its other params are generated. + const segments = supportPrompt.createSegments(promptType, params, customSupportPrompts) + const prompt = joinSegments(segments) if (command === "addToContext") { await this.postMessageToWebview({ type: "invoke", invoke: "setChatBoxMessage", text: `${prompt}\n\n`, + segments: [...segments, ...generatedSegments("\n\n")], }) await this.postMessageToWebview({ type: "action", action: "focusInput" }) return } - await this.createTask(prompt) + await this.createTask(prompt, undefined, undefined, { segments }) } public static async handleTerminalAction( @@ -1097,20 +1101,22 @@ export class ClineProvider } const { customSupportPrompts } = await visibleProvider.getState() - const prompt = supportPrompt.create(promptType, params, customSupportPrompts) + const segments = supportPrompt.createSegments(promptType, params, customSupportPrompts) + const prompt = joinSegments(segments) if (command === "terminalAddToContext") { await visibleProvider.postMessageToWebview({ type: "invoke", invoke: "setChatBoxMessage", text: `${prompt}\n\n`, + segments: [...segments, ...generatedSegments("\n\n")], }) await visibleProvider.postMessageToWebview({ type: "action", action: "focusInput" }) return } try { - await visibleProvider.createTask(prompt) + await visibleProvider.createTask(prompt, undefined, undefined, { segments }) } catch (error) { if (error instanceof OrganizationAllowListViolationError) { // Errors from terminal commands seem to get swallowed / ignored. @@ -4268,7 +4274,10 @@ export class ClineProvider text?: string, images?: string[], parentTask?: Task, - options: CreateTaskOptions & { handoffExecutionContext?: DelegatedChildContext } = {}, + options: CreateTaskOptions & { + handoffExecutionContext?: DelegatedChildContext + segments?: MentionSegment[] + } = {}, configuration: RooCodeSettings = {}, ): Promise { if (configuration) { diff --git a/src/core/webview/__tests__/ClineProvider.spec.ts b/src/core/webview/__tests__/ClineProvider.spec.ts index 464889418e..255af5788c 100644 --- a/src/core/webview/__tests__/ClineProvider.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.spec.ts @@ -425,6 +425,13 @@ afterAll(() => { vi.restoreAllMocks() }) +// Matches a non-empty segment list that holds only generated text. +const noUserSegments = { + asymmetricMatch: (value: unknown) => + Array.isArray(value) && value.length > 0 && value.every((segment) => segment?.source === "generated"), + toString: () => "onlyGeneratedSegments", +} + describe("ClineProvider", () => { beforeAll(() => { vi.mocked(Task).mockImplementation(function (options: any) { @@ -743,6 +750,8 @@ describe("ClineProvider", () => { type: "invoke", invoke: "setChatBoxMessage", text: expect.stringContaining("selected code"), + // Generated context has no user segment. + segments: noUserSegments, }) expect(sidebar.postMessage).toHaveBeenCalledWith({ type: "action", action: "focusInput" }) expect(tab.postMessage).not.toHaveBeenCalled() @@ -933,6 +942,8 @@ describe("ClineProvider", () => { type: "invoke", invoke: "setChatBoxMessage", text: expect.stringContaining("selected code"), + // Generated context is never the authorized substring. + segments: noUserSegments, }) expect(visible.postMessage).toHaveBeenCalledWith({ type: "action", action: "focusInput" }) expect(vscode.commands.executeCommand).not.toHaveBeenCalled() @@ -961,6 +972,8 @@ describe("ClineProvider", () => { type: "invoke", invoke: "setChatBoxMessage", text: expect.stringContaining("selected code"), + // Generated context is never the authorized substring. + segments: noUserSegments, }) expect(tab.postMessage).not.toHaveBeenCalled() expect(panel.reveal).not.toHaveBeenCalled() @@ -1017,7 +1030,12 @@ describe("ClineProvider", () => { const target = provider.webviewFocusTracker.getLastActiveProvider()! await target.handleCodeAction(command, promptType, { selectedText: "selected code" }) expect(sidebarTask).toHaveBeenCalledOnce() - expect(sidebarTask).toHaveBeenCalledWith(expect.stringContaining("selected code")) + expect(sidebarTask).toHaveBeenCalledWith( + expect.stringContaining("selected code"), + undefined, + undefined, + { segments: noUserSegments }, + ) expect(tabTask).not.toHaveBeenCalled() } finally { sidebarTask.mockRestore() @@ -1036,6 +1054,52 @@ describe("ClineProvider", () => { sidebarTask.mockRestore() tabTask.mockRestore() }) + + test.each([ + ["explainCode", "EXPLAIN"], + ["fixCode", "FIX"], + ["improveCode", "IMPROVE"], + ] as const)( + "forwards user-authored code action input as a user segment for %s", + async (command, promptType) => { + const task = new Task(defaultTaskOptions) + const createTask = vi.spyOn(provider, "createTask").mockResolvedValue(task) + try { + await provider.handleCodeAction(command, promptType, { + selectedText: "selected code", + userInput: "@/allowed", + }) + // A regression that always sends "" would fail this assertion. + expect(createTask).toHaveBeenCalledWith(expect.any(String), undefined, undefined, { + segments: expect.arrayContaining([{ text: "@/allowed", source: "user" }]), + }) + } finally { + createTask.mockRestore() + } + }, + ) + + test("forwards user-authored terminal action input as a user segment", async () => { + const task = new Task(defaultTaskOptions) + const sidebarTask = vi.spyOn(provider, "createTask").mockResolvedValue(task) + const tabTask = vi.spyOn(tabProvider, "createTask").mockResolvedValue(task) + try { + // No chat focus fired, so the last registered visible provider + // (the tab) receives the task, as in the fallback test above. + await ClineProvider.handleTerminalAction("terminalFixCommand", "TERMINAL_FIX", { + userInput: "@/allowed", + }) + // A regression that always sends "" would fail this assertion. + expect(tabTask).toHaveBeenCalledOnce() + expect(tabTask).toHaveBeenCalledWith(expect.any(String), undefined, undefined, { + segments: expect.arrayContaining([{ text: "@/allowed", source: "user" }]), + }) + expect(sidebarTask).not.toHaveBeenCalled() + } finally { + sidebarTask.mockRestore() + tabTask.mockRestore() + } + }) }) test("loads full model details when preparing an LM Studio task", async () => { @@ -7188,7 +7252,13 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => { expect(mockCline.overwriteClineMessages).toHaveBeenCalledWith([mockMessages[0]]) expect(mockCline.overwriteApiConversationHistory).toHaveBeenCalledWith([{ ts: 1000 }]) // Verify submitUserMessage was called with the edited content - expect(mockCline.submitUserMessage).toHaveBeenCalledWith("Edited message with preserved images", []) + expect(mockCline.submitUserMessage).toHaveBeenCalledWith( + "Edited message with preserved images", + [], + undefined, + undefined, + [{ text: "Edited message with preserved images", source: "generated" }], + ) }) test("handles editing messages with file attachments", async () => { @@ -7241,7 +7311,13 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => { }) expect(mockCline.overwriteClineMessages).toHaveBeenCalled() - expect(mockCline.submitUserMessage).toHaveBeenCalledWith("Edited message with file attachment", []) + expect(mockCline.submitUserMessage).toHaveBeenCalledWith( + "Edited message with file attachment", + [], + undefined, + undefined, + [{ text: "Edited message with file attachment", source: "generated" }], + ) }) }) @@ -7772,7 +7848,9 @@ describe("ClineProvider - Comprehensive Edit/Delete Edge Cases", () => { await messageHandler({ type: "editMessageConfirm", messageTs: 2000, text: largeEditedContent }) expect(mockCline.overwriteClineMessages).toHaveBeenCalled() - expect(mockCline.submitUserMessage).toHaveBeenCalledWith(largeEditedContent, []) + expect(mockCline.submitUserMessage).toHaveBeenCalledWith(largeEditedContent, [], undefined, undefined, [ + { text: largeEditedContent, source: "generated" }, + ]) }) test("handles deleting messages with large payloads", async () => { diff --git a/src/core/webview/__tests__/webviewMessageHandler.newTask.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.newTask.spec.ts new file mode 100644 index 0000000000..ba0f3625ec --- /dev/null +++ b/src/core/webview/__tests__/webviewMessageHandler.newTask.spec.ts @@ -0,0 +1,160 @@ +import { describe, it, expect, beforeEach, vi } from "vitest" +import { webviewMessageHandler } from "../webviewMessageHandler" +import type { ClineProvider } from "../ClineProvider" +import type { MentionSegment } from "@roo-code/types" +import type { WebviewMessage } from "../../../shared/WebviewMessage" + +vi.mock("../../../i18n", () => ({ + t: vi.fn((key: string) => key), + changeLanguage: vi.fn(), +})) + +vi.mock("vscode", () => ({ + window: { + showErrorMessage: vi.fn(), + showWarningMessage: vi.fn(), + showInformationMessage: vi.fn(), + }, + workspace: { + workspaceFolders: undefined, + getConfiguration: vi.fn(() => ({ + get: vi.fn(), + update: vi.fn(), + })), + }, + ConfigurationTarget: { Global: 1, Workspace: 2, WorkspaceFolder: 3 }, + Uri: { + parse: vi.fn((str: string) => ({ toString: () => str })), + file: vi.fn((path: string) => ({ fsPath: path })), + }, + env: { openExternal: vi.fn() }, + commands: { executeCommand: vi.fn() }, +})) + +const { resolveImageMentions } = vi.hoisted(() => ({ + resolveImageMentions: vi.fn(async ({ images }: { images?: string[] }) => ({ text: "", images: images ?? [] })), +})) +vi.mock("../../mentions/resolveImageMentions", () => ({ resolveImageMentions })) + +vi.mock("../../task-persistence", async (importOriginal) => ({ + ...(await importOriginal()), + saveTaskMessages: vi.fn(), +})) + +const generated = (text: string): MentionSegment => ({ text, source: "generated" }) +const user = (text: string): MentionSegment => ({ text, source: "user" }) + +describe("webviewMessageHandler mention segment forwarding", () => { + // Narrow double: the handler under test only touches these members. + // The single cast at construction keeps every use site type-checked. + let provider: { + createTask: ReturnType + postMessageToWebview: ReturnType + getCurrentTask: ReturnType + getState: ReturnType + postStateToWebview: ReturnType + contextProxy: { globalStorageUri: { fsPath: string } } + } + let providerAsClineProvider: ClineProvider + let currentTask: { + handleWebviewAskResponse: ReturnType + messageQueueService: { addMessage: ReturnType } + } + + beforeEach(() => { + vi.clearAllMocks() + currentTask = { + handleWebviewAskResponse: vi.fn(), + messageQueueService: { addMessage: vi.fn() }, + } + provider = { + createTask: vi.fn().mockResolvedValue({ taskId: "task-1" }), + postMessageToWebview: vi.fn().mockResolvedValue(undefined), + getCurrentTask: vi.fn(() => currentTask), + getState: vi.fn().mockResolvedValue({ maxImageFileSize: 5, maxTotalImageSize: 20 }), + postStateToWebview: vi.fn().mockResolvedValue(undefined), + contextProxy: { globalStorageUri: { fsPath: "/tmp/storage" } }, + } + providerAsClineProvider = provider as unknown as ClineProvider + }) + + const text = "context @/secret\nuser line" + const segments = [generated("context @/secret\n"), user("user line")] + // Segments that the handler must replace with all generated text. + const rejected = [ + { label: "omitted", segments: undefined }, + { label: "not matching the text", segments: [user("other")] }, + { label: "invalid", segments: [{ text, source: "system" }] }, + ] + + describe.each([ + { + type: "newTask" as const, + read: () => provider.createTask.mock.calls[0], + segmentsOf: (call: unknown[]) => (call[3] as { segments: MentionSegment[] }).segments, + }, + { + type: "askResponse" as const, + read: () => currentTask.handleWebviewAskResponse.mock.calls[0], + segmentsOf: (call: unknown[]) => call[3] as MentionSegment[], + }, + { + type: "queueMessage" as const, + read: () => currentTask.messageQueueService.addMessage.mock.calls[0], + segmentsOf: (call: unknown[]) => call[2] as MentionSegment[], + }, + ])("$type", ({ type, read, segmentsOf }) => { + const send = (extra: Record) => + webviewMessageHandler(providerAsClineProvider, { + type, + askResponse: "messageResponse", + ...extra, + } as WebviewMessage) + + it("forwards matching segments", async () => { + await send({ text, segments }) + expect(segmentsOf(read()!)).toEqual(segments) + }) + + it.each(rejected)("treats segments that are $label as all generated", async ({ segments: bad }) => { + await send({ text, segments: bad }) + expect(segmentsOf(read()!)).toEqual([generated(text)]) + }) + + it.each([ + { label: "user segment", segments: [user("see @/a.png")], expected: "see @/a.png" }, + { label: "generated segment", segments: [generated("see @/a.png")], expected: "" }, + ])("scans only user text for images with a $label", async ({ segments: given, expected }) => { + await send({ text: "see @/a.png", segments: given }) + expect(resolveImageMentions).toHaveBeenCalledWith(expect.objectContaining({ text: expected })) + }) + }) + + it("keeps the original text and images when it forwards", async () => { + const message: WebviewMessage = { type: "newTask", text, segments } + await webviewMessageHandler(providerAsClineProvider, message) + + expect(provider.createTask).toHaveBeenCalledWith(text, [], undefined, expect.anything(), undefined) + }) + + it("authorizes nothing in an edited message", async () => { + const submitUserMessage = vi.fn() + const editTask = { + taskId: "task-1", + clineMessages: [{ ts: 1, type: "say", say: "user_feedback", text: "old" }], + apiConversationHistory: [], + overwriteClineMessages: vi.fn(), + overwriteApiConversationHistory: vi.fn(), + submitUserMessage, + messageManager: { rewindToTimestamp: vi.fn() }, + } + provider.getCurrentTask.mockReturnValue(editTask) + const message: WebviewMessage = { type: "editMessageConfirm", messageTs: 1, text: "see @/a.png" } + await webviewMessageHandler(providerAsClineProvider, message) + + expect(resolveImageMentions).toHaveBeenCalledWith(expect.objectContaining({ text: "" })) + expect(submitUserMessage).toHaveBeenCalledWith("see @/a.png", expect.anything(), undefined, undefined, [ + generated("see @/a.png"), + ]) + }) +}) diff --git a/src/core/webview/__tests__/webviewMessageHandler.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.spec.ts index 98c306a4f3..d49a9f9c1a 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.spec.ts @@ -964,9 +964,13 @@ describe("webviewMessageHandler - image mentions", () => { }) expect(vi.mocked(resolveImageMentions)).toHaveBeenCalled() - expect(mockHandleWebviewAskResponse).toHaveBeenCalledWith("messageResponse", "See @/img.png", [ - "data:image/png;base64,from-mention", - ]) + // Omitted segments fail closed, so the text is all generated. + expect(mockHandleWebviewAskResponse).toHaveBeenCalledWith( + "messageResponse", + "See @/img.png", + ["data:image/png;base64,from-mention"], + [{ text: "See @/img.png", source: "generated" }], + ) }) }) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index c5abbc0a12..c5f67e164e 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -30,6 +30,7 @@ import { RouterModelsMessageType, VsCodeLmModelsMessageType, isTelemetryOptedIn, + type MentionSegment, } from "@roo-code/types" import { customToolRegistry } from "@roo-code/core" import { CloudService } from "@roo-code/cloud" @@ -90,6 +91,7 @@ import { openGitHubAccountManagement, } from "../../api/providers/github-copilot" import { openMention } from "../mentions" +import { generatedSegments, normalizeSegments, userText } from "../../shared/mentionSegments" import { resolveImageMentions } from "../mentions/resolveImageMentions" import { RooIgnoreController } from "../ignore/RooIgnoreController" import { getWorkspacePath } from "../../utils/path" @@ -226,21 +228,25 @@ export const webviewMessageHandler = async ( * Resolves image file mentions in incoming messages. * Matches read_file behavior: respects size limits and model capabilities. */ - const resolveIncomingImages = async (payload: { text?: string; images?: string[] }) => { + const resolveIncomingImages = async (payload: { text?: string; images?: string[]; segments: MentionSegment[] }) => { const text = payload.text ?? "" const images = payload.images const currentTask = provider.getCurrentTask() const state = await provider.getState() + // Scan only user segments, so generated mentions attach no images. const resolved = await resolveImageMentions({ - text, + text: userText(payload.segments), images, cwd: getCurrentCwd(), rooIgnoreController: currentTask?.rooIgnoreController, maxImageFileSize: state.maxImageFileSize, maxTotalImageSize: state.maxTotalImageSize, }) - return resolved + return { text, images: resolved.images } } + // Fails closed: segments that are missing or do not match the text are all generated. + const segmentsOf = (message: { text?: string; segments?: unknown }) => + normalizeSegments(message.text ?? "", message.segments) /** * Shared utility to find message indices based on timestamp. * When multiple messages share the same timestamp (e.g., after condense), @@ -437,6 +443,7 @@ export const webviewMessageHandler = async ( editedContent: string, restoreCheckpoint?: boolean, images?: string[], + segments?: MentionSegment[], ): Promise => { const currentCline = provider.getCurrentTask() if (!currentCline) { @@ -556,7 +563,7 @@ export const webviewMessageHandler = async ( // Update the UI to reflect the deletion await provider.postStateToWebview() - await currentCline.submitUserMessage(editedContent, images) + await currentCline.submitUserMessage(editedContent, images, undefined, undefined, segments) } catch (error) { console.error("Error in edit message:", error) vscode.window.showErrorMessage( @@ -711,12 +718,19 @@ export const webviewMessageHandler = async ( // agentically running promises in old instance don't affect our new // task. This essentially creates a fresh slate for the new task. try { - const resolved = await resolveIncomingImages({ text: message.text, images: message.images }) + const resolved = await resolveIncomingImages({ + text: message.text, + images: message.images, + segments: segmentsOf(message), + }) await provider.createTask( resolved.text, resolved.images, undefined, - { taskId: message.taskId }, + { + taskId: message.taskId, + segments: segmentsOf(message), + }, message.taskConfiguration, ) // Task created successfully - notify the UI to reset @@ -736,10 +750,15 @@ export const webviewMessageHandler = async ( case "askResponse": { - const resolved = await resolveIncomingImages({ text: message.text, images: message.images }) + const segments = segmentsOf(message) + const resolved = await resolveIncomingImages({ + text: message.text, + images: message.images, + segments, + }) provider .getCurrentTask() - ?.handleWebviewAskResponse(message.askResponse!, resolved.text, resolved.images) + ?.handleWebviewAskResponse(message.askResponse!, resolved.text, resolved.images, segments) } break @@ -2612,12 +2631,19 @@ export const webviewMessageHandler = async ( break case "editMessageConfirm": if (message.messageTs && message.text) { - const resolved = await resolveIncomingImages({ text: message.text, images: message.images }) + // Stored messages carry no segments, so fail closed: nothing is expanded. + const editSegments = generatedSegments(message.text) + const resolved = await resolveIncomingImages({ + text: message.text, + images: message.images, + segments: editSegments, + }) await handleEditMessageConfirm( message.messageTs, resolved.text, message.restoreCheckpoint, resolved.images, + editSegments, ) } break @@ -4026,8 +4052,9 @@ export const webviewMessageHandler = async ( */ case "queueMessage": { - const resolved = await resolveIncomingImages({ text: message.text, images: message.images }) - provider.getCurrentTask()?.messageQueueService.addMessage(resolved.text, resolved.images) + const segments = segmentsOf(message) + const resolved = await resolveIncomingImages({ text: message.text, images: message.images, segments }) + provider.getCurrentTask()?.messageQueueService.addMessage(resolved.text, resolved.images, segments) break } case "removeQueuedMessage": { diff --git a/src/eslint-suppressions.json b/src/eslint-suppressions.json index ae3d504003..f06ec638eb 100644 --- a/src/eslint-suppressions.json +++ b/src/eslint-suppressions.json @@ -589,11 +589,6 @@ "count": 1 } }, - "core/mentions/__tests__/processUserContentMentions.spec.ts": { - "@typescript-eslint/no-explicit-any": { - "count": 1 - } - }, "core/mentions/index.ts": { "@typescript-eslint/no-explicit-any": { "count": 1 diff --git a/src/extension/__tests__/api-send-message.spec.ts b/src/extension/__tests__/api-send-message.spec.ts index dc51d06445..1804315c85 100644 --- a/src/extension/__tests__/api-send-message.spec.ts +++ b/src/extension/__tests__/api-send-message.spec.ts @@ -2,6 +2,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest" import * as vscode from "vscode" import { API } from "../api" +import { userSegments } from "../../shared/mentionSegments" import { ClineProvider } from "../../core/webview/ClineProvider" import { makeClineProviderFactory } from "../../test-utils/provider" @@ -53,6 +54,7 @@ describe("API - SendMessage Command", () => { invoke: "sendMessage", text: messageText, images: undefined, + segments: userSegments(messageText ?? ""), }) }) @@ -72,6 +74,7 @@ describe("API - SendMessage Command", () => { invoke: "sendMessage", text: messageText, images, + segments: userSegments(messageText ?? ""), }) }) @@ -90,6 +93,7 @@ describe("API - SendMessage Command", () => { invoke: "sendMessage", text: undefined, images, + segments: userSegments(""), }) }) @@ -103,6 +107,7 @@ describe("API - SendMessage Command", () => { invoke: "sendMessage", text: undefined, images: undefined, + segments: userSegments(""), }) }) @@ -129,6 +134,7 @@ describe("API - SendMessage Command", () => { invoke: "sendMessage", text: messageText, images: undefined, + segments: userSegments(messageText ?? ""), }) }) @@ -150,7 +156,38 @@ describe("API - SendMessage Command", () => { invoke: "sendMessage", text: messageText, images, + segments: userSegments(messageText ?? ""), }) expect(mockPostMessageToWebview).toHaveBeenCalledTimes(1) }) + + describe("when the webview is not launched", () => { + const user = (text: string) => ({ text, source: "user" as const }) + + it.each([ + { label: "text and images", text: "see @/file.ts", images: ["img"], expected: [user("see @/file.ts")] }, + { label: "empty text", text: "", images: undefined, expected: [] }, + { label: "missing text", text: undefined, images: undefined, expected: [] }, + ])("submits the user text directly with $label", async ({ text, images, expected }) => { + const submitUserMessage = vi.fn().mockResolvedValue(undefined) + Object.assign(mockProvider, { + viewLaunched: false, + getCurrentTask: vi.fn().mockReturnValue({ submitUserMessage }), + }) + + await api.sendMessage(text, images) + + expect(submitUserMessage).toHaveBeenCalledWith(text ?? "", images, undefined, undefined, expected) + expect(mockPostMessageToWebview).not.toHaveBeenCalled() + }) + + it("logs and drops the message when there is no current task", async () => { + Object.assign(mockProvider, { viewLaunched: false }) + + await api.sendMessage("see @/file.ts") + + expect(mockLog).toHaveBeenCalledWith(expect.stringContaining("no current task")) + expect(mockPostMessageToWebview).not.toHaveBeenCalled() + }) + }) }) diff --git a/src/extension/__tests__/api-start-new-task.spec.ts b/src/extension/__tests__/api-start-new-task.spec.ts index e33f2b3553..d1a473d1d0 100644 --- a/src/extension/__tests__/api-start-new-task.spec.ts +++ b/src/extension/__tests__/api-start-new-task.spec.ts @@ -86,7 +86,11 @@ describe("API - startNewTask routing", () => { "new task", images, undefined, - { consecutiveMistakeLimit: Number.MAX_SAFE_INTEGER, toolRepetitionSoftLimit: 0 }, + { + consecutiveMistakeLimit: Number.MAX_SAFE_INTEGER, + toolRepetitionSoftLimit: 0, + segments: [{ text: "new task", source: "user" }], + }, configuration, ) expect(taskId).toBe(newTab ? "editor-task" : "sidebar-task") diff --git a/src/extension/api.ts b/src/extension/api.ts index 606bb7725d..e34a7fdf97 100644 --- a/src/extension/api.ts +++ b/src/extension/api.ts @@ -24,9 +24,11 @@ import { IpcMessageType, TASK_START_FAILURE_ERROR_CODE, TASK_START_FAILURE_ERROR_MESSAGE, + type MentionSegment, } from "@roo-code/types" import { IpcServer } from "@roo-code/ipc" +import { userSegments } from "../shared/mentionSegments" import { Package } from "../shared/package" import type { Mode } from "../shared/modes" @@ -250,9 +252,11 @@ export class API extends EventEmitter implements RooCodeAPI { // effectively unlimited. The soft repetition tier is disabled too: // without a reachable hard stop it would reject every identical call // with no escalation. - const options: CreateTaskOptions = { + const options: CreateTaskOptions & { segments?: MentionSegment[] } = { consecutiveMistakeLimit: Number.MAX_SAFE_INTEGER, toolRepetitionSoftLimit: 0, + // The API caller authored the task text. + segments: userSegments(text ?? ""), } const task = await provider.createTask(text, images, undefined, options, configuration) @@ -409,11 +413,17 @@ export class API extends EventEmitter implements RooCodeAPI { return } - await currentTask.submitUserMessage(text ?? "", images) + await currentTask.submitUserMessage(text ?? "", images, undefined, undefined, userSegments(text ?? "")) return } - await this.sidebarProvider.postMessageToWebview({ type: "invoke", invoke: "sendMessage", text, images }) + await this.sidebarProvider.postMessageToWebview({ + type: "invoke", + invoke: "sendMessage", + text, + images, + segments: userSegments(text ?? ""), + }) } public deleteQueuedMessage(messageId: string) { diff --git a/src/shared/__tests__/mentionSegments.spec.ts b/src/shared/__tests__/mentionSegments.spec.ts new file mode 100644 index 0000000000..10ca7f3b82 --- /dev/null +++ b/src/shared/__tests__/mentionSegments.spec.ts @@ -0,0 +1,89 @@ +import type { MentionSegment } from "@roo-code/types" + +import { + applyEdit, + generatedSegments, + joinSegments, + normalizeSegments, + trimSegments, + userSegments, + userText, +} from "../mentionSegments" + +const user = (text: string): MentionSegment => ({ text, source: "user" }) +const gen = (text: string): MentionSegment => ({ text, source: "generated" }) + +describe("mentionSegments", () => { + describe("normalizeSegments", () => { + it("keeps valid segments that match the text", () => { + expect(normalizeSegments("ab", [gen("a"), user("b")])).toEqual([gen("a"), user("b")]) + }) + + it.each([ + { label: "omitted", segments: undefined }, + { label: "not a list", segments: "user" }, + { label: "bad source", segments: [{ text: "ab", source: "system" }] }, + { label: "text mismatch", segments: [user("abc")] }, + ])("fails closed when segments are $label", ({ segments }) => { + expect(normalizeSegments("ab", segments)).toEqual([gen("ab")]) + }) + + it("returns no segments for empty text", () => { + expect(normalizeSegments("", undefined)).toEqual([]) + }) + }) + + it("joins text and selects only user text", () => { + const segments = [gen("a @/x "), user("b @/y"), gen(" c"), user("d")] + expect(joinSegments(segments)).toBe("a @/x b @/y cd") + expect(userText(segments)).toBe("b @/y\nd") + }) + + it("trims both ends across segments", () => { + expect(trimSegments([gen(" "), user(" a "), gen(" b ")])).toEqual([user("a "), gen(" b")]) + }) + + describe("applyEdit", () => { + it("marks inserted text as user and keeps the rest", () => { + expect(applyEdit([gen("context @/s")], "context @/s", "context @/s mine")).toEqual([ + gen("context @/s"), + user(" mine"), + ]) + }) + + it("marks text inserted in the middle as user and keeps both sides generated", () => { + expect(applyEdit([gen("ab")], "ab", "aXb")).toEqual([gen("a"), user("X"), gen("b")]) + }) + + it("removes deleted text and keeps its neighbors", () => { + expect(applyEdit([user("ab"), gen("cd")], "abcd", "acd")).toEqual([user("a"), gen("cd")]) + }) + + it("marks replaced generated text as user", () => { + expect(applyEdit([gen("abc")], "abc", "aZc")).toEqual([gen("a"), user("Z"), gen("c")]) + }) + + it("keeps all user text user", () => { + expect(applyEdit(userSegments("ab"), "ab", "abc")).toEqual([user("abc")]) + }) + + it("fails closed when the segments do not match the old value", () => { + expect(applyEdit([gen("x")], "ab", "abc")).toEqual(generatedSegments("abc")) + }) + + it.each([ + ["aa", "aaa"], + ["aaa", "aa"], + ["abab", "ab"], + ["ab", "abab"], + ])("keeps the joined text equal to the new value for %s -> %s", (oldValue, newValue) => { + const result = applyEdit([gen(oldValue)], oldValue, newValue) + expect(joinSegments(result)).toBe(newValue) + }) + + it("does not authorize a generated mention after an unrelated edit", () => { + const result = applyEdit([gen("see @/secret")], "see @/secret", "see @/secret now") + expect(userText(result)).toBe(" now") + }) + }) +}) diff --git a/src/shared/__tests__/support-prompts.spec.ts b/src/shared/__tests__/support-prompts.spec.ts index ea6a193d5a..2a003eb662 100644 --- a/src/shared/__tests__/support-prompts.spec.ts +++ b/src/shared/__tests__/support-prompts.spec.ts @@ -265,3 +265,27 @@ describe("Code Action Prompts", () => { }) }) }) + +describe("createSegments", () => { + it("marks only the user input as a user segment and matches create()", () => { + const params = { filePath: "a.ts", selectedText: "code @/secret", userInput: "mine @/ok" } + const segments = supportPrompt.createSegments("EXPLAIN", params) + + expect(segments.map((s) => s.text).join("")).toBe(supportPrompt.create("EXPLAIN", params)) + expect(segments.filter((s) => s.source === "user")).toEqual([{ text: "mine @/ok", source: "user" }]) + expect(segments.filter((s) => s.source === "generated").some((s) => s.text.includes("@/secret"))).toBe(true) + }) + + it("has no user segment without user input", () => { + const params = { filePath: "a.ts", selectedText: "code @/secret" } + const segments = supportPrompt.createSegments("EXPLAIN", params) + + expect(segments.some((s) => s.source === "user")).toBe(false) + expect(segments.map((s) => s.text).join("")).toBe(supportPrompt.create("EXPLAIN", params)) + }) + + it("does not treat a parameter that holds user-like text as a user segment", () => { + const segments = supportPrompt.createSegments("EXPLAIN", { selectedText: "@/secret", userInput: "" }) + expect(segments.some((s) => s.source === "user")).toBe(false) + }) +}) diff --git a/src/shared/mentionSegments.ts b/src/shared/mentionSegments.ts new file mode 100644 index 0000000000..5c46fb9443 --- /dev/null +++ b/src/shared/mentionSegments.ts @@ -0,0 +1,93 @@ +import { mentionSegmentSchema, type MentionSegment } from "@roo-code/types" + +const segmentsSchema = mentionSegmentSchema.array() + +export const joinSegments = (segments: MentionSegment[]): string => segments.map((s) => s.text).join("") + +/** The text that may have mentions expanded, one user part per line. */ +export const userText = (segments: MentionSegment[]): string => + segments + .filter((s) => s.source === "user") + .map((s) => s.text) + .join("\n") + +export const userSegments = (text: string): MentionSegment[] => (text ? [{ text, source: "user" }] : []) + +export const generatedSegments = (text: string): MentionSegment[] => (text ? [{ text, source: "generated" }] : []) + +/** Drops empty parts and joins neighbors that have the same source. */ +const merge = (segments: MentionSegment[]): MentionSegment[] => { + const result: MentionSegment[] = [] + for (const segment of segments) { + if (!segment.text) continue + const last = result[result.length - 1] + if (last?.source === segment.source) { + last.text += segment.text + } else { + result.push({ ...segment }) + } + } + return result +} + +/** + * Checks segments that arrive from outside. Fails closed: anything that is + * not a valid list whose text matches `text` becomes all generated. + */ +export const normalizeSegments = (text: string, segments: unknown): MentionSegment[] => { + const parsed = segmentsSchema.safeParse(segments) + return parsed.success && joinSegments(parsed.data) === text ? merge(parsed.data) : generatedSegments(text) +} + +export const trimSegments = (segments: MentionSegment[]): MentionSegment[] => { + const result = merge(segments) + while (result.length > 0) { + const first = result[0]! + first.text = first.text.trimStart() + if (first.text) break + result.shift() + } + while (result.length > 0) { + const last = result[result.length - 1]! + last.text = last.text.trimEnd() + if (last.text) break + result.pop() + } + return result +} + +const slice = (segments: MentionSegment[], start: number, end: number): MentionSegment[] => { + const result: MentionSegment[] = [] + let offset = 0 + for (const segment of segments) { + const from = Math.max(start - offset, 0) + const to = Math.min(end - offset, segment.text.length) + if (from < to) result.push({ text: segment.text.slice(from, to), source: segment.source }) + offset += segment.text.length + } + return result +} + +/** + * Records a manual edit. The changed range is the part between the common + * prefix and the common suffix. The new text there is `user`. The rest keeps + * its source. If the segments do not match `oldValue`, the result is all + * generated. + */ +export const applyEdit = (segments: MentionSegment[], oldValue: string, newValue: string): MentionSegment[] => { + if (joinSegments(segments) !== oldValue) return generatedSegments(newValue) + + const max = Math.min(oldValue.length, newValue.length) + let prefix = 0 + while (prefix < max && oldValue[prefix] === newValue[prefix]) prefix++ + let suffix = 0 + while (suffix < max - prefix && oldValue[oldValue.length - 1 - suffix] === newValue[newValue.length - 1 - suffix]) { + suffix++ + } + + return merge([ + ...slice(segments, 0, prefix), + ...userSegments(newValue.slice(prefix, newValue.length - suffix)), + ...slice(segments, oldValue.length - suffix, oldValue.length), + ]) +} diff --git a/src/shared/support-prompt.ts b/src/shared/support-prompt.ts index 6fe48fa364..7acd9bb9a9 100644 --- a/src/shared/support-prompt.ts +++ b/src/shared/support-prompt.ts @@ -1,3 +1,5 @@ +import type { MentionSegment } from "@roo-code/types" + // Support prompts type PromptDiagnostic = { source?: string @@ -261,6 +263,26 @@ export const supportPrompt = { const template = supportPrompt.get(customSupportPrompts, type) return createPrompt(template, params) }, + /** + * Builds the prompt as segments. Only the `userInput` parameter is `user`. + * Every other part of the template and its parameters is `generated`. + */ + createSegments: ( + type: SupportPromptType, + params: PromptParams, + customSupportPrompts?: Record, + ): MentionSegment[] => { + const template = supportPrompt.get(customSupportPrompts, type) + const mark = `\u0000${globalThis.crypto.randomUUID()}\u0000` + const userInput = typeof params["userInput"] === "string" ? params["userInput"] : "" + const parts = createPrompt(template, { ...params, userInput: mark }).split(mark) + const segments: MentionSegment[] = [] + parts.forEach((part, index) => { + if (part) segments.push({ text: part, source: "generated" }) + if (index < parts.length - 1 && userInput) segments.push({ text: userInput, source: "user" }) + }) + return segments + }, } as const export type { SupportPromptType } diff --git a/webview-ui/src/components/chat/ChatTextArea.tsx b/webview-ui/src/components/chat/ChatTextArea.tsx index 8ba0df7ab2..cf28e4d4b9 100644 --- a/webview-ui/src/components/chat/ChatTextArea.tsx +++ b/webview-ui/src/components/chat/ChatTextArea.tsx @@ -3,7 +3,7 @@ import { useEvent } from "react-use" import DynamicTextArea from "react-textarea-autosize" import { VolumeX, Image, WandSparkles, SendHorizontal, X, ListEnd, Square } from "lucide-react" -import type { ExtensionMessage } from "@roo-code/types" +import type { ExtensionMessage, MentionSegment } from "@roo-code/types" import { mentionRegex, mentionRegexGlobal, commandRegexGlobal, unescapeSpaces } from "@roo/context-mentions" import { WebviewMessage } from "@roo/WebviewMessage" @@ -59,6 +59,13 @@ const isFileDrag = (dataTransfer: DataTransfer): boolean => interface ChatTextAreaProps { inputValue: string setInputValue: (value: string) => void + /** Programmatic updates preserve the chat box authorization scope. */ + /** Sets generated text, such as an enhanced prompt. Falls back to `setInputValue`. */ + setInputValueGenerated?: (value: string) => void + /** Reads the draft's segments. Null means all user text. */ + getInputSegments?: () => MentionSegment[] | null + /** Sets the value with its segments, for prompt history. */ + setInputValueWithSegments?: (value: string, segments: MentionSegment[] | null) => void sendingDisabled: boolean selectApiConfigDisabled: boolean placeholderText: string @@ -85,6 +92,9 @@ export const ChatTextArea = forwardRef( { inputValue, setInputValue, + setInputValueGenerated, + getInputSegments, + setInputValueWithSegments, selectApiConfigDisabled, placeholderText, selectedImages, @@ -155,6 +165,11 @@ export const ChatTextArea = forwardRef( return () => document.removeEventListener("mousedown", handleClickOutside) }, [showDropdown]) + // Generated text routes through this setter so it is never marked as user text. + const setInputValueProgrammatic = setInputValueGenerated ?? setInputValue + + const programmaticInsertRef = useRef(false) + // Handle enhanced prompt response and search results. useEffect(() => { const messageHandler = (event: MessageEvent) => { @@ -173,12 +188,14 @@ export const ChatTextArea = forwardRef( // Select all text first textarea.select() + programmaticInsertRef.current = true document.execCommand("insertText", false, message.text) } else { - setInputValue(message.text) + setInputValueProgrammatic(message.text) } } catch { - setInputValue(message.text) + programmaticInsertRef.current = false + setInputValueProgrammatic(message.text) } } @@ -202,6 +219,7 @@ export const ChatTextArea = forwardRef( message.text + " " + currentValue.slice(cursorPos) + // The user chose this text, so it is a normal edit. setInputValue(newValue) // Set cursor position after the inserted text @@ -233,7 +251,7 @@ export const ChatTextArea = forwardRef( window.addEventListener("message", messageHandler) return () => window.removeEventListener("message", messageHandler) - }, [setInputValue, searchRequestId, normalizedInputValue]) + }, [setInputValue, setInputValueProgrammatic, searchRequestId, normalizedInputValue]) const [isDraggingOver, setIsDraggingOver] = useState(false) const [textAreaBaseHeight, setTextAreaBaseHeight] = useState(undefined) @@ -258,6 +276,8 @@ export const ChatTextArea = forwardRef( cwd, inputValue: normalizedInputValue, setInputValue, + getSegments: getInputSegments, + setInputValueWithSegments, }) // Fetch git commits when Git is selected or when typing a hash. @@ -278,9 +298,9 @@ export const ChatTextArea = forwardRef( setIsEnhancingPrompt(true) vscode.postMessage({ type: "enhancePrompt" as const, text: trimmedInput }) } else { - setInputValue(t("chat:enhancePromptDescription")) + setInputValueProgrammatic(t("chat:enhancePromptDescription")) } - }, [normalizedInputValue, setInputValue, t]) + }, [normalizedInputValue, setInputValueProgrammatic, t]) const allModes = useMemo(() => getAllModes(customModes), [customModes]) @@ -613,7 +633,12 @@ export const ChatTextArea = forwardRef( const handleInputChange = useCallback( (e: React.ChangeEvent) => { const newValue = e.target.value - setInputValue(newValue) + if (programmaticInsertRef.current) { + programmaticInsertRef.current = false + setInputValueProgrammatic(newValue) + } else { + setInputValue(newValue) + } // Reset history navigation when user types resetOnInputChange() @@ -675,7 +700,14 @@ export const ChatTextArea = forwardRef( setFileSearchResults([]) // Clear file search results. } }, - [setInputValue, setSearchRequestId, setFileSearchResults, setSearchLoading, resetOnInputChange], + [ + setInputValue, + setInputValueProgrammatic, + setSearchRequestId, + setFileSearchResults, + setSearchLoading, + resetOnInputChange, + ], ) useEffect(() => { diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 820e30849d..321f6d0070 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -20,10 +20,26 @@ import { getCostBreakdownIfNeeded } from "@src/utils/costFormatting" import { batchNearby } from "@src/utils/batchNearby" import { isBoundary, isIgnorableBetweenTargets } from "@src/utils/chatBatchingPredicates" -import type { ClineAsk, ClineSayTool, ClineMessage, ExtensionMessage, AudioType, SuggestionItem } from "@roo-code/types" +import type { + ClineAsk, + ClineSayTool, + ClineMessage, + ExtensionMessage, + AudioType, + SuggestionItem, + MentionSegment, +} from "@roo-code/types" import { getCompletionCheckpoint, getSuggestionMode, hasUsableAnswer, isRetiredProvider } from "@roo-code/types" import { findLast } from "@roo/array" +import { + applyEdit, + generatedSegments, + joinSegments, + normalizeSegments, + trimSegments, + userSegments, +} from "@roo/mentionSegments" import { combineApiRequests } from "@roo/combineApiRequests" import { combineCommandSequences } from "@roo/combineCommandSequences" import { getApiMetrics } from "@roo/getApiMetrics" @@ -202,6 +218,13 @@ const ChatViewComponent: React.ForwardRefRenderFunction(null) const userRespondedRef = useRef(false) + // Origin of each part of the chat box value. Null means all of it is user text. + const segmentsRef = useRef(null) + // Trimmed chat box segments for the text. A value that does not match the text is all generated. + const segmentsFor = useCallback((text: string) => { + const trimmed = segmentsRef.current ? trimSegments(segmentsRef.current) : userSegments(text) + return joinSegments(trimmed) === text ? trimmed : generatedSegments(text) + }, []) const [currentFollowUpTs, setCurrentFollowUpTs] = useState(null) const [aggregatedCostsMap, setAggregatedCostsMap] = useState< Map< @@ -629,6 +652,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction { + (text: string, images: string[], segmentsOverride?: MentionSegment[]) => { text = text.trim() + const segments = segmentsOverride ? trimSegments(segmentsOverride) : segmentsFor(text) if (text || images.length > 0) { // Intercept when the active provider is retired — show a @@ -664,9 +689,15 @@ const ChatViewComponent: React.ForwardRefRenderFunction { - // Avoid nested template literals by breaking down the logic - let newValue = text - - if (inputValue !== "") { - newValue = inputValue + " " + text + (text: string, images: string[], inserted?: MentionSegment[]) => { + // Segments that are missing or do not match the text are all generated. + const insertedSegments = normalizeSegments(text, inserted) + if (inputValue === "") { + segmentsRef.current = insertedSegments + setInputValue(text) + } else { + segmentsRef.current = [ + ...(segmentsRef.current ?? userSegments(inputValue)), + ...generatedSegments(" "), + ...insertedSegments, + ] + setInputValue(inputValue + " " + text) } - - setInputValue(newValue) setSelectedImages([...selectedImages, ...images]) }, [inputValue, selectedImages], @@ -745,6 +795,35 @@ const ChatViewComponent: React.ForwardRefRenderFunction string)) => { + const previous = inputValueRef.current + const next = typeof value === "function" ? value(previous) : value + if (segmentsRef.current) { + segmentsRef.current = applyEdit(segmentsRef.current, previous, next) + } + inputValueRef.current = next + setInputValue(next) + }, + [setInputValue], + ) + + // Generated text replaces the whole value and authorizes nothing. + const setInputValueGenerated = useCallback((value: string) => { + segmentsRef.current = generatedSegments(value) + inputValueRef.current = value + setInputValue(value) + }, []) + + // Prompt history restores a draft together with its segments. + const getInputSegments = useCallback(() => segmentsRef.current, []) + const setInputValueWithSegments = useCallback((value: string, segments: MentionSegment[] | null) => { + segmentsRef.current = segments + inputValueRef.current = value + setInputValue(value) + }, []) + // Handle stop button click from textarea const handleStopTask = useCallback(() => { vscode.postMessage({ type: "cancelTask" }) @@ -759,11 +838,13 @@ const ChatViewComponent: React.ForwardRefRenderFunction