From 6913559d7582d861fe8def589922f235d7719b89 Mon Sep 17 00:00:00 2001 From: akhileshrangani4 Date: Fri, 20 Feb 2026 11:11:35 -0800 Subject: [PATCH 1/3] feat: show tool call UI in chat messages Add file tools (readFile, listFiles, searchFiles) to the AI backend with multi-step support via stopWhen, and render tool invocations as collapsible inline elements in assistant messages. - Add createFileTools() with readFile, listFiles, searchFiles tools - Enable multi-step tool calling with stopWhen: stepCountIs(5) - Preserve UIMessage parts through the Message type and chat provider - Create ToolInvocation component for inline tool call display - Render tool parts interleaved with text in MessageContent --- .../project/chat/components/message.tsx | 102 ++++++++++++-- .../chat/components/tool-invocation.tsx | 130 ++++++++++++++++++ web/components/project/chat/index.tsx | 4 +- web/components/project/chat/lib/types.ts | 3 + .../project/chat/providers/chat-provider.tsx | 14 +- web/lib/ai/tools.ts | 94 +++++++++++++ web/server/routes/ai.ts | 21 ++- 7 files changed, 347 insertions(+), 21 deletions(-) create mode 100644 web/components/project/chat/components/tool-invocation.tsx diff --git a/web/components/project/chat/components/message.tsx b/web/components/project/chat/components/message.tsx index bbfa2e51..f1374fc0 100644 --- a/web/components/project/chat/components/message.tsx +++ b/web/components/project/chat/components/message.tsx @@ -17,10 +17,12 @@ import { RefreshCcw, } from "lucide-react" import * as React from "react" +import type { UIMessage } from "ai" import { CodeApplyProvider } from "../contexts/code-apply-context" import { stringifyContent } from "../lib/utils" import { useChat } from "../providers/chat-provider" import { ContextTab } from "./context-tab" +import { ToolInvocation } from "./tool-invocation" export type MessageProps = { messageId?: string @@ -114,12 +116,14 @@ const MessageAvatar = ({ src, alt, className }: MessageAvatarProps) => { export type MessageContentProps = { children: React.ReactNode className?: string + parts?: UIMessage["parts"] } & React.ComponentProps & React.HTMLProps const MessageContent = ({ children, className, + parts, ...props }: MessageContentProps) => { const { role, context, onOpenFile } = useMessage() @@ -128,19 +132,91 @@ const MessageContent = ({ const stringifiedContent = React.useMemo(() => { return stringifyContent(children) }, [children]) + + const classNames = cn( + "text-sm rounded-lg p-2 break-words whitespace-normal w-full rounded-lg p-2", + isAssistant + ? "bg-background text-foreground" + : "bg-secondary text-secondary-foreground", + className, + ) + + // For assistant messages with parts, render each part individually + const isToolPart = (type: string) => + type === "dynamic-tool" || type.startsWith("tool-") + const hasToolParts = + isAssistant && + parts && + parts.some((p) => isToolPart(p.type)) + + const renderedContent = React.useMemo(() => { - const classNames = cn( - "text-sm rounded-lg p-2 break-words whitespace-normal w-full rounded-lg p-2", - isAssistant - ? "bg-background text-foreground" - : "bg-secondary text-secondary-foreground", - className, - ) - return isAssistant ? ( - - {children as string} - - ) : ( + if (hasToolParts && parts) { + // Separate tool parts from text parts so we can render them distinctly + const elements: React.ReactNode[] = [] + let toolGroup: React.ReactNode[] = [] + + const flushToolGroup = () => { + if (toolGroup.length > 0) { + elements.push( +
+ {toolGroup} +
, + ) + toolGroup = [] + } + } + + for (let i = 0; i < parts.length; i++) { + const part = parts[i] + if (part.type === "text") { + flushToolGroup() + if (!part.text) continue + elements.push( + + {part.text} + , + ) + } else if (isToolPart(part.type)) { + const toolPart = part as unknown as { + type: string + toolName: string + toolCallId: string + state: string + input?: Record + output?: unknown + errorText?: string + } + const toolName = + toolPart.toolName ?? part.type.replace(/^tool-/, "") + toolGroup.push( + , + ) + } + // skip step-start and other non-renderable parts + } + flushToolGroup() + + return
{elements}
+ } + + if (isAssistant) { + return ( + + {children as string} + + ) + } + + return (
) - }, [role, className, children, onOpenFile, props]) + }, [hasToolParts, parts, classNames, isAssistant, onOpenFile, props, children]) return (
+ output?: unknown + errorText?: string +} + +function getToolLabel(toolName: string, input?: Record) { + switch (toolName) { + case "readFile": + return input?.filePath ? String(input.filePath) : "file" + case "listFiles": + return "project files" + case "searchFiles": + return String(input?.pattern ?? input?.query ?? "files") + case "webSearch": + return String(input?.query ?? "the web") + default: + return toolName + } +} + +function getToolVerb(toolName: string, isDone: boolean) { + switch (toolName) { + case "readFile": + return isDone ? "Read" : "Reading" + case "listFiles": + return isDone ? "Listed" : "Listing" + case "searchFiles": + return isDone ? "Searched" : "Searching" + case "webSearch": + return isDone ? "Searched" : "Searching" + default: + return isDone ? "Ran" : "Running" + } +} + +function getToolIcon(toolName: string) { + switch (toolName) { + case "readFile": + return FileText + case "listFiles": + return FolderTree + case "searchFiles": + return Search + case "webSearch": + return Globe + default: + return FileText + } +} + +export function ToolInvocation({ part }: { part: ToolPart }) { + const [expanded, setExpanded] = useState(false) + const Icon = getToolIcon(part.toolName) + const label = getToolLabel(part.toolName, part.input) + const isLoading = + part.state === "input-available" || part.state === "input-streaming" + const isError = part.state === "output-error" + const isDone = part.state === "output-available" + const verb = getToolVerb(part.toolName, isDone || isError) + + return ( +
+ + {expanded && ( +
+ {part.input && ( + 
+              {JSON.stringify(part.input, null, 2)}
+
+ )} + {isDone && part.output != null && ( + 
+              {typeof part.output === "string"
+                ? part.output
+                : JSON.stringify(part.output, null, 2)}
+
+ )} + {isError && part.errorText && ( +

{part.errorText}

+ )} +
+ )} +
+ ) +} diff --git a/web/components/project/chat/index.tsx b/web/components/project/chat/index.tsx index b886fe27..f560de2b 100644 --- a/web/components/project/chat/index.tsx +++ b/web/components/project/chat/index.tsx @@ -203,7 +203,9 @@ function MainChatContent({ onRejectCode={onRejectCode} onOpenFile={onOpenFile} > - {message.content} + + {message.content} + ) })} diff --git a/web/components/project/chat/lib/types.ts b/web/components/project/chat/lib/types.ts index 0cf99c98..d47d271f 100644 --- a/web/components/project/chat/lib/types.ts +++ b/web/components/project/chat/lib/types.ts @@ -1,8 +1,11 @@ +import type { UIMessage } from "ai" + export interface Message { id?: string role: "user" | "assistant" content: string context?: ContextTab[] + parts?: UIMessage["parts"] } export type ContextTab = diff --git a/web/components/project/chat/providers/chat-provider.tsx b/web/components/project/chat/providers/chat-provider.tsx index 65dc58da..fe4457a2 100644 --- a/web/components/project/chat/providers/chat-provider.tsx +++ b/web/components/project/chat/providers/chat-provider.tsx @@ -100,6 +100,7 @@ function fromUIMessage( role: aiMsg.role as "user" | "assistant", content: getTextContent(aiMsg), context: contextMap.get(aiMsg.id), + parts: aiMsg.parts, } } @@ -145,6 +146,7 @@ function ChatProvider({ children }: ChatProviderProps) { // so prepareSendMessagesRequest always reads fresh values const requestBodyRef = useRef({ contextContent: "", + projectId, projectType, activeFileContent, fileTree, @@ -153,6 +155,7 @@ function ChatProvider({ children }: ChatProviderProps) { }) requestBodyRef.current = { contextContent: requestBodyRef.current.contextContent, + projectId, projectType, activeFileContent, fileTree, @@ -178,11 +181,14 @@ function ChatProvider({ children }: ChatProviderProps) { // message from `msgs` before calling this. return { body: { - messages: msgs.map((m) => ({ - role: m.role, - content: getTextContent(m), - })), + messages: msgs + .map((m) => ({ + role: m.role, + content: getTextContent(m), + })) + .filter((m) => m.content.length > 0), context: { + projectId: ref.projectId, templateType: ref.projectType, activeFileContent: ref.activeFileContent, fileTree: ref.fileTree, diff --git a/web/lib/ai/tools.ts b/web/lib/ai/tools.ts index 351c1628..46a143eb 100644 --- a/web/lib/ai/tools.ts +++ b/web/lib/ai/tools.ts @@ -1,4 +1,5 @@ import { tool, jsonSchema } from "ai" +import type { Project } from "@gitwit/lib/services/Project" interface SerperResponse { organic?: Array<{ @@ -70,6 +71,99 @@ export const webSearchTool = tool({ }, }) +/** Create file inspection tools bound to a project's container */ +export function createFileTools(project: Project) { + return { + readFile: tool({ + description: + "Read the contents of a file in the project. Use the file path relative to the project root (e.g. 'src/App.tsx', 'package.json').", + inputSchema: jsonSchema<{ filePath: string }>({ + type: "object", + properties: { + filePath: { + type: "string", + description: "Relative path to the file from the project root", + }, + }, + required: ["filePath"], + }), + execute: async ({ filePath }) => { + try { + const content = await project.fileManager!.getFile(filePath) + if (content === undefined) { + return { error: `File not found: ${filePath}` } + } + return { filePath, content } + } catch (error) { + return { + error: `Failed to read file: ${error instanceof Error ? error.message : "Unknown error"}`, + } + } + }, + }), + + listFiles: tool({ + description: + "List all files and folders in the project. Returns the full file tree excluding node_modules.", + inputSchema: jsonSchema>({ + type: "object", + properties: {}, + }), + execute: async () => { + try { + const tree = await project.fileManager!.getFileTree() + return { tree } + } catch (error) { + return { + error: `Failed to list files: ${error instanceof Error ? error.message : "Unknown error"}`, + } + } + }, + }), + + searchFiles: tool({ + description: + "Search for a text pattern across project files using grep. Returns matching lines with file paths and line numbers.", + inputSchema: jsonSchema<{ + pattern: string + fileGlob?: string + }>({ + type: "object", + properties: { + pattern: { + type: "string", + description: "The text or regex pattern to search for", + }, + fileGlob: { + type: "string", + description: + "Optional glob to filter files (e.g. '*.tsx', '*.css'). Defaults to all files.", + }, + }, + required: ["pattern"], + }), + execute: async ({ pattern, fileGlob }) => { + try { + const includeFlag = fileGlob ? `--include='${fileGlob}'` : "" + const result = await project.container!.commands.run( + `cd /home/user/project && grep -rn ${includeFlag} --exclude-dir=node_modules --exclude-dir=.git -- ${JSON.stringify(pattern)} . | head -50`, + { timeoutMs: 10000 }, + ) + const matches = result.stdout.trim() + return { + pattern, + matches: matches || "No matches found.", + } + } catch (error) { + return { + error: `Search failed: ${error instanceof Error ? error.message : "Unknown error"}`, + } + } + }, + }), + } +} + export const defaultTools = { webSearch: webSearchTool, } diff --git a/web/server/routes/ai.ts b/web/server/routes/ai.ts index 0cefe6f7..4a4e5e35 100644 --- a/web/server/routes/ai.ts +++ b/web/server/routes/ai.ts @@ -1,10 +1,11 @@ import { createRouter } from "@/lib/api/create-app" import { getUserProviderConfig } from "@/lib/ai/helpers" -import { defaultTools } from "@/lib/ai/tools" +import { defaultTools, createFileTools } from "@/lib/ai/tools" import { createModel, buildPrompt, mergeAiderDiff } from "@gitwit/ai" import type { FileTree } from "@gitwit/ai" +import { Project } from "@gitwit/lib/services/Project" import { templateConfigs } from "@gitwit/templates" -import { generateText, streamText } from "ai" +import { generateText, streamText, stepCountIs } from "ai" import { zValidator } from "@hono/zod-validator" import z from "zod" @@ -53,11 +54,25 @@ export const aiRouter = createRouter() contextContent: context?.contextContent, }) + // Initialize project for file tools when projectId is available + let project: Project | null = null + let fileTools = {} + if (context?.projectId) { + try { + project = new Project(context.projectId) + await project.initialize() + fileTools = createFileTools(project) + } catch (error) { + console.error("Failed to initialize project for file tools:", error) + } + } + const result = streamText({ model, system, messages, - tools: defaultTools, + tools: { ...defaultTools, ...fileTools }, + stopWhen: stepCountIs(5), }) return result.toUIMessageStreamResponse() From 4d9e2895bcc99f4d73e4ed295a7281dc3851a580 Mon Sep 17 00:00:00 2001 From: akhileshrangani4 Date: Fri, 20 Feb 2026 11:11:49 -0800 Subject: [PATCH 2/3] feat: render aider diff markers as visual diffs in code blocks Parse <<<<<<< SEARCH / ======= / >>>>>>> REPLACE markers into color-coded added/removed lines with gutter indicators instead of showing raw diff syntax. --- web/app/globals.css | 2 +- web/components/ui/code-block/body.tsx | 83 +++++++++++++++++------- web/components/ui/code-block/context.tsx | 2 + web/components/ui/code-block/index.tsx | 42 ++++++++++-- 4 files changed, 98 insertions(+), 31 deletions(-) diff --git a/web/app/globals.css b/web/app/globals.css index edad99a3..ad7e2c97 100644 --- a/web/app/globals.css +++ b/web/app/globals.css @@ -351,7 +351,7 @@ /* Shiki Themes */ html.dark .shiki, -html.dark .shiki span { +html.dark .shiki span:not([data-diff-line]) { color: var(--shiki-dark) !important; background-color: var(--shiki-dark-bg) !important; /* Optional, if you also want font styles */ diff --git a/web/components/ui/code-block/body.tsx b/web/components/ui/code-block/body.tsx index b2f500e9..2be3466a 100644 --- a/web/components/ui/code-block/body.tsx +++ b/web/components/ui/code-block/body.tsx @@ -1,12 +1,18 @@ import { cn } from "@/lib/utils" import type { HighlightResult } from "@streamdown/code" import { type ComponentProps, memo, useMemo } from "react" +import type { DiffLineType } from "./context" type CodeBlockBodyProps = ComponentProps<"pre"> & { result: HighlightResult language: string + diffLineTypes?: DiffLineType[] } + +// For diff blocks: no CSS counter line numbers, just block display +const DIFF_LINE_CLASSES_BLOCK = "block" + // Memoize line numbers class string since it's constant const LINE_NUMBER_CLASSES = cn( "block", @@ -22,8 +28,21 @@ const LINE_NUMBER_CLASSES = cn( "before:select-none", ) +const DIFF_LINE_STYLES: Record = { + context: {}, + added: { backgroundColor: "rgba(34, 197, 94, 0.15)" }, + removed: { backgroundColor: "rgba(239, 68, 68, 0.15)", textDecoration: "line-through", textDecorationColor: "rgba(239, 68, 68, 0.4)" }, +} + +const DIFF_GUTTER_STYLES: Record = { + context: { color: "inherit" }, + added: { color: "rgba(34, 197, 94, 0.8)" }, + removed: { color: "rgba(239, 68, 68, 0.8)" }, +} + export const CodeBlockBody = memo( - ({ children, result, language, className, ...rest }: CodeBlockBodyProps) => { + ({ children, result, language, diffLineTypes, className, ...rest }: CodeBlockBodyProps) => { + const isDiff = !!diffLineTypes return ( 
         
-          {result.tokens.map((row, index) => (
-            
-              {row.map((token, tokenIndex) => (
-                
-                  {token.content}
-                
-              ))}
-            
-          ))}
+          {result.tokens.map((row, index) => {
+            const lineType = diffLineTypes?.[index] ?? "context"
+            return (
+              
+                {isDiff && (
+                  
+                    {lineType === "added" ? "+" : lineType === "removed" ? "-" : " "}
+                  
+                )}
+                {row.map((token, tokenIndex) => (
+                  
+                    {token.content}
+                  
+                ))}
+              
+            )
+          })}
) }, (prevProps, nextProps) => { - // Custom comparison: only re-render if result tokens actually changed return ( prevProps.result === nextProps.result && prevProps.language === nextProps.language && - prevProps.className === nextProps.className + prevProps.className === nextProps.className && + prevProps.diffLineTypes === nextProps.diffLineTypes ) }, ) diff --git a/web/components/ui/code-block/context.tsx b/web/components/ui/code-block/context.tsx index 15f73c0c..d1ce7149 100644 --- a/web/components/ui/code-block/context.tsx +++ b/web/components/ui/code-block/context.tsx @@ -1,5 +1,7 @@ import { createContext, useContext } from "react" +export type DiffLineType = "context" | "added" | "removed" + interface CodeBlockContextType { code: string } diff --git a/web/components/ui/code-block/index.tsx b/web/components/ui/code-block/index.tsx index c93cc53e..2fb63544 100644 --- a/web/components/ui/code-block/index.tsx +++ b/web/components/ui/code-block/index.tsx @@ -14,7 +14,7 @@ import { StreamdownContext } from "streamdown" import { CodePluginContext } from "../markdown" import { CodeBlockBody } from "./body" import { CodeBlockContainer } from "./container" -import { CodeBlockContext } from "./context" +import { CodeBlockContext, type DiffLineType } from "./context" import { CodeBlockHeader } from "./header" type CodeBlockProps = HTMLAttributes & { @@ -26,6 +26,32 @@ type CodeBlockProps = HTMLAttributes & { onOpenFile?: (filePath: string) => void } +/** + * Parse aider diff markers into displayable code with per-line diff annotations. + * SEARCH lines become "removed", REPLACE lines become "added", rest is "context". + * Returns null if no markers found (not a diff block). + */ +function parseAiderDiff(code: string): { displayCode: string; lineTypes: DiffLineType[] } | null { + if (!code.includes("<<<<<<< SEARCH")) return null + + const lines = code.split("\n") + const output: string[] = [] + const lineTypes: DiffLineType[] = [] + let section: "outside" | "search" | "replace" = "outside" + + for (const line of lines) { + const trimmed = line.trimEnd() + if (trimmed === "<<<<<<< SEARCH") { section = "search"; continue } + if (trimmed === "=======" && section === "search") { section = "replace"; continue } + if (trimmed === ">>>>>>> REPLACE") { section = "outside"; continue } + + output.push(line) + lineTypes.push(section === "search" ? "removed" : section === "replace" ? "added" : "context") + } + + return { displayCode: output.join("\n"), lineTypes } +} + export const CodeBlock = ({ code, language, @@ -40,12 +66,17 @@ export const CodeBlock = ({ const { shikiTheme } = useContext(StreamdownContext) const { codePlugin } = use(CodePluginContext)! + // Parse diff markers for display — raw code stays in context for apply/copy + const diffParsed = useMemo(() => parseAiderDiff(code), [code]) + const displayCode = diffParsed?.displayCode ?? code + const diffLineTypes = diffParsed?.lineTypes + // Memoize the raw fallback tokens to avoid recomputing on every render const raw: HighlightResult = useMemo( () => ({ bg: "transparent", fg: "inherit", - tokens: code.split("\n").map((line) => [ + tokens: displayCode.split("\n").map((line) => [ { content: line, color: "inherit", @@ -55,7 +86,7 @@ export const CodeBlock = ({ }, ]), }), - [code], + [displayCode], ) // Use raw as initial state @@ -71,7 +102,7 @@ export const CodeBlock = ({ const cachedResult = codePlugin.highlight( { - code, + code: displayCode, language: language as BundledLanguage, themes: shikiTheme, }, @@ -89,7 +120,7 @@ export const CodeBlock = ({ // Not cached - reset to raw tokens while waiting for highlighting // This is critical for streaming: ensures we show current code, not stale tokens setResult(raw) - }, [code, language, shikiTheme, codePlugin, raw]) + }, [displayCode, language, shikiTheme, codePlugin, raw]) return ( @@ -107,6 +138,7 @@ export const CodeBlock = ({ className={className} language={language} result={result} + diffLineTypes={diffLineTypes} {...rest} /> From 6c38c25b4bfb06c173c681765cc38b0c3a4be763 Mon Sep 17 00:00:00 2001 From: akhileshrangani4 Date: Fri, 20 Feb 2026 11:39:54 -0800 Subject: [PATCH 3/3] feat: collapsible code blocks in chat, fix paste-as-context heuristic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Code blocks in chat collapse to ~4 lines with gradient fade and expand toggle, with a shimmer bar while streaming - Fix shouldTreatAsContext being too aggressive — now requires 2+ signals instead of any single match like a bare parenthesis - Modularize code-block into diff.ts, use-syntax-highlighting.ts, and collapsible-code.tsx --- web/app/globals.css | 9 ++ .../project/chat/components/message.tsx | 3 +- web/components/project/chat/lib/utils.ts | 21 ++- web/components/ui/code-block/body.tsx | 34 ++--- .../ui/code-block/collapsible-code.tsx | 81 +++++++++++ web/components/ui/code-block/diff.ts | 109 +++++++++++++++ web/components/ui/code-block/index.tsx | 127 ++++-------------- .../ui/code-block/use-syntax-highlighting.ts | 69 ++++++++++ web/components/ui/markdown.tsx | 10 +- 9 files changed, 332 insertions(+), 131 deletions(-) create mode 100644 web/components/ui/code-block/collapsible-code.tsx create mode 100644 web/components/ui/code-block/diff.ts create mode 100644 web/components/ui/code-block/use-syntax-highlighting.ts diff --git a/web/app/globals.css b/web/app/globals.css index ad7e2c97..0966ea2e 100644 --- a/web/app/globals.css +++ b/web/app/globals.css @@ -433,3 +433,12 @@ html.dark .shiki span:not([data-diff-line]) { opacity: 0 !important; transition: opacity 0.15s ease-out; } + +/* Shimmer animation for streaming code blocks */ +@keyframes shimmer { + 0% { transform: translateX(-100%); } + 100% { transform: translateX(100%); } +} +.animate-shimmer { + animation: shimmer 1.5s ease-in-out infinite; +} diff --git a/web/components/project/chat/components/message.tsx b/web/components/project/chat/components/message.tsx index f1374fc0..0d2ff1ea 100644 --- a/web/components/project/chat/components/message.tsx +++ b/web/components/project/chat/components/message.tsx @@ -177,6 +177,7 @@ const MessageContent = ({ key={i} className={classNames} onOpenFile={onOpenFile} + collapsibleCodeBlocks {...props} > {part.text} @@ -210,7 +211,7 @@ const MessageContent = ({ if (isAssistant) { return ( - + {children as string} ) diff --git a/web/components/project/chat/lib/utils.ts b/web/components/project/chat/lib/utils.ts index 947122f8..7bb8310e 100644 --- a/web/components/project/chat/lib/utils.ts +++ b/web/components/project/chat/lib/utils.ts @@ -217,13 +217,22 @@ const pathMatchesTab = ( } function shouldTreatAsContext(text: string) { - const long = text.length > 400 - const paragraphs = text.split("\n").length > 2 - const markdown = /[#>*`]|-{2,}|```/.test(text) - const code = /(function|\{|}|\(|\)|=>|class )/.test(text) - const lists = /^[0-9]+\./m.test(text) || /^[-*•]\s/m.test(text) + // Very long text is always context + if (text.length > 1500) return true - return long || paragraphs || markdown || code || lists + // Fenced code blocks are a strong signal + if (/```[\s\S]*```/.test(text)) return true + + // Count weak signals — require multiple to trigger + let signals = 0 + if (text.length > 500) signals++ + if (text.split("\n").length > 8) signals++ + if (/^#{1,3}\s/m.test(text)) signals++ // markdown headings (not bare #) + if (/^[-*•]\s.+\n[-*•]\s/m.test(text)) signals++ // actual list with 2+ items + if (/\b(function|const|let|var|class|import|export)\b.*[{(]/.test(text)) signals++ // code patterns + if (/=>\s*[{(]/.test(text)) signals++ // arrow functions + + return signals >= 2 } export { diff --git a/web/components/ui/code-block/body.tsx b/web/components/ui/code-block/body.tsx index 2be3466a..1ba007aa 100644 --- a/web/components/ui/code-block/body.tsx +++ b/web/components/ui/code-block/body.tsx @@ -10,9 +10,6 @@ type CodeBlockBodyProps = ComponentProps<"pre"> & { } -// For diff blocks: no CSS counter line numbers, just block display -const DIFF_LINE_CLASSES_BLOCK = "block" - // Memoize line numbers class string since it's constant const LINE_NUMBER_CLASSES = cn( "block", @@ -28,16 +25,20 @@ const LINE_NUMBER_CLASSES = cn( "before:select-none", ) +// Match the editor's diff decoration styles const DIFF_LINE_STYLES: Record = { context: {}, - added: { backgroundColor: "rgba(34, 197, 94, 0.15)" }, - removed: { backgroundColor: "rgba(239, 68, 68, 0.15)", textDecoration: "line-through", textDecorationColor: "rgba(239, 68, 68, 0.4)" }, -} - -const DIFF_GUTTER_STYLES: Record = { - context: { color: "inherit" }, - added: { color: "rgba(34, 197, 94, 0.8)" }, - removed: { color: "rgba(239, 68, 68, 0.8)" }, + added: { + backgroundColor: "rgba(0, 255, 0, 0.1)", + borderLeft: "3px solid #28a745", + paddingLeft: "8px", + }, + removed: { + backgroundColor: "rgba(255, 0, 0, 0.1)", + borderLeft: "3px solid #dc3545", + paddingLeft: "8px", + opacity: 0.7, + }, } export const CodeBlockBody = memo( @@ -59,21 +60,12 @@ export const CodeBlockBody = memo( const lineType = diffLineTypes?.[index] ?? "context" return ( - {isDiff && ( - - {lineType === "added" ? "+" : lineType === "removed" ? "-" : " "} - - )} {row.map((token, tokenIndex) => ( MIN_LINES_TO_COLLAPSE + const isCollapsed = shouldCollapse && collapsed + + // Detect streaming: code is still changing + const prevCodeRef = useRef(code) + const [isStreaming, setIsStreaming] = useState(false) + useEffect(() => { + if (enabled && code !== prevCodeRef.current) { + setIsStreaming(true) + prevCodeRef.current = code + } + const timer = setTimeout(() => setIsStreaming(false), STREAMING_DEBOUNCE_MS) + return () => clearTimeout(timer) + }, [code, enabled]) + + return ( +
+ {/* Code body — clipped when collapsed */} +
+ {children} +
+ + {/* Gradient fade overlay */} + {isCollapsed && ( +
+ )} + + {/* Streaming shimmer bar */} + {isCollapsed && isStreaming && ( +
+
+
+ )} + + {/* Expand / Collapse toggle */} + {shouldCollapse && ( + + )} +
+ ) +} diff --git a/web/components/ui/code-block/diff.ts b/web/components/ui/code-block/diff.ts new file mode 100644 index 00000000..a3521f11 --- /dev/null +++ b/web/components/ui/code-block/diff.ts @@ -0,0 +1,109 @@ +import type { DiffLineType } from "./context" + +// ── LCS Algorithm ────────────────────────────────────────────────────────── +// Standard Longest Common Subsequence DP table used to compute minimal diffs +// between the old (SEARCH) and new (REPLACE) lines within each hunk. + +function lcsTable(a: string[], b: string[]): number[][] { + const m = a.length + const n = b.length + const dp: number[][] = Array.from({ length: m + 1 }, () => Array(n + 1).fill(0)) + for (let i = 1; i <= m; i++) { + for (let j = 1; j <= n; j++) { + dp[i][j] = a[i - 1] === b[j - 1] + ? dp[i - 1][j - 1] + 1 + : Math.max(dp[i - 1][j], dp[i][j - 1]) + } + } + return dp +} + +/** + * Backtrack the LCS table to produce a minimal diff. + * Matching lines → "context", only actual changes → "removed" / "added". + */ +function diffLines( + oldLines: string[], + newLines: string[], + output: string[], + lineTypes: DiffLineType[], +) { + const dp = lcsTable(oldLines, newLines) + const ops: Array<{ type: DiffLineType; line: string }> = [] + let i = oldLines.length + let j = newLines.length + + while (i > 0 || j > 0) { + if (i > 0 && j > 0 && oldLines[i - 1] === newLines[j - 1]) { + ops.push({ type: "context", line: oldLines[i - 1] }) + i-- + j-- + } else if (j > 0 && (i === 0 || dp[i][j - 1] >= dp[i - 1][j])) { + ops.push({ type: "added", line: newLines[j - 1] }) + j-- + } else { + ops.push({ type: "removed", line: oldLines[i - 1] }) + i-- + } + } + + ops.reverse() + for (const op of ops) { + output.push(op.line) + lineTypes.push(op.type) + } +} + +// ── Aider Diff Parser ────────────────────────────────────────────────────── +// Parses aider-style SEARCH/REPLACE markers into displayable code with +// per-line diff annotations. Uses LCS within each hunk so only lines that +// actually changed get marked — matching the editor's inline diff view. +// +// Format: +// <<<<<<< SEARCH +// old code lines... +// ======= +// new code lines... +// >>>>>>> REPLACE + +export function parseAiderDiff( + code: string, +): { displayCode: string; lineTypes: DiffLineType[] } | null { + if (!code.includes("<<<<<<< SEARCH")) return null + + const lines = code.split("\n") + const output: string[] = [] + const lineTypes: DiffLineType[] = [] + let section: "outside" | "search" | "replace" = "outside" + let searchLines: string[] = [] + let replaceLines: string[] = [] + + const flushHunk = () => { + diffLines(searchLines, replaceLines, output, lineTypes) + searchLines = [] + replaceLines = [] + } + + for (const line of lines) { + const trimmed = line.trimEnd() + if (trimmed === "<<<<<<< SEARCH") { section = "search"; continue } + if (trimmed === "=======" && section === "search") { section = "replace"; continue } + if (trimmed === ">>>>>>> REPLACE") { + flushHunk() + section = "outside" + continue + } + + if (section === "search") { + searchLines.push(line) + } else if (section === "replace") { + replaceLines.push(line) + } else { + output.push(line) + lineTypes.push("context") + } + } + flushHunk() + + return { displayCode: output.join("\n"), lineTypes } +} diff --git a/web/components/ui/code-block/index.tsx b/web/components/ui/code-block/index.tsx index 2fb63544..68e5dbe1 100644 --- a/web/components/ui/code-block/index.tsx +++ b/web/components/ui/code-block/index.tsx @@ -1,21 +1,12 @@ // Acknowledgment: This code is adapted from the Streamdown project(stremadown.ai). -import type { HighlightResult } from "@streamdown/code" -import {} from "@streamdown/code" -import { - type HTMLAttributes, - use, - useContext, - useEffect, - useMemo, - useState, -} from "react" -import type { BundledLanguage } from "shiki" -import { StreamdownContext } from "streamdown" -import { CodePluginContext } from "../markdown" +import { type HTMLAttributes, useMemo } from "react" import { CodeBlockBody } from "./body" +import { CollapsibleCode } from "./collapsible-code" import { CodeBlockContainer } from "./container" -import { CodeBlockContext, type DiffLineType } from "./context" +import { CodeBlockContext } from "./context" +import { parseAiderDiff } from "./diff" import { CodeBlockHeader } from "./header" +import { useSyntaxHighlighting } from "./use-syntax-highlighting" type CodeBlockProps = HTMLAttributes & { code: string @@ -24,32 +15,7 @@ type CodeBlockProps = HTMLAttributes & { filePath?: string | null isNewFile?: boolean onOpenFile?: (filePath: string) => void -} - -/** - * Parse aider diff markers into displayable code with per-line diff annotations. - * SEARCH lines become "removed", REPLACE lines become "added", rest is "context". - * Returns null if no markers found (not a diff block). - */ -function parseAiderDiff(code: string): { displayCode: string; lineTypes: DiffLineType[] } | null { - if (!code.includes("<<<<<<< SEARCH")) return null - - const lines = code.split("\n") - const output: string[] = [] - const lineTypes: DiffLineType[] = [] - let section: "outside" | "search" | "replace" = "outside" - - for (const line of lines) { - const trimmed = line.trimEnd() - if (trimmed === "<<<<<<< SEARCH") { section = "search"; continue } - if (trimmed === "=======" && section === "search") { section = "replace"; continue } - if (trimmed === ">>>>>>> REPLACE") { section = "outside"; continue } - - output.push(line) - lineTypes.push(section === "search" ? "removed" : section === "replace" ? "added" : "context") - } - - return { displayCode: output.join("\n"), lineTypes } + collapsible?: boolean } export const CodeBlock = ({ @@ -59,69 +25,23 @@ export const CodeBlock = ({ filePath, isNewFile, onOpenFile, + collapsible, className, children, ...rest }: CodeBlockProps) => { - const { shikiTheme } = useContext(StreamdownContext) - const { codePlugin } = use(CodePluginContext)! - - // Parse diff markers for display — raw code stays in context for apply/copy + // ── Diff parsing ───────────────────────────────────────────────────────── + // Parse aider SEARCH/REPLACE markers into per-line diff annotations. + // The raw `code` is kept in context for apply/copy actions. const diffParsed = useMemo(() => parseAiderDiff(code), [code]) const displayCode = diffParsed?.displayCode ?? code const diffLineTypes = diffParsed?.lineTypes - // Memoize the raw fallback tokens to avoid recomputing on every render - const raw: HighlightResult = useMemo( - () => ({ - bg: "transparent", - fg: "inherit", - tokens: displayCode.split("\n").map((line) => [ - { - content: line, - color: "inherit", - bgColor: "transparent", - htmlStyle: {}, - offset: 0, - }, - ]), - }), - [displayCode], - ) - - // Use raw as initial state - const [result, setResult] = useState(raw) - - // Try to get cached result or subscribe to highlighting - useEffect(() => { - // If no code plugin, just use raw tokens (plain text) - if (!codePlugin) { - setResult(raw) - return - } - - const cachedResult = codePlugin.highlight( - { - code: displayCode, - language: language as BundledLanguage, - themes: shikiTheme, - }, - (highlightedResult) => { - setResult(highlightedResult) - }, - ) - - if (cachedResult) { - // Already cached, use it immediately - setResult(cachedResult) - return - } - - // Not cached - reset to raw tokens while waiting for highlighting - // This is critical for streaming: ensures we show current code, not stale tokens - setResult(raw) - }, [displayCode, language, shikiTheme, codePlugin, raw]) + // ── Syntax highlighting ────────────────────────────────────────────────── + // Returns Shiki-highlighted tokens when ready, plain-text tokens while loading. + const result = useSyntaxHighlighting(displayCode, language) + // ── Render ─────────────────────────────────────────────────────────────── return ( @@ -134,13 +54,20 @@ export const CodeBlock = ({ > {children} - + > + + ) diff --git a/web/components/ui/code-block/use-syntax-highlighting.ts b/web/components/ui/code-block/use-syntax-highlighting.ts new file mode 100644 index 00000000..bef25b06 --- /dev/null +++ b/web/components/ui/code-block/use-syntax-highlighting.ts @@ -0,0 +1,69 @@ +import type { HighlightResult } from "@streamdown/code" +import { use, useContext, useEffect, useMemo, useState } from "react" +import type { BundledLanguage } from "shiki" +import { StreamdownContext } from "streamdown" +import { CodePluginContext } from "../markdown" + +/** + * Manages syntax highlighting for a code block via the Streamdown code plugin. + * + * Returns a HighlightResult with tokenized + colored spans. Falls back to + * plain-text tokens while Shiki loads the grammar asynchronously, ensuring + * streaming code is never blank. + */ +export function useSyntaxHighlighting( + displayCode: string, + language: string, +): HighlightResult { + const { shikiTheme } = useContext(StreamdownContext) + const { codePlugin } = use(CodePluginContext)! + + // Plain-text fallback tokens — used until Shiki finishes highlighting + const raw: HighlightResult = useMemo( + () => ({ + bg: "transparent", + fg: "inherit", + tokens: displayCode.split("\n").map((line) => [ + { + content: line, + color: "inherit", + bgColor: "transparent", + htmlStyle: {}, + offset: 0, + }, + ]), + }), + [displayCode], + ) + + const [result, setResult] = useState(raw) + + useEffect(() => { + if (!codePlugin) { + setResult(raw) + return + } + + // Ask the plugin for cached tokens, or subscribe for async delivery + const cachedResult = codePlugin.highlight( + { + code: displayCode, + language: language as BundledLanguage, + themes: shikiTheme, + }, + (highlightedResult) => { + setResult(highlightedResult) + }, + ) + + if (cachedResult) { + setResult(cachedResult) + return + } + + // Not cached — show raw tokens while waiting (prevents stale content during streaming) + setResult(raw) + }, [displayCode, language, shikiTheme, codePlugin, raw]) + + return result +} diff --git a/web/components/ui/markdown.tsx b/web/components/ui/markdown.tsx index 25838e10..4ce84717 100644 --- a/web/components/ui/markdown.tsx +++ b/web/components/ui/markdown.tsx @@ -31,6 +31,7 @@ const CodeBlock = lazy(() => filePath?: string | null isNewFile?: boolean onOpenFile?: (filePath: string) => void + collapsible?: boolean } > > @@ -38,6 +39,7 @@ const CodeBlock = lazy(() => // Types type MarkdownProps = ComponentProps & { onOpenFile?: (filePath: string) => void + collapsibleCodeBlocks?: boolean } interface ExtractedFileInfo { @@ -49,6 +51,7 @@ interface ExtractedFileInfo { interface MarkdownContextType { fileInfoMap: Map onOpenFile?: (filePath: string) => void + collapsibleCodeBlocks?: boolean } // Constants @@ -187,6 +190,7 @@ const CodeComponent = ({ filePath={fileInfo?.filePath ?? null} isNewFile={fileInfo?.isNewFile} onOpenFile={onOpenFile} + collapsible={markdownCtx?.collapsibleCodeBlocks} > {showCodeControls && ( <> @@ -211,7 +215,7 @@ const MemoCode = memo( MemoCode.displayName = "MarkdownCode" export const Markdown = memo( - ({ className, children, onOpenFile, ...props }: MarkdownProps) => { + ({ className, children, onOpenFile, collapsibleCodeBlocks, ...props }: MarkdownProps) => { const rawMarkdown = typeof children === "string" ? children : "" const { fileInfoMap, strippedMarkdown } = useMemo( @@ -220,7 +224,7 @@ export const Markdown = memo( ) return ( - + - prev.children === next.children && prev.onOpenFile === next.onOpenFile, + prev.children === next.children && prev.onOpenFile === next.onOpenFile && prev.collapsibleCodeBlocks === next.collapsibleCodeBlocks, ) Markdown.displayName = "Markdown"