Desktop: Add Subagents Mention Support (#6540)

3 files changed, 231 insertions, 75 deletions

diff --git a/packages/app/src/components/prompt-input.tsx b/packages/app/src/components/prompt-input.tsx
index ab2a017ca..855eb31e1 100644
--- a/packages/app/src/components/prompt-input.tsx
+++ b/packages/app/src/components/prompt-input.tsx
@@ -3,7 +3,15 @@ import { createEffect, on, Component, Show, For, onMount, onCleanup, Switch, Mat
 import { createStore, produce } from "solid-js/store"
 import { createFocusSignal } from "@solid-primitives/active-element"
 import { useLocal } from "@/context/local"
-import { ContentPart, DEFAULT_PROMPT, isPromptEqual, Prompt, usePrompt, ImageAttachmentPart } from "@/context/prompt"
+import {
+  ContentPart,
+  DEFAULT_PROMPT,
+  isPromptEqual,
+  Prompt,
+  usePrompt,
+  ImageAttachmentPart,
+  AgentPart,
+} from "@/context/prompt"
 import { useLayout } from "@/context/layout"
 import { useSDK } from "@/context/sdk"
 import { useNavigate, useParams } from "@solidjs/router"
@@ -128,7 +136,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
   const working = createMemo(() => status()?.type !== "idle")
 
   const [store, setStore] = createStore<{
-    popover: "file" | "slash" | null
+    popover: "at" | "slash" | null
     historyIndex: number
     savedPrompt: Prompt | null
     placeholder: number
@@ -171,6 +179,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     prompt.map((part) => {
       if (part.type === "text") return { ...part }
       if (part.type === "image") return { ...part }
+      if (part.type === "agent") return { ...part }
       return {
         ...part,
         selection: part.selection ? { ...part.selection } : undefined,
@@ -321,15 +330,43 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     if (!isFocused()) setStore("popover", null)
   })
 
-  const handleFileSelect = (path: string | undefined) => {
-    if (!path) return
-    addPart({ type: "file", path, content: "@" + path, start: 0, end: 0 })
+  type AtOption = { type: "agent"; name: string; display: string } | { type: "file"; path: string; display: string }
+
+  const agentList = createMemo(() =>
+    sync.data.agent
+      .filter((agent) => !agent.hidden && agent.mode !== "primary")
+      .map((agent): AtOption => ({ type: "agent", name: agent.name, display: agent.name })),
+  )
+
+  const handleAtSelect = (option: AtOption | undefined) => {
+    if (!option) return
+    if (option.type === "agent") {
+      addPart({ type: "agent", name: option.name, content: "@" + option.name, start: 0, end: 0 })
+    } else {
+      addPart({ type: "file", path: option.path, content: "@" + option.path, start: 0, end: 0 })
+    }
   }
 
-  const { flat, active, onInput, onKeyDown } = useFilteredList<string>({
-    items: local.file.searchFilesAndDirectories,
-    key: (x) => x,
-    onSelect: handleFileSelect,
+  const atKey = (x: AtOption | undefined) => {
+    if (!x) return ""
+    return x.type === "agent" ? `agent:${x.name}` : `file:${x.path}`
+  }
+
+  const {
+    flat: atFlat,
+    active: atActive,
+    onInput: atOnInput,
+    onKeyDown: atOnKeyDown,
+  } = useFilteredList<AtOption>({
+    items: async (query) => {
+      const agents = agentList()
+      const files = await local.file.searchFilesAndDirectories(query)
+      const fileOptions: AtOption[] = files.map((path) => ({ type: "file", path, display: path }))
+      return [...agents, ...fileOptions]
+    },
+    key: atKey,
+    filterKeys: ["display"],
+    onSelect: handleAtSelect,
   })
 
   const slashCommands = createMemo<SlashCommand[]>(() => {
@@ -415,6 +452,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
           if (node.nodeType !== Node.ELEMENT_NODE) return false
           const el = node as HTMLElement
           if (el.dataset.type === "file") return true
+          if (el.dataset.type === "agent") return true
           return el.tagName === "BR"
         })
         if (normalized && isPromptEqual(currentParts, domParts)) return
@@ -438,6 +476,15 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
             pill.style.userSelect = "text"
             pill.style.cursor = "default"
             editorRef.appendChild(pill)
+          } else if (part.type === "agent") {
+            const pill = document.createElement("span")
+            pill.textContent = part.content
+            pill.setAttribute("data-type", "agent")
+            pill.setAttribute("data-name", part.name)
+            pill.setAttribute("contenteditable", "false")
+            pill.style.userSelect = "text"
+            pill.style.cursor = "default"
+            editorRef.appendChild(pill)
           }
         })
 
@@ -473,6 +520,18 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       position += content.length
     }
 
+    const pushAgent = (agent: HTMLElement) => {
+      const content = agent.textContent ?? ""
+      parts.push({
+        type: "agent",
+        name: agent.dataset.name!,
+        content,
+        start: position,
+        end: position + content.length,
+      })
+      position += content.length
+    }
+
     const visit = (node: Node) => {
       if (node.nodeType === Node.TEXT_NODE) {
         buffer += node.textContent ?? ""
@@ -486,6 +545,11 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
         pushFile(el)
         return
       }
+      if (el.dataset.type === "agent") {
+        flushText()
+        pushAgent(el)
+        return
+      }
       if (el.tagName === "BR") {
         buffer += "\n"
         return
@@ -539,8 +603,8 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       const slashMatch = rawText.match(/^\/(\S*)$/)
 
       if (atMatch) {
-        onInput(atMatch[1])
-        setStore("popover", "file")
+        atOnInput(atMatch[1])
+        setStore("popover", "at")
       } else if (slashMatch) {
         slashOnInput(slashMatch[1])
         setStore("popover", "slash")
@@ -560,6 +624,36 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     queueScroll()
   }
 
+  const setRangeEdge = (range: Range, edge: "start" | "end", offset: number) => {
+    let remaining = offset
+    const nodes = Array.from(editorRef.childNodes)
+
+    for (const node of nodes) {
+      const length = getNodeLength(node)
+      const isText = node.nodeType === Node.TEXT_NODE
+      const isPill =
+        node.nodeType === Node.ELEMENT_NODE &&
+        ((node as HTMLElement).dataset.type === "file" || (node as HTMLElement).dataset.type === "agent")
+      const isBreak = node.nodeType === Node.ELEMENT_NODE && (node as HTMLElement).tagName === "BR"
+
+      if (isText && remaining <= length) {
+        if (edge === "start") range.setStart(node, remaining)
+        if (edge === "end") range.setEnd(node, remaining)
+        return
+      }
+
+      if ((isPill || isBreak) && remaining <= length) {
+        if (edge === "start" && remaining === 0) range.setStartBefore(node)
+        if (edge === "start" && remaining > 0) range.setStartAfter(node)
+        if (edge === "end" && remaining === 0) range.setEndBefore(node)
+        if (edge === "end" && remaining > 0) range.setEndAfter(node)
+        return
+      }
+
+      remaining -= length
+    }
+  }
+
   const addPart = (part: ContentPart) => {
     const selection = window.getSelection()
     if (!selection || selection.rangeCount === 0) return
@@ -582,38 +676,35 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       const gap = document.createTextNode(" ")
       const range = selection.getRangeAt(0)
 
-      const setEdge = (edge: "start" | "end", offset: number) => {
-        let remaining = offset
-        const nodes = Array.from(editorRef.childNodes)
-
-        for (const node of nodes) {
-          const length = getNodeLength(node)
-          const isText = node.nodeType === Node.TEXT_NODE
-          const isFile = node.nodeType === Node.ELEMENT_NODE && (node as HTMLElement).dataset.type === "file"
-          const isBreak = node.nodeType === Node.ELEMENT_NODE && (node as HTMLElement).tagName === "BR"
-
-          if (isText && remaining <= length) {
-            if (edge === "start") range.setStart(node, remaining)
-            if (edge === "end") range.setEnd(node, remaining)
-            return
-          }
+      if (atMatch) {
+        const start = atMatch.index ?? cursorPosition - atMatch[0].length
+        setRangeEdge(range, "start", start)
+        setRangeEdge(range, "end", cursorPosition)
+      }
 
-          if ((isFile || isBreak) && remaining <= length) {
-            if (edge === "start" && remaining === 0) range.setStartBefore(node)
-            if (edge === "start" && remaining > 0) range.setStartAfter(node)
-            if (edge === "end" && remaining === 0) range.setEndBefore(node)
-            if (edge === "end" && remaining > 0) range.setEndAfter(node)
-            return
-          }
+      range.deleteContents()
+      range.insertNode(gap)
+      range.insertNode(pill)
+      range.setStartAfter(gap)
+      range.collapse(true)
+      selection.removeAllRanges()
+      selection.addRange(range)
+    } else if (part.type === "agent") {
+      const pill = document.createElement("span")
+      pill.textContent = part.content
+      pill.setAttribute("data-type", "agent")
+      pill.setAttribute("data-name", part.name)
+      pill.setAttribute("contenteditable", "false")
+      pill.style.userSelect = "text"
+      pill.style.cursor = "default"
 
-          remaining -= length
-        }
-      }
+      const gap = document.createTextNode(" ")
+      const range = selection.getRangeAt(0)
 
       if (atMatch) {
         const start = atMatch.index ?? cursorPosition - atMatch[0].length
-        setEdge("start", start)
-        setEdge("end", cursorPosition)
+        setRangeEdge(range, "start", start)
+        setRangeEdge(range, "end", cursorPosition)
       }
 
       range.deleteContents()
@@ -834,8 +925,8 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     }
 
     if (store.popover && (event.key === "ArrowUp" || event.key === "ArrowDown" || event.key === "Enter")) {
-      if (store.popover === "file") {
-        onKeyDown(event)
+      if (store.popover === "at") {
+        atOnKeyDown(event)
       } else {
         slashOnKeyDown(event)
       }
@@ -1075,11 +1166,12 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     if (!existing) return
 
     const toAbsolutePath = (path: string) => (path.startsWith("/") ? path : sync.absolute(path))
-    const attachments = currentPrompt.filter(
+    const fileAttachments = currentPrompt.filter(
       (part) => part.type === "file",
     ) as import("@/context/prompt").FileAttachmentPart[]
+    const agentAttachments = currentPrompt.filter((part) => part.type === "agent") as AgentPart[]
 
-    const fileAttachmentParts = attachments.map((attachment) => {
+    const fileAttachmentParts = fileAttachments.map((attachment) => {
       const absolute = toAbsolutePath(attachment.path)
       const query = attachment.selection
         ? `?start=${attachment.selection.startLine}&end=${attachment.selection.endLine}`
@@ -1102,6 +1194,17 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       }
     })
 
+    const agentAttachmentParts = agentAttachments.map((attachment) => ({
+      id: Identifier.ascending("part"),
+      type: "agent" as const,
+      name: attachment.name,
+      source: {
+        value: attachment.content,
+        start: attachment.start,
+        end: attachment.end,
+      },
+    }))
+
     const imageAttachmentParts = store.imageAttachments.map((attachment) => ({
       id: Identifier.ascending("part"),
       type: "file" as const,
@@ -1171,7 +1274,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       type: "text" as const,
       text,
     }
-    const requestParts = [textPart, ...fileAttachmentParts, ...imageAttachmentParts]
+    const requestParts = [textPart, ...fileAttachmentParts, ...agentAttachmentParts, ...imageAttachmentParts]
     const optimisticParts = requestParts.map((part) => ({
       ...part,
       sessionID: existing.id,
@@ -1209,24 +1312,46 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
                  border border-border-base bg-surface-raised-stronger-non-alpha shadow-md"
         >
           <Switch>
-            <Match when={store.popover === "file"}>
-              <Show when={flat().length > 0} fallback={<div class="text-text-weak px-2 py-1">No matching files</div>}>
-                <For each={flat()}>
-                  {(i) => (
+            <Match when={store.popover === "at"}>
+              <Show
+                when={atFlat().length > 0}
+                fallback={<div class="text-text-weak px-2 py-1">No matching results</div>}
+              >
+                <For each={atFlat().slice(0, 10)}>
+                  {(item) => (
                     <button
                       classList={{
                         "w-full flex items-center gap-x-2 rounded-md px-2 py-0.5": true,
-                        "bg-surface-raised-base-hover": active() === i,
+                        "bg-surface-raised-base-hover": atActive() === atKey(item),
                       }}
-                      onClick={() => handleFileSelect(i)}
+                      onClick={() => handleAtSelect(item)}
                     >
-                      <FileIcon node={{ path: i, type: "file" }} class="shrink-0 size-4" />
-                      <div class="flex items-center text-14-regular min-w-0">
-                        <span class="text-text-weak whitespace-nowrap truncate min-w-0">{getDirectory(i)}</span>
-                        <Show when={!i.endsWith("/")}>
-                          <span class="text-text-strong whitespace-nowrap">{getFilename(i)}</span>
-                        </Show>
-                      </div>
+                      <Show
+                        when={item.type === "agent"}
+                        fallback={
+                          <>
+                            <FileIcon
+                              node={{ path: (item as { type: "file"; path: string }).path, type: "file" }}
+                              class="shrink-0 size-4"
+                            />
+                            <div class="flex items-center text-14-regular min-w-0">
+                              <span class="text-text-weak whitespace-nowrap truncate min-w-0">
+                                {getDirectory((item as { type: "file"; path: string }).path)}
+                              </span>
+                              <Show when={!(item as { type: "file"; path: string }).path.endsWith("/")}>
+                                <span class="text-text-strong whitespace-nowrap">
+                                  {getFilename((item as { type: "file"; path: string }).path)}
+                                </span>
+                              </Show>
+                            </div>
+                          </>
+                        }
+                      >
+                        <Icon name="brain" size="small" class="text-icon-info-active shrink-0" />
+                        <span class="text-14-regular text-text-strong whitespace-nowrap">
+                          @{(item as { type: "agent"; name: string }).name}
+                        </span>
+                      </Show>
                     </button>
                   )}
                 </For>
@@ -1335,7 +1460,8 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
             classList={{
               "select-text": true,
               "w-full px-5 py-3 pr-12 text-14-regular text-text-strong focus:outline-none whitespace-pre-wrap": true,
-              "[&_[data-type=file]]:text-icon-info-active": true,
+              "[&_[data-type=file]]:text-syntax-property": true,
+              "[&_[data-type=agent]]:text-syntax-type": true,
               "font-mono!": store.mode === "shell",
             }}
           />
@@ -1533,7 +1659,9 @@ function setCursorPosition(parent: HTMLElement, position: number) {
   while (node) {
     const length = getNodeLength(node)
     const isText = node.nodeType === Node.TEXT_NODE
-    const isFile = node.nodeType === Node.ELEMENT_NODE && (node as HTMLElement).dataset.type === "file"
+    const isPill =
+      node.nodeType === Node.ELEMENT_NODE &&
+      ((node as HTMLElement).dataset.type === "file" || (node as HTMLElement).dataset.type === "agent")
     const isBreak = node.nodeType === Node.ELEMENT_NODE && (node as HTMLElement).tagName === "BR"
 
     if (isText && remaining <= length) {
@@ -1546,13 +1674,13 @@ function setCursorPosition(parent: HTMLElement, position: number) {
       return
     }
 
-    if ((isFile || isBreak) && remaining <= length) {
+    if ((isPill || isBreak) && remaining <= length) {
       const range = document.createRange()
       const selection = window.getSelection()
       if (remaining === 0) {
         range.setStartBefore(node)
       }
-      if (remaining > 0 && isFile) {
+      if (remaining > 0 && isPill) {
         range.setStartAfter(node)
       }
       if (remaining > 0 && isBreak) {
diff --git a/packages/app/src/context/prompt.tsx b/packages/app/src/context/prompt.tsx
index 8d3590cd9..25d8146ea 100644
--- a/packages/app/src/context/prompt.tsx
+++ b/packages/app/src/context/prompt.tsx
@@ -21,6 +21,11 @@ export interface FileAttachmentPart extends PartBase {
   selection?: TextSelection
 }
 
+export interface AgentPart extends PartBase {
+  type: "agent"
+  name: string
+}
+
 export interface ImageAttachmentPart {
   type: "image"
   id: string
@@ -29,7 +34,7 @@ export interface ImageAttachmentPart {
   dataUrl: string
 }
 
-export type ContentPart = TextPart | FileAttachmentPart | ImageAttachmentPart
+export type ContentPart = TextPart | FileAttachmentPart | AgentPart | ImageAttachmentPart
 export type Prompt = ContentPart[]
 
 export const DEFAULT_PROMPT: Prompt = [{ type: "text", content: "", start: 0, end: 0 }]
@@ -46,6 +51,9 @@ export function isPromptEqual(promptA: Prompt, promptB: Prompt): boolean {
     if (partA.type === "file" && partA.path !== (partB as FileAttachmentPart).path) {
       return false
     }
+    if (partA.type === "agent" && partA.name !== (partB as AgentPart).name) {
+      return false
+    }
     if (partA.type === "image" && partA.id !== (partB as ImageAttachmentPart).id) {
       return false
     }
@@ -61,6 +69,7 @@ function cloneSelection(selection?: TextSelection) {
 function clonePart(part: ContentPart): ContentPart {
   if (part.type === "text") return { ...part }
   if (part.type === "image") return { ...part }
+  if (part.type === "agent") return { ...part }
   return {
     ...part,
     selection: cloneSelection(part.selection),
diff --git a/packages/ui/src/components/message-part.tsx b/packages/ui/src/components/message-part.tsx
index dc8c645de..ac80dada7 100644
--- a/packages/ui/src/components/message-part.tsx
+++ b/packages/ui/src/components/message-part.tsx
@@ -12,6 +12,7 @@ import {
 } from "solid-js"
 import { Dynamic } from "solid-js/web"
 import {
+  AgentPart,
   AssistantMessage,
   FilePart,
   Message as MessageType,
@@ -300,6 +301,8 @@ export function UserMessageDisplay(props: { message: UserMessage; parts: PartTyp
     }),
   )
 
+  const agents = createMemo(() => (props.parts?.filter((p) => p.type === "agent") as AgentPart[]) ?? [])
+
   const openImagePreview = (url: string, alt?: string) => {
     dialog.show(() => <ImagePreview src={url} alt={alt} />)
   }
@@ -337,33 +340,40 @@ export function UserMessageDisplay(props: { message: UserMessage; parts: PartTyp
       </Show>
       <Show when={text()}>
         <div data-slot="user-message-text">
-          <HighlightedText text={text()} references={inlineFiles()} />
+          <HighlightedText text={text()} references={inlineFiles()} agents={agents()} />
         </div>
       </Show>
     </div>
   )
 }
 
-function HighlightedText(props: { text: string; references: FilePart[] }) {
+type HighlightSegment = { text: string; type?: "file" | "agent" }
+
+function HighlightedText(props: { text: string; references: FilePart[]; agents: AgentPart[] }) {
   const segments = createMemo(() => {
     const text = props.text
-    const refs = [...props.references].sort((a, b) => (a.source?.text?.start ?? 0) - (b.source?.text?.start ?? 0))
 
-    const result: { text: string; highlight?: boolean }[] = []
-    let lastIndex = 0
+    const allRefs: { start: number; end: number; type: "file" | "agent" }[] = [
+      ...props.references
+        .filter((r) => r.source?.text?.start !== undefined && r.source?.text?.end !== undefined)
+        .map((r) => ({ start: r.source!.text!.start, end: r.source!.text!.end, type: "file" as const })),
+      ...props.agents
+        .filter((a) => a.source?.start !== undefined && a.source?.end !== undefined)
+        .map((a) => ({ start: a.source!.start, end: a.source!.end, type: "agent" as const })),
+    ].sort((a, b) => a.start - b.start)
 
-    for (const ref of refs) {
-      const start = ref.source?.text?.start
-      const end = ref.source?.text?.end
+    const result: HighlightSegment[] = []
+    let lastIndex = 0
 
-      if (start === undefined || end === undefined || start < lastIndex) continue
+    for (const ref of allRefs) {
+      if (ref.start < lastIndex) continue
 
-      if (start > lastIndex) {
-        result.push({ text: text.slice(lastIndex, start) })
+      if (ref.start > lastIndex) {
+        result.push({ text: text.slice(lastIndex, ref.start) })
       }
 
-      result.push({ text: text.slice(start, end), highlight: true })
-      lastIndex = end
+      result.push({ text: text.slice(ref.start, ref.end), type: ref.type })
+      lastIndex = ref.end
     }
 
     if (lastIndex < text.length) {
@@ -375,7 +385,16 @@ function HighlightedText(props: { text: string; references: FilePart[] }) {
 
   return (
     <For each={segments()}>
-      {(segment) => <span classList={{ "text-text-strong font-medium": segment.highlight }}>{segment.text}</span>}
+      {(segment) => (
+        <span
+          classList={{
+            "text-syntax-property": segment.type === "file",
+            "text-syntax-type": segment.type === "agent",
+          }}
+        >
+          {segment.text}
+        </span>
+      )}
     </For>
   )
 }